Skip to content

Commit e047501

Browse files
CopilotArzelaAscoIi
andcommitted
Refactor: Extract rate limiting logic and add error handling
Co-authored-by: ArzelaAscoIi <37148029+ArzelaAscoIi@users.noreply.github.com>
1 parent 85422f9 commit e047501

1 file changed

Lines changed: 21 additions & 12 deletions

File tree

deepset_cloud_sdk/_s3/upload.py

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,11 @@
2929
# Check if pyrate-limiter supports the new API (v4.0.0+)
3030
# In v4.0.0+, Limiter.__init__ no longer accepts raise_when_fail and max_delay
3131
# and try_acquire accepts a blocking parameter
32-
_LIMITER_SUPPORTS_BLOCKING = "blocking" in inspect.signature(Limiter.try_acquire).parameters
32+
try:
33+
_LIMITER_SUPPORTS_BLOCKING = "blocking" in inspect.signature(Limiter.try_acquire).parameters
34+
except (AttributeError, ValueError):
35+
# Fallback to False if we can't inspect the signature
36+
_LIMITER_SUPPORTS_BLOCKING = False
3337

3438

3539
class RetryableHttpError(Exception):
@@ -133,6 +137,20 @@ async def __aexit__(
133137
except AttributeError:
134138
pass
135139

140+
def _rate_limit_acquire(self, name: str = "") -> None:
141+
"""
142+
Acquire a rate limit token.
143+
144+
Uses blocking=False for pyrate-limiter 4.0.0+ to maintain non-blocking behavior,
145+
or the default behavior for 3.x which was non-blocking with raise_when_fail=False.
146+
147+
:param name: The name of the item to acquire (default: "")
148+
"""
149+
if _LIMITER_SUPPORTS_BLOCKING:
150+
self.limiter.try_acquire(name, blocking=False)
151+
else:
152+
self.limiter.try_acquire(name)
153+
136154
async def _upload_file_with_retries(
137155
self,
138156
file_name: str,
@@ -172,11 +190,7 @@ async def _upload_file(
172190
file_data = self._build_file_data(content, aws_safe_name, aws_config)
173191

174192
try:
175-
# Rate limit requests - use blocking=False for pyrate-limiter 4.0.0+
176-
if _LIMITER_SUPPORTS_BLOCKING:
177-
self.limiter.try_acquire("", blocking=False)
178-
else:
179-
self.limiter.try_acquire("")
193+
self._rate_limit_acquire()
180194

181195
async with client_session.post(
182196
aws_config.url,
@@ -191,12 +205,7 @@ async def _upload_file(
191205
# for example during automatic redirects. See https://github.com/aio-libs/aiohttp/issues/5577
192206
redirect_url = response.headers["Location"]
193207
file_data = self._build_file_data(content, aws_safe_name, aws_config)
194-
195-
# Rate limit requests - use blocking=False for pyrate-limiter 4.0.0+
196-
if _LIMITER_SUPPORTS_BLOCKING:
197-
self.limiter.try_acquire("", blocking=False)
198-
else:
199-
self.limiter.try_acquire("")
208+
self._rate_limit_acquire()
200209

201210
async with client_session.post(
202211
redirect_url,

0 commit comments

Comments
 (0)