black fmt

Author: hectorcast-db

databricks/sdk/data_plane.py

@@ -19,17 +19,14 @@ class DataPlaneTokenSource:
     """
 
     # TODO: Enable async once its stable. @oauth_credentials_provider must also have async enabled.
-    def __init__(self,
-                 token_exchange_host: str,
-                 cpts: Callable[[], Token],
-                 disable_async: Optional[bool] = True):
+    def __init__(self, token_exchange_host: str, cpts: Callable[[], Token], disable_async: Optional[bool] = True):
         self._cpts = cpts
         self._token_exchange_host = token_exchange_host
         self._token_sources = {}
         self._disable_async = disable_async
         self._lock = threading.Lock()
 
-    def get_token(self, endpoint, auth_details):
+    def token(self, endpoint, auth_details):
         key = f"{endpoint}:{auth_details}"
 
         # First, try to read without acquiring the lock to avoid contention.
@@ -43,8 +40,9 @@ def get_token(self, endpoint, auth_details):
             # Another thread might have created it while we were waiting for the lock.
             token_source = self._token_sources.get(key)
             if not token_source:
-                token_source = DataPlaneEndpointTokenSource(self._token_exchange_host, self._cpts,
-                                                            auth_details, self._disable_async)
+                token_source = DataPlaneEndpointTokenSource(
+                    self._token_exchange_host, self._cpts, auth_details, self._disable_async
+                )
                 self._token_sources[key] = token_source
 
         return token_source.token()
@@ -55,8 +53,7 @@ class DataPlaneEndpointTokenSource(oauth.Refreshable):
     EXPERIMENTAL A token source for a specific DataPlane endpoint.
     """
 
-    def __init__(self, token_exchange_host: str, cpts: Callable[[], Token], auth_details: str,
-                 disable_async: bool):
+    def __init__(self, token_exchange_host: str, cpts: Callable[[], Token], auth_details: str, disable_async: bool):
         super().__init__(disable_async=disable_async)
         self._auth_details = auth_details
         self._cpts = cpts
@@ -65,16 +62,20 @@ def __init__(self, token_exchange_host: str, cpts: Callable[[], Token], auth_det
     def refresh(self) -> Token:
         control_plane_token = self._cpts()
         headers = {"Content-Type": URL_ENCODED_CONTENT_TYPE}
-        params = parse.urlencode({
-            "grant_type": JWT_BEARER_GRANT_TYPE,
-            "authorization_details": self._auth_details,
-            "assertion": control_plane_token.access_token
-        })
-        return oauth.retrieve_token(client_id="",
-                                    client_secret="",
-                                    token_url=self._token_exchange_host + OIDC_TOKEN_PATH,
-                                    params=params,
-                                    headers=headers)
+        params = parse.urlencode(
+            {
+                "grant_type": JWT_BEARER_GRANT_TYPE,
+                "authorization_details": self._auth_details,
+                "assertion": control_plane_token.access_token,
+            }
+        )
+        return oauth.retrieve_token(
+            client_id="",
+            client_secret="",
+            token_url=self._token_exchange_host + OIDC_TOKEN_PATH,
+            params=params,
+            headers=headers,
+        )
 
 
 @dataclass

tests/integration/test_dataplane.py renamed to tests/integration/test_data_plane.py

@@ -7,12 +7,8 @@
 from databricks.sdk.oauth import Token
 from databricks.sdk.service.serving import DataPlaneInfo
 
-cp_token = Token(access_token="control plane token",
-                 token_type="type",
-                 expiry=datetime.now() + timedelta(hours=1))
-dp_token = Token(access_token="data plane token",
-                 token_type="type",
-                 expiry=datetime.now() + timedelta(hours=1))
+cp_token = Token(access_token="control plane token", token_type="type", expiry=datetime.now() + timedelta(hours=1))
+dp_token = Token(access_token="data plane token", token_type="type", expiry=datetime.now() + timedelta(hours=1))
 
 
 def success_callable(token: oauth.Token):
@@ -24,10 +20,9 @@ def success() -> oauth.Token:
 
 
 def test_endpoint_token_source_get_token(config):
-    token_source = data_plane.DataPlaneEndpointTokenSource(config.host,
-                                                           success_callable(cp_token),
-                                                           "authDetails",
-                                                           disable_async=True)
+    token_source = data_plane.DataPlaneEndpointTokenSource(
+        config.host, success_callable(cp_token), "authDetails", disable_async=True
+    )
 
     with patch("databricks.sdk.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
         token_source.token()
@@ -36,18 +31,18 @@ def test_endpoint_token_source_get_token(config):
     args, kwargs = retrieve_token.call_args
 
     assert kwargs["token_url"] == config.host + "/oidc/v1/token"
-    assert kwargs["params"] == parse.urlencode({
-        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
-        "authorization_details": "authDetails",
-        "assertion": cp_token.access_token
-    })
+    assert kwargs["params"] == parse.urlencode(
+        {
+            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+            "authorization_details": "authDetails",
+            "assertion": cp_token.access_token,
+        }
+    )
     assert kwargs["headers"] == {"Content-Type": "application/x-www-form-urlencoded"}
 
 
 def test_token_source_get_token_not_existing(config):
-    token_source = data_plane.DataPlaneTokenSource(config.host,
-                                                   success_callable(cp_token),
-                                                   disable_async=True)
+    token_source = data_plane.DataPlaneTokenSource(config.host, success_callable(cp_token), disable_async=True)
 
     with patch("databricks.sdk.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
         result_token = token_source.token(endpoint="endpoint", auth_details="authDetails")
@@ -67,9 +62,7 @@ def token(self):
 
 
 def test_token_source_get_token_existing(config):
-    another_token = Token(access_token="another token",
-                          token_type="type",
-                          expiry=datetime.now() + timedelta(hours=1))
+    another_token = Token(access_token="another token", token_type="type", expiry=datetime.now() + timedelta(hours=1))
     token_source = data_plane.DataPlaneTokenSource(config.host, success_callable(token), disable_async=True)
     token_source._token_sources["endpoint:authDetails"] = MockEndpointTokenSource(another_token)