diff --git a/firecrest/AsyncClient.py b/firecrest/AsyncClient.py index c479d12..540d7cf 100644 --- a/firecrest/AsyncClient.py +++ b/firecrest/AsyncClient.py @@ -91,6 +91,12 @@ def _retry_requests(func): async def wrapper(*args, **kwargs): client = args[0] num_retries = 0 + try: + f = kwargs["files"]["file"] + file_original_position = f[1].tell() if isinstance(f, tuple) else f.tell() + except KeyError: + file_original_position = None + resp = await func(*args, **kwargs) while True: if resp.status_code != client.TOO_MANY_REQUESTS_CODE: @@ -110,6 +116,20 @@ async def wrapper(*args, **kwargs): default=resp.headers.get("RateLimit-Reset", default=10), ) reset = int(reset) + try: + f = kwargs["files"]["file"] + logger.debug( + f"Resetting the file pointer of the uploaded file " + f"to {file_original_position}" + ) + if isinstance(f, tuple): + f[1].seek(file_original_position) + else: + f.seek(file_original_position) + + except KeyError: + pass + microservice = kwargs["endpoint"].split("/")[1] client = args[0] logger.info( diff --git a/firecrest/BasicClient.py b/firecrest/BasicClient.py index 8fa7f5e..1b0ba18 100644 --- a/firecrest/BasicClient.py +++ b/firecrest/BasicClient.py @@ -64,6 +64,12 @@ def _retry_requests(func): def wrapper(*args, **kwargs): client = args[0] num_retries = 0 + try: + f = kwargs["files"]["file"] + file_original_position = f[1].tell() if isinstance(f, tuple) else f.tell() + except KeyError: + file_original_position = None + resp = func(*args, **kwargs) while True: if resp.status_code != client.TOO_MANY_REQUESTS_CODE: @@ -87,6 +93,20 @@ def wrapper(*args, **kwargs): f"{reset} seconds and try again" ) reset = int(reset) + try: + f = kwargs["files"]["file"] + logger.debug( + f"Resetting the file pointer of the uploaded file " + f"to {file_original_position}" + ) + if isinstance(f, tuple): + f[1].seek(file_original_position) + else: + f.seek(file_original_position) + + except KeyError: + pass + time.sleep(reset) resp = func(*args, **kwargs) num_retries += 1