Reset file pointer when uploading file is retried (#86)

eth-cscs · Feb 21, 2024 · 781430a · 781430a
1 parent 659c320
commit 781430a
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git 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
@@ -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