feat: databricks-bundles Python resource loader — typed DAB resources
Some checks failed
CI / skinny-install (aco) (push) Successful in 49s
CI / skinny-install (api) (push) Successful in 30s
CI / lint-test (push) Successful in 1m27s
CI / skinny-install (bcda) (push) Successful in 33s
CI / skinny-install (bib) (push) Successful in 30s
CI / skinny-install (ccw) (push) Successful in 27s
CI / skinny-install (bls) (push) Successful in 30s
CI / skinny-install (cms) (push) Successful in 26s
CI / skinny-install (cli) (push) Successful in 34s
CI / skinny-install (conf) (push) Successful in 25s
CI / skinny-install (perf) (push) Successful in 29s
CI / skinny-install (pfs) (push) Successful in 32s
CI / skinny-install (rex) (push) Successful in 26s
Infra CI / notebooks (push) Successful in 6s
Infra CI / zotero (push) Successful in 7s
Infra CI / docs (push) Successful in 8s
Infra CI / api (push) Successful in 9s
Infra CI / mc (push) Successful in 7s
Package Supply Chain / pkg-supply-chain (push) Failing after 52s
Deploy / build-scan-report (push) Successful in 3m21s
Some checks failed
CI / skinny-install (aco) (push) Successful in 49s
CI / skinny-install (api) (push) Successful in 30s
CI / lint-test (push) Successful in 1m27s
CI / skinny-install (bcda) (push) Successful in 33s
CI / skinny-install (bib) (push) Successful in 30s
CI / skinny-install (ccw) (push) Successful in 27s
CI / skinny-install (bls) (push) Successful in 30s
CI / skinny-install (cms) (push) Successful in 26s
CI / skinny-install (cli) (push) Successful in 34s
CI / skinny-install (conf) (push) Successful in 25s
CI / skinny-install (perf) (push) Successful in 29s
CI / skinny-install (pfs) (push) Successful in 32s
CI / skinny-install (rex) (push) Successful in 26s
Infra CI / notebooks (push) Successful in 6s
Infra CI / zotero (push) Successful in 7s
Infra CI / docs (push) Successful in 8s
Infra CI / api (push) Successful in 9s
Infra CI / mc (push) Successful in 7s
Package Supply Chain / pkg-supply-chain (push) Failing after 52s
Deploy / build-scan-report (push) Successful in 3m21s
Replace 375-line OrderedDict YAML generator with typed Python resources via databricks-bundles 0.295.0. databricks.yml shrinks from 278 lines to 23 — delegates all resource definitions to bundle/resources/__init__.py via python.resources entry point. - bundle/resources/__init__.py: load_resources() returns Resources with Job, Schema, Volume from pipeline registry - databricks.yml: minimal (bundle name, targets, python entry point) - Replaced databricks-cli>=0.18.0 with databricks-bundles>=0.295.0 - 30 DAB tests passing (12 new Python resource tests) - 12,067 total tests passing Refs #224
This commit is contained in:
140
bundle/resources/__init__.py
Normal file
140
bundle/resources/__init__.py
Normal file
@@ -0,0 +1,140 @@
|
||||
"""Databricks Asset Bundle resources — Python-defined.
|
||||
|
||||
Called by the DAB CLI at deploy time via::
|
||||
|
||||
# databricks.yml
|
||||
python:
|
||||
resources:
|
||||
- "resources:load_resources"
|
||||
|
||||
Generates Job, Schema, and Volume resources from the pipeline
|
||||
registry, eliminating the need for YAML generation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from databricks.bundles.core import Resources, Variable
|
||||
from databricks.bundles.jobs import (
|
||||
CronSchedule,
|
||||
Job,
|
||||
JobCluster,
|
||||
Library,
|
||||
PythonWheelTask,
|
||||
Task,
|
||||
TaskDependency,
|
||||
)
|
||||
from databricks.bundles.schemas import Schema
|
||||
from databricks.bundles.volumes import Volume, VolumeType
|
||||
|
||||
|
||||
def _pipeline_registry() -> dict:
|
||||
"""Import the pipeline registry. Returns empty dict if unavailable."""
|
||||
try:
|
||||
from aco.pipe import registry
|
||||
|
||||
return registry
|
||||
except ImportError:
|
||||
return {}
|
||||
|
||||
|
||||
def _pipeline_schemas(registry: dict) -> list[str]:
|
||||
"""Derive schema names from pipeline expression outputs."""
|
||||
schemas: set[str] = set()
|
||||
for pipe in registry.values():
|
||||
for expr in pipe.exprs:
|
||||
name = expr.name if isinstance(expr.name, str) else expr[0]
|
||||
if "." in name:
|
||||
schemas.add(name.split(".")[0])
|
||||
return sorted(schemas)
|
||||
|
||||
|
||||
def _config() -> dict:
|
||||
"""Read [databricks] section from stack.toml."""
|
||||
try:
|
||||
from conf import cfg
|
||||
|
||||
return cfg.databricks.to_dict()
|
||||
except (ImportError, AttributeError):
|
||||
return {}
|
||||
|
||||
|
||||
def load_resources() -> Resources:
|
||||
"""Generate all DAB resources from the pipeline registry."""
|
||||
registry = _pipeline_registry()
|
||||
cfg = _config()
|
||||
schema_names = _pipeline_schemas(registry)
|
||||
|
||||
catalog = Variable(path="catalog", type=str)
|
||||
|
||||
# Schemas — one per pipeline-produced schema
|
||||
schemas = {
|
||||
name: Schema(
|
||||
name=name,
|
||||
catalog_name=catalog,
|
||||
)
|
||||
for name in schema_names
|
||||
}
|
||||
|
||||
# Volume — staging area for data uploads
|
||||
volume = Volume(
|
||||
name="staging",
|
||||
catalog_name=catalog,
|
||||
schema_name="default",
|
||||
volume_type=VolumeType.MANAGED,
|
||||
)
|
||||
|
||||
# Job cluster
|
||||
cluster_cfg = cfg.get("clusters", {}).get("default", {})
|
||||
job_cluster = JobCluster(
|
||||
job_cluster_key="default_cluster",
|
||||
new_cluster={
|
||||
"spark_version": cluster_cfg.get("spark_version", "15.4.x-scala2.12"),
|
||||
"node_type_id": cluster_cfg.get("node_type_id", "i3.xlarge"),
|
||||
"num_workers": cluster_cfg.get("num_workers", 1),
|
||||
},
|
||||
)
|
||||
|
||||
# Tasks — one per pipeline
|
||||
package_name = cfg.get("package_name", "stack")
|
||||
entry_point = cfg.get("entry_point", "cli")
|
||||
tasks = []
|
||||
for pipe_name, pipe in sorted(registry.items()):
|
||||
task = Task(
|
||||
task_key=pipe_name,
|
||||
depends_on=[TaskDependency(task_key=dep) for dep in pipe.upstream] or None,
|
||||
job_cluster_key="default_cluster",
|
||||
python_wheel_task=PythonWheelTask(
|
||||
package_name=package_name,
|
||||
entry_point=entry_point,
|
||||
named_parameters={
|
||||
"run": pipe_name,
|
||||
"--target": "spark",
|
||||
"--catalog": catalog,
|
||||
"--save": "",
|
||||
},
|
||||
),
|
||||
libraries=[Library(whl="dbfs:/FileStore/wheels/stack-latest.whl")],
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# Schedule
|
||||
schedule_expr = cfg.get("schedule", "0 0 6 * * ?")
|
||||
timezone = cfg.get("timezone", "America/New_York")
|
||||
|
||||
job = Job(
|
||||
name=f"{cfg.get('bundle_name', 'stack')}-pipelines",
|
||||
description=f"Run all {len(registry)} ACO pipelines in dependency order",
|
||||
schedule=CronSchedule(
|
||||
quartz_cron_expression=schedule_expr,
|
||||
timezone_id=timezone,
|
||||
),
|
||||
job_clusters=[job_cluster],
|
||||
tasks=tasks,
|
||||
)
|
||||
|
||||
resources = Resources()
|
||||
resources.add_job("stack_pipelines", job)
|
||||
for name, schema in schemas.items():
|
||||
resources.add_schema(name, schema)
|
||||
resources.add_volume("staging", volume)
|
||||
return resources
|
||||
261
databricks.yml
261
databricks.yml
@@ -3,13 +3,10 @@
|
||||
|
||||
bundle:
|
||||
name: stack
|
||||
sync:
|
||||
include:
|
||||
- bundle/**
|
||||
variables:
|
||||
catalog:
|
||||
description: Unity Catalog name for table references
|
||||
default: aco
|
||||
python:
|
||||
venv_path: .venv
|
||||
resources:
|
||||
- resources:load_resources
|
||||
workspace:
|
||||
host: ${DATABRICKS_HOST}
|
||||
targets:
|
||||
@@ -25,253 +22,3 @@ targets:
|
||||
staging:
|
||||
variables:
|
||||
catalog: aco_staging
|
||||
resources:
|
||||
jobs:
|
||||
stack_pipelines:
|
||||
name: stack-pipelines
|
||||
description: Run all 13 ACO pipelines in dependency order
|
||||
schedule:
|
||||
quartz_cron_expression: 0 0 6 * * ?
|
||||
timezone_id: America/New_York
|
||||
job_clusters:
|
||||
- job_cluster_key: default_cluster
|
||||
new_cluster:
|
||||
spark_version: 15.4.x-scala2.12
|
||||
node_type_id: i3.xlarge
|
||||
num_workers: 1
|
||||
tasks:
|
||||
- task_key: ahrq_measures
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: ahrq_measures
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: cclf
|
||||
depends_on:
|
||||
- task_key: input_layer
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: cclf
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: claims_preprocessing
|
||||
depends_on:
|
||||
- task_key: input_layer
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: claims_preprocessing
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: cms_quality_measures
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: cms_quality_measures
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: core
|
||||
depends_on:
|
||||
- task_key: claims_preprocessing
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: core
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: data_quality
|
||||
depends_on:
|
||||
- task_key: core
|
||||
- task_key: readmissions
|
||||
- task_key: pharmacy
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: data_quality
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: hcc_suspecting
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: hcc_suspecting
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: input_layer
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: input_layer
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: main
|
||||
depends_on:
|
||||
- task_key: readmissions
|
||||
- task_key: pharmacy
|
||||
- task_key: hcc_suspecting
|
||||
- task_key: provider_attribution
|
||||
- task_key: quality_measures
|
||||
- task_key: cms_quality_measures
|
||||
- task_key: ahrq_measures
|
||||
- task_key: data_quality
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: main
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: pharmacy
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: pharmacy
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: provider_attribution
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: provider_attribution
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: quality_measures
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: quality_measures
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
- task_key: readmissions
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
named_parameters:
|
||||
run: readmissions
|
||||
--target: spark
|
||||
--catalog: ${var.catalog}
|
||||
--save: ''
|
||||
libraries:
|
||||
- whl: dbfs:/FileStore/wheels/stack-latest.whl
|
||||
schemas:
|
||||
ahrq_measures:
|
||||
name: ahrq_measures
|
||||
catalog_name: ${var.catalog}
|
||||
cclf:
|
||||
name: cclf
|
||||
catalog_name: ${var.catalog}
|
||||
claims_preprocessing:
|
||||
name: claims_preprocessing
|
||||
catalog_name: ${var.catalog}
|
||||
cms_quality_measures:
|
||||
name: cms_quality_measures
|
||||
catalog_name: ${var.catalog}
|
||||
core:
|
||||
name: core
|
||||
catalog_name: ${var.catalog}
|
||||
data_quality:
|
||||
name: data_quality
|
||||
catalog_name: ${var.catalog}
|
||||
hcc_suspecting:
|
||||
name: hcc_suspecting
|
||||
catalog_name: ${var.catalog}
|
||||
input_layer:
|
||||
name: input_layer
|
||||
catalog_name: ${var.catalog}
|
||||
main:
|
||||
name: main
|
||||
catalog_name: ${var.catalog}
|
||||
pharmacy:
|
||||
name: pharmacy
|
||||
catalog_name: ${var.catalog}
|
||||
provider_attribution:
|
||||
name: provider_attribution
|
||||
catalog_name: ${var.catalog}
|
||||
quality_measures:
|
||||
name: quality_measures
|
||||
catalog_name: ${var.catalog}
|
||||
readmissions:
|
||||
name: readmissions
|
||||
catalog_name: ${var.catalog}
|
||||
volumes:
|
||||
staging:
|
||||
name: staging
|
||||
catalog_name: ${var.catalog}
|
||||
schema_name: default
|
||||
volume_type: MANAGED
|
||||
|
||||
@@ -238,6 +238,9 @@ _DEFAULT_PLUGINS: list[DabPlugin] = [
|
||||
VolumesPlugin(),
|
||||
]
|
||||
|
||||
# Plugins used only for SQL generation (--dab-sql flag)
|
||||
_SQL_PLUGINS: list[DabPlugin] = [SQLPlugin()]
|
||||
|
||||
|
||||
def _get_catalog_schemas() -> list[str]:
|
||||
"""Return only schemas produced by registered pipelines."""
|
||||
@@ -249,54 +252,44 @@ def _get_catalog_schemas() -> list[str]:
|
||||
return []
|
||||
|
||||
|
||||
def emit(cfg_data: dict) -> dict[str, str]:
|
||||
"""Generate databricks.yml + SQL files from config + pipeline registry.
|
||||
def emit(cfg_data: dict, *, sql: bool = False) -> dict[str, str]:
|
||||
"""Generate databricks.yml from config + pipeline registry.
|
||||
|
||||
The generated databricks.yml delegates resource definitions to
|
||||
``bundle/resources/__init__.py`` via the ``python.resources`` key.
|
||||
The YAML only contains bundle metadata, targets, and the Python
|
||||
entry point — all resources (jobs, schemas, volumes) are defined
|
||||
in typed Python code.
|
||||
|
||||
When ``sql=True``, also generates transpiled SQL files via SQLPlugin
|
||||
(for audit/validation, not deployment).
|
||||
|
||||
Returns {relative_path: content} for all generated files.
|
||||
"""
|
||||
from aco.pipe import registry
|
||||
|
||||
dab_cfg = cfg_data.get("databricks", {})
|
||||
bundle_name = dab_cfg.get("bundle_name", "stack")
|
||||
targets_cfg = dab_cfg.get("targets", {})
|
||||
catalog_schemas = _get_catalog_schemas()
|
||||
|
||||
# Collect all files from plugins
|
||||
all_files: dict[str, str] = {}
|
||||
all_resources: dict = {}
|
||||
|
||||
for plugin in _DEFAULT_PLUGINS:
|
||||
# Resource contributions
|
||||
plugin_resources = plugin.resources(dab_cfg, registry, catalog_schemas)
|
||||
for rtype, rdefs in plugin_resources.items():
|
||||
if rtype not in all_resources:
|
||||
all_resources[rtype] = OrderedDict()
|
||||
all_resources[rtype].update(rdefs)
|
||||
# SQL generation (on demand only)
|
||||
if sql:
|
||||
from aco.pipe import registry
|
||||
|
||||
# File contributions (SQL, DDL)
|
||||
plugin_files = plugin.files(dab_cfg, registry)
|
||||
all_files.update(plugin_files)
|
||||
for plugin in _SQL_PLUGINS:
|
||||
plugin_files = plugin.files(dab_cfg, registry)
|
||||
all_files.update(plugin_files)
|
||||
|
||||
# Build bundle YAML
|
||||
# Build minimal bundle YAML — resources come from Python
|
||||
bundle = OrderedDict()
|
||||
bundle["bundle"] = OrderedDict([("name", bundle_name)])
|
||||
|
||||
bundle["sync"] = OrderedDict(
|
||||
[
|
||||
("include", ["bundle/**"]),
|
||||
]
|
||||
)
|
||||
|
||||
bundle["variables"] = OrderedDict(
|
||||
bundle["python"] = OrderedDict(
|
||||
[
|
||||
("venv_path", ".venv"),
|
||||
(
|
||||
"catalog",
|
||||
OrderedDict(
|
||||
[
|
||||
("description", "Unity Catalog name for table references"),
|
||||
("default", targets_cfg.get("prod", {}).get("catalog", "aco")),
|
||||
]
|
||||
),
|
||||
"resources",
|
||||
["resources:load_resources"],
|
||||
),
|
||||
]
|
||||
)
|
||||
@@ -316,8 +309,6 @@ def emit(cfg_data: dict) -> dict[str, str]:
|
||||
("catalog", tcfg.get("catalog", "aco")),
|
||||
]
|
||||
)
|
||||
if tcfg.get("warehouse_id"):
|
||||
target["variables"]["warehouse_id"] = tcfg["warehouse_id"]
|
||||
if tcfg.get("mode") == "production":
|
||||
run_as = dab_cfg.get("run_as", {})
|
||||
if run_as.get("service_principal_name"):
|
||||
@@ -329,13 +320,5 @@ def emit(cfg_data: dict) -> dict[str, str]:
|
||||
targets[tname] = target
|
||||
bundle["targets"] = targets
|
||||
|
||||
# Permissions
|
||||
perms = dab_cfg.get("permissions", [])
|
||||
if perms:
|
||||
bundle["permissions"] = perms
|
||||
|
||||
# Resources
|
||||
bundle["resources"] = all_resources
|
||||
|
||||
all_files["databricks.yml"] = _HEADER + "\n" + _dump(dict(bundle))
|
||||
return all_files
|
||||
|
||||
@@ -140,11 +140,7 @@ def _emit_dab(*, sql: bool = False) -> dict[str, str]:
|
||||
return {}
|
||||
from backends.databricks import emit
|
||||
|
||||
all_files = emit(cfg._data)
|
||||
if not sql:
|
||||
# Return only the YAML — skip bundle/sql/ and bundle/ddl/
|
||||
return {k: v for k, v in all_files.items() if not k.startswith("bundle/")}
|
||||
return all_files
|
||||
return emit(cfg._data, sql=sql)
|
||||
|
||||
|
||||
def generate(backend: str | None = None, *, dab_sql: bool = False) -> dict[str, str]:
|
||||
|
||||
@@ -82,8 +82,8 @@ perf = [
|
||||
]
|
||||
lake = [
|
||||
"stack[aco]",
|
||||
"databricks-cli>=0.18.0",
|
||||
"databricks-sdk>=0.85.0",
|
||||
"databricks-bundles>=0.295.0",
|
||||
]
|
||||
aws = [
|
||||
"boto3>=1.35.0",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Tests for the Databricks Asset Bundle generator.
|
||||
|
||||
Covers: round-trip, completeness, dependency ordering, idempotency,
|
||||
target isolation, plugin extensibility, SQL transpilation, and DDL.
|
||||
Covers: YAML structure, Python resource loader, pipeline metadata,
|
||||
dependency ordering, schema filtering, idempotency, and targets.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -16,7 +16,7 @@ sys.path.insert(0, str(ROOT / "dev" / "scripts"))
|
||||
|
||||
|
||||
class TestPipelineMetadata:
|
||||
"""#170 — Pipeline.upstream and cluster_profile."""
|
||||
"""Pipeline.upstream and cluster_profile."""
|
||||
|
||||
def test_all_pipelines_have_upstream(self):
|
||||
from aco.pipe import registry
|
||||
@@ -68,18 +68,15 @@ class TestPipelineMetadata:
|
||||
assert registry["input_layer"].upstream == []
|
||||
|
||||
|
||||
class TestDabEmitter:
|
||||
"""#172 — Core generator produces valid databricks.yml + SQL files."""
|
||||
class TestDabYaml:
|
||||
"""Generated databricks.yml has correct structure."""
|
||||
|
||||
def _emit(self) -> dict[str, str]:
|
||||
def _yaml(self) -> dict:
|
||||
from backends.databricks import emit
|
||||
|
||||
from conf import cfg
|
||||
|
||||
return emit(cfg._data)
|
||||
|
||||
def _yaml(self) -> dict:
|
||||
files = self._emit()
|
||||
files = emit(cfg._data)
|
||||
return yaml.safe_load(files["databricks.yml"])
|
||||
|
||||
def test_generates_valid_yaml(self):
|
||||
@@ -88,64 +85,112 @@ class TestDabEmitter:
|
||||
|
||||
def test_has_required_top_level_keys(self):
|
||||
data = self._yaml()
|
||||
for key in ("bundle", "variables", "workspace", "targets", "resources"):
|
||||
for key in ("bundle", "python", "workspace", "targets"):
|
||||
assert key in data, f"Missing top-level key: {key}"
|
||||
|
||||
def test_bundle_name(self):
|
||||
assert self._yaml()["bundle"]["name"] == "stack"
|
||||
|
||||
def test_all_registry_pipelines_appear_as_tasks(self):
|
||||
def test_python_resources_entry_point(self):
|
||||
data = self._yaml()
|
||||
assert "resources" in data["python"]
|
||||
assert "resources:load_resources" in data["python"]["resources"]
|
||||
|
||||
def test_no_inline_resources(self):
|
||||
"""Resources come from Python, not inline YAML."""
|
||||
data = self._yaml()
|
||||
assert "resources" not in data or data.get("resources") is None
|
||||
|
||||
def test_no_variables_in_yaml(self):
|
||||
"""Variables defined in Python resource loader, not YAML."""
|
||||
data = self._yaml()
|
||||
assert "variables" not in data
|
||||
|
||||
|
||||
class TestPythonResources:
|
||||
"""Python resource loader produces correct typed resources."""
|
||||
|
||||
def _resources(self):
|
||||
from bundle.resources import load_resources
|
||||
|
||||
return load_resources()
|
||||
|
||||
def test_load_resources_returns_resources(self):
|
||||
from databricks.bundles.core import Resources
|
||||
|
||||
r = self._resources()
|
||||
assert isinstance(r, Resources)
|
||||
|
||||
def test_has_job(self):
|
||||
r = self._resources()
|
||||
assert "stack_pipelines" in r.jobs
|
||||
|
||||
def test_job_has_all_pipeline_tasks(self):
|
||||
from aco.pipe import registry
|
||||
|
||||
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
|
||||
task_keys = {t["task_key"] for t in tasks}
|
||||
r = self._resources()
|
||||
job = r.jobs["stack_pipelines"]
|
||||
task_keys = {t.task_key for t in job.tasks}
|
||||
for name in registry:
|
||||
assert name in task_keys, f"Pipeline '{name}' missing from tasks"
|
||||
|
||||
def test_task_count_matches_registry(self):
|
||||
from aco.pipe import registry
|
||||
|
||||
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
|
||||
assert len(tasks) == len(registry)
|
||||
r = self._resources()
|
||||
assert len(r.jobs["stack_pipelines"].tasks) == len(registry)
|
||||
|
||||
def test_tasks_use_python_wheel_task(self):
|
||||
r = self._resources()
|
||||
for task in r.jobs["stack_pipelines"].tasks:
|
||||
assert task.python_wheel_task is not None, (
|
||||
f"Task {task.task_key} missing python_wheel_task"
|
||||
)
|
||||
assert task.python_wheel_task.package_name == "stack"
|
||||
assert task.python_wheel_task.entry_point == "cli"
|
||||
|
||||
def test_tasks_have_libraries(self):
|
||||
r = self._resources()
|
||||
for task in r.jobs["stack_pipelines"].tasks:
|
||||
assert task.libraries, f"Task {task.task_key} missing libraries"
|
||||
|
||||
def test_dependencies_match_upstream(self):
|
||||
from aco.pipe import registry
|
||||
|
||||
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
|
||||
task_map = {t["task_key"]: t for t in tasks}
|
||||
r = self._resources()
|
||||
task_map = {t.task_key: t for t in r.jobs["stack_pipelines"].tasks}
|
||||
for name, pipe in registry.items():
|
||||
task = task_map[name]
|
||||
if pipe.upstream:
|
||||
deps = [d["task_key"] for d in task.get("depends_on", [])]
|
||||
assert sorted(deps) == sorted(pipe.upstream)
|
||||
deps = sorted(d.task_key for d in task.depends_on)
|
||||
assert deps == sorted(pipe.upstream)
|
||||
else:
|
||||
assert "depends_on" not in task
|
||||
assert not task.depends_on
|
||||
|
||||
def test_tasks_use_python_wheel_task(self):
|
||||
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
|
||||
for task in tasks:
|
||||
assert "python_wheel_task" in task, (
|
||||
f"Task {task['task_key']} missing python_wheel_task"
|
||||
def test_schemas_count(self):
|
||||
r = self._resources()
|
||||
assert len(r.schemas) == 13, f"Expected 13 schemas, got {len(r.schemas)}"
|
||||
|
||||
def test_schemas_include_core(self):
|
||||
r = self._resources()
|
||||
assert "core" in r.schemas
|
||||
|
||||
def test_no_legacy_schemas(self):
|
||||
r = self._resources()
|
||||
for legacy in ("alr", "ccsr", "ccw", "pfs", "ssp", "reach"):
|
||||
assert legacy not in r.schemas, (
|
||||
f"Legacy schema '{legacy}' should not be deployed"
|
||||
)
|
||||
pwt = task["python_wheel_task"]
|
||||
assert pwt["package_name"] == "stack"
|
||||
assert pwt["entry_point"] == "cli"
|
||||
assert "run" in pwt["named_parameters"]
|
||||
assert "${var.catalog}" in str(pwt["named_parameters"])
|
||||
|
||||
def test_tasks_have_libraries(self):
|
||||
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
|
||||
for task in tasks:
|
||||
assert "libraries" in task, f"Task {task['task_key']} missing libraries"
|
||||
def test_has_staging_volume(self):
|
||||
r = self._resources()
|
||||
assert "staging" in r.volumes
|
||||
|
||||
def test_no_sql_tasks(self):
|
||||
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
|
||||
for task in tasks:
|
||||
assert "sql_task" not in task, f"Task {task['task_key']} still has sql_task"
|
||||
|
||||
def test_no_warehouse_id_variable(self):
|
||||
data = self._yaml()
|
||||
assert "warehouse_id" not in data.get("variables", {})
|
||||
def test_job_has_schedule(self):
|
||||
r = self._resources()
|
||||
job = r.jobs["stack_pipelines"]
|
||||
assert job.schedule is not None
|
||||
assert job.schedule.quartz_cron_expression == "0 0 6 * * ?"
|
||||
|
||||
|
||||
class TestTargets:
|
||||
@@ -177,68 +222,27 @@ class TestTargets:
|
||||
)
|
||||
|
||||
|
||||
class TestSchemas:
|
||||
"""Unity Catalog schemas from table catalog."""
|
||||
|
||||
def test_schemas_generated(self):
|
||||
from backends.databricks import emit
|
||||
|
||||
from conf import cfg
|
||||
|
||||
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
|
||||
schemas = data["resources"].get("schemas", {})
|
||||
assert len(schemas) == 13, f"Expected 13 schemas, got {len(schemas)}"
|
||||
assert "core" in schemas
|
||||
# No legacy schemas
|
||||
for legacy in ("alr", "ccsr", "ccw", "pfs", "ssp", "reach"):
|
||||
assert legacy not in schemas, (
|
||||
f"Legacy schema '{legacy}' should not be deployed"
|
||||
)
|
||||
|
||||
def test_schema_references_catalog_variable(self):
|
||||
from backends.databricks import emit
|
||||
|
||||
from conf import cfg
|
||||
|
||||
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
|
||||
for name, schema in data["resources"]["schemas"].items():
|
||||
assert schema["catalog_name"] == "${var.catalog}"
|
||||
|
||||
|
||||
class TestIdempotency:
|
||||
"""Running the generator twice produces identical output."""
|
||||
|
||||
def test_idempotent(self):
|
||||
def test_yaml_idempotent(self):
|
||||
from backends.databricks import emit
|
||||
|
||||
from conf import cfg
|
||||
|
||||
first = emit(cfg._data)
|
||||
second = emit(cfg._data)
|
||||
assert first.keys() == second.keys()
|
||||
assert first["databricks.yml"] == second["databricks.yml"]
|
||||
|
||||
def test_resources_idempotent(self):
|
||||
"""Python resource loader produces identical structure twice."""
|
||||
from bundle.resources import load_resources
|
||||
|
||||
class TestPluginArchitecture:
|
||||
"""Plugin extensibility."""
|
||||
|
||||
def test_custom_plugin(self):
|
||||
from backends.databricks import _DEFAULT_PLUGINS, DabPlugin, emit
|
||||
|
||||
from conf import cfg
|
||||
|
||||
class TestPlugin(DabPlugin):
|
||||
def resources(self, cfg, registry, catalog_schemas):
|
||||
return {"experiments": {"test_exp": {"name": "test-experiment"}}}
|
||||
|
||||
def files(self, cfg, registry):
|
||||
return {"bundle/test/hello.sql": "SELECT 1"}
|
||||
|
||||
_DEFAULT_PLUGINS.append(TestPlugin())
|
||||
try:
|
||||
result = emit(cfg._data)
|
||||
data = yaml.safe_load(result["databricks.yml"])
|
||||
assert "experiments" in data["resources"]
|
||||
assert "bundle/test/hello.sql" in result
|
||||
finally:
|
||||
_DEFAULT_PLUGINS.pop()
|
||||
first = load_resources()
|
||||
second = load_resources()
|
||||
# Compare job task keys as a proxy for structural equality
|
||||
t1 = [t.task_key for t in first.jobs["stack_pipelines"].tasks]
|
||||
t2 = [t.task_key for t in second.jobs["stack_pipelines"].tasks]
|
||||
assert t1 == t2
|
||||
assert set(first.schemas.keys()) == set(second.schemas.keys())
|
||||
assert set(first.volumes.keys()) == set(second.volumes.keys())
|
||||
|
||||
32
uv.lock
generated
32
uv.lock
generated
@@ -642,21 +642,12 @@ wheels = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "databricks-cli"
|
||||
version = "0.18.0"
|
||||
name = "databricks-bundles"
|
||||
version = "0.295.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
{ name = "oauthlib" },
|
||||
{ name = "pyjwt" },
|
||||
{ name = "requests" },
|
||||
{ name = "six" },
|
||||
{ name = "tabulate" },
|
||||
{ name = "urllib3" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c7/6c/ee7a98f22ba6e4d39cdf7f3a7cd9461fcc562625ccbca58d94bf35fe598b/databricks-cli-0.18.0.tar.gz", hash = "sha256:87569709eda9af3e9db8047b691e420b5e980c62ef01675575c0d2b9b4211eb7", size = 95375, upload-time = "2023-10-05T09:36:27.75Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/a1/b1/75b890d39b3db744dca14b61c3ec9d26996f225c86481cb0d1220e19bc9a/databricks_bundles-0.295.0.tar.gz", hash = "sha256:d895cc4c4f31e99f8a6678790231cb035afa4a83d0ba3fc0d4de18c42ccae921", size = 87457, upload-time = "2026-03-18T14:21:43.165Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ae/a3/d56f8382c40899301f327d1c881278b09c9b8bc301c2c111633a0346d06e/databricks_cli-0.18.0-py2.py3-none-any.whl", hash = "sha256:1176a5f42d3e8af4abfc915446fb23abc44513e325c436725f5898cbb9e3384b", size = 150329, upload-time = "2023-10-05T09:36:25.745Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d5/87/299719b0921b2313667514d38fd45a7c6ea9a638d018d2674193adad6550/databricks_bundles-0.295.0-py3-none-any.whl", hash = "sha256:e9bd64c53526d3ddaffd015d918c0ff846008d169e6eee1168791896b17a5a83", size = 178690, upload-time = "2026-03-18T14:21:41.661Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3101,7 +3092,7 @@ aco = [
|
||||
]
|
||||
all = [
|
||||
{ name = "cryptography" },
|
||||
{ name = "databricks-cli" },
|
||||
{ name = "databricks-bundles" },
|
||||
{ name = "databricks-sdk" },
|
||||
{ name = "duckdb" },
|
||||
{ name = "fastapi" },
|
||||
@@ -3182,7 +3173,7 @@ gcp = [
|
||||
{ name = "google-cloud-storage" },
|
||||
]
|
||||
lake = [
|
||||
{ name = "databricks-cli" },
|
||||
{ name = "databricks-bundles" },
|
||||
{ name = "databricks-sdk" },
|
||||
{ name = "duckdb" },
|
||||
{ name = "httpx" },
|
||||
@@ -3245,7 +3236,7 @@ requires-dist = [
|
||||
{ name = "azure-storage-blob", marker = "extra == 'azure'", specifier = ">=12.23.0" },
|
||||
{ name = "boto3", marker = "extra == 'aws'", specifier = ">=1.35.0" },
|
||||
{ name = "cryptography", marker = "extra == 'api'", specifier = ">=46.0.5" },
|
||||
{ name = "databricks-cli", marker = "extra == 'lake'", specifier = ">=0.18.0" },
|
||||
{ name = "databricks-bundles", marker = "extra == 'lake'", specifier = ">=0.295.0" },
|
||||
{ name = "databricks-sdk", marker = "extra == 'lake'", specifier = ">=0.85.0" },
|
||||
{ name = "duckdb", marker = "extra == 'aco'", specifier = ">=1.0.0" },
|
||||
{ name = "duckdb", marker = "extra == 'conf'", specifier = ">=1.0.0" },
|
||||
@@ -3343,15 +3334,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/81/0d/13d1d239a25cbfb19e740db83143e95c772a1fe10202dda4b76792b114dd/starlette-0.52.1-py3-none-any.whl", hash = "sha256:0029d43eb3d273bc4f83a08720b4912ea4b071087a3b48db01b7c839f7954d74", size = 74272, upload-time = "2026-01-18T13:34:09.188Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tabulate"
|
||||
version = "0.9.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/ec/fe/802052aecb21e3797b8f7902564ab6ea0d60ff8ca23952079064155d1ae1/tabulate-0.9.0.tar.gz", hash = "sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c", size = 81090, upload-time = "2022-10-06T17:21:48.54Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/40/44/4a5f08c96eb108af5cb50b41f76142f0afa346dfa99d5296fe7202a11854/tabulate-0.9.0-py3-none-any.whl", hash = "sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f", size = 35252, upload-time = "2022-10-06T17:21:44.262Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "text-unidecode"
|
||||
version = "1.3"
|
||||
|
||||
Reference in New Issue
Block a user