refactor(dvc_webdav): ensure BearerAuthClient is initialized only once

Author: GreenHatHG

dvc_webdav/__init__.py

@@ -23,14 +23,9 @@ def ask_password(host, user):
 
 @wrap_with(threading.Lock())
 @memoize
-def get_bearer_auth_client(bearer_token_command: str, token: Optional[str] = None, save_token_cb=None):
-    logger.debug(
-        "Bearer token command provided, using BearerAuthClient, command: %s",
-        bearer_token_command,
-    )
-    return BearerAuthClient(
-        bearer_token_command, token=token, save_token_cb=save_token_cb
-    )
+def get_bearer_auth_client(bearer_token_command: str):
+    logger.debug("Bearer token command provided, using BearerAuthClient, command: %s", bearer_token_command, )
+    return BearerAuthClient(bearer_token_command)
 
 
 class WebDAVFileSystem(FileSystem):  # pylint:disable=abstract-method
@@ -54,9 +49,11 @@ def __init__(self, **config):
             }
         )
         if bearer_token_command := config.get("bearer_token_command"):
-            self.fs_args["http_client"] = get_bearer_auth_client(
-                bearer_token_command, token=config.get('token'), save_token_cb=self._save_token
-            )
+            client = get_bearer_auth_client(bearer_token_command)
+            client.save_token_cb = self._save_token
+            if token := config.get("token"):
+                client.update_token(token)
+            self.fs_args["http_client"] = client
 
     def unstrip_protocol(self, path: str) -> str:
         return self.fs_args["base_url"] + "/" + path

dvc_webdav/bearer_auth_client.py

@@ -92,26 +92,16 @@ def __init__(
             self,
             bearer_token_command: str,
             save_token_cb: Optional[TokenSaver] = None,
-            token: Optional[str] = None,
             **kwargs,
     ):
         super().__init__(**kwargs)
         if not isinstance(bearer_token_command, str) or not bearer_token_command.strip():
             raise ValueError("[BearerAuthClient] bearer_token_command must be a non-empty string")
         self.bearer_token_command = bearer_token_command
         self.save_token_cb = save_token_cb
-        self._token: Optional[str] = token
+        self._token: Optional[str] = None
         self._lock = threading.Lock()
 
-        if not self._token:
-            auth_header = self.headers.get("Authorization")
-            if auth_header and auth_header.startswith("Bearer "):
-                self._token = auth_header.split(" ", 1)[1]
-
-        if self._token:
-            logger.debug("[BearerAuthClient] Initial token found, setting Authorization header.")
-            self.headers["Authorization"] = f"Bearer {self._token}"
-
     def _refresh_token_locked(self) -> None:
         """Execute token command and update state."""
         _log_with_thread(logging.DEBUG, "[BearerAuthClient] Refreshing token via command...")
@@ -142,6 +132,16 @@ def _ensure_token(self) -> None:
             if not self._token:
                 self._refresh_token_locked()
 
+    def update_token(self, token: Optional[str]) -> None:
+        """Update the token with a new one"""
+        if not token:
+            return
+
+        with self._lock:
+            if self._token != token:
+                self._token = token
+                self.headers["Authorization"] = f"Bearer {token}"
+
     def request(self, *args, **kwargs) -> httpx.Response:
         """Wraps httpx.request with auto-refresh logic for 401 Unauthorized."""
         self._ensure_token()