From 4a428ab0840310cab4016e0735262458f23fc9de Mon Sep 17 00:00:00 2001 From: sonhmai <14060682+sonhmai@users.noreply.github.com> Date: Fri, 12 Jun 2026 19:02:06 +0700 Subject: [PATCH] feat(databricks-volume): add token provider authentication --- docs/python/resource/databricks_volume.mdx | 72 +++-- docs/python/setup/databricks.mdx | 53 ++- .../langchain/databricks_volume_deepagent.py | 21 +- .../databricks_volume/databricks_volume.py | 22 +- python/mirage/accessor/databricks_volume.py | 38 +-- .../mirage/core/databricks_volume/client.py | 278 ++++++++++++++++ python/mirage/core/databricks_volume/copy.py | 51 +-- python/mirage/core/databricks_volume/mkdir.py | 13 +- python/mirage/core/databricks_volume/read.py | 49 +-- .../mirage/core/databricks_volume/readdir.py | 14 +- python/mirage/core/databricks_volume/rm.py | 46 +-- python/mirage/core/databricks_volume/rmdir.py | 21 +- python/mirage/core/databricks_volume/stat.py | 7 +- .../mirage/core/databricks_volume/stream.py | 29 +- .../mirage/core/databricks_volume/unlink.py | 10 +- python/mirage/core/databricks_volume/write.py | 23 +- .../resource/databricks_volume/__init__.py | 10 +- .../resource/databricks_volume/config.py | 4 +- .../databricks_volume/databricks_volume.py | 39 ++- .../databricks_volume/token_provider.py | 64 ++++ python/mirage/workspace/snapshot/state.py | 4 + .../tests/core/databricks_volume/conftest.py | 107 +++--- .../core/databricks_volume/test_client.py | 304 ++++++++++++++++++ .../tests/core/databricks_volume/test_head.py | 6 +- .../tests/core/databricks_volume/test_path.py | 1 + .../tests/core/databricks_volume/test_read.py | 35 +- .../core/databricks_volume/test_readdir.py | 25 +- .../tests/core/databricks_volume/test_stat.py | 40 +-- .../core/databricks_volume/test_stream.py | 34 +- .../databricks_volume/test_accessor.py | 42 +-- .../test_databricks_volume.py | 141 +++++--- .../databricks_volume/test_token_provider.py | 44 +++ .../tests/resource/test_state_round_trip.py | 52 +++ 33 files changed, 1124 insertions(+), 575 deletions(-) create mode 100644 python/mirage/core/databricks_volume/client.py create mode 100644 python/mirage/resource/databricks_volume/token_provider.py create mode 100644 python/tests/core/databricks_volume/test_client.py create mode 100644 python/tests/resource/databricks_volume/test_token_provider.py diff --git a/docs/python/resource/databricks_volume.mdx b/docs/python/resource/databricks_volume.mdx index 85223cb2d..13d31b2c2 100644 --- a/docs/python/resource/databricks_volume.mdx +++ b/docs/python/resource/databricks_volume.mdx @@ -17,14 +17,19 @@ from mirage import MountMode, Workspace from mirage.resource.databricks_volume import ( DatabricksVolumeConfig, DatabricksVolumeResource, + StaticTokenProvider, ) -resource = DatabricksVolumeResource(DatabricksVolumeConfig( - catalog="main", - schema="default", - volume="agent_files", - root_path="/reports", -)) +resource = DatabricksVolumeResource( + DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", + catalog="main", + schema="default", + volume="agent_files", + root_path="/reports", + ), + token_provider=StaticTokenProvider("token"), +) ws = Workspace({"/dbx": resource}, mode=MountMode.READ) ``` @@ -33,12 +38,15 @@ ws = Workspace({"/dbx": resource}, mode=MountMode.READ) | `catalog` | required | Unity Catalog catalog name. | | `schema` | required | Unity Catalog schema name. | | `volume` | required | Unity Catalog volume name. | +| `host` | required | Databricks workspace host. | | `root_path` | `/` | Subdirectory inside the volume to expose. | -| `host` | `None` | Optional workspace host override. | -| `token` | `None` | Optional PAT override. Redacted in snapshots. | -| `profile` | `None` | Optional Databricks SDK profile name. | | `timeout` | `30` | Request timeout in seconds. | +Credentials are supplied separately through a token provider. Mirage includes +`StaticTokenProvider` for fixed tokens and `DatabricksProfileTokenProvider` for +Databricks CLI profiles. Applications can implement `get_token()` to supply +rotating or request-scoped tokens. + ## Mount mode `read` or `write`. @@ -51,6 +59,7 @@ Given: ```python DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main", schema="default", volume="agent_files", @@ -134,21 +143,38 @@ or local-only utilities). ## Snapshot behavior -`token` is redacted in resource state. Loading a snapshot back requires an -override config that provides fresh credentials if the runtime auth chain does -not already supply them. +Snapshots contain the volume location config but never the token provider. +Loading a snapshot requires a full resource override for the mount: + +```python +workspace = Workspace.load( + snapshot, + resources={"/dbx/": resource}, +) +``` ## Databricks Apps -For Databricks Apps, prefer SDK-default auth and keep Mirage in-process: +For Databricks Apps, pass an application-owned provider. Its `get_token()` +method can read the current request's OBO token or return a service-principal +token managed by the application: ```python +class AppTokenProvider: + def get_token(self) -> str: + return current_request_token.get() + + config = DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main", schema="default", volume="agent_files", ) -resource = DatabricksVolumeResource(config) +resource = DatabricksVolumeResource( + config, + token_provider=AppTokenProvider(), +) ws = Workspace({"/dbx/": resource}, mode=MountMode.READ) ``` @@ -162,16 +188,22 @@ import asyncio from mirage import MountMode, Workspace from mirage.resource.databricks_volume import ( + DatabricksProfileTokenProvider, DatabricksVolumeConfig, DatabricksVolumeResource, ) -resource = DatabricksVolumeResource(DatabricksVolumeConfig( - catalog="main", - schema="default", - volume="agent_files", - root_path="/reports", -)) +host = "https://example.cloud.databricks.com" +resource = DatabricksVolumeResource( + DatabricksVolumeConfig( + host=host, + catalog="main", + schema="default", + volume="agent_files", + root_path="/reports", + ), + token_provider=DatabricksProfileTokenProvider(host), +) async def main() -> None: diff --git a/docs/python/setup/databricks.mdx b/docs/python/setup/databricks.mdx index 27c3e7c12..8df42cdbf 100644 --- a/docs/python/setup/databricks.mdx +++ b/docs/python/setup/databricks.mdx @@ -15,54 +15,81 @@ guide. ## Configuration -### Databricks Apps or SDK-default auth +### Static token ```python +import os + from mirage import Workspace, MountMode from mirage.resource.databricks_volume import ( DatabricksVolumeConfig, DatabricksVolumeResource, + StaticTokenProvider, ) config = DatabricksVolumeConfig( + host=os.environ["DATABRICKS_HOST"], catalog="main", schema="default", volume="agent_files", + root_path="/reports", +) +resource = DatabricksVolumeResource( + config, + token_provider=StaticTokenProvider(os.environ["DATABRICKS_TOKEN"]), ) -resource = DatabricksVolumeResource(config=config) ws = Workspace({"/dbx/": resource}, mode=MountMode.READ) ``` -### Explicit host and token +### Profile-based auth ```python import os +from mirage.resource.databricks_volume import DatabricksProfileTokenProvider + +host = os.environ["DATABRICKS_HOST"] config = DatabricksVolumeConfig( + host=host, catalog="main", schema="default", volume="agent_files", - host=os.environ["DATABRICKS_HOST"], - token=os.environ["DATABRICKS_TOKEN"], - root_path="/reports", ) -resource = DatabricksVolumeResource(config=config) +resource = DatabricksVolumeResource( + config, + token_provider=DatabricksProfileTokenProvider(host, profile="DEV"), +) ws = Workspace({"/dbx/": resource}, mode=MountMode.READ) ``` -### Profile-based auth +### Custom token strategy + +Implement `get_token()` when the application owns token caching, refresh, or +request-scoped OBO credentials. It may return a string or an awaitable string. ```python +class AppTokenProvider: + async def get_token(self) -> str: + return await application_credentials.current_token() + + config = DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main", schema="default", volume="agent_files", - profile="DEV", ) -resource = DatabricksVolumeResource(config=config) +resource = DatabricksVolumeResource( + config, + token_provider=AppTokenProvider(), +) ws = Workspace({"/dbx/": resource}, mode=MountMode.READ) ``` +Mirage asks the provider for a token at each Files API operation. The provider +decides whether to return a cached token or refresh it. Mirage does not retry a +rejected token. + ## Config Reference | Field | Required | Default | Description | @@ -70,14 +97,12 @@ ws = Workspace({"/dbx/": resource}, mode=MountMode.READ) | `catalog` | Yes | | Unity Catalog catalog name. | | `schema` | Yes | | Unity Catalog schema name. | | `volume` | Yes | | Unity Catalog volume name. | +| `host` | Yes | | Databricks workspace host. | | `root_path` | No | `/` | Subdirectory inside the volume to expose. | -| `host` | No | | Databricks workspace host. | -| `token` | No | | Databricks personal access token. Redacted in snapshots. | -| `profile` | No | | Databricks SDK profile name. | | `timeout` | No | `30` | Request timeout in seconds. | ## Notes - Supports both read and write mount modes (`MountMode.READ` / `MountMode.WRITE`). -- Auth falls through to the Databricks SDK defaults when `host`, `token`, and `profile` are omitted. +- Credentials are supplied separately through a token provider and are never serialized. - `root_path` is normalized and cannot contain `..`. diff --git a/examples/python/agents/langchain/databricks_volume_deepagent.py b/examples/python/agents/langchain/databricks_volume_deepagent.py index ad128fb16..9b31e9067 100644 --- a/examples/python/agents/langchain/databricks_volume_deepagent.py +++ b/examples/python/agents/langchain/databricks_volume_deepagent.py @@ -21,21 +21,30 @@ from mirage import MountMode, Workspace from mirage.agents.langchain import (LangchainWorkspace, build_system_prompt, extract_text) -from mirage.resource.databricks_volume import (DatabricksVolumeConfig, - DatabricksVolumeResource) +from mirage.resource.databricks_volume import (DatabricksProfileTokenProvider, + DatabricksVolumeConfig, + DatabricksVolumeResource, + StaticTokenProvider) load_dotenv(".env.development") +host = os.environ["DATABRICKS_HOST"] +token = os.environ.get("DATABRICKS_TOKEN") +token_provider = (StaticTokenProvider(token) + if token else DatabricksProfileTokenProvider( + host, + os.environ.get("DATABRICKS_CONFIG_PROFILE", "DEFAULT"), + )) resource = DatabricksVolumeResource( DatabricksVolumeConfig( + host=host, catalog=os.environ["DATABRICKS_VOLUME_CATALOG"], schema=os.environ["DATABRICKS_VOLUME_SCHEMA"], volume=os.environ["DATABRICKS_VOLUME_NAME"], root_path=os.environ.get("DATABRICKS_VOLUME_ROOT_PATH", "/"), - host=os.environ.get("DATABRICKS_HOST"), - token=os.environ.get("DATABRICKS_TOKEN"), - profile=os.environ.get("DATABRICKS_CONFIG_PROFILE"), - )) + ), + token_provider=token_provider, +) ws = Workspace({"/dbx/": resource}, mode=MountMode.READ) diff --git a/examples/python/databricks_volume/databricks_volume.py b/examples/python/databricks_volume/databricks_volume.py index fe9a41e4b..89a1a5c8c 100644 --- a/examples/python/databricks_volume/databricks_volume.py +++ b/examples/python/databricks_volume/databricks_volume.py @@ -18,21 +18,31 @@ from dotenv import load_dotenv from mirage import MountMode, Workspace -from mirage.resource.databricks_volume import (DatabricksVolumeConfig, - DatabricksVolumeResource) +from mirage.resource.databricks_volume import (DatabricksProfileTokenProvider, + DatabricksVolumeConfig, + DatabricksVolumeResource, + StaticTokenProvider) load_dotenv(".env.development") +host = os.environ["DATABRICKS_HOST"] +token = os.environ.get("DATABRICKS_TOKEN") +token_provider = (StaticTokenProvider(token) + if token else DatabricksProfileTokenProvider( + host, + os.environ.get("DATABRICKS_CONFIG_PROFILE", "DEFAULT"), + )) config = DatabricksVolumeConfig( + host=host, catalog=os.environ["DATABRICKS_VOLUME_CATALOG"], schema=os.environ["DATABRICKS_VOLUME_SCHEMA"], volume=os.environ["DATABRICKS_VOLUME_NAME"], root_path=os.environ.get("DATABRICKS_VOLUME_ROOT_PATH", "/"), - host=os.environ.get("DATABRICKS_HOST"), - token=os.environ.get("DATABRICKS_TOKEN"), - profile=os.environ.get("DATABRICKS_CONFIG_PROFILE"), ) -resource = DatabricksVolumeResource(config=config) +resource = DatabricksVolumeResource( + config=config, + token_provider=token_provider, +) async def _run(ws, cmd): diff --git a/python/mirage/accessor/databricks_volume.py b/python/mirage/accessor/databricks_volume.py index 26fc64ecf..38bb80a7c 100644 --- a/python/mirage/accessor/databricks_volume.py +++ b/python/mirage/accessor/databricks_volume.py @@ -12,49 +12,17 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -from typing import Any - from mirage.accessor.base import Accessor +from mirage.core.databricks_volume.client import DatabricksFilesClient from mirage.resource.databricks_volume.config import DatabricksVolumeConfig -try: - from databricks.sdk import WorkspaceClient - from databricks.sdk.config import Config as WorkspaceConfig -except ImportError: - WorkspaceConfig = None - WorkspaceClient = None - class DatabricksVolumeAccessor(Accessor): def __init__( self, config: DatabricksVolumeConfig, - client: Any | None = None, + client: DatabricksFilesClient, ) -> None: self.config = config - self._client = client - - @property - def client(self) -> Any: - if self._client is None: - if WorkspaceClient is None or WorkspaceConfig is None: - raise ImportError("DatabricksVolumeResource requires the " - "'databricks' extra. Install with: " - "pip install mirage-ai[databricks]") - kwargs = { - "host": self.config.host, - "token": self.config.token, - "profile": self.config.profile, - "http_timeout_seconds": self.config.timeout, - } - sdk_config = WorkspaceConfig(**{ - k: v - for k, v in kwargs.items() if v is not None - }) - self._client = WorkspaceClient(config=sdk_config) - return self._client - - @property - def files(self) -> Any: - return self.client.files + self.client = client diff --git a/python/mirage/core/databricks_volume/client.py b/python/mirage/core/databricks_volume/client.py new file mode 100644 index 000000000..8f606696e --- /dev/null +++ b/python/mirage/core/databricks_volume/client.py @@ -0,0 +1,278 @@ +# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= + +import asyncio +import inspect +from contextvars import ContextVar +from io import BytesIO +from typing import Any, BinaryIO, Callable, Protocol, TypeVar +from urllib.parse import quote + +from mirage.resource.databricks_volume.config import DatabricksVolumeConfig +from mirage.resource.databricks_volume.token_provider import TokenProvider + +try: + from databricks.sdk import WorkspaceClient + from databricks.sdk.config import Config as WorkspaceConfig + from databricks.sdk.credentials_provider import CredentialsStrategy +except ImportError: + WorkspaceClient = None + WorkspaceConfig = None + + class CredentialsStrategy: + pass + + +_operation_token: ContextVar[str | None] = ContextVar( + "databricks_volume_operation_token", + default=None, +) +T = TypeVar("T") + + +def _response_contents(response: object) -> object: + if isinstance(response, dict): + return response.get("contents", response) + return getattr(response, "contents", response) + + +def _response_bytes(response: object) -> bytes: + contents = _response_contents(response) + if isinstance(contents, bytes): + return contents + if hasattr(contents, "read"): + return contents.read() + return bytes(contents) + + +def _response_stream(response: object) -> BinaryIO: + contents = _response_contents(response) + if not hasattr(contents, "read"): + raise RuntimeError("Databricks download response has no readable body") + return contents + + +class DatabricksFilesClient(Protocol): + + async def read_bytes( + self, + path: str, + range_header: str | None = None, + ) -> bytes: + ... + + async def open_read(self, path: str) -> "DatabricksReadStream": + ... + + async def get_metadata(self, path: str) -> object: + ... + + async def get_directory_metadata(self, path: str) -> object: + ... + + async def list_directory(self, path: str) -> list[object]: + ... + + async def upload(self, path: str, data: bytes) -> None: + ... + + async def delete(self, path: str) -> None: + ... + + async def create_directory(self, path: str) -> None: + ... + + async def delete_directory(self, path: str) -> None: + ... + + +class DatabricksReadStream(Protocol): + + async def read(self, size: int = -1) -> bytes: + ... + + async def close(self) -> None: + ... + + +class _OperationCredentials: + + def __call__(self) -> dict[str, str]: + token = _operation_token.get() + if token is None: + raise RuntimeError("Databricks request has no operation token") + return {"Authorization": f"Bearer {token}"} + + +class _TokenProviderCredentialsStrategy(CredentialsStrategy): + + def auth_type(self) -> str: + return "mirage-token-provider" + + def __call__(self, config: Any) -> _OperationCredentials: + return _OperationCredentials() + + +class _SdkDatabricksReadStream: + + def __init__( + self, + client: "SdkDatabricksFilesClient", + contents: BinaryIO, + ) -> None: + self._client = client + self._contents = contents + + async def read(self, size: int = -1) -> bytes: + return await self._client._run_sdk(self._contents.read, size) + + async def close(self) -> None: + await asyncio.to_thread(self._contents.close) + + +class SdkDatabricksFilesClient: + + def __init__( + self, + config: DatabricksVolumeConfig, + token_provider: TokenProvider, + ) -> None: + self.config = config + self.token_provider = token_provider + self._workspace: Any | None = None + + @property + def _workspace_client(self) -> Any: + if self._workspace is None: + if WorkspaceClient is None or WorkspaceConfig is None: + raise ImportError("DatabricksVolumeResource requires the " + "'databricks' extra. Install with: " + "pip install mirage-ai[databricks]") + strategy = _TokenProviderCredentialsStrategy() + sdk_config = WorkspaceConfig( + host=self.config.host, + auth_type=strategy.auth_type(), + credentials_strategy=strategy, + http_timeout_seconds=self.config.timeout, + ) + self._workspace = WorkspaceClient(config=sdk_config) + return self._workspace + + async def _resolve_token(self) -> str: + value = await asyncio.to_thread(self.token_provider.get_token) + if inspect.isawaitable(value): + value = await value + if not isinstance(value, str) or not value: + raise ValueError("token provider returned an empty token") + return value + + async def _run_sdk( + self, + fn: Callable[..., T], + *args: object, + **kwargs: object, + ) -> T: + token = await self._resolve_token() + marker = _operation_token.set(token) + try: + return await asyncio.to_thread(fn, *args, **kwargs) + finally: + _operation_token.reset(marker) + + def _read_bytes_sync( + self, + path: str, + range_header: str | None, + ) -> bytes: + if range_header is None: + return _response_bytes(self._workspace_client.files.download(path)) + headers = { + "Accept": "application/octet-stream", + "Range": range_header, + } + workspace_id = getattr(self._workspace_client.config, "workspace_id", + None) + if workspace_id: + headers["X-Databricks-Org-Id"] = workspace_id + response = self._workspace_client.api_client.do( + "GET", + f"/api/2.0/fs/files{quote(path)}", + headers=headers, + response_headers=[ + "content-length", + "content-range", + "accept-ranges", + "content-type", + "last-modified", + ], + raw=True, + ) + return _response_bytes(response) + + def _list_directory_sync(self, path: str) -> list[object]: + entries = self._workspace_client.files.list_directory_contents(path) + return list(entries) + + async def read_bytes( + self, + path: str, + range_header: str | None = None, + ) -> bytes: + return await self._run_sdk(self._read_bytes_sync, path, range_header) + + async def open_read(self, path: str) -> DatabricksReadStream: + response = await self._run_sdk( + self._workspace_client.files.download, + path, + ) + return _SdkDatabricksReadStream(self, _response_stream(response)) + + async def get_metadata(self, path: str) -> object: + return await self._run_sdk( + self._workspace_client.files.get_metadata, + path, + ) + + async def get_directory_metadata(self, path: str) -> object: + return await self._run_sdk( + self._workspace_client.files.get_directory_metadata, + path, + ) + + async def list_directory(self, path: str) -> list[object]: + return await self._run_sdk(self._list_directory_sync, path) + + async def upload(self, path: str, data: bytes) -> None: + await self._run_sdk( + self._workspace_client.files.upload, + path, + BytesIO(data), + overwrite=True, + use_parallel=False, + ) + + async def delete(self, path: str) -> None: + await self._run_sdk(self._workspace_client.files.delete, path) + + async def create_directory(self, path: str) -> None: + await self._run_sdk( + self._workspace_client.files.create_directory, + path, + ) + + async def delete_directory(self, path: str) -> None: + await self._run_sdk( + self._workspace_client.files.delete_directory, + path, + ) diff --git a/python/mirage/core/databricks_volume/copy.py b/python/mirage/core/databricks_volume/copy.py index dd4e6c509..82607d354 100644 --- a/python/mirage/core/databricks_volume/copy.py +++ b/python/mirage/core/databricks_volume/copy.py @@ -12,9 +12,6 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio -from io import BytesIO - from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.cache.index import IndexCacheStore from mirage.core.databricks_volume._helpers import ensure_path_spec @@ -25,53 +22,20 @@ from mirage.types import FileType, PathSpec -def _download_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> bytes: - response = accessor.files.download(remote_path) - contents = getattr(response, "contents", response) - if hasattr(contents, "read"): - return contents.read() - return bytes(contents) - - -def _upload_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, - data: bytes, -) -> None: - accessor.files.upload(remote_path, BytesIO(data), overwrite=True) - - -def _create_directory_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> None: - accessor.files.create_directory(remote_path) - - -def _list_directory_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> list: - return list(accessor.files.list_directory_contents(remote_path)) - - -def _copy_tree_sync( +async def _copy_tree( accessor: DatabricksVolumeAccessor, remote_src: str, remote_dst: str, ) -> None: - _create_directory_sync(accessor, remote_dst) - for entry in _list_directory_sync(accessor, remote_src): + await accessor.client.create_directory(remote_dst) + for entry in await accessor.client.list_directory(remote_src): name = entry.path.rstrip("/").rsplit("/", 1)[-1] child_dst = remote_dst.rstrip("/") + "/" + name if getattr(entry, "is_directory", False): - _copy_tree_sync(accessor, entry.path, child_dst) + await _copy_tree(accessor, entry.path, child_dst) else: - _upload_sync(accessor, child_dst, - _download_sync(accessor, entry.path)) + data = await accessor.client.read_bytes(entry.path) + await accessor.client.upload(child_dst, data) async def copy( @@ -101,8 +65,7 @@ async def copy( # forever. Refuse before any create_directory/upload. raise ValueError(f"cannot copy a directory, '{src.strip_prefix}', " f"into itself, '{dst.strip_prefix}'") - await asyncio.to_thread(_copy_tree_sync, accessor, remote_src, - remote_dst) + await _copy_tree(accessor, remote_src, remote_dst) return if same_path: # Copying a file onto itself would re-upload it; skip. diff --git a/python/mirage/core/databricks_volume/mkdir.py b/python/mirage/core/databricks_volume/mkdir.py index baf25b64a..dc51faa68 100644 --- a/python/mirage/core/databricks_volume/mkdir.py +++ b/python/mirage/core/databricks_volume/mkdir.py @@ -12,8 +12,6 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio - from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.cache.index import IndexCacheStore from mirage.core.databricks_volume._helpers import (ensure_path_spec, @@ -25,13 +23,6 @@ from mirage.types import FileType, PathSpec -def _create_directory_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> None: - accessor.files.create_directory(remote_path) - - async def mkdir( accessor: DatabricksVolumeAccessor, path: PathSpec, @@ -41,7 +32,7 @@ async def mkdir( path = ensure_path_spec(path) remote_path = backend_path(accessor.config, path) if parents: - await asyncio.to_thread(_create_directory_sync, accessor, remote_path) + await accessor.client.create_directory(remote_path) return if await exists(accessor, path): raise FileExistsError(path.strip_prefix) @@ -50,7 +41,7 @@ async def mkdir( if parent_stat.type != FileType.DIRECTORY: raise NotADirectoryError(path.strip_prefix) try: - await asyncio.to_thread(_create_directory_sync, accessor, remote_path) + await accessor.client.create_directory(remote_path) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(path.strip_prefix) from exc diff --git a/python/mirage/core/databricks_volume/read.py b/python/mirage/core/databricks_volume/read.py index 788ab6a0b..30ce3cfbb 100644 --- a/python/mirage/core/databricks_volume/read.py +++ b/python/mirage/core/databricks_volume/read.py @@ -12,9 +12,7 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio import time -from urllib.parse import quote from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.cache.index import IndexCacheStore @@ -24,18 +22,6 @@ from mirage.types import PathSpec -def _read_response_bytes(response) -> bytes: - if isinstance(response, dict): - contents = response.get("contents", response) - else: - contents = getattr(response, "contents", response) - if isinstance(contents, bytes): - return contents - if hasattr(contents, "read"): - return contents.read() - return bytes(contents) - - def _range_header(offset: int, size: int | None) -> str | None: if offset < 0: raise ValueError("offset must be non-negative") @@ -48,37 +34,6 @@ def _range_header(offset: int, size: int | None) -> str | None: return f"bytes={offset}-{offset + size - 1}" -def _download_bytes_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, - range_header: str | None, -) -> bytes: - if range_header is None: - return _read_response_bytes(accessor.files.download(remote_path)) - headers = { - "Accept": "application/octet-stream", - "Range": range_header, - } - cfg = getattr(accessor.client.api_client, "_cfg", None) - workspace_id = getattr(cfg, "workspace_id", None) - if workspace_id: - headers["X-Databricks-Org-Id"] = workspace_id - response = accessor.client.api_client.do( - "GET", - f"/api/2.0/fs/files{quote(remote_path)}", - headers=headers, - response_headers=[ - "content-length", - "content-range", - "accept-ranges", - "content-type", - "last-modified", - ], - raw=True, - ) - return _read_response_bytes(response) - - async def read_bytes( accessor: DatabricksVolumeAccessor, path: PathSpec, @@ -95,9 +50,7 @@ async def read_bytes( record("read", virtual, "databricks_volume", 0, start_ms) return b"" try: - data = await asyncio.to_thread( - _download_bytes_sync, - accessor, + data = await accessor.client.read_bytes( remote_path, _range_header(offset, size), ) diff --git a/python/mirage/core/databricks_volume/readdir.py b/python/mirage/core/databricks_volume/readdir.py index 987fc184d..34fb0afb0 100644 --- a/python/mirage/core/databricks_volume/readdir.py +++ b/python/mirage/core/databricks_volume/readdir.py @@ -12,7 +12,6 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio import logging from mirage.accessor.databricks_volume import DatabricksVolumeAccessor @@ -25,13 +24,6 @@ SCOPE_ERROR = 10_000 -def _list_directory_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> list[object]: - return list(accessor.files.list_directory_contents(remote_path)) - - async def readdir( accessor: DatabricksVolumeAccessor, path: PathSpec, @@ -46,11 +38,7 @@ async def readdir( return listing.entries remote_path = backend_path(accessor.config, list_path) try: - entries = await asyncio.to_thread( - _list_directory_sync, - accessor, - remote_path, - ) + entries = await accessor.client.list_directory(remote_path) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(list_path.strip_prefix) from exc diff --git a/python/mirage/core/databricks_volume/rm.py b/python/mirage/core/databricks_volume/rm.py index 255b7c846..0c904b4ed 100644 --- a/python/mirage/core/databricks_volume/rm.py +++ b/python/mirage/core/databricks_volume/rm.py @@ -12,8 +12,6 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio - from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.cache.index import IndexCacheStore from mirage.core.databricks_volume._helpers import ensure_path_spec @@ -24,51 +22,21 @@ from mirage.types import FileType, PathSpec -def _list_directory_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> list: - return list(accessor.files.list_directory_contents(remote_path)) - - -def _delete_file_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> None: - accessor.files.delete(remote_path) - - -def _delete_directory_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> None: - accessor.files.delete_directory(remote_path) - - -def _remove_tree_recurse( +async def _remove_tree( accessor: DatabricksVolumeAccessor, remote_dir: str, removed: list[str], ) -> None: - for entry in _list_directory_sync(accessor, remote_dir): + for entry in await accessor.client.list_directory(remote_dir): if getattr(entry, "is_directory", False): - _remove_tree_recurse(accessor, entry.path, removed) + await _remove_tree(accessor, entry.path, removed) else: - _delete_file_sync(accessor, entry.path) + await accessor.client.delete(entry.path) removed.append(entry.path) - _delete_directory_sync(accessor, remote_dir) + await accessor.client.delete_directory(remote_dir) removed.append(remote_dir) -def _remove_tree_sync( - accessor: DatabricksVolumeAccessor, - remote_root: str, -) -> list[str]: - removed: list[str] = [] - _remove_tree_recurse(accessor, remote_root, removed) - return removed - - async def rm_recursive( accessor: DatabricksVolumeAccessor, path: PathSpec, @@ -81,8 +49,8 @@ async def rm_recursive( return [path.strip_prefix] remote_root = backend_path(accessor.config, path) try: - removed = await asyncio.to_thread(_remove_tree_sync, accessor, - remote_root) + removed: list[str] = [] + await _remove_tree(accessor, remote_root, removed) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(path.strip_prefix) from exc diff --git a/python/mirage/core/databricks_volume/rmdir.py b/python/mirage/core/databricks_volume/rmdir.py index f1d0cc192..9ba1e072e 100644 --- a/python/mirage/core/databricks_volume/rmdir.py +++ b/python/mirage/core/databricks_volume/rmdir.py @@ -12,8 +12,6 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio - from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.cache.index import IndexCacheStore from mirage.core.databricks_volume._helpers import ensure_path_spec @@ -23,20 +21,6 @@ from mirage.types import FileType, PathSpec -def _list_directory_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> list: - return list(accessor.files.list_directory_contents(remote_path)) - - -def _delete_directory_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> None: - accessor.files.delete_directory(remote_path) - - async def rmdir( accessor: DatabricksVolumeAccessor, path: PathSpec, @@ -48,8 +32,7 @@ async def rmdir( raise NotADirectoryError(path.strip_prefix) remote_path = backend_path(accessor.config, path) try: - entries = await asyncio.to_thread(_list_directory_sync, accessor, - remote_path) + entries = await accessor.client.list_directory(remote_path) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(path.strip_prefix) from exc @@ -57,7 +40,7 @@ async def rmdir( if entries: raise OSError(f"directory not empty: {path.strip_prefix}") try: - await asyncio.to_thread(_delete_directory_sync, accessor, remote_path) + await accessor.client.delete_directory(remote_path) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(path.strip_prefix) from exc diff --git a/python/mirage/core/databricks_volume/stat.py b/python/mirage/core/databricks_volume/stat.py index 62ec1502c..2cbb5bbdf 100644 --- a/python/mirage/core/databricks_volume/stat.py +++ b/python/mirage/core/databricks_volume/stat.py @@ -12,7 +12,6 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio from datetime import datetime, timezone from mirage.accessor.databricks_volume import DatabricksVolumeAccessor @@ -54,8 +53,7 @@ async def _directory_stat_or_raise( path: PathSpec, ) -> FileStat: try: - await asyncio.to_thread(accessor.files.get_directory_metadata, - remote_path) + await accessor.client.get_directory_metadata(remote_path) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(path.strip_prefix) from exc @@ -76,8 +74,7 @@ async def stat( return FileStat(name="/", type=FileType.DIRECTORY) remote_path = backend_path(accessor.config, path) try: - metadata = await asyncio.to_thread(accessor.files.get_metadata, - remote_path) + metadata = await accessor.client.get_metadata(remote_path) except Exception as exc: if is_not_found(exc): return await _directory_stat_or_raise(accessor, remote_path, path) diff --git a/python/mirage/core/databricks_volume/stream.py b/python/mirage/core/databricks_volume/stream.py index cfff4cc5f..26486cde3 100644 --- a/python/mirage/core/databricks_volume/stream.py +++ b/python/mirage/core/databricks_volume/stream.py @@ -12,9 +12,7 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio from collections.abc import AsyncIterator -from typing import BinaryIO from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.cache.index import IndexCacheStore @@ -25,23 +23,6 @@ from mirage.types import PathSpec -def _download_contents(response) -> BinaryIO: - if isinstance(response, dict): - contents = response.get("contents") - else: - contents = getattr(response, "contents", None) - if contents is None: - raise RuntimeError("Databricks download response has no contents") - return contents - - -def _open_download_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> BinaryIO: - return _download_contents(accessor.files.download(remote_path)) - - async def read_stream( accessor: DatabricksVolumeAccessor, path: PathSpec, @@ -56,13 +37,9 @@ async def read_stream( remote_path = backend_path(accessor.config, path) contents = None try: - contents = await asyncio.to_thread( - _open_download_sync, - accessor, - remote_path, - ) + contents = await accessor.client.open_read(remote_path) while True: - chunk = await asyncio.to_thread(contents.read, chunk_size) + chunk = await contents.read(chunk_size) if not chunk: return if rec is not None: @@ -74,7 +51,7 @@ async def read_stream( raise finally: if contents is not None: - await asyncio.to_thread(contents.close) + await contents.close() async def range_read( diff --git a/python/mirage/core/databricks_volume/unlink.py b/python/mirage/core/databricks_volume/unlink.py index d0128995b..b5b50b924 100644 --- a/python/mirage/core/databricks_volume/unlink.py +++ b/python/mirage/core/databricks_volume/unlink.py @@ -12,7 +12,6 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio import time from mirage.accessor.databricks_volume import DatabricksVolumeAccessor @@ -25,13 +24,6 @@ from mirage.types import FileType, PathSpec -def _delete_file_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, -) -> None: - accessor.files.delete(remote_path) - - async def unlink( accessor: DatabricksVolumeAccessor, path: PathSpec, @@ -44,7 +36,7 @@ async def unlink( remote_path = backend_path(accessor.config, path) start_ms = int(time.monotonic() * 1000) try: - await asyncio.to_thread(_delete_file_sync, accessor, remote_path) + await accessor.client.delete(remote_path) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(path.strip_prefix) from exc diff --git a/python/mirage/core/databricks_volume/write.py b/python/mirage/core/databricks_volume/write.py index 4f4865361..f8bb090e5 100644 --- a/python/mirage/core/databricks_volume/write.py +++ b/python/mirage/core/databricks_volume/write.py @@ -12,9 +12,7 @@ # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= -import asyncio import time -from io import BytesIO from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.cache.index import IndexCacheStore @@ -36,20 +34,20 @@ def _is_directory_metadata(metadata: object) -> bool: return str(object_type).lower().endswith("directory") -def _ensure_parent_directory_sync( +async def _ensure_parent_directory( accessor: DatabricksVolumeAccessor, remote_parent: str, virtual_target: str, ) -> None: try: - accessor.files.get_directory_metadata(remote_parent) + await accessor.client.get_directory_metadata(remote_parent) return except Exception as exc: if not is_not_found(exc): raise not_found = exc try: - metadata = accessor.files.get_metadata(remote_parent) + metadata = await accessor.client.get_metadata(remote_parent) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(virtual_target) from not_found @@ -58,14 +56,6 @@ def _ensure_parent_directory_sync( raise NotADirectoryError(virtual_target) -def _upload_bytes_sync( - accessor: DatabricksVolumeAccessor, - remote_path: str, - data: bytes, -) -> None: - accessor.files.upload(remote_path, BytesIO(data), overwrite=True) - - async def write_bytes( accessor: DatabricksVolumeAccessor, path: PathSpec, @@ -77,16 +67,13 @@ async def write_bytes( remote_parent = backend_path(accessor.config, parent) remote_path = backend_path(accessor.config, path) start_ms = int(time.monotonic() * 1000) - # TODO native async client calling HTTP API as databricks sdk is sync - await asyncio.to_thread( - _ensure_parent_directory_sync, + await _ensure_parent_directory( accessor, remote_parent, path.strip_prefix, ) try: - await asyncio.to_thread(_upload_bytes_sync, accessor, remote_path, - data) + await accessor.client.upload(remote_path, data) except Exception as exc: if is_not_found(exc): raise FileNotFoundError(path.strip_prefix) from exc diff --git a/python/mirage/resource/databricks_volume/__init__.py b/python/mirage/resource/databricks_volume/__init__.py index 7ad99b42d..f6660db2b 100644 --- a/python/mirage/resource/databricks_volume/__init__.py +++ b/python/mirage/resource/databricks_volume/__init__.py @@ -13,8 +13,16 @@ # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= from mirage.resource.databricks_volume.config import DatabricksVolumeConfig +from mirage.resource.databricks_volume.token_provider import ( + DatabricksProfileTokenProvider, StaticTokenProvider, TokenProvider) -__all__ = ["DatabricksVolumeConfig", "DatabricksVolumeResource"] +__all__ = [ + "DatabricksProfileTokenProvider", + "DatabricksVolumeConfig", + "DatabricksVolumeResource", + "StaticTokenProvider", + "TokenProvider", +] def __getattr__(name: str): diff --git a/python/mirage/resource/databricks_volume/config.py b/python/mirage/resource/databricks_volume/config.py index f53f28d1c..491b467b8 100644 --- a/python/mirage/resource/databricks_volume/config.py +++ b/python/mirage/resource/databricks_volume/config.py @@ -27,10 +27,8 @@ class DatabricksVolumeConfig(BaseModel): # internal name avoids warning overwriting BaseModel.schema. schema_name: str = Field(alias="schema") volume: str + host: str root_path: str = "/" - host: str | None = None - token: str | None = None - profile: str | None = None timeout: int = 30 @field_validator("catalog", "schema_name", "volume") diff --git a/python/mirage/resource/databricks_volume/databricks_volume.py b/python/mirage/resource/databricks_volume/databricks_volume.py index 4b64f2bc4..9feb9c29d 100644 --- a/python/mirage/resource/databricks_volume/databricks_volume.py +++ b/python/mirage/resource/databricks_volume/databricks_volume.py @@ -13,11 +13,13 @@ # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= import dataclasses -from typing import Any +from typing import Any, Self from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.commands.builtin.databricks_volume import \ COMMANDS as DATABRICKS_VOLUME_COMMANDS +from mirage.core.databricks_volume.client import (DatabricksFilesClient, + SdkDatabricksFilesClient) from mirage.core.databricks_volume.copy import copy from mirage.core.databricks_volume.create import create from mirage.core.databricks_volume.exists import exists @@ -36,6 +38,7 @@ from mirage.resource.base import BaseResource from mirage.resource.databricks_volume.config import DatabricksVolumeConfig from mirage.resource.databricks_volume.prompt import PROMPT +from mirage.resource.databricks_volume.token_provider import TokenProvider from mirage.types import PathSpec, ResourceName _DATABRICKS_VOLUME_OPS = { @@ -65,11 +68,31 @@ class DatabricksVolumeResource(BaseResource): def __init__( self, config: DatabricksVolumeConfig, - client: Any | None = None, + token_provider: TokenProvider, + ) -> None: + self._initialize( + config, + SdkDatabricksFilesClient(config, token_provider), + ) + + @classmethod + def _from_files_client( + cls, + config: DatabricksVolumeConfig, + files_client: DatabricksFilesClient, + ) -> Self: + resource = cls.__new__(cls) + resource._initialize(config, files_client) + return resource + + def _initialize( + self, + config: DatabricksVolumeConfig, + files_client: DatabricksFilesClient, ) -> None: super().__init__() self.config = config - self.accessor = DatabricksVolumeAccessor(self.config, client) + self.accessor = DatabricksVolumeAccessor(self.config, files_client) for fn in DATABRICKS_VOLUME_COMMANDS: self.register(fn) @@ -86,16 +109,12 @@ async def resolve_glob(self, paths, prefix: str = ""): return await _resolve_glob(self.accessor, paths, self._index) def get_state(self) -> dict: - redacted = ["token"] - cfg = self.config.model_dump() - for field in redacted: - if cfg.get(field) is not None: - cfg[field] = "" return { "type": self.name, + # Token providers are runtime-only and cannot be reconstructed + # from the serialized volume location. "needs_override": True, - "redacted_fields": redacted, - "config": cfg, + "config": self.config.model_dump(), } def load_state(self, state: dict) -> None: diff --git a/python/mirage/resource/databricks_volume/token_provider.py b/python/mirage/resource/databricks_volume/token_provider.py new file mode 100644 index 000000000..10d98563e --- /dev/null +++ b/python/mirage/resource/databricks_volume/token_provider.py @@ -0,0 +1,64 @@ +# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= + +from collections.abc import Awaitable +from threading import Lock +from typing import Any, Protocol + +try: + from databricks.sdk.config import Config as WorkspaceConfig +except ImportError: + WorkspaceConfig = None + + +class TokenProvider(Protocol): + + def get_token(self) -> str | Awaitable[str]: + ... + + +class StaticTokenProvider: + + def __init__(self, token: str) -> None: + self._token = token + + def get_token(self) -> str: + return self._token + + +class DatabricksProfileTokenProvider: + + def __init__(self, host: str, profile: str = "DEFAULT") -> None: + self._host = host + self._profile = profile + self._config: Any | None = None + self._config_lock = Lock() + + def get_token(self) -> str: + if WorkspaceConfig is None: + raise ImportError("DatabricksProfileTokenProvider requires the " + "'databricks' extra. Install with: " + "pip install mirage-ai[databricks]") + with self._config_lock: + if self._config is None: + self._config = WorkspaceConfig( + host=self._host, + profile=self._profile, + ) + authorization = self._config.authenticate().get("Authorization", "") + scheme, separator, token = authorization.partition(" ") + if not separator or scheme.lower() != "bearer" or not token: + raise ValueError( + "Databricks profile did not provide a bearer token") + return token diff --git a/python/mirage/workspace/snapshot/state.py b/python/mirage/workspace/snapshot/state.py index c495241d6..7db25b6b1 100644 --- a/python/mirage/workspace/snapshot/state.py +++ b/python/mirage/workspace/snapshot/state.py @@ -324,6 +324,10 @@ def _construct_resource(mount_state: dict): def requires_resource_override(mount_state: dict) -> bool: resource_state = mount_state[MountKey.RESOURCE_STATE] + # Resources can explicitly omit runtime-only constructor dependencies + # from snapshot state even when their serialized config has no secrets. + if resource_state.get("needs_override") is True: + return True config = resource_state.get(ResourceStateKey.CONFIG) config_cls = _config_class_for(_resource_class_for(mount_state)) return has_redacted_secret(config, config_cls) diff --git a/python/tests/core/databricks_volume/conftest.py b/python/tests/core/databricks_volume/conftest.py index 2605ef8a6..5595d4f24 100644 --- a/python/tests/core/databricks_volume/conftest.py +++ b/python/tests/core/databricks_volume/conftest.py @@ -1,7 +1,6 @@ import posixpath from io import BytesIO from types import SimpleNamespace -from urllib.parse import unquote import pytest @@ -15,16 +14,6 @@ class NotFoundError(Exception): status_code = 404 -class ToThreadRecorder: - - def __init__(self) -> None: - self.calls = [] - - async def __call__(self, fn, *args, **kwargs): - self.calls.append((fn, args, kwargs)) - return fn(*args, **kwargs) - - class FakeDownload: def __init__(self, data: bytes) -> None: @@ -152,71 +141,70 @@ def _apply_range_header(data: bytes, range_header: str) -> bytes: return data[start:end] -class FakeApiClient: +class FakeDatabricksFilesClient: def __init__(self, files: FakeFiles) -> None: self.files = files - self.do_calls: list[dict[str, object]] = [] + self.read_calls: list[tuple[str, str | None]] = [] - def do( + async def read_bytes( self, - method: str, - path: str | None = None, - url: str | None = None, - query: dict | None = None, - headers: dict | None = None, - body: dict | None = None, - raw: bool = False, - files: object = None, - data: object = None, - auth: object = None, - response_headers: list[str] | None = None, - ) -> dict: - call = { - "method": method, - "path": path, - "url": url, - "query": query, - "headers": headers or {}, - "body": body, - "raw": raw, - "files": files, - "data": data, - "auth": auth, - "response_headers": response_headers, - } - self.do_calls.append(call) - if method != "GET" or path is None: - raise ValueError(f"unsupported fake API call: {method} {path}") - remote_path = unquote(path.removeprefix("/api/2.0/fs/files")) - if remote_path not in self.files.downloads: - raise NotFoundError(remote_path) - payload = self.files.downloads[remote_path] - range_header = (headers or {}).get("Range") + path: str, + range_header: str | None = None, + ) -> bytes: + self.read_calls.append((path, range_header)) + response = self.files.download(path) + payload = response.contents.read() if range_header is not None: payload = _apply_range_header(payload, range_header) - return { - "contents": BytesIO(payload), - "content-length": str(len(payload)), - "accept-ranges": "bytes", - } + return payload -class FakeClient: + async def open_read(self, path: str): + return FakeReadStream(self.files.download(path).contents) - def __init__(self, files: FakeFiles) -> None: - self.files = files - self.api_client = FakeApiClient(files) + async def get_metadata(self, path: str) -> object: + return self.files.get_metadata(path) + + async def get_directory_metadata(self, path: str) -> object: + return self.files.get_directory_metadata(path) + + async def list_directory(self, path: str) -> list[object]: + return list(self.files.list_directory_contents(path)) + + async def upload(self, path: str, data: bytes) -> None: + self.files.upload(path, BytesIO(data), overwrite=True) + + async def delete(self, path: str) -> None: + self.files.delete(path) + + async def create_directory(self, path: str) -> None: + self.files.create_directory(path) + + async def delete_directory(self, path: str) -> None: + self.files.delete_directory(path) + + +class FakeReadStream: + + def __init__(self, contents) -> None: + self.contents = contents + + async def read(self, size: int = -1) -> bytes: + return self.contents.read(size) + + async def close(self) -> None: + self.contents.close() @pytest.fixture def databricks_config() -> DatabricksVolumeConfig: return DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main", schema="default", volume="agent_files", root_path="/root", - token="secret", ) @@ -235,7 +223,10 @@ def accessor( databricks_config: DatabricksVolumeConfig, files: FakeFiles, ) -> DatabricksVolumeAccessor: - return DatabricksVolumeAccessor(databricks_config, FakeClient(files)) + return DatabricksVolumeAccessor( + databricks_config, + FakeDatabricksFilesClient(files), + ) @pytest.fixture diff --git a/python/tests/core/databricks_volume/test_client.py b/python/tests/core/databricks_volume/test_client.py new file mode 100644 index 000000000..405b5eaff --- /dev/null +++ b/python/tests/core/databricks_volume/test_client.py @@ -0,0 +1,304 @@ +import asyncio +from contextvars import ContextVar +from io import BytesIO +from types import SimpleNamespace + +import pytest + +from mirage.core.databricks_volume import client as client_module +from mirage.core.databricks_volume.client import SdkDatabricksFilesClient +from mirage.resource.databricks_volume import (DatabricksVolumeConfig, + StaticTokenProvider) + +current_token: ContextVar[str] = ContextVar("current_token") + + +class AsyncRotatingProvider: + + def __init__(self, tokens: list[str]) -> None: + self._tokens = iter(tokens) + + async def get_token(self) -> str: + await asyncio.sleep(0) + return next(self._tokens) + + +class ContextTokenProvider: + + def get_token(self) -> str: + return current_token.get() + + +class FakeWorkspaceConfig: + calls: list[dict] = [] + + def __init__(self, **kwargs) -> None: + self.calls.append(kwargs) + self._header_factory = kwargs["credentials_strategy"](self) + + def authenticate(self) -> dict[str, str]: + return self._header_factory() + + +class FakeSdkFiles: + + def __init__(self, config: FakeWorkspaceConfig) -> None: + self._config = config + self.auth_calls: list[tuple[str, str]] = [] + self.calls: list[tuple] = [] + self.uploaded: bytes | None = None + + def _authorize(self, operation: str) -> None: + authorization = self._config.authenticate()["Authorization"] + self.auth_calls.append((operation, authorization)) + + def download(self, path: str) -> object: + self._authorize(path) + self.calls.append(("download", path)) + return SimpleNamespace(contents=FakeSdkStream(self._config, + b"payload"), ) + + def get_metadata(self, path: str) -> object: + self._authorize(path) + self.calls.append(("get_metadata", path)) + return SimpleNamespace( + path=path, + is_directory=False, + file_size=1, + modification_time=2, + ) + + def get_directory_metadata(self, path: str) -> None: + self._authorize(path) + self.calls.append(("get_directory_metadata", path)) + + def list_directory_contents(self, path: str): + self.calls.append(("list_directory_contents", path)) + return self._directory_entries(path) + + def _directory_entries(self, path: str): + self._authorize(path) + yield SimpleNamespace( + path=f"{path}/report.md", + is_directory=False, + file_size=7, + ) + + def upload( + self, + path: str, + contents, + *, + overwrite: bool = False, + use_parallel: bool = True, + ) -> None: + self._authorize(path) + self.uploaded = contents.read() + self.calls.append(("upload", path, overwrite, use_parallel)) + + def delete(self, path: str) -> None: + self._authorize(path) + self.calls.append(("delete", path)) + + def create_directory(self, path: str) -> None: + self._authorize(path) + self.calls.append(("create_directory", path)) + + def delete_directory(self, path: str) -> None: + self._authorize(path) + self.calls.append(("delete_directory", path)) + + +class FakeSdkStream: + + def __init__(self, config: FakeWorkspaceConfig, data: bytes) -> None: + self._config = config + self._contents = BytesIO(data) + + def read(self, size: int = -1) -> bytes: + self._config.authenticate() + return self._contents.read(size) + + def close(self) -> None: + self._contents.close() + + +class FakeApiClient: + + def __init__(self, config: FakeWorkspaceConfig) -> None: + self._config = config + self.calls: list[dict] = [] + + def do(self, method: str, path: str, **kwargs) -> dict: + authorization = self._config.authenticate()["Authorization"] + self.calls.append({ + "method": method, + "path": path, + "authorization": authorization, + **kwargs, + }) + return { + "contents": BytesIO(b"ayl"), + "content-length": "3", + "content-range": "bytes 1-3/7", + "accept-ranges": "bytes", + } + + +class FakeWorkspaceClient: + calls: list[dict] = [] + + def __init__(self, **kwargs) -> None: + self.calls.append(kwargs) + self.config = kwargs["config"] + self.files = FakeSdkFiles(self.config) + self.api_client = FakeApiClient(self.config) + + +def _config() -> DatabricksVolumeConfig: + return DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", + catalog="main", + schema="default", + volume="documents", + timeout=17, + ) + + +async def _get_metadata_with_token( + client: SdkDatabricksFilesClient, + path: str, + token: str, +) -> None: + marker = current_token.set(token) + try: + await client.get_metadata(path) + finally: + current_token.reset(marker) + + +@pytest.fixture(autouse=True) +def fake_sdk(monkeypatch): + FakeWorkspaceConfig.calls = [] + FakeWorkspaceClient.calls = [] + monkeypatch.setattr( + client_module, + "WorkspaceConfig", + FakeWorkspaceConfig, + raising=False, + ) + monkeypatch.setattr( + client_module, + "WorkspaceClient", + FakeWorkspaceClient, + raising=False, + ) + + +@pytest.mark.asyncio +async def test_sdk_client_resolves_async_provider_for_each_operation(): + client = SdkDatabricksFilesClient( + _config(), + AsyncRotatingProvider(["token-a", "token-b"]), + ) + + await client.get_metadata("/a") + await client.get_metadata("/b") + + sdk_config = FakeWorkspaceConfig.calls[0] + assert sdk_config["host"] == "https://example.cloud.databricks.com" + assert sdk_config["http_timeout_seconds"] == 17 + assert sdk_config["auth_type"] == "mirage-token-provider" + assert "token" not in sdk_config + assert "profile" not in sdk_config + assert client._workspace_client.files.auth_calls == [ + ("/a", "Bearer token-a"), + ("/b", "Bearer token-b"), + ] + + +@pytest.mark.asyncio +async def test_sdk_client_keeps_concurrent_operation_tokens_isolated(): + client = SdkDatabricksFilesClient(_config(), ContextTokenProvider()) + + await asyncio.gather( + _get_metadata_with_token(client, "/a", "token-a"), + _get_metadata_with_token(client, "/b", "token-b"), + ) + + sdk_files = client._workspace_client.files + assert sorted(sdk_files.auth_calls) == [ + ("/a", "Bearer token-a"), + ("/b", "Bearer token-b"), + ] + + +@pytest.mark.asyncio +async def test_sdk_client_maps_files_api_operations(): + client = SdkDatabricksFilesClient( + _config(), + StaticTokenProvider("token"), + ) + + assert await client.read_bytes("/file") == b"payload" + assert await client.read_bytes("/file", "bytes=1-3") == b"ayl" + stream = await client.open_read("/file") + assert await stream.read() == b"payload" + await stream.close() + assert (await client.get_metadata("/file")).file_size == 1 + assert await client.get_directory_metadata("/dir") is None + entries = await client.list_directory("/dir") + assert [entry.path for entry in entries] == ["/dir/report.md"] + await client.upload("/new", b"new") + await client.delete("/file") + await client.create_directory("/new-dir") + await client.delete_directory("/old-dir") + + workspace = client._workspace_client + assert workspace.files.uploaded == b"new" + assert ("upload", "/new", True, False) in workspace.files.calls + assert workspace.api_client.calls == [{ + "method": + "GET", + "path": + "/api/2.0/fs/files/file", + "authorization": + "Bearer token", + "headers": { + "Accept": "application/octet-stream", + "Range": "bytes=1-3", + }, + "response_headers": [ + "content-length", + "content-range", + "accept-ranges", + "content-type", + "last-modified", + ], + "raw": + True, + }] + + +@pytest.mark.asyncio +async def test_sdk_client_rejects_empty_token(): + client = SdkDatabricksFilesClient( + _config(), + StaticTokenProvider(""), + ) + + with pytest.raises(ValueError, match="empty token"): + await client.get_metadata("/file") + + +@pytest.mark.asyncio +async def test_sdk_stream_read_reestablishes_operation_token(): + client = SdkDatabricksFilesClient( + _config(), + AsyncRotatingProvider(["open-token", "read-token"]), + ) + + stream = await client.open_read("/file") + + assert await stream.read() == b"payload" + await stream.close() diff --git a/python/tests/core/databricks_volume/test_head.py b/python/tests/core/databricks_volume/test_head.py index bac7957a5..7e5d24d29 100644 --- a/python/tests/core/databricks_volume/test_head.py +++ b/python/tests/core/databricks_volume/test_head.py @@ -20,7 +20,5 @@ async def test_head_bytes_mode_uses_single_small_range_request( source, _io = await head(accessor, [path], c="3") assert await collect_bytes(source) == b"abc" - assert files.download_calls == [] - assert len(accessor.client.api_client.do_calls) == 1 - assert accessor.client.api_client.do_calls[0]["headers"]["Range"] == ( - "bytes=0-2") + assert accessor.client.read_calls == [(f"{remote_root}/reports/latest.md", + "bytes=0-2")] diff --git a/python/tests/core/databricks_volume/test_path.py b/python/tests/core/databricks_volume/test_path.py index b223bda6f..e288a2fc8 100644 --- a/python/tests/core/databricks_volume/test_path.py +++ b/python/tests/core/databricks_volume/test_path.py @@ -42,6 +42,7 @@ def test_backend_path_rejects_escape_above_configured_root(databricks_config): def test_config_rejects_parent_segments_in_root_path(): with pytest.raises(ValidationError): DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main", schema="default", volume="agent_files", diff --git a/python/tests/core/databricks_volume/test_read.py b/python/tests/core/databricks_volume/test_read.py index 9267ceb91..116e9bf6b 100644 --- a/python/tests/core/databricks_volume/test_read.py +++ b/python/tests/core/databricks_volume/test_read.py @@ -1,12 +1,8 @@ -import asyncio - import pytest from mirage.core.databricks_volume.read import read_bytes from mirage.types import PathSpec -from .conftest import ToThreadRecorder - @pytest.mark.asyncio async def test_read_file(accessor, files, remote_root): @@ -32,24 +28,6 @@ async def test_read_slice(accessor, files, remote_root): assert result == b"bcd" -@pytest.mark.asyncio -async def test_read_file_runs_blocking_download_off_event_loop( - accessor, - files, - remote_root, - monkeypatch, -): - to_thread = ToThreadRecorder() - monkeypatch.setattr(asyncio, "to_thread", to_thread) - files.downloads[f"{remote_root}/reports/latest.md"] = b"hello" - path = PathSpec.from_str_path("/volume/reports/latest.md", "/volume") - - result = await read_bytes(accessor, path) - - assert result == b"hello" - assert len(to_thread.calls) == 1 - - @pytest.mark.asyncio async def test_read_slice_uses_databricks_range_request( accessor, @@ -62,10 +40,8 @@ async def test_read_slice_uses_databricks_range_request( result = await read_bytes(accessor, path, offset=1, size=3) assert result == b"bcd" - assert files.download_calls == [] - assert accessor.client.api_client.do_calls[0]["headers"]["Range"] == ( - "bytes=1-3") - assert accessor.client.api_client.do_calls[0]["raw"] is True + assert accessor.client.read_calls == [(f"{remote_root}/reports/latest.md", + "bytes=1-3")] @pytest.mark.asyncio @@ -80,9 +56,8 @@ async def test_read_from_offset_uses_open_ended_range( result = await read_bytes(accessor, path, offset=3) assert result == b"def" - assert files.download_calls == [] - assert accessor.client.api_client.do_calls[0]["headers"]["Range"] == ( - "bytes=3-") + assert accessor.client.read_calls == [(f"{remote_root}/reports/latest.md", + "bytes=3-")] @pytest.mark.asyncio @@ -98,4 +73,4 @@ async def test_read_zero_size_returns_empty_without_network( assert result == b"" assert files.download_calls == [] - assert accessor.client.api_client.do_calls == [] + assert accessor.client.read_calls == [] diff --git a/python/tests/core/databricks_volume/test_readdir.py b/python/tests/core/databricks_volume/test_readdir.py index 16f79f318..78d29945f 100644 --- a/python/tests/core/databricks_volume/test_readdir.py +++ b/python/tests/core/databricks_volume/test_readdir.py @@ -1,11 +1,9 @@ -import asyncio - import pytest from mirage.core.databricks_volume.readdir import readdir from mirage.types import PathSpec -from .conftest import ToThreadRecorder, directory_entry, file_entry +from .conftest import directory_entry, file_entry @pytest.mark.asyncio @@ -47,24 +45,3 @@ async def test_readdir_missing_directory_raises(accessor, index): path = PathSpec.from_str_path("/volume/missing", "/volume") with pytest.raises(FileNotFoundError): await readdir(accessor, path, index) - - -@pytest.mark.asyncio -async def test_readdir_runs_blocking_list_off_event_loop( - accessor, - files, - index, - remote_root, - monkeypatch, -): - to_thread = ToThreadRecorder() - monkeypatch.setattr(asyncio, "to_thread", to_thread) - files.directories[f"{remote_root}/reports"] = [ - file_entry(f"{remote_root}/reports/latest.md", size=6), - ] - path = PathSpec.from_str_path("/volume/reports", "/volume") - - result = await readdir(accessor, path, index) - - assert result == ["/volume/reports/latest.md"] - assert len(to_thread.calls) == 1 diff --git a/python/tests/core/databricks_volume/test_stat.py b/python/tests/core/databricks_volume/test_stat.py index 7597df67d..8f5e1c71f 100644 --- a/python/tests/core/databricks_volume/test_stat.py +++ b/python/tests/core/databricks_volume/test_stat.py @@ -1,11 +1,9 @@ -import asyncio - import pytest from mirage.core.databricks_volume.stat import _name_from_backend_path, stat from mirage.types import FileType, PathSpec -from .conftest import ToThreadRecorder, file_metadata +from .conftest import file_metadata def test_name_from_backend_path_file(): @@ -110,39 +108,3 @@ async def test_stat_directory_metadata_error_propagates( await stat(accessor, path) assert files.get_metadata_calls == [remote_path] assert files.get_directory_metadata_calls == [remote_path] - - -@pytest.mark.asyncio -async def test_stat_runs_blocking_metadata_off_event_loop( - accessor, - files, - remote_root, - monkeypatch, -): - to_thread = ToThreadRecorder() - monkeypatch.setattr(asyncio, "to_thread", to_thread) - files.metadata[f"{remote_root}/reports/latest.md"] = file_metadata(size=6) - path = PathSpec.from_str_path("/volume/reports/latest.md", "/volume") - - result = await stat(accessor, path) - - assert result.name == "latest.md" - assert len(to_thread.calls) == 1 - - -@pytest.mark.asyncio -async def test_stat_directory_fallback_runs_off_event_loop( - accessor, - files, - remote_root, - monkeypatch, -): - to_thread = ToThreadRecorder() - monkeypatch.setattr(asyncio, "to_thread", to_thread) - files.directory_metadata.add(f"{remote_root}/reports") - path = PathSpec.from_str_path("/volume/reports", "/volume") - - result = await stat(accessor, path) - - assert result.type == FileType.DIRECTORY - assert len(to_thread.calls) == 2 diff --git a/python/tests/core/databricks_volume/test_stream.py b/python/tests/core/databricks_volume/test_stream.py index e71defdd6..497ceed3d 100644 --- a/python/tests/core/databricks_volume/test_stream.py +++ b/python/tests/core/databricks_volume/test_stream.py @@ -12,7 +12,7 @@ def __init__(self, data: bytes) -> None: self.read_sizes: list[int] = [] self.closed = False - def read(self, size: int = -1) -> bytes: + async def read(self, size: int = -1) -> bytes: self.read_sizes.append(size) if size < 0: size = len(self.data) - self.offset @@ -20,25 +20,19 @@ def read(self, size: int = -1) -> bytes: self.offset += len(chunk) return chunk - def close(self) -> None: + async def close(self) -> None: self.closed = True -class TrackingDownload: +class TrackingClient: def __init__(self, contents: TrackingContents) -> None: self.contents = contents + self.open_read_calls: list[str] = [] - -class TrackingFiles: - - def __init__(self, contents: TrackingContents) -> None: - self.contents = contents - self.download_calls: list[str] = [] - - def download(self, path: str) -> TrackingDownload: - self.download_calls.append(path) - return TrackingDownload(self.contents) + async def open_read(self, path: str) -> TrackingContents: + self.open_read_calls.append(path) + return self.contents @pytest.mark.asyncio @@ -51,7 +45,7 @@ async def test_read_stream_chunks_file(accessor, files, remote_root): assert chunks == [b"ab", b"cd", b"ef"] # Streaming should use one download body, not one Range GET per chunk. assert files.download_calls == [f"{remote_root}/reports/latest.md"] - assert accessor.client.api_client.do_calls == [] + assert accessor.client.read_calls == [] @pytest.mark.asyncio @@ -74,10 +68,8 @@ async def test_range_read_uses_single_databricks_range_request( result = await range_read(accessor, path, 1, 4) assert result == b"bcd" - assert files.download_calls == [] - assert len(accessor.client.api_client.do_calls) == 1 - assert accessor.client.api_client.do_calls[0]["headers"]["Range"] == ( - "bytes=1-3") + assert accessor.client.read_calls == [(f"{remote_root}/reports/latest.md", + "bytes=1-3")] @pytest.mark.asyncio @@ -86,8 +78,8 @@ async def test_read_stream_reads_single_download_body_in_chunks( remote_root, ): contents = TrackingContents(b"abcdef") - tracking_files = TrackingFiles(contents) - accessor.client.files = tracking_files + tracking_client = TrackingClient(contents) + accessor.client = tracking_client path = PathSpec.from_str_path("/volume/reports/latest.md", "/volume") stream = read_stream(accessor, path, chunk_size=2) @@ -96,7 +88,7 @@ async def test_read_stream_reads_single_download_body_in_chunks( assert first == b"ab" assert second == b"cd" - assert tracking_files.download_calls == [ + assert tracking_client.open_read_calls == [ f"{remote_root}/reports/latest.md" ] assert contents.read_sizes == [2, 2] diff --git a/python/tests/resource/databricks_volume/test_accessor.py b/python/tests/resource/databricks_volume/test_accessor.py index bc7a49d6e..e32f0813a 100644 --- a/python/tests/resource/databricks_volume/test_accessor.py +++ b/python/tests/resource/databricks_volume/test_accessor.py @@ -1,46 +1,22 @@ -from mirage.accessor import databricks_volume as accessor_module from mirage.accessor.databricks_volume import DatabricksVolumeAccessor from mirage.resource.databricks_volume import DatabricksVolumeConfig -class FakeWorkspaceClient: - calls: list[dict] = [] +class FakeFilesClient: + pass - def __init__(self, **kwargs) -> None: - self.calls.append(kwargs) - self.files = object() - -class FakeWorkspaceConfig: - - def __init__(self, **kwargs) -> None: - for key, value in kwargs.items(): - setattr(self, key, value) - - -def test_accessor_passes_timeout_to_workspace_client(monkeypatch): - FakeWorkspaceClient.calls = [] - monkeypatch.setattr( - accessor_module, - "WorkspaceClient", - FakeWorkspaceClient, - ) - monkeypatch.setattr( - accessor_module, - "WorkspaceConfig", - FakeWorkspaceConfig, - ) +def test_accessor_stores_config_and_files_client(): config = DatabricksVolumeConfig( catalog="main", schema="default", volume="agent_files", host="https://example.cloud.databricks.com", - token="secret", timeout=17, ) - accessor = DatabricksVolumeAccessor(config) - assert accessor.files is not None - sdk_config = FakeWorkspaceClient.calls[0]["config"] - assert sdk_config.host == "https://example.cloud.databricks.com" - assert sdk_config.token == "secret" - assert sdk_config.http_timeout_seconds == 17 + client = FakeFilesClient() + + accessor = DatabricksVolumeAccessor(config, client) + + assert accessor.config is config + assert accessor.client is client diff --git a/python/tests/resource/databricks_volume/test_databricks_volume.py b/python/tests/resource/databricks_volume/test_databricks_volume.py index fcf7756f6..33cf98bb2 100644 --- a/python/tests/resource/databricks_volume/test_databricks_volume.py +++ b/python/tests/resource/databricks_volume/test_databricks_volume.py @@ -15,7 +15,6 @@ import posixpath from io import BytesIO from types import SimpleNamespace -from urllib.parse import unquote import pytest from pydantic import ValidationError @@ -24,7 +23,8 @@ from mirage.cache.index import IndexEntry, LookupStatus from mirage.core.databricks_volume.path import backend_path from mirage.resource.databricks_volume import (DatabricksVolumeConfig, - DatabricksVolumeResource) + DatabricksVolumeResource, + StaticTokenProvider) from mirage.types import PathSpec, ResourceName @@ -155,58 +155,69 @@ def _apply_range_header(data: bytes, range_header: str) -> bytes: return data[start:end] -class FakeApiClient: +class FakeDatabricksFilesClient: def __init__(self, files: FakeFiles) -> None: self.files = files - def do( + async def read_bytes( self, - method: str, - path: str | None = None, - url: str | None = None, - query: dict | None = None, - headers: dict | None = None, - body: dict | None = None, - raw: bool = False, - files: object = None, - data: object = None, - auth: object = None, - response_headers: list[str] | None = None, - ) -> dict: - if method != "GET" or path is None: - raise ValueError(f"unsupported fake API call: {method} {path}") - remote_path = unquote(path.removeprefix("/api/2.0/fs/files")) - if remote_path not in self.files.downloads: - raise NotFoundError(remote_path) - payload = self.files.downloads[remote_path] - range_header = (headers or {}).get("Range") + path: str, + range_header: str | None = None, + ) -> bytes: + response = self.files.download(path) + payload = response.contents.read() if range_header is not None: payload = _apply_range_header(payload, range_header) - return { - "contents": BytesIO(payload), - "content-length": str(len(payload)), - "accept-ranges": "bytes", - } + return payload + async def open_read(self, path: str): + return FakeReadStream(self.files.download(path).contents) -class FakeClient: + async def get_metadata(self, path: str) -> object: + return self.files.get_metadata(path) - def __init__(self, files: FakeFiles) -> None: - self.files = files - self.api_client = FakeApiClient(files) + async def get_directory_metadata(self, path: str) -> object: + return self.files.get_directory_metadata(path) + + async def list_directory(self, path: str) -> list[object]: + return list(self.files.list_directory_contents(path)) + + async def upload(self, path: str, data: bytes) -> None: + self.files.upload(path, BytesIO(data), overwrite=True) + + async def delete(self, path: str) -> None: + self.files.delete(path) + + async def create_directory(self, path: str) -> None: + self.files.create_directory(path) + + async def delete_directory(self, path: str) -> None: + self.files.delete_directory(path) + + +class FakeReadStream: + + def __init__(self, contents) -> None: + self.contents = contents + + async def read(self, size: int = -1) -> bytes: + return self.contents.read(size) + + async def close(self) -> None: + self.contents.close() def make_resource(files: FakeFiles) -> DatabricksVolumeResource: - return DatabricksVolumeResource( + return DatabricksVolumeResource._from_files_client( DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main", schema="default", volume="agent_files", root_path="/root", - token="secret", ), - client=FakeClient(files), + FakeDatabricksFilesClient(files), ) @@ -234,6 +245,7 @@ def seed_file(files: FakeFiles, path: str, data: bytes) -> None: def test_config_validation_and_normalization(): config = DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main", schema="default", volume="agent_files", @@ -242,14 +254,66 @@ def test_config_validation_and_normalization(): assert config.root_path == "/nested/path" with pytest.raises(ValidationError): DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main/other", schema="default", volume="agent_files", ) +def test_config_requires_host(): + with pytest.raises(ValidationError): + DatabricksVolumeConfig( + catalog="main", + schema="default", + volume="agent_files", + ) + + +def test_config_contains_no_credentials(): + config = DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", + catalog="main", + schema="default", + volume="agent_files", + ) + + assert "token" not in config.model_fields_set + assert not hasattr(config, "token") + assert not hasattr(config, "profile") + + +def test_resource_accepts_token_provider_and_serializes_location_only(): + config = DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", + catalog="main", + schema="default", + volume="agent_files", + ) + resource = DatabricksVolumeResource( + config, + token_provider=StaticTokenProvider("secret"), + ) + + state = resource.get_state() + + assert state == { + "type": ResourceName.DATABRICKS_VOLUME, + "needs_override": True, + "config": { + "host": "https://example.cloud.databricks.com", + "catalog": "main", + "schema": "default", + "volume": "agent_files", + "root_path": "/", + "timeout": 30, + }, + } + + def test_backend_path_uses_volume_root_and_strips_mount_prefix(): config = DatabricksVolumeConfig( + host="https://example.cloud.databricks.com", catalog="main", schema="default", volume="agent_files", @@ -265,15 +329,14 @@ def test_backend_path_uses_volume_root_and_strips_mount_prefix(): path) == ("/Volumes/main/default/agent_files/root/reports/latest.md") -def test_resource_state_redacts_token(): +def test_resource_state_requires_runtime_override(): resource = make_resource(FakeFiles()) state = resource.get_state() assert state["type"] == ResourceName.DATABRICKS_VOLUME assert state["needs_override"] is True - assert state["config"]["token"] == "" - assert state["config"]["host"] is None + assert state["config"]["host"] == ("https://example.cloud.databricks.com") assert state["config"]["catalog"] == "main" - assert "token" in state["redacted_fields"] + assert "token" not in state["config"] def test_resource_registers_ops(): diff --git a/python/tests/resource/databricks_volume/test_token_provider.py b/python/tests/resource/databricks_volume/test_token_provider.py new file mode 100644 index 000000000..0faab8bdd --- /dev/null +++ b/python/tests/resource/databricks_volume/test_token_provider.py @@ -0,0 +1,44 @@ +import mirage.resource.databricks_volume as databricks_volume +from mirage.resource.databricks_volume import token_provider as provider_module + + +class FakeWorkspaceConfig: + calls: list[dict] = [] + + def __init__(self, **kwargs) -> None: + self.calls.append(kwargs) + + def authenticate(self) -> dict[str, str]: + return {"Authorization": "Bearer profile-token"} + + +def test_token_provider_types_are_public(): + assert hasattr(databricks_volume, "TokenProvider") + assert hasattr(databricks_volume, "StaticTokenProvider") + assert hasattr(databricks_volume, "DatabricksProfileTokenProvider") + + +def test_static_token_provider_returns_token(): + provider = databricks_volume.StaticTokenProvider("token") + + assert provider.get_token() == "token" + + +def test_profile_token_provider_uses_host_and_profile(monkeypatch): + FakeWorkspaceConfig.calls = [] + monkeypatch.setattr( + provider_module, + "WorkspaceConfig", + FakeWorkspaceConfig, + ) + provider = databricks_volume.DatabricksProfileTokenProvider( + "https://example.cloud.databricks.com", + profile="DEV", + ) + + assert provider.get_token() == "profile-token" + assert provider.get_token() == "profile-token" + assert FakeWorkspaceConfig.calls == [{ + "host": "https://example.cloud.databricks.com", + "profile": "DEV", + }] diff --git a/python/tests/resource/test_state_round_trip.py b/python/tests/resource/test_state_round_trip.py index 34c9d9c9a..88c9f8c60 100644 --- a/python/tests/resource/test_state_round_trip.py +++ b/python/tests/resource/test_state_round_trip.py @@ -21,6 +21,9 @@ from mirage.resource.ram import RAMResource from mirage.resource.redis import RedisResource from mirage.resource.s3 import S3Config, S3Resource +from mirage.types import MountKey, StateKey +from mirage.workspace.snapshot.state import (build_mount_args, + requires_resource_override) REDIS_URL = os.environ.get("REDIS_URL", "") @@ -51,6 +54,55 @@ def test_ram_round_trip(): assert "/sub" in dst._store.dirs +def test_resource_state_needs_override_requires_resource_override(): + mount_state = { + MountKey.RESOURCE_CLASS: + ("mirage.resource.databricks_volume.databricks_volume." + "DatabricksVolumeResource"), + MountKey.RESOURCE_STATE: { + "type": "databricks_volume", + "needs_override": True, + "config": { + "host": "https://example.cloud.databricks.com", + "catalog": "main", + "schema": "default", + "volume": "documents", + "root_path": "/", + "timeout": 30, + }, + }, + } + + assert requires_resource_override(mount_state) is True + + +def test_resource_state_needs_override_fails_without_override(): + state = { + StateKey.MOUNTS: [{ + MountKey.PREFIX: + "/dbx/", + MountKey.MODE: + "read", + MountKey.RESOURCE_CLASS: + ("mirage.resource.databricks_volume.databricks_volume." + "DatabricksVolumeResource"), + MountKey.RESOURCE_STATE: { + "type": "databricks_volume", + "needs_override": True, + "config": { + "host": "https://example.cloud.databricks.com", + "catalog": "main", + "schema": "default", + "volume": "documents", + }, + }, + }], + } + + with pytest.raises(ValueError, match="resources=.*dbx"): + build_mount_args(state) + + # ── Disk ───────────────────────────────────────────────────────────────