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
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:
247
databricks.yml
247
databricks.yml
@@ -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
|
||||
|
||||
@@ -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"),
|
||||
]
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user