vibedata-dlt-databricks-notebook 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,9 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ dist/
5
+ build/
6
+ .venv/
7
+ .env
8
+ .pytest_cache/
9
+ uv.lock
@@ -0,0 +1,34 @@
1
+ Metadata-Version: 2.5
2
+ Name: vibedata-dlt-databricks-notebook
3
+ Version: 0.1.0
4
+ Summary: VibeData dlt runtime — Databricks notebook execution unit (runs inside a Databricks notebook): a dlt provider for the Domain's Databricks secret scope, and a databricks destination that stages load files through the Files API. dbutils is provided by the Databricks runtime.
5
+ Author-email: Vibedata <eng@acceleratedata.ai>
6
+ License: MIT
7
+ Requires-Python: >=3.11
8
+ Requires-Dist: dlt[databricks]<1.31,>=1.30.0
9
+ Requires-Dist: vibedata-dlt-core>=0.11.0
10
+ Provides-Extra: dev
11
+ Requires-Dist: pytest-mock>=3.12; extra == 'dev'
12
+ Requires-Dist: pytest>=8.0; extra == 'dev'
13
+ Description-Content-Type: text/markdown
14
+
15
+ # vibedata-dlt-databricks-notebook
16
+
17
+ Databricks notebook execution unit for the VibeData dlt runtime — the installed distribution for Databricks domains running a dlt pipeline inside a Databricks notebook. `dbutils` is provided by the Databricks runtime, not a pip dependency.
18
+
19
+ ```python
20
+ from vibedata.dlt.databricks import DatabricksSecretScopeProvider, databricks_files_api, setup_environment
21
+ import dlt
22
+
23
+ setup_environment()
24
+ dlt.config.register_provider(DatabricksSecretScopeProvider())
25
+ pipeline = dlt.pipeline(pipeline_name="<name>", destination=databricks_files_api(credentials={...}), dataset_name="<schema>")
26
+ ```
27
+
28
+ - **Import order.** Import this package before anything imports `dlt`: on Databricks compute `import dlt` otherwise binds Databricks' own Lakeflow Declarative Pipelines module.
29
+ - **`setup_environment()`** makes the Domain's `ingestion/` (the notebook's `ingestion_dir` widget) the working directory and an import root, names dlt's scratch schema `<dataset>_dlt_staging` so it cannot collide with a dbt-modelled `staging` layer, and caps load files below the SDK's single-upload limit.
30
+ - **`DatabricksSecretScopeProvider`** resolves dlt secrets from the Databricks secret scope that `[providers.databricks_secret_scope] scope` names in the committed `.dlt/config.toml`. Each key is the dotted dlt key, lowercased. It resolves nothing when no scope is configured.
31
+ - **`databricks_files_api`** is dlt's `databricks` destination with one change: it stages load files on the volume through the Files API instead of the SQL connector's `PUT`. For a catalog on Databricks-managed storage, `PUT` uploads to the region's storage gateway, whose name resolves to a refusing private address on serverless compute. dlt has no hook for the upload, so the class overrides dlt's private staging methods; raise the package's dlt cap only after re-checking them against the new release.
32
+ - **No `finalize()`.** The unit writes no audit tables yet, so a Databricks pipeline makes the one `setup_environment()` call; the other units' audit pair arrives with Databricks run history.
33
+
34
+ Install this package; it pulls in `vibedata-dlt-core` and `dlt[databricks]`.
@@ -0,0 +1,20 @@
1
+ # vibedata-dlt-databricks-notebook
2
+
3
+ Databricks notebook execution unit for the VibeData dlt runtime — the installed distribution for Databricks domains running a dlt pipeline inside a Databricks notebook. `dbutils` is provided by the Databricks runtime, not a pip dependency.
4
+
5
+ ```python
6
+ from vibedata.dlt.databricks import DatabricksSecretScopeProvider, databricks_files_api, setup_environment
7
+ import dlt
8
+
9
+ setup_environment()
10
+ dlt.config.register_provider(DatabricksSecretScopeProvider())
11
+ pipeline = dlt.pipeline(pipeline_name="<name>", destination=databricks_files_api(credentials={...}), dataset_name="<schema>")
12
+ ```
13
+
14
+ - **Import order.** Import this package before anything imports `dlt`: on Databricks compute `import dlt` otherwise binds Databricks' own Lakeflow Declarative Pipelines module.
15
+ - **`setup_environment()`** makes the Domain's `ingestion/` (the notebook's `ingestion_dir` widget) the working directory and an import root, names dlt's scratch schema `<dataset>_dlt_staging` so it cannot collide with a dbt-modelled `staging` layer, and caps load files below the SDK's single-upload limit.
16
+ - **`DatabricksSecretScopeProvider`** resolves dlt secrets from the Databricks secret scope that `[providers.databricks_secret_scope] scope` names in the committed `.dlt/config.toml`. Each key is the dotted dlt key, lowercased. It resolves nothing when no scope is configured.
17
+ - **`databricks_files_api`** is dlt's `databricks` destination with one change: it stages load files on the volume through the Files API instead of the SQL connector's `PUT`. For a catalog on Databricks-managed storage, `PUT` uploads to the region's storage gateway, whose name resolves to a refusing private address on serverless compute. dlt has no hook for the upload, so the class overrides dlt's private staging methods; raise the package's dlt cap only after re-checking them against the new release.
18
+ - **No `finalize()`.** The unit writes no audit tables yet, so a Databricks pipeline makes the one `setup_environment()` call; the other units' audit pair arrives with Databricks run history.
19
+
20
+ Install this package; it pulls in `vibedata-dlt-core` and `dlt[databricks]`.
@@ -0,0 +1,27 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "vibedata-dlt-databricks-notebook"
7
+ version = "0.1.0"
8
+ description = "VibeData dlt runtime — Databricks notebook execution unit (runs inside a Databricks notebook): a dlt provider for the Domain's Databricks secret scope, and a databricks destination that stages load files through the Files API. dbutils is provided by the Databricks runtime."
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ license = { text = "MIT" }
12
+ authors = [{ name = "Vibedata", email = "eng@acceleratedata.ai" }]
13
+
14
+ dependencies = [
15
+ "vibedata-dlt-core>=0.11.0",
16
+ "dlt[databricks]>=1.30.0,<1.31",
17
+ ]
18
+
19
+ [project.optional-dependencies]
20
+ dev = ["pytest>=8.0", "pytest-mock>=3.12"]
21
+
22
+ [tool.hatch.build.targets.wheel]
23
+ packages = ["src/vibedata"]
24
+
25
+ [tool.pytest.ini_options]
26
+ testpaths = ["tests"]
27
+ pythonpath = ["src"]
@@ -0,0 +1,25 @@
1
+ """VibeData dlt runtime — Databricks notebook execution unit.
2
+
3
+ The installed unit for Databricks domains running inside a Databricks notebook:
4
+
5
+ from vibedata.dlt.databricks import DatabricksSecretScopeProvider, databricks_files_api, setup_environment
6
+ import dlt
7
+
8
+ setup_environment()
9
+ dlt.config.register_provider(DatabricksSecretScopeProvider())
10
+ pipeline = dlt.pipeline(pipeline_name="<name>", destination=databricks_files_api(...), dataset_name="<schema>")
11
+
12
+ Import this package before anything imports `dlt`: on Databricks compute, `import dlt`
13
+ otherwise binds Databricks' own Lakeflow Declarative Pipelines module, and this package
14
+ imports dlthub's in its place. `dbutils` is provided by the Databricks runtime.
15
+ """
16
+
17
+ from ._dlt_import import import_dlthub_dlt
18
+
19
+ import_dlthub_dlt()
20
+
21
+ from .bootstrap import setup_environment # noqa: E402
22
+ from .destination import databricks_files_api # noqa: E402
23
+ from .providers.secret_scope import DatabricksSecretScopeProvider # noqa: E402
24
+
25
+ __all__ = ["DatabricksSecretScopeProvider", "databricks_files_api", "setup_environment"]
@@ -0,0 +1,36 @@
1
+ """Import dlthub's `dlt` on Databricks compute.
2
+
3
+ Databricks registers a post-import hook that answers `import dlt` with its own Lakeflow
4
+ Declarative Pipelines module. Importing with that hook set aside binds dlthub's package,
5
+ the approach dlt's own Databricks destination docs give for serverless compute.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import importlib
11
+ import sys
12
+ from types import ModuleType
13
+
14
+ POST_IMPORT_HOOK = "PostImportHook"
15
+
16
+
17
+ def _is_dlthub(module: ModuleType) -> bool:
18
+ return hasattr(module, "pipeline")
19
+
20
+
21
+ def import_dlthub_dlt() -> ModuleType:
22
+ """Return dlthub's `dlt`, importing it past Databricks' hook when nothing has yet."""
23
+ loaded = sys.modules.get("dlt")
24
+ if loaded is None:
25
+ hooks = list(sys.meta_path)
26
+ sys.meta_path = [hook for hook in hooks if POST_IMPORT_HOOK not in repr(hook)]
27
+ try:
28
+ loaded = importlib.import_module("dlt")
29
+ finally:
30
+ sys.meta_path = hooks
31
+ if not _is_dlthub(loaded):
32
+ raise ImportError(
33
+ "`dlt` resolved to Databricks' Lakeflow Declarative Pipelines module, not dlthub's. "
34
+ "Import vibedata.dlt.databricks before anything imports dlt."
35
+ )
36
+ return loaded
@@ -0,0 +1,53 @@
1
+ """Prepare a Databricks notebook for a dlt run of the Domain's committed pipeline."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ import sys
7
+ from typing import Optional
8
+
9
+ INGESTION_DIR_WIDGET = "ingestion_dir"
10
+
11
+ # dlt's default staging layout, `%s_staging`, is the schema dbt derives for a `staging`
12
+ # custom schema, so the two would share one schema.
13
+ STAGING_LAYOUT_KEY = "DESTINATION__STAGING_DATASET_NAME_LAYOUT"
14
+ STAGING_LAYOUT = "%s_dlt_staging"
15
+
16
+ # databricks_files_api uploads a load file in one Files API call only below the SDK's
17
+ # multipart threshold (50 MiB); a larger file goes through the unreachable storage gateway.
18
+ FILE_MAX_BYTES_KEY = "DATA_WRITER__FILE_MAX_BYTES"
19
+ FILE_MAX_BYTES = 40 * 1024 * 1024
20
+
21
+
22
+ def setup_environment(ingestion_dir: Optional[str] = None) -> None:
23
+ """Run from the Domain's `ingestion/`, so `sources.*` imports and `.dlt/` resolve.
24
+
25
+ `ingestion_dir` defaults to the notebook's `ingestion_dir` widget, which the job sets.
26
+ """
27
+ if ingestion_dir is None:
28
+ from databricks.sdk.runtime import dbutils
29
+
30
+ ingestion_dir = dbutils.widgets.get(INGESTION_DIR_WIDGET)
31
+ file_max_bytes = _file_max_bytes(os.environ.get(FILE_MAX_BYTES_KEY))
32
+ ingestion_dir = os.path.abspath(ingestion_dir)
33
+ os.chdir(ingestion_dir)
34
+ if ingestion_dir not in sys.path:
35
+ sys.path.insert(0, ingestion_dir)
36
+ os.environ[STAGING_LAYOUT_KEY] = STAGING_LAYOUT
37
+ os.environ[FILE_MAX_BYTES_KEY] = str(file_max_bytes)
38
+
39
+
40
+ def _file_max_bytes(configured: Optional[str]) -> int:
41
+ """The cap for `DATA_WRITER__FILE_MAX_BYTES`, the one setting this guards.
42
+
43
+ A lower value already in that variable is kept; unset, empty, 0 (dlt's "unlimited"), negative
44
+ or higher gets FILE_MAX_BYTES. dlt rotates files at buffer flushes, so the cap bounds a file's
45
+ size closely, not exactly; FILE_MAX_BYTES leaves the margin below the SDK's threshold.
46
+ """
47
+ if not configured:
48
+ return FILE_MAX_BYTES
49
+ try:
50
+ value = int(configured)
51
+ except ValueError:
52
+ raise ValueError(f"{FILE_MAX_BYTES_KEY} must be a whole number of bytes; got {configured!r}") from None
53
+ return FILE_MAX_BYTES if value <= 0 else min(value, FILE_MAX_BYTES)
@@ -0,0 +1,76 @@
1
+ """A `databricks` destination that stages load files through the Files API.
2
+
3
+ dlt stages each local load file on a Unity Catalog volume with the SQL connector's `PUT`,
4
+ which uploads to a signed URL. For a catalog on Databricks-managed storage that URL is on
5
+ the region's storage gateway, and on serverless compute the gateway's name resolves to a
6
+ private address that refuses the connection. The Files API reaches the same volume
7
+ through the workspace host. Everything else is dlt's own: the volume, the file layout and
8
+ the `COPY INTO` that loads it.
9
+
10
+ dlt has no hook for the upload, so this overrides the two private methods that stage and
11
+ remove the file. The package caps dlt to the minor release they were checked against.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ from typing import Any, Optional
17
+
18
+ from databricks.sdk import WorkspaceClient
19
+ from dlt.common.storages import FileStorage
20
+ from dlt.common.utils import uniq_id
21
+ from dlt.destinations import databricks
22
+ from dlt.destinations.impl.databricks.databricks import DatabricksClient, DatabricksLoadJob
23
+
24
+ DEFAULT_STAGING_VOLUME = "_dlt_staging_load_volume"
25
+
26
+
27
+ def _files() -> Any:
28
+ return WorkspaceClient().files
29
+
30
+
31
+ class FilesApiLoadJob(DatabricksLoadJob):
32
+ def _handle_local_file_upload(self, local_file_path: str) -> tuple[str, str, str, str]:
33
+ file_name = FileStorage.get_file_name_from_file_path(local_file_path)
34
+ # Databricks refuses to load a file whose name starts with `_` or `.`.
35
+ volume_file_name = "valid" + file_name if file_name.startswith(("_", ".")) else file_name
36
+ config = self._job_client.config
37
+ if config.staging_volume_name:
38
+ catalog, schema, volume = config.staging_volume_name.split(".")
39
+ else:
40
+ catalog, schema, volume = self._sql_client.database_name, self._sql_client.dataset_name, DEFAULT_STAGING_VOLUME
41
+ self._sql_client.execute_sql(f"CREATE VOLUME IF NOT EXISTS {catalog}.{schema}.{volume}")
42
+ volume_path = f"/Volumes/{catalog}/{schema}/{volume}/{uniq_id()}"
43
+ volume_file_path = f"{volume_path}/{volume_file_name}"
44
+ with open(local_file_path, "rb") as source:
45
+ _files().upload(volume_file_path, source, overwrite=True)
46
+ return f"FROM '{volume_path}'", file_name, volume_path, volume_file_path
47
+
48
+ def _handle_staged_file_remove(self, volume_path: str, volume_file_path: str) -> None:
49
+ files = _files()
50
+ files.delete(volume_file_path)
51
+ files.delete_directory(volume_path)
52
+
53
+
54
+ class FilesApiDatabricksClient(DatabricksClient):
55
+ def create_load_job(self, table: Any, file_path: str, load_id: str, restore: bool = False) -> Any:
56
+ job = super().create_load_job(table, file_path, load_id, restore)
57
+ # dlt builds this job only for the exact class, so it is swapped after the fact;
58
+ # reference, merge and Zerobus jobs stay as dlt built them.
59
+ if type(job) is DatabricksLoadJob:
60
+ job = FilesApiLoadJob(file_path, staging_config=job._staging_config)
61
+ return job
62
+
63
+
64
+ class databricks_files_api(databricks): # noqa: N801 - dlt names destination factories in lowercase
65
+ """dlt's `databricks` destination, staging through the Files API.
66
+
67
+ Named `databricks` by default, so `[destination.databricks]` config still applies.
68
+ """
69
+
70
+ def __init__(self, *args: Any, destination_name: Optional[str] = None, **kwargs: Any) -> None:
71
+ # The default stays None: dlt drops an argument equal to its signature default as unset.
72
+ super().__init__(*args, destination_name=destination_name or "databricks", **kwargs)
73
+
74
+ @property
75
+ def client_class(self) -> type[DatabricksClient]:
76
+ return FilesApiDatabricksClient
@@ -0,0 +1,105 @@
1
+ """The Domain's Databricks secret scope as a dlt secrets provider.
2
+
3
+ A `databricks_secret_scope` Secret Store keeps each credential as a key in one
4
+ Databricks-backed secret scope, named by its dotted dlt key, lowercased
5
+ (`sources.sf_prod.credentials.client_id`). The committed `.dlt/config.toml` names the scope:
6
+
7
+ [providers.databricks_secret_scope]
8
+ scope = "<scope>"
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import logging
14
+ import tomllib
15
+ from pathlib import Path
16
+ from typing import Any, Optional, Sequence, Tuple, Type
17
+
18
+ from dlt.common.configuration.providers import ConfigProvider
19
+
20
+ logger = logging.getLogger(__name__)
21
+
22
+ PROVIDER_NAME = "databricks_secret_scope"
23
+ CONFIG_FILE = Path(".dlt") / "config.toml"
24
+
25
+ # dlt retries a key with its sections popped from the right, down to none. Below two
26
+ # (`sources` and the connection) a key names nothing connection-specific, and answering
27
+ # it would serve one source another's credential.
28
+ MINIMUM_SECTIONS = 2
29
+ MAX_KEY_LENGTH = 128
30
+
31
+
32
+ def scope_key(key: str, *sections: str) -> str:
33
+ """The secret scope key for a dlt key. Raises for a path that cannot be one."""
34
+ if len(sections) < MINIMUM_SECTIONS or not all(sections):
35
+ raise ValueError(f"a scope key needs at least {MINIMUM_SECTIONS} sections, none empty; got {sections!r}")
36
+ name = ".".join((*sections, key)).lower()
37
+ if len(name) > MAX_KEY_LENGTH:
38
+ raise ValueError(f"{name!r} is longer than a secret scope key's {MAX_KEY_LENGTH} characters")
39
+ return name
40
+
41
+
42
+ def configured_scope(config_file: Path = CONFIG_FILE) -> Optional[str]:
43
+ """The scope the committed config names, or None when it names none."""
44
+ if not config_file.is_file():
45
+ return None
46
+ config = tomllib.loads(config_file.read_text(encoding="utf-8"))
47
+ return config.get("providers", {}).get(PROVIDER_NAME, {}).get("scope") or None
48
+
49
+
50
+ class DatabricksSecretScopeProvider(ConfigProvider):
51
+ """Resolve dlt secrets from the Domain's Databricks secret scope.
52
+
53
+ Construct it after `setup_environment()`, which makes `ingestion/` the working directory
54
+ its config is read from. With no scope configured it resolves nothing.
55
+ """
56
+
57
+ def __init__(self, scope: Optional[str] = None, secrets: Any = None) -> None:
58
+ self._scope = scope if scope is not None else configured_scope()
59
+ self._secrets = secrets
60
+ self._keys: Optional[set[str]] = None
61
+ if not self._scope:
62
+ logger.info("No [providers.%s] scope configured; resolving no secrets from it", PROVIDER_NAME)
63
+
64
+ @property
65
+ def name(self) -> str:
66
+ return PROVIDER_NAME
67
+
68
+ @property
69
+ def supports_secrets(self) -> bool:
70
+ return True
71
+
72
+ @property
73
+ def supports_sections(self) -> bool:
74
+ return True
75
+
76
+ @property
77
+ def locations(self) -> Sequence[str]:
78
+ # dlt renders locations through `os.path.abspath` in a resolution diagnostic.
79
+ return []
80
+
81
+ def get_value(
82
+ self, key: str, hint: Type[Any], pipeline_name: Optional[str], *sections: str
83
+ ) -> Tuple[Optional[Any], str]:
84
+ del hint, pipeline_name
85
+ try:
86
+ name = scope_key(key, *sections)
87
+ except ValueError:
88
+ return None, f"<no secret scope key for sections {sections!r}>"
89
+ if not self._scope or name not in self._available_keys():
90
+ return None, name
91
+ return self._client().get(self._scope, name), name
92
+
93
+ def _client(self) -> Any:
94
+ if self._secrets is None:
95
+ from databricks.sdk.runtime import dbutils
96
+
97
+ self._secrets = dbutils.secrets
98
+ return self._secrets
99
+
100
+ def _available_keys(self) -> set[str]:
101
+ # Listed once: `get` on an absent key raises a generic runtime error, which would
102
+ # leave a missing secret indistinguishable from a failing scope.
103
+ if self._keys is None:
104
+ self._keys = {secret.key for secret in self._client().list(self._scope)}
105
+ return self._keys
@@ -0,0 +1,92 @@
1
+ import os
2
+ import sys
3
+ import types
4
+
5
+ import pytest
6
+
7
+ from vibedata.dlt.databricks import _dlt_import, bootstrap
8
+
9
+
10
+ @pytest.fixture(autouse=True)
11
+ def restore_environ(monkeypatch):
12
+ # setup_environment writes os.environ; snapshot it so no test leaks into the next.
13
+ monkeypatch.setattr(os, "environ", os.environ.copy())
14
+
15
+
16
+ @pytest.mark.parametrize("configured, expected", [("0", bootstrap.FILE_MAX_BYTES), ("-5", bootstrap.FILE_MAX_BYTES), (str(10**12), bootstrap.FILE_MAX_BYTES)])
17
+ def test_an_unlimited_or_higher_file_cap_gets_the_upload_limit(configured, expected):
18
+ assert bootstrap._file_max_bytes(configured) == expected
19
+
20
+
21
+ def test_a_file_cap_that_is_not_a_number_names_the_variable():
22
+ with pytest.raises(ValueError, match=bootstrap.FILE_MAX_BYTES_KEY):
23
+ bootstrap._file_max_bytes("40MB")
24
+
25
+
26
+ def test_setup_environment_runs_from_the_ingestion_dir(tmp_path, monkeypatch):
27
+ monkeypatch.chdir(os.getcwd())
28
+ monkeypatch.setattr(sys, "path", list(sys.path))
29
+ monkeypatch.delenv(bootstrap.STAGING_LAYOUT_KEY, raising=False)
30
+ monkeypatch.delenv(bootstrap.FILE_MAX_BYTES_KEY, raising=False)
31
+
32
+ bootstrap.setup_environment(str(tmp_path))
33
+
34
+ assert os.getcwd() == str(tmp_path)
35
+ assert sys.path[0] == str(tmp_path)
36
+ assert os.environ[bootstrap.STAGING_LAYOUT_KEY] == "%s_dlt_staging"
37
+ assert int(os.environ[bootstrap.FILE_MAX_BYTES_KEY]) < 50 * 1024 * 1024
38
+
39
+
40
+ def test_import_returns_dlthub_dlt_when_already_loaded():
41
+ import dlt
42
+
43
+ assert _dlt_import.import_dlthub_dlt() is dlt
44
+
45
+
46
+ def test_import_refuses_databricks_own_dlt(monkeypatch):
47
+ monkeypatch.setitem(sys.modules, "dlt", types.ModuleType("dlt"))
48
+
49
+ with pytest.raises(ImportError, match="Lakeflow"):
50
+ _dlt_import.import_dlthub_dlt()
51
+
52
+
53
+ def test_import_sets_the_post_import_hook_aside_only_for_the_import(monkeypatch):
54
+ class PostImportHook:
55
+ def find_spec(self, *args, **kwargs):
56
+ raise AssertionError("the hook answered the import")
57
+
58
+ hook = PostImportHook()
59
+ seen = []
60
+ dlt_module = sys.modules["dlt"]
61
+ monkeypatch.delitem(sys.modules, "dlt")
62
+ monkeypatch.setattr(sys, "meta_path", [hook, *sys.meta_path])
63
+
64
+ def fake_import(name):
65
+ seen.append(list(sys.meta_path))
66
+ return dlt_module
67
+
68
+ monkeypatch.setattr(_dlt_import.importlib, "import_module", fake_import)
69
+
70
+ assert _dlt_import.import_dlthub_dlt() is dlt_module
71
+ assert hook not in seen[0]
72
+ assert sys.meta_path[0] is hook
73
+
74
+
75
+ def test_a_lower_file_cap_the_domain_set_is_kept(tmp_path, monkeypatch):
76
+ monkeypatch.chdir(os.getcwd())
77
+ monkeypatch.setattr(sys, "path", list(sys.path))
78
+ monkeypatch.setenv(bootstrap.FILE_MAX_BYTES_KEY, "1000")
79
+
80
+ bootstrap.setup_environment(str(tmp_path))
81
+
82
+ assert os.environ[bootstrap.FILE_MAX_BYTES_KEY] == "1000"
83
+
84
+
85
+ def test_a_relative_ingestion_dir_stays_importable(tmp_path, monkeypatch):
86
+ monkeypatch.chdir(tmp_path)
87
+ monkeypatch.setattr(sys, "path", list(sys.path))
88
+ (tmp_path / "ingestion").mkdir()
89
+
90
+ bootstrap.setup_environment("ingestion")
91
+
92
+ assert sys.path[0] == str(tmp_path / "ingestion")
@@ -0,0 +1,119 @@
1
+ from types import SimpleNamespace
2
+
3
+ import pytest
4
+
5
+ from vibedata.dlt.databricks import destination
6
+ from vibedata.dlt.databricks.destination import (
7
+ DatabricksLoadJob,
8
+ FilesApiDatabricksClient,
9
+ FilesApiLoadJob,
10
+ databricks_files_api,
11
+ )
12
+
13
+
14
+ class FakeFiles:
15
+ def __init__(self):
16
+ self.calls = []
17
+
18
+ def upload(self, path, source, overwrite):
19
+ self.calls.append(("upload", path, source.read(), overwrite))
20
+
21
+ def delete(self, path):
22
+ self.calls.append(("delete", path))
23
+
24
+ def delete_directory(self, path):
25
+ self.calls.append(("delete_directory", path))
26
+
27
+
28
+ class FakeSqlClient:
29
+ database_name = "cat"
30
+ dataset_name = "raw"
31
+
32
+ def __init__(self):
33
+ self.statements = []
34
+
35
+ def execute_sql(self, sql):
36
+ self.statements.append(sql)
37
+
38
+
39
+ @pytest.fixture
40
+ def files(monkeypatch):
41
+ fake = FakeFiles()
42
+ monkeypatch.setattr(destination, "_files", lambda: fake)
43
+ monkeypatch.setattr(destination, "uniq_id", lambda: "u1")
44
+ return fake
45
+
46
+
47
+ def load_job(staging_volume_name=None):
48
+ job = object.__new__(FilesApiLoadJob)
49
+ job._job_client = SimpleNamespace(config=SimpleNamespace(staging_volume_name=staging_volume_name))
50
+ job._sql_client = FakeSqlClient()
51
+ return job
52
+
53
+
54
+ def test_upload_goes_through_the_files_api_into_the_named_volume(tmp_path, files):
55
+ local = tmp_path / "pokemon.abc.0.parquet"
56
+ local.write_bytes(b"rows")
57
+ job = load_job("cat.raw.vol")
58
+
59
+ result = job._handle_local_file_upload(str(local))
60
+
61
+ assert files.calls == [("upload", "/Volumes/cat/raw/vol/u1/pokemon.abc.0.parquet", b"rows", True)]
62
+ assert result == ("FROM '/Volumes/cat/raw/vol/u1'", "pokemon.abc.0.parquet", "/Volumes/cat/raw/vol/u1", "/Volumes/cat/raw/vol/u1/pokemon.abc.0.parquet")
63
+ assert job._sql_client.statements == []
64
+
65
+
66
+ def test_without_a_named_volume_it_creates_dlts_default_one(tmp_path, files):
67
+ local = tmp_path / "_dlt_pipeline_state.abc.0.parquet"
68
+ local.write_bytes(b"state")
69
+ job = load_job()
70
+
71
+ _, file_name, _, volume_file_path = job._handle_local_file_upload(str(local))
72
+
73
+ assert job._sql_client.statements == ["CREATE VOLUME IF NOT EXISTS cat.raw._dlt_staging_load_volume"]
74
+ assert file_name == "_dlt_pipeline_state.abc.0.parquet"
75
+ assert volume_file_path == "/Volumes/cat/raw/_dlt_staging_load_volume/u1/valid_dlt_pipeline_state.abc.0.parquet"
76
+
77
+
78
+ def test_remove_deletes_the_file_then_its_folder(files):
79
+ load_job()._handle_staged_file_remove("/Volumes/c/s/v/u1", "/Volumes/c/s/v/u1/f.parquet")
80
+
81
+ assert files.calls == [("delete", "/Volumes/c/s/v/u1/f.parquet"), ("delete_directory", "/Volumes/c/s/v/u1")]
82
+
83
+
84
+ def test_the_client_swaps_only_dlts_own_copy_into_job(monkeypatch):
85
+ staging_config = object()
86
+ built = object.__new__(DatabricksLoadJob)
87
+ built._staging_config = staging_config
88
+ other = object()
89
+ jobs = iter([built, other])
90
+ monkeypatch.setattr(destination.DatabricksClient, "create_load_job", lambda *a, **k: next(jobs))
91
+ monkeypatch.setattr(FilesApiLoadJob, "__init__", lambda self, path, staging_config: setattr(self, "args", (path, staging_config)))
92
+ client = object.__new__(FilesApiDatabricksClient)
93
+
94
+ swapped = client.create_load_job({}, "t.abc.0.parquet", "load")
95
+
96
+ assert type(swapped) is FilesApiLoadJob
97
+ assert swapped.args == ("t.abc.0.parquet", staging_config)
98
+ assert client.create_load_job({}, "t.abc.0.parquet", "load") is other
99
+
100
+
101
+ def test_the_destination_keeps_dlts_databricks_name_and_config_section():
102
+ factory = databricks_files_api(credentials={"catalog": "cat"})
103
+
104
+ assert factory.destination_name == "databricks"
105
+ assert factory.client_class is FilesApiDatabricksClient
106
+ assert databricks_files_api(destination_name="other").destination_name == "other"
107
+
108
+
109
+ @pytest.mark.parametrize("method", ["_handle_local_file_upload", "_handle_staged_file_remove"])
110
+ def test_the_overridden_methods_still_exist_on_dlts_job(method):
111
+ # The overrides only take effect while dlt calls these names; a dlt release that renames
112
+ # one would silently send the upload back through PUT.
113
+ assert method in vars(DatabricksLoadJob)
114
+
115
+
116
+ def test_the_files_api_job_builds_with_dlts_constructor(tmp_path):
117
+ job = FilesApiLoadJob(str(tmp_path / "pokemon.abc.0.parquet"), staging_config=None)
118
+
119
+ assert job._staging_config is None
@@ -0,0 +1,93 @@
1
+ from types import SimpleNamespace
2
+
3
+ import pytest
4
+
5
+ from vibedata.dlt.databricks.providers.secret_scope import (
6
+ DatabricksSecretScopeProvider,
7
+ configured_scope,
8
+ scope_key,
9
+ )
10
+
11
+
12
+ class FakeSecrets:
13
+ def __init__(self, values):
14
+ self.values = values
15
+ self.gets = []
16
+
17
+ def list(self, scope):
18
+ return [SimpleNamespace(key=key) for (s, key) in self.values if s == scope]
19
+
20
+ def get(self, scope, key):
21
+ self.gets.append(key)
22
+ return self.values[(scope, key)]
23
+
24
+
25
+ def test_scope_key_is_the_dotted_key_lowercased():
26
+ assert scope_key("Client_ID", "sources", "SF_prod", "credentials") == "sources.sf_prod.credentials.client_id"
27
+
28
+
29
+ @pytest.mark.parametrize("sections", [(), ("sources",), ("sources", "")])
30
+ def test_scope_key_refuses_a_path_that_names_no_connection(sections):
31
+ with pytest.raises(ValueError):
32
+ scope_key("password", *sections)
33
+
34
+
35
+ def test_scope_key_refuses_a_name_longer_than_a_scope_key():
36
+ with pytest.raises(ValueError):
37
+ scope_key("k", "sources", "x" * 130)
38
+
39
+
40
+ def test_resolves_a_key_present_in_the_scope():
41
+ secrets = FakeSecrets({("domain", "sources.crm.credentials.password"): "s3cret"})
42
+ provider = DatabricksSecretScopeProvider(scope="domain", secrets=secrets)
43
+
44
+ assert provider.get_value("password", str, None, "sources", "crm", "credentials") == (
45
+ "s3cret",
46
+ "sources.crm.credentials.password",
47
+ )
48
+
49
+
50
+ def test_an_absent_key_resolves_to_none_without_a_get():
51
+ secrets = FakeSecrets({})
52
+ provider = DatabricksSecretScopeProvider(scope="domain", secrets=secrets)
53
+
54
+ assert provider.get_value("password", str, None, "sources", "crm")[0] is None
55
+ assert secrets.gets == []
56
+
57
+
58
+ def test_a_shallow_path_is_never_answered():
59
+ secrets = FakeSecrets({("domain", "sources.password"): "other"})
60
+ provider = DatabricksSecretScopeProvider(scope="domain", secrets=secrets)
61
+
62
+ assert provider.get_value("password", str, None, "sources")[0] is None
63
+
64
+
65
+ def test_no_configured_scope_resolves_nothing(tmp_path, monkeypatch):
66
+ monkeypatch.chdir(tmp_path)
67
+ provider = DatabricksSecretScopeProvider(secrets=FakeSecrets({}))
68
+
69
+ assert provider.get_value("password", str, None, "sources", "crm")[0] is None
70
+
71
+
72
+ def test_configured_scope_reads_the_committed_config(tmp_path):
73
+ config = tmp_path / "config.toml"
74
+ config.write_text('[providers.databricks_secret_scope]\nscope = "domain"\n')
75
+
76
+ assert configured_scope(config) == "domain"
77
+ assert configured_scope(tmp_path / "absent.toml") is None
78
+
79
+
80
+ def test_the_scope_is_listed_once():
81
+ calls = []
82
+
83
+ class CountingSecrets(FakeSecrets):
84
+ def list(self, scope):
85
+ calls.append(scope)
86
+ return super().list(scope)
87
+
88
+ secrets = CountingSecrets({("domain", "sources.crm.credentials.password"): "s3cret"})
89
+ provider = DatabricksSecretScopeProvider(scope="domain", secrets=secrets)
90
+ provider.get_value("password", str, None, "sources", "crm", "credentials")
91
+ provider.get_value("user", str, None, "sources", "crm", "credentials")
92
+
93
+ assert calls == ["domain"]