Merge branch 'main' into renaud.hartert/error-details

Author: renaudhartert-db

databricks/sdk/data_plane.py

@@ -1,9 +1,82 @@
+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 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, 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 token(self, endpoint, auth_details):
+        key = f"{endpoint}:{auth_details}"
+
+        # First, try to read without acquiring the lock to avoid contention.
+        # Reads are atomic, so this is safe.
+        token_source = self._token_sources.get(key)
+        if token_source:
+            return token_source.token()
+
+        # If token_source is not found, acquire the lock and check again.
+        with self._lock:
+            # 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
+                )
+                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, 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
+        self._token_exchange_host = token_exchange_host
+
+    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,
+        )
+
 
 @dataclass
 class DataPlaneDetails:
@@ -17,6 +90,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."""
 

tests/integration/conftest.py

@@ -7,7 +7,7 @@
 
 import pytest
 
-from databricks.sdk import AccountClient, WorkspaceClient
+from databricks.sdk import AccountClient, FilesAPI, FilesExt, WorkspaceClient
 from databricks.sdk.service.catalog import VolumeType
 
 
@@ -125,6 +125,18 @@ def volume(ucws, schema):
     ucws.volumes.delete(volume.full_name)
 
 
+@pytest.fixture(scope="session", params=[False, True])
+def files_api(request, ucws) -> FilesAPI:
+    if request.param:
+        # ensure new Files API client is used for files of any size
+        ucws.config.multipart_upload_min_stream_size = 0
+        # enable new Files API client
+        return FilesExt(ucws.api_client, ucws.config)
+    else:
+        # use the default client
+        return ucws.files
+
+
 @pytest.fixture()
 def workspace_dir(w, random):
     directory = f"/Users/{w.current_user.me().user_name}/dir-{random(12)}"

tests/integration/test_auth.py

@@ -108,12 +108,12 @@ def _get_lts_versions(w) -> typing.List[SparkVersion]:
     return lts_runtimes
 
 
-def test_runtime_auth_from_jobs_volumes(ucws, fresh_wheel_file, env_or_skip, random, volume):
+def test_runtime_auth_from_jobs_volumes(ucws, files_api, fresh_wheel_file, env_or_skip, random, volume):
     dbr_versions = [v for v in _get_lts_versions(ucws) if int(v.key.split(".")[0]) >= 15]
 
     volume_wheel = f"{volume}/tmp/wheels/{random(10)}/{fresh_wheel_file.name}"
     with fresh_wheel_file.open("rb") as f:
-        ucws.files.upload(volume_wheel, f)
+        files_api.upload(volume_wheel, f)
 
     lib = Library(whl=volume_wheel)
     return _test_runtime_auth_from_jobs_inner(ucws, env_or_skip, random, dbr_versions, lib)

tests/integration/test_data_plane.py

@@ -0,0 +1,15 @@
+from databricks.sdk.data_plane import DataPlaneTokenSource
+
+
+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.host, ucws._config.oauth_token)
+    dp_token = ts.token(info.endpoint_url, info.authorization_details)
+
+    assert dp_token.valid

tests/integration/test_files.py

@@ -218,30 +218,30 @@ def create_volume(w, catalog, schema, volume):
         return ResourceWithCleanup(lambda: w.volumes.delete(res.full_name))
 
 
-def test_files_api_upload_download(ucws, random):
+def test_files_api_upload_download(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
     with ResourceWithCleanup.create_schema(w, "main", schema):
         with ResourceWithCleanup.create_volume(w, "main", schema, volume):
             f = io.BytesIO(b"some text data")
             target_file = f"/Volumes/main/{schema}/{volume}/filesit-with-?-and-#-{random()}.txt"
-            w.files.upload(target_file, f)
-            with w.files.download(target_file).contents as f:
+            files_api.upload(target_file, f)
+            with files_api.download(target_file).contents as f:
                 assert f.read() == b"some text data"
 
 
-def test_files_api_read_twice_from_one_download(ucws, random):
+def test_files_api_read_twice_from_one_download(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
     with ResourceWithCleanup.create_schema(w, "main", schema):
         with ResourceWithCleanup.create_volume(w, "main", schema, volume):
             f = io.BytesIO(b"some text data")
             target_file = f"/Volumes/main/{schema}/{volume}/filesit-{random()}.txt"
-            w.files.upload(target_file, f)
+            files_api.upload(target_file, f)
 
-            res = w.files.download(target_file).contents
+            res = files_api.download(target_file).contents
 
             with res:
                 assert res.read() == b"some text data"
@@ -251,82 +251,82 @@ def test_files_api_read_twice_from_one_download(ucws, random):
                     res.read()
 
 
-def test_files_api_delete_file(ucws, random):
+def test_files_api_delete_file(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
     with ResourceWithCleanup.create_schema(w, "main", schema):
         with ResourceWithCleanup.create_volume(w, "main", schema, volume):
             f = io.BytesIO(b"some text data")
             target_file = f"/Volumes/main/{schema}/{volume}/filesit-{random()}.txt"
-            w.files.upload(target_file, f)
-            w.files.delete(target_file)
+            files_api.upload(target_file, f)
+            files_api.delete(target_file)
 
 
-def test_files_api_get_metadata(ucws, random):
+def test_files_api_get_metadata(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
     with ResourceWithCleanup.create_schema(w, "main", schema):
         with ResourceWithCleanup.create_volume(w, "main", schema, volume):
             f = io.BytesIO(b"some text data")
             target_file = f"/Volumes/main/{schema}/{volume}/filesit-{random()}.txt"
-            w.files.upload(target_file, f)
-            m = w.files.get_metadata(target_file)
+            files_api.upload(target_file, f)
+            m = files_api.get_metadata(target_file)
             assert m.content_type == "application/octet-stream"
             assert m.content_length == 14
             assert m.last_modified is not None
 
 
-def test_files_api_create_directory(ucws, random):
+def test_files_api_create_directory(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
     with ResourceWithCleanup.create_schema(w, "main", schema):
         with ResourceWithCleanup.create_volume(w, "main", schema, volume):
             target_directory = f"/Volumes/main/{schema}/{volume}/filesit-{random()}/"
-            w.files.create_directory(target_directory)
+            files_api.create_directory(target_directory)
 
 
-def test_files_api_list_directory_contents(ucws, random):
+def test_files_api_list_directory_contents(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
     with ResourceWithCleanup.create_schema(w, "main", schema):
         with ResourceWithCleanup.create_volume(w, "main", schema, volume):
             target_directory = f"/Volumes/main/{schema}/{volume}/filesit-{random()}"
-            w.files.upload(target_directory + "/file1.txt", io.BytesIO(b"some text data"))
-            w.files.upload(target_directory + "/file2.txt", io.BytesIO(b"some text data"))
-            w.files.upload(target_directory + "/file3.txt", io.BytesIO(b"some text data"))
+            files_api.upload(target_directory + "/file1.txt", io.BytesIO(b"some text data"))
+            files_api.upload(target_directory + "/file2.txt", io.BytesIO(b"some text data"))
+            files_api.upload(target_directory + "/file3.txt", io.BytesIO(b"some text data"))
 
-            result = list(w.files.list_directory_contents(target_directory))
+            result = list(files_api.list_directory_contents(target_directory))
             assert len(result) == 3
 
 
-def test_files_api_delete_directory(ucws, random):
+def test_files_api_delete_directory(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
     with ResourceWithCleanup.create_schema(w, "main", schema):
         with ResourceWithCleanup.create_volume(w, "main", schema, volume):
             target_directory = f"/Volumes/main/{schema}/{volume}/filesit-{random()}/"
-            w.files.create_directory(target_directory)
-            w.files.delete_directory(target_directory)
+            files_api.create_directory(target_directory)
+            files_api.delete_directory(target_directory)
 
 
-def test_files_api_get_directory_metadata(ucws, random):
+def test_files_api_get_directory_metadata(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
     with ResourceWithCleanup.create_schema(w, "main", schema):
         with ResourceWithCleanup.create_volume(w, "main", schema, volume):
             target_directory = f"/Volumes/main/{schema}/{volume}/filesit-{random()}/"
-            w.files.create_directory(target_directory)
-            w.files.get_directory_metadata(target_directory)
+            files_api.create_directory(target_directory)
+            files_api.get_directory_metadata(target_directory)
 
 
 @pytest.mark.benchmark
-def test_files_api_download_benchmark(ucws, random):
+def test_files_api_download_benchmark(ucws, files_api, random):
     w = ucws
     schema = "filesit-" + random()
     volume = "filesit-" + random()
@@ -335,7 +335,7 @@ def test_files_api_download_benchmark(ucws, random):
             # Create a 50 MB file
             f = io.BytesIO(bytes(range(256)) * 200000)
             target_file = f"/Volumes/main/{schema}/{volume}/filesit-benchmark-{random()}.txt"
-            w.files.upload(target_file, f)
+            files_api.upload(target_file, f)
 
             totals = {}
             for chunk_size_kb in [
@@ -357,7 +357,7 @@ def test_files_api_download_benchmark(ucws, random):
                 count = 10
                 for i in range(count):
                     start = time.time()
-                    f = w.files.download(target_file).contents
+                    f = files_api.download(target_file).contents
                     f.set_chunk_size(chunk_size)
                     with f as vf:
                         vf.read()

tests/test_data_plane.py

@@ -1,9 +1,80 @@
 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):
+    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()
+
+    retrieve_token.assert_called_once()
+    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["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)
+
+    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):
+    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)
+
+    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(