Cleanup

Author: hectorcast-db

databricks/sdk/oauth.py

@@ -9,12 +9,12 @@
 import urllib.parse
 import webbrowser
 from abc import abstractmethod
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from enum import Enum
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from typing import Any, Dict, List, Optional
-from concurrent.futures import ThreadPoolExecutor
 
 import requests
 import requests.auth
@@ -188,67 +188,67 @@ def retrieve_token(client_id,
     except Exception as e:
         raise NotImplementedError(f"Not supported yet: {e}")
 
+
 class _TokenState(Enum):
     """
-    tokenState represents the state of the token. Each token can be in one of
+    Represents the state of a token. Each token can be in one of
     the following three states:
       - FRESH: The token is valid.
       - STALE: The token is valid but will expire soon.
       - EXPIRED: The token has expired and cannot be used.
-
-    Token state through time:
-     issue time     expiry time
-    	v               v
-    	| fresh | stale | expired -> time
-    	|     valid     |
     """
     FRESH = 1 # The token is valid.
     STALE = 2 # The token is valid but will expire soon.
     EXPIRED = 3 # The token has expired and cannot be used.
 
 
 class Refreshable(TokenSource):
-    _executor = ThreadPoolExecutor(max_workers=10)
-    _default_stale_duration = 3
+    """A token source that supports refreshing expired tokens."""
+
+    _EXECUTOR = ThreadPoolExecutor(max_workers=10)
+    _DEFAULT_STALE_DURATION = timedelta(minutes=3)
 
-    def __init__(self, token=None, disable_async = True, stale_duration=timedelta(minutes=_default_stale_duration)):
-        self._lock = threading.Lock() # to guard _token
+    def __init__(self,
+                 token: Token = None,
+                 disable_async: bool = True,
+                 stale_duration: timedelta = _DEFAULT_STALE_DURATION):
+        self._lock = threading.Lock()
         self._token = token
         self._stale_duration = stale_duration
         self._disable_async = disable_async
         self._is_refreshing = False
         self._refresh_err = False
 
-    def token(self, blocking=False) -> Token:
+    def token(self) -> Token:
+        """Returns a valid token, blocking if async refresh is disabled."""
         if self._disable_async:
             return self._blocking_token()
         return self._async_token()
 
     def _async_token(self) -> Token:
-        self._lock.acquire()
-        token_state = self._token_state()
-        token = self._token
-        self._lock.release()
-        match token_state:
-            case _TokenState.FRESH:
-                return token
-            case _TokenState.STALE:
-                self._trigger_async_refresh()
-                return token
-            case _: #Expired
-                return self._blocking_token()
+        """
+        Returns a token.
+        If the token is stale, triggers an asynchronous refresh.
+        If the token is expired, refreshes it synchronously, blocking until the refresh is complete.
+        """
+        with self._lock:
+            token_state = self._token_state()
+            token = self._token
 
+        if token_state == _TokenState.FRESH:
+            return token
+        if token_state == _TokenState.STALE:
+            self._trigger_async_refresh()
+            return token
+        return self._blocking_token()
 
     def _token_state(self) -> _TokenState:
-        """
-        Returns the state of the token.
-        """
-        # Invalid tokens are considered expired.
+        """Returns the current state of the token."""
         if not self._token or not self._token.valid:
             return _TokenState.EXPIRED
-        # Tokens without an expiry are considered always.
         if not self._token.expiry:
             return _TokenState.FRESH
+
         lifespan = self._token.expiry - datetime.now()
         if lifespan < timedelta(seconds=0):
             return _TokenState.EXPIRED
@@ -257,13 +257,10 @@ def _token_state(self) -> _TokenState:
         return _TokenState.FRESH
 
     def _blocking_token(self) -> Token:
-
-        # The lock is kept for the entire operation to ensure that only one
-        # refresh operation is running at a time.
+        """Returns a token, blocking if necessary to refresh it."""
         with self._lock:
             # This is important to recover from potential previous failed attempts
-            # to refresh the token asynchronously, see declaration of refresh_err for
-            # more information.
+            # to refresh the token asynchronously.
             self._refresh_err = False
             self._is_refreshing = False
 
@@ -273,13 +270,13 @@ def _blocking_token(self) -> Token:
             if self._token_state() != _TokenState.EXPIRED:
                 return self._token
 
-            # Refresh the token
             self._token = self.refresh()
             return self._token
 
-
     def _trigger_async_refresh(self):
-        # Note: this is not thread safe.
+        """Starts an asynchronous refresh if none is in progress."""
+
+        # Note: _refresh_internal function is not thread safe.
         # Only call it inside the lock.
         def _refresh_internal():
             try:
@@ -288,12 +285,11 @@ def _refresh_internal():
                 self._refresh_err = True
             finally:
                 self._is_refreshing = False
-        # The lock is kept for the entire operation to ensure that only one
-        # refresh operation is running at a time.
+
         with self._lock:
             if not self._is_refreshing and not self._refresh_err:
                 self._is_refreshing = True
-                self._executor.submit(_refresh_internal)
+                self._EXECUTOR.submit(_refresh_internal)
 
     @abstractmethod
     def refresh(self) -> Token: