[Internal] Add DataPlane token source

Author: hectorcast-db

Makefile

@@ -27,6 +27,9 @@ lint:
 test:
 	pytest -m 'not integration and not benchmark' --cov=databricks --cov-report html tests
 
+integration-dataplane:
+	pytest -v --cov=databricks --cov-report html tests/integration/test_dataplane.py
+
 integration:
 	pytest -n auto -m 'integration and not benchmark' --reruns 2 --dist loadgroup --cov=databricks --cov-report html tests
 

databricks/sdk/data_plane.py

@@ -1,9 +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 import config, oauth
 from databricks.sdk.oauth import Token
 
+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:
+    """
+    EXPERIMENTAL 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, disable_async: Optional[bool] = True):
+        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, auth_details, self._disable_async)
+            self._token_sources[key] = token_source
+        return token_source.token()
+
+
+class DataPlaneEndpointTokenSource(oauth.Refreshable):
+    """
+    EXPERIMENTAL A token source for a specific DataPlane endpoint.
+    """
+
+    def __init__(self, cfg: config.Config, auth_details: str, disable_async: bool):
+        super().__init__(disable_async=disable_async)
+        self._auth_details = auth_details
+        self._cpts = cfg.oauth_token
+        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
 class DataPlaneDetails:
@@ -16,6 +76,9 @@ 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/integration/test_dataplane.py

@@ -0,0 +1,35 @@
+import base64
+import io
+import json
+import re
+import shutil
+import subprocess
+import sys
+import typing
+import urllib.parse
+from functools import partial
+from pathlib import Path
+
+import pytest
+
+from databricks.sdk.data_plane import DataPlaneTokenSource
+from databricks.sdk.service.compute import (ClusterSpec, DataSecurityMode,
+                                            Library, ResultType, SparkVersion)
+from databricks.sdk.service.jobs import NotebookTask, Task, ViewType
+from databricks.sdk.service.workspace import ImportFormat
+
+
+def test_data_plane_token_source(ucws, env_or_skip):
+    endpoint = env_or_skip("SERVING_ENDPOINT_NAME")
+    serving_endpoint = ucws.serving_endpoints.get(endpoint)
+    assert serving_endpoint.data_plane_info is not None
+    assert serving_endpoint.data_plane_info.query_info is not None
+
+    info = serving_endpoint.data_plane_info.query_info
+
+    ts = DataPlaneTokenSource(ucws.config)
+    dp_token = ts.token(info.endpoint_url, info.authorization_details)
+
+    assert dp_token.valid 
+
+

tests/test_data_plane.py

@@ -1,9 +1,93 @@
 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 test_endpoint_token_source_get_token(config):
+    config.client_id = "client_id"
+    config.client_secret = "client_secret"
+    config.oauth_token = success_callable(cp_token)
+    token_source = data_plane.DataPlaneEndpointTokenSource(config, "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"
+    config.oauth_token = success_callable(cp_token)
+    token_source = data_plane.DataPlaneTokenSource(config, 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"
+    config.oauth_token = success_callable(cp_token)
+    another_token = Token(access_token="another token",
+                          token_type="type",
+                          expiry=datetime.now() + timedelta(hours=1))
+    token_source = data_plane.DataPlaneTokenSource(config, 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))