From 7b64e60d82450d7f0350235371c46e4ca0404141 Mon Sep 17 00:00:00 2001
From: Steven Rybicki <stevenrybicki@anthropic.com>
Date: Thu, 1 Feb 2024 15:24:48 -0800
Subject: [PATCH] Add retries to allow for rate limits on WIF

---
 blobfile/_gcp.py | 12 +++++++++++-
 env.yaml         |  1 +
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/blobfile/_gcp.py b/blobfile/_gcp.py
index e8eb80b..54528a8 100644
--- a/blobfile/_gcp.py
+++ b/blobfile/_gcp.py
@@ -33,6 +33,12 @@
     strip_slashes,
 )
 
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+)
+
 MAX_EXPIRATION = 7 * 24 * 60 * 60
 OAUTH_TOKEN = "oauth_token"
 ANONYMOUS = "anonymous"
@@ -474,7 +480,11 @@ def build_req() -> Request:
     result = json.loads(resp.data)
     return (OAUTH_TOKEN, result["access_token"]), now + float(result["expires_in"])
 
-
+@retry(
+    retry=retry_if_exception_type(ValueError),
+    before=get_wif_credentials,
+    stop=stop_after_attempt(1),
+)
 def _get_access_token(conf: Config, key: Any) -> Tuple[Any, float]:
     if os.environ.get("BLOBFILE_FORCE_GOOGLE_ANONYMOUS_AUTH", "0") == "1":
         return (ANONYMOUS, ""), float("inf")
diff --git a/env.yaml b/env.yaml
index 448c201..1d89b22 100644
--- a/env.yaml
+++ b/env.yaml
@@ -19,3 +19,4 @@ dependencies:
     - boto3==1.15.18
     - lxml-stubs==0.4.0
     - xmltodict==0.13.0
+    - tenacity==8.2.2