Merge pull request #1890 from neptune-ai/pg/oauth-fix

chore: Use retry mechanism while refreshing oauth token.
neptune-ai · Dec 19, 2024 · 59e1daa · 59e1daa
2 parents d1fa7b2 + 454a1ed
commit 59e1daa
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/neptune/internal/oauth.py b/src/neptune/internal/oauth.py
@@ -52,7 +52,7 @@ def __call__(self, r):
         try:
             return self._add_token(r)
         except TokenExpiredError:
-            self._refresh_token()
+            self.refresh_token_if_needed(force=True)
             return self._add_token(r)
 
     def _add_token(self, r):
@@ -62,8 +62,8 @@ def _add_token(self, r):
         return r
 
     @with_api_exceptions_handler
-    def refresh_token_if_needed(self):
-        if self.token_expires_at - time.time() < 30:
+    def refresh_token_if_needed(self, force=False):
+        if (self.token_expires_at - time.time()) < 30 or force:
             self._refresh_token()
 
     def _refresh_token(self):