homelab/stack
1
0
Fork 0
Code Issues 110 Pull Requests 1 Actions Packages Projects 1 Releases Wiki Activity

Merge pull request 'feat: DAB generates transpiled SQL + DDL instead of wheel wrapper' (#181) from feat/dab-generator into main
All checks were successful
Deploy / build-scan-report (push) Successful in 3m55s
Details
CI / lint-test (push) Successful in 7m9s
Details

Browse Source
This commit was merged in pull request #181.
This commit is contained in:
kert
2026-03-24 06:21:29 +00:00
parent fbede316d9 6b02a804c8
commit c5d0f89c52
5 changed files with 380 additions and 404 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
.gitignore vendored
View File

@@ -8,6 +8,10 @@ zotero/profiles/
tuva/
logs/**
# Generated DAB SQL/DDL (regenerate via gen_config.py)
bundle/sql/
bundle/ddl/
# Marimo cache
notebooks/__marimo__/
.deploy_key

181
databricks.yml
View File

@@ -3,15 +3,15 @@
bundle:
name: stack
artifacts:
default:
type: whl
build: uv build
path: .
sync:
include:
- bundle/**
variables:
catalog:
description: Unity Catalog name for table references
default: aco
warehouse_id:
description: SQL warehouse ID for query execution
workspace:
host: ${DATABRICKS_HOST}
targets:
@@ -41,113 +41,71 @@ resources:
spark_version: 15.4.x-scala2.12
node_type_id: i3.xlarge
num_workers: 1
spark_conf:
spark.sql.catalog.${var.catalog}: org.apache.iceberg.spark.SparkCatalog
tasks:
- task_key: ahrq_measures
depends_on:
- task_key: core
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- ahrq_measures
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/ahrq_measures/
warehouse_id: ${var.warehouse_id}
- task_key: cclf
depends_on:
- task_key: input_layer
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- cclf
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/cclf/
warehouse_id: ${var.warehouse_id}
- task_key: claims_preprocessing
depends_on:
- task_key: input_layer
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- claims_preprocessing
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/claims_preprocessing/
warehouse_id: ${var.warehouse_id}
- task_key: cms_quality_measures
depends_on:
- task_key: core
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- cms_quality_measures
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/cms_quality_measures/
warehouse_id: ${var.warehouse_id}
- task_key: core
depends_on:
- task_key: claims_preprocessing
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- core
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/core/
warehouse_id: ${var.warehouse_id}
- 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
parameters:
- run
- data_quality
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/data_quality/
warehouse_id: ${var.warehouse_id}
- task_key: hcc_suspecting
depends_on:
- task_key: core
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- hcc_suspecting
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/hcc_suspecting/
warehouse_id: ${var.warehouse_id}
- task_key: input_layer
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- input_layer
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/input_layer/
warehouse_id: ${var.warehouse_id}
- task_key: main
depends_on:
- task_key: readmissions
@@ -159,67 +117,42 @@ resources:
- task_key: ahrq_measures
- task_key: data_quality
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- main
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/main/
warehouse_id: ${var.warehouse_id}
- task_key: pharmacy
depends_on:
- task_key: core
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- pharmacy
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/pharmacy/
warehouse_id: ${var.warehouse_id}
- task_key: provider_attribution
depends_on:
- task_key: core
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- provider_attribution
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/provider_attribution/
warehouse_id: ${var.warehouse_id}
- task_key: quality_measures
depends_on:
- task_key: core
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- quality_measures
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/quality_measures/
warehouse_id: ${var.warehouse_id}
- task_key: readmissions
depends_on:
- task_key: core
job_cluster_key: default_cluster
python_wheel_task:
package_name: stack
entry_point: cli
parameters:
- run
- readmissions
- --target
- databricks
- --save
sql_task:
file:
path: bundle/sql/readmissions/
warehouse_id: ${var.warehouse_id}
schemas:
ahrq_measures:
name: ahrq_measures

383
dev/scripts/backends/databricks.py
View File

@@ -1,20 +1,25 @@
"""Databricks Asset Bundle (DAB) emitter.
Generates databricks.yml from the pipeline registry, table catalog,
and stack.toml [databricks] config — no hardcoded pipeline names,
task lists, or dependencies.
Generates databricks.yml + transpiled SQL files from the pipeline
registry, table catalog, and stack.toml [databricks] config.
Each pipeline expression is transpiled to Databricks SQL via
``aco.lake.transpile`` and written to ``bundle/sql/<pipeline>/<expr>.sql``.
Table DDL is generated from ``aco.table`` Pydantic models.
The bundle YAML references these SQL files as ``sql_task`` definitions —
Databricks executes native SQL, not a Python wheel wrapper.
Follows the same emit() → dict[str, str] pattern as gitea.py / github.py.
"""
from __future__ import annotations
import yaml
from collections import OrderedDict
import yaml
# ── YAML helpers ──────────────────────────────────────────────────
# Use OrderedDict + custom representer to control key order in output
yaml.add_representer(
OrderedDict,
lambda dumper, data: dumper.represent_mapping(
@@ -24,13 +29,9 @@ yaml.add_representer(
def _dump(data: dict) -> str:
"""Dump to YAML with clean formatting."""
return yaml.dump(
data,
default_flow_style=False,
sort_keys=False,
width=120,
allow_unicode=True,
data, default_flow_style=False, sort_keys=False,
width=120, allow_unicode=True,
)
@@ -38,131 +39,177 @@ def _dump(data: dict) -> str:
class DabPlugin:
"""Protocol for DAB resource plugins.
"""Protocol for DAB resource plugins."""
Subclass and implement resources() to contribute resource blocks.
"""
def resources(self, cfg: dict, registry: dict, catalog_schemas: list[str]) -> dict:
return {}
def resources(
self, cfg: dict, registry: dict, catalog_schemas: list[str]
) -> dict:
"""Return a dict of resource type -> resource definitions."""
def files(self, cfg: dict, registry: dict) -> dict[str, str]:
"""Return {relative_path: content} for extra files."""
return {}
# ── Built-in plugins ─────────────────────────────────────────────
class SQLPlugin(DabPlugin):
"""Transpile every pipeline expression to Databricks SQL files."""
def files(self, cfg: dict, registry: dict) -> dict[str, str]:
from aco.lake.transpile import transpile
result: dict[str, str] = {}
for pipe_name, pipe in sorted(registry.items()):
try:
sql_map = transpile(
pipe,
target_dialect="databricks",
catalog="${var.catalog}",
output_mode="ctas",
)
except Exception as e:
result[f"bundle/sql/{pipe_name}/_error.txt"] = (
f"Transpile failed: {e}"
)
continue
for expr_name, sql in sql_map.items():
# expr_name is like "core._stg_claims_member_months"
# filename: bundle/sql/core/_stg_claims_member_months.sql
schema, table = expr_name.split(".", 1)
path = f"bundle/sql/{pipe_name}/{table}.sql"
header = (
f"-- {expr_name}\n"
f"-- Generated by gen_config.py — DO NOT EDIT\n\n"
)
result[path] = header + sql + "\n"
return result
class DDLPlugin(DabPlugin):
"""Generate CREATE TABLE DDL from SQLTable models for each output."""
def files(self, cfg: dict, registry: dict) -> dict[str, str]:
from aco.lake.catalog import Catalog
cat = Catalog()
result: dict[str, str] = {}
for schema_name in cat.schemas():
for table_ref in cat.tables(schema_name):
model = cat.model(table_ref)
schema, table = table_ref.split(".", 1)
# Generate DDL with catalog variable
ddl = model.to_ddl()
# Rewrite the table ref to include catalog variable
ddl = ddl.replace(
f"CREATE TABLE {table_ref}",
f"CREATE TABLE IF NOT EXISTS ${{var.catalog}}.{table_ref}",
)
path = f"bundle/ddl/{schema}/{table}.sql"
result[path] = (
f"-- {table_ref}\n"
f"-- Generated by gen_config.py — DO NOT EDIT\n\n"
f"{ddl}\n"
)
return result
class JobsPlugin(DabPlugin):
"""Generate jobs from the pipeline registry."""
"""Generate jobs with sql_task references to transpiled SQL files."""
def resources(self, cfg: dict, registry: dict, catalog_schemas: list) -> dict:
package_name = cfg.get("package_name", "stack")
entry_point = cfg.get("entry_point", "cli")
# Collect which cluster profiles are actually used
used_profiles: set[str] = set()
for pipe in registry.values():
used_profiles.add(getattr(pipe, "cluster_profile", "default"))
# Build job_clusters from config — only those referenced
# Build job_clusters
clusters_cfg = cfg.get("clusters", {"default": {}})
job_clusters = []
for profile_name in sorted(used_profiles):
cluster_spec = dict(clusters_cfg.get(profile_name, {}))
spark_conf = OrderedDict()
spark_conf[
"spark.sql.catalog.${var.catalog}"
] = "org.apache.iceberg.spark.SparkCatalog"
cluster_spec["spark_conf"] = spark_conf
job_clusters.append(
OrderedDict(
[
("job_cluster_key", f"{profile_name}_cluster"),
(
"new_cluster",
OrderedDict(
[
(
"spark_version",
cluster_spec.get(
"spark_version", "15.4.x-scala2.12"
),
),
(
"node_type_id",
cluster_spec.get("node_type_id", "i3.xlarge"),
),
(
"num_workers",
cluster_spec.get("num_workers", 1),
),
("spark_conf", spark_conf),
]
),
),
]
)
)
job_clusters.append(OrderedDict([
("job_cluster_key", f"{profile_name}_cluster"),
("new_cluster", OrderedDict([
("spark_version", cluster_spec.get("spark_version", "15.4.x-scala2.12")),
("node_type_id", cluster_spec.get("node_type_id", "i3.xlarge")),
("num_workers", cluster_spec.get("num_workers", 1)),
])),
]))
# Build tasks from registry
tasks = []
for name, pipe in sorted(registry.items()):
# DDL tasks — one per schema to create tables before SQL runs
ddl_tasks = []
for schema_name in sorted(catalog_schemas):
ddl_tasks.append(OrderedDict([
("task_key", f"ddl_{schema_name}"),
("job_cluster_key", "default_cluster"),
("sql_task", OrderedDict([
("file", OrderedDict([
("path", f"bundle/ddl/{schema_name}/"),
])),
])),
]))
# Pipeline tasks — each expression becomes a sql_task
pipeline_tasks = []
for pipe_name, pipe in sorted(registry.items()):
profile = getattr(pipe, "cluster_profile", "default")
task = OrderedDict()
task["task_key"] = name
if pipe.upstream:
task["depends_on"] = [
OrderedDict([("task_key", dep)]) for dep in pipe.upstream
]
task["job_cluster_key"] = f"{profile}_cluster"
task["python_wheel_task"] = OrderedDict(
[
("package_name", package_name),
("entry_point", entry_point),
(
"parameters",
["run", name, "--target", "databricks", "--save"],
),
]
)
tasks.append(task)
# Build the job
# Build depends_on from upstream pipeline names
depends_on = []
for dep in pipe.upstream:
depends_on.append(OrderedDict([("task_key", dep)]))
task = OrderedDict()
task["task_key"] = pipe_name
if depends_on:
task["depends_on"] = depends_on
task["job_cluster_key"] = f"{profile}_cluster"
# Each pipeline task runs all its SQL files in order
# We use a for_each or sequential notebook; simplest is
# a single SQL file per pipeline that sources all expressions
task["sql_task"] = OrderedDict([
("file", OrderedDict([
("path", f"bundle/sql/{pipe_name}/"),
])),
("warehouse_id", "${var.warehouse_id}"),
])
pipeline_tasks.append(task)
schedule = cfg.get("schedule", "0 0 6 * * ?")
timezone = cfg.get("timezone", "America/New_York")
job = OrderedDict()
job["name"] = f"{cfg.get('bundle_name', 'stack')}-pipelines"
job["description"] = (
f"Run all {len(registry)} ACO pipelines in dependency order"
)
job["schedule"] = OrderedDict(
[
("quartz_cron_expression", schedule),
("timezone_id", timezone),
]
)
job["description"] = f"Run all {len(registry)} ACO pipelines in dependency order"
job["schedule"] = OrderedDict([
("quartz_cron_expression", schedule),
("timezone_id", timezone),
])
job["job_clusters"] = job_clusters
job["tasks"] = tasks
job["tasks"] = pipeline_tasks
return {"jobs": OrderedDict([("stack_pipelines", job)])}
class SchemasPlugin(DabPlugin):
"""Generate Unity Catalog schema resources from the table catalog."""
"""Generate Unity Catalog schema resources."""
def resources(self, cfg: dict, registry: dict, catalog_schemas: list) -> dict:
if not catalog_schemas:
return {}
schemas = OrderedDict()
for schema_name in sorted(catalog_schemas):
schemas[schema_name] = OrderedDict(
[
("name", schema_name),
("catalog_name", "${var.catalog}"),
]
)
schemas[schema_name] = OrderedDict([
("name", schema_name),
("catalog_name", "${var.catalog}"),
])
return {"schemas": schemas}
@@ -170,23 +217,14 @@ class VolumesPlugin(DabPlugin):
"""Generate a staging volume for data upload."""
def resources(self, cfg: dict, registry: dict, catalog_schemas: list) -> dict:
return {
"volumes": OrderedDict(
[
(
"staging",
OrderedDict(
[
("name", "staging"),
("catalog_name", "${var.catalog}"),
("schema_name", "default"),
("volume_type", "MANAGED"),
]
),
)
]
)
}
return {"volumes": OrderedDict([
("staging", OrderedDict([
("name", "staging"),
("catalog_name", "${var.catalog}"),
("schema_name", "default"),
("volume_type", "MANAGED"),
])),
])}
# ── Core generator ───────────────────────────────────────────────
@@ -196,8 +234,9 @@ _HEADER = (
"# Re-generate: uv run python dev/scripts/gen_config.py\n"
)
# Default plugins
_DEFAULT_PLUGINS: list[DabPlugin] = [
SQLPlugin(),
DDLPlugin(),
JobsPlugin(),
SchemasPlugin(),
VolumesPlugin(),
@@ -205,87 +244,63 @@ _DEFAULT_PLUGINS: list[DabPlugin] = [
def _get_catalog_schemas() -> list[str]:
"""Discover schemas from the table catalog."""
try:
from aco.lake.catalog import Catalog
cat = Catalog()
return cat.schemas()
return Catalog().schemas()
except Exception:
return []
def emit(cfg_data: dict) -> dict[str, str]:
"""Generate databricks.yml from config + pipeline registry.
"""Generate databricks.yml + SQL files from config + pipeline registry.
Returns {relative_path: content}.
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()
# Merge resources from all plugins
# 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)
# Build bundle structure
bundle = OrderedDict()
# File contributions (SQL, DDL)
plugin_files = plugin.files(dab_cfg, registry)
all_files.update(plugin_files)
# bundle:
# Build bundle YAML
bundle = OrderedDict()
bundle["bundle"] = OrderedDict([("name", bundle_name)])
# artifacts:
bundle["artifacts"] = OrderedDict(
[
(
"default",
OrderedDict(
[
("type", "whl"),
("build", "uv build"),
("path", "."),
]
),
)
]
)
# Sync SQL files to workspace
bundle["sync"] = OrderedDict([
("include", ["bundle/**"]),
])
# variables:
bundle["variables"] = OrderedDict(
[
(
"catalog",
OrderedDict(
[
("description", "Unity Catalog name for table references"),
(
"default",
dab_cfg.get(
"default_catalog",
targets_cfg.get("prod", {}).get("catalog", "aco"),
),
),
]
),
)
]
)
bundle["variables"] = OrderedDict([
("catalog", OrderedDict([
("description", "Unity Catalog name for table references"),
("default", targets_cfg.get("prod", {}).get("catalog", "aco")),
])),
("warehouse_id", OrderedDict([
("description", "SQL warehouse ID for query execution"),
])),
])
# workspace:
bundle["workspace"] = OrderedDict(
[("host", "${DATABRICKS_HOST}")]
)
bundle["workspace"] = OrderedDict([("host", "${DATABRICKS_HOST}")])
# targets:
# Targets
targets = OrderedDict()
for tname, tcfg in sorted(targets_cfg.items()):
target = OrderedDict()
@@ -293,31 +308,27 @@ def emit(cfg_data: dict) -> dict[str, str]:
target["mode"] = tcfg["mode"]
if tcfg.get("default"):
target["default"] = True
target["variables"] = OrderedDict(
[("catalog", tcfg.get("catalog", "aco"))]
)
# Production targets get run_as
target["variables"] = OrderedDict([
("catalog", tcfg.get("catalog", "aco")),
])
if tcfg.get("warehouse_id"):
target["variables"]["warehouse_id"] = tcfg["warehouse_id"]
if tcfg.get("mode") == "production":
run_as_cfg = dab_cfg.get("run_as", {})
if run_as_cfg.get("service_principal_name"):
target["run_as"] = OrderedDict(
[
(
"service_principal_name",
run_as_cfg["service_principal_name"],
)
]
)
run_as = dab_cfg.get("run_as", {})
if run_as.get("service_principal_name"):
target["run_as"] = OrderedDict([
("service_principal_name", run_as["service_principal_name"]),
])
targets[tname] = target
bundle["targets"] = targets
# permissions:
perms_cfg = dab_cfg.get("permissions", [])
if perms_cfg:
bundle["permissions"] = perms_cfg
# Permissions
perms = dab_cfg.get("permissions", [])
if perms:
bundle["permissions"] = perms
# resources:
# Resources
bundle["resources"] = all_resources
content = _HEADER + "\n" + _dump(dict(bundle))
return {"databricks.yml": content}
all_files["databricks.yml"] = _HEADER + "\n" + _dump(dict(bundle))
return all_files

26
dev/scripts/gen_config.py
View File

@@ -129,23 +129,32 @@ def _clean_stale(backend: str, generated_files: dict[str, str]) -> list[str]:
# ── Main ─────────────────────────────────────────────────────────
def _emit_dab() -> dict[str, str]:
"""Generate databricks.yml from pipeline registry + stack.toml."""
def _emit_dab(*, sql: bool = False) -> dict[str, str]:
"""Generate databricks.yml (and optionally SQL/DDL) from pipeline registry.
When *sql* is False (default), only databricks.yml is emitted — fast
enough for --check in CI. When True, also transpiles all pipeline
expressions to SQL and generates DDL (1500+ files, ~4 minutes).
"""
if "databricks" not in cfg._data:
return {}
from backends.databricks import emit
return emit(cfg._data)
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
def generate(backend: str | None = None) -> dict[str, str]:
def generate(backend: str | None = None, *, dab_sql: bool = False) -> dict[str, str]:
if backend is None:
backend = cfg._data.get("ci", {}).get("backend", "woodpecker")
files: dict[str, str] = {}
files.update(gen_coredns())
files.update(_emit_backend(backend))
files.update(_emit_dab())
files.update(_emit_dab(sql=dab_sql))
return files
@@ -180,9 +189,14 @@ def main() -> int:
parser.add_argument(
"--backend", default=None, help="Override CI backend (woodpecker|github)"
)
parser.add_argument(
"--dab-sql",
action="store_true",
help="Also generate bundle/sql/ and bundle/ddl/ (slow, ~4 min)",
)
args = parser.parse_args()
files = generate(args.backend)
files = generate(args.backend, dab_sql=args.dab_sql)
if args.check:
diffs = check_all(files)

190
tests/test_dab_generator.py
View File

@@ -1,7 +1,7 @@
"""Tests for the Databricks Asset Bundle generator.
Covers: round-trip, completeness, dependency ordering, idempotency,
target isolation, plugin extensibility, and pipeline metadata.
target isolation, plugin extensibility, SQL transpilation, and DDL.
"""
from __future__ import annotations
@@ -69,35 +69,35 @@ class TestPipelineMetadata:
class TestDabEmitter:
"""#172 — Core generator produces valid databricks.yml."""
"""#172 — Core generator produces valid databricks.yml + SQL files."""
def _generate(self) -> str:
def _emit(self) -> dict[str, str]:
from backends.databricks import emit
from conf import cfg
files = emit(cfg._data)
return files["databricks.yml"]
return emit(cfg._data)
def _yaml(self) -> dict:
files = self._emit()
return yaml.safe_load(files["databricks.yml"])
def test_generates_valid_yaml(self):
content = self._generate()
data = yaml.safe_load(content)
data = self._yaml()
assert isinstance(data, dict)
def test_has_required_top_level_keys(self):
data = yaml.safe_load(self._generate())
data = self._yaml()
for key in ("bundle", "variables", "workspace", "targets", "resources"):
assert key in data, f"Missing top-level key: {key}"
def test_bundle_name(self):
data = yaml.safe_load(self._generate())
assert data["bundle"]["name"] == "stack"
assert self._yaml()["bundle"]["name"] == "stack"
def test_all_registry_pipelines_appear_as_tasks(self):
from aco.pipe import registry
data = yaml.safe_load(self._generate())
tasks = data["resources"]["jobs"]["stack_pipelines"]["tasks"]
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
task_keys = {t["task_key"] for t in tasks}
for name in registry:
assert name in task_keys, f"Pipeline '{name}' missing from tasks"
@@ -105,44 +105,93 @@ class TestDabEmitter:
def test_task_count_matches_registry(self):
from aco.pipe import registry
data = yaml.safe_load(self._generate())
tasks = data["resources"]["jobs"]["stack_pipelines"]["tasks"]
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
assert len(tasks) == len(registry)
def test_dependencies_match_upstream(self):
from aco.pipe import registry
data = yaml.safe_load(self._generate())
tasks = data["resources"]["jobs"]["stack_pipelines"]["tasks"]
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
task_map = {t["task_key"]: t for t in 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), (
f"{name}: depends_on {deps} != upstream {pipe.upstream}"
)
assert sorted(deps) == sorted(pipe.upstream)
else:
assert "depends_on" not in task, (
f"{name} has no upstream but has depends_on"
)
assert "depends_on" not in task
def test_task_parameters_include_pipeline_name(self):
data = yaml.safe_load(self._generate())
tasks = data["resources"]["jobs"]["stack_pipelines"]["tasks"]
def test_tasks_use_sql_task(self):
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
for task in tasks:
params = task["python_wheel_task"]["parameters"]
assert params[0] == "run"
assert params[1] == task["task_key"]
assert "--target" in params
assert "--save" in params
assert "sql_task" in task, f"Task {task['task_key']} missing sql_task"
assert "file" in task["sql_task"]
class TestSQLFiles:
"""SQL transpilation produces actual Databricks SQL."""
def _emit(self) -> dict[str, str]:
from backends.databricks import emit
from conf import cfg
return emit(cfg._data)
def test_sql_files_generated(self):
files = self._emit()
sql_files = [
k for k in files if k.startswith("bundle/sql/") and k.endswith(".sql")
]
assert len(sql_files) > 100, f"Only {len(sql_files)} SQL files generated"
def test_sql_files_contain_sql(self):
files = self._emit()
sql_files = [
k for k in files if k.startswith("bundle/sql/") and k.endswith(".sql")
]
for path in sql_files[:10]:
content = files[path]
assert "SELECT" in content.upper() or "CREATE" in content.upper(), (
f"{path} doesn't contain SQL"
)
def test_sql_references_catalog_variable(self):
files = self._emit()
sql_files = [
k for k in files if k.startswith("bundle/sql/") and k.endswith(".sql")
]
# At least some SQL files should reference the catalog variable
has_catalog_ref = any("${var.catalog}" in files[p] for p in sql_files)
assert has_catalog_ref, "No SQL files reference ${var.catalog}"
class TestDDLFiles:
"""DDL generation from SQLTable models."""
def _emit(self) -> dict[str, str]:
from backends.databricks import emit
from conf import cfg
return emit(cfg._data)
def test_ddl_files_generated(self):
files = self._emit()
ddl_files = [k for k in files if k.startswith("bundle/ddl/")]
assert len(ddl_files) > 100, f"Only {len(ddl_files)} DDL files generated"
def test_ddl_contains_create_table(self):
files = self._emit()
ddl_files = [k for k in files if k.startswith("bundle/ddl/")]
for path in ddl_files[:10]:
assert "CREATE TABLE" in files[path], f"{path} missing CREATE TABLE"
class TestTargets:
"""#171 — Target isolation."""
"""Target isolation."""
def _data(self) -> dict:
def _yaml(self) -> dict:
from backends.databricks import emit
from conf import cfg
@@ -150,46 +199,26 @@ class TestTargets:
return yaml.safe_load(emit(cfg._data)["databricks.yml"])
def test_three_targets(self):
data = self._data()
assert set(data["targets"].keys()) == {"dev", "staging", "prod"}
assert set(self._yaml()["targets"].keys()) == {"dev", "staging", "prod"}
def test_dev_is_default(self):
data = self._data()
assert data["targets"]["dev"].get("default") is True
assert self._yaml()["targets"]["dev"].get("default") is True
def test_dev_mode(self):
data = self._data()
assert data["targets"]["dev"]["mode"] == "development"
assert self._yaml()["targets"]["dev"]["mode"] == "development"
def test_prod_mode(self):
data = self._data()
assert data["targets"]["prod"]["mode"] == "production"
assert self._yaml()["targets"]["prod"]["mode"] == "production"
def test_each_target_has_catalog(self):
data = self._data()
for tname, tcfg in data["targets"].items():
for tname, tcfg in self._yaml()["targets"].items():
assert "catalog" in tcfg.get("variables", {}), (
f"Target '{tname}' missing catalog variable"
)
class TestArtifacts:
"""#175 — Artifact configuration."""
def test_has_artifact(self):
from backends.databricks import emit
from conf import cfg
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
assert "artifacts" in data
assert "default" in data["artifacts"]
assert data["artifacts"]["default"]["type"] == "whl"
assert data["artifacts"]["default"]["build"] == "uv build"
class TestSchemas:
"""#173 — Unity Catalog schemas from table catalog."""
"""Unity Catalog schemas from table catalog."""
def test_schemas_generated(self):
from backends.databricks import emit
@@ -198,8 +227,7 @@ class TestSchemas:
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
schemas = data["resources"].get("schemas", {})
assert len(schemas) > 0, "No schemas generated"
# Core schema should always exist
assert len(schemas) > 0
assert "core" in schemas
def test_schema_references_catalog_variable(self):
@@ -209,40 +237,25 @@ class TestSchemas:
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
for name, schema in data["resources"]["schemas"].items():
assert schema["catalog_name"] == "${var.catalog}", (
f"Schema '{name}' doesn't reference catalog variable"
)
class TestVolumes:
"""#173 — Staging volume."""
def test_staging_volume(self):
from backends.databricks import emit
from conf import cfg
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
volumes = data["resources"].get("volumes", {})
assert "staging" in volumes
assert volumes["staging"]["volume_type"] == "MANAGED"
assert schema["catalog_name"] == "${var.catalog}"
class TestIdempotency:
"""#179 — Running the generator twice produces identical output."""
"""Running the generator twice produces identical output."""
def test_idempotent(self):
from backends.databricks import emit
from conf import cfg
first = emit(cfg._data)["databricks.yml"]
second = emit(cfg._data)["databricks.yml"]
assert first == second
first = emit(cfg._data)
second = emit(cfg._data)
assert first.keys() == second.keys()
assert first["databricks.yml"] == second["databricks.yml"]
class TestPluginArchitecture:
"""#178 — Plugin extensibility."""
"""Plugin extensibility."""
def test_custom_plugin(self):
from backends.databricks import _DEFAULT_PLUGINS, DabPlugin, emit
@@ -253,13 +266,14 @@ class TestPluginArchitecture:
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:
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
result = emit(cfg._data)
data = yaml.safe_load(result["databricks.yml"])
assert "experiments" in data["resources"]
assert (
data["resources"]["experiments"]["test_exp"]["name"]
== "test-experiment"
)
assert "bundle/test/hello.sql" in result
finally:
_DEFAULT_PLUGINS.pop()