[Internal] Add DataPlane token source

Author: hectorcast-db

databricks/sdk/data_plane.py

@@ -1,8 +1,69 @@
+from __future__ import annotations
+
 import threading
 from dataclasses import dataclass
-from typing import Callable, List
+from typing import Callable, List, Optional
+
+from urllib import parse
 
 from databricks.sdk.oauth import Token
+from databricks.sdk import oauth, config, credentials_provider
+
+
+URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
+JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
+OIDC_TOKEN_PATH = "/oidc/v1/token"
+
+class DataPlaneTokenSource:
+    """
+    Manages token sources for multiple DataPlane endpoints.
+    """
+    # TODO: Enable async once its stable. @oauth_credentials_provider must also have async enabled.
+    def __init__(self, cfg: config.Config, cpts: Callable[[], Token], disable_async: Optional[bool] = True):
+        self._cpts = cpts
+        self._cfg = cfg
+        self._token_sources = {}
+        self._disable_async = disable_async
+
+    def token(self, endpoint: str, auth_details: str):
+        """
+        Get a token for a specific DataPlane endpoint.
+        :param endpoint: endpoint URL for which to get a token
+        :param auth_details:  authorization details used to generate the token
+        :return: a token for the specified endpoint
+        """
+        key = f"{endpoint}:{auth_details}"
+        token_source = self._token_sources.get(key)
+        if not token_source:
+            token_source = DataPlaneEndpointTokenSource(self._cfg, self._cpts, auth_details, self._disable_async)
+            self._token_sources[key] = token_source
+        return token_source.token()
+
+
+class DataPlaneEndpointTokenSource(oauth.Refreshable):
+    """
+    A token source for a specific DataPlane endpoint.
+    """
+    def __init__(self, cfg: config.Config, cpts: Callable[[], Token], auth_details: str, disable_async: bool):
+        super().__init__(disable_async=disable_async)
+        self._auth_details = auth_details
+        self._cpts = cpts
+        self._cfg = cfg
+
+    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=self._cfg.client_id,
+                              client_secret=self._cfg.client_secret,
+                              token_url=self._cfg.host + OIDC_TOKEN_PATH,
+                              params=params,
+                              headers=headers)
+
 
 
 @dataclass
@@ -16,6 +77,8 @@ class DataPlaneDetails:
     """Token to query the DataPlane endpoint."""
 
 
+## Old implementation. #TODO: Remove after the new implementation is used
+
 class DataPlaneService:
     """Helper class to fetch and manage DataPlane details."""
     from .service.serving import DataPlaneInfo

tests/test_data_plane.py

@@ -1,14 +1,84 @@
 from datetime import datetime, timedelta
+from unittest.mock import patch
+from urllib import parse
 
+from databricks.sdk import data_plane, oauth
 from databricks.sdk.data_plane import DataPlaneService
 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))
+
+def success_callable(token: oauth.Token):
+    def success() -> oauth.Token:
+        return token
+    return success
+
+def patch_retrieve_token(token: oauth.Token):
+    return patch("databricks.sdk.oauth.retrieve_token", return_value=token)
+
+def test_endpoint_token_source_get_token(config):
+    config.client_id = "client_id"
+    config.client_secret = "client_secret"
+    token_source = data_plane.DataPlaneEndpointTokenSource(config, 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()
+
+    retrieve_token.assert_called_once()
+    args, kwargs = retrieve_token.call_args
+
+    assert kwargs["client_id"] == config.client_id
+    assert kwargs["client_secret"] == config.client_secret
+    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["headers"] == {"Content-Type": "application/x-www-form-urlencoded"}
+
+def test_token_source_get_token_not_existing(config):
+    config.client_id = "client_id"
+    config.client_secret = "client_secret"
+    token_source = data_plane.DataPlaneTokenSource(config, 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")
+
+    retrieve_token.assert_called_once()
+    assert result_token.access_token == dp_token.access_token
+    assert "endpoint:authDetails" in token_source._token_sources
+
+class MockEndpointTokenSource:
+
+    def __init__(self, token: oauth.Token):
+        self._token = token
+
+    def token(self):
+        return self._token
+
+def test_token_source_get_token_existing(config):
+    config.client_id = "client_id"
+    config.client_secret = "client_secret"
+    another_token = Token(access_token="another token", token_type="type", expiry=datetime.now() + timedelta(hours=1))
+    token_source = data_plane.DataPlaneTokenSource(config, success_callable(cp_token),  disable_async=True)
+    token_source._token_sources["endpoint:authDetails"] = MockEndpointTokenSource(another_token)
+
+    with patch("databricks.sdk.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
+            result_token = token_source.token(endpoint="endpoint", auth_details="authDetails")
+
+    retrieve_token.assert_not_called()
+    assert result_token.access_token == another_token.access_token
+
+
+
+
+
+## These tests are for the old implementation. #TODO: Remove after the new implementation is used
+
 info = DataPlaneInfo(authorization_details="authDetails", endpoint_url="url")
 
 token = Token(access_token="token", token_type="type", expiry=datetime.now() + timedelta(hours=1))
 
-
 class MockRefresher:
 
     def __init__(self, expected: str):