Add retries to allow for rate limits on WIF

blobfile/_common.py

@@ -735,7 +735,7 @@ def _upload_buf(self, buf: memoryview, finalize: bool = False) -> int:
             self._upload_chunk(chunk, finalize)
             self._offset += len(chunk)
         finally:
-            del chunk, buf  # pyright: ignore[reportUnboundVariable]
+            del chunk, buf  # pyright: ignore[reportPossiblyUnboundVariable]
         return size
 
     def close(self) -> None:
@@ -769,7 +769,7 @@ def write(self, b: bytes) -> int:  # type: ignore
                     size = self._upload_buf(mv)
                     self._buf = bytearray(mv[size:])
                 finally:
-                    del mv  # pyright: ignore[reportUnboundVariable]
+                    del mv  # pyright: ignore[reportPossiblyUnboundVariable]
         assert len(self._buf) < self._chunk_size
         return len(b)
 

blobfile/_gcp.py

@@ -33,6 +33,8 @@
     strip_slashes,
 )
 
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed, wait_random
+
 MAX_EXPIRATION = 7 * 24 * 60 * 60
 OAUTH_TOKEN = "oauth_token"
 ANONYMOUS = "anonymous"
@@ -150,6 +152,13 @@ def _create_access_token_request(creds: Dict[str, Any], scopes: List[str]) -> Re
 # Derived from:
 # - https://github.com/googleapis/google-auth-library-python/blob/0afc61a3a9c3d5035c913ad4a7568e8a888e2ec9/google/auth/external_account.py#L361
 # - https://github.com/googleapis/google-auth-library-python/blob/0afc61a3a9c3d5035c913ad4a7568e8a888e2ec9/google/oauth2/sts.py#L95
+
+
+@retry(
+    retry=retry_if_exception_type(ValueError),
+    stop=stop_after_attempt(1),
+    wait=wait_fixed(3) + wait_random(0, 2),
+)
 def _create_sts_token_request(creds: Dict[str, Any]) -> Request:
     data = {
         "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",

env.yaml

@@ -19,3 +19,4 @@ dependencies:
     - boto3==1.15.18
     - lxml-stubs==0.4.0
     - xmltodict==0.13.0
+    - tenacity==8.2.2

setup.py

@@ -21,7 +21,13 @@
     author="Christopher Hesse",
     license="Public Domain",
     packages=find_packages(),
-    install_requires=["pycryptodomex~=3.8", "urllib3>=1.25.3,<3", "lxml~=4.9", "filelock~=3.0"],
+    install_requires=[
+        "pycryptodomex~=3.8",
+        "urllib3>=1.25.3,<3",
+        "lxml~=4.9",
+        "filelock~=3.0",
+        "tenacity~=8.2.2",
+    ],
     python_requires=">=3.8.0",
     # indicate that we have type information
     package_data={"blobfile": ["py.typed", "VERSION"]},