diff --git a/experimental/bundletest/examples/orders_bundle/databricks.yml b/experimental/bundletest/examples/orders_bundle/databricks.yml index caedc4eac99..3beb27fa643 100644 --- a/experimental/bundletest/examples/orders_bundle/databricks.yml +++ b/experimental/bundletest/examples/orders_bundle/databricks.yml @@ -1,8 +1,18 @@ # Minimal illustrative bundle. The local backend runs a job's real sql_task artifact # against DuckDB; the cloud backend (follow-up PR) deploys and runs it on a real workspace. +# +# It also declares one instance of every user-authored resource kind so the example gallery +# can show a handle for each. Only jobs run locally; every other kind is read (not run) for +# static wiring assertions via env.resource(kind, name) / the typed env.() handles. +# The dashboard, genie space, and quality monitor all reference shop.gold.order_summary — +# the gold table aggregate_orders produces — so those doubles as cross-resource wiring checks. bundle: name: orders +variables: + warehouse_id: + description: SQL warehouse the transform runs on + resources: jobs: # SQL job — runnable on the local backend. @@ -33,6 +43,213 @@ resources: notebook_task: notebook_path: src/score_model -variables: - warehouse_id: - description: SQL warehouse the transform runs on + pipelines: + # Lakeflow pipeline — declared for wiring assertions; running it is cloud-only. + enrich_orders: + name: enrich_orders + catalog: shop + schema: enriched + libraries: + - notebook: + path: src/enrich_orders + + dashboards: + # Lakeview dashboard with an inline serialized definition (inlined as YAML), so + # source_tables() can read which tables its datasets query without a workspace. + orders_overview: + display_name: Orders Overview + warehouse_id: ${var.warehouse_id} + serialized_dashboard: + datasets: + - name: summary + displayName: Order summary + queryLines: + - "SELECT order_count, total_revenue\n" + - "FROM shop.gold.order_summary\n" + pages: + - name: main + displayName: Overview + layout: [] + + genie_spaces: + # Genie Space with an inline serialized definition. A space declares its tables + # explicitly under data_sources.tables[].identifier (a different schema from a dashboard). + orders_genie: + title: Orders Genie + warehouse_id: ${var.warehouse_id} + serialized_space: + version: 1 + data_sources: + tables: + - identifier: shop.gold.order_summary + + quality_monitors: + # Monitors the gold table the aggregate job produces. + orders_quality: + table_name: shop.gold.order_summary + output_schema_name: shop.monitoring + assets_dir: /Shared/orders/monitoring + snapshot: {} + + vector_search_indexes: + orders_index: + name: shop.ml.orders_index + endpoint_name: orders-vs-endpoint + index_type: DELTA_SYNC + delta_sync_index_spec: + source_table: shop.gold.order_summary + pipeline_type: TRIGGERED + + vector_search_endpoints: + orders_vs_endpoint: + name: orders-vs-endpoint + endpoint_type: STANDARD + + model_serving_endpoints: + orders_model_endpoint: + name: orders-model-endpoint + config: + served_entities: + - name: orders + entity_name: shop.ml.orders_model + entity_version: "1" + + apps: + orders_app: + name: orders-app + source_code_path: ./app + config: + command: + - python + - app.py + + models: + orders_model: + name: orders_model + + experiments: + orders_experiment: + name: /Shared/orders_experiment + + registered_models: + orders_registered: + name: orders + catalog_name: shop + schema_name: ml + + catalogs: + orders_catalog: + name: shop + grants: + - principal: users + privileges: + - USE_CATALOG + + schemas: + orders_schema: + name: gold + catalog_name: shop + + volumes: + raw_data: + name: raw_data + catalog_name: shop + schema_name: bronze + volume_type: MANAGED + + external_locations: + orders_location: + name: orders_location + url: s3://example-bucket/orders + credential_name: orders_credential + + clusters: + orders_cluster: + cluster_name: orders + spark_version: 15.4.x-scala2.12 + node_type_id: i3.xlarge + num_workers: 2 + + instance_pools: + orders_pool: + instance_pool_name: orders-pool + node_type_id: i3.xlarge + idle_instance_autotermination_minutes: 10 + + secret_scopes: + orders_scope: + name: orders-scope + + secrets: + orders_secret: + name: token + catalog_name: shop + schema_name: secrets + + cluster_policies: + orders_policy: + name: orders-policy + definition: "{}" + + sql_warehouses: + orders_warehouse: + name: orders-warehouse + cluster_size: Small + warehouse_type: PRO + + alerts: + orders_alert: + display_name: Orders alert + query_id: q-orders-123 + + database_instances: + orders_db_instance: + name: orders-db-instance + capacity: CU_1 + + database_catalogs: + orders_db_catalog: + name: orders-db-catalog + database_instance_name: orders-db-instance + + synced_database_tables: + orders_synced: + name: shop.gold.orders_synced + + postgres_projects: + orders_pg_project: + project_id: orders-proj + display_name: Orders + + postgres_branches: + orders_pg_branch: + branch_id: main + parent: projects/orders-proj + + postgres_endpoints: + orders_pg_endpoint: + endpoint_id: ep1 + parent: projects/orders-proj/branches/main + + postgres_catalogs: + orders_pg_catalog: + catalog_id: shop_pg + database_connection_id: conn-1 + + postgres_databases: + orders_pg_database: + database_id: orders + parent: projects/orders-proj/branches/main + + postgres_roles: + orders_pg_role: + role_id: app + parent: projects/orders-proj/branches/main + + postgres_synced_tables: + orders_pg_synced: + synced_table_id: shop.gold.orders_synced_pg + + postgres_snapshot_schedules: + orders_pg_snapshot: + branch: projects/orders-proj/branches/main diff --git a/experimental/bundletest/examples/orders_bundle/tests/README.md b/experimental/bundletest/examples/orders_bundle/tests/README.md index d66ff1fbf41..9627709db0b 100644 --- a/experimental/bundletest/examples/orders_bundle/tests/README.md +++ b/experimental/bundletest/examples/orders_bundle/tests/README.md @@ -1,18 +1,44 @@ # Example gallery -One test file per resource / capability the framework supports today. Each file's header -lists precisely what it CAN and CANNOT test locally. New files land here as the framework -grows, so this folder doubles as a record of capability over time. +One test file per resource / capability the framework supports. Each file's header lists +precisely what it CAN and CANNOT test locally. New files land here as the framework grows, +so this folder doubles as a record of capability over time. -Supported now (local DuckDB backend): +## Execution (local DuckDB backend) - `test_job_sql.py` — SQL job: run the real `.sql`, assert on output tables -- `test_job_config.py` — read a resource's declared wiring (no execution) - `test_job_nonsql.py` — non-SQL job: skips loudly (boundary demo) +- `test_job_config.py` — read a job's declared wiring (no execution) - `test_volume.py` — volume upload + read the file back (row count, columns) - `test_pipeline_end_to_end.py` — chain two SQL jobs (bronze → silver → gold), assert final table -Not yet (need the cloud backend or new handles): -- pipelines (Lakeflow/DLT) — run + assert on output tables -- dashboards — assert source tables / wiring -- alerts, permissions, clusters — config + live state +## Resource handles — static wiring, read from `databricks.yml` (no workspace) + +Every resource kind is reachable via the generic `env.resource(kind, name)` handle +(`.exists()`, `.config`, `.permissions()`, `.grants()`). Kinds that reference other +tables / artifacts / resources get a typed handle with a special accessor on top: + +| Resource | typed handle | special accessor | example | +|---|---|---|---| +| pipelines | `env.pipeline` | `.catalog` / `.schema` / `.libraries()` | `test_pipeline.py` | +| dashboards | `env.dashboard` | `.source_tables()` | `test_dashboard.py` | +| genie_spaces | `env.genie_space` | `.source_tables()` | `test_genie_space.py` | +| quality_monitors | `env.quality_monitor` | `.monitored_table()` | `test_quality_monitor.py` | +| vector_search_indexes | `env.vector_search_index` | `.source_table()` / `.endpoint_name` | `test_vector_search_index.py` | +| model_serving_endpoints | `env.model_serving_endpoint` | `.served_models()` | `test_model_serving_endpoint.py` | +| apps | `env.app` | `.command()` / `.source_code_path` | `test_app.py` | +| jobs | `env.jobs[...]` | `.run()` / `.last_run()` | `test_job_config.py` | +| volumes | `env.volume` | `.upload()` / `.file()` | `test_volume.py` | + +Every remaining user-authored kind (models, experiments, registered_models, catalogs, +schemas, external_locations, clusters, instance_pools, secret_scopes, secrets, +cluster_policies, sql_warehouses, alerts, vector_search_endpoints, database_instances, +database_catalogs, synced_database_tables, and the `postgres_*` family) is covered through +the generic handle — one example each in `test_config_resources.py`. `test_resource_generic.py` +shows the generic handle and the KeyError-safe `.exists()` for an undeclared resource. + +## Not yet — needs the cloud backend +- running pipelines / dashboards / model endpoints and asserting on their live output - running Python / Scala / R / notebook jobs +- server-defaulted or normalized config, and any live/deployed state +- introspecting a dashboard / genie space defined only by `file_path` (no inline queries) → + skips loudly with `LocalUnsupported` diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_app.py b/experimental/bundletest/examples/orders_bundle/tests/test_app.py new file mode 100644 index 00000000000..532ac23ee94 --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_app.py @@ -0,0 +1,15 @@ +"""App wiring — its launch command and source path. + +CAN test locally: +- the app exists, its launch command, and source_code_path + +CANNOT test locally — needs the cloud backend: +- deploying / starting the app +""" + + +def test_app_command(env): + app = env.app("orders_app") + assert app.exists() + assert app.command() == ["python", "app.py"] + assert app.source_code_path == "./app" diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_config_resources.py b/experimental/bundletest/examples/orders_bundle/tests/test_config_resources.py new file mode 100644 index 00000000000..0aab7711066 --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_config_resources.py @@ -0,0 +1,129 @@ +"""Config/wiring coverage for every resource kind without a dedicated typed handle. + +These kinds have no cross-resource references to encode, so they're covered uniformly through +the generic env.resource(kind, name) handle — one example per kind, each asserting the +resource exists and a real declared field (never a no-op). The kinds with special wiring +(pipelines, dashboards, genie_spaces, quality_monitors, vector_search_indexes, +model_serving_endpoints, apps) have their own dedicated test files. + +CAN test locally: existence + any statically declared field. +CANNOT test locally — needs the cloud backend: that the workspace accepted the config, +server-defaulted/normalized values, and any live state. +""" + + +# --- MLflow / models --- +def test_model(env): + assert env.resource("models", "orders_model").config["name"] == "orders_model" + + +def test_experiment(env): + assert env.resource("experiments", "orders_experiment").config["name"] == "/Shared/orders_experiment" + + +def test_registered_model(env): + rm = env.resource("registered_models", "orders_registered") + assert rm.config["catalog_name"] == "shop" + assert rm.config["schema_name"] == "ml" + + +# --- Unity Catalog --- +def test_schema(env): + assert env.resource("schemas", "orders_schema").config["catalog_name"] == "shop" + + +def test_volume(env): + vol = env.volume("raw_data") + assert vol.exists() + assert vol.config["catalog_name"] == "shop" + assert vol.config["volume_type"] == "MANAGED" + + +def test_external_location(env): + assert env.resource("external_locations", "orders_location").config["url"] == "s3://example-bucket/orders" + + +def test_secret(env): + assert env.resource("secrets", "orders_secret").config["schema_name"] == "secrets" + + +# --- Compute --- +def test_cluster(env): + assert env.resource("clusters", "orders_cluster").config["node_type_id"] == "i3.xlarge" + + +def test_instance_pool(env): + assert env.resource("instance_pools", "orders_pool").config["node_type_id"] == "i3.xlarge" + + +def test_cluster_policy(env): + assert env.resource("cluster_policies", "orders_policy").config["name"] == "orders-policy" + + +def test_sql_warehouse(env): + assert env.resource("sql_warehouses", "orders_warehouse").config["cluster_size"] == "Small" + + +# --- SQL / secrets / alerts --- +def test_secret_scope(env): + assert env.resource("secret_scopes", "orders_scope").config["name"] == "orders-scope" + + +def test_alert(env): + assert env.resource("alerts", "orders_alert").config["display_name"] == "Orders alert" + + +def test_vector_search_endpoint(env): + assert env.resource("vector_search_endpoints", "orders_vs_endpoint").config["endpoint_type"] == "STANDARD" + + +# --- Databases (Lakebase) --- +def test_database_instance(env): + assert env.resource("database_instances", "orders_db_instance").config["name"] == "orders-db-instance" + + +def test_database_catalog(env): + assert env.resource("database_catalogs", "orders_db_catalog").config["name"] == "orders-db-catalog" + + +def test_synced_database_table(env): + assert env.resource("synced_database_tables", "orders_synced").config["name"] == "shop.gold.orders_synced" + + +# --- Postgres family: same compact shape, keyed on each kind's own id field --- +def test_postgres_project(env): + assert env.resource("postgres_projects", "orders_pg_project").config["project_id"] == "orders-proj" + + +def test_postgres_branch(env): + assert env.resource("postgres_branches", "orders_pg_branch").config["parent"] == "projects/orders-proj" + + +def test_postgres_endpoint(env): + assert env.resource("postgres_endpoints", "orders_pg_endpoint").config["endpoint_id"] == "ep1" + + +def test_postgres_catalog(env): + assert env.resource("postgres_catalogs", "orders_pg_catalog").config["catalog_id"] == "shop_pg" + + +def test_postgres_database(env): + assert env.resource("postgres_databases", "orders_pg_database").config["database_id"] == "orders" + + +def test_postgres_role(env): + assert env.resource("postgres_roles", "orders_pg_role").config["role_id"] == "app" + + +def test_postgres_synced_table(env): + assert ( + env.resource("postgres_synced_tables", "orders_pg_synced").config["synced_table_id"] + == "shop.gold.orders_synced_pg" + ) + + +def test_postgres_snapshot_schedule(env): + assert ( + env.resource("postgres_snapshot_schedules", "orders_pg_snapshot").config["branch"] + == "projects/orders-proj/branches/main" + ) diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_dashboard.py b/experimental/bundletest/examples/orders_bundle/tests/test_dashboard.py new file mode 100644 index 00000000000..cd2ee35136f --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_dashboard.py @@ -0,0 +1,21 @@ +"""Dashboard wiring — which tables the dashboard's datasets read from. + +source_tables() parses the dashboard's inline serialized definition and pulls the qualified +table each dataset query reads. This catches a dashboard pointed at a stale or misspelled +table without opening a workspace. + +CAN test locally: +- the source tables a dashboard with an inline serialized_dashboard reads from +- that it's wired to the gold table the jobs produce (cross-resource wiring) + +CANNOT test locally — needs the cloud backend: +- a dashboard defined only by file_path -> skips loudly (no inline queries in databricks.yml) +- that the dashboard renders / its queries actually run +""" + + +def test_dashboard_reads_gold_table(env): + dashboard = env.dashboard("orders_overview") + assert dashboard.exists() + # aggregate_orders.sql writes shop.gold.order_summary — the dashboard reads it. + assert dashboard.source_tables() == ["shop.gold.order_summary"] diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_genie_space.py b/experimental/bundletest/examples/orders_bundle/tests/test_genie_space.py new file mode 100644 index 00000000000..377692ac99f --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_genie_space.py @@ -0,0 +1,18 @@ +"""Genie Space wiring — which tables the space is grounded on. + +A Genie space declares its tables explicitly under data_sources.tables[].identifier — a +different serialized schema from a Lakeview dashboard, so it has its own parser. + +CAN test locally: +- the source tables a space with an inline serialized_space declares + +CANNOT test locally — needs the cloud backend: +- a space defined only by file_path -> skips loudly (no inline definition in databricks.yml) +- asking the space a question / running its queries +""" + + +def test_genie_space_reads_gold_table(env): + space = env.genie_space("orders_genie") + assert space.exists() + assert space.source_tables() == ["shop.gold.order_summary"] diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_job_config.py b/experimental/bundletest/examples/orders_bundle/tests/test_job_config.py index 7781ff5489a..1f14477ca6a 100644 --- a/experimental/bundletest/examples/orders_bundle/tests/test_job_config.py +++ b/experimental/bundletest/examples/orders_bundle/tests/test_job_config.py @@ -12,8 +12,8 @@ - that the workspace actually accepted / deployed the config - server-defaulted or normalized values -NOTE: only 'jobs' exist in this bundle today; the same get_resource call works for any -resource kind (pipelines, volumes, ...) once they're added. +NOTE: env.resource(kind, name) wraps this same get_resource call for any resource kind, and +the typed handles (env.pipeline, env.dashboard, ...) add resource-specific accessors on top. """ diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_model_serving_endpoint.py b/experimental/bundletest/examples/orders_bundle/tests/test_model_serving_endpoint.py new file mode 100644 index 00000000000..3a0de2bde03 --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_model_serving_endpoint.py @@ -0,0 +1,14 @@ +"""Model serving endpoint wiring — which models it serves. + +CAN test locally: +- the endpoint exists and the model(s) it serves + +CANNOT test locally — needs the cloud backend: +- querying the endpoint / checking it is READY +""" + + +def test_endpoint_serves_model(env): + endpoint = env.model_serving_endpoint("orders_model_endpoint") + assert endpoint.exists() + assert endpoint.served_models() == ["shop.ml.orders_model"] diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_pipeline.py b/experimental/bundletest/examples/orders_bundle/tests/test_pipeline.py new file mode 100644 index 00000000000..99899bce5d8 --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_pipeline.py @@ -0,0 +1,21 @@ +"""Pipeline (Lakeflow) wiring — read a pipeline's declared config without running it. + +Running a pipeline and asserting on the tables it materializes needs a real workspace, so +it's cloud-only. Locally we read the wiring straight from databricks.yml. + +CAN test locally: +- a pipeline exists and its target catalog / schema +- which notebook / file libraries it runs, in declaration order + +CANNOT test locally — needs the cloud backend: +- running the pipeline and asserting on its output tables (@pytest.mark.cloud_only) +- server-normalized / defaulted config +""" + + +def test_pipeline_wiring(env): + pipeline = env.pipeline("enrich_orders") + assert pipeline.exists() + assert pipeline.catalog == "shop" + assert pipeline.schema == "enriched" + assert pipeline.libraries() == ["src/enrich_orders"] diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_quality_monitor.py b/experimental/bundletest/examples/orders_bundle/tests/test_quality_monitor.py new file mode 100644 index 00000000000..342f7e6e8da --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_quality_monitor.py @@ -0,0 +1,15 @@ +"""Quality monitor wiring — which table the monitor is attached to. + +CAN test locally: +- the monitor exists and the table it monitors +- that it monitors the gold table the jobs produce (cross-resource wiring) + +CANNOT test locally — needs the cloud backend: +- refreshing the monitor / reading its computed metrics +""" + + +def test_monitor_targets_gold_table(env): + monitor = env.quality_monitor("orders_quality") + assert monitor.exists() + assert monitor.monitored_table() == "shop.gold.order_summary" diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_resource_generic.py b/experimental/bundletest/examples/orders_bundle/tests/test_resource_generic.py new file mode 100644 index 00000000000..da7c380e49b --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_resource_generic.py @@ -0,0 +1,27 @@ +"""The generic env.resource(kind, name) handle — reaches ANY resource kind. + +Every resource kind is reachable through env.resource(kind, name), which reads its declared +config from the bundle. The typed handles (env.pipeline, env.dashboard, ...) add resource- +specific accessors on top of this same base for the kinds that reference tables/artifacts. + +CAN test locally: +- any declared resource's existence and its config fields +- KeyError-safe exists() for a resource that isn't declared + +CANNOT test locally — needs the cloud backend: +- that the workspace accepted / deployed the config, and server-defaulted values +""" + + +def test_generic_handle_reads_any_kind(env): + # A kind with no dedicated typed handle is still fully reachable. + catalog = env.resource("catalogs", "orders_catalog") + assert catalog.exists() + assert catalog.config["name"] == "shop" + assert catalog.grants() == [{"principal": "users", "privileges": ["USE_CATALOG"]}] + + +def test_missing_resource_does_not_exist(env): + # get_resource raises KeyError for an undeclared resource; exists() must not propagate it. + assert not env.resource("catalogs", "nope").exists() + assert not env.pipeline("nope").exists() diff --git a/experimental/bundletest/examples/orders_bundle/tests/test_vector_search_index.py b/experimental/bundletest/examples/orders_bundle/tests/test_vector_search_index.py new file mode 100644 index 00000000000..5442805ee08 --- /dev/null +++ b/experimental/bundletest/examples/orders_bundle/tests/test_vector_search_index.py @@ -0,0 +1,16 @@ +"""Vector search index wiring — its endpoint and the table it syncs from. + +CAN test locally: +- the index exists, its endpoint, and the source table it delta-syncs from +- that it syncs the gold table the jobs produce (cross-resource wiring) + +CANNOT test locally — needs the cloud backend: +- building / querying the index +""" + + +def test_index_syncs_gold_table(env): + index = env.vector_search_index("orders_index") + assert index.exists() + assert index.endpoint_name == "orders-vs-endpoint" + assert index.source_table() == "shop.gold.order_summary" diff --git a/experimental/bundletest/src/bundletest/env.py b/experimental/bundletest/src/bundletest/env.py index b97fccc7ce1..63e895ea9b6 100644 --- a/experimental/bundletest/src/bundletest/env.py +++ b/experimental/bundletest/src/bundletest/env.py @@ -2,11 +2,13 @@ from __future__ import annotations +import json import os +import re from contextlib import contextmanager from typing import TYPE_CHECKING, Any, Iterator -from bundletest.backend import Backend, JobRunFailed, RunResult +from bundletest.backend import Backend, JobRunFailed, LocalUnsupported, RunResult from bundletest.table import FileHandle, TableHandle if TYPE_CHECKING: @@ -15,12 +17,47 @@ DEFAULT_BACKEND = "local" -class JobHandle: - """A single job resource.""" +class ResourceHandle: + """A single declared bundle resource, keyed by (kind, name). - def __init__(self, backend: Backend, name: str): + Reads the resource's config through the backend seam (``get_resource``), so every handle + works unchanged on the local and cloud backends. ``kind`` is the plural key under + ``resources:`` in databricks.yml (e.g. ``jobs``, ``pipelines``, ``quality_monitors``). + Subclasses add resource-specific accessors on top of this base. + """ + + def __init__(self, backend: Backend, kind: str, name: str): self._backend = backend + self.kind = kind self.name = name + + @property + def config(self) -> dict[str, Any]: + return self._backend.get_resource(self.kind, self.name) + + def exists(self) -> bool: + """Whether the resource is DECLARED in the bundle config — not whether it is deployed. + Nothing is deployed on the local backend, and even on cloud this reflects the bundle's + declared resources, not live workspace state.""" + # get_resource does `...[name]`, which raises KeyError for an undeclared resource. + try: + self._backend.get_resource(self.kind, self.name) + return True + except KeyError: + return False + + def permissions(self) -> list: + return self.config.get("permissions", []) + + def grants(self) -> list: + return self.config.get("grants", []) + + +class JobHandle(ResourceHandle): + """A single job resource. Runs one job by name (not the whole DAG).""" + + def __init__(self, backend: Backend, name: str): + super().__init__(backend, "jobs", name) self._last: RunResult | None = None def run(self, params: dict[str, Any] | None = None, check: bool = True) -> RunResult: @@ -48,12 +85,11 @@ def __getitem__(self, name: str) -> JobHandle: return self._cache.setdefault(name, JobHandle(self._backend, name)) -class VolumeHandle: +class VolumeHandle(ResourceHandle): """A single volume resource.""" def __init__(self, backend: Backend, name: str): - self._backend = backend - self.name = name + super().__init__(backend, "volumes", name) def upload(self, src: str, dst: str | None = None) -> None: self._backend.put_file(dst or f"/Volumes/{self.name}/{os.path.basename(src)}", src) @@ -62,6 +98,149 @@ def file(self, filename: str) -> FileHandle: return FileHandle(self._backend, self.name, filename) +class PipelineHandle(ResourceHandle): + """A pipeline (Lakeflow) resource — wiring only; running it is cloud-only.""" + + def __init__(self, backend: Backend, name: str): + super().__init__(backend, "pipelines", name) + + @property + def catalog(self) -> str | None: + return self.config.get("catalog") + + @property + def schema(self) -> str | None: + # Newer pipelines set `schema`; older (DLT) pipelines set `target` for the same thing. + return self.config.get("schema") or self.config.get("target") + + def libraries(self) -> list[str]: + """Notebook/file paths the pipeline runs, in declaration order.""" + paths = [] + for lib in self.config.get("libraries", []): + if "notebook" in lib: + paths.append(lib["notebook"]["path"]) + elif "file" in lib: + paths.append(lib["file"]["path"]) + return paths + + +class DashboardHandle(ResourceHandle): + """A Lakeview dashboard resource.""" + + def __init__(self, backend: Backend, name: str): + super().__init__(backend, "dashboards", name) + + def source_tables(self) -> list[str]: + return _dashboard_source_tables(self.name, self.config.get("serialized_dashboard")) + + +class GenieSpaceHandle(ResourceHandle): + """A Genie Space resource.""" + + def __init__(self, backend: Backend, name: str): + super().__init__(backend, "genie_spaces", name) + + def source_tables(self) -> list[str]: + return _genie_source_tables(self.name, self.config.get("serialized_space")) + + +class QualityMonitorHandle(ResourceHandle): + """A quality monitor resource. Keyed on the table it monitors.""" + + def __init__(self, backend: Backend, name: str): + super().__init__(backend, "quality_monitors", name) + + def monitored_table(self) -> str: + return self.config["table_name"] + + +class VectorSearchIndexHandle(ResourceHandle): + """A vector search index resource.""" + + def __init__(self, backend: Backend, name: str): + super().__init__(backend, "vector_search_indexes", name) + + @property + def endpoint_name(self) -> str | None: + return self.config.get("endpoint_name") + + def source_table(self) -> str | None: + return self.config.get("delta_sync_index_spec", {}).get("source_table") + + +class ModelServingEndpointHandle(ResourceHandle): + """A model serving endpoint resource.""" + + def __init__(self, backend: Backend, name: str): + super().__init__(backend, "model_serving_endpoints", name) + + def served_models(self) -> list[str]: + """Model names served by the endpoint.""" + cfg = self.config.get("config", {}) + # `served_entities` is the current shape; `served_models` is the deprecated one. + entities = cfg.get("served_entities") or cfg.get("served_models", []) + return [e.get("entity_name") or e.get("model_name") for e in entities] + + +class AppHandle(ResourceHandle): + """A Databricks App resource.""" + + def __init__(self, backend: Backend, name: str): + super().__init__(backend, "apps", name) + + def command(self) -> list[str]: + return self.config.get("config", {}).get("command", []) + + @property + def source_code_path(self) -> str | None: + return self.config.get("source_code_path") + + +def _load_serialized(name: str, serialized: Any) -> dict: + """Parse an inline serialized definition: a dict when inlined as YAML, or a JSON string. + + A file_path-only resource carries no inline definition in databricks.yml, so it can't be + introspected locally — that's a loud skip, not a failure. (On the cloud backend + `bundle summary` inlines file_path into the serialized field, so it resolves there.)""" + if serialized is None: + raise LocalUnsupported(f"{name!r} defined by file_path — no inline definition to parse locally") + return serialized if isinstance(serialized, dict) else json.loads(serialized) + + +# Tables a query reads: the qualified (dotted) identifier right after FROM / JOIN. Matching +# only dotted names skips CTE names and aliases (which are unqualified); an outer backtick +# pair is unwrapped. Good enough for wiring, not a SQL parser — it won't unwrap per-segment +# backticks (`a`.`b`) and would match a name inside a `-- FROM ...` comment. +_FROM_JOIN = re.compile(r"\b(?:FROM|JOIN)\s+`?([A-Za-z_]\w*(?:\.\w+)+)`?", re.IGNORECASE) + + +def _dashboard_source_tables(name: str, serialized: Any) -> list[str]: + """Qualified source tables read by a Lakeview dashboard's dataset queries.""" + spec = _load_serialized(name, serialized) + tables: list[str] = [] + for dataset in spec.get("datasets", []): + # Lakeview stores a query as queryLines (current) or a single query string (older). + query = "".join(dataset.get("queryLines", [])) or dataset.get("query", "") + for m in _FROM_JOIN.finditer(query): + if m.group(1) not in tables: + tables.append(m.group(1)) + return tables + + +def _genie_source_tables(name: str, serialized: Any) -> list[str]: + """Source tables a Genie space is grounded on. + + A space's serialized_space uses a different schema from a dashboard: it declares its + tables explicitly under data_sources.tables[].identifier rather than in SQL queries.""" + spec = _load_serialized(name, serialized) + tables: list[str] = [] + for table in spec.get("data_sources", {}).get("tables", []): + identifier = table.get("identifier") + if identifier and identifier not in tables: + tables.append(identifier) + return tables + + class BundleEnv: """A deployed bundle under test, backed by a single ``Backend``.""" @@ -82,12 +261,37 @@ def seed(self, table: str, rows: list[dict[str, Any]]) -> None: self.backend.seed_table(table, rows) # --- resource handles --- + def resource(self, kind: str, name: str) -> ResourceHandle: + """Generic handle for any resource kind (the plural `resources:` key).""" + return ResourceHandle(self.backend, kind, name) + def table(self, fqn: str) -> TableHandle: return TableHandle(self.backend, fqn) def volume(self, name: str) -> VolumeHandle: return VolumeHandle(self.backend, name) + def pipeline(self, name: str) -> PipelineHandle: + return PipelineHandle(self.backend, name) + + def dashboard(self, name: str) -> DashboardHandle: + return DashboardHandle(self.backend, name) + + def genie_space(self, name: str) -> GenieSpaceHandle: + return GenieSpaceHandle(self.backend, name) + + def quality_monitor(self, name: str) -> QualityMonitorHandle: + return QualityMonitorHandle(self.backend, name) + + def vector_search_index(self, name: str) -> VectorSearchIndexHandle: + return VectorSearchIndexHandle(self.backend, name) + + def model_serving_endpoint(self, name: str) -> ModelServingEndpointHandle: + return ModelServingEndpointHandle(self.backend, name) + + def app(self, name: str) -> AppHandle: + return AppHandle(self.backend, name) + def run_job(self, name: str, params: dict[str, Any] | None = None, check: bool = True) -> RunResult: return self.jobs[name].run(params, check=check)