feat: Retry http requests in auth flow (#49)

* feat: Retry http requests in auth flow

* Update src/neptune_api/client.py

Co-authored-by: Krzysztof Godlewski <krzysztof.godlewski@neptune.ai>

---------

Co-authored-by: Krzysztof Godlewski <krzysztof.godlewski@neptune.ai>

Author: PatrykGala

pyproject.toml

@@ -15,6 +15,7 @@ attrs = ">=21.3.0"
 python-dateutil = "^2.8.0"
 PyJWT = "^2.0.0"
 protobuf = ">=4,<6"
+backoff = "^2.2.1"
 
 [tool.poetry]
 name = "neptune-api"
@@ -95,3 +96,4 @@ check_untyped_defs = "True"
 warn_return_any = "True"
 show_error_codes = "True"
 warn_unused_ignores = "True"
+ignore_missing_imports = "True"

src/neptune_api/client.py

@@ -1,5 +1,6 @@
 __all__ = ["Client", "AuthenticatedClient", "NeptuneAuthenticator"]
 
+import logging
 import ssl
 import threading
 import typing
@@ -12,6 +13,7 @@
     Union,
 )
 
+import backoff
 import httpx
 from attrs import (
     define,
@@ -23,6 +25,9 @@
 from neptune_api.errors import UnableToRefreshTokenError
 from neptune_api.types import OAuthToken
 
+# Disable httpx logging, httpx logs requests at INFO level
+logging.getLogger("httpx").setLevel(logging.WARN)
+
 
 @define
 class Client:
@@ -349,12 +354,9 @@ def __init__(
         self._client = client
         self._token: Optional[OAuthToken] = None
 
-    def _refresh_existing_token(self) -> None:
+    def _refresh_existing_token(self) -> OAuthToken:
         if self._token is None:
-            # This should never happen, but just in case
-            self._token = self._api_key_exchange_factory(self._client, self._credentials)
-            return
-
+            raise ValueError("Cannot refresh an empty token")
         try:
             response = self._client.get_httpx_client().post(
                 url=self._token_refreshing_endpoint,
@@ -366,21 +368,26 @@ def _refresh_existing_token(self) -> None:
                 },
             )
             data = response.json()
+            return OAuthToken.from_tokens(access=data["access_token"], refresh=data["refresh_token"])
         except Exception as e:
             raise UnableToRefreshTokenError("Unable to refresh token") from e
 
-        self._token = OAuthToken.from_tokens(access=data["access_token"], refresh=data["refresh_token"])
-
+    @backoff.on_exception(backoff.expo, Exception, max_time=30, max_tries=3)
     def _refresh_token(self) -> None:
         with self.__LOCK:
             if self._token is not None:
-                self._refresh_existing_token()
+                self._token = self._refresh_existing_token()
 
             if self._token is None:
                 self._token = self._api_key_exchange_factory(self._client, self._credentials)
 
     def _refresh_token_if_expired(self) -> None:
-        if self._token is None or self._token.is_expired:
+        try:
+            if self._token is None or self._token.is_expired:
+                self._refresh_token()
+        except Exception:
+            # Reset the token to None to force a new token retrieval
+            self._token = None
             self._refresh_token()
 
     def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: