feat: modernize DAB — Python wheel tasks, schema cleanup, dead config removal
All checks were successful
CI / skinny-install (aco) (push) Successful in 50s
CI / skinny-install (api) (push) Successful in 30s
CI / lint-test (push) Successful in 1m27s
CI / skinny-install (bcda) (push) Successful in 30s
CI / skinny-install (bib) (push) Successful in 31s
CI / skinny-install (bls) (push) Successful in 27s
CI / skinny-install (ccw) (push) Successful in 31s
CI / skinny-install (cli) (push) Successful in 31s
CI / skinny-install (cms) (push) Successful in 28s
CI / skinny-install (conf) (push) Successful in 31s
CI / skinny-install (pfs) (push) Successful in 31s
CI / skinny-install (rex) (push) Successful in 27s
Infra CI / notebooks (push) Successful in 7s
Infra CI / zotero (push) Successful in 7s
Infra CI / docs (push) Successful in 39s
Infra CI / api (push) Successful in 5s
Infra CI / mc (push) Successful in 6s
Deploy / build-scan-report (push) Successful in 3m42s

Replace SQL task deployment with PythonWheelTask execution. Same
narwhals code runs on DuckDB (local) and Spark (Databricks) via
make_databricks_loader — eliminates 208 transpiled SQL files from
the deployment path.

- JobsPlugin: python_wheel_task with package_name=stack, entry_point=cli
- CLI: added --catalog flag and spark target with _SparkContext
- Catalog.pipeline_schemas(): 13 schemas from registry (was 34 from aco.table)
- Removed: large cluster config, warehouse_id variable, DDL tasks, SQLPlugin
- SQL transpilation retained as validation-only tool (112 tests)
- 24 DAB tests updated and passing

Closes #219, closes #220, closes #221, closes #222, closes #223,
closes #225, closes #226
This commit is contained in:
kert
2026-03-24 23:40:10 -04:00
parent 6d291334ca
commit 46f05db0ee
6 changed files with 244 additions and 212 deletions

View File

@@ -10,8 +10,6 @@ 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:
@@ -46,66 +44,114 @@ resources:
depends_on:
- task_key: core
job_cluster_key: default_cluster
sql_task:
file:
path: bundle/sql/ahrq_measures/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/cclf/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/claims_preprocessing/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/cms_quality_measures/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/core/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/data_quality/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/hcc_suspecting/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/input_layer/
warehouse_id: ${var.warehouse_id}
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
@@ -117,85 +163,82 @@ resources:
- task_key: ahrq_measures
- task_key: data_quality
job_cluster_key: default_cluster
sql_task:
file:
path: bundle/sql/main/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/pharmacy/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/provider_attribution/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/quality_measures/
warehouse_id: ${var.warehouse_id}
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
sql_task:
file:
path: bundle/sql/readmissions/
warehouse_id: ${var.warehouse_id}
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}
alr:
name: alr
catalog_name: ${var.catalog}
alr_codes:
name: alr_codes
catalog_name: ${var.catalog}
asr:
name: asr
catalog_name: ${var.catalog}
bcda:
name: bcda
catalog_name: ${var.catalog}
cclf:
name: cclf
catalog_name: ${var.catalog}
ccsr:
name: ccsr
catalog_name: ${var.catalog}
ccw:
name: ccw
catalog_name: ${var.catalog}
chronic_conditions:
name: chronic_conditions
catalog_name: ${var.catalog}
claims_preprocessing:
name: claims_preprocessing
catalog_name: ${var.catalog}
clinical_concept_library:
name: clinical_concept_library
catalog_name: ${var.catalog}
cms:
name: cms
catalog_name: ${var.catalog}
cms_hcc:
name: cms_hcc
catalog_name: ${var.catalog}
cms_provider_attribution:
name: cms_provider_attribution
catalog_name: ${var.catalog}
cms_quality_measures:
name: cms_quality_measures
catalog_name: ${var.catalog}
@@ -205,15 +248,6 @@ resources:
data_quality:
name: data_quality
catalog_name: ${var.catalog}
ed_classification:
name: ed_classification
catalog_name: ${var.catalog}
financial_pmpm:
name: financial_pmpm
catalog_name: ${var.catalog}
hcc_recapture:
name: hcc_recapture
catalog_name: ${var.catalog}
hcc_suspecting:
name: hcc_suspecting
catalog_name: ${var.catalog}
@@ -223,12 +257,6 @@ resources:
main:
name: main
catalog_name: ${var.catalog}
metadata:
name: metadata
catalog_name: ${var.catalog}
pfs:
name: pfs
catalog_name: ${var.catalog}
pharmacy:
name: pharmacy
catalog_name: ${var.catalog}
@@ -238,24 +266,9 @@ resources:
quality_measures:
name: quality_measures
catalog_name: ${var.catalog}
reach:
name: reach
catalog_name: ${var.catalog}
readmissions:
name: readmissions
catalog_name: ${var.catalog}
reference_data:
name: reference_data
catalog_name: ${var.catalog}
ssp:
name: ssp
catalog_name: ${var.catalog}
terminology:
name: terminology
catalog_name: ${var.catalog}
tuva_provider_attribution:
name: tuva_provider_attribution
catalog_name: ${var.catalog}
volumes:
staging:
name: staging

View File

@@ -128,39 +128,11 @@ class JobsPlugin(DabPlugin):
)
)
# 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 — each pipeline runs as a Python wheel task
pipeline_tasks = []
for pipe_name, pipe in sorted(registry.items()):
profile = getattr(pipe, "cluster_profile", "default")
# Build depends_on from upstream pipeline names
depends_on = []
for dep in pipe.upstream:
depends_on.append(OrderedDict([("task_key", dep)]))
@@ -171,22 +143,26 @@ class JobsPlugin(DabPlugin):
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(
task["python_wheel_task"] = OrderedDict(
[
("package_name", cfg.get("package_name", "stack")),
("entry_point", cfg.get("entry_point", "cli")),
(
"file",
"named_parameters",
OrderedDict(
[
("path", f"bundle/sql/{pipe_name}/"),
("run", pipe_name),
("--target", "spark"),
("--catalog", "${var.catalog}"),
("--save", ""),
]
),
),
("warehouse_id", "${var.warehouse_id}"),
]
)
task["libraries"] = [
OrderedDict([("whl", "dbfs:/FileStore/wheels/stack-latest.whl")])
]
pipeline_tasks.append(task)
schedule = cfg.get("schedule", "0 0 6 * * ?")
@@ -257,7 +233,6 @@ _HEADER = (
)
_DEFAULT_PLUGINS: list[DabPlugin] = [
SQLPlugin(),
JobsPlugin(),
SchemasPlugin(),
VolumesPlugin(),
@@ -265,10 +240,11 @@ _DEFAULT_PLUGINS: list[DabPlugin] = [
def _get_catalog_schemas() -> list[str]:
"""Return only schemas produced by registered pipelines."""
try:
from aco.lake.catalog import Catalog
return Catalog().schemas()
return Catalog.pipeline_schemas()
except Exception:
return []
@@ -305,7 +281,6 @@ def emit(cfg_data: dict) -> dict[str, str]:
bundle = OrderedDict()
bundle["bundle"] = OrderedDict([("name", bundle_name)])
# Sync SQL files to workspace
bundle["sync"] = OrderedDict(
[
("include", ["bundle/**"]),
@@ -323,14 +298,6 @@ def emit(cfg_data: dict) -> dict[str, str]:
]
),
),
(
"warehouse_id",
OrderedDict(
[
("description", "SQL warehouse ID for query execution"),
]
),
),
]
)

View File

@@ -320,6 +320,32 @@ class Catalog:
results.append((ns, modname))
return results
# ── Pipeline-derived schemas ─────────────────────────────────
@staticmethod
def pipeline_schemas() -> list[str]:
"""Return schemas actually produced by registered pipelines.
Derives schema names from the output names of all ``Expr``
objects across all pipelines in the registry. This is the
authoritative set for Databricks deployment — excludes legacy
schemas from ``aco.table`` that no pipeline produces.
Returns
-------
list[str]
Sorted unique schema names like ``['core', 'readmissions', ...]``.
"""
from aco.pipe import registry
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)
# ── Iceberg catalog (REST API via PyIceberg) ─────────────────
def iceberg_namespaces(self) -> list[str]:

View File

@@ -9,7 +9,11 @@ def run(
name: str = typer.Argument(help="Pipeline name (e.g. readmissions)."),
target: str = typer.Option(
"",
help="Target context: local, lake, trino, databricks. Defaults to stack.toml.",
help="Target context: local, lake, trino, databricks, spark.",
),
catalog: str = typer.Option(
"",
help="Unity Catalog name (overrides config). Implies target=spark.",
),
save: bool = typer.Option(
False, "--save", help="Save pipeline outputs back to the database."
@@ -23,6 +27,10 @@ def run(
typer.echo(f"Unknown pipeline '{name}'. Available: {names}", err=True)
raise typer.Exit(1)
# --catalog implies spark target
if catalog and not target:
target = "spark"
if not target:
from conf import cfg
@@ -30,7 +38,7 @@ def run(
target = getattr(target, "target", "local") if target else "local"
pipeline = registry[name]
ctx = _make_context(target, read_only=not save)
ctx = _make_context(target, read_only=not save, catalog_override=catalog)
typer.echo(f"Running pipeline '{name}' ({len(pipeline)} steps, target={target})")
cache = pipeline.run(ctx.load)
@@ -45,7 +53,7 @@ def run(
typer.echo("Done.")
def _make_context(target: str, *, read_only: bool = True): # noqa: ANN202
def _make_context(target: str, *, read_only: bool = True, catalog_override: str = ""): # noqa: ANN202
"""Build the appropriate Context from the target name."""
from conf import cfg, path
@@ -83,6 +91,12 @@ def _make_context(target: str, *, read_only: bool = True): # noqa: ANN202
dialect="databricks",
)
if target == "spark":
# For PythonWheelTask execution on Databricks runtime.
# Uses make_databricks_loader with the active SparkSession.
cat = catalog_override or getattr(cfg.lake.databricks, "catalog", "aco")
return _SparkContext(catalog=cat)
typer.echo(
f"Unknown target '{target}'. Use: local, lake, trino, databricks.",
err=True,
@@ -90,6 +104,36 @@ def _make_context(target: str, *, read_only: bool = True): # noqa: ANN202
raise typer.Exit(1)
class _SparkContext:
"""Lightweight context for Databricks PythonWheelTask execution.
Uses ``make_databricks_loader`` with the active SparkSession.
No DuckDB or Iceberg — reads/writes via ``spark.table()``.
"""
def __init__(self, catalog: str = "aco") -> None:
self.catalog = catalog
self._spark = None
def _get_spark(self): # noqa: ANN202
if self._spark is None:
from pyspark.sql import SparkSession
self._spark = SparkSession.builder.getOrCreate()
return self._spark
def load(self, table_ref: str): # noqa: ANN202
from aco.pipe.runner import make_databricks_loader
loader = make_databricks_loader(self._get_spark(), self.catalog)
return loader(table_ref)
def save(self, table_ref: str, df, mode: str = "replace") -> None: # noqa: ANN001
self._get_spark() # ensure session is active
full_ref = f"{self.catalog}.{table_ref}"
df.write.mode("overwrite" if mode == "replace" else mode).saveAsTable(full_ref)
def _save_outputs(ctx, cache: dict, pipeline_name: str) -> None: # noqa: ANN001
"""Save pipeline outputs that have qualified table names."""
for table_ref, df in cache.items():

View File

@@ -139,11 +139,6 @@ spark_version = "15.4.x-scala2.12"
node_type_id = "i3.xlarge"
num_workers = 1
[databricks.clusters.large]
spark_version = "15.4.x-scala2.12"
node_type_id = "i3.2xlarge"
num_workers = 4
[databricks.targets.dev]
mode = "development"
default = true

View File

@@ -121,49 +121,31 @@ class TestDabEmitter:
else:
assert "depends_on" not in task
def test_tasks_use_sql_task(self):
def test_tasks_use_python_wheel_task(self):
tasks = self._yaml()["resources"]["jobs"]["stack_pipelines"]["tasks"]
for task in tasks:
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"
assert "python_wheel_task" in task, (
f"Task {task['task_key']} missing python_wheel_task"
)
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_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}"
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_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", {})
class TestTargets:
@@ -205,8 +187,13 @@ class TestSchemas:
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
schemas = data["resources"].get("schemas", {})
assert len(schemas) > 0
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