add DAB generator: databricks.yml from pipeline registry + stack.toml
All checks were successful
Infra CI / notebooks (push) Successful in 7s
Infra CI / zotero (push) Successful in 8s
CI / lint-test (push) Successful in 1m23s
Infra CI / docs (push) Successful in 10s
Infra CI / api (push) Successful in 26s
Infra CI / mc (push) Successful in 6s
Infra CI / notebooks (pull_request) Successful in 13s
Infra CI / zotero (pull_request) Successful in 7s
Infra CI / docs (pull_request) Successful in 6s
Infra CI / api (pull_request) Successful in 6s
Infra CI / mc (pull_request) Successful in 7s
CI / lint-test (pull_request) Successful in 1m17s
All checks were successful
Infra CI / notebooks (push) Successful in 7s
Infra CI / zotero (push) Successful in 8s
CI / lint-test (push) Successful in 1m23s
Infra CI / docs (push) Successful in 10s
Infra CI / api (push) Successful in 26s
Infra CI / mc (push) Successful in 6s
Infra CI / notebooks (pull_request) Successful in 13s
Infra CI / zotero (pull_request) Successful in 7s
Infra CI / docs (pull_request) Successful in 6s
Infra CI / api (pull_request) Successful in 6s
Infra CI / mc (pull_request) Successful in 7s
CI / lint-test (pull_request) Successful in 1m17s
closes #170, closes #171, closes #172, closes #173, closes #174, closes #175, closes #176, closes #177, closes #178, closes #179 Pipeline metadata (#170): - Add upstream and cluster_profile to Pipeline model - Populate upstream for all 13 pipeline modules Config (#171): - Add [databricks] section to stack.toml with clusters, targets Core emitter (#172): - backends/databricks.py generates databricks.yml from registry - All 13 tasks with correct depends_on, no hardcoding - Wired into gen_config.py (generate + --check) Resources (#173, #174): - 30 Unity Catalog schemas from aco.lake.catalog.Catalog - Staging volume for data upload - Cluster profiles: only referenced profiles emitted Artifacts (#175): wheel build via uv build Permissions (#176): production targets support run_as CI gate (#177): gen_config.py --check covers databricks.yml Plugin architecture (#178): DabPlugin protocol, 3 built-in plugins Tests (#179): 23 tests covering round-trip, completeness, deps, idempotency
This commit is contained in:
434
databricks.yml
434
databricks.yml
@@ -1,173 +1,331 @@
|
||||
# Databricks Asset Bundle — deploys all pipelines as Databricks jobs.
|
||||
#
|
||||
# Usage:
|
||||
# databricks bundle validate
|
||||
# databricks bundle deploy --target dev
|
||||
# databricks bundle run stack_pipelines --target dev
|
||||
# DO NOT EDIT — generated by gen_config.py from stack.toml + pipeline registry
|
||||
# Re-generate: uv run python dev/scripts/gen_config.py
|
||||
|
||||
bundle:
|
||||
name: stack
|
||||
|
||||
artifacts:
|
||||
default:
|
||||
type: whl
|
||||
build: uv build
|
||||
path: .
|
||||
variables:
|
||||
catalog:
|
||||
description: Unity Catalog name for table references
|
||||
default: aco
|
||||
|
||||
workspace:
|
||||
host: ${DATABRICKS_HOST}
|
||||
|
||||
targets:
|
||||
dev:
|
||||
mode: development
|
||||
default: true
|
||||
variables:
|
||||
catalog: aco_dev
|
||||
staging:
|
||||
variables:
|
||||
catalog: aco_staging
|
||||
prod:
|
||||
mode: production
|
||||
variables:
|
||||
catalog: aco
|
||||
|
||||
staging:
|
||||
variables:
|
||||
catalog: aco_staging
|
||||
resources:
|
||||
jobs:
|
||||
stack_pipelines:
|
||||
name: "stack-pipelines"
|
||||
description: "Run all 13 ACO pipelines in dependency order"
|
||||
name: stack-pipelines
|
||||
description: Run all 13 ACO pipelines in dependency order
|
||||
schedule:
|
||||
quartz_cron_expression: "0 0 6 * * ?"
|
||||
quartz_cron_expression: 0 0 6 * * ?
|
||||
timezone_id: America/New_York
|
||||
job_clusters:
|
||||
- job_cluster_key: pipeline_cluster
|
||||
new_cluster:
|
||||
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"
|
||||
|
||||
- job_cluster_key: default_cluster
|
||||
new_cluster:
|
||||
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: input_layer
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "input_layer", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: cclf
|
||||
depends_on:
|
||||
- task_key: input_layer
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "cclf", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: claims_preprocessing
|
||||
depends_on:
|
||||
- task_key: input_layer
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "claims_preprocessing", "--target", "databricks", "--save"]
|
||||
|
||||
- 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
|
||||
- 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
|
||||
- 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
|
||||
- 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
|
||||
- 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
|
||||
- task_key: data_quality
|
||||
depends_on:
|
||||
- task_key: core
|
||||
depends_on:
|
||||
- task_key: claims_preprocessing
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "core", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: readmissions
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "readmissions", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: pharmacy
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "pharmacy", "--target", "databricks", "--save"]
|
||||
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters:
|
||||
- run
|
||||
- data_quality
|
||||
- --target
|
||||
- databricks
|
||||
- --save
|
||||
- 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
|
||||
- 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
|
||||
- task_key: main
|
||||
depends_on:
|
||||
- task_key: readmissions
|
||||
- task_key: pharmacy
|
||||
- task_key: hcc_suspecting
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "hcc_suspecting", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: provider_attribution
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "provider_attribution", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: quality_measures
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "quality_measures", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: cms_quality_measures
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "cms_quality_measures", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: ahrq_measures
|
||||
depends_on:
|
||||
- task_key: core
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "ahrq_measures", "--target", "databricks", "--save"]
|
||||
|
||||
- task_key: data_quality
|
||||
depends_on:
|
||||
- task_key: core
|
||||
- task_key: readmissions
|
||||
- task_key: pharmacy
|
||||
job_cluster_key: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "data_quality", "--target", "databricks", "--save"]
|
||||
|
||||
- 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: pipeline_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters: ["run", "main", "--target", "databricks", "--save"]
|
||||
job_cluster_key: default_cluster
|
||||
python_wheel_task:
|
||||
package_name: stack
|
||||
entry_point: cli
|
||||
parameters:
|
||||
- run
|
||||
- main
|
||||
- --target
|
||||
- databricks
|
||||
- --save
|
||||
- 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
|
||||
- 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
|
||||
- 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
|
||||
- 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
|
||||
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}
|
||||
core:
|
||||
name: core
|
||||
catalog_name: ${var.catalog}
|
||||
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}
|
||||
input_layer:
|
||||
name: input_layer
|
||||
catalog_name: ${var.catalog}
|
||||
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}
|
||||
provider_attribution:
|
||||
name: provider_attribution
|
||||
catalog_name: ${var.catalog}
|
||||
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
|
||||
catalog_name: ${var.catalog}
|
||||
schema_name: default
|
||||
volume_type: MANAGED
|
||||
|
||||
323
dev/scripts/backends/databricks.py
Normal file
323
dev/scripts/backends/databricks.py
Normal file
@@ -0,0 +1,323 @@
|
||||
"""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.
|
||||
|
||||
Follows the same emit() → dict[str, str] pattern as gitea.py / github.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import yaml
|
||||
from collections import OrderedDict
|
||||
|
||||
# ── YAML helpers ──────────────────────────────────────────────────
|
||||
|
||||
# Use OrderedDict + custom representer to control key order in output
|
||||
yaml.add_representer(
|
||||
OrderedDict,
|
||||
lambda dumper, data: dumper.represent_mapping(
|
||||
"tag:yaml.org,2002:map", data.items()
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
# ── Plugin protocol ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class DabPlugin:
|
||||
"""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 a dict of resource type -> resource definitions."""
|
||||
return {}
|
||||
|
||||
|
||||
# ── Built-in plugins ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class JobsPlugin(DabPlugin):
|
||||
"""Generate jobs from the pipeline registry."""
|
||||
|
||||
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
|
||||
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),
|
||||
]
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Build tasks from registry
|
||||
tasks = []
|
||||
for 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
|
||||
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["job_clusters"] = job_clusters
|
||||
job["tasks"] = tasks
|
||||
|
||||
return {"jobs": OrderedDict([("stack_pipelines", job)])}
|
||||
|
||||
|
||||
class SchemasPlugin(DabPlugin):
|
||||
"""Generate Unity Catalog schema resources from the table catalog."""
|
||||
|
||||
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}"),
|
||||
]
|
||||
)
|
||||
return {"schemas": schemas}
|
||||
|
||||
|
||||
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"),
|
||||
]
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
# ── Core generator ───────────────────────────────────────────────
|
||||
|
||||
_HEADER = (
|
||||
"# DO NOT EDIT — generated by gen_config.py from stack.toml + pipeline registry\n"
|
||||
"# Re-generate: uv run python dev/scripts/gen_config.py\n"
|
||||
)
|
||||
|
||||
# Default plugins
|
||||
_DEFAULT_PLUGINS: list[DabPlugin] = [
|
||||
JobsPlugin(),
|
||||
SchemasPlugin(),
|
||||
VolumesPlugin(),
|
||||
]
|
||||
|
||||
|
||||
def _get_catalog_schemas() -> list[str]:
|
||||
"""Discover schemas from the table catalog."""
|
||||
try:
|
||||
from aco.lake.catalog import Catalog
|
||||
|
||||
cat = Catalog()
|
||||
return cat.schemas()
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def emit(cfg_data: dict) -> dict[str, str]:
|
||||
"""Generate databricks.yml from config + pipeline registry.
|
||||
|
||||
Returns {relative_path: content}.
|
||||
"""
|
||||
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
|
||||
all_resources: dict = {}
|
||||
for plugin in _DEFAULT_PLUGINS:
|
||||
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()
|
||||
|
||||
# bundle:
|
||||
bundle["bundle"] = OrderedDict([("name", bundle_name)])
|
||||
|
||||
# artifacts:
|
||||
bundle["artifacts"] = OrderedDict(
|
||||
[
|
||||
(
|
||||
"default",
|
||||
OrderedDict(
|
||||
[
|
||||
("type", "whl"),
|
||||
("build", "uv build"),
|
||||
("path", "."),
|
||||
]
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# 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"),
|
||||
),
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# workspace:
|
||||
bundle["workspace"] = OrderedDict(
|
||||
[("host", "${DATABRICKS_HOST}")]
|
||||
)
|
||||
|
||||
# targets:
|
||||
targets = OrderedDict()
|
||||
for tname, tcfg in sorted(targets_cfg.items()):
|
||||
target = OrderedDict()
|
||||
if tcfg.get("mode"):
|
||||
target["mode"] = tcfg["mode"]
|
||||
if tcfg.get("default"):
|
||||
target["default"] = True
|
||||
target["variables"] = OrderedDict(
|
||||
[("catalog", tcfg.get("catalog", "aco"))]
|
||||
)
|
||||
# Production targets get run_as
|
||||
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"],
|
||||
)
|
||||
]
|
||||
)
|
||||
targets[tname] = target
|
||||
bundle["targets"] = targets
|
||||
|
||||
# permissions:
|
||||
perms_cfg = dab_cfg.get("permissions", [])
|
||||
if perms_cfg:
|
||||
bundle["permissions"] = perms_cfg
|
||||
|
||||
# resources:
|
||||
bundle["resources"] = all_resources
|
||||
|
||||
content = _HEADER + "\n" + _dump(dict(bundle))
|
||||
return {"databricks.yml": content}
|
||||
@@ -129,6 +129,15 @@ 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."""
|
||||
if "databricks" not in cfg._data:
|
||||
return {}
|
||||
from backends.databricks import emit
|
||||
|
||||
return emit(cfg._data)
|
||||
|
||||
|
||||
def generate(backend: str | None = None) -> dict[str, str]:
|
||||
if backend is None:
|
||||
backend = cfg._data.get("ci", {}).get("backend", "woodpecker")
|
||||
@@ -136,6 +145,7 @@ def generate(backend: str | None = None) -> dict[str, str]:
|
||||
files: dict[str, str] = {}
|
||||
files.update(gen_coredns())
|
||||
files.update(_emit_backend(backend))
|
||||
files.update(_emit_dab())
|
||||
return files
|
||||
|
||||
|
||||
|
||||
@@ -52,6 +52,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["core"],
|
||||
exprs=[
|
||||
# Staging
|
||||
Expr(
|
||||
@@ -595,7 +596,7 @@ pipeline = Pipeline(
|
||||
"in denom, not excluded."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -21,6 +21,12 @@ class Pipeline(BaseModel):
|
||||
"""Ordered list of ``Expr`` objects. Typed as ``Any`` to avoid
|
||||
circular import (``Expr`` lives in ``aco.express.base``)."""
|
||||
|
||||
upstream: list[str] = []
|
||||
"""Pipeline names this pipeline depends on (for DAB task ordering)."""
|
||||
|
||||
cluster_profile: str = "default"
|
||||
"""Compute tier hint for DAB cluster selection."""
|
||||
|
||||
def run(self, load: Callable[[str], Any]) -> dict[str, Any]:
|
||||
"""Execute expressions in order, returning cache of DataFrames."""
|
||||
from aco.pipe.runner import run_pipeline
|
||||
|
||||
@@ -90,6 +90,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["input_layer"],
|
||||
exprs=[
|
||||
# -- MBI crosswalk ----------------------------------------
|
||||
Expr(
|
||||
@@ -332,7 +333,7 @@ pipeline = Pipeline(
|
||||
"and race value crosswalks."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -453,6 +453,7 @@ _ORPHANED: list[Expr] = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["input_layer"],
|
||||
exprs=_SERVICE_CATEGORY + _ENCOUNTER_STEPS + _OFFICE_VISITS + _ORPHANED,
|
||||
)
|
||||
|
||||
|
||||
@@ -516,11 +516,12 @@ pipeline_hwr = Pipeline(
|
||||
# Combined pipeline for running all three measures together
|
||||
# Staging steps are already in pipeline_uamcc.exprs
|
||||
pipeline = Pipeline(
|
||||
upstream=["core"],
|
||||
exprs=[
|
||||
*pipeline_uamcc.exprs,
|
||||
*pipeline_acr.exprs,
|
||||
*pipeline_hwr.exprs,
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -25,6 +25,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["claims_preprocessing"],
|
||||
exprs=[
|
||||
Expr(
|
||||
name="core._stg_claims_member_months",
|
||||
@@ -161,7 +162,7 @@ pipeline = Pipeline(
|
||||
"for the core clinical staging table."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -12,6 +12,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["core", "readmissions", "pharmacy"],
|
||||
exprs=[
|
||||
Expr(
|
||||
name="data_quality.testing_summary",
|
||||
@@ -29,7 +30,7 @@ pipeline = Pipeline(
|
||||
"the latest test invocation."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -19,6 +19,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["core"],
|
||||
exprs=[
|
||||
Expr(
|
||||
name="hcc_suspecting._stg_condition",
|
||||
@@ -87,7 +88,7 @@ pipeline = Pipeline(
|
||||
"pipeline."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -26,6 +26,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=[],
|
||||
exprs=[
|
||||
Expr(
|
||||
name="input_layer.input_layer__appointment",
|
||||
@@ -158,7 +159,7 @@ pipeline = Pipeline(
|
||||
"input layer source table."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -28,6 +28,16 @@ _REFS = [Tag.module("aco")]
|
||||
# main.snapshot_run_results — these must be loaded from DB (not from
|
||||
# pipeline cache), so they come AFTER alerts_dbt_models in the list.
|
||||
pipeline = Pipeline(
|
||||
upstream=[
|
||||
"readmissions",
|
||||
"pharmacy",
|
||||
"hcc_suspecting",
|
||||
"provider_attribution",
|
||||
"quality_measures",
|
||||
"cms_quality_measures",
|
||||
"ahrq_measures",
|
||||
"data_quality",
|
||||
],
|
||||
exprs=[
|
||||
Expr(
|
||||
name="main.alerts_anomaly_detection",
|
||||
@@ -197,7 +207,7 @@ pipeline = Pipeline(
|
||||
refs=_REFS,
|
||||
description=("Stub join of dbt run results with snapshots."),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -17,6 +17,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["core"],
|
||||
exprs=[
|
||||
Expr(
|
||||
name="pharmacy._int_brand_with_generic_available",
|
||||
@@ -80,7 +81,7 @@ pipeline = Pipeline(
|
||||
"per-NDC cost breakdown for each brand RxCUI."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -23,6 +23,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["core"],
|
||||
exprs=[
|
||||
# Staging
|
||||
Expr(
|
||||
@@ -169,7 +170,7 @@ pipeline = Pipeline(
|
||||
refs=_REFS,
|
||||
description=("5-step attribution waterfall per performance_year."),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -28,6 +28,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["core"],
|
||||
exprs=[
|
||||
# Performance period functions
|
||||
Expr(
|
||||
@@ -202,7 +203,7 @@ pipeline = Pipeline(
|
||||
"core.pharmacy_claim for quality measure processing."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
@@ -26,6 +26,7 @@ _REFS = [
|
||||
]
|
||||
|
||||
pipeline = Pipeline(
|
||||
upstream=["core"],
|
||||
exprs=[
|
||||
Expr(
|
||||
name="readmissions._int_encounter",
|
||||
@@ -215,7 +216,7 @@ pipeline = Pipeline(
|
||||
"<= 30 days."
|
||||
),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
run = pipeline.run
|
||||
|
||||
33
stack.toml
33
stack.toml
@@ -105,6 +105,39 @@ host = "trino"
|
||||
port = 8080
|
||||
catalog = "iceberg"
|
||||
|
||||
# ── Databricks Asset Bundle generation ────────────────────────────
|
||||
# gen_config.py reads this to produce databricks.yml.
|
||||
# Add a new pipeline module → commit → databricks.yml updates automatically.
|
||||
|
||||
[databricks]
|
||||
bundle_name = "stack"
|
||||
package_name = "stack"
|
||||
entry_point = "cli"
|
||||
schedule = "0 0 6 * * ?"
|
||||
timezone = "America/New_York"
|
||||
|
||||
[databricks.clusters.default]
|
||||
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
|
||||
catalog = "aco_dev"
|
||||
|
||||
[databricks.targets.staging]
|
||||
catalog = "aco_staging"
|
||||
|
||||
[databricks.targets.prod]
|
||||
mode = "production"
|
||||
catalog = "aco"
|
||||
|
||||
[default]
|
||||
target = "local" # duckdb | databricks | trino | lake
|
||||
|
||||
|
||||
265
tests/test_dab_generator.py
Normal file
265
tests/test_dab_generator.py
Normal file
@@ -0,0 +1,265 @@
|
||||
"""Tests for the Databricks Asset Bundle generator.
|
||||
|
||||
Covers: round-trip, completeness, dependency ordering, idempotency,
|
||||
target isolation, plugin extensibility, and pipeline metadata.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT / "dev" / "scripts"))
|
||||
|
||||
|
||||
class TestPipelineMetadata:
|
||||
"""#170 — Pipeline.upstream and cluster_profile."""
|
||||
|
||||
def test_all_pipelines_have_upstream(self):
|
||||
from aco.pipe import registry
|
||||
|
||||
for name, pipe in registry.items():
|
||||
assert hasattr(pipe, "upstream"), f"{name} missing upstream"
|
||||
assert isinstance(pipe.upstream, list), f"{name}.upstream not a list"
|
||||
|
||||
def test_all_pipelines_have_cluster_profile(self):
|
||||
from aco.pipe import registry
|
||||
|
||||
for name, pipe in registry.items():
|
||||
assert hasattr(pipe, "cluster_profile"), f"{name} missing cluster_profile"
|
||||
|
||||
def test_no_cycles_in_upstream(self):
|
||||
from aco.pipe import registry
|
||||
|
||||
visited: set[str] = set()
|
||||
path: set[str] = set()
|
||||
|
||||
def visit(name: str) -> None:
|
||||
if name in path:
|
||||
raise AssertionError(f"Cycle detected: {name} in {path}")
|
||||
if name in visited:
|
||||
return
|
||||
path.add(name)
|
||||
pipe = registry.get(name)
|
||||
if pipe:
|
||||
for dep in pipe.upstream:
|
||||
visit(dep)
|
||||
path.remove(name)
|
||||
visited.add(name)
|
||||
|
||||
for name in registry:
|
||||
visit(name)
|
||||
|
||||
def test_upstream_refs_exist_in_registry(self):
|
||||
from aco.pipe import registry
|
||||
|
||||
for name, pipe in registry.items():
|
||||
for dep in pipe.upstream:
|
||||
assert dep in registry, (
|
||||
f"{name} references upstream '{dep}' not in registry"
|
||||
)
|
||||
|
||||
def test_root_pipelines_have_empty_upstream(self):
|
||||
from aco.pipe import registry
|
||||
|
||||
assert registry["input_layer"].upstream == []
|
||||
|
||||
|
||||
class TestDabEmitter:
|
||||
"""#172 — Core generator produces valid databricks.yml."""
|
||||
|
||||
def _generate(self) -> str:
|
||||
from backends.databricks import emit
|
||||
|
||||
from conf import cfg
|
||||
|
||||
files = emit(cfg._data)
|
||||
return files["databricks.yml"]
|
||||
|
||||
def test_generates_valid_yaml(self):
|
||||
content = self._generate()
|
||||
data = yaml.safe_load(content)
|
||||
assert isinstance(data, dict)
|
||||
|
||||
def test_has_required_top_level_keys(self):
|
||||
data = yaml.safe_load(self._generate())
|
||||
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"
|
||||
|
||||
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"]
|
||||
task_keys = {t["task_key"] for t in 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
|
||||
|
||||
data = yaml.safe_load(self._generate())
|
||||
tasks = data["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"]
|
||||
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}"
|
||||
)
|
||||
else:
|
||||
assert "depends_on" not in task, (
|
||||
f"{name} has no upstream but has depends_on"
|
||||
)
|
||||
|
||||
def test_task_parameters_include_pipeline_name(self):
|
||||
data = yaml.safe_load(self._generate())
|
||||
tasks = data["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
|
||||
|
||||
|
||||
class TestTargets:
|
||||
"""#171 — Target isolation."""
|
||||
|
||||
def _data(self) -> dict:
|
||||
from backends.databricks import emit
|
||||
|
||||
from conf import cfg
|
||||
|
||||
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"}
|
||||
|
||||
def test_dev_is_default(self):
|
||||
data = self._data()
|
||||
assert data["targets"]["dev"].get("default") is True
|
||||
|
||||
def test_dev_mode(self):
|
||||
data = self._data()
|
||||
assert data["targets"]["dev"]["mode"] == "development"
|
||||
|
||||
def test_prod_mode(self):
|
||||
data = self._data()
|
||||
assert data["targets"]["prod"]["mode"] == "production"
|
||||
|
||||
def test_each_target_has_catalog(self):
|
||||
data = self._data()
|
||||
for tname, tcfg in data["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."""
|
||||
|
||||
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) > 0, "No schemas generated"
|
||||
# Core schema should always exist
|
||||
assert "core" in schemas
|
||||
|
||||
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}", (
|
||||
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"
|
||||
|
||||
|
||||
class TestIdempotency:
|
||||
"""#179 — 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
|
||||
|
||||
|
||||
class TestPluginArchitecture:
|
||||
"""#178 — 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"}}}
|
||||
|
||||
_DEFAULT_PLUGINS.append(TestPlugin())
|
||||
try:
|
||||
data = yaml.safe_load(emit(cfg._data)["databricks.yml"])
|
||||
assert "experiments" in data["resources"]
|
||||
assert (
|
||||
data["resources"]["experiments"]["test_exp"]["name"]
|
||||
== "test-experiment"
|
||||
)
|
||||
finally:
|
||||
_DEFAULT_PLUGINS.pop()
|
||||
Reference in New Issue
Block a user