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.
- vibedata_dlt_databricks_notebook-0.1.0/.gitignore +9 -0
- vibedata_dlt_databricks_notebook-0.1.0/PKG-INFO +34 -0
- vibedata_dlt_databricks_notebook-0.1.0/README.md +20 -0
- vibedata_dlt_databricks_notebook-0.1.0/pyproject.toml +27 -0
- vibedata_dlt_databricks_notebook-0.1.0/src/vibedata/dlt/databricks/__init__.py +25 -0
- vibedata_dlt_databricks_notebook-0.1.0/src/vibedata/dlt/databricks/_dlt_import.py +36 -0
- vibedata_dlt_databricks_notebook-0.1.0/src/vibedata/dlt/databricks/bootstrap.py +53 -0
- vibedata_dlt_databricks_notebook-0.1.0/src/vibedata/dlt/databricks/destination.py +76 -0
- vibedata_dlt_databricks_notebook-0.1.0/src/vibedata/dlt/databricks/providers/__init__.py +0 -0
- vibedata_dlt_databricks_notebook-0.1.0/src/vibedata/dlt/databricks/providers/secret_scope.py +105 -0
- vibedata_dlt_databricks_notebook-0.1.0/tests/test_bootstrap.py +92 -0
- vibedata_dlt_databricks_notebook-0.1.0/tests/test_destination.py +119 -0
- vibedata_dlt_databricks_notebook-0.1.0/tests/test_secret_scope.py +93 -0
|
@@ -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
|
|
File without changes
|
|
@@ -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"]
|