diff --git a/Makefile b/Makefile index a4009383a..d2fd9bfe4 100644 --- a/Makefile +++ b/Makefile @@ -118,19 +118,77 @@ app-test: $(if $(COV),--cov=src/databricks_labs_dqx_app/backend --cov-report=term-missing --cov-report=xml:coverage-app.xml) # Grant Unity Catalog permissions after bundle deploy. +# # Usage: make app-grant-permissions PROFILE=my-profile +# make app-grant-permissions PROFILE=my-profile TARGET=dev \ +# BUNDLE_VARS='--var=catalog_name=foo' +# +# BUNDLE_VARS must match the overrides used at deploy time. The script +# reads bundle variables to discover which catalog / schema / volume to +# grant on, so an unforwarded override would point the GRANTs at the +# wrong objects. app-grant-permissions: @test -n "$(PROFILE)" || (echo "Usage: make app-grant-permissions PROFILE= [TARGET=]"; exit 1) - app/scripts/post_deploy_grants.sh -p $(PROFILE) $(if $(TARGET),-t $(TARGET)) + app/scripts/post_deploy_grants.sh -p $(PROFILE) $(if $(TARGET),-t $(TARGET)) $(if $(BUNDLE_VARS),-- $(BUNDLE_VARS)) -# Full deploy: build, bundle deploy, grant permissions, and start the app. +# Adopt pre-existing storage resources into bundle management. +# +# Use ONCE per target on workspaces where the schemas / volume / +# Lakebase instance already exist (e.g. from a previous bootstrap- +# script deploy, or from manual creation). The app's Postgres schema +# inside ``databricks_postgres`` is not bundle-managed and does not +# need binding — the app re-creates it on first connection if missing. +# Without binding, ``databricks bundle deploy`` would try to CREATE +# them and fail with "already exists" / "Instance name is not unique". +# +# Bind is idempotent at the CLI level — re-binding the same resource +# is a no-op. Skip this target on fresh workspaces; ``bundle deploy`` +# creates the resources directly. +# +# Usage: make app-bind PROFILE=my-profile TARGET=dev +# make app-bind PROFILE=my-profile TARGET=dev \ +# BUNDLE_VARS='--var=lakebase_instance_name=' +# +# BUNDLE_VARS forwards arbitrary ``--var key=value`` arguments through +# to the bundle CLI. Use the same override here as you intend to pass +# to ``make app-deploy`` — the bind step reads the resolved bundle +# variables to know which instance/schema name to bind to, so an +# override applied only at deploy time would bind the wrong resource. +app-bind: + @test -n "$(PROFILE)" || (echo "Usage: make app-bind PROFILE= TARGET="; exit 1) + @test -n "$(TARGET)" || (echo "Usage: make app-bind PROFILE= TARGET="; exit 1) + app/scripts/bind_resources.sh -p $(PROFILE) -t $(TARGET) $(if $(BUNDLE_VARS),-- $(BUNDLE_VARS)) + +# Full deploy: build, bundle deploy (creates storage on fresh +# workspaces, updates managed resources otherwise), grant permissions +# to the app SP, and start the app. Run ``make app-bind`` once before +# the FIRST deploy on a workspace where the storage was previously +# provisioned out-of-band — otherwise the bundle will try to CREATE +# the existing resources and fail. +# # Usage: make app-deploy PROFILE=my-profile TARGET=dev +# make app-deploy PROFILE=my-profile TARGET=dev \ +# BUNDLE_VARS='--var=lakebase_instance_name=' +# +# BUNDLE_VARS forwards arbitrary ``--var key=value`` arguments to +# every step of the deploy: ``bundle deploy``, ``post_deploy_grants.sh``, +# and ``bundle run``. Threading the same overrides through all three is +# load-bearing — the grants script re-runs ``bundle validate`` to +# discover the deployed catalog / schema / volume, and without the +# overrides it would issue GRANTs on the bundle defaults instead of +# the resources actually deployed. +# +# The common use case is overriding ``lakebase_instance_name`` when the +# original name is locked by Lakebase's ~7-day soft-delete retention +# after a manual delete (the deploy errors out with "Instance name is +# not unique" otherwise). Per-deploy CLI overrides keep ``databricks.yml`` +# clean. app-deploy: app-build @test -n "$(PROFILE)" || (echo "Usage: make app-deploy PROFILE= TARGET="; exit 1) @test -n "$(TARGET)" || (echo "Usage: make app-deploy PROFILE= TARGET="; exit 1) - cd app && databricks bundle deploy -p $(PROFILE) -t $(TARGET) - app/scripts/post_deploy_grants.sh -p $(PROFILE) -t $(TARGET) - cd app && databricks bundle run $(APP_NAME) -p $(PROFILE) -t $(TARGET) + cd app && databricks bundle deploy -p $(PROFILE) -t $(TARGET) $(BUNDLE_VARS) + app/scripts/post_deploy_grants.sh -p $(PROFILE) -t $(TARGET) $(if $(BUNDLE_VARS),-- $(BUNDLE_VARS)) + cd app && databricks bundle run $(APP_NAME) -p $(PROFILE) -t $(TARGET) $(BUNDLE_VARS) APP_NAME ?= dqx-studio @@ -166,4 +224,4 @@ lock-dependencies: perl -pi -e 's|registry = "https://[^"]*"|registry = "https://pypi.org/simple"|g' uv.lock .DEFAULT: all -.PHONY: all clean dev lint fmt test integration e2e perf anomaly coverage combine-coverage docs-build docs-serve-dev docs-install docs-serve docs-clean app-install app-build app-start-dev app-stop-dev app-check app-test app-grant-permissions app-deploy fork-sync build lock-dependencies lock-app-dependencies +.PHONY: all clean dev lint fmt test integration e2e perf anomaly coverage combine-coverage docs-build docs-serve-dev docs-install docs-serve docs-clean app-install app-build app-start-dev app-stop-dev app-check app-test app-grant-permissions app-bind app-deploy fork-sync build lock-dependencies lock-app-dependencies diff --git a/app/.build-constraints.txt b/app/.build-constraints.txt index c1518967e..022b0933a 100644 --- a/app/.build-constraints.txt +++ b/app/.build-constraints.txt @@ -101,15 +101,15 @@ markupsafe==3.0.3 \ --hash=sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a \ --hash=sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50 # via jinja2 -packaging==26.1 \ - --hash=sha256:5d9c0669c6285e491e0ced2eee587eaf67b670d94a19e94e3984a481aba6802f \ - --hash=sha256:f042152b681c4bfac5cae2742a55e103d27ab2ec0f3d88037136b6bfe7c9c5de +packaging==26.2 \ + --hash=sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e \ + --hash=sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661 # via # dunamai # hatchling -pathspec==1.0.4 \ - --hash=sha256:0210e2ae8a21a9137c0d470578cb0e595af87edaa6ebf12ff176f14a02e0e645 \ - --hash=sha256:fb6ae2fd4e7c921a165808a552060e722767cfa526f99ca5156ed2ce45a5c723 +pathspec==1.1.1 \ + --hash=sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a \ + --hash=sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189 # via hatchling pluggy==1.6.0 \ --hash=sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3 \ @@ -119,9 +119,9 @@ tomlkit==0.14.0 \ --hash=sha256:592064ed85b40fa213469f81ac584f67a4f2992509a7c3ea2d632208623a3680 \ --hash=sha256:cf00efca415dbd57575befb1f6634c4f42d2d87dbba376128adb42c121b87064 # via uv-dynamic-versioning -trove-classifiers==2026.1.14.14 \ - --hash=sha256:00492545a1402b09d4858605ba190ea33243d361e2b01c9c296ce06b5c3325f3 \ - --hash=sha256:1f9553927f18d0513d8e5ff80ab8980b8202ce37ecae0e3274ed2ef11880e74d +trove-classifiers==2026.4.28.13 \ + --hash=sha256:8f4b1eb4e16296b57d612965444f87a83861cc989a0451ac97fe4265ddef03b8 \ + --hash=sha256:c85bb8a53c3de7330d1699b844ed9fb809a602a09ac15dc79ad6d1a509be0676 # via hatchling uv-dynamic-versioning==0.14.0 \ --hash=sha256:574fbc07e87ace45c01d55967ad3b864871257b98ff5b8ac87c261227ac8db5b \ diff --git a/app/CLAUDE.md b/app/CLAUDE.md index ed734dcf2..7dec24c75 100644 --- a/app/CLAUDE.md +++ b/app/CLAUDE.md @@ -36,36 +36,72 @@ RBAC is enforced — routes use `require_role(*roles)` from `backend/dependencie ## Internal Storage -App uses a dedicated catalog selected at install time, with two schemas (managed by `MigrationRunner` in `backend/migrations/`): +App uses a **hybrid backend** — analytical/append tables in Delta, OLTP +tables in Lakebase Postgres. Both backends are managed by their own +migration runner in `backend/migrations/`. Schemas, volume, and Lakebase +instance are declared as bundle resources in `databricks.yml` with +`lifecycle.prevent_destroy: true`, so `databricks bundle destroy` cannot +drop them — see "Bundle conventions" below. The app's `dqx_studio` +Postgres schema (inside the `databricks_postgres` admin database on the +Lakebase instance) is created at startup, not provisioned by the bundle, +but is protected transitively by the instance-level guard. ``` {user_catalog} - ├── dqx_app ← main schema (SP-managed) - │ ├── dq_app_settings ← key/value app configuration - │ ├── dq_quality_rules ← active/approved rules - │ ├── dq_quality_rules_history ← rule change audit log - │ ├── dq_role_mappings ← role → workspace group mappings (RBAC) - │ ├── dq_comments ← comment threads on rules/runs - │ ├── dq_profiling_results ← profiler run results (suggestions in generated_rules_json) - │ ├── dq_validation_runs ← dryrun + scheduled run history - │ ├── dq_quarantine_records ← invalid rows captured by runs - │ ├── dq_metrics ← per-run quality metrics for trend tracking - │ ├── dq_schedule_configs ← per-schedule config (cron/interval, target rules) - │ ├── dq_schedule_configs_history ← schedule config change audit log - │ ├── dq_schedule_runs ← scheduler last/next run state (survives restarts) - │ └── dq_migrations ← migration version tracker - └── dqx_app_tmp ← temp views created via OBO for profiler/dryrun jobs + ├── dqx_studio ← main schema (SP-managed) + │ ├── dq_profiling_results (Delta) profiler run results + │ ├── dq_validation_runs (Delta) dryrun + scheduled run history + │ ├── dq_quarantine_records (Delta) invalid rows captured by runs + │ ├── dq_metrics (Delta) per-run quality metrics for trend tracking + │ ├── dq_app_settings (OLTP*) key/value app configuration + │ ├── dq_quality_rules (OLTP*) active/approved rules + │ ├── dq_quality_rules_history (OLTP*) rule change audit log + │ ├── dq_role_mappings (OLTP*) role → workspace group mappings (RBAC) + │ ├── dq_comments (OLTP*) comment threads on rules/runs + │ ├── dq_schedule_configs (OLTP*) per-schedule config (cron/interval, target rules) + │ ├── dq_schedule_configs_history (OLTP*) schedule config change audit log + │ ├── dq_schedule_runs (OLTP*) scheduler last/next run state (survives restarts) + │ └── dq_migrations (Delta) Delta migration version tracker + ├── dqx_studio_tmp ← temp views created via OBO for profiler/dryrun jobs + └── dqx_studio.wheels (volume) ← DQX + task-runner wheels uploaded at app startup + +Lakebase database (when enabled, default = `dqx-studio-lakebase`): + └── dqx_studio (database) + └── public (schema, configurable via DQX_LAKEBASE_SCHEMA) + ├── dq_app_settings, dq_role_mappings, dq_quality_rules, + │ dq_quality_rules_history, dq_comments, dq_schedule_configs, + │ dq_schedule_configs_history, dq_schedule_runs + └── dq_migrations (Postgres migration version tracker) ``` -A separate UC volume (`{catalog}.dqx_app.wheels` by default) holds the DQX + task-runner wheels uploaded at app startup. +`(OLTP*)` = lives in **Lakebase Postgres** when +`lakebase_instance_name` is set, otherwise **Delta** (the +`v2: Delta OLTP fallback` migration). ## Key Decisions -- **No config.yaml** — all settings stored in Delta tables -- **Dedicated catalog** — user selects at install, `dqx_app` schema created by the app -- **Rule promotion** — export rules then deploy separately to prod; or save directly to prod checks table -- **Internal storage** — Delta table (preferred); Lakebase also supported as option at install time -- **Target environments** — Dev, UAT/QA (prod-like data); app is not intended for production rule execution +- **No config.yaml** — all settings stored in Delta or Lakebase tables. +- **Dedicated catalog** — user selects at install; `dqx_studio` and `dqx_studio_tmp` schemas are declared as bundle resources and created by `databricks bundle deploy`. +- **Hybrid storage** — high-volume append tables in Delta; transactional/low-latency tables in Lakebase Postgres. +- **Rule promotion** — export rules then deploy separately to prod; or save directly to prod checks table. +- **Target environments** — Dev, UAT/QA (prod-like data); app is not intended for production rule execution. + +## Bundle conventions + +Stateful resources declared in `databricks.yml`: + +- `resources.schemas.main_schema` — `dqx_studio` schema +- `resources.schemas.tmp_schema` — `dqx_studio_tmp` schema +- `resources.volumes.wheels` — wheels volume +- `resources.database_instances.lakebase` — Lakebase Postgres instance (autoscaling) + +Each carries `lifecycle.prevent_destroy: true` (Databricks CLI 0.268+), which blocks `databricks bundle destroy` and any deploy that would force-replace the resource. To intentionally tear something down: drop the flag, `databricks bundle deployment unbind -t `, then destroy. + +The app connects to the always-present `databricks_postgres` admin database on the Lakebase instance (set as the default `lakebase_database_name`) and creates its own `dqx_studio` Postgres schema there on first start. No DAB resource is needed to provision a per-app logical database; the bundle stays fully declarative. We deliberately do not use `database_catalogs` because it also creates a Unity Catalog catalog and therefore requires `CREATE CATALOG` on the metastore — a permission most app deployers don't hold. + +For workspaces where the schemas / volume / Lakebase instance already exist (e.g. created out-of-band before this layout existed), run `make app-bind PROFILE=... TARGET=...` once per target to adopt them — otherwise `databricks bundle deploy` errors out with "already exists" / "Instance name is not unique". + +Privileges on UC objects for the auto-created app SP are still reapplied with `scripts/post_deploy_grants.sh` after each deploy, because the app SP's UUID isn't known at bundle-write time. ## Architecture diff --git a/app/DEPLOYMENT.md b/app/DEPLOYMENT.md index 5dc182529..2276eb945 100644 --- a/app/DEPLOYMENT.md +++ b/app/DEPLOYMENT.md @@ -4,21 +4,48 @@ Production deployment uses [Declarative Automation Bundles](https://docs.databri ## Prerequisites -Before starting, make sure you have all of the following — several steps require elevated permissions, so confirm access before you begin. +Before you start, confirm you have **all** of the items below. The single most common deployment failure is missing one permission — and the error you see is almost always downstream of the missing grant, not on the grant itself. -**Tooling** -- **Databricks CLI** installed and authenticated against your workspace -- **`jq`** (used by the post-deploy grants script) -- **`make`** (used by the one-command deploy target) +### Tooling -**Access** -- **Workspace admin** — required to create service principals, grant catalog permissions, and enable workspace-level features +- **Databricks CLI** v0.268+ installed and authenticated against your workspace (`databricks auth login -p `). v0.268 is the minimum that supports `lifecycle.prevent_destroy` on bundle resources. +- **`jq`** (used by the post-deploy grants script and the resource-bind helper) +- **`make`** (drives the one-command deploy target) -**Workspace configuration** -- An **existing Unity Catalog catalog** where the app's schemas and volumes will be created — the bundle does not create the catalog itself -- **Databricks Apps** feature enabled on the workspace -- **User token passthrough** enabled for Databricks Apps (see [Step 2](#step-2-enable-user-token-passthrough)) -- **Serverless compute** enabled on the workspace (the task-runner job runs on serverless) +### Required permissions + +The deploying user (you) needs the permissions below. They are **all** consumed by `make app-deploy`; deployment will halt the first time it hits a missing one. We've listed which step in the flow each permission unblocks so you can debug surgically if a grant gets missed. + +| # | Permission | Granted on | Used by | What fails without it | +|---|---|---|---|---| +| 1 | **Workspace access** entitlement | You, in the workspace | All CLI calls | `databricks` CLI can't reach the workspace | +| 2 | **Databricks SQL access** entitlement | You, in the workspace | `bundle deploy` (creates the X-Small SQL warehouse) | `Error: not authorized to create SQL Endpoint` | +| 3 | **Allow cluster create** entitlement | You, in the workspace | `bundle deploy` (warehouse + job clusters) | Warehouse / job creation rejected | +| 4 | **Databricks Apps: Can Manage** workspace permission | You, in the workspace | `bundle deploy` of the App resource | App creation rejected | +| 5 | **Databricks Database (Lakebase): Manager** entitlement | You, in the workspace | `bundle deploy` of the `database_instances` resource | `Error: User does not have permission to create database instances` | +| 6 | **USE CATALOG** + **CREATE SCHEMA** on `` | Your user or an admin group you're in | `bundle deploy` of the `schemas` and `volumes` resources | `Error: User does not have CREATE_SCHEMA on catalog ''` | +| 7 | **MANAGE** on `` (or be the catalog owner) | Your user or an admin group you're in | `post_deploy_grants.sh` (issues `GRANT USE CATALOG / ALL PRIVILEGES … TO ` and `… TO account users`) | `Error: User does not have privilege MANAGE on catalog ''` | +| 8 | **Service Principal: User** role on the task-runner SP | Your user, on the SP you'll use as `dqx_service_principal_application_id` | `bundle deploy` of the `jobs.dqx_task_runner` resource (sets `run_as.service_principal_name`) | `Error: User is not authorized to use this service principal` | +| 9 | **Service Principal: Manager** role on the task-runner SP, *or* a pre-shared OAuth client secret | Your user, on the same SP | Only needed if you want to **mint a fresh OAuth secret yourself** for the task-runner (e.g. via `databricks service-principal-secrets-proxy create `) | `Error: User is not authorized to perform this operation` when minting a new secret | +| 10 | **Account admin** (one-time, post-deploy) | Account level | Updating the app's OAuth custom-app integration to include the `all-apis` scope (see [Expand OAuth Scopes](#optional-expand-oauth-scopes)) | Some app features (job submission, advanced SCIM lookups) return 403 | + +**Two convenience patterns** that reduce the per-user grants in rows 6 and 7: + +- **Make the catalog ownership easy:** ask an admin to add you to an existing UC-admin group that already holds `MANAGE` (or `ALL PRIVILEGES`) on ``. This unlocks rows 6 and 7 in one membership change instead of two per-object grants. +- **Workspace admin shortcut:** if you become workspace admin, rows 1–5 + 8–9 collapse automatically. Rows 6 and 7 (UC) and 10 (account admin) still need to be granted explicitly — workspace admin does **not** confer Unity Catalog or account-level rights. + +### Workspace features that must be enabled + +These are configured at the workspace or account level — not by you, not by the bundle. Confirm with your admin before the first deploy: + +- **Databricks Apps** is enabled on the workspace +- **User token passthrough** (a.k.a. user authorization / OBO) is enabled for Databricks Apps — see [Step 2](#step-2-enable-user-token-passthrough). Without this the app can't make OBO calls and Unity Catalog browsing fails. +- **Serverless compute** is enabled on the workspace — the task-runner job runs exclusively on serverless +- **Lakebase Postgres** is enabled on the workspace (default OLTP backend). The Lakebase instance is declared as a bundle resource (`resources.database_instances.lakebase`) with `lifecycle.prevent_destroy: true` so a `bundle destroy` cannot drop it and wipe OLTP state — see [Stateful storage and destroy protection](#step-3-stateful-storage-and-destroy-protection). The app connects to the always-present `databricks_postgres` admin database on the instance and creates its own `dqx_studio` Postgres schema inside it on first connection — no separate logical-DB provisioning step. + +### The catalog must already exist + +The bundle **does not create the catalog itself** — that's deliberate. Catalogs are typically owned by a governance team and creating them requires `CREATE CATALOG` on the metastore. Pick an existing catalog you (or an admin group you're in) have rights on, and set `catalog_name` in [Step 4](#step-4-configure-databricksyml). The bundle creates the schemas (`dqx_studio`, `dqx_studio_tmp`) and the wheels volume *inside* that catalog — no `CREATE CATALOG` permission required at the metastore level. ## Step 1: Create a Service Principal @@ -28,7 +55,7 @@ The bundle requires a service principal to run the task-runner job. This is sepa 1. Go to **Settings → Identity and Access → Service Principals** 2. Click **Add service principal → Create new** 3. Give it a name (e.g., `dqx-task-runner-sp`) -4. Note the **Application ID** — you'll use it in Step 3 as `dqx_service_principal_application_id` +4. Note the **Application ID** — you'll use it in [Step 4](#step-4-configure-databricksyml) as `dqx_service_principal_application_id` 5. **Grant yourself (or the identity you'll deploy the bundle with) the `User` role on this new SP.** Open the SP you just created, go to the **Permissions** tab, click **Add permissions**, search for your user (or deploy-time principal), and assign the role **`User`** (equivalent to `servicePrincipal.user` in the SCIM API). This lets your deploying identity configure jobs with `run_as: service_principal_name` pointing at this SP. Without it, `databricks bundle deploy` will fail with a permission error when it tries to set up the task-runner job. @@ -44,7 +71,36 @@ The app uses On-Behalf-Of (OBO) tokens to access Unity Catalog resources with th Contact your workspace admin or enable it via the workspace settings if not already active. -## Step 3: Configure `databricks.yml` +## Step 3: Stateful storage and destroy protection + +DQX Studio's stateful resources — the two schemas (`dqx_studio`, `dqx_studio_tmp`), the wheels volume, and the Lakebase instance — are all declared as bundle resources in `app/databricks.yml`. Each one carries `lifecycle.prevent_destroy: true` (Databricks CLI 0.268+), which **blocks `databricks bundle destroy` from dropping the resource** and wiping the data. Use this command line to verify: + +```bash +grep -A1 'lifecycle:' app/databricks.yml | head +``` + +You'll see one `prevent_destroy: true` for each of: `schemas.main_schema`, `schemas.tmp_schema`, `volumes.wheels`, `database_instances.lakebase`. + +> **The app's `dqx_studio` Postgres schema** (inside the `databricks_postgres` admin database on the Lakebase instance) is created by the app at first start. It's stateful but lives below the resource layer DABs models, so `prevent_destroy` doesn't apply to it directly. The instance-level guard above is what protects it: as long as `database_instances.lakebase` survives, the schema and its tables survive. + +What this means in practice: + +- **Fresh workspace** — `make app-deploy` does everything in one command: `databricks bundle deploy` provisions the schemas → volume → Lakebase instance → SQL warehouse → job → app in dependency order, then `post_deploy_grants.sh` issues catalog/schema/volume GRANTs. +- **Existing workspace** where these resources were created out-of-band (e.g. from a previous version of this app that used a bootstrap script) — you must **bind** them into bundle management once per target. See [Migrating an existing workspace](#migrating-an-existing-workspace). +- **Schema drift** — if you change `catalog_name`, `schema_name`, or `lakebase_instance_name` in a way that would force the bundle to delete and recreate the resource, `prevent_destroy` blocks the destroy step and the deploy fails fast (good — the alternative is silent data loss). Treat those names as immutable. +- **Intentional teardown** — to drop a protected resource, remove `lifecycle.prevent_destroy: true` from `databricks.yml`, run `databricks bundle deployment unbind -t ` to detach it from bundle state, then destroy it manually. + +### Migrating an existing workspace + +If the workspace was previously deployed with the old bootstrap-script flow (or if the resources were created manually), `databricks bundle deploy` will fail with "already exists" / "Instance name is not unique" on the first run because it's trying to CREATE resources that already exist. Fix this in one command: + +```bash +make app-bind PROFILE= TARGET= +``` + +`make app-bind` invokes `app/scripts/bind_resources.sh`, which calls `databricks bundle deployment bind` for each stateful resource (one bind per target). After bind, `bundle deploy` sees the resource as already-managed and does diff-and-update instead of CREATE. Bind is a one-time operation per target — once bound, subsequent deploys don't need it. Re-running `make app-bind` on a fully-bound target is a safe no-op. + +## Step 4: Configure `databricks.yml` Update a deploy target. The minimum required is a `catalog_name` and `dqx_service_principal_application_id`; everything else has a sensible default and can be overridden per target. In `app/databricks.yml`: @@ -54,15 +110,8 @@ targets: workspace: profile: variables: - # Required catalog_name: dqx_service_principal_application_id: - - # Optional — uncomment and override defaults per target as needed - # admin_group: - # app_name: - # sql_warehouse_name: - # schema_name: presets: trigger_pause_status: PAUSED ``` @@ -78,25 +127,35 @@ All target-level variables, their defaults, and what they control: | `admin_group` | `admins` | No | Workspace group whose members get the in-app `ADMIN` role unconditionally (bootstrap admin path). The default `admins` is the built-in workspace admins group — every workspace admin becomes a DQX admin automatically. Override with a dedicated group (e.g. `dqx-admins-prod`) for narrower bootstrap access. Additional roles are assigned at runtime via the in-app Role Management UI. | | `app_name` | `dqx-studio` | No | Deployed Databricks App name. Override per target (e.g. `dqx-studio-dev`, `dqx-studio-prod`) when deploying multiple targets to the same workspace, or for personal sandboxes. | | `sql_warehouse_name` | `dqx-studio-sql-warehouse` | No | Deployed SQL warehouse name (the bundle creates an X-Small serverless warehouse for app queries). Override per target to avoid duplicates in shared workspaces. | -| `schema_name` | `dqx_app` | No | Main schema inside the catalog — holds rules, run history, role mappings, and other app state. Override only if you need a non-default schema layout. | +| `schema_name` | `dqx_studio` | No | Main schema — holds run history, profiling, metrics, quarantine, and OLTP fallback tables. Declared as `resources.schemas.main_schema` in the bundle with `lifecycle.prevent_destroy: true`. | +| `tmp_schema_name` | `dqx_studio_tmp` | No | Per-user temp-view schema. Declared as `resources.schemas.tmp_schema` with `lifecycle.prevent_destroy: true`. | +| `wheels_volume_name` | `wheels` | No | UC volume under `.` for the DQX + task-runner wheels. Declared as `resources.volumes.wheels` with `lifecycle.prevent_destroy: true`. | +| `lakebase_instance_name` | `dqx-studio-lakebase` | No | Lakebase Postgres instance for OLTP state. Declared as `resources.database_instances.lakebase` with `lifecycle.prevent_destroy: true`. Autoscaling by default per [Lakebase Autoscaling](https://docs.databricks.com/aws/en/oltp/upgrade-to-autoscaling). | +| `lakebase_database_name` | `databricks_postgres` | No | Logical Postgres database inside the Lakebase instance the app connects to. Defaults to `databricks_postgres` (always present, no provisioning step). All DQX tables live in a dedicated `dqx_studio` Postgres schema inside this database, so multiple apps can safely share the same `databricks_postgres` on one Lakebase instance. Override only if you've manually created a different logical DB you want to use. | +| `lakebase_capacity` | `CU_1` | No | Lakebase compute capacity. Valid values: `CU_1`, `CU_2`, `CU_4`, `CU_8`. To resize an existing instance, change this value and redeploy. Bump up if Lakebase queries queue in the app logs. | -> **Note on duplicate names in Databricks:** SQL warehouses, jobs, and apps within the same workspace are tracked by ID, not by name, so technically duplicates are allowed. But operators browse the Jobs / Apps / Warehouses UI by name, so distinct names per target are strongly recommended when you deploy more than one target to the same workspace. +> **Note on duplicate names in Databricks:** SQL warehouses, jobs, and apps within the same workspace are tracked by ID, not by name, so technically duplicates are allowed. Lakebase database instances are looked up by name by the app at runtime, so they're effectively unique-per-workspace. Operators browse the Jobs / Apps / Warehouses / Databases UI by name, so distinct names per target are strongly recommended when you deploy more than one target to the same workspace. -## Step 4: One-Command Deploy (recommended) +## Step 5: One-Command Deploy (recommended) Build, deploy, grant permissions, and start the app in a single command: ```bash +# Existing workspace where the storage was previously bootstrapped +# out-of-band (e.g. with the older bootstrap script): bind once. +make app-bind PROFILE= TARGET= + +# Every deploy (fresh or otherwise): make app-deploy PROFILE= TARGET= ``` -This runs the following steps automatically: -1. `make app-build` — builds the frontend and wheels -2. `databricks bundle deploy` — creates the warehouse, schemas, volume, job, and app -3. `app/scripts/post_deploy_grants.sh` — discovers both SPs and executes all `GRANT` statements -4. `databricks bundle run` — starts the app +`make app-deploy` runs the following steps automatically: +1. `make app-build` — builds the frontend and wheels. +2. `databricks bundle deploy` — provisions or updates the schemas, wheels volume, Lakebase instance, SQL warehouse, task-runner job, and Databricks App in dependency order. Stateful resources carry `lifecycle.prevent_destroy: true` so a future destroy can't drop them — see [Step 3](#step-3-stateful-storage-and-destroy-protection). +3. `app/scripts/post_deploy_grants.sh` — discovers both service principals and executes the `GRANT` statements on the catalog, schemas, and volume (the auto-created app SP's UUID isn't known at bundle-write time, which is why grants live in a post-deploy script). Lakebase grants are handled by the bundle's `database` resource binding. +4. `databricks bundle run` — starts the app. -> **First start**: The app runs database migrations and uploads DQX wheels to the UC volume. If the task-runner job runs before the app has started at least once, it will fail to find its wheels. Wait for `"Uploaded databricks_labs_dqx-..."` in the logs before triggering runs. +> **First start**: The app runs both Delta and Lakebase database migrations on startup, and uploads DQX wheels to the UC volume. If the task-runner job runs before the app has started at least once, it will fail to find its wheels. Wait for `"Uploaded databricks_labs_dqx-..."` in the logs before triggering runs. If Lakebase is enabled, also wait for `"Lakebase OLTP routing enabled"` before opening the UI — when Lakebase is configured and init fails, the app refuses to start (logged as `"Lakebase initialisation failed ... Refusing to start"`) and the Apps platform will restart the container. Silent fallback to Delta is intentionally disallowed because it would split OLTP writes across two physical stores and orphan prior Lakebase data on every flap. To intentionally run on Delta only, unset `DQX_LAKEBASE_INSTANCE_NAME`. ### Step-by-step alternative @@ -106,11 +165,15 @@ If you prefer to run each step individually: # Build make app-build -# Deploy the bundle +# (One-time, only on a workspace whose storage was created out-of-band) +make app-bind PROFILE= TARGET= + +# Deploy the bundle (creates / updates schemas, volume, Lakebase +# instance, SQL warehouse, task-runner job, app) cd app && databricks bundle deploy -p -t -# Grant permissions (auto-discovers SP IDs and catalog from bundle config) -make app-grant-permissions PROFILE= +# Grant permissions to the app SP (auto-discovered after deploy) +make app-grant-permissions PROFILE= TARGET= # Start the app cd app && databricks bundle run dqx-studio -p -t @@ -126,20 +189,22 @@ The grant script discovers both SPs automatically. If you need to run the SQL ma -- App service principal GRANT USE CATALOG ON CATALOG TO ``; -GRANT ALL PRIVILEGES ON SCHEMA .dqx_app TO ``; -GRANT ALL PRIVILEGES ON SCHEMA .dqx_app_tmp TO ``; -GRANT ALL PRIVILEGES ON VOLUME .dqx_app.wheels TO ``; +GRANT ALL PRIVILEGES ON SCHEMA .dqx_studio TO ``; +GRANT ALL PRIVILEGES ON SCHEMA .dqx_studio_tmp TO ``; +GRANT ALL PRIVILEGES ON VOLUME .dqx_studio.wheels TO ``; -- Job service principal (task runner) GRANT USE CATALOG ON CATALOG TO ``; -GRANT ALL PRIVILEGES ON SCHEMA .dqx_app TO ``; -GRANT ALL PRIVILEGES ON SCHEMA .dqx_app_tmp TO ``; -GRANT ALL PRIVILEGES ON VOLUME .dqx_app.wheels TO ``; +GRANT ALL PRIVILEGES ON SCHEMA .dqx_studio TO ``; +GRANT ALL PRIVILEGES ON SCHEMA .dqx_studio_tmp TO ``; +GRANT ALL PRIVILEGES ON VOLUME .dqx_studio.wheels TO ``; -- End users need USE CATALOG to create temporary views for dry runs GRANT USE CATALOG ON CATALOG TO `account users`; ``` +> **Lakebase grants are handled differently.** When Lakebase is enabled, the bundle binds the database to the app via a `database` resource block (`permission: CAN_CONNECT_AND_CREATE`). DABs translates that into the equivalent Postgres role grants automatically — there is no separate SQL to run. The first time the app connects, `PgMigrationRunner` creates its own schema and tables inside the Lakebase database. + To grant app access to end users, go to **Apps → `` → Permissions** and assign `Can Use`. Replace `` with the value of `app_name` configured for your target (default `dqx-studio`). Access the app at: @@ -147,6 +212,21 @@ Access the app at: https:///apps/ ``` +## Lakebase backend + +DQX Studio stores its **OLTP state** — rules catalog, app settings, RBAC, comments, schedule configs, and scheduler bookkeeping — in a Lakebase Postgres instance for sub-millisecond reads and to avoid SQL warehouse cold starts. **Append-mostly observability tables** (`dq_validation_runs`, `dq_profiling_results`, `dq_metrics`, `dq_quarantine_records`) live in Delta because they're written by the Spark task runner and queried by AI/BI dashboards. + +| Backend | Tables | Why | +|---|---|---| +| Delta Lake | `dq_validation_runs`, `dq_profiling_results`, `dq_quarantine_records`, `dq_metrics` | High-volume append; Spark task runner writes them; columnar reads. | +| Lakebase Postgres | `dq_app_settings`, `dq_role_mappings`, `dq_quality_rules`, `dq_quality_rules_history`, `dq_comments`, `dq_schedule_configs`, `dq_schedule_configs_history`, `dq_schedule_runs` | OLTP — sub-ms reads from FastAPI handlers, row-level upserts, primary keys. | + +The Lakebase instance is declared as a bundle resource (`database_instances.lakebase`) and provisioned by `databricks bundle deploy` with `lifecycle.prevent_destroy: true` — see [Step 3](#step-3-stateful-storage-and-destroy-protection). The app connects to the always-present `databricks_postgres` admin database on the instance and creates its own `dqx_studio` Postgres schema there on first start; nothing else needs to be provisioned. + +### Lakebase token rotation + +Lakebase OAuth tokens expire after one hour. The app's `PgExecutor` runs a background daemon thread that refreshes the password every `DQX_LAKEBASE_TOKEN_REFRESH_MINUTES` minutes (default 50). Existing connections age out via `psycopg_pool.ConnectionPool.max_lifetime` so a long-running app can stay up indefinitely without reconnecting. + ## (Optional) Expand OAuth Scopes > **Most deployments don't need this step.** The OAuth scopes configured automatically by DABs (`sql`, `catalog.catalogs:read`, `catalog.schemas:read`, `catalog.tables:read`, `serving.serving-endpoints`) plus the identity scopes Databricks Apps grants implicitly are sufficient for all DQX Studio features on a standard workspace. @@ -246,3 +326,37 @@ The task-runner job installs wheels from the UC volume. If the volume is empty t databricks apps logs -p # Look for: "Uploaded databricks_labs_dqx--py3-none-any.whl" ``` + +**App says `"schema dqx_studio does not exist"` (or similar) on first start:** +The schemas didn't deploy, or the bundle is pointing at a different catalog than the app. Confirm with `databricks bundle validate -p -t ` that `catalog_name` and `schema_name` resolve to the expected values, then redeploy: +```bash +make app-deploy PROFILE= TARGET= +``` + +**App logs `"Lakebase initialisation failed ... Refusing to start"` and the container restart-loops:** +The app deliberately refuses to start when Lakebase is configured (`DQX_LAKEBASE_INSTANCE_NAME` non-empty) and init fails — silently falling back to Delta would split OLTP writes across two physical stores and orphan prior Lakebase data. Diagnose with the steps below; the Apps platform will pick up the next successful start automatically. + +1. Confirm the Lakebase instance exists and is `AVAILABLE`: + ```bash + databricks database list-database-instances -p + ``` + If the instance is missing, re-run `databricks bundle deploy`. If the instance is there but the state is `STARTING` / `UPDATING`, wait for it to reach `AVAILABLE` and the next restart will succeed. +2. Confirm the app SP has `CAN_CONNECT_AND_CREATE` on the bound logical database. Check the bundle's `database` resource block under `resources.apps.dqx-studio.resources` (it should bind `database_name: ${var.lakebase_database_name}`) and redeploy. +3. Confirm OAuth token issuance is healthy — Lakebase tokens currently expire after one hour; a misconfigured OAuth integration or revoked SP credential will surface here. +4. If you intentionally want to run on Delta only (no Lakebase), redeploy with `DQX_LAKEBASE_INSTANCE_NAME=""` (or remove the override from your target in `databricks.yml`). The app will start in legacy UC-only mode and OLTP tables will live on Delta. + +**`databricks bundle deploy` fails with `"already exists"` / `"Instance name is not unique"` on the first deploy of a target:** +The schemas, volume, or Lakebase instance were created out-of-band before this version of the bundle (e.g. by the older bootstrap script). Run the bind step once per target to adopt them into bundle management: +```bash +make app-bind PROFILE= TARGET= +make app-deploy PROFILE= TARGET= +``` +See [Migrating an existing workspace](#migrating-an-existing-workspace). + +If the conflict is specifically `"Instance name is not unique"` for the Lakebase instance and the instance does NOT appear in `databricks database list-database-instances`, it's likely in the ~7-day soft-delete retention window (the name stays reserved). Edit your target in `databricks.yml` and override `lakebase_instance_name: `, then deploy. + +**`databricks bundle destroy` fails with `"cannot destroy resource: prevent_destroy is set"`:** +This is the safety guard doing its job — see [Step 3](#step-3-stateful-storage-and-destroy-protection). To intentionally tear down a stateful resource, remove `lifecycle.prevent_destroy: true` from the relevant block in `databricks.yml`, run `databricks bundle deployment unbind -t ` to detach it from bundle state, then destroy it manually with `databricks schemas delete` / `databricks volumes delete` / `databricks database delete-database-instance`. + +**Lakebase queries time out / app logs show pool exhaustion:** +Bump `lakebase_capacity` from `CU_1` to `CU_2` (or higher) in `databricks.yml` and redeploy. You can also raise `DQX_LAKEBASE_POOL_MAX_SIZE` (default 10) on the app's environment if many concurrent requests are hitting the OLTP path. diff --git a/app/DEVELOPMENT.md b/app/DEVELOPMENT.md index 46ad7531f..88003503b 100644 --- a/app/DEVELOPMENT.md +++ b/app/DEVELOPMENT.md @@ -40,13 +40,29 @@ Create a file at `app/.env` (git-ignored) with the variables below, filling in y DATABRICKS_CONFIG_PROFILE= # matches a profile in ~/.databrickscfg DATABRICKS_WAREHOUSE_ID= DQX_CATALOG=dqx # Unity Catalog catalog name -DQX_SCHEMA=dqx_app # schema inside the catalog +DQX_SCHEMA=dqx_studio # schema inside the catalog DQX_JOB_ID= # required for profiler/dry-run -DQX_WHEELS_VOLUME=/Volumes/dqx/dqx_app/wheels # UC volume path; auto-set by DABs in production +DQX_WHEELS_VOLUME=/Volumes/dqx/dqx_studio/wheels # UC volume path; auto-set by DABs in production DQX_ADMIN_GROUP=admins # workspace group granted bootstrap Admin access + +# Lakebase (optional — leave DQX_LAKEBASE_INSTANCE_NAME empty to run OLTP tables on Delta locally) +DQX_LAKEBASE_INSTANCE_NAME= # e.g. dqx-studio-lakebase; empty = Delta-only mode +DQX_LAKEBASE_DATABASE_NAME=databricks_postgres # logical Postgres DB; defaults to the always-present admin DB +DQX_LAKEBASE_SCHEMA=public # Postgres schema (default: public) +DQX_LAKEBASE_POOL_MIN_SIZE=1 # psycopg connection pool floor +DQX_LAKEBASE_POOL_MAX_SIZE=10 # psycopg connection pool ceiling +DQX_LAKEBASE_TOKEN_REFRESH_MINUTES=50 # OAuth token refresh cadence (token expires at 60) ``` -`DQX_JOB_ID` and `DQX_WHEELS_VOLUME` are injected automatically when deployed via DABs. For local dev, set them manually if you want profiler and dry-run to work. +`DQX_JOB_ID`, `DQX_WHEELS_VOLUME`, `DQX_LAKEBASE_INSTANCE_NAME`, and `DQX_LAKEBASE_DATABASE_NAME` are injected automatically when deployed via DABs. For local dev, set them manually only if you want to exercise the corresponding feature locally: + +| Want to test... | Set... | +|---|---| +| Profiler / dry-run | `DQX_JOB_ID` (and the wheel volume must exist) | +| Lakebase OLTP path | `DQX_LAKEBASE_INSTANCE_NAME` (empty = falls back to Delta — fine for most local dev) | +| Wheel sync | `DQX_WHEELS_VOLUME` | + +> **Lakebase locally:** The same OAuth token-refresh logic that runs in production also runs locally. The app authenticates as your CLI user (via `databricks-sdk` default auth chain), so your CLI principal must have `CAN_CONNECT_AND_CREATE` on the Lakebase database. Easiest path: run the bundle once against your dev workspace so the bundle's `database` resource binds the permissions, then point `DQX_LAKEBASE_INSTANCE_NAME` at the deployed instance. ## 2. Install Dependencies @@ -139,7 +155,7 @@ The profiler creates a temporary view using your OBO token and submits a Databri If the wheel upload fails locally with a `403`, grant your user write access: ```bash -databricks volumes grant .dqx_app.wheels WRITE_VOLUME --user -p +databricks volumes grant .dqx_studio.wheels WRITE_VOLUME --user -p ``` ## Troubleshooting diff --git a/app/README.md b/app/README.md index 2e19d7270..af2720dac 100644 --- a/app/README.md +++ b/app/README.md @@ -29,9 +29,9 @@ Operations that must respect the logged-in user's permissions use the `X-Forward Operations the app owns and manages run as the app's own service principal: - **Job submission** for profiler and dry-run tasks -- **Rules catalog CRUD** (reading and writing the rules Delta table) -- **Schema migrations** (creating and evolving Delta tables) -- **App settings** (reading and writing settings from the Delta table) +- **Rules catalog CRUD** (reading and writing rules — Lakebase Postgres by default, Delta in fallback mode) +- **Schema migrations** (creating and evolving both Delta and Lakebase tables) +- **App settings** (reading and writing settings — Lakebase Postgres by default) - **Wheel upload** — on startup the app uploads DQX wheels to the UC volume and patches the task-runner job environment This ensures: @@ -73,32 +73,45 @@ On every cold start the FastAPI lifespan (`backend/app.py`) hashes the locally b - **`/api/v1/*`** — FastAPI handles all API requests - **`/*`** — FastAPI serves the compiled React SPA; TanStack Router handles client-side navigation -### Internal Storage +### Internal Storage (Hybrid Backend) -The app uses a dedicated catalog (selected at install time) with two schemas plus a wheels volume: +The app uses a **hybrid storage architecture**: high-volume append/analytical tables stay on Delta in Unity Catalog, while OLTP tables (rules catalog, app settings, RBAC, comments, schedule configs) live in **Lakebase Postgres** for sub-millisecond reads (see [DEPLOYMENT.md → Lakebase backend](DEPLOYMENT.md#lakebase-backend)). + +The schemas, wheels volume, and Lakebase instance are declared as bundle resources in `databricks.yml` with `lifecycle.prevent_destroy: true`. The bundle creates them on first deploy; `databricks bundle destroy` is blocked from dropping them. The app's `dqx_studio` Postgres schema (inside the `databricks_postgres` admin database on the Lakebase instance) is created at startup and is not itself a bundle resource, but is protected transitively by the instance-level guard. For workspaces where these resources were created out-of-band, run `make app-bind` once per target to adopt them into bundle management before the first deploy (see [DEPLOYMENT.md → Migrating an existing workspace](DEPLOYMENT.md#migrating-an-existing-workspace)). ``` -{catalog} - ├── dqx_app ← main schema (SP-managed via MigrationRunner) - │ ├── dq_app_settings ← key/value app configuration - │ ├── dq_quality_rules ← active/approved rules - │ ├── dq_quality_rules_history ← rule change audit log - │ ├── dq_role_mappings ← role → workspace group mappings (RBAC) - │ ├── dq_comments ← comment threads on rules/runs - │ ├── dq_profiling_results ← profiler runs (suggestions in generated_rules_json) - │ ├── dq_validation_runs ← dryrun + scheduled run lifecycle (1 row/run) - │ ├── dq_quarantine_records ← invalid rows captured by runs - │ ├── dq_metrics ← long-format observability events (N rows/run, - │ │ matches DQX OBSERVATION_TABLE_SCHEMA so - │ │ AI/BI dashboards target the spec directly) - │ ├── dq_schedule_configs ← per-schedule config (cron/interval, target rules) - │ ├── dq_schedule_configs_history ← schedule change audit log - │ ├── dq_schedule_runs ← scheduler last/next run state - │ └── dq_migrations ← migration version tracker - ├── dqx_app_tmp ← temp views created via OBO for profiler/dryrun jobs +{catalog} (Unity Catalog) + ├── dqx_studio ← main schema (provisioned out-of-band; tables managed by MigrationRunner) + │ ├── dq_profiling_results (Delta, always) profiler runs (suggestions in generated_rules_json) + │ ├── dq_validation_runs (Delta, always) dryrun + scheduled run lifecycle (1 row/run) + │ ├── dq_quarantine_records (Delta, always) invalid rows captured by runs + │ ├── dq_metrics (Delta, always) long-format observability events + │ │ (matches DQX OBSERVATION_TABLE_SCHEMA so AI/BI + │ │ dashboards target the spec directly) + │ ├── dq_app_settings (OLTP*) key/value app configuration + │ ├── dq_quality_rules (OLTP*) active/approved rules + │ ├── dq_quality_rules_history (OLTP*) rule change audit log + │ ├── dq_role_mappings (OLTP*) role → workspace group mappings (RBAC) + │ ├── dq_comments (OLTP*) comment threads on rules/runs + │ ├── dq_schedule_configs (OLTP*) per-schedule config (cron/interval, target rules) + │ ├── dq_schedule_configs_history (OLTP*) schedule change audit log + │ ├── dq_schedule_runs (OLTP*) scheduler last/next run state + │ └── dq_migrations ← Delta migration version tracker + ├── dqx_studio_tmp ← temp views created via OBO for profiler/dryrun jobs └── wheels (UC volume) ← DQX + task-runner wheels uploaded at app startup + +Lakebase (Postgres) — when enabled (default): + dqx-studio-lakebase (database_instance) + └── dqx_studio (database) + └── public (schema) ← provisioned by PgMigrationRunner on first start + ├── dq_app_settings, dq_role_mappings, dq_quality_rules, + ├── dq_quality_rules_history, dq_comments, dq_schedule_configs, + ├── dq_schedule_configs_history, dq_schedule_runs + └── dq_migrations ← Lakebase migration version tracker ``` +`(OLTP*)` = lives in **Lakebase Postgres** when `lakebase_instance_name` is set in `databricks.yml`, otherwise in **Delta** (the `v2: Delta OLTP fallback` migration). The split is invisible to service code: `SqlExecutor` (Delta) and `PgExecutor` (Lakebase) share an identical public surface — `execute`, `query`, `query_dicts`, `upsert`, plus the dialect helpers `q(identifier)`, `json_literal_expr(json_str)`, and `ts_text(col)` that emit dialect-correct SQL fragments. + ### Role-Based Access Control Roles (`ADMIN`, `RULE_APPROVER`, `RULE_AUTHOR`, `VIEWER`, plus the orthogonal `RUNNER`) are defined in `backend/common/authorization.py` and resolved from Databricks workspace-group membership in `dq_role_mappings` (plus the bootstrap `DQX_ADMIN_GROUP`). Routes enforce roles via `require_role(*roles)` from `backend/dependencies.py`. diff --git a/app/databricks.yml b/app/databricks.yml index 6b36c387e..c17447c18 100644 --- a/app/databricks.yml +++ b/app/databricks.yml @@ -8,10 +8,10 @@ variables: default: "dqx" schema_name: description: "Name of the schema for DQX Studio storage" - default: "dqx_app" + default: "dqx_studio" tmp_schema_name: description: "Name of the schema for temporary views and intermediate processing" - default: "dqx_app_tmp" + default: "dqx_studio_tmp" admin_group: description: "Databricks workspace group name for bootstrap Admin access" default: "admins" @@ -27,6 +27,47 @@ variables: dqx_service_principal_application_id: description: "Application ID of the service principal for DQX tasks" default: "00000000-0000-0000-0000-000000000000" + # ------------------------------------------------------------------ + # Lakebase backend + # ------------------------------------------------------------------ + # The app stores its OLTP state (rules catalog, app settings, RBAC, + # comments, schedule configs, scheduler bookkeeping) in a Lakebase + # Postgres instance for sub-millisecond reads and to avoid SQL + # warehouse cold starts. Append-mostly observability tables + # (dq_validation_runs, dq_profiling_results, dq_metrics, + # dq_quarantine_records) stay in Delta — those are written by the + # Spark task runner and consumed by AI/BI dashboards. + # + # The Lakebase instance is declared as a bundle resource below with + # ``lifecycle.prevent_destroy: true``. New Lakebase instances are + # autoscaling by default (see + # https://docs.databricks.com/aws/en/oltp/upgrade-to-autoscaling). + # + # Connection target: ``databricks_postgres`` — the always-present + # admin database every Lakebase instance ships with. We don't create + # a per-app logical Postgres database because the only DAB resource + # that creates one (``database_catalogs``) also creates a Unity + # Catalog catalog and therefore requires ``CREATE CATALOG`` on the + # metastore — a permission most app deployers don't hold. Using + # ``databricks_postgres`` keeps the bundle fully declarative with no + # external bootstrap steps, and we still get per-app isolation via a + # dedicated Postgres schema (``DQX_LAKEBASE_SCHEMA_NAME`` in the app + # config, defaults to ``dqx_studio``). + lakebase_instance_name: + description: "Name of the Lakebase Postgres instance for OLTP state." + default: "dqx-studio-lakebase" + lakebase_database_name: + description: > + Logical Postgres database the app connects to inside the Lakebase + instance. Defaults to ``databricks_postgres`` (always present, no + creation required). All DQX tables live inside a dedicated + ``dqx_studio`` schema within this database for isolation, so + multiple apps can safely share ``databricks_postgres`` on the + same Lakebase instance. + default: "databricks_postgres" + lakebase_capacity: + description: "Lakebase compute capacity (CU_1, CU_2, CU_4, CU_8)" + default: "CU_1" app_config: type: complex description: "Configuration for DQX Studio" @@ -55,6 +96,16 @@ variables: value: "${var.admin_group}" - name: "DQX_WHEELS_VOLUME" value: "/Volumes/${var.catalog_name}/${var.schema_name}/${var.wheels_volume_name}" + # Lakebase wiring. The instance is declared as a bundle + # resource further down (with ``lifecycle.prevent_destroy: true`` + # so ``bundle destroy`` cannot drop it). The logical database + # defaults to the always-present ``databricks_postgres``; the + # app creates its ``dqx_studio`` Postgres schema inside it on + # first connection. + - name: "DQX_LAKEBASE_INSTANCE_NAME" + value: "${var.lakebase_instance_name}" + - name: "DQX_LAKEBASE_DATABASE_NAME" + value: "${var.lakebase_database_name}" sync: include: @@ -62,8 +113,20 @@ sync: artifacts: default: + # ``apx build`` injects a build timestamp into the application + # wheel's *filename* but leaves the wheel's internal METADATA + # untouched. pip reads the version from METADATA, so successive + # rebuilds (without a new git commit) all advertise the same + # version and pip silently skips reinstall in the Databricks App + # container's persistent venv — the running app keeps executing + # stale code. ``_align_wheel_version.py`` rewrites METADATA, + # RECORD and the dist-info dir to match the timestamped filename + # so every rebuild produces a unique package version and pip + # always installs fresh. build: > uv run apx build && + uv run python scripts/_align_wheel_version.py + $(ls -1 .build/databricks_labs_dqx_app-*.whl | head -1) && uv build ../ --wheel --out-dir .build/ && uv build tasks/ --wheel --out-dir .build/tasks/ && dqx_whl=$(ls -1 .build/databricks_labs_dqx-*.whl | head -1 | xargs basename) && @@ -108,6 +171,21 @@ resources: job: id: ${resources.jobs.dqx_task_runner.id} permission: "CAN_MANAGE" + # Lakebase Postgres for OLTP state (rules catalog, app + # settings, RBAC, comments, schedule configs, scheduler + # bookkeeping). The instance is declared as a bundle resource + # below; the database defaults to the always-present + # ``databricks_postgres`` admin DB. This entry binds the app + # to the (instance, database) pair so DABs configures + # ``CAN_CONNECT_AND_CREATE`` automatically — that lets the app + # SP create its ``dqx_studio`` schema inside the bound DB on + # first connection. + - name: "dqx-lakebase" + description: "Lakebase Postgres backend for app OLTP state" + database: + database_name: ${var.lakebase_database_name} + instance_name: ${resources.database_instances.lakebase.name} + permission: "CAN_CONNECT_AND_CREATE" permissions: - group_name: "account users" @@ -125,50 +203,51 @@ resources: - group_name: "users" level: "CAN_USE" - volumes: - dqx_wheels: - catalog_name: ${var.catalog_name} - schema_name: ${var.schema_name} - name: ${var.wheels_volume_name} - volume_type: MANAGED - grants: - - principal: ${resources.apps.dqx-studio.id} - privileges: - - ALL_PRIVILEGES - - principal: ${var.dqx_service_principal_application_id} - privileges: - - ALL_PRIVILEGES - + # ------------------------------------------------------------------ + # Stateful storage (schemas, volume, Lakebase instance + database). + # + # Every entry below is tagged ``lifecycle.prevent_destroy: true`` so + # that ``databricks bundle destroy`` (or a deploy that would + # otherwise replace the resource) fails fast instead of dropping the + # backing data. To intentionally tear something down: remove the + # flag, ``databricks bundle deployment unbind ``, then destroy. + # + # For an existing workspace whose resources were provisioned + # out-of-band (e.g. from the previous bootstrap script), one-time + # binding is required before the first deploy — see ``make app-bind`` + # and DEPLOYMENT.md → "Migrating an existing workspace". + # ------------------------------------------------------------------ schemas: - dqx_schema: + main_schema: catalog_name: ${var.catalog_name} name: ${var.schema_name} - grants: - - principal: ${resources.apps.dqx-studio.id} - privileges: - - ALL_PRIVILEGES - - principal: ${var.dqx_service_principal_application_id} - privileges: - - ALL_PRIVILEGES - - dqx_tmp_schema: + comment: "DQX Studio main schema (rules, run history, profiling, metrics, quarantine, OLTP fallback)" + lifecycle: + prevent_destroy: true + + tmp_schema: catalog_name: ${var.catalog_name} name: ${var.tmp_schema_name} - grants: - - principal: ${resources.apps.dqx-studio.id} - privileges: - - ALL_PRIVILEGES - - MANAGE - - principal: ${var.dqx_service_principal_application_id} - privileges: - - ALL_PRIVILEGES - - MANAGE - # account users need CREATE_TABLE so the app can create per-user - # temporary views in this schema via OBO tokens. - - principal: "account users" - privileges: - - USE_SCHEMA - - CREATE_TABLE + comment: "DQX Studio temp schema for per-user dry-run views" + lifecycle: + prevent_destroy: true + + volumes: + wheels: + catalog_name: ${var.catalog_name} + schema_name: ${resources.schemas.main_schema.name} + name: ${var.wheels_volume_name} + volume_type: MANAGED + comment: "DQX Studio wheel storage for the task runner" + lifecycle: + prevent_destroy: true + + database_instances: + lakebase: + name: ${var.lakebase_instance_name} + capacity: ${var.lakebase_capacity} + lifecycle: + prevent_destroy: true jobs: dqx_task_runner: @@ -223,18 +302,26 @@ resources: default: "unknown" targets: - dev: - workspace: - profile: - variables: - # Required - catalog_name: - dqx_service_principal_application_id: - - # Optional — uncomment and override defaults per target as needed - # admin_group: - # app_name: - # sql_warehouse_name: - # schema_name: - presets: - trigger_pause_status: PAUSED + # Template for adding a new target. Required variables are + # ``catalog_name`` and ``dqx_service_principal_application_id``; + # everything else has a sensible default in the ``variables:`` block + # at the top of this file and can be overridden per target below. + # + # example-target: + # workspace: + # profile: + # variables: + # catalog_name: + # dqx_service_principal_application_id: + # # Optional overrides: + # # admin_group: + # # app_name: + # # sql_warehouse_name: + # # schema_name: + # # tmp_schema_name: + # # wheels_volume_name: + # # lakebase_instance_name: + # # lakebase_database_name: + # # lakebase_capacity: CU_1 | CU_2 | CU_4 | CU_8 + # presets: + # trigger_pause_status: PAUSED diff --git a/app/pyproject.toml b/app/pyproject.toml index 81c37e46e..9f9718feb 100644 --- a/app/pyproject.toml +++ b/app/pyproject.toml @@ -12,6 +12,11 @@ dependencies = [ "databricks-sdk~=0.73", "databricks-sql-connector[pyarrow]==4.2.5", "openpyxl>=3.1", + # Lakebase (Postgres) backend for OLTP state. ``binary`` ships a + # pre-built libpq so we don't need a system Postgres install at + # runtime; ``pool`` adds the threaded ConnectionPool used by + # PgExecutor. Both are optional features of the same package. + "psycopg[binary,pool]>=3.2", ] [dependency-groups] @@ -76,6 +81,49 @@ exclude = ["src/databricks_labs_dqx_app/ui"] line-length = 120 target-version = "py311" +[tool.ruff.lint] +# --------------------------------------------------------------------------- +# Project-wide policy on broad-except (BLE001). +# --------------------------------------------------------------------------- +# BLE001 ("Do not catch blind exception: ``Exception``") is intentionally +# NOT in the active rule set. Several backend surfaces deliberately catch +# the broad ``Exception`` because best-effort resilience is the contract, +# not a code smell: +# +# * Lifespan teardown — :meth:`pg_executor.PgExecutor.close` and the +# ``lifespan`` cleanup branch in ``backend.app``. A close() that +# raises would leak Postgres connections on every restart flap and +# could hang the worker indefinitely; the pool is already being torn +# down so the only meaningful action is "log and move on". +# +# * Background-thread resilience — +# :meth:`pg_executor.PgExecutor._token_refresh_loop`. The refresher +# is a daemon thread that MUST NOT die: if a single SDK / network / +# OAuth failure killed it, the pool's tokens would silently expire +# and every subsequent connect would fail with confusing SCRAM +# errors. The loop logs and backs off instead. +# +# * Route-level resilience wrappers — endpoints that swallow per-item +# exceptions so a single bad row doesn't 500 the whole list endpoint +# (config label-definition migration, optional check-function +# resolution, etc.). +# +# A narrower catch (e.g. ``(psycopg.Error, OSError)``) cannot cover +# every plausible failure mode at these surfaces. The guarantee we need +# is "this MUST NOT propagate", and that's exactly what ``except +# Exception`` provides. +# +# Convention for any new broad-except site: +# 1. Add an inline comment explaining the resilience contract — why +# narrowing would re-introduce a real failure mode. If you can't +# write the comment, you probably shouldn't use the broad catch. +# 2. Do NOT add ``# noqa: BLE001``. Per AGENTS.md rule #6, per-line +# noqa is the wrong escape hatch. +# 3. If a future change makes BLE001 worth enforcing project-wide, +# enable it here and add ``[tool.ruff.lint.per-file-ignores]`` +# entries for the resilience files listed above — never per-line +# noqa. + [tool.basedpyright] # ``apx dev check`` runs basedpyright but only syncs the default + dev # dependency groups, not ``test``. That means pytest (and the rest of the diff --git a/app/scripts/_align_wheel_version.py b/app/scripts/_align_wheel_version.py new file mode 100644 index 000000000..aa7ed3375 --- /dev/null +++ b/app/scripts/_align_wheel_version.py @@ -0,0 +1,169 @@ +"""Re-align a wheel's METADATA / RECORD / dist-info dir to its filename. + +Background +---------- +``apx build`` post-processes the freshly-built application wheel by +injecting a build timestamp into the **filename** (e.g. renaming +``databricks_labs_dqx_app-0.13.0.post38.dev0+65c9602-py3-none-any.whl`` +to +``databricks_labs_dqx_app-0.13.0.post38.dev0+65c9602.post20260506102508-py3-none-any.whl``) +but does **not** update the wheel's internal ``METADATA``, ``RECORD`` or +``*.dist-info`` directory name. + +The result is a wheel whose filename advertises a new version each +build but whose package metadata still reports the unchanged +git-derived version. ``pip`` consults the metadata, sees the version +is already installed, and silently skips reinstall — so a Databricks +App container's persistent venv keeps running stale code on every +redeploy. + +This script repairs the wheel produced by ``apx build`` by: + +1. parsing the version segment out of the filename +2. unpacking the wheel +3. renaming the ``*.dist-info`` directory to use that version +4. rewriting ``METADATA`` so ``Version:`` matches +5. regenerating ``RECORD`` checksums +6. repacking the wheel in place + +After running this each rebuild produces a wheel with a unique +metadata version, forcing pip to install fresh code on every deploy. + +Usage +----- + python _align_wheel_version.py path/to/wheel.whl + +The script overwrites the input wheel; safe to run idempotently. +""" + +from __future__ import annotations + +import base64 +import hashlib +import re +import shutil +import sys +import tempfile +import zipfile +from pathlib import Path + +# Wheel filename format (PEP 427): +# {distribution}-{version}(-{build tag})?-{python tag}-{abi tag}-{platform tag}.whl +# Distribution and version may both contain dots; the version segment +# starts right after the leading "{distribution}-" and ends at the +# next "-py" / "-cp" / "-pp" tag. Everything after that until the +# trailing ``.whl`` is the compatibility tag triple. +_WHEEL_FILENAME_RE = re.compile( + r"^(?P[A-Za-z0-9_]+)-(?P[^-]+(?:\+[^-]+)?)-(?P(?:py|cp|pp|ip)\d.*?)\.whl$" +) + + +def _filename_version(wheel: Path) -> tuple[str, str]: + """Return ``(distribution, version)`` parsed from the wheel filename. + + Uses a regex rather than ``packaging.utils.parse_wheel_filename`` + because the latter normalises the distribution name and rejects + local versions that contain dots — both of which are present in the + wheels apx produces. + """ + m = _WHEEL_FILENAME_RE.match(wheel.name) + if not m: + raise SystemExit(f"unrecognised wheel filename: {wheel.name}") + return m.group("dist"), m.group("version") + + +def _record_entry(rel_path: str, data: bytes) -> str: + digest = base64.urlsafe_b64encode(hashlib.sha256(data).digest()).rstrip(b"=").decode() + return f"{rel_path},sha256={digest},{len(data)}" + + +def align(wheel_path: Path) -> None: + """Rewrite ``wheel_path`` so its metadata version matches its filename.""" + dist, filename_version = _filename_version(wheel_path) + + with tempfile.TemporaryDirectory() as raw_tmp: + tmp = Path(raw_tmp) + with zipfile.ZipFile(wheel_path, "r") as zf: + zf.extractall(tmp) + + old_dist_infos = [p for p in tmp.iterdir() if p.is_dir() and p.name.endswith(".dist-info")] + if not old_dist_infos: + raise SystemExit(f"no .dist-info dir in {wheel_path.name}") + if len(old_dist_infos) > 1: + raise SystemExit(f"multiple .dist-info dirs in {wheel_path.name}: {old_dist_infos}") + old_dist_info = old_dist_infos[0] + + # Hatchling preserves underscores in the distribution prefix + # of the dist-info directory, so reuse that prefix verbatim + # rather than recomputing from ``dist`` (which packaging.utils + # would normalise to dashes). The directory name is + # ``{prefix}-{version}.dist-info``; strip the literal suffix + # then split off the version. + stem = old_dist_info.name[: -len(".dist-info")] + existing_prefix = stem.rsplit("-", 1)[0] + new_dist_info = tmp / f"{existing_prefix}-{filename_version}.dist-info" + # Skip the rename when the dist-info dir is already correctly + # named (e.g. apx didn't inject a build tag on this rebuild so + # the filename version matches what's already inside the + # wheel). Without this guard, the ``rmtree`` below would delete + # the source directory — both paths point at the same dir — and + # the subsequent ``rename`` would crash with FileNotFoundError. + # The METADATA/RECORD rewrites further down are themselves + # idempotent, so we still run them to repair any drift. + if new_dist_info != old_dist_info: + if new_dist_info.exists(): + shutil.rmtree(new_dist_info) + old_dist_info.rename(new_dist_info) + + metadata = new_dist_info / "METADATA" + if not metadata.exists(): + raise SystemExit(f"missing METADATA in {wheel_path.name}") + text = metadata.read_text(encoding="utf-8") + new_text, n = re.subn( + r"^Version: .*$", + f"Version: {filename_version}", + text, + count=1, + flags=re.MULTILINE, + ) + if n == 0: + raise SystemExit(f"no Version: line in METADATA of {wheel_path.name}") + metadata.write_text(new_text, encoding="utf-8") + + record = new_dist_info / "RECORD" + record_rel = record.relative_to(tmp).as_posix() + entries: list[str] = [] + for f in sorted(tmp.rglob("*")): + if not f.is_file(): + continue + rel = f.relative_to(tmp).as_posix() + if rel == record_rel: + continue + entries.append(_record_entry(rel, f.read_bytes())) + entries.append(f"{record_rel},,") + record.write_text("\n".join(entries) + "\n", encoding="utf-8") + + tmp_wheel = wheel_path.with_suffix(wheel_path.suffix + ".tmp") + with zipfile.ZipFile(tmp_wheel, "w", zipfile.ZIP_DEFLATED) as zf: + for f in sorted(tmp.rglob("*")): + if f.is_file(): + zf.write(f, f.relative_to(tmp).as_posix()) + + tmp_wheel.replace(wheel_path) + + +def main(argv: list[str]) -> int: + if len(argv) != 2: + print(f"usage: {Path(argv[0]).name} ", file=sys.stderr) + return 2 + wheel = Path(argv[1]) + if not wheel.is_file(): + print(f"not a file: {wheel}", file=sys.stderr) + return 1 + align(wheel) + print(f" aligned wheel metadata to filename version: {wheel.name}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv)) diff --git a/app/scripts/bind_resources.sh b/app/scripts/bind_resources.sh new file mode 100755 index 000000000..a105ef427 --- /dev/null +++ b/app/scripts/bind_resources.sh @@ -0,0 +1,150 @@ +#!/usr/bin/env bash +# +# One-time binding of pre-existing storage resources to the bundle. +# +# Use this when the schemas / volume / Lakebase instance already +# exist in the target workspace (e.g. from the previous bootstrap- +# script flow, or from manual creation) and you're adopting them into +# the bundle for the first time. Without binding, ``databricks bundle +# deploy`` tries to CREATE the resources and fails with "already +# exists". +# +# The app connects to the always-present ``databricks_postgres`` admin +# database on the Lakebase instance (no separate logical database to +# provision or bind), and creates its ``dqx_studio`` Postgres schema +# there on first start. +# +# Bind is idempotent at the CLI level; re-running this script on a +# fully-bound workspace is a no-op (the CLI replies "already bound"). +# +# Skip this script for fresh workspaces — ``databricks bundle deploy`` +# creates the resources directly. +# +# Usage: +# ./scripts/bind_resources.sh -p -t +# +# Requirements: +# - databricks CLI v0.268+ (lifecycle.prevent_destroy support) +# - jq installed + +set -euo pipefail + +PROFILE="" +TARGET="" + +usage() { + echo "Usage: $0 -p -t " + exit 1 +} + +while getopts "p:t:" opt; do + case $opt in + p) PROFILE="$OPTARG" ;; + t) TARGET="$OPTARG" ;; + *) usage ;; + esac +done +shift $((OPTIND - 1)) + +[[ -z "$PROFILE" || -z "$TARGET" ]] && usage + +# Everything after a ``--`` separator is forwarded to every bundle +# subcommand as extra ``--var key=value`` overrides. Threading them +# into ``bundle validate`` matters: that call is what produces the +# JSON we parse below to learn which instance/schema name to bind to. +# Without forwarding, a deploy-time override would bind the wrong +# resource and the next deploy would still see a state-vs-config drift. +EXTRA_VARS=("$@") + +CLI="databricks -p $PROFILE" +BUNDLE_FLAGS=(-t "$TARGET") + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +BUNDLE_DIR="$(dirname "$SCRIPT_DIR")" +cd "$BUNDLE_DIR" + +# --------------------------------------------------------------------------- +# Read resource identifiers from the bundle config. Identifiers must +# match the ``name`` / ``catalog_name`` / etc. fields rendered for the +# selected target, so we ask the CLI for the resolved bundle. +# --------------------------------------------------------------------------- +BUNDLE_VALIDATE_STDERR=$(mktemp) +trap 'rm -f "$BUNDLE_VALIDATE_STDERR"' EXIT +if ! BUNDLE_JSON=$($CLI bundle validate "${BUNDLE_FLAGS[@]}" "${EXTRA_VARS[@]}" -o json 2>"$BUNDLE_VALIDATE_STDERR"); then + echo "ERROR: 'databricks bundle validate' failed:" >&2 + cat "$BUNDLE_VALIDATE_STDERR" >&2 + exit 1 +fi + +CATALOG=$(echo "$BUNDLE_JSON" | jq -r '.variables.catalog_name.value // .variables.catalog_name.default // empty') +SCHEMA=$(echo "$BUNDLE_JSON" | jq -r '.variables.schema_name.value // .variables.schema_name.default // "dqx_studio"') +TMP_SCHEMA=$(echo "$BUNDLE_JSON" | jq -r '.variables.tmp_schema_name.value // .variables.tmp_schema_name.default // "dqx_studio_tmp"') +VOLUME=$(echo "$BUNDLE_JSON" | jq -r '.variables.wheels_volume_name.value // .variables.wheels_volume_name.default // "wheels"') +LB_INSTANCE=$(echo "$BUNDLE_JSON" | jq -r '.variables.lakebase_instance_name.value // .variables.lakebase_instance_name.default // empty') + +if [[ -z "$CATALOG" ]]; then + echo "ERROR: catalog_name is not configured in the bundle target." >&2 + exit 1 +fi + +echo "==> Binding pre-existing resources (target=$TARGET, profile=$PROFILE)" +echo " Catalog: $CATALOG" +echo " Main schema: $SCHEMA" +echo " Tmp schema: $TMP_SCHEMA" +echo " Volume: $VOLUME" +echo " Lakebase: $LB_INSTANCE" +echo "" + +# ``databricks bundle deployment bind`` requires explicit ``--auto-approve`` +# when stdin is not a TTY (newer CLI versions stopped accepting piped +# ``yes`` and now error out with "current console does not support +# prompting"). The confirmation it would otherwise prompt for is just +# "apply bundle resource updates to the existing remote resource on the +# next deploy" — which is exactly what we want, so auto-approving is +# safe. +# +# Bind is idempotent on the CLI side: re-binding an already-bound +# resource exits 0. Any non-zero exit from ``bundle deployment bind`` +# therefore signals a genuine problem (the resource doesn't exist, the +# principal lacks permission, the target points at the wrong workspace, +# …) rather than a benign "already bound" state. We surface it loudly +# instead of warning-and-continuing so the user sees the real cause +# now, not a confusing "resource already exists" failure inside the +# next ``bundle deploy``. +bind() { + local key="$1" + local id="$2" + echo " binding ${key} -> ${id}" + if ! $CLI bundle deployment bind "$key" "$id" --auto-approve "${BUNDLE_FLAGS[@]}" "${EXTRA_VARS[@]}"; then + echo "" >&2 + echo "ERROR: failed to bind ${key} -> ${id}" >&2 + echo "" >&2 + echo "Common causes:" >&2 + echo " - The remote resource does not exist yet. This script adopts" >&2 + echo " pre-existing resources; on a fresh workspace, skip it and" >&2 + echo " run 'make app-deploy' directly — the bundle creates the" >&2 + echo " resources for you." >&2 + echo " - The current principal (profile=${PROFILE}) lacks permission" >&2 + echo " to read or bind the resource." >&2 + echo " - The bundle target (${TARGET}) points at a different" >&2 + echo " workspace than the one where the resource was created." >&2 + echo "" >&2 + echo "Re-run after fixing the underlying cause, or run this single" >&2 + echo "bind manually if you're certain the resource exists and is" >&2 + echo "accessible:" >&2 + echo " databricks -p ${PROFILE} bundle deployment bind \\" >&2 + echo " ${key} ${id} --auto-approve -t ${TARGET}" >&2 + exit 1 + fi +} + +bind main_schema "${CATALOG}.${SCHEMA}" +bind tmp_schema "${CATALOG}.${TMP_SCHEMA}" +bind wheels "${CATALOG}.${SCHEMA}.${VOLUME}" + +if [[ -n "$LB_INSTANCE" ]]; then + bind lakebase "$LB_INSTANCE" +fi + +echo "" +echo "==> Bind complete. Run 'make app-deploy PROFILE=$PROFILE TARGET=$TARGET' next." diff --git a/app/scripts/post_deploy_grants.sh b/app/scripts/post_deploy_grants.sh index a09d8b23f..62c43b653 100755 --- a/app/scripts/post_deploy_grants.sh +++ b/app/scripts/post_deploy_grants.sh @@ -7,13 +7,20 @@ # GRANT statements via the Statement Execution API. # # Usage: -# ./scripts/post_deploy_grants.sh -p [-t ] +# ./scripts/post_deploy_grants.sh -p [-t ] [-- ] # # The bundle target is required when the bundle defines more than one # target and none is marked as default. Without it, ``databricks bundle # validate`` errors out and we have no way to discover the catalog name # or job SP from variables. # +# Everything after a ``--`` separator is forwarded to ``bundle validate`` +# as extra ``--var key=value`` overrides. Pass the SAME overrides used +# at ``bundle deploy`` time — otherwise the script reads the default +# catalog/schema/volume names from the bundle and issues GRANTs on the +# wrong objects (the overridden resources receive no permissions and +# the app SP cannot read them). +# # Requirements: # - databricks CLI authenticated # - jq installed @@ -25,7 +32,7 @@ PROFILE="" TARGET="" usage() { - echo "Usage: $0 -p [-t ]" + echo "Usage: $0 -p [-t ] [-- ]" exit 1 } @@ -36,9 +43,18 @@ while getopts "p:t:" opt; do *) usage ;; esac done +shift $((OPTIND - 1)) [[ -z "$PROFILE" ]] && usage +# Forwarded ``--var key=value`` overrides. Threading them into +# ``bundle validate`` matters because that call produces the JSON we +# parse below for catalog / schema / volume / job-SP names. Without +# forwarding, a deploy-time override (e.g. ``--var=catalog_name=foo``) +# would issue GRANTs on the bundle's default catalog instead of the +# one actually deployed. +EXTRA_VARS=("$@") + CLI="databricks -p $PROFILE" # Bundle commands need the target unless one is marked default. BUNDLE_FLAGS=() @@ -54,7 +70,10 @@ cd "$BUNDLE_DIR" # instead of an empty $BUNDLE_JSON and a confusing downstream error. BUNDLE_VALIDATE_STDERR=$(mktemp) trap 'rm -f "$BUNDLE_VALIDATE_STDERR"' EXIT -if ! BUNDLE_JSON=$($CLI bundle validate "${BUNDLE_FLAGS[@]}" -o json 2>"$BUNDLE_VALIDATE_STDERR"); then +# macOS bash 3.2 + ``set -u`` treats expansion of an empty array as +# unbound. Guard with the ``${arr[@]+...}`` idiom so a run with no +# extra ``--var`` overrides (or no ``-t``) does not abort here. +if ! BUNDLE_JSON=$($CLI bundle validate ${BUNDLE_FLAGS[@]+"${BUNDLE_FLAGS[@]}"} ${EXTRA_VARS[@]+"${EXTRA_VARS[@]}"} -o json 2>"$BUNDLE_VALIDATE_STDERR"); then echo "ERROR: 'databricks bundle validate' failed:" >&2 cat "$BUNDLE_VALIDATE_STDERR" >&2 if [[ -z "$TARGET" ]]; then @@ -91,8 +110,8 @@ fi echo " Warehouse: $WH_ID" CATALOG=$(echo "$BUNDLE_JSON" | jq -r '.variables.catalog_name.value // .variables.catalog_name.default // empty') -SCHEMA=$(echo "$BUNDLE_JSON" | jq -r '.variables.schema_name.value // .variables.schema_name.default // "dqx_app"') -TMP_SCHEMA=$(echo "$BUNDLE_JSON" | jq -r '.variables.tmp_schema_name.value // .variables.tmp_schema_name.default // "dqx_app_tmp"') +SCHEMA=$(echo "$BUNDLE_JSON" | jq -r '.variables.schema_name.value // .variables.schema_name.default // "dqx_studio"') +TMP_SCHEMA=$(echo "$BUNDLE_JSON" | jq -r '.variables.tmp_schema_name.value // .variables.tmp_schema_name.default // "dqx_studio_tmp"') VOLUME=$(echo "$BUNDLE_JSON" | jq -r '.variables.wheels_volume_name.value // .variables.wheels_volume_name.default // "wheels"') JOB_SP=$(echo "$BUNDLE_JSON" | jq -r '.variables.dqx_service_principal_application_id.value // .variables.dqx_service_principal_application_id.default // empty') diff --git a/app/src/databricks_labs_dqx_app/backend/CLAUDE.md b/app/src/databricks_labs_dqx_app/backend/CLAUDE.md index 0dd2651b4..7b6b46fe7 100644 --- a/app/src/databricks_labs_dqx_app/backend/CLAUDE.md +++ b/app/src/databricks_labs_dqx_app/backend/CLAUDE.md @@ -12,11 +12,12 @@ backend/ ├── cache.py # CacheFactory — async in-memory TTL cache + @cached decorator ├── config.py # AppConfig (Pydantic BaseSettings, DQX_ env prefix) ├── dependencies.py # FastAPI Depends() — OBO/SP auth, RBAC, services -├── migrations/ # MigrationRunner — versioned DDL applied at startup +├── migrations/ # MigrationRunner (Delta) + PgMigrationRunner (Lakebase) ├── models.py # Pydantic request/response models ├── run_status_manager.py # Helpers for reading/updating dq_validation_runs status ├── settings.py # SettingsManager — per-user prefs in ~/.dqx/app.yml ├── sql_executor.py # SqlExecutor — Databricks Statement Execution API wrapper +├── pg_executor.py # PgExecutor — Lakebase Postgres wrapper (parity API w/ SqlExecutor) ├── sql_utils.py # Shared SQL helpers: escape_sql_string, validate_fqn, quote_fqn ├── runtime.py # Runtime singleton (lazy WorkspaceClient) ├── logger.py # Custom logging formatter @@ -168,3 +169,120 @@ uv run uvicorn databricks_labs_dqx_app.backend.app:app --reload # Dev server - **Scheduler:** runs in-process as an asyncio task, gated by an exclusive file lock (`/tmp/.dqx_scheduler.lock`) so only one uvicorn worker drives it. Disable with `DQX_SCHEDULER_DISABLED=1`. - **Caches:** `app_cache` (`cache.py`) is per-process in-memory with TTL. SP `WorkspaceClient`, OBO `WorkspaceClient`, and per-user catalog list are all cached. Use the `MISS` sentinel — never `is None` — to detect cache absence. - **SPA static files:** `spa_static.py` falls through to `index.html` only for non-asset paths (positive allowlist of asset extensions), so SPA routes containing dots still work. + +## Hybrid Storage Backend (Delta + Lakebase) + +The DQX Studio data model is split across two physical backends and the +choice is driven entirely by `databricks.yml`: + +| Backend | Tables | Why | +|---------|--------|-----| +| **Delta Lake** (always) | `dq_validation_runs`, `dq_profiling_results`, `dq_quarantine_records`, `dq_metrics` | Spark task runner writes these; high-volume append-mostly; columnar reads. | +| **Lakebase Postgres** *(default — opt-out via `lakebase_instance_name=""`)* | `dq_app_settings`, `dq_role_mappings`, `dq_quality_rules`, `dq_quality_rules_history`, `dq_comments`, `dq_schedule_configs`, `dq_schedule_configs_history`, `dq_schedule_runs` | Low-latency point reads/writes from FastAPI request handlers; row-level upserts; primary-key/foreign-key semantics. | + +When Lakebase is **disabled** (no `lakebase_instance_name` set), the OLTP +tables fall back to Delta — `MigrationRunner` runs both +`v1: Delta analytical baseline` *and* `v2: Delta OLTP fallback`. When +Lakebase is **enabled**, only `v1` runs on Delta and `PgMigrationRunner` +provisions the OLTP tables in Postgres. + +### Key types + +- `SqlExecutor` (`sql_executor.py`) wraps the Databricks Statement + Execution API for Delta. +- `PgExecutor` (`pg_executor.py`) wraps `psycopg` + a `psycopg_pool.ConnectionPool` + for Lakebase. It mirrors `SqlExecutor`'s public surface: `execute`, + `query`, `query_dicts`, `upsert`, plus the dialect helpers + `q(identifier)`, `json_literal_expr(json_str)`, `ts_text(col)`. A + background daemon thread refreshes the OAuth password every + `DQX_LAKEBASE_TOKEN_REFRESH_MINUTES` minutes (default 50; tokens + expire at 60). The pool's `kwargs["password"]` is mutated in place + so subsequent connects pick up the new credential, and existing + connections age out via `max_lifetime`. +- Services keep their `sql: SqlExecutor` annotation; the dependency + injection layer (`dependencies.get_sp_oltp_executor`) hands back + whichever executor is registered, casting to `SqlExecutor` because + the two classes share an identical method surface. +- The `SchedulerService` accepts `oltp_sql: SqlExecutor | PgExecutor | None` + and routes OLTP-table SQL (schedule configs, settings, rules) to + the OLTP executor while keeping retention/GC against the Delta + executor. + +### Retention sweep (daily) + +The scheduler runs a `DELETE` pass against the analytical tables once +per `_RETENTION_INTERVAL_HOURS` (24h). Two knobs, both stored in +`dq_app_settings` and surfaced via `GET/PUT /api/v1/config/retention`: + +| Setting key | Default | Tables affected | +|------------------------------|--------:|-----------------| +| `retention_days` | 90 | `dq_validation_runs`, `dq_profiling_results`, `dq_metrics`, plus the OLTP history tables (`dq_quality_rules_history`, `dq_schedule_configs_history`). Picked to match what trend dashboards expect. | +| `quarantine_retention_days` | 30 | `dq_quarantine_records` only. Tighter because that table holds the full source row payload (PII surface). | + +Both resolvers share a `_RETENTION_DAYS_MIN = 7` floor so a +mis-typed setting can never wipe data inside the safety window. Reads +swallow exceptions and fall back to the compiled-in default so a +SQL-warehouse hiccup never crashes the scheduler tick. + +### Writing portable SQL inside services + +Always go through the executor's dialect helpers — never hard-code +backticks, `parse_json(...)`, or `CAST(... AS STRING)`: + +```python +self._sql.q("check") # `check` (Delta) | "check" (Postgres) +self._sql.json_literal_expr(j) # parse_json('...') | '...'::jsonb +self._sql.ts_text("created_at") # CAST(created_at AS STRING) | created_at +``` + +For upserts, `SqlExecutor.upsert(table, key_cols, value_cols)` and +`PgExecutor.upsert` take the same arguments. Pass +`RawSql("current_timestamp()")` for timestamps — both backends rewrite +to their native syntax. + +### Bundle / DAB conventions + +Stateful resources declared in `databricks.yml` with +`lifecycle.prevent_destroy: true` (Databricks CLI 0.268+): + +* `resources.schemas.main_schema` — `dqx_studio` schema +* `resources.schemas.tmp_schema` — `dqx_studio_tmp` schema +* `resources.volumes.wheels` — wheels volume +* `resources.database_instances.lakebase` — Lakebase Postgres instance + (autoscaling by default per [Lakebase Autoscaling](https://docs.databricks.com/aws/en/oltp/upgrade-to-autoscaling)) + +The app connects to the always-present `databricks_postgres` admin +database on the Lakebase instance — that's the default value of +`lakebase_database_name` and the value the app→database binding +wires up. On first start, the app creates its own `dqx_studio` +Postgres schema inside `databricks_postgres` and runs migrations +against it. Multiple apps can therefore share the same +`databricks_postgres` on one Lakebase instance safely; each gets its +own schema namespace. + +The bundle deliberately does NOT use `database_catalogs`. That DAB +resource is the only way to *create* a custom logical Postgres +database, but it also creates a Unity Catalog catalog as a side +effect and therefore requires `CREATE CATALOG` on the metastore — a +permission most app deployers don't hold. Connecting to the +pre-existing `databricks_postgres` instead keeps the bundle fully +declarative with no out-of-band bootstrap step and no metastore-level +permissions assumed. + +`prevent_destroy` blocks `databricks bundle destroy` and any deploy +that would force-replace a bundle-managed resource — the alternative +is silent data loss. To intentionally tear one down: remove the flag, +run `databricks bundle deployment unbind `, then destroy. The +app's `dqx_studio` Postgres schema lives below the resource layer +DABs models, so `prevent_destroy` doesn't apply to it directly; the +instance-level guard is what protects it. + +For workspaces where the schemas, volume, or Lakebase instance were +provisioned out-of-band (e.g. by the legacy bootstrap script), +one-time binding is required: `make app-bind PROFILE=... TARGET=...`. +After bind, `bundle deploy` adopts the existing resources instead of +trying to CREATE them. + +Privileges on UC objects for the auto-created app SP are still applied +by `scripts/post_deploy_grants.sh` after each deploy — the app SP's +UUID isn't known at bundle-write time. diff --git a/app/src/databricks_labs_dqx_app/backend/app.py b/app/src/databricks_labs_dqx_app/backend/app.py index 8ef44d5d7..762568344 100644 --- a/app/src/databricks_labs_dqx_app/backend/app.py +++ b/app/src/databricks_labs_dqx_app/backend/app.py @@ -10,9 +10,19 @@ from ._scheduler_registry import get_scheduler, set_scheduler from .config import conf -from .dependencies import get_sp_ws +from .dependencies import get_sp_ws, set_oltp_executor from .logger import logger from .migrations import MigrationRunner + +# ``PgMigrationRunner`` is safe to import at module load: its module +# (``migrations.postgres``) now imports the trust-boundary helpers +# from the psycopg-free :mod:`backend.pg_cursor_helpers` module +# rather than from :mod:`backend.pg_executor`, so loading it does +# NOT transitively pull in :mod:`psycopg`. The remaining lazy import +# (``build_pg_executor``) inside :func:`lifespan` is the one that +# genuinely needs psycopg at runtime, and stays lazy for the +# Delta-only test environments that don't install it. +from .migrations.postgres import PgMigrationRunner from .routes import api_router from .services.scheduler_service import SchedulerService from .services.view_service import mark_tmp_schema_ready @@ -21,6 +31,16 @@ _SCHEDULER_LOCK_PATH = Path("/tmp/.dqx_scheduler.lock") # noqa: S108 +# Module-level reference that pins the lock file descriptor for the +# process lifetime. The fd MUST stay live as long as we hold the +# advisory lock — letting it get garbage-collected would silently +# close the fd and release the flock, after which a second uvicorn +# worker could grab the lease and we'd end up with two schedulers +# fighting over the same job. Stored as a plain module global rather +# than via ``globals()[...]`` so the dependency is greppable and the +# type-checker can see it. +_scheduler_lock_fd: int | None = None + def _try_acquire_scheduler_lease() -> bool: """Use an exclusive file lock so only one uvicorn worker runs the scheduler. @@ -30,12 +50,11 @@ def _try_acquire_scheduler_lease() -> bool: """ import fcntl + global _scheduler_lock_fd try: fd = os.open(str(_SCHEDULER_LOCK_PATH), os.O_CREAT | os.O_RDWR) fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - # Keep fd open (and thus the lock held) for the process lifetime. - # Store on module so it isn't garbage-collected. - globals()["_scheduler_lock_fd"] = fd + _scheduler_lock_fd = fd return True except OSError: return False @@ -173,12 +192,113 @@ async def lifespan(app: FastAPI): sp_ws = await get_sp_ws() wh_id = os.environ.get("DATABRICKS_WAREHOUSE_ID") or os.environ.get("DATABRICKS_SQL_WAREHOUSE_ID") or "" sp_sql = SqlExecutor(ws=sp_ws, warehouse_id=wh_id, catalog=conf.catalog, schema=conf.schema_name) + + # ------------------------------------------------------------------ + # Lakebase (optional) — open the pool, run Postgres migrations, and + # register the executor as the OLTP backend used by service DI. + # + # When ``lakebase_enabled`` is true the operator has explicitly + # chosen Lakebase as the OLTP store: rules, schedules, RBAC, + # comments, and app settings live in the Postgres schema, NOT in + # Delta. Silently falling back to Delta on a transient init + # failure (network blip, OAuth issuance, momentary + # CAN_CONNECT_AND_CREATE drop) would split-brain the deployment — + # writes during the flap land in the empty Delta fallback tables + # while the canonical Postgres rows become invisible, and prior + # data reappears after the next restart with anything written in + # the interim orphaned. That's silent data loss, not graceful + # degradation, so we follow the same fail-loud-and-let-the- + # platform-restart-us pattern as the SP/migrations block above: + # raise, let the Databricks Apps platform restart the container, + # and surface the underlying problem via restart-loop alerting. + # The opt-out is intentional and explicit (unset + # ``DQX_LAKEBASE_INSTANCE_NAME``); a flap is not an opt-out. + # The legitimate Delta-only path runs in the ``else`` branch below. + # ------------------------------------------------------------------ + pg_executor = None + if conf.lakebase_enabled: + try: + # ``build_pg_executor`` lives in :mod:`backend.pg_executor`, + # which imports :mod:`psycopg` at module load. That import + # is fine in production (the Lakebase-enabled path) but + # would break Delta-only test environments that don't + # install psycopg, so we defer it to this branch. + # ``PgMigrationRunner`` itself is already imported at + # module load — it routes through the psycopg-free + # :mod:`backend.pg_cursor_helpers` module. + from .pg_executor import build_pg_executor + + pg_executor = await asyncio.to_thread( + build_pg_executor, + sp_ws, + instance_name=conf.lakebase_instance_name, + database=conf.lakebase_database_name, + schema=conf.lakebase_schema_name, + token_refresh_minutes=conf.lakebase_token_refresh_minutes, + token_refresh_retry_seconds=conf.lakebase_token_refresh_retry_seconds, + token_refresh_retry_jitter=conf.lakebase_token_refresh_retry_jitter, + token_refresh_max_failures=conf.lakebase_token_refresh_max_failures, + pool_min_size=conf.lakebase_pool_min_size, + pool_max_size=conf.lakebase_pool_max_size, + ) + pg_runner = PgMigrationRunner(pg_executor) + pg_applied = await asyncio.to_thread(pg_runner.run_all) + if pg_applied: + logger.info("Applied %d Lakebase migration(s)", pg_applied) + else: + logger.info("Lakebase schema is up to date") + set_oltp_executor(pg_executor) + logger.info( + "Lakebase OLTP routing enabled (instance=%s, database=%s, schema=%s)", + conf.lakebase_instance_name, + conf.lakebase_database_name, + conf.lakebase_schema_name, + ) + except Exception: + # Close any partially-built pool before re-raising so a + # restart loop doesn't accumulate orphaned server-side + # Postgres connections on every flap. ``close()`` is + # idempotent and best-effort (see pg_executor.close). + if pg_executor is not None: + try: + await asyncio.to_thread(pg_executor.close) + except Exception: + # Best-effort cleanup inside the OUTER ``except`` — + # we're already about to re-raise the init failure, + # so a close-time exception here would mask the real + # root cause from the operator. Log and let the + # outer raise propagate. (Same resilience contract + # as :meth:`pg_executor.PgExecutor.close`; see the + # BLE001 policy in pyproject.toml.) + logger.warning("Error closing Lakebase pool during init failure", exc_info=True) + set_oltp_executor(None) + logger.exception( + "Lakebase initialisation failed (instance=%s, database=%s, schema=%s). " + "Refusing to start — silent fallback to Delta would split OLTP writes across " + "two physical stores and orphan prior Lakebase data on every flap. " + "Common causes: the database_instance is not provisioned, the app SP lacks " + "CAN_CONNECT_AND_CREATE on the bound database, OAuth token issuance is failing, " + "or the Lakebase endpoint is transiently unreachable. Fix the underlying issue " + "and the platform will restart this container automatically. To intentionally " + "run on Delta only, unset DQX_LAKEBASE_INSTANCE_NAME.", + conf.lakebase_instance_name, + conf.lakebase_database_name, + conf.lakebase_schema_name, + ) + raise + else: + logger.info("Lakebase not configured (DQX_LAKEBASE_INSTANCE_NAME is empty). OLTP tables will live on Delta.") + set_oltp_executor(None) + + # Delta migrations always run, but the OLTP fallback DDL is + # skipped when Lakebase owns those tables — the same data model + # is created in Postgres above. runner = MigrationRunner(sql=sp_sql) - applied = runner.run_all() + applied = runner.run_all(include_oltp_fallback=pg_executor is None) if applied: - logger.info("Applied %d database migration(s)", applied) + logger.info("Applied %d Delta migration(s)", applied) else: - logger.info("Database schema is up to date") + logger.info("Delta schema is up to date") # Best-effort below — the app can recover from these failing. @@ -234,6 +354,7 @@ async def lifespan(app: FastAPI): schema=conf.schema_name, tmp_schema=conf.tmp_schema_name, job_id=conf.job_id, + oltp_sql=pg_executor, ) set_scheduler(_scheduler) _scheduler.start() @@ -252,6 +373,23 @@ async def lifespan(app: FastAPI): await sched.stop() set_scheduler(None) + # Close the Lakebase pool last so any in-flight writes from + # ``sched.stop()`` finish first. + if pg_executor is not None: + try: + await asyncio.to_thread(pg_executor.close) + logger.info("Lakebase connection pool closed") + except Exception: + # Lifespan shutdown: same resilience contract as + # :meth:`pg_executor.PgExecutor.close` itself. A raise here + # would prevent the ``set_oltp_executor(None)`` reset below + # from running, leaving a closed pool wired in for the next + # request and producing a confusing "operation on closed + # pool" error instead of the underlying close failure. + # See the BLE001 policy block in pyproject.toml. + logger.warning("Error closing Lakebase pool", exc_info=True) + set_oltp_executor(None) + app = FastAPI(title=f"{conf.app_name}", lifespan=lifespan) diff --git a/app/src/databricks_labs_dqx_app/backend/config.py b/app/src/databricks_labs_dqx_app/backend/config.py index c0c79c1d3..03aa9b2c9 100644 --- a/app/src/databricks_labs_dqx_app/backend/config.py +++ b/app/src/databricks_labs_dqx_app/backend/config.py @@ -27,8 +27,8 @@ class AppConfig(BaseSettings): app_name: str = Field(default=app_name) api_prefix: str = Field(default="/api") catalog: str = Field(default="dqx") - schema_name: str = Field(default="dqx_app", validation_alias="DQX_SCHEMA") - tmp_schema_name: str = Field(default="dqx_app_tmp", validation_alias="DQX_TMP_SCHEMA") + schema_name: str = Field(default="dqx_studio", validation_alias="DQX_SCHEMA") + tmp_schema_name: str = Field(default="dqx_studio_tmp", validation_alias="DQX_TMP_SCHEMA") job_id: str = Field(default="", validation_alias="DQX_JOB_ID") wheels_volume: str = Field(default="", validation_alias="DQX_WHEELS_VOLUME") llm_endpoint: str = Field(default="databricks-claude-sonnet-4-5", validation_alias="DQX_LLM_ENDPOINT") @@ -42,10 +42,109 @@ class AppConfig(BaseSettings): dryrun_max_sample_size: int = Field(default=10_000) dryrun_default_sample_size: int = Field(default=1_000) + # ------------------------------------------------------------------ + # Lakebase (Postgres) backend + # ------------------------------------------------------------------ + # When ``lakebase_instance_name`` is set the OLTP-style tables + # (rules catalog, app settings, RBAC, comments, schedule configs, + # scheduler bookkeeping) are routed to a Lakebase Postgres instance + # instead of Delta. Bulk/append-only tables (validation runs, + # profiling results, metrics, quarantine records) always stay in + # Delta because they are written by the Spark task runner. + # + # Leaving these empty keeps the legacy "everything on Delta" + # behaviour, so existing deployments continue to work without + # changes. See ``app/databricks.yml`` for the deploy-time toggle. + lakebase_instance_name: str = Field( + default="", + validation_alias="DQX_LAKEBASE_INSTANCE_NAME", + description="Lakebase instance name. Empty disables Lakebase routing.", + ) + # Default must match the ``lakebase_database_name`` bundle var in + # ``app/databricks.yml`` (``databricks_postgres`` — the always-present + # admin database every Lakebase instance ships with). The bundle + # intentionally does NOT create a per-app logical Postgres database; + # per-app isolation lives in ``lakebase_schema_name`` below. Keep + # the two defaults in sync so contexts where the DABs env-var + # injection isn't in effect — local dev, unit tests, debug shells — + # don't try to connect to a logical DB that was never created. + lakebase_database_name: str = Field( + default="databricks_postgres", + validation_alias="DQX_LAKEBASE_DATABASE_NAME", + description="Database within the Lakebase instance the app connects to.", + ) + lakebase_schema_name: str = Field( + default="dqx_studio", + validation_alias="DQX_LAKEBASE_SCHEMA", + description="Postgres schema for app tables. Created at startup if missing.", + ) + lakebase_pool_min_size: int = Field(default=1, validation_alias="DQX_LAKEBASE_POOL_MIN_SIZE") + lakebase_pool_max_size: int = Field(default=10, validation_alias="DQX_LAKEBASE_POOL_MAX_SIZE") + # Lakebase OAuth tokens currently expire after one hour; refresh + # well before that so in-flight queries never see a 401. + lakebase_token_refresh_minutes: int = Field( + default=50, + validation_alias="DQX_LAKEBASE_TOKEN_REFRESH_MINUTES", + ) + # ------------------------------------------------------------------ + # Token-refresh resilience (see PgExecutor._token_refresh_loop). + # ------------------------------------------------------------------ + # When a refresh attempt fails we sleep ``retry_seconds`` ± jitter + # and try again *fast* — NOT the full 50-minute scheduled interval + # — because the pool's ``max_lifetime`` is also 50 minutes. If we + # only retried at the scheduled cadence, a refresh failure window + # of 50-60 minutes would silently drain the pool: every in-pool + # connection dies at ``max_lifetime``, no valid replacement is + # ever minted, and the app stops accepting work without any + # explicit "token rotation failed" error visible to operators. + # + # After ``max_failures`` consecutive failures (default 12 × 10s ≈ + # 2 minutes of sustained failure) the executor escalates by + # exiting the process so the supervisor (uvicorn / Databricks + # Apps) restarts the worker. Crashing loud at 2 minutes is much + # better than silently degrading at 50 minutes. + lakebase_token_refresh_retry_seconds: int = Field( + default=10, + validation_alias="DQX_LAKEBASE_TOKEN_REFRESH_RETRY_SECONDS", + description=( + "Base back-off between failed token-refresh attempts. " "Jittered by ±retry_jitter on each retry." + ), + ) + lakebase_token_refresh_retry_jitter: float = Field( + default=0.3, + validation_alias="DQX_LAKEBASE_TOKEN_REFRESH_RETRY_JITTER", + description=( + "Fractional jitter applied to retry back-off so multiple " + "workers don't thunder-herd against the SDK on the same " + "second after a shared transient failure." + ), + ) + lakebase_token_refresh_max_failures: int = Field( + default=12, + validation_alias="DQX_LAKEBASE_TOKEN_REFRESH_MAX_FAILURES", + description=( + "Consecutive refresh failures before the executor escalates " + "by exiting the process so the supervisor restarts the " + "worker. Default 12 × 10s ≈ 2 minutes of sustained " + "failure, which leaves ~48 minutes of pool buffer before " + "max_lifetime would otherwise drain the pool silently." + ), + ) + @property def static_assets_path(self) -> Path: return Path(str(resources.files(app_slug))).joinpath("__dist__") + @property + def lakebase_enabled(self) -> bool: + """``True`` when the deployment was provisioned with Lakebase. + + Falls back to ``False`` (legacy UC-only mode) when the + instance name is empty so existing tests and dev setups keep + working with no Postgres dependency. + """ + return bool(self.lakebase_instance_name.strip()) + conf = AppConfig() diff --git a/app/src/databricks_labs_dqx_app/backend/dependencies.py b/app/src/databricks_labs_dqx_app/backend/dependencies.py index 9521af2d5..e1ef0f5cb 100644 --- a/app/src/databricks_labs_dqx_app/backend/dependencies.py +++ b/app/src/databricks_labs_dqx_app/backend/dependencies.py @@ -29,7 +29,37 @@ from .services.comments_service import CommentsService from .services.schedule_config_service import ScheduleConfigService from .services.view_service import ViewService -from .sql_executor import SqlExecutor +from .sql_executor import OltpExecutorProtocol, SqlExecutor + +# Process-wide OLTP executor. Constructed once at app startup by +# ``app.lifespan`` and reused across all requests so the psycopg pool +# isn't rebuilt per call. ``None`` means Lakebase is not configured and +# the legacy Delta executor handles OLTP traffic instead. +# +# Annotated with :class:`OltpExecutorProtocol` so neither this module +# nor downstream callers need to import :class:`PgExecutor` directly +# — the Protocol is the only contract the OLTP plumbing depends on. +# Avoids dragging psycopg into the import graph for Delta-only +# deployments. +_pg_executor: OltpExecutorProtocol | None = None + + +def set_oltp_executor(executor: OltpExecutorProtocol | None) -> None: + """Register (or clear) the process-wide OLTP executor. + + Called from :func:`backend.app.lifespan` after the connection pool + is open. Keeping this in module state (rather than passing it + through every request) lets the FastAPI ``Depends`` graph stay + request-local while still sharing the pool. + """ + global _pg_executor + _pg_executor = executor + + +def get_oltp_executor() -> OltpExecutorProtocol | None: + """Return the registered OLTP executor or ``None`` if Lakebase is off.""" + return _pg_executor + _SP_TTL = 45 * 60 # 45 minutes _OBO_TTL = 45 * 60 # 45 minutes @@ -116,6 +146,31 @@ async def get_obo_sql_executor( ) +async def get_sp_oltp_executor( + sp_sql: Annotated[SqlExecutor, Depends(get_sp_sql_executor)], +) -> OltpExecutorProtocol: + """Return the executor that owns the OLTP tables. + + When Lakebase is configured the lifespan handler registers a + :class:`backend.pg_executor.PgExecutor` via :func:`set_oltp_executor` + and we hand it back to every OLTP service. Otherwise we fall back + to the legacy Delta executor (``get_sp_sql_executor``) so existing + deployments keep working with no code changes on their side. + + The return type is :class:`OltpExecutorProtocol` so every + downstream service annotation type-checks against the structural + surface shared by both executors. No ``cast`` is needed: both + :class:`SqlExecutor` and :class:`PgExecutor` structurally satisfy + the Protocol, so the type-checker validates every ``.execute()``, + ``.query()`` etc. call against a single source of truth instead + of being muted by a Delta-class cast around the Postgres path. + """ + pg = get_oltp_executor() + if pg is None: + return sp_sql + return pg + + # --------------------------------------------------------------------------- # Service factories # --------------------------------------------------------------------------- @@ -124,21 +179,26 @@ async def get_obo_sql_executor( async def get_migration_runner( sql: Annotated[SqlExecutor, Depends(get_sp_sql_executor)], ) -> MigrationRunner: - """Create a MigrationRunner using app (SP) credentials.""" + """Create the Delta MigrationRunner using app (SP) credentials. + + The Postgres :class:`PgMigrationRunner` is constructed separately + in :func:`backend.app.lifespan` because it needs the running + :class:`PgExecutor`, not a SQL warehouse executor. + """ return MigrationRunner(sql=sql) async def get_app_settings_service( - sql: Annotated[SqlExecutor, Depends(get_sp_sql_executor)], + sql: Annotated[OltpExecutorProtocol, Depends(get_sp_oltp_executor)], ) -> AppSettingsService: - """Create an AppSettingsService using app (SP) credentials.""" + """Create an AppSettingsService routed at the OLTP executor.""" return AppSettingsService(sql=sql) async def get_role_service( - sql: Annotated[SqlExecutor, Depends(get_sp_sql_executor)], + sql: Annotated[OltpExecutorProtocol, Depends(get_sp_oltp_executor)], ) -> RoleService: - """Create a RoleService using app (SP) credentials.""" + """Create a RoleService routed at the OLTP executor.""" return RoleService(sql=sql) @@ -155,9 +215,9 @@ async def get_ai_rules_service( async def get_rules_catalog_service( - sql: Annotated[SqlExecutor, Depends(get_sp_sql_executor)], + sql: Annotated[OltpExecutorProtocol, Depends(get_sp_oltp_executor)], ) -> RulesCatalogService: - """Create a RulesCatalogService using app (SP) credentials.""" + """Create a RulesCatalogService routed at the OLTP executor.""" return RulesCatalogService(sql=sql) @@ -184,16 +244,16 @@ async def get_view_service( async def get_comments_service( - sql: Annotated[SqlExecutor, Depends(get_sp_sql_executor)], + sql: Annotated[OltpExecutorProtocol, Depends(get_sp_oltp_executor)], ) -> CommentsService: - """Create a CommentsService using app (SP) credentials.""" + """Create a CommentsService routed at the OLTP executor.""" return CommentsService(sql=sql) async def get_schedule_config_service( - sql: Annotated[SqlExecutor, Depends(get_sp_sql_executor)], + sql: Annotated[OltpExecutorProtocol, Depends(get_sp_oltp_executor)], ) -> ScheduleConfigService: - """Create a ScheduleConfigService using app (SP) credentials.""" + """Create a ScheduleConfigService routed at the OLTP executor.""" return ScheduleConfigService(sql=sql) @@ -372,6 +432,9 @@ async def get_user_catalog_names( "get_obo_ws", "get_sp_sql_executor", "get_obo_sql_executor", + "get_sp_oltp_executor", + "get_oltp_executor", + "set_oltp_executor", "get_conf", "get_check_validator", "get_migration_runner", diff --git a/app/src/databricks_labs_dqx_app/backend/migrations/__init__.py b/app/src/databricks_labs_dqx_app/backend/migrations/__init__.py index e2f164f9f..f14f6ca80 100644 --- a/app/src/databricks_labs_dqx_app/backend/migrations/__init__.py +++ b/app/src/databricks_labs_dqx_app/backend/migrations/__init__.py @@ -1,16 +1,105 @@ -"""Database migration runner for DQX Studio. +"""Database migration runner for DQX Studio (Delta). Migrations are versioned DDL statements applied in order against the -configured catalog/schema. The runner tracks every applied version in a -``dq_migrations`` meta-table, so re-starting the app never re-applies a -migration that already succeeded. - -Schema baseline ---------------- -The app has not shipped to external users, so the schema is delivered -as a **single consolidated baseline** rather than 36 incremental -migrations. Every table is created at its final shape, with liquid -clustering inlined into the ``CREATE TABLE`` statement. +configured Unity Catalog catalog/schema. The runner tracks every +applied version in a ``dq_migrations`` meta-table, so re-starting the +app never re-applies a migration that already succeeded. + +Hybrid backend split +-------------------- +The schema is delivered in **two parts** so the OLTP-style tables can +optionally live on Lakebase Postgres while the high-volume analytical +tables stay in Delta: + +- **v1 — Delta analytical baseline** (always applied). Holds the + Spark-written tables: ``dq_validation_runs``, + ``dq_profiling_results``, ``dq_quarantine_records``, + ``dq_metrics``. +- **v2 — Delta OLTP fallback** (only applied when Lakebase is + disabled, i.e. ``include_oltp_fallback=True``). Holds the + FastAPI-served tables: ``dq_app_settings``, ``dq_quality_rules``, + ``dq_quality_rules_history``, ``dq_role_mappings``, ``dq_comments``, + ``dq_schedule_configs``, ``dq_schedule_configs_history``, + ``dq_schedule_runs``. + +When Lakebase is enabled the same OLTP tables are created via +:mod:`backend.migrations.postgres` against the Postgres schema and v2 +is skipped on the Delta side. + +Atomicity model — Delta vs Postgres asymmetry +--------------------------------------------- +The two backends have **different** failure-recovery semantics. Read +this before adding a new Delta migration: + +- **Postgres path** (:mod:`backend.migrations.postgres`) — every DDL + statement in a migration *and* the ``dq_migrations`` version-row + ``INSERT`` execute inside a **single transaction** (``BEGIN`` … + ``COMMIT``). If any statement fails the whole migration rolls back; + the next run retries cleanly from the beginning. +- **Delta path** (this module) — the Databricks Statement Execution + API auto-commits **per call**; there is no ``BEGIN``/``COMMIT`` + primitive for DDL on a SQL warehouse, so true atomicity is not + available. A mid-migration failure therefore leaves earlier + statements committed but the version row **uninserted**. Recovery + relies on every DDL statement being individually re-runnable. + +That recovery contract has three concrete invariants every Delta +migration author MUST satisfy: + +1. **Idempotent DDL.** Use ``CREATE TABLE IF NOT EXISTS`` and + ``CREATE SCHEMA IF NOT EXISTS``; never plain ``CREATE``. Plain + ``ALTER TABLE ADD COLUMN`` and ``ALTER TABLE ADD CONSTRAINT`` are + OK because their "already-exists" error fragments are swallowed + (see ``_IDEMPOTENT_ERROR_FRAGMENTS``). +2. **Every "already exists" error a statement can raise must appear in + ``_IDEMPOTENT_ERROR_FRAGMENTS``.** If you introduce a new DDL kind + (``CREATE INDEX``, ``CREATE VIEW``, ``CREATE FUNCTION`` …) and + half-applied recovery surfaces an unfamiliar error message, the + migration becomes un-replayable in production until someone + appends the fragment and ships a fix. Validate by checking the + error catalog for the DDL kind you're adding. +3. **Template scanner invariants** — checked at module import by + ``_validate_template_safe`` against every entry in :data:`MIGRATIONS`, + so any violation crashes the app at boot rather than mid-deploy: + + - **No ``;`` inside a single-quoted string literal.** The runner + splits compound migration bodies on ``;`` because the Statement + Execution API only accepts one statement per call. A ``;`` + inside a CHECK constraint string, default value, or comment + would split a statement mid-literal and silently corrupt it. + If you ever genuinely need ``;`` in a literal, replace + ``str.split(";")`` inside ``_apply`` with a real SQL statement + splitter (e.g. ``sqlparse.split``). + - **No ``{catalog}`` / ``{schema}`` placeholder inside a single- + quoted string literal.** ``_apply`` substitutes the + :meth:`SqlExecutor.q`-quoted form into these placeholders so + hyphenated Databricks catalog names like ``prod-east`` parse + correctly when used as object identifiers. The quoted form + contains backticks, which would be nonsense inside a string + literal. If you need the raw identifier inside a literal, + introduce a separate placeholder rather than reusing + ``{catalog}`` / ``{schema}``. + +The Postgres runner does not need invariants 1 and 2 (the transaction +guarantees it) but does still split on ``;`` and so applies invariant 3 +to its own templates. + +Status casing convention +------------------------ +Two status families intentionally use different casing: + +- **Run-lifecycle** (``dq_validation_runs.status``, + ``dq_profiling_results.status``) — UPPERCASE + (``RUNNING``/``SUCCESS``/``FAILED``/``CANCELED``). These mirror the + Databricks Jobs SDK ``life_cycle_state`` / ``result_state`` values + that are passed straight through ``RunStatusOut`` to the frontend. +- **App-domain workflow** (``dq_quality_rules.status``, + ``dq_schedule_runs.status``) — lowercase. These are pure DQX + vocabulary (``draft``/``approved``, ``pending``/``partial_failure``) + with no SDK counterpart. + +CHECK constraints enforce the agreed values per domain — see each +table's ``chk_*_status`` constraint below. Adding a new table or schema change after baseline -------------------------------------------------- @@ -21,22 +110,16 @@ SQL warehouse versions; ``_apply`` instead catches and tolerates ``COLUMN_ALREADY_EXISTS`` so re-running is safe). -Example:: - - Migration( - version=2, - description="Add description column to dq_role_mappings", - sql_template=( - "ALTER TABLE {catalog}.{schema}.dq_role_mappings " - "ADD COLUMN description STRING" - ), - ) +If the change touches an OLTP table, mirror it in +:mod:`backend.migrations.postgres` so Lakebase deployments stay in +sync. Upgrading an existing dev workspace ----------------------------------- -A workspace that previously ran the legacy v1–v36 sequence will have -``dq_migrations`` rows for versions that no longer exist. The cleanest -path is:: +A workspace that previously ran the legacy migration sequence will have +``dq_migrations`` rows for versions that no longer exist, and tables +whose column types or constraints predate this baseline revision. The +cleanest path is:: DROP SCHEMA . CASCADE; @@ -53,6 +136,100 @@ logger = logging.getLogger(__name__) + +# Migration templates have two invariants that ``_apply`` relies on but +# the language doesn't enforce. ``_validate_template_safe`` checks both +# at module import time against every entry in :data:`MIGRATIONS` so a +# future author hits an :class:`AssertionError` on app boot instead of +# a half-applied or mis-quoted migration in production: +# +# 1. **No ``;`` inside a single-quoted string literal.** +# ``_apply`` splits the template on ``;`` to feed one statement per +# Statement Execution API call. A semicolon inside a string literal +# would split the literal in half and corrupt both halves. +# +# 2. **No ``{catalog}`` or ``{schema}`` placeholder inside a single- +# quoted string literal.** ``_apply`` substitutes the executor- +# *quoted* form into those placeholders (so hyphenated catalogs like +# ``prod-east`` parse correctly when used as identifiers). The +# quoted form contains backticks, which would be nonsense inside a +# string literal context. Today no migration uses placeholders in a +# literal — this guard prevents a future author from introducing one +# without adding a separate raw placeholder first. +# +# Doubled single-quotes (``''``) are the standard SQL escape for a +# literal apostrophe inside a single-quoted string; the scanner handles +# them correctly by keeping the in-literal flag set across the pair. +_SINGLE_QUOTE = "'" +_PLACEHOLDER_TOKENS = ("{catalog}", "{schema}") + + +def _validate_template_safe(template: str) -> None: + """Assert *template* satisfies the runner's substitution invariants. + + See the comment block above this function for the full list. In + short: no ``;`` and no ``{catalog}``/``{schema}`` placeholder may + appear inside a single-quoted string literal. + + The scanner is a tiny state machine: it walks the template character + by character, flips an ``in_literal`` flag on each unescaped ``'``, + and asserts neither forbidden pattern is seen while the flag is set. + + Raises: + AssertionError: if any forbidden pattern appears inside a + single-quoted string literal in *template*. + """ + in_literal = False + i = 0 + n = len(template) + while i < n: + ch = template[i] + if ch == _SINGLE_QUOTE: + # Doubled single-quote ('') inside a literal is the SQL + # escape for a literal apostrophe; keep the flag as-is and + # skip both characters. + if in_literal and i + 1 < n and template[i + 1] == _SINGLE_QUOTE: + i += 2 + continue + in_literal = not in_literal + i += 1 + continue + if in_literal: + if ch == ";": + _raise_template_violation(template, i, "';' inside a single-quoted string literal", split_hint=True) + for tok in _PLACEHOLDER_TOKENS: + if template.startswith(tok, i): + _raise_template_violation( + template, + i, + f"placeholder {tok!r} inside a single-quoted string literal", + split_hint=False, + ) + i += 1 + + +def _raise_template_violation(template: str, offset: int, what: str, *, split_hint: bool) -> None: + """Raise an actionable :class:`AssertionError` for a template scanner failure.""" + start = max(0, offset - 40) + end = min(len(template), offset + 40) + excerpt = template[start:end].replace("\n", " ") + msg = f"Migration template contains {what} at offset {offset} (... {excerpt!r} ...). " + if split_hint: + msg += ( + "MigrationRunner._apply uses str.split(';') which would silently corrupt this " + "statement. Either move the ';' out of the literal, or replace the splitter in " + "_apply with a real SQL statement splitter (e.g. sqlparse.split)." + ) + else: + msg += ( + "MigrationRunner._apply substitutes the executor-quoted form (containing " + "backticks) into these placeholders so identifiers like 'prod-east' stay " + "parseable in object-name positions — but that quoted form would be nonsense " + "inside a string literal. Use a separate raw placeholder for literal contexts." + ) + raise AssertionError(msg) + + # --------------------------------------------------------------------------- # Migration definitions # --------------------------------------------------------------------------- @@ -70,186 +247,334 @@ class Migration: # Order is significant. Never change or remove existing entries — only # append new ones. # -# v1 is the consolidated baseline created on 2026-05-03 by collapsing -# the original v1–v36 incremental sequence. Each table is defined at -# its final shape with liquid clustering inlined; the legacy -# wide-format ``dq_metrics`` (renamed to ``dq_metrics_v1_legacy`` in -# the original v32) is dropped from the baseline because it has no -# clean-install consumer. +# v1 is the consolidated baseline. Each table is defined at its final +# shape with liquid clustering, primary keys, and CHECK constraints +# inlined. Revisions to the baseline are allowed (and encouraged) until +# the app ships externally; existing dev workspaces upgrade by +# ``DROP SCHEMA … CASCADE`` and re-running migrations from scratch. +# +# Notes on column choices: +# - PRIMARY KEY constraints are informational (``NOT ENFORCED RELY``); +# Delta uses them for query optimization, lineage, AI/BI tooling. +# - CHECK constraints ARE enforced — picking the right value-set on +# day one is cheap; loosening later is just an ALTER TABLE. +# - ``VARIANT`` (DBR 15.3+ / serverless) replaces ad-hoc JSON-in-string +# for the largest blob columns (``dq_quality_rules.check``, +# ``dq_quarantine_records.row_data``/``errors``). +# - Run-lifecycle ``status`` columns use UPPERCASE values to mirror the +# Databricks Jobs SDK; app-domain ``status`` columns use lowercase. +# v1 always runs against Delta and only contains the high-volume +# analytical tables that the Spark task runner writes. Keeping these +# in Delta lets AI/BI dashboards consume them directly via SQL +# warehouse without round-tripping through Postgres. +_V1_ANALYTICAL_BASELINE = ( + # Profiler runs — one row per profile job. Mutable lifecycle + # (RUNNING → SUCCESS/FAILED/CANCELED). Status values mirror + # Databricks Jobs SDK convention. + # + # NOTE: Delta only allows PRIMARY KEY and FOREIGN KEY constraints + # inline in CREATE TABLE — every CHECK constraint must be added + # via a separate ALTER TABLE … ADD CONSTRAINT statement after the + # table exists. The migration runner swallows + # ``DELTA_CONSTRAINT_ALREADY_EXISTS`` so re-runs are idempotent. + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_profiling_results (" + " run_id STRING NOT NULL," + " requesting_user STRING," + " source_table_fqn STRING NOT NULL," + " view_fqn STRING," + " sample_limit INT," + " rows_profiled INT," + " columns_profiled INT," + " duration_seconds DOUBLE," + " summary_json STRING," + " generated_rules_json STRING," + " status STRING NOT NULL," + " error_message STRING," + " canceled_by STRING," + " updated_at TIMESTAMP," + " job_run_id BIGINT," + " rule_set_fingerprint STRING," + " created_at TIMESTAMP" + ") CLUSTER BY (source_table_fqn, run_id, created_at);" + f"ALTER TABLE {_PLACEHOLDER}.dq_profiling_results " + f" ADD CONSTRAINT chk_dq_profiling_results_status " + f" CHECK (status IN ('RUNNING','SUCCESS','FAILED','CANCELED'));" + # + # Validation (dryrun + scheduled) runs — one row per run, mutable + # lifecycle. Joins to ``dq_metrics`` on + # ``(run_id, rule_set_fingerprint)``. + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_validation_runs (" + " run_id STRING NOT NULL," + " requesting_user STRING," + " source_table_fqn STRING NOT NULL," + " view_fqn STRING," + " checks_json STRING," + " sample_size INT," + " total_rows INT," + " valid_rows INT," + " invalid_rows INT," + " error_rows INT," + " warning_rows INT," + " error_summary_json STRING," + " sample_invalid_json STRING," + " status STRING NOT NULL," + " error_message STRING," + " canceled_by STRING," + " updated_at TIMESTAMP," + " run_type STRING," + " job_run_id BIGINT," + " rule_set_fingerprint STRING," + " created_at TIMESTAMP" + ") CLUSTER BY (source_table_fqn, run_id, created_at);" + f"ALTER TABLE {_PLACEHOLDER}.dq_validation_runs " + f" ADD CONSTRAINT chk_dq_validation_runs_status " + f" CHECK (status IN ('RUNNING','SUCCESS','FAILED','CANCELED'));" + f"ALTER TABLE {_PLACEHOLDER}.dq_validation_runs " + f" ADD CONSTRAINT chk_dq_validation_runs_run_type " + f" CHECK (run_type IS NULL OR run_type IN ('dryrun','scheduled','preview'));" + # + # Quarantined invalid rows captured during validation. ``row_data`` + # and ``errors`` are VARIANT for native JSON predicate pushdown and + # ~10x compression vs. STRING. + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_quarantine_records (" + " quarantine_id STRING NOT NULL," + " run_id STRING NOT NULL," + " source_table_fqn STRING NOT NULL," + " requesting_user STRING," + " row_data VARIANT," + " errors VARIANT," + " warnings VARIANT," + " created_at TIMESTAMP," + " CONSTRAINT pk_dq_quarantine_records PRIMARY KEY (quarantine_id) RELY" + ") CLUSTER BY (run_id);" + # + # Long-format observability events written by DQMetricsObserver. + # Schema mirrors the public DQX OBSERVATION_TABLE_SCHEMA so AI/BI + # dashboard templates targeting the spec drop straight in. + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_metrics (" + " run_id STRING NOT NULL," + " run_name STRING," + " input_location STRING," + " output_location STRING," + " quarantine_location STRING," + " checks_location STRING," + " rule_set_fingerprint STRING," + " metric_name STRING NOT NULL," + " metric_value STRING," + " run_time TIMESTAMP NOT NULL," + " error_column_name STRING," + " warning_column_name STRING," + " user_metadata MAP" + ") CLUSTER BY (input_location, run_id, run_time)" +) + + +# v2 is the Delta-only OLTP fallback. It is **only** applied when +# Lakebase is disabled (``include_oltp_fallback=True`` in +# :meth:`MigrationRunner.run_all`). When Lakebase is enabled, the same +# tables are created via :mod:`backend.migrations.postgres` against the +# Postgres backend. +_V2_OLTP_FALLBACK = ( + # Settings — single-row-per-key key/value store (workspace config, + # label catalog, custom metrics, timezone, ...). + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_app_settings (" + " setting_key STRING NOT NULL," + " setting_value STRING," + " updated_at TIMESTAMP," + " updated_by STRING," + " CONSTRAINT pk_dq_app_settings PRIMARY KEY (setting_key) RELY" + ") CLUSTER BY (setting_key);" + # + # Active rule catalog. ``rule_id`` is a per-check stable identifier; + # each row holds exactly ONE check serialized as a VARIANT object + # (no array wrapper). ``source`` records which authoring path + # produced the rule. + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_quality_rules (" + " rule_id STRING NOT NULL," + " table_fqn STRING NOT NULL," + " `check` VARIANT NOT NULL," + " version INT NOT NULL," + " status STRING NOT NULL," + " source STRING NOT NULL," + " created_by STRING," + " created_at TIMESTAMP," + " updated_by STRING," + " updated_at TIMESTAMP," + " CONSTRAINT pk_dq_quality_rules PRIMARY KEY (rule_id) RELY" + ") CLUSTER BY (table_fqn, status, rule_id);" + f"ALTER TABLE {_PLACEHOLDER}.dq_quality_rules " + f" ADD CONSTRAINT chk_dq_quality_rules_status " + f" CHECK (status IN ('draft','pending_approval','approved','rejected'));" + f"ALTER TABLE {_PLACEHOLDER}.dq_quality_rules " + f" ADD CONSTRAINT chk_dq_quality_rules_source " + f" CHECK (source IN ('ui','sql','profiler','import','ai'));" + # + # Append-only audit trail for rule changes. Carries the post-state + # ``check`` payload on every row plus an explicit + # ``prev_status``/``new_status`` pair for status transitions. + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_quality_rules_history (" + " rule_id STRING," + " table_fqn STRING NOT NULL," + " `check` VARIANT," + " version INT," + " source STRING," + " action STRING NOT NULL," + " prev_status STRING," + " new_status STRING," + " changed_by STRING," + " changed_at TIMESTAMP" + ") CLUSTER BY (table_fqn, changed_at);" + # + # RBAC: maps app roles (admin/rule_approver/rule_author/viewer/ + # runner) to Databricks workspace groups. Tiny table — no + # clustering needed. + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_role_mappings (" + " role STRING NOT NULL," + " group_name STRING NOT NULL," + " created_by STRING," + " created_at TIMESTAMP," + " updated_by STRING," + " updated_at TIMESTAMP," + " CONSTRAINT pk_dq_role_mappings PRIMARY KEY (role, group_name) RELY" + ");" + f"ALTER TABLE {_PLACEHOLDER}.dq_role_mappings " + f" ADD CONSTRAINT chk_dq_role_mappings_role " + f" CHECK (role IN ('admin','rule_approver','rule_author','viewer','runner'));" + # + # Per-entity comment threads (rules, runs, profiles, ...). + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_comments (" + " comment_id STRING NOT NULL," + " entity_type STRING NOT NULL," + " entity_id STRING NOT NULL," + " user_email STRING NOT NULL," + " comment STRING NOT NULL," + " created_at TIMESTAMP," + " CONSTRAINT pk_dq_comments PRIMARY KEY (comment_id) RELY" + ") CLUSTER BY (entity_type, entity_id);" + f"ALTER TABLE {_PLACEHOLDER}.dq_comments " + f" ADD CONSTRAINT chk_dq_comments_entity_type " + f" CHECK (entity_type IN ('run','rule'));" + # + # Scheduler bookkeeping: last/next run pointer per schedule. + # ``status`` is app-domain (lowercase). + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_schedule_runs (" + " schedule_name STRING NOT NULL," + " last_run_at TIMESTAMP," + " next_run_at TIMESTAMP," + " last_run_id STRING," + " status STRING," + " updated_at TIMESTAMP," + " CONSTRAINT pk_dq_schedule_runs PRIMARY KEY (schedule_name) RELY" + ") CLUSTER BY (schedule_name);" + f"ALTER TABLE {_PLACEHOLDER}.dq_schedule_runs " + f" ADD CONSTRAINT chk_dq_schedule_runs_status " + f" CHECK (status IS NULL OR status IN ('pending','success','partial_failure','failed'));" + # + # Per-schedule live config (cron/interval, scope filters). + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_schedule_configs (" + " schedule_name STRING NOT NULL," + " config_json STRING NOT NULL," + " version INT NOT NULL," + " created_by STRING," + " created_at TIMESTAMP," + " updated_by STRING," + " updated_at TIMESTAMP," + " CONSTRAINT pk_dq_schedule_configs PRIMARY KEY (schedule_name) RELY" + ") CLUSTER BY (schedule_name);" + # + # Append-only audit trail for schedule changes. + f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_schedule_configs_history (" + " schedule_name STRING NOT NULL," + " config_json STRING," + " version INT," + " action STRING NOT NULL," + " changed_by STRING," + " changed_at TIMESTAMP" + ") CLUSTER BY (schedule_name, changed_at)" +) + + +# Backfills ``warning_rows`` on workspaces deployed before v1 added it. +# On fresh deploys ``_apply`` swallows the ``COLUMN_ALREADY_EXISTS`` error +# per the column-addition rule documented at the top of this module. +_V3_VALIDATION_RUNS_WARNING_ROWS = f"ALTER TABLE {_PLACEHOLDER}.dq_validation_runs " f" ADD COLUMN warning_rows INT" + + +# Quarantine rows that fail only warning-level checks would otherwise +# show an empty ``errors`` column in the UI. We mirror DQX's row-level +# ``_warnings`` map into its own VARIANT so warnings can be rendered +# alongside errors in the dry-run sample table. +_V4_QUARANTINE_WARNINGS = f"ALTER TABLE {_PLACEHOLDER}.dq_quarantine_records " f" ADD COLUMN warnings VARIANT" + + +# ``invalid_rows`` (set from ``invalid_df.count()``) conflated "rows that +# failed any check" with "rows with errors" — and could over-count when +# certain DQX checks fan out internally. ``error_rows`` is the +# authoritative count from the DQX observer (``error_row_count``), so the +# UI now surfaces it as the primary "Errors" stat. ``invalid_rows`` is +# kept for backwards compatibility but no longer drives the UI. +_V5_VALIDATION_RUNS_ERROR_ROWS = f"ALTER TABLE {_PLACEHOLDER}.dq_validation_runs " f" ADD COLUMN error_rows INT" + + +# OLTP fallback migration is identified by ``oltp_fallback=True`` so +# the runner can skip it when Lakebase is enabled. Keeping the flag on +# the migration itself (rather than e.g. a hard-coded version number) +# makes it easy to add follow-up Delta-only OLTP migrations later +# without re-discovering the rule. +@dataclass(frozen=True) +class DeltaMigration(Migration): + """Migration variant that knows whether it carries OLTP fallback DDL. + + A subclass (rather than a flag on :class:`Migration`) keeps + backwards compatibility for any callers that still hand-build + ``Migration`` instances and don't care about the flag. + """ + + oltp_fallback: bool = False + + MIGRATIONS: list[Migration] = [ - Migration( + DeltaMigration( version=1, - description="Baseline schema — all DQX Studio tables at their final shape", - sql_template=( - # Settings — single-row-per-key key/value store (workspace - # config, label catalog, custom metrics, timezone, ...). - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_app_settings (" - " setting_key STRING NOT NULL," - " setting_value STRING," - " updated_at TIMESTAMP," - " updated_by STRING" - ") CLUSTER BY (setting_key);" - # - # Active rule catalog. ``rule_id`` is a per-check stable - # identifier; ``source`` records which authoring path - # produced the rule (single-table, sql, profiler, import). - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_quality_rules (" - " table_fqn STRING NOT NULL," - " checks STRING NOT NULL," - " version INT," - " status STRING," - " source STRING," - " rule_id STRING," - " created_by STRING," - " created_at TIMESTAMP," - " updated_by STRING," - " updated_at TIMESTAMP" - ") CLUSTER BY (table_fqn, status);" - # - # Append-only audit trail for rule changes. - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_quality_rules_history (" - " table_fqn STRING NOT NULL," - " checks STRING," - " version INT," - " source STRING," - " rule_id STRING," - " action STRING NOT NULL," - " changed_by STRING," - " changed_at TIMESTAMP" - ") CLUSTER BY (table_fqn, changed_at);" - # - # Profiler runs — one row per profile job. Mutable - # lifecycle (RUNNING → SUCCESS/FAILED/CANCELED). - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_profiling_results (" - " run_id STRING NOT NULL," - " requesting_user STRING," - " source_table_fqn STRING NOT NULL," - " view_fqn STRING," - " sample_limit INT," - " rows_profiled INT," - " columns_profiled INT," - " duration_seconds DOUBLE," - " summary_json STRING," - " generated_rules_json STRING," - " status STRING," - " error_message STRING," - " canceled_by STRING," - " updated_at STRING," - " job_run_id BIGINT," - " rule_set_fingerprint STRING," - " created_at STRING" - ") CLUSTER BY (source_table_fqn, created_at);" - # - # Validation (dryrun + scheduled) runs — one row per run, - # mutable lifecycle. Joins to ``dq_metrics`` on - # ``(run_id, rule_set_fingerprint)``. - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_validation_runs (" - " run_id STRING NOT NULL," - " requesting_user STRING," - " source_table_fqn STRING NOT NULL," - " view_fqn STRING," - " checks_json STRING," - " sample_size INT," - " total_rows INT," - " valid_rows INT," - " invalid_rows INT," - " error_summary_json STRING," - " sample_invalid_json STRING," - " status STRING," - " error_message STRING," - " canceled_by STRING," - " updated_at STRING," - " run_type STRING," - " job_run_id BIGINT," - " rule_set_fingerprint STRING," - " created_at STRING" - ") CLUSTER BY (source_table_fqn, created_at);" - # - # RBAC: maps app roles (ADMIN/RULE_APPROVER/RULE_AUTHOR/ - # VIEWER/RUNNER) to Databricks workspace groups. - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_role_mappings (" - " role STRING NOT NULL," - " group_name STRING NOT NULL," - " created_by STRING," - " created_at TIMESTAMP," - " updated_by STRING," - " updated_at TIMESTAMP" - ") CLUSTER BY (role);" - # - # Per-entity comment threads (rules, runs, profiles, ...). - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_comments (" - " comment_id STRING NOT NULL," - " entity_type STRING NOT NULL," - " entity_id STRING NOT NULL," - " user_email STRING NOT NULL," - " comment STRING NOT NULL," - " created_at TIMESTAMP" - ") CLUSTER BY (entity_type, entity_id);" - # - # Quarantined invalid rows captured during validation. - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_quarantine_records (" - " quarantine_id STRING NOT NULL," - " run_id STRING NOT NULL," - " source_table_fqn STRING NOT NULL," - " requesting_user STRING," - " row_data STRING," - " errors STRING," - " created_at STRING" - ") CLUSTER BY (run_id, source_table_fqn);" - # - # Long-format observability events written by - # DQMetricsObserver. Schema mirrors the public DQX - # OBSERVATION_TABLE_SCHEMA so AI/BI dashboard templates - # targeting the spec drop straight in. - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_metrics (" - " run_id STRING NOT NULL," - " run_name STRING," - " input_location STRING," - " output_location STRING," - " quarantine_location STRING," - " checks_location STRING," - " rule_set_fingerprint STRING," - " metric_name STRING NOT NULL," - " metric_value STRING," - " run_time TIMESTAMP NOT NULL," - " error_column_name STRING," - " warning_column_name STRING," - " user_metadata MAP" - ") CLUSTER BY (input_location, run_time);" - # - # Scheduler bookkeeping: last/next run pointer per schedule. - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_schedule_runs (" - " schedule_name STRING NOT NULL," - " last_run_at TIMESTAMP," - " next_run_at TIMESTAMP," - " last_run_id STRING," - " status STRING" - ");" - # - # Per-schedule live config (cron/interval, scope filters). - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_schedule_configs (" - " schedule_name STRING NOT NULL," - " config_json STRING NOT NULL," - " version INT," - " created_by STRING," - " created_at TIMESTAMP," - " updated_by STRING," - " updated_at TIMESTAMP" - ") CLUSTER BY (schedule_name);" - # - # Append-only audit trail for schedule changes. - f"CREATE TABLE IF NOT EXISTS {_PLACEHOLDER}.dq_schedule_configs_history (" - " schedule_name STRING NOT NULL," - " config_json STRING," - " version INT," - " action STRING NOT NULL," - " changed_by STRING," - " changed_at TIMESTAMP" - ") CLUSTER BY (schedule_name, changed_at)" - ), + description="Delta analytical baseline (validation, profiling, quarantine, metrics)", + sql_template=_V1_ANALYTICAL_BASELINE, + oltp_fallback=False, + ), + DeltaMigration( + version=2, + description="Delta OLTP fallback (rules, app settings, RBAC, schedules) — used only when Lakebase is disabled", + sql_template=_V2_OLTP_FALLBACK, + oltp_fallback=True, + ), + DeltaMigration( + version=3, + description="Add warning_rows column to dq_validation_runs (backfill for pre-v3 deploys)", + sql_template=_V3_VALIDATION_RUNS_WARNING_ROWS, + oltp_fallback=False, + ), + DeltaMigration( + version=4, + description="Add warnings VARIANT column to dq_quarantine_records (mirror DQX _warnings map)", + sql_template=_V4_QUARANTINE_WARNINGS, + oltp_fallback=False, + ), + DeltaMigration( + version=5, + description="Add error_rows column to dq_validation_runs (DQX error_row_count, replaces invalid_rows for UI)", + sql_template=_V5_VALIDATION_RUNS_ERROR_ROWS, + oltp_fallback=False, ), ] + +# Fail loudly at app import if any migration template violates the +# scanner invariants (no ``;`` and no ``{catalog}``/``{schema}`` +# placeholder inside a string literal). Converts latent correctness +# risks into deploy-time AssertionErrors; see the "Atomicity model" +# section of the module docstring for the full recovery contract. +for _m in MIGRATIONS: + _validate_template_safe(_m.sql_template) + # --------------------------------------------------------------------------- # Runner # --------------------------------------------------------------------------- @@ -274,14 +599,37 @@ def __init__(self, sql: SqlExecutor) -> None: self._sql = sql self._catalog = sql.catalog self._schema = sql.schema - self._meta_table = _META_TABLE.format(catalog=sql.catalog, schema=sql.schema) + # Pre-compute the executor-quoted forms once. Every place the + # catalog/schema appears in a SQL statement uses the quoted + # form so a hyphenated Databricks catalog name (``prod-east``, + # ``team-data-platform``) — which is parse-invalid raw — stays + # safe end-to-end. Templates use ``{catalog}``/``{schema}`` + # placeholders in object-name positions only (never inside a + # string literal), so substituting the quoted form is always + # correct. See :meth:`SqlExecutor.q` for the dialect contract. + self._catalog_q = sql.q(sql.catalog) + self._schema_q = sql.q(sql.schema) + self._meta_table = _META_TABLE.format(catalog=self._catalog_q, schema=self._schema_q) # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ - def run_all(self) -> int: - """Ensure the schema exists and apply all pending migrations. + def run_all(self, *, include_oltp_fallback: bool = True) -> int: + """Ensure the schema exists and apply all pending Delta migrations. + + Parameters + ---------- + include_oltp_fallback: + When ``True`` (legacy mode, no Lakebase) all migrations + run including the OLTP fallback DDL (v2 in the baseline). + When ``False`` (Lakebase enabled) migrations marked with + ``oltp_fallback=True`` are skipped — the same tables are + created in Postgres via :class:`PgMigrationRunner` instead. + + Skipped migrations are *not* recorded as applied, so + disabling Lakebase later will cause them to run on the + next deploy and create the Delta-side tables on demand. Returns: The number of migrations applied in this invocation. @@ -300,6 +648,15 @@ def run_all(self) -> int: ) continue + if not include_oltp_fallback and isinstance(migration, DeltaMigration) and migration.oltp_fallback: + logger.info( + "Skipping Delta OLTP fallback migration v%d " + "(Lakebase enabled — these tables live in Postgres): %s", + migration.version, + migration.description, + ) + continue + logger.info( "Applying migration v%d: %s", migration.version, @@ -340,8 +697,13 @@ def _ensure_schema(self) -> None: Must use the bootstrap executor (catalog-only context) because the schema does not exist yet and passing ``schema=`` to the statement-execution API would cause it to fail before the DDL runs. + + Catalog and schema are interpolated through :attr:`_catalog_q` / + :attr:`_schema_q` so a hyphenated catalog (``prod-east``) — which + is parse-invalid raw — gets correctly backtick-quoted by + :meth:`SqlExecutor.q`. """ - sql = f"CREATE SCHEMA IF NOT EXISTS {self._catalog}.{self._schema}" + sql = f"CREATE SCHEMA IF NOT EXISTS {self._catalog_q}.{self._schema_q}" self._sql.execute_no_schema(sql) logger.debug("Ensured schema exists: %s.%s", self._catalog, self._schema) @@ -373,14 +735,65 @@ def _applied_at_map(self) -> dict[int, str]: # the whole migration to abort. _IDEMPOTENT_ERROR_FRAGMENTS = ( "COLUMN_ALREADY_EXISTS", + # Databricks Delta surfaces an ``ADD COLUMN`` that targets an + # already-present column as ``FIELD_ALREADY_EXISTS`` (singular). + # Older versions used the plural ``FIELDS_ALREADY_EXISTS``; + # keep both to defend against runtime wording drift across + # workspaces / DBR versions. + "FIELD_ALREADY_EXISTS", "FIELDS_ALREADY_EXISTS", "TABLE_OR_VIEW_ALREADY_EXISTS", "TABLE_ALREADY_EXISTS", "already has liquid clustering defined", + # CHECK constraints are added via ``ALTER TABLE … ADD CONSTRAINT`` + # in a separate statement after CREATE TABLE (Delta only allows + # PK/FK inline). When a previous migration attempt got past + # CREATE TABLE but failed before recording the version, the + # next run sees the constraint and emits this error — safe to + # swallow because the desired state is already in place. The + # second fragment guards against future error-message tweaks + # since Databricks has used both the SQLSTATE-prefixed code and + # plain English wording at different times. + "DELTA_CONSTRAINT_ALREADY_EXISTS", + "constraint already exists", ) def _apply(self, migration: Migration) -> None: - formatted = migration.sql_template.format(catalog=self._catalog, schema=self._schema) + """Apply *migration* statement-by-statement against the warehouse. + + Unlike the Postgres runner, this method **cannot** wrap a + multi-statement migration + the ``dq_migrations`` version-row + ``INSERT`` in a single transaction: the Databricks Statement + Execution API auto-commits per call and there is no ``BEGIN``/ + ``COMMIT`` primitive for DDL on a SQL warehouse. A mid-migration + failure therefore leaves earlier statements committed but the + version row uninserted, and the next ``run_all`` re-runs the + whole migration from the beginning. + + Recovery is safe **only** because every DDL kind we currently + emit raises an "already exists"-style error that + ``_IDEMPOTENT_ERROR_FRAGMENTS`` swallows, letting the rerun + skip past committed statements and reach the + previously-failed one. See the module docstring's + "Atomicity model" section for the full contract that any new + migration author has to satisfy. + + Templates are split on ``;`` so each statement reaches the API + as a single call (the API rejects compound statements). The + ``_validate_template_split_safe`` check at module import time + guarantees no template embeds a ``;`` inside a string literal, + which would otherwise silently corrupt the split. + + ``{catalog}`` / ``{schema}`` placeholders receive the + :meth:`SqlExecutor.q`-quoted form so identifiers containing + characters outside ``[A-Za-z0-9_]`` (most commonly a hyphen + in a Databricks catalog name) stay parseable. Placeholders + are only used in object-name positions in the shipped + templates — substituting the quoted form would be incorrect + inside a string literal, but :func:`_validate_template_split_safe` + rules out the dangerous shape ahead of time. + """ + formatted = migration.sql_template.format(catalog=self._catalog_q, schema=self._schema_q) for stmt in formatted.split(";"): stmt = stmt.strip() if stmt: diff --git a/app/src/databricks_labs_dqx_app/backend/migrations/postgres.py b/app/src/databricks_labs_dqx_app/backend/migrations/postgres.py new file mode 100644 index 000000000..9a3ff8f04 --- /dev/null +++ b/app/src/databricks_labs_dqx_app/backend/migrations/postgres.py @@ -0,0 +1,375 @@ +"""Lakebase (Postgres) migration runner for DQX Studio OLTP tables. + +The OLTP subset of the schema (rules catalog, app settings, RBAC, +comments, schedule configs, scheduler bookkeeping) lives in Lakebase +Postgres when ``conf.lakebase_enabled`` is true. Append-mostly +analytical tables (``dq_validation_runs``, ``dq_profiling_results``, +``dq_metrics``, ``dq_quarantine_records``) stay in Delta because they +are written by the Spark task runner. + +The runner mirrors :class:`backend.migrations.MigrationRunner` so the +operational story (versioned, idempotent, recorded in a +``dq_migrations`` meta table) is identical regardless of backend. + +Schema mapping highlights +------------------------- +- Delta ``VARIANT`` (``check``, future blob columns) → Postgres + ``JSONB``. Both store JSON natively with predicate pushdown; JSONB + also supports GIN indexes for low-latency lookups. +- Delta ``TIMESTAMP`` → Postgres ``TIMESTAMPTZ``. Postgres' + timezone-naive ``TIMESTAMP`` would silently drop offsets and break + cross-region debugging. +- Delta ``BIGINT`` / ``INT`` / ``DOUBLE`` → ``BIGINT`` / ``INT`` / + ``DOUBLE PRECISION`` respectively. +- Delta ``CLUSTER BY`` is not exposed; Postgres uses indexes + declaratively where the access pattern justifies them. Each table + gets the small set of indexes the FastAPI services actually need. + +Adding a new migration +---------------------- +Append a new :class:`PgMigration` entry with the next monotonically +increasing version number. Postgres supports ``ALTER TABLE ... ADD +COLUMN IF NOT EXISTS`` natively so re-running is safe out of the box. +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from contextlib import AbstractContextManager +from dataclasses import dataclass +from typing import Any, Protocol + +# Import directly from the psycopg-free helpers module — NOT from +# :mod:`..pg_executor`, which would transitively pull in :mod:`psycopg` +# at module load. The migration runner itself only needs the trust- +# boundary wrappers; deferring the heavy psycopg import to the +# executor's own users keeps this module importable in Delta-only +# environments (e.g. the dqx-library integration test rig). +from ..pg_cursor_helpers import run_parameterized_sql, run_trusted_sql +from ..sql_executor import OltpExecutorProtocol + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Migration protocol — extends the shared :class:`OltpExecutorProtocol` +# rather than re-declaring ``schema`` / ``q`` / ``execute`` / ``query`` +# fresh. Keeps the executor surface single-sourced so a future addition +# to the OLTP contract (e.g. a new ``upsert_*`` variant) doesn't drift +# between callers. +# +# The migration runner needs ONE thing beyond the OLTP contract: a raw +# :meth:`connection` context manager so it can wrap multiple DDL +# statements + the meta-table INSERT in a single Postgres transaction +# (the auto-committing :meth:`OltpExecutorProtocol.execute` would split +# them across separate transactions and leave the meta row out of sync +# on partial failure). Only :class:`PgExecutor` exposes that today — +# Delta runs through its own :class:`MigrationRunner` because Spark SQL +# DDL is not transactional anyway. +# --------------------------------------------------------------------------- + + +class _Executor(OltpExecutorProtocol, Protocol): + """OLTP executor with transactional ``connection()`` access. + + Structurally satisfied only by :class:`backend.pg_executor.PgExecutor` + (and matching test doubles). The return type of :meth:`connection` + is typed loosely as ``AbstractContextManager[Any]`` so this Protocol + stays free of a psycopg dependency — Delta-only deployments can + import this module without dragging psycopg into the graph. + """ + + def connection(self) -> AbstractContextManager[Any]: ... + + +@dataclass(frozen=True) +class PgMigration: + version: int + description: str + sql: str + + +# ``{schema}`` is substituted at apply-time so a single migration list +# can be re-targeted for tests/dev/prod without copy-paste. +_S = "{schema}" + + +PG_MIGRATIONS: list[PgMigration] = [ + PgMigration( + version=1, + description="Lakebase OLTP baseline (app_settings, rules, role mappings, comments, schedules)", + sql=( + # ---------------------------------------------------------- + # dq_app_settings — single-row-per-key KV store. + # ---------------------------------------------------------- + f"CREATE TABLE IF NOT EXISTS {_S}.dq_app_settings (" + " setting_key TEXT PRIMARY KEY," + " setting_value TEXT," + " updated_at TIMESTAMPTZ," + " updated_by TEXT" + ");" + # ---------------------------------------------------------- + # dq_quality_rules — active rule catalog. + # ---------------------------------------------------------- + f"CREATE TABLE IF NOT EXISTS {_S}.dq_quality_rules (" + " rule_id TEXT PRIMARY KEY," + " table_fqn TEXT NOT NULL," + ' "check" JSONB NOT NULL,' + " version INTEGER NOT NULL," + " status TEXT NOT NULL," + " source TEXT NOT NULL," + " created_by TEXT," + " created_at TIMESTAMPTZ," + " updated_by TEXT," + " updated_at TIMESTAMPTZ," + " CONSTRAINT chk_dq_quality_rules_status " + " CHECK (status IN ('draft','pending_approval','approved','rejected'))," + " CONSTRAINT chk_dq_quality_rules_source " + " CHECK (source IN ('ui','sql','profiler','import','ai'))" + ");" + # Two read-paths dominate: by table_fqn (rules-list page) and + # by status filter (review queue). One composite index covers + # both since Postgres can use a leading-column-only scan. + f"CREATE INDEX IF NOT EXISTS idx_dq_quality_rules_table_status " + f" ON {_S}.dq_quality_rules (table_fqn, status);" + # ---------------------------------------------------------- + # dq_quality_rules_history — append-only audit trail. + # ---------------------------------------------------------- + f"CREATE TABLE IF NOT EXISTS {_S}.dq_quality_rules_history (" + " history_id BIGSERIAL PRIMARY KEY," + " rule_id TEXT," + " table_fqn TEXT NOT NULL," + ' "check" JSONB,' + " version INTEGER," + " source TEXT," + " action TEXT NOT NULL," + " prev_status TEXT," + " new_status TEXT," + " changed_by TEXT," + " changed_at TIMESTAMPTZ" + ");" + f"CREATE INDEX IF NOT EXISTS idx_dq_quality_rules_history_rule_changed_at " + f" ON {_S}.dq_quality_rules_history (rule_id, changed_at DESC);" + f"CREATE INDEX IF NOT EXISTS idx_dq_quality_rules_history_table_changed_at " + f" ON {_S}.dq_quality_rules_history (table_fqn, changed_at DESC);" + # ---------------------------------------------------------- + # dq_role_mappings — RBAC. + # ---------------------------------------------------------- + f"CREATE TABLE IF NOT EXISTS {_S}.dq_role_mappings (" + " role TEXT NOT NULL," + " group_name TEXT NOT NULL," + " created_by TEXT," + " created_at TIMESTAMPTZ," + " updated_by TEXT," + " updated_at TIMESTAMPTZ," + " PRIMARY KEY (role, group_name)," + " CONSTRAINT chk_dq_role_mappings_role " + " CHECK (role IN ('admin','rule_approver','rule_author','viewer','runner'))" + ");" + # ---------------------------------------------------------- + # dq_comments — per-entity comment threads. + # ---------------------------------------------------------- + f"CREATE TABLE IF NOT EXISTS {_S}.dq_comments (" + " comment_id TEXT PRIMARY KEY," + " entity_type TEXT NOT NULL," + " entity_id TEXT NOT NULL," + " user_email TEXT NOT NULL," + " comment TEXT NOT NULL," + " created_at TIMESTAMPTZ," + " CONSTRAINT chk_dq_comments_entity_type " + " CHECK (entity_type IN ('run','rule'))" + ");" + f"CREATE INDEX IF NOT EXISTS idx_dq_comments_entity_created_at " + f" ON {_S}.dq_comments (entity_type, entity_id, created_at DESC);" + # ---------------------------------------------------------- + # dq_schedule_runs — last/next-run pointer per schedule. + # ---------------------------------------------------------- + f"CREATE TABLE IF NOT EXISTS {_S}.dq_schedule_runs (" + " schedule_name TEXT PRIMARY KEY," + " last_run_at TIMESTAMPTZ," + " next_run_at TIMESTAMPTZ," + " last_run_id TEXT," + " status TEXT," + " updated_at TIMESTAMPTZ," + " CONSTRAINT chk_dq_schedule_runs_status " + " CHECK (status IS NULL OR status IN " + " ('pending','success','partial_failure','failed'))" + ");" + # The scheduler loop polls "next_run_at <= now() AND status + # IS NOT 'pending'" every tick; a btree index on + # next_run_at keeps that scan cheap as the schedule list + # grows. + f"CREATE INDEX IF NOT EXISTS idx_dq_schedule_runs_next_run_at " + f" ON {_S}.dq_schedule_runs (next_run_at);" + # ---------------------------------------------------------- + # dq_schedule_configs — live config per schedule. + # ---------------------------------------------------------- + f"CREATE TABLE IF NOT EXISTS {_S}.dq_schedule_configs (" + " schedule_name TEXT PRIMARY KEY," + " config_json TEXT NOT NULL," + " version INTEGER NOT NULL," + " created_by TEXT," + " created_at TIMESTAMPTZ," + " updated_by TEXT," + " updated_at TIMESTAMPTZ" + ");" + # ---------------------------------------------------------- + # dq_schedule_configs_history — append-only audit trail. + # ---------------------------------------------------------- + f"CREATE TABLE IF NOT EXISTS {_S}.dq_schedule_configs_history (" + " history_id BIGSERIAL PRIMARY KEY," + " schedule_name TEXT NOT NULL," + " config_json TEXT," + " version INTEGER," + " action TEXT NOT NULL," + " changed_by TEXT," + " changed_at TIMESTAMPTZ" + ");" + f"CREATE INDEX IF NOT EXISTS idx_dq_schedule_configs_history_schedule_changed_at " + f" ON {_S}.dq_schedule_configs_history (schedule_name, changed_at DESC);" + ), + ), +] + + +_META_TABLE = f"{_S}.dq_migrations" + + +class PgMigrationRunner: + """Applies pending Postgres migrations to the Lakebase OLTP schema. + + Mirrors :class:`backend.migrations.MigrationRunner` so the lifespan + handler can call ``run_all`` against either backend with the same + contract. The schema is created (``CREATE SCHEMA IF NOT EXISTS``) + so the very first deploy doesn't need a separate bootstrap step on + the Postgres side — Lakebase is provisioned bare with just a + ``public`` schema by default. + """ + + def __init__( + self, + executor: _Executor, + *, + migrations: Sequence[PgMigration] | None = None, + ) -> None: + """Construct the runner. + + Parameters + ---------- + executor: + Postgres executor (PgExecutor in production, a mock in unit + tests) satisfying :class:`_Executor`. + migrations: + Optional migration list. Defaults to the module-level + :data:`PG_MIGRATIONS` catalogue that production deploys + use. Exposed as a constructor seam so tests can inject a + deterministic fake list without monkey-patching the + module-level constant — monkey-patching couples tests to + an import path and creates surprise interactions when two + tests run in the same process. + """ + self._exec = executor + self._schema = executor.schema + # Pre-compute the executor-quoted schema once and feed it to + # every DDL/DML site (templates, meta-table reference, + # ``CREATE SCHEMA``). The Delta runner's + # :class:`MigrationRunner` does the same with backticks; this + # keeps the two backends symmetric and removes the previous + # foot-gun where ``_ensure_schema`` was the only quoted site + # while ``_meta_table`` and the migration templates raw- + # interpolated the schema name. + self._schema_q = executor.q(self._schema) + self._meta_table = _META_TABLE.format(schema=self._schema_q) + # Take a defensive tuple snapshot so callers can't mutate the + # runner's view of the catalogue after construction. + self._migrations: tuple[PgMigration, ...] = tuple(migrations if migrations is not None else PG_MIGRATIONS) + + def run_all(self) -> int: + self._ensure_schema() + self._ensure_meta_table() + applied = self._applied_versions() + count = 0 + for migration in self._migrations: + if migration.version in applied: + logger.debug( + "Postgres migration v%d (%s) already applied", + migration.version, + migration.description, + ) + continue + logger.info( + "Applying Postgres migration v%d: %s", + migration.version, + migration.description, + ) + self._apply(migration) + count += 1 + return count + + def _ensure_schema(self) -> None: + # ``self._schema_q`` is already the quoted form (see __init__). + # Calling :meth:`_Executor.q` rather than hardcoding the ANSI + # double-quote keeps the dialect contract on the executor. + self._exec.execute(f"CREATE SCHEMA IF NOT EXISTS {self._schema_q}") + + def _ensure_meta_table(self) -> None: + # ``self._meta_table`` already embeds the quoted schema. + sql = ( + f"CREATE TABLE IF NOT EXISTS {self._meta_table} (" + " version INTEGER PRIMARY KEY," + " description TEXT NOT NULL," + " applied_at TIMESTAMPTZ NOT NULL" + ")" + ) + self._exec.execute(sql) + + def _applied_versions(self) -> set[int]: + rows = self._exec.query(f"SELECT version FROM {self._meta_table} ORDER BY version") + return {int(row[0]) for row in rows} + + def _apply(self, migration: PgMigration) -> None: + """Apply *migration* atomically. + + Postgres DDL is transactional (CREATE TABLE / INDEX inside + BEGIN/COMMIT all roll back together on error — modulo + ``CREATE INDEX CONCURRENTLY`` which we don't use). We run every + DDL statement in the migration **and** the ``dq_migrations`` + INSERT inside a single transaction so a half-applied migration + can never end up with committed DDL but no version row. If any + statement fails the whole migration rolls back and the next run + retries cleanly from the beginning. + + Statements are still split on ``;`` and executed one at a time + through a single cursor so an error message pinpoints the exact + failing DDL rather than a position inside a multi-kilobyte + compound string. + """ + formatted = migration.sql.format(schema=self._schema_q) + with self._exec.connection() as conn: + with conn.cursor() as cur: + for stmt in formatted.split(";"): + stmt = stmt.strip() + if stmt: + run_trusted_sql(cur, stmt) + # Same transaction: either the whole migration lands + # *and* its version row is recorded, or nothing does. + # + # The template here is trusted (``self._meta_table`` is + # built via the executor's ``q()`` quoter at __init__), + # but ``migration.version`` / ``migration.description`` + # are runtime values from a Python object — they MUST + # NOT be f-string-interpolated even with manual escaping. + # Routing them through ``run_parameterized_sql`` hands + # the binding to psycopg / libpq, which keeps the values + # out of the SQL string entirely. See the docstring on + # ``backend.pg_executor.run_trusted_sql`` for the full + # contract. + run_parameterized_sql( + cur, + f"INSERT INTO {self._meta_table} (version, description, applied_at) " + "VALUES (%s, %s, CURRENT_TIMESTAMP)", + (migration.version, migration.description), + ) + conn.commit() + logger.info("Postgres migration v%d applied", migration.version) diff --git a/app/src/databricks_labs_dqx_app/backend/models.py b/app/src/databricks_labs_dqx_app/backend/models.py index a0c176bf1..8b4d02920 100644 --- a/app/src/databricks_labs_dqx_app/backend/models.py +++ b/app/src/databricks_labs_dqx_app/backend/models.py @@ -141,7 +141,11 @@ class DryRunSubmitOut(BaseModel): class DryRunOut(BaseModel): total_rows: int valid_rows: int + # ``invalid_rows`` is kept for backwards compatibility but is no longer + # the primary count surfaced in the UI — see ``error_rows`` below. invalid_rows: int + error_rows: int = 0 + warning_rows: int = 0 error_summary: list[dict[str, Any]] sample_invalid: list[dict[str, Any]] @@ -261,6 +265,10 @@ class DryRunResultsOut(BaseModel): total_rows: int | None = None valid_rows: int | None = None invalid_rows: int | None = None + # ``error_rows`` / ``warning_rows`` are the authoritative DQX observer + # counts; ``invalid_rows`` is kept for backwards compatibility only. + error_rows: int | None = None + warning_rows: int | None = None error_summary: list[dict[str, Any]] = Field(default_factory=list) sample_invalid: list[dict[str, Any]] = Field(default_factory=list) @@ -277,6 +285,8 @@ class ValidationRunSummaryOut(BaseModel): run_type: str | None = None valid_rows: int | None = None invalid_rows: int | None = None + error_rows: int | None = None + warning_rows: int | None = None created_at: str | None = None error_message: str | None = None checks: list[dict[str, Any]] = Field(default_factory=list) @@ -328,6 +338,7 @@ class QuarantineRecordOut(BaseModel): requesting_user: str | None = None row_data: dict[str, Any] | None = None errors: list[Any] | None = None + warnings: list[Any] | None = None created_at: str | None = None diff --git a/app/src/databricks_labs_dqx_app/backend/pg_cursor_helpers.py b/app/src/databricks_labs_dqx_app/backend/pg_cursor_helpers.py new file mode 100644 index 000000000..95cd26f7d --- /dev/null +++ b/app/src/databricks_labs_dqx_app/backend/pg_cursor_helpers.py @@ -0,0 +1,148 @@ +"""Trust-boundary wrappers around :meth:`psycopg.Cursor.execute`. + +Both helpers in this module exist for **one** reason: psycopg's PEP-675 +stubs require :meth:`Cursor.execute`'s ``query`` argument to be a +:class:`typing.LiteralString`. Code that legitimately composes SQL +(validated identifiers, vendored DDL templates, parameter-binding +placeholders) cannot satisfy that constraint without a ``cast``. +Centralising the cast in named helpers — rather than sprinkling +``# pyright: ignore[reportCallIssue, reportArgumentType]`` at every +call site — is the documented project policy (AGENTS.md rule #6 + the +``TestPsycopgExecutePyrightIgnorePolicy`` test). + +Why this lives in a separate module from :mod:`backend.pg_executor` +------------------------------------------------------------------ +Per the reviewer note on the original ``app.py`` lazy-import comment: +``migrations.postgres`` only needs these helpers, but importing them +from :mod:`backend.pg_executor` transitively imports :mod:`psycopg` +(which the Delta-only test environment does not install). That made +the lazy-import dance in :func:`app.lifespan` necessary even though +:class:`PgMigrationRunner` itself doesn't touch psycopg directly. + +This module imports ``Cursor`` **only under** :data:`typing.TYPE_CHECKING`, +so :mod:`psycopg` is never loaded at runtime. The actual ``execute`` +call is duck-typed — any object with an ``execute(sql)`` / +``execute(sql, params)`` method works, which keeps mocks / fakes in +tests trivial. + +Trust contract — read the per-function docstrings +------------------------------------------------- +:func:`run_trusted_sql` is for **static** SQL. :func:`run_parameterized_sql` +is for **trusted templates + runtime values**, where values flow through +psycopg's native parameter binder so they are NEVER spliced into the SQL +string. Manual ``str.replace("'", "''")`` "escaping" is explicitly out +of contract for both helpers — see :func:`run_trusted_sql` for the full +rationale. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, LiteralString, cast + +if TYPE_CHECKING: + # Type-only import keeps this module psycopg-free at runtime so + # ``migrations.postgres`` (and any other caller that only needs the + # trust-boundary helpers) can be imported in a Delta-only environment + # without installing psycopg. The runtime call below is duck-typed. + from psycopg import Cursor + + +def run_trusted_sql(cur: Cursor[Any], sql: str) -> None: + """Execute a backend-composed STATIC SQL string against a psycopg cursor. + + psycopg's stubs (PEP 675) require :meth:`Cursor.execute`'s + ``query`` argument to be a :class:`typing.LiteralString` — a + string the type-checker can statically prove is not a runtime + concatenation of untrusted input. That's a deliberate type-level + defence against SQL injection; passing a plain ``str`` makes + basedpyright raise ``reportCallIssue`` + ``reportArgumentType``. + + Casting to ``LiteralString`` in this *one* place concentrates + that trust boundary into a single auditable wrapper instead of + sprinkling ``# pyright: ignore[reportCallIssue, + reportArgumentType]`` at every call site (which AGENTS.md rule + #6 explicitly prohibits). + + Trust contract (narrow on purpose) + ---------------------------------- + Every string reaching this helper MUST be built EXCLUSIVELY from: + + 1. **Compile-time string literals** vendored alongside the code + that consumes them — DDL templates, fixed SELECT/UPDATE + skeletons, schema-bootstrap statements, etc. + 2. **Identifiers** (catalog / schema / table / column names) + validated by :func:`backend.sql_utils.validate_fqn` and + quoted via :meth:`PgExecutor.q` / + :func:`backend.sql_utils.quote_fqn`. The quoting protects + reserved-word identifiers; the validation rejects anything + that doesn't look like an identifier in the first place. + 3. **Compile-time value literals** rendered by + :func:`sql_executor._render_value` for the well-known set of + safe sentinel values (``RawSql("current_timestamp()")``, + integer / boolean constants used as flags, etc.). + + Explicitly OUT of contract + -------------------------- + Runtime VALUE literals — anything that originates from a Python + object's mutable state, a config file, a function argument, an + `f"...{value}..."` interpolation, or a manual + ``str.replace("'", "''")`` "escape" — must NOT reach this + helper. Manual escaping is not a defence: it silently rubber- + stamps every future change that drops a new f-string into the + same template. The ``cast(LiteralString, ...)`` performed here + is a type-checker silence, not a runtime sanitiser. + + For "trusted template + runtime value" callers, use the sibling + helper :func:`run_parameterized_sql` (which routes values + through psycopg's native ``cur.execute(sql, params)`` binding), + or call ``cur.execute(template, params)`` directly. Either path + keeps the value out of the SQL string entirely so the LiteralString + cast on the template alone is sound. + """ + _ = cur.execute(cast(LiteralString, sql)) + + +def run_parameterized_sql(cur: Cursor[Any], sql: str, params: Sequence[Any]) -> None: + """Execute a trusted SQL TEMPLATE with psycopg-bound runtime values. + + Sibling to :func:`run_trusted_sql` for the common pattern where + the SQL structure (identifiers, clauses, the ``%s`` placeholders + themselves) is trusted but the values bound to those placeholders + originate at runtime — from migration metadata, an admin's + config row, a method argument, anything that isn't a compile-time + literal. + + The *template* is subject to the same trust contract as + :func:`run_trusted_sql`: it MUST be composed entirely from + compile-time literals plus identifiers routed through + :meth:`PgExecutor.q` / :func:`quote_fqn`. The contract is + enforced by the ``cast(LiteralString, sql)`` here, exactly like + in :func:`run_trusted_sql`, so the same auditing discipline + applies to the template surface. + + The *values* in ``params`` are passed through psycopg's native + parameter binding (``cur.execute(template, params)``), which + forwards them to libpq as binary parameters — never spliced into + the SQL string. That makes value-side injection structurally + impossible regardless of the value's provenance, which is the + point: ``run_trusted_sql`` cannot offer that guarantee for + f-string-interpolated values because it has no access to the + distinction between "template" and "value" once the caller has + concatenated them. + + Example:: + + run_parameterized_sql( + cur, + f"INSERT INTO {self._meta_table} (version, description) " + "VALUES (%s, %s)", + (migration.version, migration.description), + ) + + The ``self._meta_table`` is an identifier built via + :meth:`PgExecutor.q` so it satisfies the template contract; the + integer and description go through psycopg's binder so neither + needs manual escaping. + """ + _ = cur.execute(cast(LiteralString, sql), params) diff --git a/app/src/databricks_labs_dqx_app/backend/pg_executor.py b/app/src/databricks_labs_dqx_app/backend/pg_executor.py new file mode 100644 index 000000000..056af9a4a --- /dev/null +++ b/app/src/databricks_labs_dqx_app/backend/pg_executor.py @@ -0,0 +1,770 @@ +"""Lakebase Postgres executor. + +Mirrors the public surface of :class:`SqlExecutor` so that services can +target either Delta (via SQL warehouse) or Lakebase (via psycopg) with +the same call signatures. + +Why a separate class instead of a generic SQLAlchemy abstraction? +We need extremely tight control over: + +- **OAuth token refresh.** Lakebase tokens expire after one hour. A + background daemon thread refreshes the password every + ``DQX_LAKEBASE_TOKEN_REFRESH_MINUTES`` minutes and the connection + pool's ``configure`` callback hands new connections the latest + password. Existing connections in the pool are recycled when they + exceed ``max_lifetime``. +- **String-typed result rows.** The legacy :class:`SqlExecutor` returns + ``list[list[str]]`` because the Statement Execution API serialises + cells via ``Format.JSON_ARRAY``. Service code expects to call + ``json.loads`` on JSON columns and to receive ISO-string timestamps. + We coerce psycopg's natively-typed results to that shape so existing + services work unchanged. +- **Dialect helpers.** :meth:`q` and :meth:`json_literal_expr` produce + Postgres-flavoured identifiers/literals so portable service SQL + doesn't need a dialect branch. +""" + +from __future__ import annotations + +import json +import logging +import os +import random +import threading +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import date, datetime, timezone +from decimal import Decimal +from typing import Any + +from databricks.sdk import WorkspaceClient +from psycopg import Connection, Cursor +from psycopg_pool import ConnectionPool + +# ``run_trusted_sql`` / ``run_parameterized_sql`` live in the +# psycopg-free :mod:`backend.pg_cursor_helpers` module so callers +# that ONLY need the trust-boundary wrappers (notably +# :mod:`backend.migrations.postgres`) can be imported in environments +# that don't have psycopg installed. We re-export them here so existing +# imports of the form ``from backend.pg_executor import run_trusted_sql`` +# keep working — the executor is the "real" psycopg consumer and the +# natural place to look for them. +from databricks_labs_dqx_app.backend.pg_cursor_helpers import ( + run_parameterized_sql, + run_trusted_sql, +) +from databricks_labs_dqx_app.backend.sql_executor import RawSql, _render_value +from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string + +logger = logging.getLogger(__name__) + +# Re-exports so ``from backend.pg_executor import run_trusted_sql`` keeps +# working after the helpers moved to ``backend.pg_cursor_helpers``. +__all__ = [ + "PgExecutor", + "build_pg_executor", + "run_parameterized_sql", + "run_trusted_sql", +] + + +class _TokenHolder: + """Thread-safe container for the rotating Lakebase OAuth token. + + The connection pool's ``configure`` callback (called when each new + physical connection is opened) reads :attr:`token` so the *next* + connection always picks up the latest credential. Existing + connections continue to work — Postgres only validates the + password during the SCRAM handshake, not on every query. + """ + + def __init__(self, token: str) -> None: + self._token = token + self._lock = threading.Lock() + + @property + def token(self) -> str: + with self._lock: + return self._token + + @token.setter + def token(self, value: str) -> None: + with self._lock: + self._token = value + + +def _generate_token(ws: WorkspaceClient, instance_name: str) -> str: + """Generate a fresh Lakebase OAuth token (1-hour TTL).""" + cred = ws.database.generate_database_credential( + request_id=str(uuid.uuid4()), + instance_names=[instance_name], + ) + if not cred.token: + raise RuntimeError(f"Lakebase credential response had no token (instance={instance_name})") + return cred.token + + +def _to_text(value: Any) -> str | None: + """Coerce a psycopg cell value to Delta-compatible string output. + + - ``None`` stays ``None``. + - ``dict``/``list`` → compact JSON (matching Delta's ``to_json``). + - ``datetime``/``date`` → ISO 8601. + - ``bool`` → ``"true"``/``"false"`` (lowercase, JSON-style). + - ``Decimal``/``int``/``float`` → ``str(value)``. + - Everything else → ``str(value)``. + """ + if value is None: + return None + if isinstance(value, (dict, list)): + return json.dumps(value, separators=(",", ":")) + if isinstance(value, (datetime, date)): + return value.isoformat() + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, Decimal): + return str(value) + if isinstance(value, (bytes, bytearray, memoryview)): + # Binary columns are not used in any OLTP table today; surface + # them as hex if a future migration introduces one. + return bytes(value).hex() + return str(value) + + +def _pg_render_value(value: Any) -> str: + """Postgres-flavoured literal renderer for :meth:`PgExecutor.upsert`. + + Behaves identically to :func:`backend.sql_executor._render_value` + except that: + + - :class:`RawSql("current_timestamp()")` is rewritten to + ``CURRENT_TIMESTAMP`` because Postgres rejects the parenthesised + Spark SQL form. Other ``RawSql`` payloads pass through verbatim + so callers can still inject Postgres-specific helpers like + ``now()`` or ``::jsonb`` casts. + - ``bool`` renders as ``TRUE``/``FALSE`` which Postgres accepts. + """ + if isinstance(value, RawSql): + expr = value.expr.strip() + # Common Spark idiom that doesn't parse in Postgres — translate. + if expr.lower() in {"current_timestamp()", "now()"}: + return "CURRENT_TIMESTAMP" + return value.expr + return _render_value(value) + + +class PgExecutor: + """Drop-in :class:`SqlExecutor` replacement backed by Lakebase Postgres. + + Constructed at app startup when ``conf.lakebase_enabled`` is true; + the lifespan handler kicks off the token-refresh thread and runs + the Postgres migrations once before traffic arrives. + """ + + dialect: str = "postgres" + + def __init__( + self, + *, + ws: WorkspaceClient, + instance_name: str, + database: str, + schema: str, + username: str, + host: str, + port: int = 5432, + token_refresh_minutes: int = 50, + token_refresh_retry_seconds: int = 10, + token_refresh_retry_jitter: float = 0.3, + token_refresh_max_failures: int = 12, + pool_min_size: int = 1, + pool_max_size: int = 10, + ) -> None: + self._ws = ws + self._instance_name = instance_name + self._database = database + self._schema = schema + self._username = username + self._host = host + self._port = port + self._token_refresh_seconds = max(60, token_refresh_minutes * 60) + # Retry tuning: clamp to sensible floors so a mis-configured + # env var can't degenerate into a busy-loop or an unbounded + # failure window. ``retry_seconds`` floors at 1s (anything + # below that is effectively a tight loop against the SDK); + # ``jitter`` clamps to ``[0, 1]`` so the spread can't go + # negative or larger than ±100%; ``max_failures`` floors at + # 1 so we always escalate on at least one failure rather than + # spinning forever. + self._token_refresh_retry_seconds = max(1, int(token_refresh_retry_seconds)) + self._token_refresh_retry_jitter = max(0.0, min(1.0, float(token_refresh_retry_jitter))) + self._token_refresh_max_failures = max(1, int(token_refresh_max_failures)) + + # Bootstrap the first token before the pool starts so the very + # first connection has valid credentials. The bootstrap mint + # counts as a successful refresh — the metric / counter start + # in the "healthy" state, not a transient "never refreshed" + # one that would confuse a health endpoint at t=0. + self._token_holder = _TokenHolder(_generate_token(ws, instance_name)) + self._last_successful_refresh_at: datetime | None = datetime.now(timezone.utc) + self._consecutive_refresh_failures: int = 0 + + # ``kwargs`` is the dict the pool hands to ``Connection.connect`` + # every time it opens a new physical connection. Mutating + # ``password`` on token refresh means subsequent connects pick + # up the fresh credential without restarting the pool. The + # ``options`` flag sets the Postgres ``search_path`` so + # unqualified table references resolve to the app schema. + self._connect_kwargs: dict[str, Any] = { + "host": host, + "port": port, + "dbname": database, + "user": username, + "password": self._token_holder.token, + "sslmode": "require", + "options": f"-c search_path={schema}", + } + + # ``max_lifetime`` recycles connections every 50 minutes which + # ensures we never hand out a connection authenticated with a + # near-expired token. ``check`` runs ``SELECT 1`` on idle pool + # members so a server-side disconnect doesn't poison the pool. + self._pool: ConnectionPool = ConnectionPool( + conninfo="", + min_size=pool_min_size, + max_size=pool_max_size, + max_lifetime=self._token_refresh_seconds, + check=ConnectionPool.check_connection, + open=False, # opened explicitly below so failures surface eagerly + kwargs=self._connect_kwargs, + timeout=30.0, + name="dqx-lakebase", + ) + self._pool.open(wait=True, timeout=30.0) + logger.info( + "Lakebase connection pool open (host=%s db=%s schema=%s user=%s)", + host, + database, + schema, + username, + ) + + self._stop = threading.Event() + self._refresher = threading.Thread( + target=self._token_refresh_loop, + name="dqx-lakebase-token-refresh", + daemon=True, + ) + self._refresher.start() + + # ------------------------------------------------------------------ + # Public API mirrors SqlExecutor + # ------------------------------------------------------------------ + + @property + def warehouse_id(self) -> str: + # Kept for type compatibility with ``SqlExecutor``; Lakebase + # has no warehouse concept. + return "" + + @property + def catalog(self) -> str: + # Postgres has no Unity Catalog; return the database name so + # callers that build fully-qualified identifiers still get a + # 3-part name (``database.schema.table``) on Postgres. + return self._database + + @property + def schema(self) -> str: + return self._schema + + @property + def database(self) -> str: + return self._database + + # ------------------------------------------------------------------ + # Token-refresh observability + # ------------------------------------------------------------------ + # These two properties surface the refresh daemon's state so + # operators can wire a health endpoint / metric exporter to it + # WITHOUT poking at private attributes. The refresh loop updates + # them under no extra lock — they are coarse-grained gauges, not + # transactional state, so a single torn read at the boundary is + # harmless (worst case: a health check sees a stale value for + # one scrape interval and corrects on the next). + + @property + def last_successful_refresh_at(self) -> datetime | None: + """Wall-clock time of the last successful token refresh, UTC. + + Set to ``datetime.now(UTC)`` in :meth:`__init__` because the + bootstrap token mint IS a successful refresh from the pool's + perspective. Operators monitoring this metric should alert + if ``now() - last_successful_refresh_at`` exceeds + ``token_refresh_minutes`` by a meaningful margin (e.g. 2x) + — that is the leading indicator of the silent-pool-drain + failure mode this metric exists to prevent. + """ + return self._last_successful_refresh_at + + @property + def consecutive_refresh_failures(self) -> int: + """Failed-refresh streak since the last success. + + Reset to 0 on every successful refresh. Reaching + ``token_refresh_max_failures`` triggers the escalation path + in :meth:`_token_refresh_loop` (process exit so the + supervisor restarts the worker). + """ + return self._consecutive_refresh_failures + + def fqn(self, table: str) -> str: + """Return the schema-qualified path for *table*. + + Postgres only has one catalog per connection so we return + ``schema.table``; :meth:`SqlExecutor.fqn` returns three parts. + See :meth:`SqlExecutor.fqn` for the parity contract. + """ + return f"{self._schema}.{table}" + + def q(self, identifier: str) -> str: + """Quote a Postgres identifier (ANSI double quotes, doubled internal ``"``).""" + return '"' + identifier.replace('"', '""') + '"' + + def json_literal_expr(self, json_str: str) -> str: + """Return a Postgres expression that yields a JSONB value for *json_str*.""" + return f"'{escape_sql_string(json_str)}'::jsonb" + + def ts_text(self, col: str) -> str: + """Project a timestamp column as a string. + + Postgres TIMESTAMPTZ values are converted to ISO strings by + :func:`_to_text` when the row leaves the cursor, so we just + select the column verbatim and let the row-level coercion do + the work. This keeps service SQL portable: callers always + write ``executor.ts_text('created_at')`` regardless of dialect. + """ + return col + + @contextmanager + def connection(self) -> Iterator[Connection]: + """Yield a pooled connection for multi-statement transactional work. + + Use this when several statements MUST be atomic — e.g. DDL + a + bookkeeping INSERT that record their joint application. The + caller is responsible for ``conn.commit()``; if the ``with`` + block exits without committing (exception or otherwise) the + connection is rolled back and returned to the pool by psycopg- + pool. Single one-shot statements should keep using + :meth:`execute`, which commits per call. + + **Statement timeouts are the caller's responsibility here.** + Unlike :meth:`execute` / :meth:`query` / :meth:`query_dicts`, + this method does NOT auto-emit ``SET LOCAL statement_timeout`` + because callers in this branch are running multi-statement + units of work — e.g. the migration runner applying a series + of DDLs whose appropriate timeout depends on the work being + done (a 30s app query vs. a 10-minute index rebuild). Callers + that want a guard should issue ``SET LOCAL statement_timeout + = `` themselves inside the transaction (routed through + :func:`run_trusted_sql` for the LiteralString cast). + """ + with self._pool.connection() as conn: + yield conn + + @staticmethod + def _apply_statement_timeout(cur: Cursor[Any], timeout_seconds: int) -> None: + """Cap the next statement on ``cur`` at ``timeout_seconds`` via + Postgres' ``statement_timeout`` GUC. + + Without this, a misbehaving query (accidental cross join, missing + predicate, a server-side function that hangs on a lock) would pin + the pool connection forever and eventually starve the pool — + every subsequent caller would block on ``ConnectionPool.getconn`` + instead of failing fast. The legacy :class:`SqlExecutor` got this + for free via warehouse-side polling deadlines; Postgres needs us + to opt in explicitly per statement. + + Why ``SET LOCAL`` rather than ``SET``: + ``SET LOCAL`` scopes the GUC to the current transaction and + is reset on COMMIT / ROLLBACK. Pool connections are reused + across callers, so a plain ``SET`` would silently leak the + most recent caller's timeout into the next caller's + statements — a 1-second timeout set by a healthcheck would + kill a 60-second migration on the next checkout. ``SET + LOCAL`` is the only correct primitive for a pooled + executor. + + psycopg3 starts an implicit transaction on the first statement + issued through a cursor (autocommit defaults to ``False``), so + firing this helper as the first ``cur.execute`` call BEFORE + the user SQL puts both statements in the same transaction and + the ``LOCAL`` scope covers the user SQL correctly. + + The ``timeout_seconds * 1000`` is built from Python ``int`` + arithmetic — no user input flows in — which is why + :func:`run_trusted_sql` can safely cast it to + :class:`LiteralString`. + """ + ms = max(1, int(timeout_seconds) * 1000) + run_trusted_sql(cur, f"SET LOCAL statement_timeout = {ms}") + + def execute(self, sql: str, *, timeout_seconds: int = 120) -> None: + """Run a single non-result-returning statement and commit it. + + ``timeout_seconds`` is enforced via + :meth:`_apply_statement_timeout` — a runaway query gets + cancelled rather than pinning a pool connection forever. + """ + with self._pool.connection() as conn: + with conn.cursor() as cur: + self._apply_statement_timeout(cur, timeout_seconds) + run_trusted_sql(cur, sql) + conn.commit() + + def execute_no_schema(self, sql: str) -> None: + """Parity stub. Postgres has no per-statement catalog/schema context. + + Schemas are created via ``CREATE SCHEMA IF NOT EXISTS`` like any + other DDL, so we simply delegate to :meth:`execute`. + """ + self.execute(sql) + + def query(self, sql: str, *, timeout_seconds: int = 120) -> list[list[str]]: + """Run a query and return rows as lists of strings (Delta-compatible). + + The return type is annotated as ``list[list[str]]`` to mirror + :meth:`SqlExecutor.query`, but at runtime NULL cells surface + as ``None`` (just like the JSON_ARRAY response format used by + the Statement Execution API). Services already handle both + — e.g. ``int(row[0]) if row[0] else 0`` — so we keep the + Optional shape rather than coercing NULL to ``""``. + + ``timeout_seconds`` is enforced via + :meth:`_apply_statement_timeout`. + """ + with self._pool.connection() as conn: + with conn.cursor() as cur: + self._apply_statement_timeout(cur, timeout_seconds) + run_trusted_sql(cur, sql) + rows = cur.fetchall() + return [[_to_text(cell) for cell in row] for row in rows] # pyright: ignore[reportReturnType] + + def query_dicts(self, sql: str, *, timeout_seconds: int = 120) -> list[dict[str, str | None]]: + """Run a query and return rows as ``{column: stringified value}`` dicts. + + ``timeout_seconds`` is enforced via + :meth:`_apply_statement_timeout`. + """ + with self._pool.connection() as conn: + with conn.cursor() as cur: + self._apply_statement_timeout(cur, timeout_seconds) + run_trusted_sql(cur, sql) + rows = cur.fetchall() + cols = [d.name for d in (cur.description or [])] + return [{col: _to_text(cell) for col, cell in zip(cols, row)} for row in rows] + + def upsert( + self, + table: str, + key_cols: dict[str, Any], + value_cols: dict[str, Any], + *, + timeout_seconds: int = 120, + ) -> None: + """Postgres ``INSERT ... ON CONFLICT ... DO UPDATE`` upsert. + + Keeps the same call shape as :meth:`SqlExecutor.upsert` so a + service can switch backends without code changes. The natural + key composing ``key_cols`` MUST have a UNIQUE/PRIMARY KEY index + in the migration. + """ + if not key_cols: + raise ValueError("upsert requires at least one key column") + + all_cols = list(key_cols.keys()) + list(value_cols.keys()) + all_vals = [_pg_render_value(v) for v in list(key_cols.values()) + list(value_cols.values())] + + # Natural-key columns get quoted via q() so reserved words like + # ``check`` survive. Service-provided keys are already validated + # in higher layers, but using q() also makes them dialect-safe. + quoted_cols = [self.q(c) for c in all_cols] + quoted_keys = [self.q(c) for c in key_cols] + update_set = ", ".join(f"{self.q(c)} = EXCLUDED.{self.q(c)}" for c in value_cols) + + if value_cols: + conflict_clause = f"ON CONFLICT ({', '.join(quoted_keys)}) DO UPDATE SET {update_set}" + else: + # Pure existence check — keys-only row, no update payload. + conflict_clause = f"ON CONFLICT ({', '.join(quoted_keys)}) DO NOTHING" + + sql = f"INSERT INTO {table} ({', '.join(quoted_cols)}) " f"VALUES ({', '.join(all_vals)}) " f"{conflict_clause}" + self.execute(sql, timeout_seconds=timeout_seconds) + + def upsert_with_audit( + self, + table: str, + key_cols: dict[str, Any], + value_cols: dict[str, Any], + *, + preserve_created: bool = True, + increment_on_update: str | None = None, + timeout_seconds: int = 120, + ) -> None: + """Postgres ``INSERT ... ON CONFLICT ... DO UPDATE`` with audit semantics. + + See :meth:`backend.sql_executor.OltpExecutorProtocol.upsert_with_audit` + for the contract. The increment self-reference uses the bare + column form (``"col" + 1``) which Postgres resolves against + the existing row inside DO UPDATE — no table prefix needed. + """ + if not key_cols: + raise ValueError("upsert_with_audit requires at least one key column") + if increment_on_update is not None and increment_on_update not in value_cols: + raise ValueError( + f"increment_on_update={increment_on_update!r} must be present in value_cols " + "with its initial INSERT value (e.g. {'version': 1})" + ) + + all_cols = list(key_cols.keys()) + list(value_cols.keys()) + all_vals = [_pg_render_value(v) for v in list(key_cols.values()) + list(value_cols.values())] + quoted_cols = [self.q(c) for c in all_cols] + quoted_keys = [self.q(c) for c in key_cols] + + # DO UPDATE SET excludes created_* columns when preserve_created; + # the increment column (if any) gets the bare-column + # self-reference form Postgres resolves inside ON CONFLICT DO + # UPDATE — no table prefix needed. + update_pairs: list[str] = [] + for col, val in value_cols.items(): + if preserve_created and col.startswith("created_"): + continue + qcol = self.q(col) + if col == increment_on_update: + # Bare column reference resolves to the existing row. + update_pairs.append(f"{qcol} = {qcol} + 1") + else: + update_pairs.append(f"{qcol} = {_pg_render_value(val)}") + + if update_pairs: + conflict_clause = f"ON CONFLICT ({', '.join(quoted_keys)}) DO UPDATE SET {', '.join(update_pairs)}" + else: + # All value_cols were created_* — INSERT-only with + # existence semantics. Rare but valid (e.g. a pure + # "create-if-missing" audit row). + conflict_clause = f"ON CONFLICT ({', '.join(quoted_keys)}) DO NOTHING" + + sql = f"INSERT INTO {table} ({', '.join(quoted_cols)}) " f"VALUES ({', '.join(all_vals)}) " f"{conflict_clause}" + self.execute(sql, timeout_seconds=timeout_seconds) + + def select_json_text(self, col: str) -> str: + """Postgres JSONB cells are coerced to text by :func:`_to_text` already.""" + return col + + def interval_days_expr(self, days: int) -> str: + """Postgres interval literal — single-quoted, lowercase ``days`` (plural).""" + return f"INTERVAL '{int(days)} days'" + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def close(self) -> None: + self._stop.set() + try: + self._pool.close() + except Exception: + # Best-effort lifespan teardown. close() can fail for a wide + # variety of reasons (pool-worker shutdown timing, libpq + # socket close errors, in-flight queries that the pool + # decides to abort) and the resilience contract here is + # "this MUST NOT propagate" — we're already tearing the pool + # down, and a raise would leak Postgres connections on every + # restart flap and could hang the worker. Logged for + # observability; ruff BLE001 is intentionally not enforced + # here per the policy block in ``pyproject.toml``. + logger.warning("Error closing Lakebase pool", exc_info=True) + + def _next_wait_seconds(self) -> float: + """Pick the next sleep duration before the refresh loop wakes. + + Two modes: + + * **Healthy** (``consecutive_refresh_failures == 0``): wait the + full scheduled interval (``token_refresh_seconds``) so we + don't burn SDK quota on a working token. + * **Failing** (counter > 0): wait the *retry* interval, jittered + by ±``token_refresh_retry_jitter`` so multiple workers + recovering from the same transient SDK blip don't thunder- + herd against the credential endpoint on the same second. + + Why the distinction matters: the pool's ``max_lifetime`` is + the same as ``token_refresh_seconds`` (50 min by default), so + retrying only at the scheduled cadence after a failure would + give the pool a 50-60 minute window to drain silently before + the next refresh attempt. Retrying at 10s on failure keeps + us responsive to transient failures and gives the escalation + path (``max_failures``) a sane budget to fire inside. + """ + if self._consecutive_refresh_failures == 0: + return float(self._token_refresh_seconds) + spread = self._token_refresh_retry_seconds * self._token_refresh_retry_jitter + # Lower bound is 0 in principle but never below 0.1s; without + # the floor a jitter-of-1.0 + retry-of-1s could degenerate + # into a busy-loop on the few microseconds either side of 0. + return max( + 0.1, + random.uniform( # noqa: S311 — jitter, not crypto + self._token_refresh_retry_seconds - spread, + self._token_refresh_retry_seconds + spread, + ), + ) + + def _escalate_refresh_failure(self) -> None: + """Last-resort hand-off to the process supervisor. + + Called when ``consecutive_refresh_failures`` reaches + ``token_refresh_max_failures``. The default budget (12 × + ~10s ≈ 2 minutes) leaves the pool with comfortable headroom + before ``max_lifetime`` would otherwise start expiring + connections, so escalating loudly here is strictly better + than continuing to retry silently while the pool drains. + + We log CRITICAL with full context and then exit the process + via :func:`os._exit` (not ``sys.exit``: that raises + :class:`SystemExit`, which a daemon thread swallows without + affecting the main process). The supervisor — uvicorn in + local dev, the Databricks Apps runtime in production — will + restart the worker and the new process will mint a fresh + token at startup, recovering automatically once the + underlying SDK / network / OAuth issue clears. + """ + logger.critical( + "Lakebase token refresh failed %d times in a row; exiting so the " + "supervisor restarts the worker (instance=%s, last_success=%s). " + "Continuing would silently drain the pool when ``max_lifetime`` " + "starts recycling connections without a valid replacement token.", + self._consecutive_refresh_failures, + self._instance_name, + self._last_successful_refresh_at.isoformat() if self._last_successful_refresh_at else None, + ) + # ``os._exit`` skips atexit / signal handlers so it cannot + # deadlock on a shutdown path that itself depends on a + # working Lakebase pool. The supervisor sees a non-zero exit + # and restarts us — that's the contract. + os._exit(1) # noqa: PLW1514 - deliberate hard-exit; see docstring + + def _token_refresh_loop(self) -> None: + """Background thread that rotates the Lakebase OAuth token. + + Updates both the shared :class:`_TokenHolder` and the pool's + ``kwargs`` dict so the very next physical connect uses the + fresh credential. Existing connections keep working until + ``max_lifetime`` recycles them — Postgres only validates the + password during the SCRAM handshake, not on every query. + + Sleep cadence is dynamic — see :meth:`_next_wait_seconds`: + after a success we wait the full scheduled interval; after a + failure we wait the short retry interval (jittered) so we + recover from transient failures inside the pool's + ``max_lifetime`` window rather than letting the pool drain. + + After ``token_refresh_max_failures`` consecutive failures we + hand off to :meth:`_escalate_refresh_failure` instead of + continuing to retry — losing the refresh loop is a process- + level health event, not a transient warning. + """ + while not self._stop.is_set(): + if self._stop.wait(self._next_wait_seconds()): + return + try: + fresh = _generate_token(self._ws, self._instance_name) + self._token_holder.token = fresh + # Mutating the same dict the pool was constructed with + # is the supported way to inject rotating credentials + # (psycopg-pool re-reads ``kwargs`` on every connect). + self._connect_kwargs["password"] = fresh + self._last_successful_refresh_at = datetime.now(timezone.utc) + self._consecutive_refresh_failures = 0 + logger.info("Lakebase OAuth token refreshed") + except Exception: + # Background-thread resilience: this daemon MUST NOT die. + # Token issuance can fail for SDK reasons (Databricks API + # rate limit, transient OAuth blip), network reasons + # (DNS, transient connectivity), or workspace-side reasons + # (the SP's permissions getting briefly invalidated mid- + # rotate). If any of those killed the loop, the pool's + # tokens would silently expire and every subsequent + # connect would fail with confusing SCRAM errors. The + # broad catch is the resilience contract — narrowing to + # (DatabricksError, requests.RequestException, OSError) + # cannot enumerate every plausible failure mode, and the + # cost of missing one is a silently-dead refresher. Ruff + # BLE001 is intentionally not enforced here per the + # policy block in ``pyproject.toml``. + self._consecutive_refresh_failures += 1 + logger.warning( + "Failed to refresh Lakebase token (attempt %d/%d); will retry", + self._consecutive_refresh_failures, + self._token_refresh_max_failures, + exc_info=True, + ) + if self._consecutive_refresh_failures >= self._token_refresh_max_failures: + self._escalate_refresh_failure() + return # only reached if _escalate is patched in tests + + +# --------------------------------------------------------------------------- +# Construction helper +# --------------------------------------------------------------------------- + + +def build_pg_executor( + ws: WorkspaceClient, + *, + instance_name: str, + database: str, + schema: str, + token_refresh_minutes: int = 50, + token_refresh_retry_seconds: int = 10, + token_refresh_retry_jitter: float = 0.3, + token_refresh_max_failures: int = 12, + pool_min_size: int = 1, + pool_max_size: int = 10, +) -> PgExecutor: + """Construct a :class:`PgExecutor` from a Databricks workspace client. + + Resolves the instance's read/write DNS endpoint and the calling + identity's username (service principal in production, real user + locally) before opening the pool. The token-refresh tuning + kwargs are passed straight through to :class:`PgExecutor` — see + its :meth:`__init__` and :meth:`_token_refresh_loop` for the + full back-off / escalation contract. + """ + instance = ws.database.get_database_instance(name=instance_name) + host = instance.read_write_dns + if not host: + raise RuntimeError( + f"Lakebase instance {instance_name!r} has no read_write_dns. " "Is it provisioned and running?" + ) + + me = ws.current_user.me() + username = me.user_name or me.id or "" + if not username: + raise RuntimeError("Could not determine workspace identity for Lakebase connection") + + return PgExecutor( + ws=ws, + instance_name=instance_name, + database=database, + schema=schema, + username=username, + host=host, + token_refresh_minutes=token_refresh_minutes, + token_refresh_retry_seconds=token_refresh_retry_seconds, + token_refresh_retry_jitter=token_refresh_retry_jitter, + token_refresh_max_failures=token_refresh_max_failures, + pool_min_size=pool_min_size, + pool_max_size=pool_max_size, + ) diff --git a/app/src/databricks_labs_dqx_app/backend/routes/v1/check_functions.py b/app/src/databricks_labs_dqx_app/backend/routes/v1/check_functions.py index f6c853177..195d5a686 100644 --- a/app/src/databricks_labs_dqx_app/backend/routes/v1/check_functions.py +++ b/app/src/databricks_labs_dqx_app/backend/routes/v1/check_functions.py @@ -220,7 +220,18 @@ def _resolve_callable(name: str) -> Callable[..., Any] | None: try: return resolve_check_function(name, fail_on_missing=False) - except Exception as exc: # noqa: BLE001 — defensive + except Exception as exc: + # Resolution can fail for reasons beyond "name not found": + # ``resolve_check_function`` triggers side-effect imports of + # optional check modules (PII, geo, etc.) that may not be + # installed on every deployment. A missing optional module + # raises ImportError; a syntactically-valid but broken module + # could raise AttributeError, TypeError, or anything its + # ``@register_rule`` evaluator throws. The contract here is + # "this endpoint must keep listing available functions even + # when one optional module misbehaves" — returning None lets + # the caller skip the broken entry without 500ing the whole + # registry endpoint. See the BLE001 policy in pyproject.toml. logger.warning("Failed to resolve DQX check function %r: %s", name, exc) return None diff --git a/app/src/databricks_labs_dqx_app/backend/routes/v1/config.py b/app/src/databricks_labs_dqx_app/backend/routes/v1/config.py index f353b503c..78a8e5d4e 100644 --- a/app/src/databricks_labs_dqx_app/backend/routes/v1/config.py +++ b/app/src/databricks_labs_dqx_app/backend/routes/v1/config.py @@ -4,7 +4,7 @@ from fastapi import APIRouter, Depends, HTTPException -from databricks_labs_dqx_app.backend.common.authorization import UserRole +from databricks_labs_dqx_app.backend.common.authorization import UserRole, get_user_email from databricks_labs_dqx_app.backend.dependencies import get_app_settings_service, require_role from databricks_labs_dqx_app.backend.logger import logger from pydantic import BaseModel, Field @@ -20,6 +20,17 @@ _TZ_SETTING_KEY = "display_timezone" _TZ_DEFAULT = "UTC" +# Defaults for the retention sweep — kept in sync with +# ``backend.services.scheduler_service``. Imported lazily inside the +# route to avoid pulling the scheduler module into the import graph +# of routes that have no scheduler dependency. +_RETENTION_DAYS_DEFAULT = 90 +_QUARANTINE_RETENTION_DAYS_DEFAULT = 30 +_RETENTION_DAYS_MIN = 7 +# Generous upper bound — anything past ~3 years is almost certainly a +# typo, and lets the UI render a meaningful slider/input range. +_RETENTION_DAYS_MAX = 3650 + _LABEL_DEFS_SETTING_KEY = "label_definitions" # Keys must be safe for YAML round-tripping and stable as DataFrame columns: # letters, digits, underscore, leading with a letter. @@ -101,10 +112,11 @@ def get_config( def save_config( body: ConfigIn, svc: Annotated[AppSettingsService, Depends(get_app_settings_service)], + email: Annotated[str, Depends(get_user_email)], ) -> ConfigOut: """Save workspace config to application state (admin only).""" try: - svc.save_config(body.config) + svc.save_config(body.config, user_email=email) _notify_scheduler() config = svc.get_config() return ConfigOut(config=config) @@ -140,6 +152,7 @@ def get_run_config( def save_run_config( body: RunConfigIn, svc: Annotated[AppSettingsService, Depends(get_app_settings_service)], + email: Annotated[str, Depends(get_user_email)], ) -> RunConfigOut: """Save a run config — creates or updates by name (admin only).""" config = svc.get_config() @@ -153,7 +166,7 @@ def save_run_config( if not updated: config.run_configs.append(body.config) - svc.save_config(config) + svc.save_config(config, user_email=email) _notify_scheduler() return RunConfigOut(config=body.config) @@ -167,6 +180,7 @@ def save_run_config( def delete_run_config( name: str, svc: Annotated[AppSettingsService, Depends(get_app_settings_service)], + email: Annotated[str, Depends(get_user_email)], ) -> ConfigOut: """Delete a run config by name (admin only).""" config = svc.get_config() @@ -176,7 +190,7 @@ def delete_run_config( if len(config.run_configs) == original_count: raise HTTPException(status_code=404, detail=f"Run config '{name}' not found") - svc.save_config(config) + svc.save_config(config, user_email=email) _notify_scheduler() return ConfigOut(config=config) @@ -203,12 +217,119 @@ def get_timezone( def save_timezone( body: TimezoneIn, svc: Annotated[AppSettingsService, Depends(get_app_settings_service)], + email: Annotated[str, Depends(get_user_email)], ) -> TimezoneOut: """Set the display timezone (admin only).""" - svc.save_setting(_TZ_SETTING_KEY, body.timezone) + svc.save_setting(_TZ_SETTING_KEY, body.timezone, user_email=email) return TimezoneOut(timezone=body.timezone) +# --------------------------------------------------------------------------- +# Retention — global vs. quarantine-specific DELETE windows surfaced for the +# admin UI. The scheduler reads the same keys directly from +# ``dq_app_settings`` (see ``SchedulerService._resolve_retention_days`` / +# ``_resolve_quarantine_retention_days``); these endpoints are the +# read/write surface and the only place we centralise validation. +# --------------------------------------------------------------------------- + + +class RetentionSettingsOut(BaseModel): + """Effective retention settings + the defaults the scheduler falls back to. + + ``retention_days`` / ``quarantine_retention_days`` reflect the + *current effective values* — the persisted setting if one exists, + otherwise the compiled-in default. The ``*_default`` and ``*_min`` + fields let the UI render hints and validation without duplicating + the constants on the frontend. + """ + + retention_days: int + quarantine_retention_days: int + retention_days_default: int = _RETENTION_DAYS_DEFAULT + quarantine_retention_days_default: int = _QUARANTINE_RETENTION_DAYS_DEFAULT + retention_days_min: int = _RETENTION_DAYS_MIN + retention_days_max: int = _RETENTION_DAYS_MAX + retention_days_set: bool + quarantine_retention_days_set: bool + + +class RetentionSettingsIn(BaseModel): + """Update payload — either field omitted means *leave unchanged*.""" + + retention_days: int | None = None + quarantine_retention_days: int | None = None + + +def _validate_retention_days(value: int, *, field: str) -> int: + if value < _RETENTION_DAYS_MIN: + raise HTTPException( + status_code=400, + detail=(f"{field} must be at least {_RETENTION_DAYS_MIN} days " "to protect against accidental data loss."), + ) + if value > _RETENTION_DAYS_MAX: + raise HTTPException( + status_code=400, + detail=f"{field} must be at most {_RETENTION_DAYS_MAX} days.", + ) + return value + + +@router.get( + "/retention", + response_model=RetentionSettingsOut, + operation_id="getRetentionSettings", + dependencies=[require_role(UserRole.ADMIN)], +) +def get_retention_settings( + svc: Annotated[AppSettingsService, Depends(get_app_settings_service)], +) -> RetentionSettingsOut: + """Return the current retention windows + defaults (admin only).""" + rd = svc.get_retention_days() + qd = svc.get_quarantine_retention_days() + return RetentionSettingsOut( + retention_days=rd if rd is not None else _RETENTION_DAYS_DEFAULT, + quarantine_retention_days=qd if qd is not None else _QUARANTINE_RETENTION_DAYS_DEFAULT, + retention_days_set=rd is not None, + quarantine_retention_days_set=qd is not None, + ) + + +@router.put( + "/retention", + response_model=RetentionSettingsOut, + operation_id="saveRetentionSettings", + dependencies=[require_role(UserRole.ADMIN)], +) +def save_retention_settings( + body: RetentionSettingsIn, + svc: Annotated[AppSettingsService, Depends(get_app_settings_service)], + email: Annotated[str, Depends(get_user_email)], +) -> RetentionSettingsOut: + """Update one or both retention windows (admin only). + + Either field may be omitted to leave the existing value unchanged. + Both values are validated against the safety floor and ceiling + before being persisted. + """ + if body.retention_days is None and body.quarantine_retention_days is None: + raise HTTPException( + status_code=400, + detail="At least one of retention_days or quarantine_retention_days must be provided.", + ) + + if body.retention_days is not None: + validated = _validate_retention_days(body.retention_days, field="retention_days") + svc.save_retention_days(validated, user_email=email) + logger.info("Saved global retention_days=%d", validated) + + if body.quarantine_retention_days is not None: + validated_q = _validate_retention_days(body.quarantine_retention_days, field="quarantine_retention_days") + svc.save_quarantine_retention_days(validated_q, user_email=email) + logger.info("Saved quarantine_retention_days=%d", validated_q) + + return get_retention_settings(svc) + + # --------------------------------------------------------------------------- # Label definitions — admin-managed catalog of label keys + allowed values. # Powers the constrained-mode label picker on rule authoring pages, and @@ -235,7 +356,18 @@ def _load_label_definitions(svc: AppSettingsService) -> list[LabelDefinition]: continue try: out.append(LabelDefinition.model_validate(item)) - except Exception as e: # noqa: BLE001 — best-effort migration of legacy shapes + except Exception as e: + # Per-item resilience: settings stored before the v1 label- + # definition schema may carry legacy keys or extra fields + # that ``LabelDefinition.model_validate`` rejects. Skipping + # the malformed entry preserves the rest of the user's + # configured labels; failing the whole list would 500 the + # label-definitions endpoint over a single bad row. Logged + # so the operator can fix the source row at their leisure. + # ValidationError alone would be too narrow — older payload + # shapes can trip TypeError / AttributeError inside the + # validator before Pydantic's own error wrapping fires. + # See the BLE001 policy block in pyproject.toml. logger.warning("Skipping malformed label definition %r: %s", item, e) return out @@ -265,6 +397,7 @@ def get_label_definitions( def save_label_definitions( body: LabelDefinitionsIn, svc: Annotated[AppSettingsService, Depends(get_app_settings_service)], + email: Annotated[str, Depends(get_user_email)], ) -> LabelDefinitionsOut: """Replace the full set of label definitions (admin only). @@ -306,7 +439,7 @@ def save_label_definitions( ) ) - svc.save_setting(_LABEL_DEFS_SETTING_KEY, json.dumps([d.model_dump() for d in cleaned])) + svc.save_setting(_LABEL_DEFS_SETTING_KEY, json.dumps([d.model_dump() for d in cleaned]), user_email=email) logger.info("Saved %d label definition(s)", len(cleaned)) return LabelDefinitionsOut(definitions=cleaned) @@ -386,6 +519,7 @@ def get_custom_metrics( def save_custom_metrics( body: CustomMetricsIn, svc: Annotated[AppSettingsService, Depends(get_app_settings_service)], + email: Annotated[str, Depends(get_user_email)], ) -> CustomMetricsOut: """Replace the global custom-metrics list (admin only). @@ -400,6 +534,6 @@ def save_custom_metrics( continue seen.add(expr) cleaned.append(expr) - saved = svc.save_custom_metrics(cleaned) + saved = svc.save_custom_metrics(cleaned, user_email=email) logger.info("Saved %d custom metric expression(s)", len(saved)) return CustomMetricsOut(metrics=saved) diff --git a/app/src/databricks_labs_dqx_app/backend/routes/v1/dryrun.py b/app/src/databricks_labs_dqx_app/backend/routes/v1/dryrun.py index ca47fb3fb..f98ce3856 100644 --- a/app/src/databricks_labs_dqx_app/backend/routes/v1/dryrun.py +++ b/app/src/databricks_labs_dqx_app/backend/routes/v1/dryrun.py @@ -109,6 +109,10 @@ async def list_validation_runs( total_rows=int(v) if (v := row.get("total_rows")) else None, valid_rows=int(v) if (v := row.get("valid_rows")) else None, invalid_rows=int(v) if (v := row.get("invalid_rows")) else None, + # ``is not None`` (vs the truthiness pattern above) so the UI + # can distinguish "0 errors / warnings" from pre-migration NULLs. + error_rows=int(v) if (v := row.get("error_rows")) is not None else None, + warning_rows=int(v) if (v := row.get("warning_rows")) is not None else None, created_at=row.get("created_at"), run_type=row.get("run_type"), error_message=row.get("error_message"), @@ -511,6 +515,8 @@ def get_dry_run_results( total_rows=int(v) if (v := row.get("total_rows")) else None, valid_rows=int(v) if (v := row.get("valid_rows")) else None, invalid_rows=int(v) if (v := row.get("invalid_rows")) else None, + error_rows=int(v) if (v := row.get("error_rows")) is not None else None, + warning_rows=int(v) if (v := row.get("warning_rows")) is not None else None, error_summary=json.loads(error_summary_json), sample_invalid=json.loads(sample_invalid_json), ) diff --git a/app/src/databricks_labs_dqx_app/backend/routes/v1/quarantine.py b/app/src/databricks_labs_dqx_app/backend/routes/v1/quarantine.py index eb657cd3b..ca0aab44c 100644 --- a/app/src/databricks_labs_dqx_app/backend/routes/v1/quarantine.py +++ b/app/src/databricks_labs_dqx_app/backend/routes/v1/quarantine.py @@ -36,9 +36,14 @@ def _query_quarantine( count_rows = sql.query(count_sql) total_count = int(count_rows[0][0] or 0) if count_rows and count_rows[0] else 0 + # row_data / errors / warnings are VARIANT — render as JSON strings for + # the existing _row_to_record parser. created_at is TIMESTAMP — cast to + # STRING so query_dicts returns an ISO-formatted value. data_sql = ( f"SELECT quarantine_id, run_id, source_table_fqn, requesting_user, " - f"row_data, errors, created_at " + f"to_json(row_data) AS row_data, to_json(errors) AS errors, " + f"to_json(warnings) AS warnings, " + f"CAST(created_at AS STRING) AS created_at " f"FROM {table} WHERE run_id = '{er}' " f"ORDER BY created_at DESC LIMIT {limit} OFFSET {offset}" ) @@ -61,6 +66,18 @@ def _row_to_record(row: dict[str, Any]) -> QuarantineRecordOut: except (json.JSONDecodeError, TypeError): errors = [row["errors"]] + # ``warnings`` is missing on rows written before migration v4; the + # column is ``null`` for SQL-check quarantines. + warnings: list[Any] | None = None + raw_warnings = row.get("warnings") + if raw_warnings and raw_warnings != "null": + try: + parsed = json.loads(raw_warnings) + if parsed is not None: + warnings = parsed if isinstance(parsed, list) else [parsed] + except (json.JSONDecodeError, TypeError): + warnings = [raw_warnings] + return QuarantineRecordOut( quarantine_id=row.get("quarantine_id", ""), run_id=row.get("run_id", ""), @@ -68,6 +85,7 @@ def _row_to_record(row: dict[str, Any]) -> QuarantineRecordOut: requesting_user=row.get("requesting_user"), row_data=row_data, errors=errors, + warnings=warnings, created_at=row.get("created_at"), ) @@ -132,7 +150,9 @@ def export_quarantine_records( stmt = ( f"SELECT quarantine_id, run_id, source_table_fqn, requesting_user, " - f"row_data, errors, created_at " + f"to_json(row_data) AS row_data, to_json(errors) AS errors, " + f"to_json(warnings) AS warnings, " + f"CAST(created_at AS STRING) AS created_at " f"FROM {table} WHERE run_id = '{er}' ORDER BY created_at DESC " # noqa: S608 f"LIMIT {int(max_rows)}" ) @@ -164,15 +184,18 @@ def export_quarantine_records( except (json.JSONDecodeError, TypeError): pass data_keys = sorted(all_data_keys) - headers = ["quarantine_id", "run_id", "source_table_fqn", "errors"] + data_keys + ["created_at"] + headers = ["quarantine_id", "run_id", "source_table_fqn", "errors", "warnings"] + data_keys + ["created_at"] ws.append(headers) for r in rows: + raw_warn = r.get("warnings") + warn_str = "" if not raw_warn or raw_warn == "null" else raw_warn flat: dict[str, str] = { "quarantine_id": r.get("quarantine_id", "") or "", "run_id": r.get("run_id", "") or "", "source_table_fqn": r.get("source_table_fqn", "") or "", "errors": r.get("errors", "") or "", + "warnings": warn_str or "", "created_at": r.get("created_at", "") or "", } raw_rd = r.get("row_data") @@ -203,15 +226,18 @@ def export_quarantine_records( except (json.JSONDecodeError, TypeError): pass csv_dk = sorted(csv_data_keys) - fieldnames = ["quarantine_id", "run_id", "source_table_fqn", "errors"] + csv_dk + ["created_at"] + fieldnames = ["quarantine_id", "run_id", "source_table_fqn", "errors", "warnings"] + csv_dk + ["created_at"] writer = csv.DictWriter(buf, fieldnames=fieldnames, extrasaction="ignore") writer.writeheader() for r in rows: + raw_warn_csv = r.get("warnings") + warn_str_csv = "" if not raw_warn_csv or raw_warn_csv == "null" else raw_warn_csv csv_flat: dict[str, str] = { "quarantine_id": r.get("quarantine_id", "") or "", "run_id": r.get("run_id", "") or "", "source_table_fqn": r.get("source_table_fqn", "") or "", "errors": r.get("errors", "") or "", + "warnings": warn_str_csv or "", "created_at": r.get("created_at", "") or "", } raw_rd = r.get("row_data") diff --git a/app/src/databricks_labs_dqx_app/backend/run_status_manager.py b/app/src/databricks_labs_dqx_app/backend/run_status_manager.py index 27dd85e95..23135e68f 100644 --- a/app/src/databricks_labs_dqx_app/backend/run_status_manager.py +++ b/app/src/databricks_labs_dqx_app/backend/run_status_manager.py @@ -34,7 +34,9 @@ def update_run_status( es = escape_sql_string(status) em = escape_sql_string(error_message or "") - set_clause = f"status = '{es}', error_message = '{em}', updated_at = CAST(current_timestamp() AS STRING)" + # updated_at is TIMESTAMP in the baseline; pass current_timestamp() + # directly rather than casting to STRING. + set_clause = f"status = '{es}', error_message = '{em}', updated_at = current_timestamp()" if canceled_by: ec = escape_sql_string(canceled_by) set_clause += f", canceled_by = '{ec}'" diff --git a/app/src/databricks_labs_dqx_app/backend/services/app_settings_service.py b/app/src/databricks_labs_dqx_app/backend/services/app_settings_service.py index 641d1f738..a25cd2d0f 100644 --- a/app/src/databricks_labs_dqx_app/backend/services/app_settings_service.py +++ b/app/src/databricks_labs_dqx_app/backend/services/app_settings_service.py @@ -1,79 +1,101 @@ import json import logging -from databricks.labs.blueprint.installation import Installation from databricks.labs.dqx.config import WorkspaceConfig +from pydantic import TypeAdapter, ValidationError -from databricks_labs_dqx_app.backend.sql_executor import SqlExecutor +from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol, RawSql logger = logging.getLogger(__name__) _CONFIG_KEY = "workspace_config" +# Module-level adapter so we pay the type-tree walk once at import time +# rather than on every ``get_config`` call. ``TypeAdapter`` is Pydantic's +# public v2 surface for validating non-BaseModel types against a target +# type (here, the DQX ``WorkspaceConfig`` dataclass). We use this +# instead of :func:`databricks.labs.blueprint.installation.Installation._unmarshal_type` +# — that one is a leading-underscore private API in blueprint and the +# next blueprint release can change its signature without warning. +# +# Pydantic round-trips ``WorkspaceConfig`` exactly the same way the +# blueprint helper does (verified against blueprint output during +# refactor). And it's the same path the project's ``ConfigIn`` / +# ``ConfigOut`` BaseModel response models already use to deserialize +# the user-POSTed config — so we're aligning the LOAD path with the +# trusted SAVE path rather than relying on a separate, private +# deserializer. +_WORKSPACE_CONFIG_ADAPTER: TypeAdapter[WorkspaceConfig] = TypeAdapter(WorkspaceConfig) + class AppSettingsService: - """Manages app configuration in a Delta table using app (SP) credentials. + """Manages app configuration backed by ``dq_app_settings``. + + Config is stored as a JSON blob keyed by a well-known key. All + operations use the app's service principal, not the calling + user's OBO token. - Config is stored as a JSON blob in the dq_app_settings table keyed by a - well-known key. All operations use the app's service principal, not the - calling user's OBO token. + The ``dq_app_settings`` table is one of the OLTP tables that lives + in Lakebase Postgres when ``conf.lakebase_enabled`` is true and in + Delta otherwise. The injected executor decides which. """ - def __init__(self, sql: SqlExecutor) -> None: + def __init__(self, sql: OltpExecutorProtocol) -> None: self._sql = sql - self._table = f"{sql.catalog}.{sql.schema}.dq_app_settings" + self._table = sql.fqn("dq_app_settings") # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def ensure_table(self) -> None: - """Create the settings table if it doesn't exist.""" - sql = ( - f"CREATE TABLE IF NOT EXISTS {self._table} (" - " setting_key STRING NOT NULL," - " setting_value STRING," - " updated_at TIMESTAMP," - " updated_by STRING" - ")" - ) - self._sql.execute(sql) - logger.info(f"Ensured settings table exists: {self._table}") + """No-op kept for backwards compatibility. + + The migration runner now owns table creation for both backends + (see :mod:`backend.migrations` and + :mod:`backend.migrations.postgres`); calling code that still + invokes this method gets a quiet ``DEBUG`` log and we move on. + """ + logger.debug("AppSettingsService.ensure_table() is a no-op; migrations handle DDL") def get_config(self) -> WorkspaceConfig: - """Load the workspace config from the settings table.""" + """Load the workspace config from the settings table. + + Deserialization goes through :data:`_WORKSPACE_CONFIG_ADAPTER` + (a module-level Pydantic ``TypeAdapter[WorkspaceConfig]``) so + we depend only on Pydantic — already a project dependency for + the route models — and not on any private blueprint helper. + + Defensive fallback: malformed JSON or a payload that doesn't + match the ``WorkspaceConfig`` shape returns an empty default + rather than propagating. This keeps a corrupt settings row from + bricking the admin UI; the bad row stays visible to operators + via the WARNING log without short-circuiting the rest of the + app. + """ sql = f"SELECT setting_value FROM {self._table} WHERE setting_key = '{_CONFIG_KEY}'" rows = self._sql.query(sql) if not rows: logger.info("No config found in settings table, returning default") return WorkspaceConfig(run_configs=[]) - data = json.loads(rows[0][0]) - result = Installation._unmarshal_type(data, "dq_app_settings", WorkspaceConfig) - if not isinstance(result, WorkspaceConfig): - logger.warning("Config unmarshal returned unexpected type, using default") + raw = rows[0][0] + try: + data = json.loads(raw) + except (TypeError, json.JSONDecodeError): + logger.warning("workspace_config row is not valid JSON; returning default") return WorkspaceConfig(run_configs=[]) - return result - def save_config(self, config: WorkspaceConfig) -> WorkspaceConfig: - """Save the workspace config to the settings table.""" - from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string + try: + return _WORKSPACE_CONFIG_ADAPTER.validate_python(data) + except ValidationError: + logger.warning("workspace_config row failed schema validation; returning default", exc_info=True) + return WorkspaceConfig(run_configs=[]) + def save_config(self, config: WorkspaceConfig, user_email: str | None = None) -> WorkspaceConfig: + """Save the workspace config to the settings table.""" config_dict = config.as_dict() - value = json.dumps(config_dict) - escaped = escape_sql_string(value) - - sql = ( - f"MERGE INTO {self._table} AS target " - f"USING (SELECT '{_CONFIG_KEY}' AS setting_key) AS source " - "ON target.setting_key = source.setting_key " - "WHEN MATCHED THEN UPDATE SET " - f" setting_value = '{escaped}', " - " updated_at = current_timestamp() " - "WHEN NOT MATCHED THEN INSERT (setting_key, setting_value, updated_at) " - f"VALUES ('{_CONFIG_KEY}', '{escaped}', current_timestamp())" - ) - self._sql.execute(sql) + self.save_setting(_CONFIG_KEY, json.dumps(config_dict), user_email=user_email) logger.info("Saved workspace config to settings table") return config @@ -86,24 +108,18 @@ def get_setting(self, key: str) -> str | None: rows = self._sql.query(sql) return rows[0][0] if rows else None - def save_setting(self, key: str, value: str) -> None: - """Upsert a single setting value.""" - from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string - - escaped_key = escape_sql_string(key) - escaped_val = escape_sql_string(value) - sql = ( - f"MERGE INTO {self._table} AS target " - f"USING (SELECT '{escaped_key}' AS setting_key) AS source " - "ON target.setting_key = source.setting_key " - "WHEN MATCHED THEN UPDATE SET " - f" setting_value = '{escaped_val}', " - " updated_at = current_timestamp() " - "WHEN NOT MATCHED THEN INSERT (setting_key, setting_value, updated_at) " - f"VALUES ('{escaped_key}', '{escaped_val}', current_timestamp())" + def save_setting(self, key: str, value: str, *, user_email: str | None = None) -> None: + """Upsert a single setting value, recording who wrote it.""" + self._sql.upsert( + self._table, + key_cols={"setting_key": key}, + value_cols={ + "setting_value": value, + "updated_at": RawSql("current_timestamp()"), + "updated_by": user_email, + }, ) - self._sql.execute(sql) - logger.info("Saved setting: %s", key) + logger.info("Saved setting: %s (by=%s)", key, user_email or "system") # ------------------------------------------------------------------ # Custom metrics — global SQL-expression list passed to DQMetricsObserver. @@ -127,8 +143,59 @@ def get_custom_metrics(self) -> list[str]: return [] return [s for s in parsed if isinstance(s, str) and s.strip()] - def save_custom_metrics(self, expressions: list[str]) -> list[str]: + def save_custom_metrics(self, expressions: list[str], *, user_email: str | None = None) -> list[str]: """Persist the global custom-metric list. Returns the cleaned list.""" cleaned = [s.strip() for s in expressions if isinstance(s, str) and s.strip()] - self.save_setting("custom_metrics_v1", json.dumps(cleaned)) + self.save_setting("custom_metrics_v1", json.dumps(cleaned), user_email=user_email) return cleaned + + # ------------------------------------------------------------------ + # Retention — daily DELETE sweep window for analytical tables. + # Two knobs: + # * ``retention_days`` — applied to dq_validation_runs, + # dq_profiling_results, dq_metrics + # and the OLTP history tables + # (default 90). + # * ``quarantine_retention_days`` — applied only to + # dq_quarantine_records, which + # stores the full source row + # payload (PII surface). Default + # 30 so row-level data ages out + # faster than trend tables. + # Both keys store a plain integer string. The scheduler reads them + # via ``SchedulerService._resolve_setting_days`` which floors at 7 + # days so a misconfiguration cannot wipe data inside the safety + # window. Returning ``None`` means the setting is unset and the + # consumer should fall back to its compiled-in default. + # ------------------------------------------------------------------ + + _RETENTION_KEY = "retention_days" + _QUARANTINE_RETENTION_KEY = "quarantine_retention_days" + + def get_retention_days(self) -> int | None: + """Return the configured global retention window, or ``None`` if unset.""" + return self._get_int_setting(self._RETENTION_KEY) + + def get_quarantine_retention_days(self) -> int | None: + """Return the configured quarantine retention window, or ``None`` if unset.""" + return self._get_int_setting(self._QUARANTINE_RETENTION_KEY) + + def save_retention_days(self, days: int, *, user_email: str | None = None) -> int: + """Persist the global retention window. Returns the saved value.""" + self.save_setting(self._RETENTION_KEY, str(int(days)), user_email=user_email) + return int(days) + + def save_quarantine_retention_days(self, days: int, *, user_email: str | None = None) -> int: + """Persist the quarantine retention window. Returns the saved value.""" + self.save_setting(self._QUARANTINE_RETENTION_KEY, str(int(days)), user_email=user_email) + return int(days) + + def _get_int_setting(self, key: str) -> int | None: + raw = self.get_setting(key) + if raw is None or raw == "": + return None + try: + return int(raw) + except (TypeError, ValueError): + logger.warning("Setting %s is not parseable as int (%r); ignoring", key, raw) + return None diff --git a/app/src/databricks_labs_dqx_app/backend/services/comments_service.py b/app/src/databricks_labs_dqx_app/backend/services/comments_service.py index 37f872523..afdf36202 100644 --- a/app/src/databricks_labs_dqx_app/backend/services/comments_service.py +++ b/app/src/databricks_labs_dqx_app/backend/services/comments_service.py @@ -4,7 +4,7 @@ from datetime import datetime, timezone from uuid import uuid4 -from databricks_labs_dqx_app.backend.sql_executor import SqlExecutor +from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol logger = logging.getLogger(__name__) @@ -30,9 +30,9 @@ def __init__( class CommentsService: VALID_ENTITY_TYPES = {"run", "rule"} - def __init__(self, sql: SqlExecutor) -> None: + def __init__(self, sql: OltpExecutorProtocol) -> None: self._sql = sql - self._table = f"{sql.catalog}.{sql.schema}.dq_comments" + self._table = sql.fqn("dq_comments") def add_comment(self, entity_type: str, entity_id: str, user_email: str, comment: str) -> Comment: from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string, validate_entity_type @@ -40,7 +40,6 @@ def add_comment(self, entity_type: str, entity_id: str, user_email: str, comment validate_entity_type(entity_type, self.VALID_ENTITY_TYPES) comment_id = uuid4().hex[:16] - now = datetime.now(timezone.utc).isoformat() e_comment = escape_sql_string(comment) e_email = escape_sql_string(user_email) e_entity_id = escape_sql_string(entity_id) @@ -49,7 +48,7 @@ def add_comment(self, entity_type: str, entity_id: str, user_email: str, comment sql = ( f"INSERT INTO {self._table} (comment_id, entity_type, entity_id, user_email, comment, created_at) " f"VALUES ('{comment_id}', '{e_type}', '{e_entity_id}', " - f"'{e_email}', '{e_comment}', '{now}')" + f"'{e_email}', '{e_comment}', now())" ) self._sql.execute(sql) logger.info("Added comment %s on %s/%s by %s", comment_id, entity_type, entity_id, user_email) @@ -60,7 +59,7 @@ def add_comment(self, entity_type: str, entity_id: str, user_email: str, comment entity_id=entity_id, user_email=user_email, comment=comment, - created_at=now, + created_at=datetime.now(timezone.utc).isoformat(), ) def list_comments(self, entity_type: str, entity_id: str) -> list[Comment]: @@ -71,7 +70,7 @@ def list_comments(self, entity_type: str, entity_id: str) -> list[Comment]: e_entity_id = escape_sql_string(entity_id) sql = ( f"SELECT comment_id, entity_type, entity_id, user_email, comment, " - f"CAST(created_at AS STRING) " + f"{self._sql.ts_text('created_at')} " f"FROM {self._table} " f"WHERE entity_type = '{e_type}' AND entity_id = '{e_entity_id}' " f"ORDER BY created_at ASC LIMIT 200" diff --git a/app/src/databricks_labs_dqx_app/backend/services/job_service.py b/app/src/databricks_labs_dqx_app/backend/services/job_service.py index 120c0b406..c2d15779f 100644 --- a/app/src/databricks_labs_dqx_app/backend/services/job_service.py +++ b/app/src/databricks_labs_dqx_app/backend/services/job_service.py @@ -129,18 +129,22 @@ def _record_running_placeholder( run_type: str | None = None, job_run_id: int | None = None, ) -> None: - """Insert a RUNNING placeholder row. Non-fatal on failure.""" - from datetime import datetime, timezone + """Insert a RUNNING placeholder row. Non-fatal on failure. + + ``created_at`` is now TIMESTAMP in the schema; we use + ``current_timestamp()`` rather than an ISO-string literal so the + warehouse stamps the value with its own clock and zone-mapping + works correctly on the cluster key. + """ from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string - now = datetime.now(timezone.utc).isoformat() er = escape_sql_string(run_id) eu = escape_sql_string(requesting_user) ef = escape_sql_string(source_table_fqn) ev = escape_sql_string(view_fqn) cols = f"run_id, requesting_user, source_table_fqn, view_fqn, {size_column}, status, created_at" - vals = f"'{er}', '{eu}', '{ef}', '{ev}', {int(size_value)}, 'RUNNING', '{now}'" + vals = f"'{er}', '{eu}', '{ef}', '{ev}', {int(size_value)}, 'RUNNING', current_timestamp()" if run_type: ert = escape_sql_string(run_type) cols += ", run_type" @@ -201,16 +205,23 @@ def record_dryrun_started( job_run_id=job_run_id, ) + # ``updated_at`` and ``created_at`` are TIMESTAMP — cast to STRING so + # the existing query_dicts → JSON serialization keeps producing ISO + # values for the frontend without further plumbing. _PROFILE_COLS = ( "run_id, requesting_user, source_table_fqn, view_fqn, sample_limit, " "rows_profiled, columns_profiled, duration_seconds, summary_json, " - "generated_rules_json, status, error_message, canceled_by, updated_at, created_at" + "generated_rules_json, status, error_message, canceled_by, " + "CAST(updated_at AS STRING) AS updated_at, " + "CAST(created_at AS STRING) AS created_at" ) _DRYRUN_COLS = ( "run_id, requesting_user, source_table_fqn, sample_size, " - "total_rows, valid_rows, invalid_rows, " - "status, error_message, canceled_by, updated_at, created_at, " + "total_rows, valid_rows, invalid_rows, error_rows, warning_rows, " + "status, error_message, canceled_by, " + "CAST(updated_at AS STRING) AS updated_at, " + "CAST(created_at AS STRING) AS created_at, " "COALESCE(run_type, 'dryrun') AS run_type, " "checks_json" ) diff --git a/app/src/databricks_labs_dqx_app/backend/services/role_service.py b/app/src/databricks_labs_dqx_app/backend/services/role_service.py index bc25eca91..40aaa39b5 100644 --- a/app/src/databricks_labs_dqx_app/backend/services/role_service.py +++ b/app/src/databricks_labs_dqx_app/backend/services/role_service.py @@ -10,7 +10,7 @@ from datetime import datetime, timezone from databricks_labs_dqx_app.backend.common.authorization import ROLE_PRIORITY, UserRole -from databricks_labs_dqx_app.backend.sql_executor import SqlExecutor +from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol, RawSql from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string logger = logging.getLogger(__name__) @@ -31,15 +31,17 @@ class RoleMapping: class RoleService: - """Manages role-to-group mappings in a Delta table. + """Manages role-to-group mappings in ``dq_role_mappings``. - All operations use the app's service principal, not the calling user's - OBO token. + All operations use the app's service principal, not the calling + user's OBO token. The table lives on Lakebase Postgres when + Lakebase is enabled and on Delta otherwise — the injected + executor decides which. """ - def __init__(self, sql: SqlExecutor) -> None: + def __init__(self, sql: OltpExecutorProtocol) -> None: self._sql = sql - self._table = f"{sql.catalog}.{sql.schema}.dq_role_mappings" + self._table = sql.fqn("dq_role_mappings") self._mappings_cache: list[RoleMapping] | None = None self._mappings_cache_expires: float = 0.0 @@ -50,7 +52,7 @@ def list_mappings(self, *, use_cache: bool = False) -> list[RoleMapping]: sql = ( f"SELECT role, group_name, created_by, " - f"CAST(created_at AS STRING), updated_by, CAST(updated_at AS STRING) " + f"{self._sql.ts_text('created_at')}, updated_by, {self._sql.ts_text('updated_at')} " f"FROM {self._table} ORDER BY role, group_name" ) rows = self._sql.query(sql) @@ -79,7 +81,7 @@ def get_mappings_for_role(self, role: str) -> list[RoleMapping]: escaped_role = escape_sql_string(role) sql = ( f"SELECT role, group_name, created_by, " - f"CAST(created_at AS STRING), updated_by, CAST(updated_at AS STRING) " + f"{self._sql.ts_text('created_at')}, updated_by, {self._sql.ts_text('updated_at')} " f"FROM {self._table} WHERE role = '{escaped_role}' ORDER BY group_name" ) rows = self._sql.query(sql) @@ -96,26 +98,29 @@ def get_mappings_for_role(self, role: str) -> list[RoleMapping]: ] def create_mapping(self, role: str, group_name: str, user_email: str) -> RoleMapping: - """Create or update a role-to-group mapping.""" + """Create or update a role-to-group mapping. + + Dialect-agnostic — :meth:`OltpExecutorProtocol.upsert_with_audit` + with ``preserve_created=True`` handles the + ``MERGE INTO`` (Delta) / ``INSERT ... ON CONFLICT DO UPDATE`` + (Postgres) choice so the original ``created_*`` survive on + UPDATE. + """ if role not in [r.value for r in UserRole]: raise ValueError(f"Invalid role: {role}. Must be one of {[r.value for r in UserRole]}") - escaped_role = escape_sql_string(role) - escaped_group = escape_sql_string(group_name) - escaped_user = escape_sql_string(user_email) - - sql = ( - f"MERGE INTO {self._table} AS target " - f"USING (SELECT '{escaped_role}' AS role, '{escaped_group}' AS group_name) AS source " - "ON target.role = source.role AND target.group_name = source.group_name " - "WHEN MATCHED THEN UPDATE SET " - f" updated_by = '{escaped_user}', " - " updated_at = current_timestamp() " - "WHEN NOT MATCHED THEN INSERT (role, group_name, created_by, created_at, updated_by, updated_at) " - f"VALUES ('{escaped_role}', '{escaped_group}', '{escaped_user}', current_timestamp(), " - f"'{escaped_user}', current_timestamp())" + now = RawSql("now()") + self._sql.upsert_with_audit( + table=self._table, + key_cols={"role": role, "group_name": group_name}, + value_cols={ + "created_by": user_email, + "created_at": now, + "updated_by": user_email, + "updated_at": now, + }, + preserve_created=True, ) - self._sql.execute(sql) self.invalidate_mappings_cache() logger.info(f"Created/updated role mapping: {role} -> {group_name}") diff --git a/app/src/databricks_labs_dqx_app/backend/services/rules_catalog_service.py b/app/src/databricks_labs_dqx_app/backend/services/rules_catalog_service.py index 5dd521cbc..f2a1a1008 100644 --- a/app/src/databricks_labs_dqx_app/backend/services/rules_catalog_service.py +++ b/app/src/databricks_labs_dqx_app/backend/services/rules_catalog_service.py @@ -6,14 +6,20 @@ from typing import Any from uuid import uuid4 -from databricks_labs_dqx_app.backend.sql_executor import SqlExecutor +from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string logger = logging.getLogger(__name__) class RuleCatalogEntry: - """Represents a single rule (individual check) in the rules catalog.""" + """Represents a single rule (one check) in the rules catalog. + + Note: ``checks`` is preserved as a one-element list for API/DTO + compatibility with existing routes and frontend types. The underlying + Delta column ``check`` is a VARIANT object holding the bare check — + the array wrapper is purely an in-memory representation. + """ def __init__( self, @@ -44,8 +50,9 @@ class RulesCatalogService: """Manages the rules catalog in a Delta table using app (SP) credentials. Each row represents a single rule (one check) identified by ``rule_id``. - Multiple rules can target the same ``table_fqn``. For execution the caller - aggregates all approved rules for a table into a single checks array. + The persisted ``check`` column is a VARIANT holding the bare check object + (no array wrapper); ``RuleCatalogEntry.checks`` exposes a one-element list + so existing callers don't have to change. """ VALID_STATUSES = {"draft", "pending_approval", "approved", "rejected"} @@ -57,16 +64,40 @@ class RulesCatalogService: "rejected": {"draft"}, } - _SELECT_COLS = ( - "table_fqn, checks, version, status, created_by, " - "CAST(created_at AS STRING), updated_by, CAST(updated_at AS STRING), " - "COALESCE(source, 'ui'), COALESCE(rule_id, '')" - ) - - def __init__(self, sql: SqlExecutor) -> None: + def __init__(self, sql: OltpExecutorProtocol) -> None: self._sql = sql - self._table = f"{sql.catalog}.{sql.schema}.dq_quality_rules" - self._history_table = f"{sql.catalog}.{sql.schema}.dq_quality_rules_history" + self._table = sql.fqn("dq_quality_rules") + self._history_table = sql.fqn("dq_quality_rules_history") + # ``check`` is a SQL reserved word in both Delta and Postgres, + # so quote it via the executor so we get backticks on Delta and + # double-quotes on Postgres. + self._check_col = sql.q("check") + self._select_cols = self._build_select_cols() + + def _build_select_cols(self) -> str: + """Build the column projection used by every SELECT. + + Delegates dialect-specific rendering to the executor: + + - :meth:`OltpExecutorProtocol.select_json_text` renders the + VARIANT/JSONB ``check`` column as JSON text — ``to_json(col)`` + on Delta, the bare column on Postgres (where + :func:`backend.pg_executor._to_text` ISO-encodes JSONB cells + on the way out). + - :meth:`OltpExecutorProtocol.ts_text` handles the timestamp + column projection the same way — ``CAST(... AS STRING)`` on + Delta, verbatim on Postgres. + + The column *order* is identical regardless of dialect so + :meth:`_row_to_entry` doesn't have to branch. + """ + check_text = self._sql.select_json_text(self._check_col) + created_at = self._sql.ts_text("created_at") + updated_at = self._sql.ts_text("updated_at") + return ( + f"table_fqn, {check_text} AS check_json, version, status, created_by, " + f"{created_at}, updated_by, {updated_at}, source, rule_id" + ) # ------------------------------------------------------------------ # Public API @@ -126,24 +157,36 @@ def save( if duplicates: logger.info("Skipped %d duplicate check(s) for table %s", len(duplicates), table_fqn) - now = datetime.now(timezone.utc).isoformat() e_table = escape_sql_string(table_fqn) e_source = escape_sql_string(source) e_user = escape_sql_string(user_email) created: list[RuleCatalogEntry] = [] + check_col = self._check_col for check in non_dup_checks: rule_id = uuid4().hex[:16] - single_check_json = json.dumps([check]) - e_checks = escape_sql_string(single_check_json) + check_json = json.dumps(check) + check_expr = self._sql.json_literal_expr(check_json) sql = ( f"INSERT INTO {self._table} " - "(table_fqn, checks, version, status, source, created_by, created_at, updated_by, updated_at, rule_id) " - f"VALUES ('{e_table}', '{e_checks}', 1, 'draft', '{e_source}', " - f"'{e_user}', '{now}', '{e_user}', '{now}', '{rule_id}')" + f"(rule_id, table_fqn, {check_col}, version, status, source, " + f"created_by, created_at, updated_by, updated_at) " + f"VALUES ('{rule_id}', '{e_table}', {check_expr}, 1, 'draft', '{e_source}', " + f"'{e_user}', now(), '{e_user}', now())" ) self._sql.execute(sql) - self._record_history(table_fqn, single_check_json, 1, source, user_email, now, "save", rule_id) + self._record_history( + table_fqn=table_fqn, + check_json=check_json, + version=1, + source=source, + user_email=user_email, + action="save", + rule_id=rule_id, + prev_status=None, + new_status="draft", + ) + now = datetime.now(timezone.utc).isoformat() created.append( RuleCatalogEntry( table_fqn=table_fqn, @@ -173,25 +216,37 @@ def update_rule( entry = self.get_by_rule_id(rule_id) if entry is None: raise RuntimeError(f"Rule not found: {rule_id}") - - checks_json = json.dumps(checks) - now = datetime.now(timezone.utc).isoformat() - e_checks = escape_sql_string(checks_json) + if not checks: + raise ValueError("update_rule requires exactly one check; got none") + + # update_rule is one-rule-at-a-time at the table layer; collapse to a + # single check object for the VARIANT/JSONB column. + check = checks[0] + check_json = json.dumps(check) + check_expr = self._sql.json_literal_expr(check_json) e_user = escape_sql_string(user_email) e_rule_id = escape_sql_string(rule_id) sql = ( f"UPDATE {self._table} SET " - f" checks = '{e_checks}', " - " version = version + 1, " - " status = 'draft', " + f" {self._check_col} = {check_expr}, " + f" version = version + 1, " + f" status = 'draft', " f" updated_by = '{e_user}', " - f" updated_at = '{now}' " + f" updated_at = now() " f"WHERE rule_id = '{e_rule_id}'" ) self._sql.execute(sql) self._record_history( - entry.table_fqn, checks_json, entry.version + 1, entry.source, user_email, now, "update", rule_id + table_fqn=entry.table_fqn, + check_json=check_json, + version=entry.version + 1, + source=entry.source, + user_email=user_email, + action="update", + rule_id=rule_id, + prev_status=entry.status, + new_status="draft", ) logger.info("Updated rule %s (table %s)", rule_id, entry.table_fqn) return self.get_by_rule_id(rule_id) or entry @@ -228,7 +283,7 @@ def get(self, table_fqn: str) -> RuleCatalogEntry | None: def get_by_rule_id(self, rule_id: str) -> RuleCatalogEntry | None: """Get a single rule by its rule_id.""" e_rule_id = escape_sql_string(rule_id) - sql = f"SELECT {self._SELECT_COLS} FROM {self._table} WHERE rule_id = '{e_rule_id}'" # noqa: S608 + sql = f"SELECT {self._select_cols} FROM {self._table} WHERE rule_id = '{e_rule_id}'" # noqa: S608 rows = self._sql.query(sql) if not rows: return None @@ -237,7 +292,7 @@ def get_by_rule_id(self, rule_id: str) -> RuleCatalogEntry | None: def list_rules_for_table(self, table_fqn: str, status: str | None = None) -> list[RuleCatalogEntry]: """List all individual rules for a given table, optionally filtered by status.""" e_table = escape_sql_string(table_fqn) - sql = f"SELECT {self._SELECT_COLS} FROM {self._table} WHERE table_fqn = '{e_table}'" # noqa: S608 + sql = f"SELECT {self._select_cols} FROM {self._table} WHERE table_fqn = '{e_table}'" # noqa: S608 if status: e_status = escape_sql_string(status) sql += f" AND status = '{e_status}'" @@ -247,7 +302,7 @@ def list_rules_for_table(self, table_fqn: str, status: str | None = None) -> lis def list_rules(self, status: str | None = None) -> list[RuleCatalogEntry]: """List all individual rules, optionally filtered by status.""" - sql = f"SELECT {self._SELECT_COLS} FROM {self._table}" + sql = f"SELECT {self._select_cols} FROM {self._table}" if status: e_status = escape_sql_string(status) sql += f" WHERE status = '{e_status}'" @@ -332,11 +387,24 @@ def delete(self, rule_id: str, user_email: str) -> None: e_rule_id = escape_sql_string(rule_id) sql = f"DELETE FROM {self._table} WHERE rule_id = '{e_rule_id}'" self._sql.execute(sql) - now = datetime.now(timezone.utc).isoformat() table_fqn = entry.table_fqn if entry else "unknown" version = entry.version if entry else 0 source = entry.source if entry else "ui" - self._record_history(table_fqn, None, version, source, user_email, now, "delete", rule_id) + # Preserve the post-state ``check`` payload in history so audit + # readers can reconstruct what was deleted without walking back to + # the prior save row. + check_json = json.dumps(entry.checks[0]) if entry and entry.checks else None + self._record_history( + table_fqn=table_fqn, + check_json=check_json, + version=version, + source=source, + user_email=user_email, + action="delete", + rule_id=rule_id, + prev_status=entry.status if entry else None, + new_status=None, + ) logger.info("Deleted rule %s (table %s, by %s)", rule_id, table_fqn, user_email) def delete_by_table(self, table_fqn: str, user_email: str) -> None: @@ -344,8 +412,17 @@ def delete_by_table(self, table_fqn: str, user_email: str) -> None: e_table = escape_sql_string(table_fqn) sql = f"DELETE FROM {self._table} WHERE table_fqn = '{e_table}'" self._sql.execute(sql) - now = datetime.now(timezone.utc).isoformat() - self._record_history(table_fqn, None, 0, "ui", user_email, now, "delete_all") + self._record_history( + table_fqn=table_fqn, + check_json=None, + version=0, + source="ui", + user_email=user_email, + action="delete_all", + rule_id=None, + prev_status=None, + new_status=None, + ) logger.info("Deleted all rules for table %s (by %s)", table_fqn, user_email) def set_status( @@ -378,7 +455,6 @@ def set_status( f"but current is v{entry.version}. Another user may have modified the rule." ) - now = datetime.now(timezone.utc).isoformat() e_status = escape_sql_string(status) e_user = escape_sql_string(user_email) e_rule_id = escape_sql_string(rule_id) @@ -387,12 +463,24 @@ def set_status( f"UPDATE {self._table} SET " f" status = '{e_status}', " f" updated_by = '{e_user}', " - f" updated_at = '{now}' " + f" updated_at = now() " f"WHERE rule_id = '{e_rule_id}' AND version = {entry.version}" ) self._sql.execute(sql) + # Always include the post-state ``check`` payload + explicit + # prev/new status pair so dashboards reconstructing the trail + # don't have to walk back to the prior save row. + check_json = json.dumps(entry.checks[0]) if entry.checks else None self._record_history( - entry.table_fqn, None, entry.version, entry.source, user_email, now, f"status:{status}", rule_id + table_fqn=entry.table_fqn, + check_json=check_json, + version=entry.version, + source=entry.source, + user_email=user_email, + action=f"status:{status}", + rule_id=rule_id, + prev_status=entry.status, + new_status=status, ) logger.info("Updated status for rule %s to %s (by %s)", rule_id, status, user_email) @@ -435,21 +523,14 @@ def get_approved_checks_for_table(self, table_fqn: str) -> list[dict[str, Any]]: def backfill_rule_ids(self) -> int: """Assign a rule_id to every row that currently has NULL or empty rule_id. - Returns the number of rows updated. + Retained as a no-op safety net: in the current baseline ``rule_id`` + is ``NOT NULL`` with a PK constraint, so this should always return 0. + Kept so legacy upgrade paths that ran the pre-baseline migrations + can still call it without an AttributeError. """ count_sql = f"SELECT COUNT(*) FROM {self._table} WHERE rule_id IS NULL OR rule_id = ''" rows = self._sql.query(count_sql) total = int(rows[0][0]) if rows and rows[0] else 0 - if total == 0: - return 0 - - sql = ( - f"UPDATE {self._table} " - "SET rule_id = SUBSTRING(MD5(CONCAT(table_fqn, checks, CAST(RAND() AS STRING))), 1, 16) " - "WHERE rule_id IS NULL OR rule_id = ''" - ) - self._sql.execute(sql) - logger.info("Backfilled rule_id for %d rule(s)", total) return total # ------------------------------------------------------------------ @@ -491,8 +572,27 @@ def _check_no_duplicate_pending_or_approved(self, entry: RuleCatalogEntry) -> No ) def _row_to_entry(self, row: list[str]) -> RuleCatalogEntry: - """Convert a query result row to a RuleCatalogEntry.""" - checks = json.loads(row[1], strict=False) if row[1] else [] + """Convert a query result row to a RuleCatalogEntry. + + Row layout (see :meth:`_build_select_cols`): + ``[table_fqn, check_json, version, status, created_by, created_at, + updated_by, updated_at, source, rule_id]``. + + ``check_json`` is the JSON rendering of the VARIANT/JSONB + column (via Delta's ``to_json`` or PgExecutor's automatic + text coercion). We wrap it in a one-element list so the + in-memory ``checks`` shape keeps existing callers (route DTOs, + scheduler) unchanged. + """ + check_json = row[1] + if check_json: + try: + parsed = json.loads(check_json, strict=False) + except json.JSONDecodeError: + parsed = None + checks = [parsed] if isinstance(parsed, dict) else [] + else: + checks = [] return RuleCatalogEntry( table_fqn=row[0], checks=checks, @@ -508,28 +608,44 @@ def _row_to_entry(self, row: list[str]) -> RuleCatalogEntry: def _record_history( self, + *, table_fqn: str, - checks_json: str | None, + check_json: str | None, version: int, source: str, user_email: str, - timestamp: str, action: str, - rule_id: str | None = None, + rule_id: str | None, + prev_status: str | None, + new_status: str | None, ) -> None: - """Insert an audit row into the history table (best-effort).""" + """Insert an audit row into the history table (best-effort). + + Always carries the post-state ``check`` payload (when one exists) + and an explicit ``prev_status``/``new_status`` pair so callers + querying the audit trail don't have to walk backwards through the + log to reconstruct what changed. + + ``check`` is a SQL reserved word so the column name is quoted + via the executor's :meth:`q` helper (backticks for Delta, + double quotes for Postgres). + """ try: e_table = escape_sql_string(table_fqn) - e_checks = escape_sql_string(checks_json or "") e_source = escape_sql_string(source) e_action = escape_sql_string(action) e_user = escape_sql_string(user_email) - e_rule_id = escape_sql_string(rule_id or "") + rule_id_sql = f"'{escape_sql_string(rule_id)}'" if rule_id else "NULL" + check_sql = self._sql.json_literal_expr(check_json) if check_json else "NULL" + prev_sql = f"'{escape_sql_string(prev_status)}'" if prev_status else "NULL" + new_sql = f"'{escape_sql_string(new_status)}'" if new_status else "NULL" + sql = ( f"INSERT INTO {self._history_table} " - "(table_fqn, checks, version, source, action, changed_by, changed_at, rule_id) VALUES " - f"('{e_table}', '{e_checks}', {version}, '{e_source}', '{e_action}', " - f"'{e_user}', '{timestamp}', '{e_rule_id}')" + f"(rule_id, table_fqn, {self._check_col}, version, source, action, " + f"prev_status, new_status, changed_by, changed_at) VALUES " + f"({rule_id_sql}, '{e_table}', {check_sql}, {version}, '{e_source}', " + f"'{e_action}', {prev_sql}, {new_sql}, '{e_user}', now())" ) self._sql.execute(sql) except Exception: diff --git a/app/src/databricks_labs_dqx_app/backend/services/schedule_config_service.py b/app/src/databricks_labs_dqx_app/backend/services/schedule_config_service.py index 941fde0e0..4b9b476dd 100644 --- a/app/src/databricks_labs_dqx_app/backend/services/schedule_config_service.py +++ b/app/src/databricks_labs_dqx_app/backend/services/schedule_config_service.py @@ -13,7 +13,7 @@ from datetime import datetime, timezone from typing import Any -from databricks_labs_dqx_app.backend.sql_executor import SqlExecutor +from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol, RawSql from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string, validate_schedule_name logger = logging.getLogger(__name__) @@ -33,15 +33,16 @@ class ScheduleConfigEntry: class ScheduleConfigService: """CRUD for per-schedule configuration rows in ``dq_schedule_configs``.""" - def __init__(self, sql: SqlExecutor) -> None: + def __init__(self, sql: OltpExecutorProtocol) -> None: self._sql = sql - self._table = f"{sql.catalog}.{sql.schema}.dq_schedule_configs" - self._history_table = f"{sql.catalog}.{sql.schema}.dq_schedule_configs_history" + self._table = sql.fqn("dq_schedule_configs") + self._history_table = sql.fqn("dq_schedule_configs_history") def list_schedules(self) -> list[ScheduleConfigEntry]: + ts = self._sql.ts_text sql = ( f"SELECT schedule_name, config_json, version, created_by, " - f"CAST(created_at AS STRING), updated_by, CAST(updated_at AS STRING) " + f"{ts('created_at')}, updated_by, {ts('updated_at')} " f"FROM {self._table} ORDER BY schedule_name" ) rows = self._sql.query(sql) @@ -50,9 +51,10 @@ def list_schedules(self) -> list[ScheduleConfigEntry]: def get(self, name: str) -> ScheduleConfigEntry | None: validate_schedule_name(name) escaped = escape_sql_string(name) + ts = self._sql.ts_text sql = ( f"SELECT schedule_name, config_json, version, created_by, " - f"CAST(created_at AS STRING), updated_by, CAST(updated_at AS STRING) " + f"{ts('created_at')}, updated_by, {ts('updated_at')} " f"FROM {self._table} WHERE schedule_name = '{escaped}'" ) rows = self._sql.query(sql) @@ -66,33 +68,42 @@ def save( config: dict[str, Any], user_email: str, ) -> ScheduleConfigEntry: + """Upsert a schedule config row, incrementing ``version`` on update. + + Dialect-agnostic via :meth:`OltpExecutorProtocol.upsert_with_audit`: + + - ``preserve_created=True`` keeps the original ``created_*`` + on UPDATE, + - ``increment_on_update="version"`` rewrites the version + column's UPDATE branch to ``existing + 1`` using the + dialect-appropriate self-reference (``target.version + 1`` + on Delta MERGE, bare ``"version" + 1`` on Postgres + ON CONFLICT). + """ validate_schedule_name(name) config_json = json.dumps(config) - now = datetime.now(timezone.utc).isoformat() - escaped_name = escape_sql_string(name) - escaped_json = escape_sql_string(config_json) - escaped_user = escape_sql_string(user_email) - - sql = ( - f"MERGE INTO {self._table} AS target " - f"USING (SELECT '{escaped_name}' AS schedule_name) AS source " - "ON target.schedule_name = source.schedule_name " - "WHEN MATCHED THEN UPDATE SET " - f" config_json = '{escaped_json}', " - " version = target.version + 1, " - f" updated_by = '{escaped_user}', " - f" updated_at = '{now}' " - "WHEN NOT MATCHED THEN INSERT " - "(schedule_name, config_json, version, created_by, created_at, updated_by, updated_at) " - f"VALUES ('{escaped_name}', '{escaped_json}', 1, '{escaped_user}', '{now}', '{escaped_user}', '{now}')" + now = RawSql("now()") + self._sql.upsert_with_audit( + table=self._table, + key_cols={"schedule_name": name}, + value_cols={ + "config_json": config_json, + "version": 1, # initial INSERT value; UPDATE branch increments via increment_on_update + "created_by": user_email, + "created_at": now, + "updated_by": user_email, + "updated_at": now, + }, + preserve_created=True, + increment_on_update="version", ) - self._sql.execute(sql) self._record_history(name, config_json, user_email, "save") logger.info("Saved schedule config: %s (user=%s)", name, user_email) entry = self.get(name) if entry is None: + now = datetime.now(timezone.utc).isoformat() return ScheduleConfigEntry( schedule_name=name, config=config, @@ -125,7 +136,7 @@ def get_history(self, name: str) -> list[dict[str, Any]]: escaped = escape_sql_string(name) sql = ( f"SELECT schedule_name, config_json, version, action, changed_by, " - f"CAST(changed_at AS STRING) " + f"{self._sql.ts_text('changed_at')} " f"FROM {self._history_table} " f"WHERE schedule_name = '{escaped}' " "ORDER BY changed_at DESC" @@ -162,7 +173,6 @@ def _record_history( version: int = 0, ) -> None: try: - now = datetime.now(timezone.utc).isoformat() escaped_name = escape_sql_string(name) escaped_json = escape_sql_string(config_json) escaped_user = escape_sql_string(user_email) @@ -171,7 +181,7 @@ def _record_history( f"INSERT INTO {self._history_table} " "(schedule_name, config_json, version, action, changed_by, changed_at) " f"VALUES ('{escaped_name}', '{escaped_json}', {version}, '{escaped_action}', " - f"'{escaped_user}', '{now}')" + f"'{escaped_user}', now())" ) self._sql.execute(sql) except Exception: diff --git a/app/src/databricks_labs_dqx_app/backend/services/scheduler_service.py b/app/src/databricks_labs_dqx_app/backend/services/scheduler_service.py index d788a770c..2eb6531b3 100644 --- a/app/src/databricks_labs_dqx_app/backend/services/scheduler_service.py +++ b/app/src/databricks_labs_dqx_app/backend/services/scheduler_service.py @@ -22,7 +22,7 @@ from databricks.sdk import WorkspaceClient from databricks_labs_dqx_app.backend.logger import get_logger -from databricks_labs_dqx_app.backend.sql_executor import SqlExecutor +from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol, RawSql, SqlExecutor logger = get_logger("scheduler") @@ -65,6 +65,45 @@ _GC_AGE_HOURS = 48 _GC_MAX_DROPS_PER_RUN = 500 +# Retention sweep — daily DELETE pass against the high-volume tables to +# keep them from growing without bound. Each (table, time-column) pair +# in :data:`_RETENTION_TABLES` is trimmed to ``RETENTION_DAYS`` worth of +# history. Defaults to 90 days; configurable via the ``retention_days`` +# setting in ``dq_app_settings``. The sweep runs at most once per +# ``_RETENTION_INTERVAL_HOURS`` so the warehouse isn't billed repeatedly +# for the same DELETE. +_RETENTION_DAYS_DEFAULT = 90 +_RETENTION_DAYS_MIN = 7 +_RETENTION_INTERVAL_HOURS = 24 + +# ``dq_quarantine_records`` is the only table that holds full row +# payloads (the source row + ``_errors`` / ``_warnings`` blobs). Those +# rows are PII-sensitive and tend to drive most of the Studio's storage +# growth, so we expose a *separate* retention knob with a tighter +# default (30 days) instead of subjecting them to the same window as +# trend tables like ``dq_metrics`` (which dashboards expect to look +# back ~3 months on). Set via the ``quarantine_retention_days`` key +# in ``dq_app_settings``; falls back here when unset. +_QUARANTINE_RETENTION_DAYS_DEFAULT = 30 +_QUARANTINE_TABLE_NAME = "dq_quarantine_records" + +# Retention is split per-backend: analytical (Delta) tables are +# trimmed via the SQL warehouse executor, OLTP tables via the OLTP +# executor (Lakebase if enabled, Delta otherwise). Both lists are +# walked on every retention sweep. ``dq_quarantine_records`` is in +# this list but resolves its cutoff via :meth:`_resolve_quarantine_retention_days` +# instead of the global :meth:`_resolve_retention_days`. +_DELTA_RETENTION_TABLES: tuple[tuple[str, str], ...] = ( + ("dq_validation_runs", "created_at"), + ("dq_profiling_results", "created_at"), + (_QUARANTINE_TABLE_NAME, "created_at"), + ("dq_metrics", "run_time"), +) +_OLTP_RETENTION_TABLES: tuple[tuple[str, str], ...] = ( + ("dq_quality_rules_history", "changed_at"), + ("dq_schedule_configs_history", "changed_at"), +) + class SchedulerService: """Manages a background loop that checks schedule configs and triggers runs.""" @@ -77,7 +116,23 @@ def __init__( schema: str, tmp_schema: str, job_id: str, + oltp_sql: OltpExecutorProtocol | None = None, ) -> None: + """Construct the scheduler. + + Parameters + ---------- + oltp_sql: + Executor used for OLTP-table operations (schedule + tracking, schedule configs, app settings, rule reads). + When ``None`` (legacy mode, no Lakebase) the same Delta + executor is used for everything. When Lakebase is enabled, + callers pass a :class:`backend.pg_executor.PgExecutor` so + the high-frequency reads/writes hit Postgres. Typed as + :class:`OltpExecutorProtocol` so both concrete executors + are accepted without a runtime cast — the Protocol is the + structural contract every OLTP call site relies on. + """ self._ws = ws self._job_id = job_id self._catalog = catalog @@ -85,13 +140,25 @@ def __init__( self._tmp_schema = tmp_schema self._sql = SqlExecutor(ws=ws, warehouse_id=warehouse_id, catalog=catalog, schema=schema) self._tmp_sql = SqlExecutor(ws=ws, warehouse_id=warehouse_id, catalog=catalog, schema=tmp_schema) + # OLTP executor — either a PgExecutor (Lakebase) or the same + # Delta executor (legacy mode). All schedule / settings / + # rule access goes through this; only analytical table + # operations (retention sweep, orphan view GC) use ``self._sql``. + # No cast: ``SqlExecutor`` structurally satisfies + # :class:`OltpExecutorProtocol`, so basedpyright validates + # every ``self._oltp_sql.foo()`` call against the same + # Protocol surface regardless of which concrete executor was + # injected. + self._oltp_sql: OltpExecutorProtocol = oltp_sql if oltp_sql is not None else self._sql self._task: asyncio.Task[None] | None = None self._reload_event = asyncio.Event() self._force_recalc = False - self._table = f"{catalog}.{schema}.dq_schedule_runs" - self._configs_table = f"{catalog}.{schema}.dq_schedule_configs" - self._settings_table = f"{catalog}.{schema}.dq_app_settings" - self._rules_table = f"{catalog}.{schema}.dq_quality_rules" + # Both backend layouts qualify the table differently — let the + # OLTP executor's catalog/schema decide. + self._table = self._oltp_sql.fqn("dq_schedule_runs") + self._configs_table = self._oltp_sql.fqn("dq_schedule_configs") + self._settings_table = self._oltp_sql.fqn("dq_app_settings") + self._rules_table = self._oltp_sql.fqn("dq_quality_rules") # Orphan-tmp-view GC: fires every Saturday at 01:00 UTC. Held in # process memory rather than persisted — a missed Saturday (e.g. @@ -100,6 +167,11 @@ def __init__( # and orphans only accumulate slowly. self._next_view_gc_at: datetime = self._next_saturday_01_utc(datetime.now(timezone.utc)) + # Retention sweep: fires every ``_RETENTION_INTERVAL_HOURS`` + # (default 24h). Held in process memory like the view GC; a + # missed sweep is harmless since the next one catches up. + self._next_retention_at: datetime = datetime.now(timezone.utc) + timedelta(hours=_RETENTION_INTERVAL_HOURS) + # ------------------------------------------------------------------ # Lifecycle # ------------------------------------------------------------------ @@ -140,6 +212,7 @@ async def _loop(self) -> None: self._force_recalc = False await self._tick(recalc=recalc) await self._maybe_gc_orphan_views(datetime.now(timezone.utc)) + await self._maybe_run_retention(datetime.now(timezone.utc)) except asyncio.CancelledError: raise except Exception: @@ -236,7 +309,7 @@ def _load_schedule_configs(self) -> dict[str, dict[str, Any]]: try: sql = f"SELECT schedule_name, config_json FROM {self._configs_table}" - rows = self._sql.query(sql) + rows = self._oltp_sql.query(sql) for row in rows: name = row[0] or "" if not name: @@ -256,7 +329,7 @@ def _load_schedule_configs(self) -> dict[str, dict[str, Any]]: # Legacy fallback: read from dq_app_settings blob try: sql = f"SELECT setting_value FROM {self._settings_table} WHERE setting_key = 'workspace_config'" - rows = self._sql.query(sql) + rows = self._oltp_sql.query(sql) if not rows: return {} data = json.loads(rows[0][0]) @@ -286,12 +359,13 @@ def _get_tracker(self, name: str) -> dict[str, str] | None: validate_schedule_name(name) escaped = escape_sql_string(name) + ts = self._oltp_sql.ts_text sql = ( - f"SELECT schedule_name, CAST(last_run_at AS STRING), CAST(next_run_at AS STRING), " + f"SELECT schedule_name, {ts('last_run_at')}, {ts('next_run_at')}, " f"last_run_id, status " f"FROM {self._table} WHERE schedule_name = '{escaped}'" ) - rows = self._sql.query(sql) + rows = self._oltp_sql.query(sql) if not rows: return None row = rows[0] @@ -311,28 +385,34 @@ def _upsert_tracker( last_run_id: str | None, status: str, ) -> None: - from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string, validate_schedule_name + from databricks_labs_dqx_app.backend.sql_utils import validate_schedule_name validate_schedule_name(name) if status not in _VALID_TRACKER_STATUSES: raise ValueError(f"Invalid tracker status: '{status}'. Must be one of {_VALID_TRACKER_STATUSES}") - escaped_name = escape_sql_string(name) - escaped_status = escape_sql_string(status) - last_str = f"'{last_run_at.isoformat()}'" if last_run_at else "NULL" - next_str = f"'{next_run_at.isoformat()}'" if next_run_at else "NULL" - run_id_str = f"'{escape_sql_string(last_run_id)}'" if last_run_id else "NULL" - sql = ( - f"MERGE INTO {self._table} AS target " - f"USING (SELECT '{escaped_name}' AS schedule_name) AS source " - "ON target.schedule_name = source.schedule_name " - "WHEN MATCHED THEN UPDATE SET " - f" last_run_at = {last_str}, next_run_at = {next_str}, " - f" last_run_id = {run_id_str}, status = '{escaped_status}' " - "WHEN NOT MATCHED THEN INSERT (schedule_name, last_run_at, next_run_at, last_run_id, status) " - f"VALUES ('{escaped_name}', {last_str}, {next_str}, {run_id_str}, '{escaped_status}')" + # Render datetimes as portable TIMESTAMP literals. The + # ``TIMESTAMP''`` form is ANSI SQL and works in both + # Delta and Postgres without modification. PgExecutor's + # upsert renderer treats ``RawSql("current_timestamp()")`` + # specially and rewrites it to ``CURRENT_TIMESTAMP`` so the + # same call works for both backends. + def _ts(dt: datetime | None) -> RawSql: + if dt is None: + return RawSql("NULL") + return RawSql(f"TIMESTAMP'{dt.isoformat()}'") + + self._oltp_sql.upsert( + self._table, + key_cols={"schedule_name": name}, + value_cols={ + "last_run_at": _ts(last_run_at), + "next_run_at": _ts(next_run_at), + "last_run_id": last_run_id, + "status": status, + "updated_at": RawSql("current_timestamp()"), + }, ) - self._sql.execute(sql) # ------------------------------------------------------------------ # Trigger run @@ -413,29 +493,95 @@ def _trigger_run(self, schedule_name: str, cfg: dict[str, Any], run_id_prefix: s # ------------------------------------------------------------------ def _resolve_scope(self, cfg: dict[str, Any]) -> list[str]: - """Return list of unique table_fqn matching the schedule's scope from approved rules.""" + """Return list of unique table_fqn matching the schedule's scope from approved rules. + + Two orthogonal filters intersected: + * ``scope_mode`` / ``scope_catalogs|schemas|tables`` — FQN-based. + * ``scope_labels`` — keep only FQNs that have at least one approved + check carrying a matching ``user_metadata`` label. + """ mode = cfg.get("scope_mode", "all") sql = f"SELECT DISTINCT table_fqn FROM {self._rules_table} WHERE status = 'approved'" - rows = self._sql.query(sql) + rows = self._oltp_sql.query(sql) fqns = [r[0] for r in rows if r[0]] - if mode == "all": - return fqns - if mode == "catalog": catalogs = set(cfg.get("scope_catalogs") or []) - return [f for f in fqns if self._fqn_part(f, 0) in catalogs] - - if mode == "schema": + fqns = [f for f in fqns if self._fqn_part(f, 0) in catalogs] + elif mode == "schema": schemas = set(cfg.get("scope_schemas") or []) - return [f for f in fqns if self._fqn_schema(f) in schemas] - - if mode == "tables": + fqns = [f for f in fqns if self._fqn_schema(f) in schemas] + elif mode == "tables": tables = set(cfg.get("scope_tables") or []) - return [f for f in fqns if f in tables] + fqns = [f for f in fqns if f in tables] + + label_filter = self._parse_scope_labels(cfg.get("scope_labels")) + if label_filter: + fqns = [f for f in fqns if self._fqn_has_matching_label(f, label_filter)] return fqns + @staticmethod + def _parse_scope_labels(raw: Any) -> set[tuple[str, str]]: + """Normalise the persisted ``scope_labels`` field to a set of (key, value). + + Accepts the canonical ``[{key, value}, ...]`` shape produced by the + UI plus a lenient ``["key=value", ...]`` shorthand for hand-edited + configs. Invalid entries are silently dropped — a malformed label + filter must never block a scheduled run. + """ + if not isinstance(raw, list): + return set() + out: set[tuple[str, str]] = set() + for entry in raw: + if isinstance(entry, dict): + key = entry.get("key") + if isinstance(key, str) and key: + out.add((key, str(entry.get("value") or ""))) + elif isinstance(entry, str): + if not entry: + continue + idx = entry.find("=") + if idx < 0: + out.add((entry, "")) + else: + out.add((entry[:idx], entry[idx + 1 :])) + return out + + def _fqn_has_matching_label( + self, + table_fqn: str, + label_filter: set[tuple[str, str]], + ) -> bool: + """True iff any approved check on ``table_fqn`` carries a matching label.""" + rule = self._get_approved_rule(table_fqn) + if rule is None: + return False + for check in rule.get("checks") or []: + md = self._check_user_metadata(check) + for key, value in md.items(): + if (key, value) in label_filter: + return True + return False + + @staticmethod + def _check_user_metadata(check: Any) -> dict[str, str]: + """Pull the ``user_metadata`` map off a check payload regardless of shape. + + Mirrors the front-end ``getUserMetadata`` helper — checks come in + either as a top-level dict with ``user_metadata`` directly on them, + or wrapped under a ``check`` key (legacy export shape). + """ + if not isinstance(check, dict): + return {} + candidate = check.get("user_metadata") + if not isinstance(candidate, dict): + inner = check.get("check") + candidate = inner.get("user_metadata") if isinstance(inner, dict) else None + if not isinstance(candidate, dict): + return {} + return {str(k): str(v) for k, v in candidate.items() if k} + @staticmethod def _fqn_part(fqn: str, idx: int) -> str: parts = fqn.split(".") @@ -451,24 +597,44 @@ def _fqn_schema(fqn: str) -> str: # ------------------------------------------------------------------ def _get_approved_rule(self, table_fqn: str) -> dict[str, Any] | None: - """Get merged checks from all approved rule rows for a table.""" + """Get merged checks from all approved rule rows for a table. + + After the v1 baseline split, each row stores a single check in + the VARIANT/JSONB ``check`` column rather than an array of + checks. The scheduler still presents one merged ``checks`` list + downstream (the task runner expects an array) so we collect + each row's bare object and append it. + """ from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string e_fqn = escape_sql_string(table_fqn) + check_col = self._oltp_sql.q("check") + # Dialect-agnostic JSON projection via the executor's + # :meth:`select_json_text` — ``to_json(col)`` on Delta, + # bare column on Postgres (PgExecutor._to_text JSON-encodes + # JSONB cells on the way out). + check_text = self._oltp_sql.select_json_text(check_col) sql = ( - f"SELECT table_fqn, checks FROM {self._rules_table} " f"WHERE table_fqn = '{e_fqn}' AND status = 'approved'" + f"SELECT table_fqn, {check_text} AS check_json FROM {self._rules_table} " + f"WHERE table_fqn = '{e_fqn}' AND status = 'approved'" ) - rows = self._sql.query(sql) + rows = self._oltp_sql.query(sql) if not rows: return None merged_checks: list[dict[str, Any]] = [] for row in rows: try: - parsed = json.loads(row[1], strict=False) if row[1] else [] - if isinstance(parsed, list): - merged_checks.extend(parsed) + parsed = json.loads(row[1], strict=False) if row[1] else None except json.JSONDecodeError: continue + if isinstance(parsed, dict): + merged_checks.append(parsed) + elif isinstance(parsed, list): + # Defensive: pre-baseline rows wrapped the single check + # in a one-element list. Tolerate that on read so a + # workspace that hasn't run ``DROP SCHEMA CASCADE`` + # against legacy data doesn't lose its rules. + merged_checks.extend([c for c in parsed if isinstance(c, dict)]) if not merged_checks: return None return {"table_fqn": rows[0][0], "checks": merged_checks} @@ -487,7 +653,7 @@ def _load_custom_metrics(self) -> list[str]: key = escape_sql_string("custom_metrics_v1") sql = f"SELECT setting_value FROM {self._settings_table} WHERE setting_key = '{key}'" # noqa: S608 - rows = self._sql.query(sql) + rows = self._oltp_sql.query(sql) if not rows or rows[0][0] is None: return [] parsed = json.loads(rows[0][0]) @@ -628,6 +794,132 @@ def _gc_orphan_views(self) -> None: skipped, ) + # ------------------------------------------------------------------ + # Retention — daily DELETE sweep against high-volume tables + # ------------------------------------------------------------------ + + def _resolve_retention_days(self) -> int: + """Return the configured retention window in days (>= 7). + + Looks up ``retention_days`` in ``dq_app_settings`` and falls back + to :data:`_RETENTION_DAYS_DEFAULT` (90 days) when unset or + unparseable. Capped at the lower bound :data:`_RETENTION_DAYS_MIN` + so a misconfiguration can't accidentally wipe live data. + """ + return self._resolve_setting_days("retention_days", _RETENTION_DAYS_DEFAULT) + + def _resolve_quarantine_retention_days(self) -> int: + """Return the quarantine-specific retention window in days (>= 7). + + Quarantine rows hold the full source row payload (PII surface) + so we maintain a separate, tighter default + (:data:`_QUARANTINE_RETENTION_DAYS_DEFAULT`, 30 days) than the + global retention. Configurable via ``quarantine_retention_days`` + in ``dq_app_settings``. Same min-floor protection as the global + resolver. + """ + return self._resolve_setting_days( + "quarantine_retention_days", + _QUARANTINE_RETENTION_DAYS_DEFAULT, + ) + + def _resolve_setting_days(self, key: str, default: int) -> int: + """Read an integer-day setting from ``dq_app_settings``. + + Shared parsing/floor logic for the global and quarantine + retention knobs. Any read or parse failure falls back to + *default*; the returned value is always >= :data:`_RETENTION_DAYS_MIN` + so a misconfiguration can never wipe data inside the safety floor. + """ + try: + from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string + + escaped_key = escape_sql_string(key) + sql = f"SELECT setting_value FROM {self._settings_table} WHERE setting_key = '{escaped_key}'" # noqa: S608 + rows = self._oltp_sql.query(sql) + if rows and rows[0] and rows[0][0]: + value = int(rows[0][0]) + return max(_RETENTION_DAYS_MIN, value) + except Exception: + logger.debug("Failed to read %s setting; falling back to default", key, exc_info=True) + return default + + async def _maybe_run_retention(self, now: datetime) -> None: + """Run the retention sweep if the daily timer has elapsed. + + Cheap to skip (one comparison) and runs in a background thread + so it doesn't block the loop. Failures are logged but never + fatal — the next tick re-tries. + """ + if now < self._next_retention_at: + return + + scheduled_for = self._next_retention_at + # Advance the timer first so a slow sweep can't double-fire. + self._next_retention_at = now + timedelta(hours=_RETENTION_INTERVAL_HOURS) + logger.info( + "Retention sweep: triggering daily cleanup (was due at %s); next run scheduled for %s", + scheduled_for.isoformat(), + self._next_retention_at.isoformat(), + ) + try: + await asyncio.to_thread(self._run_retention) + except Exception: + logger.exception("Retention sweep failed (non-fatal)") + + def _run_retention(self) -> None: + """DELETE rows older than ``retention_days`` from each high-volume table. + + Each table is processed independently — a failure on one + doesn't abort the others. The DELETE predicate uses an + INTERVAL literal so the backend stamps the cutoff against its + own clock (no Python-side time skew). + + Tables are split between the analytical Delta executor and + the OLTP executor (Lakebase or Delta-fallback) because the + ``INTERVAL`` syntax differs between dialects: Delta uses + ``INTERVAL N DAY`` (no quotes); Postgres uses + ``INTERVAL ' days'``. + """ + days = self._resolve_retention_days() + quarantine_days = self._resolve_quarantine_retention_days() + logger.info( + "Retention sweep: deleting rows older than %d days (quarantine: %d days)", + days, + quarantine_days, + ) + + total_deleted = 0 + # Delta tables — quoted with backticks so a future + # special-character schema name doesn't break the DELETE. + # ``dq_quarantine_records`` honours its own cutoff so PII row + # payloads can be aged out faster than the trend tables. + for table_name, time_col in _DELTA_RETENTION_TABLES: + table = f"`{self._catalog}`.`{self._schema}`.{table_name}" + cutoff = quarantine_days if table_name == _QUARANTINE_TABLE_NAME else days + stmt = f"DELETE FROM {table} " f"WHERE {time_col} < current_timestamp() - INTERVAL {cutoff} DAY" + try: + self._sql.execute(stmt) + logger.info("Retention sweep (Delta): cleaned %s (cutoff=%dd)", table_name, cutoff) + total_deleted += 1 + except Exception as exc: + logger.warning("Retention sweep: %s failed (%s); continuing", table_name, exc) + + # OLTP tables — fqn(), q(), and the INTERVAL literal are all + # delegated to the executor so the body stays dialect-agnostic. + interval = self._oltp_sql.interval_days_expr(days) + for table_name, time_col in _OLTP_RETENTION_TABLES: + table = self._oltp_sql.fqn(table_name) + stmt = f"DELETE FROM {table} " f"WHERE {time_col} < CURRENT_TIMESTAMP - {interval}" + try: + self._oltp_sql.execute(stmt) + logger.info("Retention sweep (OLTP): cleaned %s (cutoff=%dd)", table_name, days) + total_deleted += 1 + except Exception as exc: + logger.warning("Retention sweep: %s failed (%s); continuing", table_name, exc) + + logger.info("Retention sweep complete: %d table(s) processed", total_deleted) + # ------------------------------------------------------------------ # View creation (SP credentials) # ------------------------------------------------------------------ diff --git a/app/src/databricks_labs_dqx_app/backend/sql_executor.py b/app/src/databricks_labs_dqx_app/backend/sql_executor.py index 73a1e7f7b..9c2ec7505 100644 --- a/app/src/databricks_labs_dqx_app/backend/sql_executor.py +++ b/app/src/databricks_labs_dqx_app/backend/sql_executor.py @@ -10,23 +10,281 @@ import logging import time +from typing import Any, Protocol, runtime_checkable from databricks.sdk import WorkspaceClient from databricks.sdk.service.sql import Disposition, Format, StatementState +from databricks_labs_dqx_app.backend.sql_utils import escape_sql_string + logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# OLTP executor protocol — the shared surface SqlExecutor and PgExecutor +# both implement. +# --------------------------------------------------------------------------- + + +@runtime_checkable +class OltpExecutorProtocol(Protocol): + """Structural surface shared by Delta (:class:`SqlExecutor`) and Postgres + (``backend.pg_executor.PgExecutor``) executors used by OLTP services. + + Why this exists + --------------- + The previous design used ``Union[SqlExecutor, PgExecutor]`` plus + runtime ``cast(SqlExecutor, pg)`` calls at every OLTP boundary so + services with ``sql: SqlExecutor`` signatures kept typechecking + even though a :class:`PgExecutor` flowed through them at runtime. + The casts silently muted the type-checker for every Postgres path + — a method removed from :class:`PgExecutor` (or renamed only on + one side) wouldn't surface until production. + + A structural Protocol replaces both: + + - Services and the scheduler annotate their executor parameter as + :class:`OltpExecutorProtocol`. basedpyright then verifies every + ``self._sql.foo()`` call exists on the Protocol — and therefore + on both concrete classes — without any cast. + - :func:`backend.dependencies.get_sp_oltp_executor` returns the + Protocol so the FastAPI ``Depends`` graph stays type-honest. + - :class:`backend.services.scheduler_service.SchedulerService` + stores it as ``self._oltp_sql: OltpExecutorProtocol`` and drops + the ``cast(SqlExecutor, oltp_sql)`` line that previously hid the + Postgres path from type-checking. + + What the surface includes + ------------------------- + The Protocol intentionally captures only what OLTP services use. + Warehouse-specific affordances (``warehouse_id``, ``connection()``, + ``close()``) stay on the concrete classes — services that need + them keep their concrete-class annotations. + + ``@runtime_checkable`` + ---------------------- + Enables ``isinstance(executor, OltpExecutorProtocol)`` for tests + and adapters. Cheap to add; note that the runtime check verifies + method *presence* only, not signature compatibility. + """ + + # Routing identifier. ``"delta"`` on :class:`SqlExecutor`, + # ``"postgres"`` on :class:`PgExecutor`. Services that need + # dialect-specific SQL branch on this. + dialect: str + + @property + def catalog(self) -> str: + """The catalog/database the executor operates against. + + Delta: Unity Catalog name. Postgres: the logical database + (``PgExecutor.catalog`` aliases ``database`` so portable code + works on both backends). + """ + ... + + @property + def schema(self) -> str: + """The schema the executor operates against.""" + ... + + def fqn(self, table: str) -> str: + """Build a backend-specific fully-qualified path for *table*. + + Three parts on Delta (``catalog.schema.table``), two on + Postgres (``schema.table``). Always prefer this over + hand-rolled ``f"{x.catalog}.{x.schema}.{t}"`` so the dialect + branch stays in one place. + """ + ... + + def q(self, identifier: str) -> str: + """Dialect-aware identifier quoter (backticks on Delta, ANSI ``"`` on Postgres). + + Doubles the internal quoting character so reserved words and + identifiers containing the dialect's quote character (e.g. + a hyphenated UC catalog like ``team-data``) survive + interpolation into DDL. + """ + ... + + def json_literal_expr(self, json_str: str) -> str: + """Render *json_str* as a backend-native JSON literal expression. + + Delta wraps with :func:`parse_json`; Postgres uses ``::jsonb``. + Always returns a safe-to-inline expression (escaping baked in). + """ + ... + + def ts_text(self, col: str) -> str: + """Project a timestamp column as an ISO-string-safe expression. + + Delta wraps in ``CAST(... AS STRING)``; Postgres returns the + column verbatim because :func:`backend.pg_executor._to_text` + ISO-formats it during row coercion. + """ + ... + + def execute(self, sql: str, *, timeout_seconds: int = 120) -> None: + """Run a non-returning statement and commit it.""" + + def execute_no_schema(self, sql: str) -> None: + """Run DDL with catalog-only context (no default schema). + + Used for bootstrap DDL that must run before the target schema + exists. On Postgres this delegates to :meth:`execute`. + """ + + def query(self, sql: str, *, timeout_seconds: int = 120) -> list[list[str]]: + """Run a query and return rows as lists of stringified cells. + + Cells are strings (or ``None`` for NULL) to mirror the + Statement Execution API's ``JSON_ARRAY`` serialisation, so + services that ``json.loads`` JSON columns work on both + backends without dialect branches. + """ + ... + + def query_dicts(self, sql: str, *, timeout_seconds: int = 120) -> list[dict[str, str | None]]: + """Like :meth:`query` but rows are column-name-keyed dicts.""" + ... + + def upsert( + self, + table: str, + key_cols: dict[str, Any], + value_cols: dict[str, Any], + *, + timeout_seconds: int = 120, + ) -> None: + """Idempotently insert-or-update a single row keyed by *key_cols*. + + Delta uses ``MERGE INTO``; Postgres uses + ``INSERT ... ON CONFLICT ... DO UPDATE``. The signature is + identical so services don't need a dialect branch. + """ + + def upsert_with_audit( + self, + table: str, + key_cols: dict[str, Any], + value_cols: dict[str, Any], + *, + preserve_created: bool = True, + increment_on_update: str | None = None, + timeout_seconds: int = 120, + ) -> None: + """Audit-aware variant of :meth:`upsert` for tables with audit columns. + + Exists so service code stops branching on ``self._sql.dialect`` + to write ``MERGE INTO ...`` (Delta) or ``INSERT ... ON CONFLICT + ... DO UPDATE`` (Postgres) by hand. The two patterns are + equivalent — both upsert with audit semantics — so the dialect + choice lives on the executor, not on every caller. + + Behaves like :meth:`upsert` (single-row INSERT-or-UPDATE keyed + by *key_cols*) with two added knobs: + + - ``preserve_created`` (default True): any column in + *value_cols* whose name starts with ``created_`` is set on + INSERT but NOT touched on UPDATE. Mirrors the universal + audit convention: ``created_by`` and ``created_at`` are + immutable; ``updated_by`` / ``updated_at`` move on every + write. + - ``increment_on_update``: a column name. The column's + *initial value* on INSERT comes from *value_cols* (use + ``{"version": 1}`` for the conventional starting point); + on UPDATE the executor rewrites that column to ``existing + + 1`` using the dialect-appropriate self-reference syntax + (``target.col + 1`` on Delta MERGE, bare ``col + 1`` on + Postgres ON CONFLICT). Use for monotonic version counters + on history-tracked rows. + """ + + def select_json_text(self, col: str) -> str: + """Project a JSON-typed column as JSON-text on this backend. + + Delta's VARIANT column comes back as a non-string object via + the JSON_ARRAY response format, so callers that want to + ``json.loads`` need a ``to_json(col)`` wrapper. Postgres' + JSONB cells are coerced to JSON strings by + :func:`backend.pg_executor._to_text` during row conversion, + so the same projection just selects the column verbatim. + + Always prefer this over hand-writing ``to_json(col)`` — + scattering the wrapper across service code reintroduces + the dialect branch the executor is supposed to absorb. + """ + ... + + def interval_days_expr(self, days: int) -> str: + """Render an INTERVAL literal equivalent to *days* days. + + Composes into DATE/TIMESTAMP arithmetic — e.g. + ``CURRENT_TIMESTAMP - executor.interval_days_expr(90)`` is + "ninety days ago" on both backends. + + Delta speaks ``INTERVAL 90 DAY`` (no quotes, uppercase + singular); Postgres speaks ``INTERVAL '90 days'`` (single- + quoted literal, lowercase plural). The two are mutually + unparseable — Postgres throws ``syntax error at or near + "DAY"`` on the Delta form, and Delta throws on the quoted + form — so callers must never hand-build either. + """ + ... + + _TERMINAL_STATES = {StatementState.SUCCEEDED, StatementState.FAILED, StatementState.CANCELED, StatementState.CLOSED} +# Sentinel value for raw SQL expressions in upsert() — anything wrapped with +# ``RawSql("current_timestamp()")`` is interpolated verbatim instead of being +# escaped as a string literal. +class RawSql: + """Marker that wraps a raw SQL expression so :meth:`SqlExecutor.upsert` does not escape it.""" + + __slots__ = ("expr",) + + def __init__(self, expr: str) -> None: + self.expr = expr + + +def _render_value(value: Any) -> str: + """Convert a Python value to a SQL literal or raw expression. + + Strings are wrapped in single quotes and ANSI-escaped (`'` → `''`). + ``None`` becomes ``NULL``. Booleans are rendered as ``TRUE``/``FALSE``. + Numerics pass through. :class:`RawSql` is returned verbatim — this is + how callers inject ``current_timestamp()``, ``parse_json('...')``, etc. + """ + if isinstance(value, RawSql): + return value.expr + if value is None: + return "NULL" + if isinstance(value, bool): + return "TRUE" if value else "FALSE" + if isinstance(value, (int, float)): + return repr(value) + return f"'{escape_sql_string(str(value))}'" + + class SqlExecutor: """Thin wrapper around the Databricks Statement Execution API. Provides ``execute``, ``query``, ``query_dicts``, and ``execute_no_schema`` covering all usage patterns in the DQX Studio backend. + + The class also exposes a tiny dialect surface (:attr:`dialect`, + :meth:`q`, :meth:`json_literal_expr`) so that services which can + target either Delta or Lakebase Postgres can share the same SQL + while staying portable. See :class:`PgExecutor` for the Postgres + implementation. """ + # Cheap identifier for routing decisions in shared service code. + dialect: str = "delta" + def __init__( self, ws: WorkspaceClient, @@ -51,6 +309,65 @@ def catalog(self) -> str: def schema(self) -> str: return self._schema + # ------------------------------------------------------------------ + # Dialect helpers — kept identical-named on every executor so callers + # can hand-write portable SQL without an "if dialect" branch. + # ------------------------------------------------------------------ + def fqn(self, table: str) -> str: + """Return the fully-qualified path for *table* on this backend. + + Delta uses three-part names (``catalog.schema.table``) because + Unity Catalog organises tables under a catalog. Postgres has + only one catalog per connection so :class:`PgExecutor.fqn` drops + the catalog component and returns ``schema.table``. Services + that just need an addressable table identifier should call this + instead of inlining ``f"{sql.catalog}.{sql.schema}.{table}"`` — + the dialect branch lives here once instead of being repeated at + every call site. + """ + return f"{self._catalog}.{self._schema}.{table}" + + def q(self, identifier: str) -> str: + """Quote an identifier for this dialect. + + Delta uses backticks (``` `check` ```); Postgres uses double + quotes (``"check"``). Always quote columns/tables/catalogs/ + schemas that share a name with a reserved word (``check``, + ``user``, ``order``…) **or** that contain characters outside + ``[A-Za-z0-9_]`` — Databricks catalog names in particular are + commonly hyphenated (``prod-east``, ``team-data-platform``) + and would otherwise emit a parse error when interpolated raw + into DDL. + + Internal backticks in *identifier* are doubled per Databricks + SQL convention so a name like ``` weird`name ``` becomes + ``` `weird``name` ``` rather than producing two adjacent + terminators. This mirrors :meth:`PgExecutor.q` which doubles + internal double-quotes for the same reason. + """ + return "`" + identifier.replace("`", "``") + "`" + + def json_literal_expr(self, json_str: str) -> str: + """Return the SQL expression that turns *json_str* into a JSON value. + + Delta uses :func:`parse_json`; Postgres uses ``::jsonb``. The + returned expression is safe to inline into a larger statement + as it already includes the proper escaping. + """ + return f"parse_json('{escape_sql_string(json_str)}')" + + def ts_text(self, col: str) -> str: + """Project a timestamp column as an ISO-formatted string. + + Delta wraps the column in ``CAST(... AS STRING)`` because the + Statement Execution API otherwise returns timestamps as their + Spark string repr (which uses a space separator instead of + ``T``). On Postgres the column is selected verbatim and the + :class:`PgExecutor` row converter ISO-formats it on the way + out. + """ + return f"CAST({col} AS STRING)" + def execute(self, sql: str, *, timeout_seconds: int = 120) -> None: """Execute a SQL statement that does not return rows. @@ -187,6 +504,130 @@ def query_dicts(self, sql: str, *, timeout_seconds: int = 120) -> list[dict[str, ] return [dict(zip(columns, row)) for row in resp.result.data_array] + def upsert( + self, + table: str, + key_cols: dict[str, Any], + value_cols: dict[str, Any], + *, + timeout_seconds: int = 120, + ) -> None: + """Idempotently insert-or-update one row identified by *key_cols*. + + Centralises the ``MERGE INTO ... USING (SELECT)`` pattern that + otherwise gets duplicated across every service. *key_cols* are + the natural-key columns used in the ON clause; *value_cols* are + the non-key columns set on update / inserted on miss. + + Each value can be one of: + - ``str`` / ``int`` / ``float`` / ``bool`` / ``None`` — bound as + a SQL literal (strings are ANSI-escaped). + - :class:`RawSql` — interpolated verbatim, e.g. + ``RawSql("current_timestamp()")``. + + The full row written on insert is the union of *key_cols* and + *value_cols*. Updates only touch *value_cols* (the keys are + immutable by definition of the merge predicate). + """ + if not key_cols: + raise ValueError("upsert requires at least one key column") + + # Every column-name position is routed through :meth:`q` so a + # future audit column whose name is a Delta reserved word + # (``check``, ``order``, ``group``, …) survives interpolation + # into the MERGE. Mirrors the same quoting strategy that + # :meth:`PgExecutor.upsert` already applies on the Postgres + # side so the two backends stay symmetric — without this the + # Postgres path accepts a reserved-word column and Delta + # raises a parse error. + on_clause = " AND ".join(f"target.{self.q(k)} = source.{self.q(k)}" for k in key_cols) + source_select = ", ".join(f"{_render_value(v)} AS {self.q(k)}" for k, v in key_cols.items()) + update_set = ", ".join(f"{self.q(k)} = {_render_value(v)}" for k, v in value_cols.items()) + all_cols = list(key_cols.keys()) + list(value_cols.keys()) + all_vals = [_render_value(v) for v in list(key_cols.values()) + list(value_cols.values())] + insert_cols = ", ".join(self.q(c) for c in all_cols) + insert_vals = ", ".join(all_vals) + + sql = ( + f"MERGE INTO {table} AS target " + f"USING (SELECT {source_select}) AS source ON {on_clause} " + f"WHEN MATCHED THEN UPDATE SET {update_set} " + f"WHEN NOT MATCHED THEN INSERT ({insert_cols}) VALUES ({insert_vals})" + ) + self.execute(sql, timeout_seconds=timeout_seconds) + + def upsert_with_audit( + self, + table: str, + key_cols: dict[str, Any], + value_cols: dict[str, Any], + *, + preserve_created: bool = True, + increment_on_update: str | None = None, + timeout_seconds: int = 120, + ) -> None: + """Delta MERGE INTO with audit semantics. + + See :meth:`OltpExecutorProtocol.upsert_with_audit` for the + contract. ``target`` is the alias used in the MERGE so the + increment self-reference renders as ``target. + 1``. + """ + if not key_cols: + raise ValueError("upsert_with_audit requires at least one key column") + if increment_on_update is not None and increment_on_update not in value_cols: + raise ValueError( + f"increment_on_update={increment_on_update!r} must be present in value_cols " + "with its initial INSERT value (e.g. {'version': 1})" + ) + + # Every column-name position is routed through :meth:`q` so a + # future audit column whose name is a Delta reserved word + # (``check``, ``order``, ``group``, …) survives interpolation + # into the MERGE. Mirrors the same quoting strategy that + # :meth:`PgExecutor.upsert_with_audit` already applies on the + # Postgres side so the two backends stay symmetric — without + # this the Postgres path accepts a reserved-word audit column + # and Delta raises a parse error. + on_clause = " AND ".join(f"target.{self.q(k)} = source.{self.q(k)}" for k in key_cols) + source_select = ", ".join(f"{_render_value(v)} AS {self.q(k)}" for k, v in key_cols.items()) + + # UPDATE SET excludes created_* columns when preserve_created; + # the increment column (if any) gets the dialect-specific + # self-reference form instead of the literal value. + update_pairs: list[str] = [] + for col, val in value_cols.items(): + if preserve_created and col.startswith("created_"): + continue + qcol = self.q(col) + if col == increment_on_update: + # Delta MERGE — qualify with the ``target`` alias to + # disambiguate from the source row. + update_pairs.append(f"{qcol} = target.{qcol} + 1") + else: + update_pairs.append(f"{qcol} = {_render_value(val)}") + update_set = ", ".join(update_pairs) + + all_cols = list(key_cols.keys()) + list(value_cols.keys()) + all_vals = [_render_value(v) for v in list(key_cols.values()) + list(value_cols.values())] + insert_cols = ", ".join(self.q(c) for c in all_cols) + insert_vals = ", ".join(all_vals) + + sql = ( + f"MERGE INTO {table} AS target " + f"USING (SELECT {source_select}) AS source ON {on_clause} " + f"WHEN MATCHED THEN UPDATE SET {update_set} " + f"WHEN NOT MATCHED THEN INSERT ({insert_cols}) VALUES ({insert_vals})" + ) + self.execute(sql, timeout_seconds=timeout_seconds) + + def select_json_text(self, col: str) -> str: + """Delta needs ``to_json`` to serialise VARIANT through JSON_ARRAY.""" + return f"to_json({col})" + + def interval_days_expr(self, days: int) -> str: + """Delta interval literal — bare integer, uppercase ``DAY`` (singular).""" + return f"INTERVAL {int(days)} DAY" + def _wait_for_completion(self, statement_id: str, timeout_seconds: int) -> StatementState: """Poll statement status until it reaches a terminal state.""" start = time.time() diff --git a/app/src/databricks_labs_dqx_app/ui/components/DryRunResults.tsx b/app/src/databricks_labs_dqx_app/ui/components/DryRunResults.tsx index 822edfd6f..b3e7c068f 100644 --- a/app/src/databricks_labs_dqx_app/ui/components/DryRunResults.tsx +++ b/app/src/databricks_labs_dqx_app/ui/components/DryRunResults.tsx @@ -93,7 +93,20 @@ function downloadAsExcel(rows: Record[], filename: string): voi export function DryRunResults({ result }: DryRunResultsProps) { const totalRows = result.total_rows ?? 0; const validRows = result.valid_rows ?? 0; - const invalidRows = result.invalid_rows ?? 0; + // ``invalid_rows`` (rows that failed any check) is kept on the model for + // backwards compatibility, but we surface DQX's authoritative observer + // counts instead: ``error_rows`` (= ``error_row_count``) and + // ``warning_rows`` (= ``warning_row_count``). Pre-v5 history rows have + // ``error_rows = null`` — fall back to ``invalid_rows`` so the card + // doesn't show ``0`` for runs created before the rename. + const errorRows = result.error_rows ?? result.invalid_rows ?? 0; + const warningRows = result.warning_rows; + const showWarnings = warningRows != null; + // The bottom data table still drives off the quarantine endpoint, which + // contains rows that failed any check (errors OR warnings). We use + // ``hasFailedRows`` so the table renders whenever there's anything to + // show — not just when there are errors. + const hasFailedRows = errorRows > 0 || (warningRows ?? 0) > 0; const errorSummary = result.error_summary ?? []; const sampleInvalid = result.sample_invalid ?? []; @@ -111,6 +124,7 @@ export function DryRunResults({ result }: DryRunResultsProps) { return { ...rest, errors: Array.isArray(_errors) ? _errors.map(formatError).join("; ") : String(_errors ?? ""), + warnings: Array.isArray(_warnings) ? _warnings.map(formatError).join("; ") : String(_warnings ?? ""), }; }); }, [sampleInvalid]); @@ -129,10 +143,10 @@ export function DryRunResults({ result }: DryRunResultsProps) { const { data: quarantineResp, isLoading: quarantineLoading } = useListQuarantineRecords(result.run_id, { offset, limit: pageSize }, { - query: { enabled: invalidRows > 0 }, + query: { enabled: hasFailedRows }, }); const { data: countResp } = useQuarantineCount(result.run_id, { - query: { enabled: invalidRows > 0 }, + query: { enabled: hasFailedRows }, }); const quarantineRecords: QuarantineRecordOut[] = quarantineResp?.data?.records ?? []; @@ -147,11 +161,19 @@ export function DryRunResults({ result }: DryRunResultsProps) { return sampleInvalid.slice(offset, offset + pageSize); }, [hasQuarantine, sampleInvalid, offset, pageSize]); - const rows: Array<{ rowData: Record; errors: unknown[] }> = useMemo(() => { + const rows: Array<{ + rowData: Record; + errors: unknown[]; + warnings: unknown[]; + }> = useMemo(() => { if (hasQuarantine) { return displayRecords.map((r) => ({ rowData: r.row_data ?? {}, errors: r.errors ?? [], + // ``warnings`` is null for rows written before migration v4 and + // for SQL-check quarantines. Render an empty array so the column + // shows nothing rather than ``null``. + warnings: r.warnings ?? [], })); } return fallbackRows.map((r) => { @@ -159,10 +181,19 @@ export function DryRunResults({ result }: DryRunResultsProps) { return { rowData: rest, errors: Array.isArray(_errors) ? (_errors as unknown[]) : [], + warnings: Array.isArray(_warnings) ? (_warnings as unknown[]) : [], }; }); }, [hasQuarantine, displayRecords, fallbackRows]); + // Only render the Warnings column when at least one displayed row has a + // warning payload. Pre-v4 quarantine rows have ``null`` warnings — hiding + // the column keeps the table compact for runs that have only errors. + const hasAnyWarning = useMemo( + () => rows.some((r) => Array.isArray(r.warnings) && r.warnings.length > 0), + [rows], + ); + const dataColumns = useMemo(() => { const keys = new Set(); for (const r of rows) { @@ -175,8 +206,12 @@ export function DryRunResults({ result }: DryRunResultsProps) { return (
- {/* Summary stats */} -
+ {/* ``Errors`` and ``Warnings`` are independent buckets — a row can be in + both. Both come from DQX's observer (``error_row_count`` / + ``warning_row_count``) so they're bounded by the input row count + and don't suffer from the fan-out artefacts that ``invalid_rows`` + can have on certain check types. */} +
{totalRows}
Total Rows
@@ -189,12 +224,21 @@ export function DryRunResults({ result }: DryRunResultsProps) {
-
{invalidRows}
+
{errorRows}
- Invalid + Errors
+ {showWarnings && ( +
+
{warningRows}
+
+ + Warnings +
+
+ )}
{/* Pass rate bar */} @@ -213,12 +257,17 @@ export function DryRunResults({ result }: DryRunResultsProps) {
- {/* Error summary */} + {/* Check failure summary — each row is one approved check that + produced at least one error or warning. We split the total + into ``error_count`` / ``warning_count`` columns so a + warning-level check is visually distinct from an error-level + one (otherwise warning-only checks look identical to errors, + which is what users hit in practice). */} {errorSummary.length > 0 && (

- Error Summary + Failed checks ({errorSummary.length} distinct) @@ -227,23 +276,44 @@ export function DryRunResults({ result }: DryRunResultsProps) { - - + + + - {visibleErrorSummary.map((item, idx) => ( - - - - - ))} + {visibleErrorSummary.map((item, idx) => { + const errCount = Number(item.error_count ?? 0); + const warnCount = Number(item.warning_count ?? 0); + return ( + + + + + + ); + })}
ErrorCountCheckErrorsWarnings
- {summarizeErrorText(String(item.error ?? ""))} - - - {String(item.count ?? 0)} - -
+ {summarizeErrorText(String(item.error ?? ""))} + + {errCount > 0 ? ( + + {errCount} + + ) : ( + — + )} + + {warnCount > 0 ? ( + + {warnCount} + + ) : ( + — + )} +
{hiddenErrorCount > 0 && ( @@ -258,13 +328,15 @@ export function DryRunResults({ result }: DryRunResultsProps) {

)} - {/* Invalid rows data table */} - {invalidRows > 0 && ( + {/* Failed rows data table — includes both error rows and warning rows + since DQX's split puts anything that failed a check into the same + quarantine bucket. */} + {hasFailedRows && (

- - Invalid Rows + + Failed rows ({hasQuarantine ? `${quarantineTotal} quarantined` : `${sampleInvalid.length} samples`}) @@ -333,7 +405,7 @@ export function DryRunResults({ result }: DryRunResultsProps) { {quarantineLoading ? (
Loading quarantine data...
) : rows.length === 0 ? ( -
No invalid rows to display
+
No failed rows to display
) : ( @@ -345,6 +417,9 @@ export function DryRunResults({ result }: DryRunResultsProps) { ))} + {hasAnyWarning && ( + + )} @@ -365,6 +440,21 @@ export function DryRunResults({ result }: DryRunResultsProps) { ))} + {hasAnyWarning && ( + + )} ))} diff --git a/app/src/databricks_labs_dqx_app/ui/lib/api-custom.ts b/app/src/databricks_labs_dqx_app/ui/lib/api-custom.ts index a5779b944..0130fe9d8 100644 --- a/app/src/databricks_labs_dqx_app/ui/lib/api-custom.ts +++ b/app/src/databricks_labs_dqx_app/ui/lib/api-custom.ts @@ -174,6 +174,8 @@ export interface ValidationRunSummaryOut { total_rows: number | null; valid_rows: number | null; invalid_rows: number | null; + error_rows: number | null; + warning_rows: number | null; created_at: string | null; error_message: string | null; checks: Record[]; @@ -567,6 +569,7 @@ export interface QuarantineRecordOut { requesting_user: string | null; row_data: Record | null; errors: unknown[] | null; + warnings: unknown[] | null; created_at: string | null; } @@ -891,3 +894,83 @@ export const useSaveLabelDefinitions = < ...mutationOptions, }); }; + +// --------------------------------------------------------------------------- +// Retention settings — global vs. quarantine-specific DELETE windows +// surfaced for the admin Configuration page. Mirrors +// ``backend/routes/v1/config.py``. +// --------------------------------------------------------------------------- + +export interface RetentionSettingsOut { + retention_days: number; + quarantine_retention_days: number; + retention_days_default: number; + quarantine_retention_days_default: number; + retention_days_min: number; + retention_days_max: number; + retention_days_set: boolean; + quarantine_retention_days_set: boolean; +} + +export interface RetentionSettingsIn { + retention_days?: number | null; + quarantine_retention_days?: number | null; +} + +export const getRetentionSettings = ( + options?: AxiosRequestConfig, +): Promise> => + axios.default.get("/api/v1/config/retention", options); + +export const getRetentionSettingsQueryKey = () => ["retention-settings"] as const; + +export const useRetentionSettings = < + TData = Awaited>["data"], + TError = AxiosError, +>( + options?: { + query?: Partial>, TError, TData>>; + axios?: AxiosRequestConfig; + }, +): UseQueryResult => { + const { query: queryOptions, axios: axiosOptions } = options ?? {}; + return useQuery({ + queryKey: queryOptions?.queryKey ?? getRetentionSettingsQueryKey(), + queryFn: () => getRetentionSettings(axiosOptions), + select: ((resp: Awaited>) => resp.data) as never, + staleTime: 5 * 60 * 1000, + ...queryOptions, + }) as UseQueryResult; +}; + +export const saveRetentionSettings = ( + body: RetentionSettingsIn, + options?: AxiosRequestConfig, +): Promise> => + axios.default.put("/api/v1/config/retention", body, options); + +export const useSaveRetentionSettings = < + TError = AxiosError, + TContext = unknown, +>( + options?: { + mutation?: UseMutationOptions< + Awaited>, + TError, + { data: RetentionSettingsIn }, + TContext + >; + axios?: AxiosRequestConfig; + }, +): UseMutationResult< + Awaited>, + TError, + { data: RetentionSettingsIn }, + TContext +> => { + const { mutation: mutationOptions, axios: axiosOptions } = options ?? {}; + return useMutation({ + mutationFn: ({ data }: { data: RetentionSettingsIn }) => saveRetentionSettings(data, axiosOptions), + ...mutationOptions, + }); +}; diff --git a/app/src/databricks_labs_dqx_app/ui/lib/api.ts b/app/src/databricks_labs_dqx_app/ui/lib/api.ts index a4eafc60a..ea29b040f 100644 --- a/app/src/databricks_labs_dqx_app/ui/lib/api.ts +++ b/app/src/databricks_labs_dqx_app/ui/lib/api.ts @@ -309,6 +309,10 @@ export type DryRunResultsOutValidRows = number | null; export type DryRunResultsOutInvalidRows = number | null; +export type DryRunResultsOutErrorRows = number | null; + +export type DryRunResultsOutWarningRows = number | null; + export type DryRunResultsOutErrorSummaryItem = { [key: string]: unknown }; export type DryRunResultsOutSampleInvalidItem = { [key: string]: unknown }; @@ -319,6 +323,8 @@ export interface DryRunResultsOut { total_rows?: DryRunResultsOutTotalRows; valid_rows?: DryRunResultsOutValidRows; invalid_rows?: DryRunResultsOutInvalidRows; + error_rows?: DryRunResultsOutErrorRows; + warning_rows?: DryRunResultsOutWarningRows; error_summary?: DryRunResultsOutErrorSummaryItem[]; sample_invalid?: DryRunResultsOutSampleInvalidItem[]; } @@ -689,6 +695,8 @@ export type QuarantineRecordOutRowData = QuarantineRecordOutRowDataAnyOf | null; export type QuarantineRecordOutErrors = unknown[] | null; +export type QuarantineRecordOutWarnings = unknown[] | null; + export type QuarantineRecordOutCreatedAt = string | null; export interface QuarantineRecordOut { @@ -698,9 +706,42 @@ export interface QuarantineRecordOut { requesting_user?: QuarantineRecordOutRequestingUser; row_data?: QuarantineRecordOutRowData; errors?: QuarantineRecordOutErrors; + warnings?: QuarantineRecordOutWarnings; created_at?: QuarantineRecordOutCreatedAt; } +export type RetentionSettingsInRetentionDays = number | null; + +export type RetentionSettingsInQuarantineRetentionDays = number | null; + +/** + * Update payload — either field omitted means *leave unchanged*. + */ +export interface RetentionSettingsIn { + retention_days?: RetentionSettingsInRetentionDays; + quarantine_retention_days?: RetentionSettingsInQuarantineRetentionDays; +} + +/** + * Effective retention settings + the defaults the scheduler falls back to. + +``retention_days`` / ``quarantine_retention_days`` reflect the +*current effective values* — the persisted setting if one exists, +otherwise the compiled-in default. The ``*_default`` and ``*_min`` +fields let the UI render hints and validation without duplicating +the constants on the frontend. + */ +export interface RetentionSettingsOut { + retention_days: number; + quarantine_retention_days: number; + retention_days_default?: number; + quarantine_retention_days_default?: number; + retention_days_min?: number; + retention_days_max?: number; + retention_days_set: boolean; + quarantine_retention_days_set: boolean; +} + export type RoleMappingOutCreatedBy = string | null; export type RoleMappingOutCreatedAt = string | null; @@ -1026,6 +1067,10 @@ export type ValidationRunSummaryOutValidRows = number | null; export type ValidationRunSummaryOutInvalidRows = number | null; +export type ValidationRunSummaryOutErrorRows = number | null; + +export type ValidationRunSummaryOutWarningRows = number | null; + export type ValidationRunSummaryOutCreatedAt = string | null; export type ValidationRunSummaryOutErrorMessage = string | null; @@ -1044,6 +1089,8 @@ export interface ValidationRunSummaryOut { run_type?: ValidationRunSummaryOutRunType; valid_rows?: ValidationRunSummaryOutValidRows; invalid_rows?: ValidationRunSummaryOutInvalidRows; + error_rows?: ValidationRunSummaryOutErrorRows; + warning_rows?: ValidationRunSummaryOutWarningRows; created_at?: ValidationRunSummaryOutCreatedAt; error_message?: ValidationRunSummaryOutErrorMessage; checks?: ValidationRunSummaryOutChecksItem[]; @@ -3173,6 +3220,373 @@ export const useSaveTimezone = < return useMutation(mutationOptions, queryClient); }; +/** + * Return the current retention windows + defaults (admin only). + * @summary Get Retention Settings + */ +export const getRetentionSettings = ( + options?: AxiosRequestConfig, +): Promise> => { + return axios.default.get(`/api/v1/config/retention`, options); +}; + +export const getGetRetentionSettingsQueryKey = () => { + return [`/api/v1/config/retention`] as const; +}; + +export const getGetRetentionSettingsQueryOptions = < + TData = Awaited>, + TError = AxiosError, +>(options?: { + query?: Partial< + UseQueryOptions< + Awaited>, + TError, + TData + > + >; + axios?: AxiosRequestConfig; +}) => { + const { query: queryOptions, axios: axiosOptions } = options ?? {}; + + const queryKey = queryOptions?.queryKey ?? getGetRetentionSettingsQueryKey(); + + const queryFn: QueryFunction< + Awaited> + > = ({ signal }) => getRetentionSettings({ signal, ...axiosOptions }); + + return { queryKey, queryFn, ...queryOptions } as UseQueryOptions< + Awaited>, + TError, + TData + > & { queryKey: DataTag }; +}; + +export type GetRetentionSettingsQueryResult = NonNullable< + Awaited> +>; +export type GetRetentionSettingsQueryError = AxiosError; + +export function useGetRetentionSettings< + TData = Awaited>, + TError = AxiosError, +>( + options: { + query: Partial< + UseQueryOptions< + Awaited>, + TError, + TData + > + > & + Pick< + DefinedInitialDataOptions< + Awaited>, + TError, + Awaited> + >, + "initialData" + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): DefinedUseQueryResult & { + queryKey: DataTag; +}; +export function useGetRetentionSettings< + TData = Awaited>, + TError = AxiosError, +>( + options?: { + query?: Partial< + UseQueryOptions< + Awaited>, + TError, + TData + > + > & + Pick< + UndefinedInitialDataOptions< + Awaited>, + TError, + Awaited> + >, + "initialData" + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): UseQueryResult & { + queryKey: DataTag; +}; +export function useGetRetentionSettings< + TData = Awaited>, + TError = AxiosError, +>( + options?: { + query?: Partial< + UseQueryOptions< + Awaited>, + TError, + TData + > + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): UseQueryResult & { + queryKey: DataTag; +}; +/** + * @summary Get Retention Settings + */ + +export function useGetRetentionSettings< + TData = Awaited>, + TError = AxiosError, +>( + options?: { + query?: Partial< + UseQueryOptions< + Awaited>, + TError, + TData + > + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): UseQueryResult & { + queryKey: DataTag; +} { + const queryOptions = getGetRetentionSettingsQueryOptions(options); + + const query = useQuery(queryOptions, queryClient) as UseQueryResult< + TData, + TError + > & { queryKey: DataTag }; + + query.queryKey = queryOptions.queryKey; + + return query; +} + +export const getGetRetentionSettingsSuspenseQueryOptions = < + TData = Awaited>, + TError = AxiosError, +>(options?: { + query?: Partial< + UseSuspenseQueryOptions< + Awaited>, + TError, + TData + > + >; + axios?: AxiosRequestConfig; +}) => { + const { query: queryOptions, axios: axiosOptions } = options ?? {}; + + const queryKey = queryOptions?.queryKey ?? getGetRetentionSettingsQueryKey(); + + const queryFn: QueryFunction< + Awaited> + > = ({ signal }) => getRetentionSettings({ signal, ...axiosOptions }); + + return { queryKey, queryFn, ...queryOptions } as UseSuspenseQueryOptions< + Awaited>, + TError, + TData + > & { queryKey: DataTag }; +}; + +export type GetRetentionSettingsSuspenseQueryResult = NonNullable< + Awaited> +>; +export type GetRetentionSettingsSuspenseQueryError = + AxiosError; + +export function useGetRetentionSettingsSuspense< + TData = Awaited>, + TError = AxiosError, +>( + options: { + query: Partial< + UseSuspenseQueryOptions< + Awaited>, + TError, + TData + > + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): UseSuspenseQueryResult & { + queryKey: DataTag; +}; +export function useGetRetentionSettingsSuspense< + TData = Awaited>, + TError = AxiosError, +>( + options?: { + query?: Partial< + UseSuspenseQueryOptions< + Awaited>, + TError, + TData + > + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): UseSuspenseQueryResult & { + queryKey: DataTag; +}; +export function useGetRetentionSettingsSuspense< + TData = Awaited>, + TError = AxiosError, +>( + options?: { + query?: Partial< + UseSuspenseQueryOptions< + Awaited>, + TError, + TData + > + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): UseSuspenseQueryResult & { + queryKey: DataTag; +}; +/** + * @summary Get Retention Settings + */ + +export function useGetRetentionSettingsSuspense< + TData = Awaited>, + TError = AxiosError, +>( + options?: { + query?: Partial< + UseSuspenseQueryOptions< + Awaited>, + TError, + TData + > + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): UseSuspenseQueryResult & { + queryKey: DataTag; +} { + const queryOptions = getGetRetentionSettingsSuspenseQueryOptions(options); + + const query = useSuspenseQuery( + queryOptions, + queryClient, + ) as UseSuspenseQueryResult & { + queryKey: DataTag; + }; + + query.queryKey = queryOptions.queryKey; + + return query; +} + +/** + * Update one or both retention windows (admin only). + +Either field may be omitted to leave the existing value unchanged. +Both values are validated against the safety floor and ceiling +before being persisted. + * @summary Save Retention Settings + */ +export const saveRetentionSettings = ( + retentionSettingsIn: RetentionSettingsIn, + options?: AxiosRequestConfig, +): Promise> => { + return axios.default.put( + `/api/v1/config/retention`, + retentionSettingsIn, + options, + ); +}; + +export const getSaveRetentionSettingsMutationOptions = < + TError = AxiosError, + TContext = unknown, +>(options?: { + mutation?: UseMutationOptions< + Awaited>, + TError, + { data: RetentionSettingsIn }, + TContext + >; + axios?: AxiosRequestConfig; +}): UseMutationOptions< + Awaited>, + TError, + { data: RetentionSettingsIn }, + TContext +> => { + const mutationKey = ["saveRetentionSettings"]; + const { mutation: mutationOptions, axios: axiosOptions } = options + ? options.mutation && + "mutationKey" in options.mutation && + options.mutation.mutationKey + ? options + : { ...options, mutation: { ...options.mutation, mutationKey } } + : { mutation: { mutationKey }, axios: undefined }; + + const mutationFn: MutationFunction< + Awaited>, + { data: RetentionSettingsIn } + > = (props) => { + const { data } = props ?? {}; + + return saveRetentionSettings(data, axiosOptions); + }; + + return { mutationFn, ...mutationOptions }; +}; + +export type SaveRetentionSettingsMutationResult = NonNullable< + Awaited> +>; +export type SaveRetentionSettingsMutationBody = RetentionSettingsIn; +export type SaveRetentionSettingsMutationError = + AxiosError; + +/** + * @summary Save Retention Settings + */ +export const useSaveRetentionSettings = < + TError = AxiosError, + TContext = unknown, +>( + options?: { + mutation?: UseMutationOptions< + Awaited>, + TError, + { data: RetentionSettingsIn }, + TContext + >; + axios?: AxiosRequestConfig; + }, + queryClient?: QueryClient, +): UseMutationResult< + Awaited>, + TError, + { data: RetentionSettingsIn }, + TContext +> => { + const mutationOptions = getSaveRetentionSettingsMutationOptions(options); + + return useMutation(mutationOptions, queryClient); +}; + /** * Return all admin-defined label definitions. diff --git a/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/config.tsx b/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/config.tsx index ec53be77c..997467709 100644 --- a/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/config.tsx +++ b/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/config.tsx @@ -3,7 +3,7 @@ import { QueryErrorResetBoundary, useQueryClient } from "@tanstack/react-query"; import { ErrorBoundary } from "react-error-boundary"; import { Button } from "@/components/ui/button"; import { PageBreadcrumb } from "@/components/apx/PageBreadcrumb"; -import { AlertCircle, Globe, Loader2, Search, Tags, Plus, Trash2, X } from "lucide-react"; +import { AlertCircle, Clock, Globe, Loader2, Search, Tags, Plus, Trash2, X } from "lucide-react"; import { FadeIn } from "@/components/anim/FadeIn"; import { ShinyText } from "@/components/anim/ShinyText"; import { RoleManagement } from "@/components/RoleManagement"; @@ -25,7 +25,11 @@ import { useLabelDefinitions, useSaveLabelDefinitions, getLabelDefinitionsQueryKey, + useRetentionSettings, + useSaveRetentionSettings, + getRetentionSettingsQueryKey, type LabelDefinition, + type RetentionSettingsOut, } from "@/lib/api-custom"; import type { AxiosError } from "axios"; import { toast } from "sonner"; @@ -590,6 +594,217 @@ function DefinitionEditorCard({ ); } +// ───────────────────────────────────────────────────────────────────────────── +// Retention Settings — admin-controlled DELETE windows for the daily sweep. +// Two knobs: a global retention applied to dq_validation_runs, dq_metrics, +// dq_profiling_results + the OLTP history tables; and a tighter +// quarantine-specific retention applied only to dq_quarantine_records (which +// holds full source row payloads + errors/warnings). The split exists so PII +// can age out faster than trend tables that the dashboards look back on. +// ───────────────────────────────────────────────────────────────────────────── + +function RetentionSettings() { + const { data, isLoading } = useRetentionSettings(); + const queryClient = useQueryClient(); + const saveMutation = useSaveRetentionSettings(); + const { data: role } = useCurrentUserRoleSuspense(); + const isAdmin = role?.data?.role === "admin"; + + const settings = data as RetentionSettingsOut | undefined; + const [global, setGlobal] = useState(""); + const [quarantine, setQuarantine] = useState(""); + const [hydrated, setHydrated] = useState(false); + + useEffect(() => { + if (settings && !hydrated) { + setGlobal(String(settings.retention_days)); + setQuarantine(String(settings.quarantine_retention_days)); + setHydrated(true); + } + }, [settings, hydrated]); + + const min = settings?.retention_days_min ?? 7; + const max = settings?.retention_days_max ?? 3650; + + const parsedGlobal = Number.parseInt(global, 10); + const parsedQuarantine = Number.parseInt(quarantine, 10); + + const validation = useMemo(() => { + const errors: string[] = []; + const check = (label: string, value: number) => { + if (Number.isNaN(value)) { + errors.push(`${label} must be a whole number of days.`); + return; + } + if (value < min) errors.push(`${label} must be at least ${min} days.`); + if (value > max) errors.push(`${label} must be at most ${max} days.`); + }; + check("Global retention", parsedGlobal); + check("Quarantine retention", parsedQuarantine); + return errors; + }, [parsedGlobal, parsedQuarantine, min, max]); + + const isDirty = useMemo(() => { + if (!settings) return false; + return ( + parsedGlobal !== settings.retention_days || + parsedQuarantine !== settings.quarantine_retention_days + ); + }, [settings, parsedGlobal, parsedQuarantine]); + + const handleSave = () => { + if (!settings || validation.length > 0) return; + const payload: { retention_days?: number; quarantine_retention_days?: number } = {}; + if (parsedGlobal !== settings.retention_days) payload.retention_days = parsedGlobal; + if (parsedQuarantine !== settings.quarantine_retention_days) { + payload.quarantine_retention_days = parsedQuarantine; + } + saveMutation.mutate( + { data: payload }, + { + onSuccess: (resp) => { + queryClient.invalidateQueries({ queryKey: getRetentionSettingsQueryKey() }); + setGlobal(String(resp.data.retention_days)); + setQuarantine(String(resp.data.quarantine_retention_days)); + toast.success("Retention settings saved."); + }, + onError: (err: unknown) => { + const axErr = err as AxiosError<{ detail?: string }>; + toast.error(axErr?.response?.data?.detail ?? "Failed to save retention settings."); + }, + }, + ); + }; + + const handleReset = () => { + if (!settings) return; + setGlobal(String(settings.retention_days)); + setQuarantine(String(settings.quarantine_retention_days)); + }; + + const resetToDefaults = () => { + if (!settings) return; + setGlobal(String(settings.retention_days_default)); + setQuarantine(String(settings.quarantine_retention_days_default)); + }; + + if (isLoading || !settings) return ; + + return ( + + + + + Data Retention + + + +

+ The scheduler runs a daily DELETE pass against the analytical tables. + Quarantine holds the full source + row payload (errors, warnings, and the row itself) so its window is kept + tighter than the trend tables by default. Both values are floored at{" "} + {min} days to protect against accidental data loss. +

+ +
+
+ + setGlobal(e.target.value)} + className="h-8" + /> +

+ Applies to dq_validation_runs, dq_profiling_results,{" "} + dq_metrics, and the OLTP history tables. +
+ Default: {settings.retention_days_default} days + {!settings.retention_days_set && " (not yet customised)"} +

+
+ +
+ + setQuarantine(e.target.value)} + className="h-8" + /> +

+ Applies only to dq_quarantine_records (the table that + stores per-row failures, including the source row payload). +
+ Default: {settings.quarantine_retention_days_default} days + {!settings.quarantine_retention_days_set && " (not yet customised)"} +

+
+
+ + {validation.length > 0 && ( +
+ {validation.map((msg, i) => ( +

+ + {msg} +

+ ))} +
+ )} + +
+ + + + {!isAdmin && ( + + Only admins can change retention. + + )} +
+ + + ); +} + function ConfigPage() { const { isAdmin } = usePermissions(); const navigate = useNavigate(); @@ -636,6 +851,13 @@ function ConfigPage() { + + }> + + + + + diff --git a/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/runs-history.tsx b/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/runs-history.tsx index 21fce4727..722a05dbc 100644 --- a/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/runs-history.tsx +++ b/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/runs-history.tsx @@ -165,7 +165,7 @@ function RunsHistoryPage() { ); } -type RunsSortKey = "table" | "type" | "status" | "requested_by" | "total" | "valid" | "invalid" | "run_date"; +type RunsSortKey = "table" | "type" | "status" | "requested_by" | "total" | "valid" | "errors" | "warnings" | "run_date"; function RunHistoryContent() { const { data: currentUser } = useCurrentUserSuspense(selector()); @@ -178,7 +178,11 @@ function RunHistoryContent() { const [tableSearch, setTableSearch] = useState(""); const [statusFilter, setStatusFilter] = useState("all"); const [runTypeFilter, setRunTypeFilter] = useState("all"); - const [invalidOnly, setInvalidOnly] = useState(false); + // ``failedOnly`` keeps a row if it has either errors or warnings — the + // user-facing "Has failures" toggle that replaced the old "Has invalid" + // filter. We still tolerate ``error_rows`` being ``null`` on pre-v5 rows + // by falling back to ``invalid_rows``. + const [failedOnly, setFailedOnly] = useState(false); const [myRunsOnly, setMyRunsOnly] = useState(false); const [labelFilter, setLabelFilter] = useState>(new Set()); const [expandedRunId, setExpandedRunId] = useState(null); @@ -285,7 +289,12 @@ function RunHistoryContent() { } if (statusFilter !== "all" && run.status !== statusFilter) return false; if (runTypeFilter !== "all" && (run.run_type ?? "dryrun") !== runTypeFilter) return false; - if (invalidOnly && !(run.invalid_rows != null && run.invalid_rows > 0)) return false; + if (failedOnly) { + const errors = run.error_rows ?? run.invalid_rows; + const warnings = run.warning_rows; + const hasFailures = (errors != null && errors > 0) || (warnings != null && warnings > 0); + if (!hasFailures) return false; + } if (myRunsOnly && currentUserEmail && run.requesting_user !== currentUserEmail) return false; if (labelFilter.size > 0) { // Match the label filter against any check captured on the run; an @@ -321,8 +330,11 @@ function RunHistoryContent() { case "valid": cmp = (a.valid_rows ?? 0) - (b.valid_rows ?? 0); break; - case "invalid": - cmp = (a.invalid_rows ?? 0) - (b.invalid_rows ?? 0); + case "errors": + cmp = (a.error_rows ?? a.invalid_rows ?? 0) - (b.error_rows ?? b.invalid_rows ?? 0); + break; + case "warnings": + cmp = (a.warning_rows ?? 0) - (b.warning_rows ?? 0); break; case "run_date": cmp = (a.created_at ?? "").localeCompare(b.created_at ?? ""); @@ -330,7 +342,7 @@ function RunHistoryContent() { } return cmp * dir; }); - }, [allRuns, catalogFilter, schemaFilter, tableFilter, tableSearch, statusFilter, runTypeFilter, invalidOnly, myRunsOnly, currentUserEmail, rSortKey, rSortDir, labelFilter]); + }, [allRuns, catalogFilter, schemaFilter, tableFilter, tableSearch, statusFilter, runTypeFilter, failedOnly, myRunsOnly, currentUserEmail, rSortKey, rSortDir, labelFilter]); // Distinct labels seen across all runs' checks. Drives the LabelFilter // dropdown content for this page. @@ -368,7 +380,7 @@ function RunHistoryContent() { setTableFilter("all"); }; - const hasActiveFilters = catalogFilter !== "all" || schemaFilter !== "all" || tableFilter !== "all" || tableSearch !== "" || statusFilter !== "all" || runTypeFilter !== "all" || invalidOnly || myRunsOnly || labelFilter.size > 0; + const hasActiveFilters = catalogFilter !== "all" || schemaFilter !== "all" || tableFilter !== "all" || tableSearch !== "" || statusFilter !== "all" || runTypeFilter !== "all" || failedOnly || myRunsOnly || labelFilter.size > 0; const PAGE_SIZE = 25; const [currentPage, setCurrentPage] = useState(1); @@ -378,7 +390,7 @@ function RunHistoryContent() { [runs, currentPage], ); - useEffect(() => { setCurrentPage(1); }, [catalogFilter, schemaFilter, tableFilter, tableSearch, statusFilter, runTypeFilter, invalidOnly, myRunsOnly, rSortKey, rSortDir]); + useEffect(() => { setCurrentPage(1); }, [catalogFilter, schemaFilter, tableFilter, tableSearch, statusFilter, runTypeFilter, failedOnly, myRunsOnly, rSortKey, rSortDir]); return (
@@ -474,13 +486,14 @@ function RunHistoryContent() {
+ {pagedRuns.map((run) => { - const invalidPct = - run.total_rows && run.invalid_rows - ? ((run.invalid_rows / run.total_rows) * 100).toFixed(1) + // ``error_rows`` is the DQX observer count (post-v5); fall back + // to ``invalid_rows`` for runs created before the rename. + const errors = run.error_rows ?? run.invalid_rows; + const errorPct = + run.total_rows && errors + ? ((errors / run.total_rows) * 100).toFixed(1) : null; const isExpanded = expandedRunId === run.run_id; return ( setExpandedRunId(isExpanded ? null : run.run_id)} /> @@ -672,12 +691,12 @@ function RunHistoryContent() { function RunHistoryRow({ run, - invalidPct, + errorPct, isExpanded, onToggle, }: { run: ValidationRunSummaryOut; - invalidPct: string | null; + errorPct: string | null; isExpanded: boolean; onToggle: () => void; }) { @@ -754,14 +773,29 @@ function RunHistoryRow({ {run.valid_rows?.toLocaleString() ?? "—"} + @@ -771,7 +805,7 @@ function RunHistoryRow({ {isExpanded && ( - - +
ErrorsWarnings
+
+ {(row.warnings ?? []).map((w, wi) => ( + + {formatError(w)} + + ))} +
+
- + + + @@ -578,16 +594,19 @@ function RunHistoryContent() {
- {run.invalid_rows != null ? ( - 0 ? "text-red-600 font-medium" : ""}> - {run.invalid_rows.toLocaleString()} - {invalidPct && run.invalid_rows > 0 && ( - ({invalidPct}%) - )} + {/* ``error_rows`` is the authoritative DQX observer count + (``error_row_count``). Pre-v5 runs only have ``invalid_rows`` + — fall back to that so historical data still renders. */} + {(() => { + const errors = run.error_rows ?? run.invalid_rows; + if (errors == null) return "—"; + return ( + 0 ? "text-red-600 font-medium" : ""}> + {errors.toLocaleString()} + {errorPct && errors > 0 && ( + ({errorPct}%) + )} + + ); + })()} + + {run.warning_rows != null ? ( + 0 ? "text-amber-600 font-medium" : ""}> + {run.warning_rows.toLocaleString()} ) : ( + // Em dash distinguishes pre-v3 history rows from "0 warnings". "—" )}
+
{run.status === "FAILED" && run.error_message && (
diff --git a/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/runs.tsx b/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/runs.tsx index 62609dac8..0d61d1db4 100644 --- a/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/runs.tsx +++ b/app/src/databricks_labs_dqx_app/ui/routes/_sidebar/runs.tsx @@ -89,10 +89,62 @@ import { useActiveRuns, type ActiveRun } from "@/hooks/use-active-runs"; import { getDryRunStatus, type RunStatusOut } from "@/lib/api"; import { cancelDryRun } from "@/lib/api-custom"; import { CircleStop, ShieldAlert } from "lucide-react"; -import { parseFqn, formatDateTime as formatDate } from "@/lib/format-utils"; +import { parseFqn, formatDateTime as formatDate, getUserMetadata, labelToken, tokenToLabel, formatLabel } from "@/lib/format-utils"; +import { LabelFilter, LabelsBadges, labelsMatchFilter } from "@/components/Labels"; import { usePermissions } from "@/hooks/use-permissions"; import { requireRunnerOrRedirect } from "@/lib/route-guards"; +// Collect every distinct ``(key, value)`` label seen on the checks of the +// supplied approved rule sets. Used to populate ```` dropdowns +// on the Execute tab and the schedule editor. +function collectAvailableLabels( + rules: RuleCatalogEntryOut[], +): { key: string; value: string }[] { + const seen = new Set(); + const out: { key: string; value: string }[] = []; + for (const r of rules) { + for (const c of r.checks) { + const md = getUserMetadata(c as Record); + for (const [k, v] of Object.entries(md)) { + const tok = labelToken(k, v); + if (!seen.has(tok)) { + seen.add(tok); + out.push({ key: k, value: v }); + } + } + } + } + return out; +} + +// True iff at least one check on ``rule`` carries a label that satisfies +// the user's selection. Empty selection always passes. +function ruleMatchesLabels( + rule: RuleCatalogEntryOut, + selected: Set, +): boolean { + if (selected.size === 0) return true; + return rule.checks.some((c) => + labelsMatchFilter(getUserMetadata(c as Record), selected), + ); +} + +// Merge every check's ``user_metadata`` for the rule set into a single +// ``Record`` for display next to the table name. +// Most rule sets use the same labels across all checks; for the rare +// case where a key has different values per check, last-write-wins — +// the LabelFilter still shows every distinct ``(key, value)`` pair. +function collectRuleLabels(rule: RuleCatalogEntryOut): Record { + const merged: Record = {}; + for (const c of rule.checks) { + const md = getUserMetadata(c as Record); + for (const [k, v] of Object.entries(md)) { + merged[k] = v; + } + } + return merged; +} + export const Route = createFileRoute("/_sidebar/runs")({ // URL-level guard: aborts the route load before the page ever mounts // when the user lacks the RUNNER role. This complements the in-component @@ -128,6 +180,11 @@ interface ScheduleConfig { scope_catalogs?: string[]; scope_schemas?: string[]; scope_tables?: string[]; + // Optional label filter intersected with the FQN-based scope above. A rule + // set is included only if at least one of its checks carries a matching + // ``user_metadata`` label. Persisted as a list of ``{key, value}`` so the + // YAML round-trips cleanly and values containing ``=`` don't collide. + scope_labels?: { key: string; value: string }[]; sample_size?: number; } @@ -348,18 +405,40 @@ function ScopePicker({ return Array.from(set).sort(); }, [approvedRules]); + const availableLabels = useMemo( + () => collectAvailableLabels(approvedRules), + [approvedRules], + ); + + const selectedLabelTokens = useMemo(() => { + const set = new Set(); + for (const { key, value } of config.scope_labels ?? []) { + set.add(labelToken(key, value)); + } + return set; + }, [config.scope_labels]); + + const onLabelFilterChange = (selected: Set) => { + const next = [...selected].map(tokenToLabel); + update({ scope_labels: next.length > 0 ? next : undefined }); + }; + const matchedCount = useMemo(() => { return approvedRules.filter((r) => { const { catalog, schema } = parseFqn(r.table_fqn); - switch (config.scope_mode) { - case "all": return true; - case "catalog": return (config.scope_catalogs ?? []).includes(catalog); - case "schema": return (config.scope_schemas ?? []).includes(`${catalog}.${schema}`); - case "tables": return (config.scope_tables ?? []).includes(r.table_fqn); - default: return true; - } + const fqnMatches = (() => { + switch (config.scope_mode) { + case "all": return true; + case "catalog": return (config.scope_catalogs ?? []).includes(catalog); + case "schema": return (config.scope_schemas ?? []).includes(`${catalog}.${schema}`); + case "tables": return (config.scope_tables ?? []).includes(r.table_fqn); + default: return true; + } + })(); + if (!fqnMatches) return false; + return ruleMatchesLabels(r, selectedLabelTokens); }).length; - }, [approvedRules, config]); + }, [approvedRules, config, selectedLabelTokens]); const toggleInList = (list: string[], item: string): string[] => { return list.includes(item) ? list.filter((x) => x !== item) : [...list, item]; @@ -442,7 +521,10 @@ function ScopePicker({ onClick={() => !disabled && update({ scope_tables: toggleInList(config.scope_tables ?? [], rule.table_fqn) })} > update({ scope_tables: toggleInList(config.scope_tables ?? [], rule.table_fqn) })} disabled={disabled} /> - {rule.display_name || rule.table_fqn} +
+ {rule.display_name || rule.table_fqn} + +
{rule.checks.length} rule{rule.checks.length !== 1 ? "s" : ""}
))} @@ -451,6 +533,23 @@ function ScopePicker({
)} +
+ +
+ +
+ {availableLabels.length === 0 && ( +

+ No labels found on approved rules — add labels to checks to use this filter. +

+ )} +
+

{matchedCount} of {approvedRules.length} approved rule set{approvedRules.length !== 1 ? "s" : ""} matched

@@ -495,16 +594,23 @@ function resolveScheduleScope( cfg: ScheduleConfig, approvedRules: RuleCatalogEntryOut[], ): string[] { + const labelTokens = new Set( + (cfg.scope_labels ?? []).map(({ key, value }) => labelToken(key, value)), + ); return approvedRules .filter((r) => { const { catalog, schema } = parseFqn(r.table_fqn); - switch (cfg.scope_mode) { - case "all": return true; - case "catalog": return (cfg.scope_catalogs ?? []).includes(catalog); - case "schema": return (cfg.scope_schemas ?? []).includes(`${catalog}.${schema}`); - case "tables": return (cfg.scope_tables ?? []).includes(r.table_fqn); - default: return true; - } + const fqnMatches = (() => { + switch (cfg.scope_mode) { + case "all": return true; + case "catalog": return (cfg.scope_catalogs ?? []).includes(catalog); + case "schema": return (cfg.scope_schemas ?? []).includes(`${catalog}.${schema}`); + case "tables": return (cfg.scope_tables ?? []).includes(r.table_fqn); + default: return true; + } + })(); + if (!fqnMatches) return false; + return ruleMatchesLabels(r, labelTokens); }) .map((r) => r.table_fqn); } @@ -719,6 +825,7 @@ function ExecuteTab({ onGoToHistory }: { onGoToHistory: () => void }) { const [searchQuery, setSearchQuery] = useState(""); const [filterCatalog, setFilterCatalog] = useState("__all__"); const [filterSchema, setFilterSchema] = useState("__all__"); + const [labelFilter, setLabelFilter] = useState>(new Set()); const [isRunning, setIsRunning] = useState(false); const [runNotification, setRunNotification] = useState(null); @@ -759,6 +866,11 @@ function ExecuteTab({ onGoToHistory }: { onGoToHistory: () => void }) { } }, [availableSchemas, filterSchema]); + const availableLabels = useMemo( + () => collectAvailableLabels(approvedRules), + [approvedRules], + ); + const filteredRules = useMemo(() => { return approvedRules.filter((r) => { const isSqlCheck = r.table_fqn.startsWith(_SQL_CHECK_PREFIX); @@ -779,9 +891,10 @@ function ExecuteTab({ onGoToHistory }: { onGoToHistory: () => void }) { const q = searchQuery.toLowerCase(); if (!(r.display_name || r.table_fqn).toLowerCase().includes(q)) return false; } + if (!ruleMatchesLabels(r, labelFilter)) return false; return true; }); - }, [approvedRules, filterCatalog, filterSchema, searchQuery]); + }, [approvedRules, filterCatalog, filterSchema, searchQuery, labelFilter]); const grouped = useMemo(() => { const groups = new Map(); @@ -1072,6 +1185,13 @@ function ExecuteTab({ onGoToHistory }: { onGoToHistory: () => void }) { + +
e.stopPropagation()} />
{rule.display_name || rule.table_fqn} +
+ {rule.display_name || rule.table_fqn} + +
+ 		{rule.checks.length}
@@ -1618,17 +1743,23 @@ function RunEditorError({ resetErrorBoundary }: { resetErrorBoundary: () => void // =========================================================================== function scopeLabel(cfg: ScheduleConfig): string { - switch (cfg.scope_mode) { - case "all": return "All approved rules"; - case "catalog": return (cfg.scope_catalogs ?? []).join(", ") || "No catalogs"; - case "schema": return (cfg.scope_schemas ?? []).join(", ") || "No schemas"; - case "tables": { - const tables = cfg.scope_tables ?? []; - if (tables.length <= 2) return tables.join(", ") || "No tables"; - return `${tables[0]}, ${tables[1]} +${tables.length - 2} more`; + const base = (() => { + switch (cfg.scope_mode) { + case "all": return "All approved rules"; + case "catalog": return (cfg.scope_catalogs ?? []).join(", ") || "No catalogs"; + case "schema": return (cfg.scope_schemas ?? []).join(", ") || "No schemas"; + case "tables": { + const tables = cfg.scope_tables ?? []; + if (tables.length <= 2) return tables.join(", ") || "No tables"; + return `${tables[0]}, ${tables[1]} +${tables.length - 2} more`; + } + default: return "All"; } - default: return "All"; - } + })(); + const labels = cfg.scope_labels ?? []; + if (labels.length === 0) return base; + const labelStr = labels.map(({ key, value }) => formatLabel(key, value)).join(", "); + return `${base} · labels: ${labelStr}`; } function SchedulesListView({ isDeleting }: { isDeleting?: boolean }) { @@ -1739,7 +1870,7 @@ function SchedulesListView({ isDeleting }: { isDeleting?: boolean }) {
- {sched.config.sample_size != null ? sched.config.sample_size.toLocaleString() : "All rows"} + {sched.config.sample_size ? sched.config.sample_size.toLocaleString() : "All rows"} {sched.updated_at ? formatDate(sched.updated_at) : sched.created_at ? formatDate(sched.created_at) : "—"} @@ -2177,6 +2308,26 @@ function FormEditor({ if (parsed.scope_catalogs) schedFromYaml.scope_catalogs = parsed.scope_catalogs; if (parsed.scope_schemas) schedFromYaml.scope_schemas = parsed.scope_schemas; if (parsed.scope_tables) schedFromYaml.scope_tables = parsed.scope_tables; + if (Array.isArray(parsed.scope_labels)) { + // Tolerate both the canonical ``[{key, value}]`` form and a + // ``["key=value"]`` shorthand that admins might hand-edit. + schedFromYaml.scope_labels = parsed.scope_labels + .map((entry: unknown) => { + if (entry && typeof entry === "object" && !Array.isArray(entry)) { + const e = entry as Record; + if (typeof e.key === "string") { + return { key: e.key, value: String(e.value ?? "") }; + } + } + if (typeof entry === "string") { + const idx = entry.indexOf("="); + if (idx < 0) return { key: entry, value: "" }; + return { key: entry.slice(0, idx), value: entry.slice(idx + 1) }; + } + return null; + }) + .filter((v: unknown): v is { key: string; value: string } => v !== null); + } if (parsed.sample_size != null) schedFromYaml.sample_size = parsed.sample_size; if (parsed.cron_expression) schedFromYaml.cron_expression = parsed.cron_expression; setScheduleConfig({ ...DEFAULT_SCHEDULE, ...schedFromYaml }); diff --git a/app/tasks/src/dqx_task_runner/runner.py b/app/tasks/src/dqx_task_runner/runner.py index 984db577f..01c7a4785 100644 --- a/app/tasks/src/dqx_task_runner/runner.py +++ b/app/tasks/src/dqx_task_runner/runner.py @@ -18,7 +18,7 @@ import sys import time import re -from datetime import datetime, timezone, date +from datetime import date, datetime from decimal import Decimal from typing import Any, cast @@ -267,7 +267,12 @@ def _run_profile( # Write result row. Profiling has no checks yet, but we still record a # null fingerprint slot so later pipeline stages that join on # rule_set_fingerprint don't have to special-case profile rows. - now = datetime.now(timezone.utc).isoformat() + # + # ``created_at`` is TIMESTAMP per the baseline schema — we materialise + # via ``current_timestamp()`` instead of an ISO string so cluster-key + # zone maps stay tight and downstream filters behave correctly. + from pyspark.sql import functions as F + result_row = spark.createDataFrame( [ ( @@ -283,7 +288,6 @@ def _run_profile( _json_dumps(rules) if rules else "[]", "SUCCESS", None, - now, None, ) ], @@ -291,10 +295,10 @@ def _run_profile( "run_id STRING, requesting_user STRING, source_table_fqn STRING, " "view_fqn STRING, sample_limit INT, rows_profiled INT, columns_profiled INT, " "duration_seconds DOUBLE, summary_json STRING, generated_rules_json STRING, " - "status STRING, error_message STRING, created_at STRING, " + "status STRING, error_message STRING, " "rule_set_fingerprint STRING" ), - ) + ).withColumn("created_at", F.current_timestamp()) result_row.writeTo(result_table).append() logger.info("Profile results written to %s (run_id=%s)", result_table, run_id) @@ -353,7 +357,11 @@ def _run_dryrun( ) engine = DQEngine(workspace_client=ws, spark=spark, observer=observer) - df = _read_view_with_retry(spark, view_fqn).limit(sample_size) + # ``sample_size = 0`` is the UI convention for "All rows" — skip + # ``.limit`` for it (passing 0 would short-circuit to an empty DataFrame). + df = _read_view_with_retry(spark, view_fqn) + if sample_size: + df = df.limit(sample_size) # ``apply_checks_by_metadata_and_split`` returns ``(valid, invalid)`` # when no observer is attached and ``(valid, invalid, observation)`` # when one is. We always construct the engine with ``observer=...`` @@ -368,21 +376,29 @@ def _run_dryrun( # observation. ``invalid_df.count()`` triggers metric collection. invalid_rows = invalid_df.count() - sample_invalid: list[dict[str, Any]] = [] - if invalid_rows > 0: - sample_rows = invalid_df.limit(10).collect() - sample_invalid = [row.asDict(recursive=True) for row in sample_rows] - - observed = observation.get if observation is not None else {} + # CRITICAL: snapshot the observation IMMEDIATELY after the first + # full-scan action and BEFORE any limit/partial action. On Spark + # Connect (Databricks), ``Observation.get`` returns a reference to a + # live dict that gets mutated by every subsequent action — including + # ``invalid_df.limit(10).collect()`` below, whose limit pushdown + # would otherwise overwrite ``input_row_count`` with ~10. We also + # ``dict(...)`` the result to detach our copy from the live dict so + # the later quarantine write (which re-fires the observer) cannot + # change what we pass to ``_persist_observed_metrics``. + observed: dict[str, Any] = dict(observation.get) if observation is not None else {} total_rows = int(observed.get("input_row_count", 0) or 0) valid_rows = int(observed.get("valid_row_count", 0) or 0) error_rows = int(observed.get("error_row_count", 0) or 0) warning_rows = int(observed.get("warning_row_count", 0) or 0) - - # Backwards-compatible error_summary derived from check_metrics. error_summary = _check_metrics_to_error_summary(observed.get("check_metrics")) - now = datetime.now(timezone.utc).isoformat() + sample_invalid: list[dict[str, Any]] = [] + if invalid_rows > 0: + sample_rows = invalid_df.limit(10).collect() + sample_invalid = [row.asDict(recursive=True) for row in sample_rows] + + from pyspark.sql import functions as F + result_row = spark.createDataFrame( [ ( @@ -395,11 +411,12 @@ def _run_dryrun( total_rows, valid_rows, invalid_rows, + error_rows, + warning_rows, _json_dumps(error_summary), _json_dumps(sample_invalid), "SUCCESS", None, - now, run_type, fingerprint, ) @@ -407,12 +424,12 @@ def _run_dryrun( schema=( "run_id STRING, requesting_user STRING, source_table_fqn STRING, " "view_fqn STRING, checks_json STRING, sample_size INT, " - "total_rows INT, valid_rows INT, invalid_rows INT, " + "total_rows INT, valid_rows INT, invalid_rows INT, error_rows INT, warning_rows INT, " "error_summary_json STRING, sample_invalid_json STRING, " - "status STRING, error_message STRING, created_at STRING, " + "status STRING, error_message STRING, " "run_type STRING, rule_set_fingerprint STRING" ), - ) + ).withColumn("created_at", F.current_timestamp()) result_row.writeTo(result_table).append() logger.info( "Dry-run results written to %s (run_id=%s, run_type=%s, errors=%d, warnings=%d)", @@ -485,9 +502,12 @@ def _run_dryrun_sql_check( sample_rows = violation_df.limit(10).collect() sample_invalid = [row.asDict(recursive=True) for row in sample_rows] - now = datetime.now(timezone.utc).isoformat() + from pyspark.sql import functions as F + if run_type == "dryrun" and sample_size == 0: run_type = "scheduled" + # SQL checks treat every violation as an error (no warning concept), so + # ``error_rows == invalid_rows`` and ``warning_rows == 0``. result_row = spark.createDataFrame( [ ( @@ -500,11 +520,12 @@ def _run_dryrun_sql_check( total_rows, valid_rows, invalid_rows, + invalid_rows, + 0, _json_dumps(error_summary), _json_dumps(sample_invalid), "SUCCESS", None, - now, run_type, fingerprint, ) @@ -512,12 +533,12 @@ def _run_dryrun_sql_check( schema=( "run_id STRING, requesting_user STRING, source_table_fqn STRING, " "view_fqn STRING, checks_json STRING, sample_size INT, " - "total_rows INT, valid_rows INT, invalid_rows INT, " + "total_rows INT, valid_rows INT, invalid_rows INT, error_rows INT, warning_rows INT, " "error_summary_json STRING, sample_invalid_json STRING, " - "status STRING, error_message STRING, created_at STRING, " + "status STRING, error_message STRING, " "run_type STRING, rule_set_fingerprint STRING" ), - ) + ).withColumn("created_at", F.current_timestamp()) result_row.writeTo(result_table).append() logger.info( "SQL-check %s results written to %s (run_id=%s, violations=%d)", run_type, result_table, run_id, invalid_rows @@ -582,7 +603,13 @@ def _write_quarantine_records( result_catalog: str, result_schema: str, ) -> None: - """Persist every invalid row to dq_quarantine_records.""" + """Persist every invalid row to dq_quarantine_records. + + ``row_data`` and ``errors`` columns are VARIANT in the baseline + schema. We materialise them via ``parse_json(to_json(...))`` so the + row payload becomes a typed JSON value (queryable with + ``variant_get``) rather than an opaque STRING. + """ from pyspark.sql import functions as F invalid_count = invalid_df.count() @@ -591,18 +618,34 @@ def _write_quarantine_records( return quarantine_table = f"{result_catalog}.{result_schema}.dq_quarantine_records" - now = datetime.now(timezone.utc).isoformat() data_cols = [c for c in invalid_df.columns if c not in ("_warnings", "_errors", "_rule_name")] + # ``_warnings`` may be absent if every check is error-level; default to a + # JSON null so the column stays well-typed in the VARIANT. + warnings_expr = ( + F.parse_json(F.to_json(F.col("_warnings"))) + if "_warnings" in invalid_df.columns + else F.parse_json(F.lit("null")) + ) quarantine_df = ( invalid_df.withColumn("quarantine_id", F.expr("uuid()")) .withColumn("run_id", F.lit(run_id)) .withColumn("source_table_fqn", F.lit(source_table_fqn)) .withColumn("requesting_user", F.lit(requesting_user)) - .withColumn("row_data", F.to_json(F.struct(*data_cols))) - .withColumn("errors", F.to_json(F.col("_errors"))) - .withColumn("created_at", F.lit(now)) - .select("quarantine_id", "run_id", "source_table_fqn", "requesting_user", "row_data", "errors", "created_at") + .withColumn("row_data", F.parse_json(F.to_json(F.struct(*data_cols)))) + .withColumn("errors", F.parse_json(F.to_json(F.col("_errors")))) + .withColumn("warnings", warnings_expr) + .withColumn("created_at", F.current_timestamp()) + .select( + "quarantine_id", + "run_id", + "source_table_fqn", + "requesting_user", + "row_data", + "errors", + "warnings", + "created_at", + ) ) quarantine_df.writeTo(quarantine_table).append() logger.info("Wrote %d quarantine rows to %s (run_id=%s)", invalid_count, quarantine_table, run_id) @@ -650,23 +693,28 @@ def _write_sql_quarantine_records( capped_df = violation_df.limit(persisted_target) if persisted_target < invalid_count else violation_df quarantine_table = f"{result_catalog}.{result_schema}.dq_quarantine_records" - now = datetime.now(timezone.utc).isoformat() # Every column on the violation view is part of the row payload — # there are no DQX-internal columns to strip (the ``not in`` filter # is just defensive in case a check author re-uses those names). + # ``row_data``/``errors`` are VARIANT — wrap with parse_json so the + # values land as typed JSON rather than opaque strings. data_cols = [c for c in capped_df.columns if c not in ("_warnings", "_errors", "_rule_name")] - row_data_expr = F.to_json(F.struct(*data_cols)) if data_cols else F.lit("{}") - errors_json = _json_dumps({check_name: "SQL check violation"}) + row_data_expr = F.parse_json(F.to_json(F.struct(*data_cols))) if data_cols else F.parse_json(F.lit("{}")) + errors_expr = F.parse_json(F.lit(_json_dumps({check_name: "SQL check violation"}))) + # SQL checks have no warning-level distinction — every violation is + # treated as an error. Persist ``warnings`` as JSON null so the + # column shape matches row-level quarantine writes. quarantine_df = ( capped_df.withColumn("quarantine_id", F.expr("uuid()")) .withColumn("run_id", F.lit(run_id)) .withColumn("source_table_fqn", F.lit(source_table_fqn)) .withColumn("requesting_user", F.lit(requesting_user)) .withColumn("row_data", row_data_expr) - .withColumn("errors", F.lit(errors_json)) - .withColumn("created_at", F.lit(now)) + .withColumn("errors", errors_expr) + .withColumn("warnings", F.parse_json(F.lit("null"))) + .withColumn("created_at", F.current_timestamp()) .select( "quarantine_id", "run_id", @@ -674,6 +722,7 @@ def _write_sql_quarantine_records( "requesting_user", "row_data", "errors", + "warnings", "created_at", ) ) @@ -923,14 +972,19 @@ def _run_scheduled( invalid_rows = invalid_df.count() # triggers the observation - observed = observation.get if observation is not None else {} + # Defensive snapshot — on Spark Connect, ``Observation.get`` returns + # a reference to a live dict that subsequent actions (the quarantine + # write below) re-fire and can mutate. ``dict(...)`` detaches our + # copy. + observed: dict[str, Any] = dict(observation.get) if observation is not None else {} total_rows = int(observed.get("input_row_count", 0) or 0) valid_rows = int(observed.get("valid_row_count", 0) or 0) error_rows = int(observed.get("error_row_count", 0) or 0) warning_rows = int(observed.get("warning_row_count", 0) or 0) error_summary = _check_metrics_to_error_summary(observed.get("check_metrics")) - now = datetime.now(timezone.utc).isoformat() + from pyspark.sql import functions as F + result_table = f"{result_catalog}.{result_schema}.dq_validation_runs" result_row = spark.createDataFrame( [ @@ -944,11 +998,12 @@ def _run_scheduled( total_rows, valid_rows, invalid_rows, + error_rows, + warning_rows, _json_dumps(error_summary), None, "SUCCESS", None, - now, "scheduled", fingerprint, ) @@ -956,12 +1011,12 @@ def _run_scheduled( schema=( "run_id STRING, requesting_user STRING, source_table_fqn STRING, " "view_fqn STRING, checks_json STRING, sample_size INT, " - "total_rows INT, valid_rows INT, invalid_rows INT, " + "total_rows INT, valid_rows INT, invalid_rows INT, error_rows INT, warning_rows INT, " "error_summary_json STRING, sample_invalid_json STRING, " - "status STRING, error_message STRING, created_at STRING, " + "status STRING, error_message STRING, " "run_type STRING, rule_set_fingerprint STRING" ), - ) + ).withColumn("created_at", F.current_timestamp()) result_row.writeTo(result_table).append() logger.info( "Scheduled run results written to %s (run_id=%s, errors=%d, warnings=%d)", @@ -1009,7 +1064,8 @@ def _write_error( skip_history: bool = False, ) -> None: """Write a FAILED status row so the app can report the error.""" - now = datetime.now(timezone.utc).isoformat() + from pyspark.sql import functions as F + checks_str = _json_dumps(checks) if checks else None fingerprint = _compute_fingerprint(checks or []) if task_type == "profile": @@ -1029,7 +1085,6 @@ def _write_error( None, "FAILED", error_message, - now, fingerprint, ) ], @@ -1037,10 +1092,10 @@ def _write_error( "run_id STRING, requesting_user STRING, source_table_fqn STRING, " "view_fqn STRING, sample_limit INT, rows_profiled INT, columns_profiled INT, " "duration_seconds DOUBLE, summary_json STRING, generated_rules_json STRING, " - "status STRING, error_message STRING, created_at STRING, " + "status STRING, error_message STRING, " "rule_set_fingerprint STRING" ), - ) + ).withColumn("created_at", F.current_timestamp()) else: table = f"{result_catalog}.{result_schema}.dq_validation_runs" if skip_history: @@ -1061,11 +1116,12 @@ def _write_error( None, None, None, + None, # error_rows + None, # warning_rows None, None, "FAILED", error_message, - now, error_run_type, fingerprint, ) @@ -1073,12 +1129,12 @@ def _write_error( schema=( "run_id STRING, requesting_user STRING, source_table_fqn STRING, " "view_fqn STRING, checks_json STRING, sample_size INT, " - "total_rows INT, valid_rows INT, invalid_rows INT, " + "total_rows INT, valid_rows INT, invalid_rows INT, error_rows INT, warning_rows INT, " "error_summary_json STRING, sample_invalid_json STRING, " - "status STRING, error_message STRING, created_at STRING, " + "status STRING, error_message STRING, " "run_type STRING, rule_set_fingerprint STRING" ), - ) + ).withColumn("created_at", F.current_timestamp()) row.writeTo(table).append() logger.info("Error result written to %s (run_id=%s)", table, run_id) diff --git a/app/tests/conftest.py b/app/tests/conftest.py index d1718e114..49f71ad41 100644 --- a/app/tests/conftest.py +++ b/app/tests/conftest.py @@ -91,3 +91,132 @@ def _make(email: str = "user@example.com", groups: list[str] | None = None) -> A return user return _make + + +@pytest.fixture +def make_scheduler(): + """Factory that constructs a real :class:`SchedulerService` for tests. + + Why this exists + --------------- + Earlier scheduler tests skipped ``__init__`` via + ``SchedulerService.__new__(SchedulerService)`` and hand-attached + private attributes (``_sql``, ``_oltp_sql``, ``_settings_table``, + ``_catalog``, ...). That pattern couples the test suite to internal + field names: a refactor of ``_oltp_sql`` → ``_oltp_executor`` would + silently turn every fixture into AttributeError at access time + rather than a loud TypeError at construction. + + This factory routes through the real public constructor instead. + The constructor parameter ``oltp_sql=`` IS the seam — renaming it + fails fast in CI; renaming the resulting ``_oltp_sql`` private + attribute only matters for the one test that introspects Delta-side + ``execute`` calls separately (which we document explicitly via the + ``distinct_sql`` opt-in). + + The constructor and ``SqlExecutor.__init__`` are both pure + attribute-assignment — no IO, no warehouse calls, no event-loop + binding — so wrapping them in a unit test is safe. + + Returns + ------- + A callable that builds a ``(service, mocks)`` pair. ``mocks.oltp`` + is the OLTP executor passed via the public ``oltp_sql=`` parameter; + ``mocks.sql`` is the analytical Delta executor. With + ``distinct_sql=True``, ``mocks.sql`` is a fresh ``MagicMock`` + swapped onto the service so tests can distinguish Delta DELETE + calls from OLTP ones. With ``distinct_sql=False`` (the default, + matching legacy single-backend deployments), ``mocks.sql`` aliases + the constructor-built ``SqlExecutor`` — tests in that mode only + interact with ``mocks.oltp``. + + Usage:: + + svc, mocks = make_scheduler(oltp_dialect="postgres") + mocks.oltp.query.return_value = [("90",)] + assert svc._resolve_retention_days() == 90 + """ + from types import SimpleNamespace + + def _make( + *, + oltp_dialect: str = "delta", + oltp_query_return: list[tuple[Any, ...]] | None = None, + catalog: str = "dqx", + schema: str = "public", + tmp_schema: str | None = None, + distinct_sql: bool = False, + distinct_tmp_sql: bool = False, + oltp_spec: list[str] | None = None, + ) -> tuple[Any, SimpleNamespace]: + from databricks_labs_dqx_app.backend.services.scheduler_service import SchedulerService + + if oltp_spec is not None: + # spec-bound mock — used by tests that need attribute + # access on the mock to be restricted (so any code path + # that touches an unlisted attribute fails loudly). The + # spec list MUST cover every attribute the constructor + + # tested method touches: ``fqn``, ``query``, ``execute`` + # at minimum. + oltp = MagicMock(spec=oltp_spec, name="oltp_sql") + else: + oltp = MagicMock(name="oltp_sql") + oltp.dialect = oltp_dialect + + # Production-realistic FQN shape: the constructor calls + # ``oltp_sql.fqn("dq_app_settings")`` etc., so a side_effect + # that returns ``catalog.schema.`` produces a meaningful + # ``_settings_table`` for tests that later inspect it. + oltp.fqn.side_effect = lambda t: f"{catalog}.{schema}.{t}" + oltp.query.return_value = [] if oltp_query_return is None else oltp_query_return + + # ``OltpExecutorProtocol`` methods that services depend on for + # dialect-agnostic SQL — wire concrete side_effects matching + # the real implementations so tests that assert on rendered + # SQL (e.g. retention's INTERVAL literal) see real strings, + # not MagicMock reprs. Only attached when the spec permits. + if oltp_spec is None or "interval_days_expr" in oltp_spec: + if oltp_dialect == "postgres": + oltp.interval_days_expr.side_effect = lambda d: f"INTERVAL '{int(d)} days'" + else: + oltp.interval_days_expr.side_effect = lambda d: f"INTERVAL {int(d)} DAY" + if oltp_spec is None or "select_json_text" in oltp_spec: + if oltp_dialect == "postgres": + oltp.select_json_text.side_effect = lambda c: c + else: + oltp.select_json_text.side_effect = lambda c: f"to_json({c})" + + svc = SchedulerService( + ws=MagicMock(name="WorkspaceClient"), + warehouse_id="test-wh", + catalog=catalog, + schema=schema, + tmp_schema=tmp_schema if tmp_schema is not None else f"{schema}_tmp", + job_id="test-job-0", + oltp_sql=oltp, + ) + + mocks = SimpleNamespace(oltp=oltp) + # The analytical and tmp executors have no constructor seam + # today (the constructor builds them from ``ws + warehouse_id + + # {schema,tmp_schema}``). Tests that need to assert on their + # calls swap them here explicitly. These two attribute writes + # are the entire remaining internal-name surface — collapsing + # them into one helper means a future rename of either field + # fails in one place rather than scattered across the suite. + if distinct_sql: + delta = MagicMock(name="delta_sql") + svc._sql = delta # noqa: SLF001 - see helper docstring + mocks.sql = delta + else: + mocks.sql = svc._sql # noqa: SLF001 - alias for legacy single-backend tests + if distinct_tmp_sql: + tmp = MagicMock(name="tmp_sql") + svc._tmp_sql = tmp # noqa: SLF001 - see helper docstring + mocks.tmp = tmp + else: + mocks.tmp = svc._tmp_sql # noqa: SLF001 + + return svc, mocks + + return _make diff --git a/app/tests/test_app_scheduler_lease.py b/app/tests/test_app_scheduler_lease.py new file mode 100644 index 000000000..9c8357a25 --- /dev/null +++ b/app/tests/test_app_scheduler_lease.py @@ -0,0 +1,148 @@ +"""Tests for ``backend.app._try_acquire_scheduler_lease``. + +The function holds an fcntl advisory lock for the lifetime of the +worker so only one uvicorn worker runs the scheduler. The lock fd +MUST stay live as long as the process runs — if it gets garbage- +collected the OS releases the flock and a sibling worker can grab +the lease, resulting in two schedulers racing on the same job. + +Pre-refactor the fd was pinned via ``globals()["_scheduler_lock_fd"] = fd`` +— a magic-string lookup that's invisible to grep, type-checkers, +and IDE rename-refactor. The fix replaces that with a proper +module-level ``_scheduler_lock_fd: int | None = None`` global plus +``global _scheduler_lock_fd`` inside the function. + +These tests pin the new convention so a regression to the +``globals()`` form would fail loudly. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + + +@pytest.fixture +def fresh_lease_state(monkeypatch, tmp_path): + """Per-test isolation: fresh lock path + reset global to None. + + Each test gets its own lock file under ``tmp_path`` so concurrent + test runs don't contend on the production ``/tmp/.dqx_scheduler.lock`` + path, and the module global is reset before AND after the test so + a successful acquisition in one test doesn't bleed into the next. + """ + from databricks_labs_dqx_app.backend import app as app_mod + + monkeypatch.setattr(app_mod, "_SCHEDULER_LOCK_PATH", tmp_path / ".lease.lock") + original_fd = app_mod._scheduler_lock_fd + app_mod._scheduler_lock_fd = None + yield app_mod + # Best-effort cleanup: close any fd this test left behind so we don't + # leak descriptors across the suite. + leftover = app_mod._scheduler_lock_fd + if isinstance(leftover, int) and leftover != original_fd: + try: + os.close(leftover) + except OSError: + pass + app_mod._scheduler_lock_fd = original_fd + + +class TestSchedulerLeaseConvention: + """The reviewer's specific anti-pattern to lock down.""" + + def test_lock_fd_exists_as_module_level_global_not_dynamic_attribute(self): + """``_scheduler_lock_fd`` MUST be a real module-level annotation. + + The pre-refactor code wrote it via ``globals()["_scheduler_lock_fd"] = fd`` + which works at runtime but is invisible to grep, type-checkers, + and IDE rename. A regression that re-introduces the dynamic form + would (a) not appear in ``vars(module)`` at import time and + (b) leave the type annotation absent. + """ + from databricks_labs_dqx_app.backend import app as app_mod + + # Present at import time — the module-level statement + # ``_scheduler_lock_fd: int | None = None`` ran. + assert "_scheduler_lock_fd" in vars(app_mod), ( + "_scheduler_lock_fd must be declared at module top level. " + "If it only appears after _try_acquire_scheduler_lease() runs, " + "someone re-introduced the globals()[...] = fd pattern." + ) + + # Has a typed declaration. ``__annotations__`` is populated by + # the ``: int | None = None`` annotation in the module body. + # The dynamic-globals form leaves no annotation behind. + assert "_scheduler_lock_fd" in app_mod.__annotations__, ( + "_scheduler_lock_fd must have a module-level type annotation " + "so basedpyright and reviewers can see the dependency." + ) + + def test_source_does_not_use_globals_dict_assignment(self): + """Belt-and-suspenders: the literal anti-pattern must not return. + + Reads the source of ``_try_acquire_scheduler_lease`` and asserts + the magic-string form is absent. Cheap to maintain — if a future + refactor moves the global into a class or context-manager, this + test still passes; it only fails on the specific regression the + reviewer flagged. + """ + import inspect + + from databricks_labs_dqx_app.backend import app as app_mod + + src = inspect.getsource(app_mod._try_acquire_scheduler_lease) + assert 'globals()["_scheduler_lock_fd"]' not in src, ( + "Regression: the magic-string globals() assignment is back. " + "Use ``global _scheduler_lock_fd; _scheduler_lock_fd = fd`` instead." + ) + assert "globals()['_scheduler_lock_fd']" not in src + + +class TestSchedulerLeaseBehaviour: + """The fd-pinning semantics that the refactor must preserve.""" + + def test_acquire_returns_true_and_pins_fd_on_module_global(self, fresh_lease_state): + """Happy path: lease is acquired, fd is non-None, lock file exists.""" + app_mod = fresh_lease_state + + acquired = app_mod._try_acquire_scheduler_lease() + + assert acquired is True + # The fd is pinned to the module global so GC can't reap it + # mid-process and silently drop the flock. + assert isinstance(app_mod._scheduler_lock_fd, int) + assert app_mod._scheduler_lock_fd >= 0 + # And the lock file exists where we monkeypatched it. + assert Path(app_mod._SCHEDULER_LOCK_PATH).exists() + + def test_acquire_failure_does_not_clobber_existing_fd(self, fresh_lease_state, monkeypatch): + """If ``flock`` raises after a prior successful acquisition the + existing fd must survive — otherwise the previously-held lease + evaporates. + + Simulates the failure path by monkeypatching ``fcntl.flock`` to + always raise OSError on the second call. + """ + app_mod = fresh_lease_state + + assert app_mod._try_acquire_scheduler_lease() is True + original_fd = app_mod._scheduler_lock_fd + assert isinstance(original_fd, int) + + import fcntl + + def boom(*args, **kwargs): + raise OSError("would block") + + monkeypatch.setattr(fcntl, "flock", boom) + acquired = app_mod._try_acquire_scheduler_lease() + + assert acquired is False + # Critical: the previously-pinned fd must still be the same int. + # A regression that did ``_scheduler_lock_fd = fd`` *before* + # ``flock(...)`` could overwrite it with an unlocked fd; a + # regression that used ``finally`` could close it. + assert app_mod._scheduler_lock_fd == original_fd diff --git a/app/tests/test_custom_metrics.py b/app/tests/test_custom_metrics.py index 1271c8e8a..0abbc4835 100644 --- a/app/tests/test_custom_metrics.py +++ b/app/tests/test_custom_metrics.py @@ -10,7 +10,6 @@ from __future__ import annotations import json -from unittest.mock import MagicMock import pytest @@ -67,17 +66,175 @@ def test_save_normalises_and_persists(self, svc): s, sql = svc result = s.save_custom_metrics([" sum(x) as total ", None, "", "count(*) as n"]) assert result == ["sum(x) as total", "count(*) as n"] - # ``save_setting`` performs a MERGE — confirm the JSON payload - # matches what we expect. - assert sql.execute.called - sql_arg = sql.execute.call_args.args[0] - assert "sum(x) as total" in sql_arg - assert "count(*) as n" in sql_arg + # ``save_setting`` now goes through SqlExecutor.upsert(); the + # JSON payload lands on the value-cols dict. + assert sql.upsert.called + kwargs = sql.upsert.call_args.kwargs + # Either positional (table, key_cols, value_cols) or keyword form. + if "value_cols" in kwargs: + payload = kwargs["value_cols"]["setting_value"] + else: + payload = sql.upsert.call_args.args[2]["setting_value"] + assert "sum(x) as total" in payload + assert "count(*) as n" in payload def test_save_empty_list_writes_empty_json(self, svc): s, sql = svc assert s.save_custom_metrics([]) == [] - assert sql.execute.called + assert sql.upsert.called + + +# --------------------------------------------------------------------------- +# AppSettingsService.get_config — WorkspaceConfig round-trip +# --------------------------------------------------------------------------- + + +class TestAppSettingsGetConfig: + """``get_config`` deserialization must not depend on private blueprint APIs. + + The pre-refactor code called :func:`Installation._unmarshal_type` + (leading underscore — private). The fix routes through Pydantic's + public :class:`~pydantic.TypeAdapter` instead. These tests pin the + new contract: + + 1. **No private-API regressions.** ``Installation._unmarshal_type`` + must not be touched even if a stale import survives — + :mod:`backend.services.app_settings_service` must not import + ``Installation`` at all. + 2. **Round-trip fidelity.** A rich, nested ``WorkspaceConfig`` + written via ``as_dict()`` re-emerges from ``get_config()`` + byte-equivalent on the ``as_dict()`` side. The point of the + fallback is to never lose user data. + 3. **Safe fallbacks.** Corrupt rows (bad JSON, wrong shape) yield + an empty-default ``WorkspaceConfig`` instead of crashing the + admin UI. A regression here means a single bad settings row + takes down the whole config page. + """ + + @pytest.fixture + def svc(self, sql_executor_mock): + from databricks_labs_dqx_app.backend.services.app_settings_service import AppSettingsService + + return AppSettingsService(sql_executor_mock), sql_executor_mock + + def test_module_does_not_import_private_blueprint_api(self): + # Belt-and-suspenders against the reviewer-flagged anti-pattern. + # If anyone re-imports Installation here, this test fails — a + # regression would re-couple us to a private blueprint API. + import importlib + + mod = importlib.import_module("databricks_labs_dqx_app.backend.services.app_settings_service") + assert not hasattr(mod, "Installation"), ( + "app_settings_service must not import blueprint's Installation — " + "deserialization goes through pydantic.TypeAdapter instead." + ) + + def test_returns_default_when_no_row(self, svc): + from databricks.labs.dqx.config import WorkspaceConfig + + s, sql = svc + sql.query.return_value = [] + result = s.get_config() + assert isinstance(result, WorkspaceConfig) + assert result.run_configs == [] + + def test_round_trips_a_rich_nested_config_via_pydantic(self, svc): + """The critical contract: nothing on the as_dict() round-trip is lost. + + Any future schema change in WorkspaceConfig must continue + working without touching THIS service's code — that's the whole + point of delegating to TypeAdapter rather than hand-rolling + ``dataclasses.fields + getattr``. + """ + from databricks.labs.dqx.config import ( + AnomalyConfig, + ExtraParams, + InputConfig, + LLMConfig, + LLMModelConfig, + OutputConfig, + ProfilerConfig, + RunConfig, + WorkspaceConfig, + ) + + s, sql = svc + original = WorkspaceConfig( + run_configs=[ + RunConfig( + name="prod", + input_config=InputConfig(location="c.s.src"), + output_config=OutputConfig(location="c.s.out"), + quarantine_config=OutputConfig(location="c.s.q"), + profiler_config=ProfilerConfig(), + checks_location="checks.yml", + warehouse_id="wh-prod", + reference_tables={"r1": InputConfig(location="c.s.ref")}, + custom_check_functions={"fn": "sql"}, + anomaly_config=AnomalyConfig(), + ) + ], + log_level="DEBUG", + serverless_clusters=False, + extra_params=ExtraParams( + result_column_names={"flag": "is_bad"}, + user_metadata={"team": "data"}, + suppress_skipped=True, + variables={}, + ), + profiler_override_clusters={"a": "c1"}, + quality_checker_override_clusters={"b": "c2"}, + e2e_override_clusters={}, + anomaly_override_clusters={}, + profiler_spark_conf={"spark.x": "1"}, + quality_checker_spark_conf={}, + e2e_spark_conf={}, + anomaly_spark_conf={}, + custom_metrics=["count(*) AS n"], + llm_config=LLMConfig(model=LLMModelConfig(model_name="m", api_key="k", api_base="b")), + ) + sql.query.return_value = [(json.dumps(original.as_dict()),)] + + loaded = s.get_config() + + assert isinstance(loaded, WorkspaceConfig) + # The exact dict shape must match — every nested + # dataclass (RunConfig, ExtraParams, LLMConfig.model, etc.) + # round-trips back to its concrete type, not to a dict. + assert loaded.as_dict() == original.as_dict() + # Spot-check that the result is fully-typed (not bare dicts). + assert isinstance(loaded.run_configs[0], RunConfig) + assert isinstance(loaded.run_configs[0].input_config, InputConfig) + assert isinstance(loaded.extra_params, ExtraParams) + assert isinstance(loaded.llm_config, LLMConfig) + assert isinstance(loaded.llm_config.model, LLMModelConfig) + + def test_invalid_json_falls_back_to_default(self, svc, caplog): + from databricks.labs.dqx.config import WorkspaceConfig + + s, sql = svc + sql.query.return_value = [("{this is not valid json",)] + with caplog.at_level("WARNING"): + result = s.get_config() + assert isinstance(result, WorkspaceConfig) + assert result.run_configs == [] + # Operator-visible signal — a regression that swallows the warning + # silently would let corrupt rows hide indefinitely. + assert any("not valid JSON" in r.getMessage() for r in caplog.records) + + def test_schema_mismatch_falls_back_to_default(self, svc, caplog): + """ValidationError path: well-formed JSON but wrong shape.""" + from databricks.labs.dqx.config import WorkspaceConfig + + s, sql = svc + # ``run_configs`` must be a list of RunConfig dicts; a bare + # string will trip Pydantic's ValidationError. + sql.query.return_value = [(json.dumps({"run_configs": "definitely-not-a-list"}),)] + with caplog.at_level("WARNING"): + result = s.get_config() + assert isinstance(result, WorkspaceConfig) + assert result.run_configs == [] + assert any("schema validation" in r.getMessage() for r in caplog.records) # --------------------------------------------------------------------------- @@ -150,33 +307,35 @@ class TestSchedulerLoadCustomMetrics: """The scheduler's hot path must never crash on a missing/blank setting.""" @pytest.fixture - def scheduler(self): - from databricks_labs_dqx_app.backend.services.scheduler_service import SchedulerService - - # __init__ wires SqlExecutor for itself; we bypass it and inject - # the minimal attributes ``_load_custom_metrics`` actually touches. - svc = SchedulerService.__new__(SchedulerService) - svc._sql = MagicMock() - svc._settings_table = "dqx.public.dq_app_settings" - return svc + def scheduler(self, make_scheduler): + # Built via the real ``SchedulerService(...)`` constructor — + # see :func:`make_scheduler` in conftest.py for the rationale. + # ``_load_custom_metrics`` reads through ``self._oltp_sql.query``, + # which the factory exposes as ``mocks.oltp``. + return make_scheduler() def test_returns_empty_when_no_rows(self, scheduler): - scheduler._sql.query.return_value = [] - assert scheduler._load_custom_metrics() == [] + svc, mocks = scheduler + mocks.oltp.query.return_value = [] + assert svc._load_custom_metrics() == [] def test_returns_empty_when_value_is_null(self, scheduler): - scheduler._sql.query.return_value = [(None,)] - assert scheduler._load_custom_metrics() == [] + svc, mocks = scheduler + mocks.oltp.query.return_value = [(None,)] + assert svc._load_custom_metrics() == [] def test_parses_json_list(self, scheduler): - scheduler._sql.query.return_value = [(json.dumps(["sum(x) as total"]),)] - assert scheduler._load_custom_metrics() == ["sum(x) as total"] + svc, mocks = scheduler + mocks.oltp.query.return_value = [(json.dumps(["sum(x) as total"]),)] + assert svc._load_custom_metrics() == ["sum(x) as total"] def test_swallows_query_exception(self, scheduler): - scheduler._sql.query.side_effect = RuntimeError("warehouse offline") + svc, mocks = scheduler + mocks.oltp.query.side_effect = RuntimeError("warehouse offline") # Must not propagate — the scheduler tick keeps running. - assert scheduler._load_custom_metrics() == [] + assert svc._load_custom_metrics() == [] def test_filters_invalid_entries(self, scheduler): - scheduler._sql.query.return_value = [(json.dumps(["good as g", 42, "", None]),)] - assert scheduler._load_custom_metrics() == ["good as g"] + svc, mocks = scheduler + mocks.oltp.query.return_value = [(json.dumps(["good as g", 42, "", None]),)] + assert svc._load_custom_metrics() == ["good as g"] diff --git a/app/tests/test_lint_policy.py b/app/tests/test_lint_policy.py new file mode 100644 index 000000000..a8df11c11 --- /dev/null +++ b/app/tests/test_lint_policy.py @@ -0,0 +1,437 @@ +"""Locks down project-wide lint policy decisions. + +These aren't *behaviour* tests — they're *convention* tests that +prevent the codebase from quietly drifting away from policies that +were decided on review (and documented in ``pyproject.toml``). + +Currently locked down: + +* **BLE001** (broad-except) — the project does NOT enforce BLE001 + because several backend surfaces require best-effort resilience + (lifespan teardown, background-thread refresh loops, route-level + per-item recovery). The documented policy lives in + ``[tool.ruff.lint]`` in ``pyproject.toml`` with mandatory inline- + comment justification at each site. + + This test guards both halves of that contract: the policy block + exists in ``pyproject.toml`` AND no source file has re-introduced + the rejected per-line ``# noqa: BLE001`` pattern. Per AGENTS.md + rule #6 ("Never disable linting to silence issues"), the + per-line escape hatch is the wrong place for this — fix the code + or document the exemption project-wide. + +* **reportCallIssue / reportArgumentType on psycopg ``execute``** — + psycopg's PEP-675 stubs require :meth:`Cursor.execute`'s ``query`` + argument to be a :class:`typing.LiteralString`. Backend code that + composes SQL programmatically (validated identifiers, escaped + values) used to silence basedpyright at every call site with + ``# pyright: ignore[reportCallIssue, reportArgumentType]``. That + also violates AGENTS.md rule #6 — the documented fix is the + ``run_trusted_sql`` helper in :mod:`backend.pg_executor`, which + performs the cast once in a single auditable wrapper. This test + pins the convention. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +SRC_ROOT = REPO_ROOT / "src" +TESTS_ROOT = REPO_ROOT / "tests" +PYPROJECT = REPO_ROOT / "pyproject.toml" + +# Matches the actual silencing pattern: an ``except`` clause carrying +# an inline noqa directive for the broad-except rule (with or without +# whitespace after the colon, with or without an exception alias). +# Deliberately anchored to the ``except`` keyword so docstrings, +# error messages, and *this very file's* policy-text references +# don't cause false positives. The regex literal below is the +# canonical form callers should look for. +_BLE_NOQA_ON_EXCEPT = re.compile(r"^\s*except\b.*#\s*noqa\s*:\s*BLE001\b") + +# Matches the actual silencing pattern for psycopg's PEP-675 +# constraint on ``Cursor.execute``: a line that calls ``.execute(`` +# AND carries an inline pyright-ignore for ``reportCallIssue`` or +# ``reportArgumentType``. Anchored on ``.execute(`` so generic +# pyright-ignore directives elsewhere in the codebase aren't false- +# positive hits. +_EXECUTE_PYRIGHT_IGNORE = re.compile( + r"\.execute\b.*#\s*pyright:\s*ignore\[[^\]]*" r"(?:reportCallIssue|reportArgumentType)" +) + + +class TestBleNoqaPolicy: + """Per-line ``# noqa: BLE001`` is banned by the documented policy.""" + + @pytest.mark.parametrize( + "root,description", + [ + (SRC_ROOT, "src/"), + (TESTS_ROOT, "tests/"), + ], + ) + def test_no_per_line_noqa_ble001_in_source_or_tests(self, root, description): + """Any ``except ...: # noqa: BLE001`` must be replaced by a + substantive inline comment + the documented project-wide + policy in ``pyproject.toml``. + + Adding one in source/tests means either (a) the code needs a + substantive comment explaining the resilience contract + instead, or (b) BLE001 needs to be enabled with a documented + per-file ignore. Per-line silencing is rejected by both the + reviewer comment that created this policy and by AGENTS.md + rule #6. + """ + offenders: list[tuple[Path, int, str]] = [] + for py in root.rglob("*.py"): + # Skip any vendor/dist trees if a future build places .py + # files under them. The src layout shouldn't ship any, but + # the rglob is permissive. + if any(part in {".venv", "__pycache__", "__dist__"} for part in py.parts): + continue + text = py.read_text(encoding="utf-8") + for lineno, line in enumerate(text.splitlines(), start=1): + if _BLE_NOQA_ON_EXCEPT.match(line): + offenders.append((py.relative_to(REPO_ROOT), lineno, line.strip())) + + assert not offenders, ( + f"Per-line BLE001 silencing re-introduced on except clauses in {description}. " + f"Per the policy block in pyproject.toml, the broad-except " + f"resilience pattern is documented project-wide; individual " + f"sites must NOT silence the rule per-line. Offenders:\n" + + "\n".join(f" {p}:{n} {snippet}" for p, n, snippet in offenders) + ) + + def test_pyproject_has_documented_ble001_policy_block(self): + """The single source of truth for the broad-except policy. + + If someone removes the policy block from ``pyproject.toml`` + AND the codebase still has broad excepts, future contributors + have no documented rationale for the pattern and the next + reviewer will (correctly) flag it again. Keep the policy + comment and the code aligned. + """ + text = PYPROJECT.read_text(encoding="utf-8") + assert "[tool.ruff.lint]" in text, ( + "pyproject.toml must declare a ``[tool.ruff.lint]`` block " "to anchor the documented BLE001 policy." + ) + assert "BLE001" in text, ( + "pyproject.toml must document the project's BLE001 policy " + "(it's the single source of truth that replaces the " + "per-line ``# noqa: BLE001`` comments)." + ) + # And the policy must explicitly forbid the per-line form so a + # future reader looking at the policy block immediately sees + # the prohibition rather than having to grep AGENTS.md. + assert "noqa: BLE001" in text or "noqa:BLE001" in text, ( + "The pyproject.toml policy block must mention the rejected " + "``# noqa: BLE001`` pattern explicitly so contributors see " + "the prohibition in the same place as the policy." + ) + + +class TestPsycopgExecutePyrightIgnorePolicy: + """``Cursor.execute(sql)`` must route through the trust-boundary helper.""" + + @pytest.mark.parametrize( + "root,description", + [ + (SRC_ROOT, "src/"), + (TESTS_ROOT, "tests/"), + ], + ) + def test_no_per_call_pyright_ignore_on_cursor_execute(self, root, description): + """Forbid inline pyright-ignore on cursor execute calls. + + The rejected pattern is a cursor execute call that suppresses + ``reportCallIssue`` or ``reportArgumentType`` on the same + line. The replacement is + :func:`backend.pg_executor.run_trusted_sql`, which performs + the ``cast(LiteralString, sql)`` once and documents the + trust boundary in one place. + + If you find yourself wanting to suppress the warning at a + new call site, route through ``run_trusted_sql`` instead. + If the SQL string contains untrusted input, parameterise + the query (pass values via the ``params`` arg) rather than + widening the helper's trust contract. + """ + offenders: list[tuple[Path, int, str]] = [] + for py in root.rglob("*.py"): + if any(part in {".venv", "__pycache__", "__dist__"} for part in py.parts): + continue + text = py.read_text(encoding="utf-8") + for lineno, line in enumerate(text.splitlines(), start=1): + if _EXECUTE_PYRIGHT_IGNORE.search(line): + offenders.append((py.relative_to(REPO_ROOT), lineno, line.strip())) + + assert not offenders, ( + f"Per-call pyright-ignore on ``.execute(...)`` re-introduced in {description}. " + f"Route the call through ``run_trusted_sql`` in pg_executor.py " + f"so the LiteralString cast lives in one auditable wrapper " + f"rather than scattered across call sites. Offenders:\n" + + "\n".join(f" {p}:{n} {snippet}" for p, n, snippet in offenders) + ) + + def test_run_trusted_sql_helper_is_importable_and_documented(self): + """The helper that replaces the per-call ignores must exist + and carry the trust-boundary docstring. If it gets renamed + or deleted, this test fails loud rather than letting the + per-call pattern silently creep back in. + + The helper now lives in :mod:`backend.pg_cursor_helpers` (a + psycopg-free module — see that module's docstring), but is + re-exported by :mod:`backend.pg_executor` for back-compat. + Both import paths are pinned so a refactor that drops either + re-export breaks loudly here. + """ + from databricks_labs_dqx_app.backend.pg_cursor_helpers import ( + run_trusted_sql as run_trusted_sql_canonical, + ) + from databricks_labs_dqx_app.backend.pg_executor import ( + run_trusted_sql as run_trusted_sql_reexport, + ) + + assert run_trusted_sql_canonical is run_trusted_sql_reexport, ( + "``backend.pg_executor.run_trusted_sql`` must be the SAME " + "object as ``backend.pg_cursor_helpers.run_trusted_sql`` — " + "a divergence means the re-export drifted from the canonical " + "definition and callers will silently use the wrong version." + ) + run_trusted_sql = run_trusted_sql_canonical + assert callable(run_trusted_sql), "run_trusted_sql must be a callable wrapper" + doc = (run_trusted_sql.__doc__ or "").lower() + # The docstring is the trust-boundary contract — these tokens + # are the load-bearing words. If any drift out, the next + # reviewer can't tell why ``cast(LiteralString, ...)`` is + # safe at the call sites. + for token in ("literalstring", "trust", "untrusted"): + assert token in doc, ( + f"run_trusted_sql.__doc__ must mention {token!r} " + "so the trust boundary is explicit in the helper " + "rather than implicit in the call sites." + ) + # The narrowed contract (post-review) must explicitly point + # callers with RUNTIME values at the parameterised sibling + # helper / psycopg's native binding — otherwise a future + # reader sees only the LiteralString cast and assumes the + # helper sanitises everything. + assert "run_parameterized_sql" in doc or "cur.execute(sql, params)" in doc, ( + "run_trusted_sql.__doc__ must point runtime-value callers " + "at ``run_parameterized_sql`` (or psycopg's native " + "``cur.execute(sql, params)`` binding) so the cast is not " + "mistaken for a value sanitiser." + ) + + def test_run_parameterized_sql_helper_is_importable_and_documented(self): + """Sibling of :func:`run_trusted_sql` for the "trusted + template + runtime values" case. It centralises the same + ``cast(LiteralString, ...)`` discipline (so call sites do + NOT need ``# pyright: ignore``) while binding values through + psycopg's native parameter binding so f-string value + interpolation is structurally impossible at the call site. + + Same dual-import pin as :func:`run_trusted_sql`: the canonical + definition lives in :mod:`backend.pg_cursor_helpers` and the + executor re-exports it. Both paths and identity-equality are + pinned.""" + from databricks_labs_dqx_app.backend.pg_cursor_helpers import ( + run_parameterized_sql as run_parameterized_sql_canonical, + ) + from databricks_labs_dqx_app.backend.pg_executor import ( + run_parameterized_sql as run_parameterized_sql_reexport, + ) + + assert run_parameterized_sql_canonical is run_parameterized_sql_reexport, ( + "``backend.pg_executor.run_parameterized_sql`` must be the " + "SAME object as the canonical definition in " + "``backend.pg_cursor_helpers`` — re-export drift means " + "callers silently bind through the wrong version." + ) + run_parameterized_sql = run_parameterized_sql_canonical + assert callable(run_parameterized_sql), "run_parameterized_sql must be a callable wrapper" + doc = (run_parameterized_sql.__doc__ or "").lower() + # Load-bearing tokens — explain the two halves of the + # contract (the template is trusted like in run_trusted_sql, + # the values go through psycopg's binder). + for token in ("literalstring", "trusted", "psycopg", "binding"): + assert token in doc, ( + f"run_parameterized_sql.__doc__ must mention {token!r} " + "so the trust contract (trusted template + bound values) " + "is explicit in the helper." + ) + + def test_helper_bodies_are_the_only_raw_cursor_execute_sites(self): + """The two trust-boundary helper bodies in + :mod:`backend.pg_cursor_helpers` are the ONLY places in the + backend that may call ``cur.execute(...)`` directly. Every + other caller — pg_executor's own methods, the migration + runner, anything new — must route through the helpers so + the ``cast(LiteralString, ...)`` happens in exactly one + auditable place. + + This is stricter than the previous form of the test (which + only checked ``pg_executor.py``): after the helpers moved + to their own psycopg-free module the executor file now + contains ZERO raw ``cur.execute`` calls, and the two + canonical raw calls live in ``pg_cursor_helpers.py``. + """ + backend = SRC_ROOT / "databricks_labs_dqx_app" / "backend" + helpers_text = (backend / "pg_cursor_helpers.py").read_text(encoding="utf-8") + executor_text = (backend / "pg_executor.py").read_text(encoding="utf-8") + + # The two helper bodies are the only legitimate raw call + # sites. Counting matches (rather than substring) keeps a + # comment that happens to mention ``cur.execute(`` from + # blowing the test up. + raw_in_helpers = sum(1 for line in helpers_text.splitlines() if re.match(r"^\s*(_ = )?cur\.execute\(", line)) + raw_in_executor = sum(1 for line in executor_text.splitlines() if re.match(r"^\s*(_ = )?cur\.execute\(", line)) + + assert raw_in_helpers == 2, ( + f"Expected exactly 2 raw ``cur.execute(...)`` in " + f"pg_cursor_helpers.py (the bodies of ``run_trusted_sql`` " + f"and ``run_parameterized_sql``), found {raw_in_helpers}. " + "If you added a third helper, update the policy test " + "deliberately rather than letting the count drift." + ) + assert raw_in_executor == 0, ( + f"Expected ZERO raw ``cur.execute(...)`` in pg_executor.py — " + f"every cursor call must route through ``run_trusted_sql`` / " + f"``run_parameterized_sql`` from ``backend.pg_cursor_helpers``. " + f"Found {raw_in_executor}." + ) + + # The executor must still USE the helpers — at minimum on its + # ``execute``, ``query``, and ``query_dicts`` data paths. + # Counting helper invocations protects against a refactor + # that silently routes a method through ``conn.execute`` or + # ``conn.cursor().execute`` directly. + helper_calls = executor_text.count("run_trusted_sql(cur,") + assert helper_calls >= 3, ( + f"Expected at least 3 ``run_trusted_sql(cur, ...)`` call sites in " + f"pg_executor.py (execute, query, query_dicts), found {helper_calls}." + ) + + def test_pg_cursor_helpers_is_psycopg_free_at_runtime(self): + """The whole point of splitting :mod:`backend.pg_cursor_helpers` + out of :mod:`backend.pg_executor` was to give consumers (notably + :mod:`backend.migrations.postgres`) a way to import the trust- + boundary helpers WITHOUT transitively pulling in :mod:`psycopg`. + + That property breaks the moment someone moves the ``Cursor`` + import out of the ``TYPE_CHECKING`` guard. This test pins both + halves of the contract — the file must (a) import ``Cursor`` + from psycopg ONLY under the guard and (b) carry zero + unconditional ``import psycopg`` / ``from psycopg import …`` + statements. + + If a future change actually needs a runtime psycopg dependency + in this module, the helpers' purpose has changed and the + caller graph (especially ``migrations.postgres``) needs to be + re-audited; bumping this test deliberately at that point is + the intended remediation. + """ + backend = SRC_ROOT / "databricks_labs_dqx_app" / "backend" + text = (backend / "pg_cursor_helpers.py").read_text(encoding="utf-8") + + # Walk lines, tracking whether we are inside the + # ``if TYPE_CHECKING:`` block. Any ``import psycopg`` / + # ``from psycopg import …`` line OUTSIDE that block breaks + # the runtime-free promise. + in_type_checking_block = False + offenders: list[tuple[int, str]] = [] + # A new top-level statement (non-indented, non-blank, non-comment) + # exits the ``if TYPE_CHECKING:`` indented block. + for lineno, line in enumerate(text.splitlines(), start=1): + stripped = line.strip() + if stripped.startswith("if TYPE_CHECKING") and stripped.endswith(":"): + in_type_checking_block = True + continue + if in_type_checking_block: + # Indented or blank/comment lines stay inside the guard; + # a non-indented, non-blank, non-comment line leaves it. + if line and not line.startswith((" ", "\t")) and not stripped.startswith("#"): + in_type_checking_block = False + if in_type_checking_block: + continue + if re.match(r"^\s*(import\s+psycopg|from\s+psycopg(\s|\.))", line): + offenders.append((lineno, stripped)) + + assert not offenders, ( + "pg_cursor_helpers.py imported psycopg outside the " + "``if TYPE_CHECKING:`` guard. That re-couples every " + "caller of these helpers (including ``migrations.postgres`` " + "and any future psycopg-free consumer) to psycopg, " + "undoing the whole point of the split. Offenders:\n" + + "\n".join(f" line {n}: {snippet}" for n, snippet in offenders) + ) + + # And, defensively, confirm the guard exists at all — a + # refactor that drops it entirely would still pass the + # negative check above (no psycopg imports anywhere) but + # would mean the module no longer carries the type hints + # the call-sites rely on. + assert "if TYPE_CHECKING:" in text, ( + "pg_cursor_helpers.py is expected to import ``Cursor`` " + "for type hints under an ``if TYPE_CHECKING:`` guard. " + "Dropping the guard means losing the type surface the " + "helpers' signatures are written against." + ) + + def test_migrations_postgres_does_not_transitively_import_psycopg(self): + """Companion of the previous test, from the consumer's side. + + :mod:`backend.migrations.postgres` MUST be importable in an + environment that does NOT have :mod:`psycopg` installed. + Before the helper split, importing it would crash with + ``ModuleNotFoundError: psycopg`` because + ``from ..pg_executor import run_trusted_sql`` pulled in + psycopg via the executor's top-level imports. + + We can't simulate "psycopg is not installed" cleanly in + unit tests, so we assert the structural precondition: the + runner imports the helpers from :mod:`pg_cursor_helpers` + (which is psycopg-free per the test above), not from + :mod:`pg_executor`. This catches a regression as soon as + anyone re-routes the import without realising the cost. + """ + text = (SRC_ROOT / "databricks_labs_dqx_app" / "backend" / "migrations" / "postgres.py").read_text( + encoding="utf-8" + ) + + # The forbidden pattern: importing the trust-boundary + # helpers from ``pg_executor`` instead of ``pg_cursor_helpers``. + # Matches both relative (``..pg_executor``) and absolute + # (``databricks_labs_dqx_app.backend.pg_executor``) forms. + forbidden = re.compile( + r"^\s*from\s+(?:\.\.|databricks_labs_dqx_app\.backend\.)pg_executor\s+import\s+.*\b" + r"(run_trusted_sql|run_parameterized_sql)\b" + ) + offenders = [(n, line) for n, line in enumerate(text.splitlines(), 1) if forbidden.match(line)] + assert not offenders, ( + "migrations/postgres.py imports a trust-boundary helper " + "from ``backend.pg_executor`` — that re-couples the " + "migration runner to ``psycopg`` and resurrects the lazy-" + "import dance in ``app.py``. Import from " + "``backend.pg_cursor_helpers`` instead. Offenders:\n" + + "\n".join(f" line {n}: {line.strip()}" for n, line in offenders) + ) + + # And confirm the migration runner DOES use the canonical + # psycopg-free import path — symmetric assertion so a + # refactor that deletes the helpers from the runner entirely + # (replacing with raw cursor calls) fails loudly here. + expected = re.compile( + r"^\s*from\s+(?:\.\.|databricks_labs_dqx_app\.backend\.)pg_cursor_helpers\s+import\s+.*\b" + r"(run_trusted_sql|run_parameterized_sql)\b" + ) + assert any(expected.match(line) for line in text.splitlines()), ( + "migrations/postgres.py is expected to import the trust-" + "boundary helpers from ``backend.pg_cursor_helpers``. " + "If the runner stopped using them, route every cursor " + "call through the helpers instead of dropping them." + ) diff --git a/app/tests/test_migration_runner.py b/app/tests/test_migration_runner.py new file mode 100644 index 000000000..d1a0fc586 --- /dev/null +++ b/app/tests/test_migration_runner.py @@ -0,0 +1,230 @@ +"""Unit tests for the Delta :class:`MigrationRunner`. + +Scope: the structural invariants that the Postgres runner gets "for +free" from transactional atomicity but the Delta runner has to enforce +manually — see the "Atomicity model" section of +:mod:`databricks_labs_dqx_app.backend.migrations` for the full +recovery contract. Plus the identifier-quoting contract that both +runners share via ``executor.q()``. + +The runner itself is exercised in integration tests against a real +SQL warehouse; here we focus on the pure-Python guards that prevent +future migration authors from silently breaking the contract: + +1. ``_validate_template_safe`` rejects any template containing + ``;`` inside a single-quoted string literal, and any + ``{catalog}``/``{schema}`` placeholder appearing in literal + position (where the quoted-form substitution would be incorrect). +2. Every entry in the live :data:`MIGRATIONS` list passes the + validator — regression coverage against a future author editing + a template into a no-longer-safe shape. +3. The runner consistently uses ``sql.q()`` for catalog/schema + identifiers — so a deployment with a hyphenated Databricks + catalog name like ``prod-east`` emits parseable DDL instead of + raw ``CREATE SCHEMA prod-east.dqx_studio``. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from databricks_labs_dqx_app.backend.migrations import ( + MIGRATIONS, + MigrationRunner, + _validate_template_safe, +) +from databricks_labs_dqx_app.backend.sql_executor import SqlExecutor + + +# --------------------------------------------------------------------------- +# Template scanner: positive + negative + live regression +# --------------------------------------------------------------------------- + + +class TestValidateTemplateSafe: + """Direct tests of the import-time scanner.""" + + @pytest.mark.parametrize( + "template", + [ + # Empty / trivial + "", + "SELECT 1", + # Multi-statement, no literals + "CREATE TABLE foo (x INT); ALTER TABLE foo ADD COLUMN y INT;", + # Literals without forbidden characters + "ALTER TABLE foo ADD CONSTRAINT chk CHECK (status IN ('a','b','c'))", + # Apostrophe escape inside literal (``''`` is one literal apostrophe) + "INSERT INTO foo VALUES ('it''s fine')", + # Adjacent literals (closing + opening) — the scanner must NOT + # treat `''` between two separate strings as an escape. + "INSERT INTO foo VALUES ('a','b')", + # Semicolon AFTER closing the literal is fine + "INSERT INTO foo VALUES ('a'); INSERT INTO foo VALUES ('b')", + # Placeholders in object-name positions are the supported use case + "CREATE TABLE IF NOT EXISTS {catalog}.{schema}.foo (x INT)", + "INSERT INTO {catalog}.{schema}.bar (note) VALUES ('description with {curly} braces')", + ], + ) + def test_safe_templates_pass(self, template: str) -> None: + # Must not raise. + _validate_template_safe(template) + + @pytest.mark.parametrize( + ("template", "expected_fragment"), + [ + # Invariant 1: ``;`` inside a CHECK constraint literal. + ( + "ALTER TABLE foo ADD CONSTRAINT chk CHECK (label IN ('one;two', 'three'))", + "';' inside", + ), + # Invariant 1: ``;`` inside a default value + ("ALTER TABLE foo ADD COLUMN note STRING DEFAULT 'a; b'", "';' inside"), + # Invariant 1: ``;`` inside an escaped-apostrophe literal — the + # scanner must keep the literal flag set across ``''``. + ("INSERT INTO foo VALUES ('it''s; bad')", "';' inside"), + # Invariant 1: trailing ``;`` inside literal + ("INSERT INTO foo VALUES ('trailing;')", "';' inside"), + # Invariant 2: ``{catalog}`` placeholder inside literal + ("INSERT INTO foo VALUES ('catalog={catalog}')", "{catalog}"), + # Invariant 2: ``{schema}`` placeholder inside literal + ("INSERT INTO foo VALUES ('schema={schema}')", "{schema}"), + ], + ) + def test_unsafe_templates_raise(self, template: str, expected_fragment: str) -> None: + with pytest.raises(AssertionError, match="inside a single-quoted string literal"): + _validate_template_safe(template) + # The specific violation should appear in the message so the + # author knows which invariant they broke. + with pytest.raises(AssertionError) as excinfo: + _validate_template_safe(template) + assert expected_fragment in str(excinfo.value) + + def test_assertion_message_includes_offset_and_excerpt(self) -> None: + """The error message must be actionable — give the offset and a snippet.""" + prefix = "ALTER TABLE foo ADD COLUMN name STRING; " + template = prefix + "INSERT INTO foo VALUES ('bad;value')" + with pytest.raises(AssertionError) as excinfo: + _validate_template_safe(template) + msg = str(excinfo.value) + assert "offset" in msg + # Excerpt should include some text near the offending ``;`` — + # we don't pin the exact window because the window size is an + # implementation detail, but the literal contents should appear. + assert "bad" in msg + + +class TestLiveMigrationsAreTemplateSafe: + """Regression: every shipped migration template must pass the validator. + + Without this, the runner's startup-time check would still catch a + violation at app boot — but a unit-test failure on PR is much + faster feedback than a deploy-time crash. + """ + + def test_no_migration_template_violates_invariants(self) -> None: + for migration in MIGRATIONS: + try: + _validate_template_safe(migration.sql_template) + except AssertionError as exc: # pragma: no cover - surfaced via fail() + pytest.fail( + f"Migration v{migration.version} ({migration.description}) " + f"violates a template scanner invariant: {exc}" + ) + + +# --------------------------------------------------------------------------- +# Identifier-quoting contract — review item #8. +# --------------------------------------------------------------------------- + + +class TestSqlExecutorQuoting: + """Direct tests of :meth:`SqlExecutor.q` (the Delta quoter).""" + + @pytest.fixture + def executor(self) -> SqlExecutor: + return SqlExecutor(ws=MagicMock(name="WorkspaceClient"), warehouse_id="wh", catalog="c", schema="s") + + @pytest.mark.parametrize( + ("identifier", "expected"), + [ + ("foo", "`foo`"), + # The exact hyphenated-catalog case the reviewer flagged + ("prod-east", "`prod-east`"), + ("team-data-platform", "`team-data-platform`"), + # Reserved-word column — backticks are mandatory + ("check", "`check`"), + # Internal backticks must be doubled per Databricks SQL convention, + # mirroring PgExecutor.q which doubles internal ``"``. + ("weird`name", "`weird``name`"), + # Pre-doubled (no over-escaping) + ("a``b", "`a````b`"), + # Unicode is fine — Databricks identifiers permit it when quoted + ("café", "`café`"), + ], + ) + def test_quotes_identifier(self, executor: SqlExecutor, identifier: str, expected: str) -> None: + assert executor.q(identifier) == expected + + +class TestMigrationRunnerUsesQuotedIdentifiers: + """End-to-end regression for the catalog/schema raw-interpolation bug. + + Builds a runner against a fake executor whose ``catalog`` is the + reviewer-flagged ``prod-east`` and asserts every SQL the runner + issues references the *quoted* form. A regression here would mean + a hyphenated-catalog deployment silently emits parse-invalid DDL. + """ + + @pytest.fixture + def runner_state(self) -> tuple[MigrationRunner, MagicMock]: + ws = MagicMock(name="WorkspaceClient") + # Use a real SqlExecutor so we exercise its real ``.q()`` path + # rather than re-mocking the quoting behaviour. + sql = SqlExecutor(ws=ws, warehouse_id="wh", catalog="prod-east", schema="dqx_studio") + runner = MigrationRunner(sql=sql) + return runner, ws + + def test_meta_table_is_quoted(self, runner_state: tuple[MigrationRunner, MagicMock]) -> None: + runner, _ = runner_state + # `prod-east`.`dqx_studio`.dq_migrations + assert runner._meta_table == "`prod-east`.`dqx_studio`.dq_migrations" + + def test_ensure_schema_emits_quoted_ddl(self, runner_state: tuple[MigrationRunner, MagicMock]) -> None: + runner, _ = runner_state + # Mock out the execute path so we can read the SQL it would send. + captured: list[str] = [] + runner._sql.execute_no_schema = lambda s: captured.append(s) # type: ignore[method-assign] + runner._ensure_schema() + assert captured == ["CREATE SCHEMA IF NOT EXISTS `prod-east`.`dqx_studio`"] + + def test_apply_substitutes_quoted_form_into_placeholders( + self, runner_state: tuple[MigrationRunner, MagicMock] + ) -> None: + """``{catalog}`` / ``{schema}`` must receive the backtick-quoted form.""" + runner, _ = runner_state + captured: list[str] = [] + runner._sql.execute = lambda s, **_: captured.append(s) # type: ignore[method-assign] + + # Use a tiny ad-hoc migration that won't trip the idempotency + # swallow list (so we can read the captured SQL directly). + from databricks_labs_dqx_app.backend.migrations import DeltaMigration + + m = DeltaMigration( + version=999, + description="test", + sql_template="CREATE TABLE IF NOT EXISTS {catalog}.{schema}.test_t (x INT)", + oltp_fallback=False, + ) + runner._apply(m) + + # The first captured SQL is the CREATE TABLE; the second is the + # INSERT INTO dq_migrations. Both must reference the quoted form. + assert captured[0] == "CREATE TABLE IF NOT EXISTS `prod-east`.`dqx_studio`.test_t (x INT)" + assert captured[1].startswith("INSERT INTO `prod-east`.`dqx_studio`.dq_migrations") + # And specifically NOT the raw form + assert "prod-east.dqx_studio" not in " ".join( + captured + ), "Found raw (un-quoted) interpolation — hyphenated catalogs would emit parse-invalid DDL" diff --git a/app/tests/test_pg_executor.py b/app/tests/test_pg_executor.py new file mode 100644 index 000000000..d1289162f --- /dev/null +++ b/app/tests/test_pg_executor.py @@ -0,0 +1,1666 @@ +"""Unit tests for :mod:`backend.pg_executor`. + +Strategy +-------- +The real :class:`PgExecutor` opens a psycopg connection pool and spawns +a daemon thread inside ``__init__``. Both side-effects are off-limits in +a unit-test environment (we don't want a live Postgres dependency, a +background thread surviving the test, or a flaky pool open). + +Two tactics keep the tests fast, hermetic, and focused: + +1. **Pure helpers are tested as pure functions.** ``_to_text``, + ``_pg_render_value``, and the SQL-shape produced by ``upsert`` need + nothing from the surrounding class state, so we test them at the + module level without instantiating ``PgExecutor``. + +2. **Method-level tests build a half-initialised instance via + ``__new__``** rather than going through ``__init__``. The + ``_make_pg_executor`` helper below sets only the attributes a given + test needs and mocks the pool / token holder / stop event, so we + exercise the behaviour-under-test in isolation without touching the + real psycopg pool or starting the refresher thread. + +The :class:`PgExecutor` end-to-end happy path (real pool open, real +SCRAM handshake, real token rotation) is covered by integration tests +that run against a live Lakebase instance — out of scope here. +""" + +from __future__ import annotations + +import datetime as dt +import logging +import threading +from decimal import Decimal +from unittest.mock import MagicMock, patch + +import pytest + +from databricks_labs_dqx_app.backend.pg_executor import ( + PgExecutor, + _generate_token, + _pg_render_value, + _to_text, + _TokenHolder, + build_pg_executor, +) +from databricks_labs_dqx_app.backend.sql_executor import RawSql + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_pg_executor( + *, + schema: str = "public", + database: str = "dqx", + instance_name: str = "test-instance", + token_refresh_seconds: int = 0, + token_refresh_retry_seconds: int = 10, + token_refresh_retry_jitter: float = 0.3, + token_refresh_max_failures: int = 12, + initial_token: str = "tok-initial", +) -> PgExecutor: + """Build a method-callable :class:`PgExecutor` without running ``__init__``. + + Skips the real ``ConnectionPool`` open, the bootstrap + ``_generate_token`` call, and the refresher-thread spawn. + + *Why ``__new__`` here, unlike* ``conftest.make_scheduler`` *which uses + the real constructor?* ``SchedulerService.__init__`` is pure + attribute assignment, so injecting dependencies through the real + constructor is straightforward. ``PgExecutor.__init__``, in + contrast, performs three unavoidable side effects: it calls + ``_generate_token`` (a network round-trip to the SDK), opens a + real ``ConnectionPool`` (a TCP connect to Postgres), and starts a + daemon thread. A unit test cannot perform any of these. Adding + constructor kwargs to gate each side effect would expose + test-only seams in production code — worse than the localised + ``__new__`` here, which is contained to this single test helper. + The ``_make_pg_executor`` factory makes the asymmetry explicit + and keeps the boundary in test code. + + The returned instance has the minimum attribute surface every + method under test reads: + + - ``_ws`` / ``_instance_name`` — used by the refresh loop and by + :func:`build_pg_executor` smoke tests. + - ``_schema`` / ``_database`` — used by ``schema`` / ``database`` / + ``fqn`` / ``catalog`` properties. + - ``_token_holder`` — real :class:`_TokenHolder` so the refresh-loop + tests can read its mutated state. + - ``_connect_kwargs`` — real ``dict`` so the refresh-loop tests can + observe the mutated ``password`` key. + - ``_pool`` — :class:`MagicMock` with a context-manager-shaped + ``connection()`` so the executor methods that go through the + pool can be exercised without psycopg. + - ``_stop`` — real :class:`threading.Event` so the refresh-loop + tests can terminate the loop deterministically. + + Tests that need to inspect the inner cursor build it inline (see + ``_cursor_of`` below) rather than baking it into the helper to + keep this factory uncluttered. + """ + inst = PgExecutor.__new__(PgExecutor) + inst._ws = MagicMock(name="WorkspaceClient") + inst._instance_name = instance_name + inst._database = database + inst._schema = schema + inst._username = "test-user" + inst._host = "test-host" + inst._port = 5432 + inst._token_refresh_seconds = token_refresh_seconds + inst._token_refresh_retry_seconds = token_refresh_retry_seconds + inst._token_refresh_retry_jitter = token_refresh_retry_jitter + inst._token_refresh_max_failures = token_refresh_max_failures + inst._token_holder = _TokenHolder(initial_token) + inst._connect_kwargs = {"password": initial_token} + inst._last_successful_refresh_at = None + inst._consecutive_refresh_failures = 0 + + # Pool is a MagicMock with the context-manager protocol wired so + # ``with self._pool.connection() as conn:`` yields a deterministic + # mock connection. + pool = MagicMock(name="ConnectionPool") + conn_cm = pool.connection.return_value + conn_cm.__exit__.return_value = None # don't swallow exceptions + inst._pool = pool + + inst._stop = threading.Event() + inst._refresher = MagicMock(name="Thread") # never started in tests + return inst + + +def _conn_of(executor: PgExecutor) -> MagicMock: + """Return the connection mock the executor sees inside ``with self._pool.connection() as conn:``.""" + return executor._pool.connection.return_value.__enter__.return_value # type: ignore[no-any-return] + + +def _cursor_of(executor: PgExecutor) -> MagicMock: + """Return the cursor mock the executor uses inside ``with conn.cursor() as cur:``.""" + conn = _conn_of(executor) + cur_cm = conn.cursor.return_value + cur_cm.__exit__.return_value = None # don't swallow exceptions + return cur_cm.__enter__.return_value # type: ignore[no-any-return] + + +# =========================================================================== +# Pure helpers +# =========================================================================== + + +class TestToText: + """Coerces psycopg-typed cells to the Delta-compatible string output. + + The runtime type ordering matters: ``isinstance(True, int)`` is + ``True`` in Python, so the ``bool`` branch MUST be checked before + the int/float branch. A regression here would render booleans as + ``"1"`` / ``"0"`` instead of ``"true"`` / ``"false"`` and break + every service that round-trips a boolean through the API. + """ + + def test_none_passes_through(self) -> None: + assert _to_text(None) is None + + @pytest.mark.parametrize( + ("value", "expected"), + [ + ({"a": 1, "b": [2, 3]}, '{"a":1,"b":[2,3]}'), + ([1, 2, 3], "[1,2,3]"), + ({}, "{}"), + ([], "[]"), + ], + ) + def test_dict_and_list_render_compact_json(self, value: object, expected: str) -> None: + assert _to_text(value) == expected + + def test_datetime_renders_iso_8601(self) -> None: + ts = dt.datetime(2026, 5, 28, 9, 30, 0, tzinfo=dt.timezone.utc) + assert _to_text(ts) == "2026-05-28T09:30:00+00:00" + + def test_naive_datetime_renders_iso_without_offset(self) -> None: + ts = dt.datetime(2026, 5, 28, 9, 30, 0) + assert _to_text(ts) == "2026-05-28T09:30:00" + + def test_date_renders_iso(self) -> None: + assert _to_text(dt.date(2026, 5, 28)) == "2026-05-28" + + def test_bool_true_renders_lowercase(self) -> None: + """Critical: bool must match BEFORE int (bool is an int subclass).""" + assert _to_text(True) == "true" + + def test_bool_false_renders_lowercase(self) -> None: + assert _to_text(False) == "false" + + @pytest.mark.parametrize( + ("value", "expected"), + [ + # Decimals within the "normal" exponent range round-trip + # as plain notation, byte-for-byte equal to what the user + # originally wrote. This is the common case and the one + # service code actually exercises. + (Decimal("0"), "0"), + (Decimal("123"), "123"), + (Decimal("123.456"), "123.456"), + (Decimal("-0.5"), "-0.5"), + # Very small / very large magnitudes flip to scientific + # notation because that's what Python's :meth:`Decimal.__str__` + # produces once the exponent leaves the normal range. We + # deliberately accept this output (see method docstring) — + # downstream Postgres and CSV consumers parse ``1E-7`` and + # ``1E+21`` exactly like the plain forms. Switching to + # ``format(d, "f")`` would force plain notation but bloat + # output for legitimately tiny / huge values and surprise + # readers used to Decimal's canonical repr. + (Decimal("0.0000001"), "1E-7"), + (Decimal("1E+21"), "1E+21"), + ], + ) + def test_decimal_uses_decimal_str_repr(self, value: Decimal, expected: str) -> None: + """Decimal values pass through ``str(value)`` verbatim. + + The function does NOT force plain notation; it relies on + :meth:`Decimal.__str__`, which switches to scientific notation + once the exponent is outside Decimal's "normal" range + (configurable via :class:`decimal.Context`, default ±28). The + test name documents this contract — an earlier version of this + test claimed "no scientific notation" while asserting on + ``"1E-7"``, which silently passed for the wrong reason. + + If a future caller cannot tolerate scientific notation (e.g. + a strict CSV importer that doesn't recognise the ``E`` form), + the fix is to apply ``format(d, "f")`` at the formatting layer + — NOT to swap :func:`_to_text`'s implementation, which would + bloat output for legitimately tiny / huge values. + """ + assert _to_text(value) == expected + + @pytest.mark.parametrize( + ("value", "expected"), + [ + (0, "0"), + (42, "42"), + (-7, "-7"), + (1.5, "1.5"), + (0.0, "0.0"), + ], + ) + def test_numeric_renders_as_str(self, value: float, expected: str) -> None: + assert _to_text(value) == expected + + @pytest.mark.parametrize( + ("value", "expected"), + [ + (b"\x00\x01\x02", "000102"), + (bytearray(b"hello"), "68656c6c6f"), + (memoryview(b"\xde\xad\xbe\xef"), "deadbeef"), + ], + ) + def test_bytes_like_renders_as_hex(self, value: object, expected: str) -> None: + assert _to_text(value) == expected + + def test_plain_string_passes_through(self) -> None: + assert _to_text("hello") == "hello" + + def test_arbitrary_object_falls_back_to_str(self) -> None: + class Custom: + def __str__(self) -> str: + return "custom-repr" + + assert _to_text(Custom()) == "custom-repr" + + +class TestPgRenderValue: + """Postgres-flavoured literal renderer for ``upsert``. + + Delegates non-:class:`RawSql` values to :func:`SqlExecutor._render_value`; + the Postgres-specific responsibility is translating Spark-style + ``current_timestamp()`` / ``now()`` to bare ``CURRENT_TIMESTAMP`` + that Postgres actually parses, while letting other ``RawSql`` + payloads through verbatim so callers can still inject Postgres- + specific helpers. + """ + + @pytest.mark.parametrize("payload", ["current_timestamp()", "now()", "CURRENT_TIMESTAMP()", "Now()"]) + def test_spark_timestamp_idioms_become_postgres_current_timestamp(self, payload: str) -> None: + assert _pg_render_value(RawSql(payload)) == "CURRENT_TIMESTAMP" + + @pytest.mark.parametrize("payload", [" current_timestamp() ", "\tnow()\n"]) + def test_spark_timestamp_idioms_tolerate_surrounding_whitespace(self, payload: str) -> None: + assert _pg_render_value(RawSql(payload)) == "CURRENT_TIMESTAMP" + + @pytest.mark.parametrize( + "payload", + [ + "CURRENT_TIMESTAMP", # already-Postgres + "uuid_generate_v4()", # postgres-specific function + "(SELECT max(version) FROM other_table)", # arbitrary subexpression + "EXCLUDED.updated_at", # ON CONFLICT helper + ], + ) + def test_other_raw_sql_passes_through_verbatim(self, payload: str) -> None: + assert _pg_render_value(RawSql(payload)) == payload + + def test_bool_true_renders_as_TRUE(self) -> None: + assert _pg_render_value(True) == "TRUE" + + def test_bool_false_renders_as_FALSE(self) -> None: + assert _pg_render_value(False) == "FALSE" + + def test_none_renders_as_NULL(self) -> None: + assert _pg_render_value(None) == "NULL" + + @pytest.mark.parametrize( + ("value", "expected"), + [ + (42, "42"), + (-7, "-7"), + (1.5, "1.5"), + ], + ) + def test_numeric_renders_as_literal(self, value: float, expected: str) -> None: + assert _pg_render_value(value) == expected + + def test_string_is_ansi_escaped(self) -> None: + assert _pg_render_value("it's fine") == "'it''s fine'" + + +# =========================================================================== +# Identifier quoting (review item #8 already tightened the Delta side; +# here we cover the Postgres side directly). +# =========================================================================== + + +class TestPgQuoting: + """Direct tests of :meth:`PgExecutor.q`.""" + + @pytest.mark.parametrize( + ("identifier", "expected"), + [ + ("foo", '"foo"'), + # Reserved-word columns the OLTP schema actually contains + ("user", '"user"'), + ("order", '"order"'), + ("check", '"check"'), + # Internal ``"`` must be doubled per ANSI SQL. + ('weird"name', '"weird""name"'), + # Pre-doubled (no over-escaping) + ('a""b', '"a""""b"'), + # Unicode survives (Postgres permits it in quoted identifiers) + ("café", '"café"'), + # Empty string — corner case for any escaping helper + ("", '""'), + ], + ) + def test_quotes_identifier(self, identifier: str, expected: str) -> None: + executor = _make_pg_executor() + assert executor.q(identifier) == expected + + +# =========================================================================== +# upsert SQL shape +# =========================================================================== + + +class TestUpsertSqlShape: + """Verifies the INSERT ... ON CONFLICT ... SQL shape. + + The reviewer's specific request: cover both the ``DO UPDATE`` + branch (when ``value_cols`` is populated) and the ``DO NOTHING`` + branch (when ``value_cols`` is empty — i.e. "ensure-this-row- + exists" semantics). All column names must be quoted via + :meth:`PgExecutor.q` so reserved words survive. + """ + + def _capture(self, executor: PgExecutor) -> list[str]: + captured: list[str] = [] + # Monkey-patch the ``execute`` method so we read the rendered SQL + # without needing a real pool. Bound to ``self`` via ``__get__`` + # so the method swap survives the call chain. + executor.execute = lambda sql, **_: captured.append(sql) # type: ignore[method-assign] + return captured + + def test_do_update_branch_when_value_cols_populated(self) -> None: + executor = _make_pg_executor() + captured = self._capture(executor) + + executor.upsert( + table='"dq"."dq_app_settings"', + key_cols={"setting_key": "max_concurrent_runs"}, + value_cols={"setting_value": "10", "updated_by": "admin"}, + ) + + assert len(captured) == 1 + sql = captured[0] + assert 'INSERT INTO "dq"."dq_app_settings"' in sql + assert 'ON CONFLICT ("setting_key") DO UPDATE SET' in sql + # Both value cols must appear in the SET clause, each referencing EXCLUDED. + assert '"setting_value" = EXCLUDED."setting_value"' in sql + assert '"updated_by" = EXCLUDED."updated_by"' in sql + # Column names in the INSERT list must be quoted (reserved-word safety). + assert '("setting_key", "setting_value", "updated_by")' in sql + # String literals must be ANSI-escaped. + assert "'max_concurrent_runs'" in sql + + def test_do_nothing_branch_when_value_cols_empty(self) -> None: + """Empty ``value_cols`` is the reviewer-flagged 'ensure-row-exists' shape.""" + executor = _make_pg_executor() + captured = self._capture(executor) + + executor.upsert( + table='"dq"."dq_role_mappings"', + key_cols={"role": "admin", "group_name": "platform-admins"}, + value_cols={}, + ) + + assert len(captured) == 1 + sql = captured[0] + assert 'INSERT INTO "dq"."dq_role_mappings"' in sql + # Both keys appear in the ON CONFLICT target. + assert 'ON CONFLICT ("role", "group_name") DO NOTHING' in sql + # No DO UPDATE / SET clause should leak in. + assert "DO UPDATE" not in sql + assert " SET " not in sql + + def test_raises_when_no_key_cols(self) -> None: + executor = _make_pg_executor() + self._capture(executor) + with pytest.raises(ValueError, match="upsert requires at least one key column"): + executor.upsert(table="t", key_cols={}, value_cols={"x": 1}) + + def test_raw_sql_in_value_cols_translates_current_timestamp(self) -> None: + """RawSql values flow through _pg_render_value, so Spark idioms are translated.""" + executor = _make_pg_executor() + captured = self._capture(executor) + + executor.upsert( + table='"dq"."dq_app_settings"', + key_cols={"setting_key": "feature_flag"}, + value_cols={"setting_value": "on", "updated_at": RawSql("current_timestamp()")}, + ) + + sql = captured[0] + # The Spark idiom must NOT appear; the Postgres bare keyword should. + assert "current_timestamp()" not in sql + assert "CURRENT_TIMESTAMP" in sql + + def test_reserved_word_key_column_is_quoted(self) -> None: + executor = _make_pg_executor() + captured = self._capture(executor) + executor.upsert( + table="t", + key_cols={"check": "1"}, # reserved-word column + value_cols={"order": "2"}, # another reserved word + ) + sql = captured[0] + assert '("check", "order")' in sql + assert 'ON CONFLICT ("check")' in sql + assert '"order" = EXCLUDED."order"' in sql + + +# =========================================================================== +# upsert_with_audit — Protocol method that absorbs the dialect branch +# =========================================================================== + + +class TestUpsertWithAuditSqlShape: + """Verifies the audit-aware upsert SQL shape on Postgres. + + The Protocol contract :meth:`OltpExecutorProtocol.upsert_with_audit` + is responsible for two behaviours service code used to hand-roll + via dialect branches: + + 1. ``preserve_created=True`` excludes ``created_*`` columns from + the DO UPDATE SET clause so the original creator/timestamp + survive after the first INSERT. + 2. ``increment_on_update`` rewrites the named column's DO UPDATE + expression to ``col = col + 1`` (bare reference, which Postgres + resolves to the existing row inside ON CONFLICT DO UPDATE). + + These tests pin the rendered SQL so a regression to the old + "include every value_cols in DO UPDATE" shape would fail loudly. + """ + + def _capture(self, executor: PgExecutor) -> list[str]: + captured: list[str] = [] + executor.execute = lambda sql, **_: captured.append(sql) # type: ignore[method-assign] + return captured + + def test_preserve_created_excludes_created_cols_from_update_set(self) -> None: + executor = _make_pg_executor() + captured = self._capture(executor) + + executor.upsert_with_audit( + table='"dq"."dq_role_mappings"', + key_cols={"role": "admin", "group_name": "platform-admins"}, + value_cols={ + "created_by": "alice@x", + "created_at": RawSql("now()"), + "updated_by": "alice@x", + "updated_at": RawSql("now()"), + }, + preserve_created=True, + ) + + sql = captured[0] + # INSERT carries ALL six cols (2 keys + 4 audit). + assert '("role", "group_name", "created_by", "created_at", "updated_by", "updated_at")' in sql + # DO UPDATE SET names ONLY the non-created cols. + assert "DO UPDATE SET" in sql + assert '"updated_by" = ' in sql + assert '"updated_at" = ' in sql + # created_* MUST be absent from the SET clause — the regression + # we're guarding against is the executor clobbering the original + # creator/timestamp on every upsert. + set_clause = sql.split("DO UPDATE SET", 1)[1] + assert '"created_by"' not in set_clause + assert '"created_at"' not in set_clause + + def test_increment_on_update_uses_bare_column_self_reference(self) -> None: + """On Postgres, ``col + 1`` inside DO UPDATE resolves to the existing row.""" + executor = _make_pg_executor() + captured = self._capture(executor) + + executor.upsert_with_audit( + table='"dq"."dq_schedule_config"', + key_cols={"schedule_name": "main"}, + value_cols={ + "config_json": '{"k":"v"}', + "version": 1, # initial INSERT value + "updated_by": "alice@x", + "updated_at": RawSql("now()"), + }, + preserve_created=False, + increment_on_update="version", + ) + + sql = captured[0] + # INSERT VALUES include the literal initial 1 (not the increment expr). + assert "VALUES ('main', '{\"k\":\"v\"}', 1, 'alice@x', CURRENT_TIMESTAMP)" in sql + # DO UPDATE rewrites version to the bare-column self-reference + # — no table prefix, no EXCLUDED reference. + assert '"version" = "version" + 1' in sql + # And specifically NOT the EXCLUDED form (which would set it to + # the proposed 1 every time, defeating the increment). + assert '"version" = EXCLUDED."version"' not in sql + + def test_combined_preserve_created_and_increment_for_schedule_config_shape(self) -> None: + """The exact call ``ScheduleConfigService.save`` makes — locked + down end-to-end so a service refactor that drops either kwarg + fails here.""" + executor = _make_pg_executor() + captured = self._capture(executor) + + executor.upsert_with_audit( + table='"dq"."dq_schedule_config"', + key_cols={"schedule_name": "main"}, + value_cols={ + "config_json": '{}', + "version": 1, + "created_by": "alice@x", + "created_at": RawSql("now()"), + "updated_by": "alice@x", + "updated_at": RawSql("now()"), + }, + preserve_created=True, + increment_on_update="version", + ) + + sql = captured[0] + set_clause = sql.split("DO UPDATE SET", 1)[1] + # SET clause has updated_*, config_json, and the version self- + # reference — but NOT created_*. + assert '"created_by"' not in set_clause + assert '"created_at"' not in set_clause + assert '"config_json"' in set_clause + assert '"updated_by"' in set_clause + assert '"updated_at"' in set_clause + assert '"version" = "version" + 1' in set_clause + + def test_all_value_cols_are_created_renders_do_nothing(self) -> None: + """Edge case: insert-only audit row with no updatable cols.""" + executor = _make_pg_executor() + captured = self._capture(executor) + + executor.upsert_with_audit( + table='"dq"."ledger"', + key_cols={"id": "1"}, + value_cols={"created_by": "system", "created_at": RawSql("now()")}, + preserve_created=True, + ) + + sql = captured[0] + # Nothing to update → DO NOTHING (not a malformed empty SET). + assert "DO NOTHING" in sql + assert "DO UPDATE" not in sql + assert " SET " not in sql + + def test_raises_when_no_key_cols(self) -> None: + executor = _make_pg_executor() + self._capture(executor) + with pytest.raises(ValueError, match="upsert_with_audit requires at least one key column"): + executor.upsert_with_audit(table="t", key_cols={}, value_cols={"x": 1}) + + def test_raises_when_increment_col_missing_from_value_cols(self) -> None: + """Guard against silently dropping a version bump.""" + executor = _make_pg_executor() + self._capture(executor) + with pytest.raises(ValueError, match="increment_on_update='version' must be present"): + executor.upsert_with_audit( + table="t", + key_cols={"k": "1"}, + value_cols={"updated_by": "u"}, # no version! + increment_on_update="version", + ) + + def test_reserved_word_increment_column_is_quoted(self) -> None: + executor = _make_pg_executor() + captured = self._capture(executor) + executor.upsert_with_audit( + table="t", + key_cols={"k": "1"}, + value_cols={"order": 1}, # reserved word + preserve_created=False, + increment_on_update="order", + ) + sql = captured[0] + # Increment must use quoted identifier on BOTH sides of the +. + assert '"order" = "order" + 1' in sql + + +# =========================================================================== +# select_json_text / interval_days_expr — Protocol render helpers +# =========================================================================== + + +class TestProtocolRenderHelpers: + """The two pure-string helpers that absorb dialect rendering. + + Both are trivial but the *exact* string shape matters — a typo + becomes a SQL syntax error against a live Postgres at the worst + possible moment (a migration window). Lock them down here so any + accidental rewrite trips the test, not the customer's Lakebase. + """ + + def test_select_json_text_returns_column_verbatim(self) -> None: + # Postgres JSONB → text coercion happens in _to_text on the way + # out, so the projection is just the bare column. No to_json(), + # no CAST — a regression to either would re-introduce the + # dialect branch the Protocol absorbs. + executor = _make_pg_executor() + assert executor.select_json_text('"check"') == '"check"' + assert executor.select_json_text("config_json") == "config_json" + + def test_interval_days_expr_uses_quoted_lowercase_plural(self) -> None: + # Postgres syntax: INTERVAL ' days' (single-quoted literal, + # lowercase plural). The Delta form (INTERVAL N DAY) would + # produce ``ERROR: syntax error at or near "DAY"`` against a + # real Postgres. + executor = _make_pg_executor() + assert executor.interval_days_expr(90) == "INTERVAL '90 days'" + assert executor.interval_days_expr(1) == "INTERVAL '1 days'" + # int() coercion is part of the contract — defends against a + # ``float`` slipping in from a config bug. + assert executor.interval_days_expr(7.0) == "INTERVAL '7 days'" # type: ignore[arg-type] + + +# =========================================================================== +# connection() context manager +# =========================================================================== + + +class TestConnectionContextManager: + """The wrapper around ``self._pool.connection()``. + + The contract documented on :meth:`PgExecutor.connection`: + + - The caller is responsible for ``conn.commit()``. + - If the ``with`` block exits without committing (exception or + otherwise) psycopg-pool rolls back and returns the connection. + + These tests verify our wrapper is a *true* pass-through: it + neither commits nor rolls back implicitly, and it propagates + exceptions raised inside the ``with`` block. + """ + + def test_yields_the_pool_connection(self) -> None: + executor = _make_pg_executor() + with executor.connection() as conn: + # The yielded object is the same conn the pool produced. + assert conn is _conn_of(executor) + + def test_does_not_implicitly_commit_or_rollback(self) -> None: + """Caller-managed transaction semantics — we don't sneak commits in.""" + executor = _make_pg_executor() + with executor.connection() as conn: + pass + # Our wrapper never calls commit() or rollback() on the conn — + # the caller (or psycopg-pool's __exit__) owns that decision. + conn.commit.assert_not_called() # type: ignore[attr-defined] + conn.rollback.assert_not_called() # type: ignore[attr-defined] + + def test_exception_inside_block_propagates(self) -> None: + """A user error inside ``with executor.connection()`` must not be swallowed.""" + executor = _make_pg_executor() + with pytest.raises(RuntimeError, match="boom"): + with executor.connection(): + raise RuntimeError("boom") + + +# =========================================================================== +# _token_refresh_loop — happy / sad / permanent-failure paths +# =========================================================================== + + +class TestTokenRefreshLoop: + """Background-thread refresh logic, tested without spawning the thread. + + We invoke ``_token_refresh_loop`` directly on a half-built executor + so the test runs synchronously. ``_stop`` is a real + :class:`threading.Event` so we can wake the loop deterministically. + All ``_escalate_refresh_failure`` calls are patched out — the + real implementation hard-exits the process via ``os._exit``, which + we cover separately in :class:`TestRefreshFailureEscalation`. + """ + + def _stop_after(self, executor: PgExecutor, n: int) -> None: + """Patch ``_stop.wait`` so the loop runs exactly ``n`` iterations. + + Each ``wait`` call increments a counter; on the ``n``-th call + we set the stop event so the next ``while not self._stop.is_set()`` + check exits the loop. ``wait`` always returns ``False`` so the + loop's "stop fired during sleep" early-return path doesn't fire + — that's covered separately. + """ + counter = {"n": 0} + + def fake_wait(timeout: float | None = None) -> bool: # noqa: ARG001 + counter["n"] += 1 + if counter["n"] >= n: + executor._stop.set() + return False + + executor._stop.wait = fake_wait # type: ignore[method-assign] + + def test_happy_path_updates_token_kwargs_and_metric(self) -> None: + executor = _make_pg_executor(initial_token="old-token") + self._stop_after(executor, 1) + + with patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + return_value="fresh-token", + ) as gen: + executor._token_refresh_loop() + + gen.assert_called_with(executor._ws, executor._instance_name) + assert executor._token_holder.token == "fresh-token" + assert executor._connect_kwargs["password"] == "fresh-token" + # Success path updates the observability surface that the + # health endpoint relies on. + assert executor.consecutive_refresh_failures == 0 + assert executor.last_successful_refresh_at is not None + + def test_stop_event_terminates_loop_without_generating_token(self) -> None: + """``wait`` returning True means the stop event fired during the sleep.""" + executor = _make_pg_executor() + executor._stop.set() # already set before the loop runs + + with patch("databricks_labs_dqx_app.backend.pg_executor._generate_token") as gen: + executor._token_refresh_loop() + + gen.assert_not_called() + + def test_failure_then_success_increments_then_resets_counter(self, caplog) -> None: + """Transient failure → counter increments; recovery → counter resets to 0.""" + executor = _make_pg_executor(initial_token="old-token", token_refresh_max_failures=10) + self._stop_after(executor, 2) + + with patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + side_effect=[RuntimeError("transient SDK error"), "recovered-token"], + ) as gen: + with caplog.at_level(logging.WARNING, logger="databricks_labs_dqx_app.backend.pg_executor"): + executor._token_refresh_loop() + + assert gen.call_count == 2 + warnings = [r for r in caplog.records if r.levelno == logging.WARNING] + assert any("Failed to refresh Lakebase token" in r.message for r in warnings) + # The warning carries the attempt counter so operators can + # tell "first failure" from "approaching escalation". + assert any("1/10" in r.message for r in warnings) + # Recovery clears the failure counter and updates the token. + assert executor._token_holder.token == "recovered-token" + assert executor._connect_kwargs["password"] == "recovered-token" + assert executor.consecutive_refresh_failures == 0 + + def test_permanent_failure_below_threshold_does_not_escalate(self) -> None: + """Sustained failures *under* ``max_failures`` keep spinning, don't escalate.""" + # Raise the threshold so 3 failures don't trip escalation — + # we want to verify the loop keeps trying without dying or + # exiting the process. + executor = _make_pg_executor(initial_token="old-token", token_refresh_max_failures=99) + self._stop_after(executor, 3) + + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + side_effect=RuntimeError("Lakebase is down"), + ) as gen, + patch.object(type(executor), "_escalate_refresh_failure") as escalate, + ): + executor._token_refresh_loop() + + assert gen.call_count >= 3 + # Existing connections keep their original credential until + # ``max_lifetime`` recycles them — we never mutated the + # token holder. + assert executor._token_holder.token == "old-token" + assert executor._connect_kwargs["password"] == "old-token" + # Failure counter accumulated across all attempts. + assert executor.consecutive_refresh_failures >= 3 + # And — critically — we did NOT escalate; the escalation + # path is reserved for ``>= max_failures``. + escalate.assert_not_called() + + +# =========================================================================== +# _next_wait_seconds — back-off cadence + jitter bounds +# =========================================================================== + + +class TestNextWaitSeconds: + """The back-off chooser is the load-bearing piece of the redesign. + + Healthy → wait the full scheduled interval (don't burn SDK quota). + Failing → wait the short retry interval, jittered, so we recover + inside the pool's ``max_lifetime`` window instead of letting it + drain silently. + """ + + def test_healthy_returns_scheduled_interval(self) -> None: + """``consecutive_refresh_failures == 0`` ⇒ full scheduled wait.""" + executor = _make_pg_executor( + token_refresh_seconds=3000, + token_refresh_retry_seconds=10, + ) + executor._consecutive_refresh_failures = 0 + assert executor._next_wait_seconds() == 3000.0 + + def test_failing_returns_retry_interval_not_scheduled(self) -> None: + """One failure flips to the short retry path — this is the whole point. + + The original loop slept 60s after a failure and *then* waited + the full scheduled interval (50 min by default) before retrying, + which silently drained the pool. Verify the new loop instead + picks the retry budget, not the scheduled one. + """ + executor = _make_pg_executor( + token_refresh_seconds=3000, + token_refresh_retry_seconds=10, + token_refresh_retry_jitter=0.0, # no jitter for a deterministic check + ) + executor._consecutive_refresh_failures = 1 + wait = executor._next_wait_seconds() + assert wait == 10.0, ( + f"After a failure the loop must pick the retry interval (10s), " + f"not the scheduled interval (3000s); got {wait}s. " + f"That regression is what silently drains the pool on sustained failure." + ) + + def test_jitter_stays_within_documented_bounds(self) -> None: + """Jitter must be uniform within ``base ± base*jitter``. + + Operators tune ``retry_jitter`` to spread thunder-herd risk — + if the jitter ever exceeded the documented bound, multiple + workers' back-offs could synchronise (worse: a bug could + produce wait=0, busy-looping the SDK). Sample a chunk of + iterations and check the empirical bounds. + """ + executor = _make_pg_executor( + token_refresh_retry_seconds=10, + token_refresh_retry_jitter=0.3, + ) + executor._consecutive_refresh_failures = 1 + samples = [executor._next_wait_seconds() for _ in range(200)] + lower = 10 - (10 * 0.3) # 7.0 + upper = 10 + (10 * 0.3) # 13.0 + assert all(lower <= s <= upper for s in samples), ( + f"Jitter escaped ±30% bound. lower={min(samples)}, upper={max(samples)}, " + f"expected within [{lower}, {upper}]." + ) + # And we actually observed some spread — not a single fixed + # value (which would mean ``random.uniform`` was patched out + # or the jitter math collapsed to a constant). + assert len(set(samples)) > 1, ( + "Jitter produced a single repeated value across 200 samples; " + "the randomisation is missing and workers will thunder-herd." + ) + + def test_jitter_zero_returns_exact_retry_seconds(self) -> None: + """``jitter=0`` ⇒ deterministic, returns ``retry_seconds`` exactly. + + Documented escape hatch for tests / deterministic env where + a flapping back-off would make logs hard to read. + """ + executor = _make_pg_executor( + token_refresh_retry_seconds=7, + token_refresh_retry_jitter=0.0, + ) + executor._consecutive_refresh_failures = 5 + for _ in range(10): + assert executor._next_wait_seconds() == 7.0 + + +# =========================================================================== +# _escalate_refresh_failure — process exit on sustained failure +# =========================================================================== + + +class TestRefreshFailureEscalation: + """After ``max_failures`` consecutive failures the loop hands off to + the supervisor by exiting the process. + + The escalation path is the user-facing fix for the silent-pool-drain + risk: rather than continuing to retry while ``max_lifetime`` quietly + expires every pooled connection, crash loud and let the supervisor + restart us with a fresh token. + """ + + def test_loop_calls_escalate_after_max_failures(self, caplog) -> None: + """Reaching the threshold triggers escalation, not another retry.""" + executor = _make_pg_executor( + initial_token="old-token", + token_refresh_max_failures=3, + ) + # Force the loop to keep iterating; the escalation path is + # what stops it, not the wait counter. + executor._stop.wait = lambda timeout=None: False # type: ignore[method-assign,assignment] + + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + side_effect=RuntimeError("Lakebase is down"), + ), + patch.object(type(executor), "_escalate_refresh_failure") as escalate, + ): + with caplog.at_level(logging.WARNING, logger="databricks_labs_dqx_app.backend.pg_executor"): + executor._token_refresh_loop() + + # Escalation fired exactly once — on the 3rd consecutive + # failure, the loop calls it and then ``return``s. + escalate.assert_called_once() + assert executor.consecutive_refresh_failures == 3 + + def test_escalate_logs_critical_with_context_then_exits(self, caplog) -> None: + """Escalation logs CRITICAL with the failure count + last-success + + instance name BEFORE exiting, so operators have something to + page on / search for in log aggregators. + + ``os._exit`` is patched out so the test runs synchronously. + """ + executor = _make_pg_executor(instance_name="lakebase-prod-us-east-1") + executor._consecutive_refresh_failures = 12 + before = executor._last_successful_refresh_at + + with patch("databricks_labs_dqx_app.backend.pg_executor.os._exit") as exit_call: + with caplog.at_level(logging.CRITICAL, logger="databricks_labs_dqx_app.backend.pg_executor"): + executor._escalate_refresh_failure() + + exit_call.assert_called_once_with(1) + criticals = [r for r in caplog.records if r.levelno == logging.CRITICAL] + assert criticals, "Escalation must log CRITICAL for ops paging." + msg = criticals[0].getMessage() + assert "12" in msg, "Log must include the failure count for context." + assert "lakebase-prod-us-east-1" in msg, ( + "Log must include the instance name so multi-instance deployments " "know which Lakebase failed." + ) + if before is not None: + assert before.isoformat() in msg, ( + "Log must include the last-successful-refresh timestamp " + "so operators know how long the failure window is." + ) + + def test_loop_invokes_os_exit_after_exactly_max_failures(self) -> None: + """End-to-end: loop → escalate → ``os._exit(1)`` with NOTHING + between ``os._exit`` and the loop patched out. + + The earlier ``test_loop_calls_escalate_after_max_failures`` + patches ``_escalate_refresh_failure`` itself, so it only + verifies the loop *reaches* the escalation hand-off. That + leaves the load-bearing claim — "the supervisor actually sees + a non-zero exit after max_failures consecutive failures" — + untested end-to-end. If a future refactor changed + ``os._exit(1)`` to ``sys.exit(1)`` (which a daemon thread + silently swallows), the existing test would still pass and + the pool would silently drain in production exactly like the + original 60s-sleep bug. + + Here we drive the loop with ``_generate_token`` raising every + time, leave ``_escalate_refresh_failure`` UN-patched, and + verify the chain reaches ``os._exit(1)`` after exactly + ``max_failures`` token-generation attempts — not earlier + (escalates one short of the budget) and not later + (escalates one past the budget). + """ + executor = _make_pg_executor(token_refresh_max_failures=4) + # Never auto-stop via the sleep counter — the loop must stop + # on its own via the ``return`` after ``_escalate_refresh_failure``. + executor._stop.wait = lambda timeout=None: False # type: ignore[method-assign,assignment] + + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + side_effect=RuntimeError("Lakebase is down"), + ) as gen, + patch("databricks_labs_dqx_app.backend.pg_executor.os._exit") as exit_call, + ): + executor._token_refresh_loop() + + # Exactly ``max_failures`` token-generation attempts before + # escalation fired. Anything less means escalation fires too + # early (operators lose retries the policy promised them); + # anything more means escalation fires too late (the pool's + # ``max_lifetime`` may start expiring connections without a + # valid replacement token, which is the original bug). + assert gen.call_count == 4, ( + f"Expected escalation after exactly 4 consecutive failures, " + f"got {gen.call_count} ``_generate_token`` calls. The " + f"threshold check in ``_token_refresh_loop`` is off-by-one " + f"or the counter increment landed on the wrong branch." + ) + exit_call.assert_called_once_with(1) + assert executor.consecutive_refresh_failures == 4 + + def test_success_grants_a_fresh_threshold_window_before_escalation(self) -> None: + """A single successful refresh resets the ENTIRE failure budget, + not just decrements the counter by one. + + Drive ``max_failures=3`` with the sequence:: + + fail, fail, (counter = 2; one short of escalate) + success, (counter must reset to 0) + fail, fail, (counter = 2; one short of escalate) + fail (counter = 3 → escalate) + + Escalation must fire on the **6th** ``_generate_token`` call + (the 3rd post-recovery failure), not the 4th (which would mean + the success only suppressed the immediate escalate but the + threshold was already poisoned), and not the 5th (which would + mean the success only decremented by one). + + This catches subtle refactors of the reset path: + - ``self._consecutive_refresh_failures -= 1`` instead of + ``= 0`` → escalation would fire on the 5th call. + - Forgetting the reset entirely (e.g. only updating + ``last_successful_refresh_at``) → escalation would fire + on the 4th call. + Without this test, either bug ships green. + """ + executor = _make_pg_executor(token_refresh_max_failures=3) + executor._stop.wait = lambda timeout=None: False # type: ignore[method-assign,assignment] + + # 6 attempts: 2 fail, 1 success, 3 fail (escalate). The list + # is sized to the *expected* count so an early escalation + # leaves entries unused (gen.call_count < 6, test fails) and + # a late one would attempt a 7th call (StopIteration → still + # caught by the loop's ``except Exception`` but escalation + # would by then have fired off the wrong counter value). + token_seq: list[object] = [ + RuntimeError("transient #1"), + RuntimeError("transient #2"), + "recovered-token", + RuntimeError("transient #3"), + RuntimeError("transient #4"), + RuntimeError("transient #5"), + ] + + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + side_effect=token_seq, + ) as gen, + patch("databricks_labs_dqx_app.backend.pg_executor.os._exit") as exit_call, + ): + executor._token_refresh_loop() + + assert gen.call_count == 6, ( + f"Expected 6 ``_generate_token`` calls (2 fail + 1 success + " + f"3 fail to re-reach the threshold of 3), got {gen.call_count}. " + f"A count of 4 means success only suppressed the immediate " + f"escalate (no counter reset). A count of 5 means success " + f"only decremented the counter by one instead of resetting it." + ) + exit_call.assert_called_once_with(1) + # Counter at escalation time is exactly the threshold — the + # success path fully reset it, the post-recovery failures + # rebuilt it from zero. + assert executor.consecutive_refresh_failures == 3 + # And the success in the middle DID update the token + the + # observability surface — proving the reset is genuine (the + # success branch ran in full), not a side-effect of some + # other branch suppressing escalation. + assert executor._token_holder.token == "recovered-token" + assert executor._connect_kwargs["password"] == "recovered-token" + + +# =========================================================================== +# Observability surfaces — last_successful_refresh_at + counter +# =========================================================================== + + +class TestRefreshObservability: + """``last_successful_refresh_at`` + ``consecutive_refresh_failures`` + are the metrics health endpoints scrape to catch the silent-drain + failure before it cascades. + """ + + def test_initial_state_after_init_is_healthy(self) -> None: + """A freshly-built executor reads as 'just refreshed, no failures'. + + The bootstrap token mint IS a successful refresh from the + pool's perspective, so the metric must start in the healthy + state — not as ``None`` / unknown, which would falsely page + every operator at t=0. + """ + # We can't run the real __init__ in unit tests (no real pool), + # so the fixture sets these to ``None`` / 0; verify the *real* + # __init__ sets the bootstrap timestamp by patching out the + # external surfaces and constructing for real. + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + return_value="bootstrap-token", + ), + patch("databricks_labs_dqx_app.backend.pg_executor.ConnectionPool") as Pool, + patch("databricks_labs_dqx_app.backend.pg_executor.threading.Thread"), + ): + Pool.check_connection = MagicMock() + executor = PgExecutor( + ws=MagicMock(), + instance_name="instance", + database="dqx", + schema="public", + username="sp", + host="host", + ) + + assert executor.consecutive_refresh_failures == 0 + assert executor.last_successful_refresh_at is not None + + def test_success_updates_timestamp_monotonic(self) -> None: + """Each successful refresh advances ``last_successful_refresh_at``.""" + executor = _make_pg_executor() + before = dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc) + executor._last_successful_refresh_at = before + executor._stop.wait = lambda timeout=None: (executor._stop.set() or False) # type: ignore[method-assign,assignment,func-returns-value] + + with patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + return_value="fresh", + ): + executor._token_refresh_loop() + + assert executor.last_successful_refresh_at is not None + assert executor.last_successful_refresh_at > before, ( + "Successful refresh must advance the timestamp so a stalled " + "loop is detectable via ``now() - last_successful_refresh_at``." + ) + + def test_failure_does_not_update_timestamp(self) -> None: + """Last-success only moves on success — never on failure. + + Otherwise a refresh attempt that raised would silently reset + the metric and hide the failure window from operators. + """ + executor = _make_pg_executor() + before = dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc) + executor._last_successful_refresh_at = before + executor._stop.wait = lambda timeout=None: (executor._stop.set() or False) # type: ignore[method-assign,assignment,func-returns-value] + + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + side_effect=RuntimeError("boom"), + ), + patch.object(type(executor), "_escalate_refresh_failure"), + ): + executor._token_refresh_loop() + + assert executor.last_successful_refresh_at == before + assert executor.consecutive_refresh_failures >= 1 + + +# =========================================================================== +# Constructor clamps — defend against mis-configured env vars +# =========================================================================== + + +class TestRefreshTuningClamps: + """Mis-configured env vars can't degenerate the loop into a tight + busy-loop or an infinite failure window.""" + + def test_retry_seconds_floors_at_one(self) -> None: + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + return_value="t", + ), + patch("databricks_labs_dqx_app.backend.pg_executor.ConnectionPool"), + patch("databricks_labs_dqx_app.backend.pg_executor.threading.Thread"), + ): + executor = PgExecutor( + ws=MagicMock(), + instance_name="i", + database="d", + schema="s", + username="u", + host="h", + token_refresh_retry_seconds=0, + ) + assert executor._token_refresh_retry_seconds == 1 + + def test_jitter_clamps_to_unit_interval(self) -> None: + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + return_value="t", + ), + patch("databricks_labs_dqx_app.backend.pg_executor.ConnectionPool"), + patch("databricks_labs_dqx_app.backend.pg_executor.threading.Thread"), + ): + high = PgExecutor( + ws=MagicMock(), + instance_name="i", + database="d", + schema="s", + username="u", + host="h", + token_refresh_retry_jitter=5.0, + ) + low = PgExecutor( + ws=MagicMock(), + instance_name="i", + database="d", + schema="s", + username="u", + host="h", + token_refresh_retry_jitter=-0.5, + ) + assert high._token_refresh_retry_jitter == 1.0 + assert low._token_refresh_retry_jitter == 0.0 + + def test_max_failures_floors_at_one(self) -> None: + """``max_failures=0`` would never escalate, the opposite of safe.""" + with ( + patch( + "databricks_labs_dqx_app.backend.pg_executor._generate_token", + return_value="t", + ), + patch("databricks_labs_dqx_app.backend.pg_executor.ConnectionPool"), + patch("databricks_labs_dqx_app.backend.pg_executor.threading.Thread"), + ): + executor = PgExecutor( + ws=MagicMock(), + instance_name="i", + database="d", + schema="s", + username="u", + host="h", + token_refresh_max_failures=0, + ) + assert executor._token_refresh_max_failures == 1 + + +# =========================================================================== +# build_pg_executor — error paths (the only ones exercisable without a +# real Postgres). The happy-path call sequence is exercised in +# integration tests. +# =========================================================================== + + +class TestBuildPgExecutor: + """Cover the two early-validation RuntimeErrors at lines 436-444.""" + + def test_raises_when_instance_has_no_read_write_dns(self) -> None: + ws = MagicMock(name="WorkspaceClient") + # get_database_instance returns an object whose .read_write_dns is empty. + ws.database.get_database_instance.return_value = MagicMock(read_write_dns=None) + + with pytest.raises(RuntimeError, match="has no read_write_dns"): + build_pg_executor( + ws, + instance_name="missing-dns-instance", + database="dqx", + schema="dqx_studio", + ) + # Error fires *before* we ever ask for the calling identity. + ws.current_user.me.assert_not_called() + + def test_raises_when_calling_identity_has_no_username_or_id(self) -> None: + ws = MagicMock(name="WorkspaceClient") + ws.database.get_database_instance.return_value = MagicMock(read_write_dns="some.host:5432") + # Both user_name AND id come back falsy — Lakebase has nothing + # to authenticate as. + ws.current_user.me.return_value = MagicMock(user_name=None, id=None) + + with pytest.raises(RuntimeError, match="Could not determine workspace identity"): + build_pg_executor( + ws, + instance_name="ok-instance", + database="dqx", + schema="dqx_studio", + ) + + def test_uses_id_as_fallback_when_user_name_is_empty(self) -> None: + """``me.user_name or me.id`` fallback — proves we don't bail on user_name=None alone.""" + ws = MagicMock(name="WorkspaceClient") + ws.database.get_database_instance.return_value = MagicMock(read_write_dns="ok.host") + ws.current_user.me.return_value = MagicMock(user_name=None, id="sp-1234") + + # We stop before the real PgExecutor opens its pool — patch + # PgExecutor.__init__ to a no-op and assert it was called with + # the resolved username and host. + with patch( + "databricks_labs_dqx_app.backend.pg_executor.PgExecutor.__init__", + return_value=None, + ) as init_mock: + build_pg_executor( + ws, + instance_name="ok-instance", + database="dqx", + schema="dqx_studio", + ) + + # __init__ is called as PgExecutor.__new__(cls).__init__(ws=..., ...) + # so the kwargs are what we want to inspect. + kwargs = init_mock.call_args.kwargs + assert kwargs["username"] == "sp-1234", "fell through user_name=None to id" + assert kwargs["host"] == "ok.host" + assert kwargs["instance_name"] == "ok-instance" + + +# =========================================================================== +# Data-path methods (execute / query / query_dicts / close / execute_no_schema) +# =========================================================================== + + +class TestDataPathMethods: + """The pool/cursor plumbing methods. + + These all share the same shape: ``with pool.connection() as conn: with + conn.cursor() as cur:`` — so the mock fixtures already in + :func:`_make_pg_executor` are enough to drive them and observe the + SQL passed in and the rows coerced out. + """ + + def test_execute_commits_after_running(self) -> None: + executor = _make_pg_executor() + executor.execute("UPDATE t SET x = 1") + cur = _cursor_of(executor) + # Now TWO cursor.execute calls: SET LOCAL guard + user SQL. + # The user SQL is the *second* call; the first is the + # statement_timeout preamble verified in TestStatementTimeout. + assert cur.execute.call_count == 2 + assert cur.execute.call_args_list[-1].args == ("UPDATE t SET x = 1",) + # One-shot statements MUST auto-commit (parity with SqlExecutor). + _conn_of(executor).commit.assert_called_once() + + def test_execute_no_schema_delegates_to_execute(self) -> None: + executor = _make_pg_executor() + executor.execute_no_schema("CREATE SCHEMA IF NOT EXISTS s") + cur = _cursor_of(executor) + # Same two-call shape as execute(): SET LOCAL + user SQL. + assert cur.execute.call_count == 2 + assert cur.execute.call_args_list[-1].args == ("CREATE SCHEMA IF NOT EXISTS s",) + _conn_of(executor).commit.assert_called_once() + + def test_query_returns_string_coerced_rows(self) -> None: + executor = _make_pg_executor() + cur = _cursor_of(executor) + cur.fetchall.return_value = [ + (1, True, dt.date(2026, 5, 28), None), + (2, False, dt.date(2026, 5, 29), "literal"), + ] + rows = executor.query("SELECT id, flag, day, note FROM t") + # Each cell is coerced via _to_text → strings (or None for NULL). + assert rows == [ + ["1", "true", "2026-05-28", None], + ["2", "false", "2026-05-29", "literal"], + ] + + def test_query_dicts_zips_column_names(self) -> None: + executor = _make_pg_executor() + cur = _cursor_of(executor) + # ``MagicMock(name=...)`` is reserved for the mock identity, not + # the synthesised attribute — set ``.name`` explicitly so the + # implementation's ``d.name`` lookup returns the column label. + col_a, col_b = MagicMock(), MagicMock() + col_a.name = "a" + col_b.name = "b" + cur.description = [col_a, col_b] + cur.fetchall.return_value = [(1, True), (2, None)] + rows = executor.query_dicts("SELECT a, b FROM t") + assert rows == [{"a": "1", "b": "true"}, {"a": "2", "b": None}] + + def test_query_dicts_handles_empty_description(self) -> None: + """``cur.description`` can be ``None`` for DDL — must not raise.""" + executor = _make_pg_executor() + cur = _cursor_of(executor) + cur.description = None + cur.fetchall.return_value = [] + assert executor.query_dicts("SELECT 1 WHERE FALSE") == [] + + # ----- statement_timeout enforcement (per-call) ----------------------- + + def _execute_calls(self, executor: PgExecutor) -> list[str]: + """Return the raw SQL strings ``cur.execute`` was called with.""" + cur = _cursor_of(executor) + return [call.args[0] for call in cur.execute.call_args_list] + + def test_execute_emits_set_local_statement_timeout_before_user_sql(self) -> None: + """Default 120s timeout ⇒ ``SET LOCAL statement_timeout = 120000`` first.""" + executor = _make_pg_executor() + executor.execute("UPDATE t SET x = 1") + calls = self._execute_calls(executor) + assert len(calls) == 2, ( + "Expected exactly 2 cursor.execute calls (SET LOCAL guard + user SQL); " f"got {len(calls)}: {calls}" + ) + assert calls[0] == "SET LOCAL statement_timeout = 120000", ( + f"First call must be the SET LOCAL guard, got {calls[0]!r}. " + "Without this, a runaway query pins the pool connection indefinitely." + ) + assert calls[1] == "UPDATE t SET x = 1" + + def test_execute_respects_custom_timeout_seconds(self) -> None: + """Per-call ``timeout_seconds`` overrides the default cleanly.""" + executor = _make_pg_executor() + executor.execute("DELETE FROM t WHERE id < 100", timeout_seconds=5) + calls = self._execute_calls(executor) + assert calls[0] == "SET LOCAL statement_timeout = 5000" + assert calls[1] == "DELETE FROM t WHERE id < 100" + + def test_query_emits_set_local_before_user_sql(self) -> None: + """``query`` must guard its statement too — a misbehaving SELECT + with a missing predicate is just as capable of pinning a pool + connection as a runaway UPDATE.""" + executor = _make_pg_executor() + cur = _cursor_of(executor) + cur.fetchall.return_value = [] + executor.query("SELECT * FROM t WHERE id > 0", timeout_seconds=30) + calls = self._execute_calls(executor) + assert calls[0] == "SET LOCAL statement_timeout = 30000" + assert calls[1] == "SELECT * FROM t WHERE id > 0" + + def test_query_dicts_emits_set_local_before_user_sql(self) -> None: + """``query_dicts`` must guard its statement too.""" + executor = _make_pg_executor() + cur = _cursor_of(executor) + cur.description = None + cur.fetchall.return_value = [] + executor.query_dicts("SELECT * FROM t", timeout_seconds=45) + calls = self._execute_calls(executor) + assert calls[0] == "SET LOCAL statement_timeout = 45000" + assert calls[1] == "SELECT * FROM t" + + def test_timeout_uses_set_local_not_set(self) -> None: + """``SET`` (no LOCAL) would leak the timeout to the next pool + caller because connections are reused. ``SET LOCAL`` scopes + the GUC to the transaction and resets on COMMIT/ROLLBACK — + the only correct primitive for a pooled executor. + """ + executor = _make_pg_executor() + executor.execute("UPDATE t SET x = 1", timeout_seconds=10) + calls = self._execute_calls(executor) + timeout_call = calls[0] + assert timeout_call.startswith("SET LOCAL "), ( + f"Timeout must use SET LOCAL (transaction-scoped); got {timeout_call!r}. " + "A plain SET would leak the previous caller's timeout to the next " + "checkout — a healthcheck's 1s cap would kill the next migration." + ) + + def test_timeout_seconds_floors_at_one_millisecond(self) -> None: + """Defends against ``timeout_seconds=0`` degenerating into 'no + timeout' (Postgres treats ``statement_timeout=0`` as unlimited). + """ + executor = _make_pg_executor() + executor.execute("UPDATE t SET x = 1", timeout_seconds=0) + calls = self._execute_calls(executor) + assert calls[0] == "SET LOCAL statement_timeout = 1", ( + f"timeout_seconds=0 must NOT translate to statement_timeout=0 " + f"(Postgres treats 0 as unlimited, which is the failure mode " + f"we're guarding against). Got {calls[0]!r}." + ) + + def test_timeout_helper_routes_through_trust_boundary(self) -> None: + """The SET LOCAL statement is dispatched via ``run_trusted_sql`` + (the documented :class:`LiteralString` cast) — proves we don't + bypass the trust-boundary helper for internally-generated SQL. + """ + from databricks_labs_dqx_app.backend import pg_executor as mod + + with patch.object(mod, "run_trusted_sql", wraps=mod.run_trusted_sql) as helper: + executor = _make_pg_executor() + executor.execute("UPDATE t SET x = 1", timeout_seconds=60) + # Two routed calls: SET LOCAL + the user SQL. + assert helper.call_count == 2 + sqls = [c.args[1] for c in helper.call_args_list] + assert sqls[0] == "SET LOCAL statement_timeout = 60000" + assert sqls[1] == "UPDATE t SET x = 1" + + def test_connection_context_manager_does_not_auto_emit_timeout(self) -> None: + """``connection()`` deliberately leaves timeout to the caller — + documented in the docstring because the migration runner uses + it for DDL whose appropriate timeout depends on the work being + done (a 30s query vs a 10-minute index rebuild). + """ + executor = _make_pg_executor() + with executor.connection(): + pass + # No cursor was opened, so no execute calls fired at all. + # We're verifying the manager itself doesn't smuggle a SET + # LOCAL onto the connection behind the caller's back. + cur = _cursor_of(executor) + assert cur.execute.call_count == 0 + + def test_close_signals_stop_and_closes_pool(self) -> None: + executor = _make_pg_executor() + assert not executor._stop.is_set() + executor.close() + assert executor._stop.is_set(), "stop event must be set so refresher loop exits" + executor._pool.close.assert_called_once() + + def test_close_does_not_raise_on_pool_close_error(self) -> None: + """``close`` is best-effort — pool close failure must be swallowed and logged.""" + executor = _make_pg_executor() + executor._pool.close.side_effect = RuntimeError("pool already shut down") + # Must not propagate — Apps platform calls this on shutdown + # and a raise here would mask the real shutdown reason. + executor.close() + assert executor._stop.is_set() + + +# =========================================================================== +# Misc — sanity coverage for adjacent surface area used by callers +# =========================================================================== + + +class TestSimpleProperties: + """The trivial property getters that callers rely on for portability.""" + + def test_dialect_is_postgres(self) -> None: + assert PgExecutor.dialect == "postgres" + + def test_schema_property(self) -> None: + e = _make_pg_executor(schema="my_schema") + assert e.schema == "my_schema" + + def test_database_property(self) -> None: + e = _make_pg_executor(database="my_db") + assert e.database == "my_db" + + def test_catalog_property_returns_database(self) -> None: + """``catalog`` aliases ``database`` so portable 3-part FQN code works.""" + e = _make_pg_executor(database="my_db") + assert e.catalog == "my_db" + + def test_warehouse_id_is_empty_string_for_type_parity(self) -> None: + """Kept on the class for SqlExecutor type-compatibility — Lakebase has no warehouse.""" + e = _make_pg_executor() + assert e.warehouse_id == "" + + def test_fqn_returns_two_part_name(self) -> None: + """Postgres has one catalog per connection — fqn drops the catalog part.""" + e = _make_pg_executor(schema="my_schema") + assert e.fqn("dq_quality_rules") == "my_schema.dq_quality_rules" + + def test_ts_text_is_identity(self) -> None: + """Postgres lets ``_to_text`` ISO-format the timestamp on the way out.""" + e = _make_pg_executor() + assert e.ts_text("created_at") == "created_at" + + def test_json_literal_expr_renders_jsonb_cast(self) -> None: + e = _make_pg_executor() + assert e.json_literal_expr('{"k":"v"}') == "'{\"k\":\"v\"}'::jsonb" + + def test_json_literal_expr_escapes_apostrophes(self) -> None: + e = _make_pg_executor() + assert e.json_literal_expr("it's") == "'it''s'::jsonb" + + +class TestTokenHolder: + """Thread-safe accessor for the rotating OAuth credential.""" + + def test_get_returns_initial(self) -> None: + h = _TokenHolder("initial") + assert h.token == "initial" + + def test_set_updates_value(self) -> None: + h = _TokenHolder("initial") + h.token = "updated" + assert h.token == "updated" + + def test_concurrent_read_write_does_not_corrupt(self) -> None: + """Smoke test: the lock prevents torn reads on the str pointer. + + Strings are atomically assigned on CPython, so this is more of + a contract demonstration than a true thread-safety check — + but it verifies the lock doesn't deadlock under contention, + which is what would actually break a long-running app. + """ + h = _TokenHolder("tok-0") + # Worker threads can't propagate exceptions to the main test + # thread automatically — anything they raise becomes a silent + # ``Thread._invoke_excepthook`` log and the test would falsely + # pass. We catch ``BaseException`` (the widest catchable) so + # even KeyboardInterrupt / SystemExit from a buggy lock impl + # surfaces as a real assertion failure below. See the BLE001 + # policy block in ``pyproject.toml`` — test thread-bridge + # collectors are the same resilience pattern as production + # background-thread loops. + errors: list[BaseException] = [] + + def reader() -> None: + try: + for _ in range(1000): + _ = h.token + except BaseException as exc: + errors.append(exc) + + def writer() -> None: + try: + for i in range(1000): + h.token = f"tok-{i}" + except BaseException as exc: + errors.append(exc) + + threads = [threading.Thread(target=reader) for _ in range(3)] + [ + threading.Thread(target=writer) for _ in range(2) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=5) + assert not t.is_alive(), "Token holder deadlocked under contention" + assert not errors, f"Concurrent access raised: {errors}" + + +class TestGenerateToken: + """Thin wrapper around ``ws.database.generate_database_credential``.""" + + def test_returns_token_from_credential_response(self) -> None: + ws = MagicMock(name="WorkspaceClient") + ws.database.generate_database_credential.return_value = MagicMock(token="fresh-tok") + assert _generate_token(ws, "my-instance") == "fresh-tok" + # Request ID should be a UUID-shaped string and instance_names a single-item list. + call_kwargs = ws.database.generate_database_credential.call_args.kwargs + assert call_kwargs["instance_names"] == ["my-instance"] + assert isinstance(call_kwargs["request_id"], str) + assert len(call_kwargs["request_id"]) >= 32 + + def test_raises_when_credential_response_has_no_token(self) -> None: + ws = MagicMock(name="WorkspaceClient") + ws.database.generate_database_credential.return_value = MagicMock(token=None) + with pytest.raises(RuntimeError, match="no token"): + _generate_token(ws, "my-instance") diff --git a/app/tests/test_pg_migration_runner.py b/app/tests/test_pg_migration_runner.py new file mode 100644 index 000000000..946771873 --- /dev/null +++ b/app/tests/test_pg_migration_runner.py @@ -0,0 +1,534 @@ +"""Unit tests for :class:`PgMigrationRunner` (Lakebase Postgres migrations). + +The runner is tested in pure isolation against a typed mock of the real +:class:`PgExecutor`. Production exercises a live Lakebase instance via +integration tests; here we focus on dispatch, atomicity, and identifier +quoting. + +Why the helpers look the way they do +------------------------------------ +* **`create_autospec(PgExecutor, instance=True)`** rather than a bare + ``MagicMock``. The spec'd mock rejects attribute access for any + attribute the real :class:`PgExecutor` doesn't expose, so a future + rename (``query`` → ``run_query``, etc.) surfaces as an + AttributeError at test time. AGENTS.md mandates this pattern: + "construct dependencies with create_autospec rather than patching + internal module state." + +* **No explicit ``__exit__`` wiring.** The auto-generated context- + manager support on the return values of method calls returns + ``False`` from ``__exit__`` by default — exceptions inside the + ``with`` block propagate out correctly without a manual override. + (The earlier version of this file pre-emptively forced ``__exit__`` + to ``None`` "just in case"; that was cargo-culted from a different + failure mode and is no longer needed once the helper is spec'd + against the real surface.) + +* **Migration list is injected via the constructor seam.** The runner + accepts a ``migrations=`` kwarg defaulting to the production + catalogue. Tests pass a deterministic fake list directly rather + than monkey-patching the module-level ``PG_MIGRATIONS`` — the + monkey-patch pattern coupled tests to an import path and produced + surprise ordering issues when multiple tests ran in the same + process. +""" + +from __future__ import annotations + +import dataclasses +from unittest.mock import MagicMock, create_autospec + +import pytest + +from databricks_labs_dqx_app.backend.migrations.postgres import ( + PG_MIGRATIONS, + PgMigration, + PgMigrationRunner, +) +from databricks_labs_dqx_app.backend.pg_executor import PgExecutor + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_executor(*, applied_versions: tuple[int, ...] = (), schema: str = "public") -> MagicMock: + """Build a typed mock of :class:`PgExecutor` pre-seeded with applied versions. + + Uses ``create_autospec`` so the mock rejects access to attributes + that aren't on the real class — a future rename of any method the + runner calls produces a loud test-time AttributeError rather than + a silent passthrough. + """ + exec_mock = create_autospec(PgExecutor, instance=True) + exec_mock.schema = schema + exec_mock.database = "test_db" + # Mirror :meth:`PgExecutor.q` — ANSI double-quotes with internal + # ``"`` doubled. The runner uses ``q()`` for every schema reference + # so a hyphenated schema name stays parseable end-to-end. + exec_mock.q.side_effect = lambda ident: '"' + ident.replace('"', '""') + '"' + # ``_applied_versions`` issues one SELECT and parses ``[[str(v)], ...]``. + exec_mock.query.return_value = [[str(v)] for v in applied_versions] + return exec_mock + + +def _connection_of(exec_mock: MagicMock) -> MagicMock: + """Return the ``conn`` object the runner sees inside its ``with`` block.""" + return exec_mock.connection.return_value.__enter__.return_value + + +def _cursor_of(exec_mock: MagicMock) -> MagicMock: + """Return the ``cur`` object the runner uses inside :meth:`_apply`.""" + return _connection_of(exec_mock).cursor.return_value.__enter__.return_value + + +def _executed_sqls(cursor_mock: MagicMock) -> list[str]: + """All SQL strings that hit the inner cursor, in call order.""" + return [c.args[0] for c in cursor_mock.execute.call_args_list] + + +def _insert_meta_sqls(cursor_mock: MagicMock) -> list[str]: + """Just the ``INSERT INTO ...dq_migrations`` SQL templates.""" + return [s for s in _executed_sqls(cursor_mock) if "INSERT INTO" in s and "dq_migrations" in s] + + +def _insert_meta_calls(cursor_mock: MagicMock) -> list[tuple[str, tuple[object, ...]]]: + """Return ``(sql_template, params)`` pairs for every meta-table INSERT. + + Post-review the runner uses psycopg native parameter binding + (:func:`run_parameterized_sql` -> ``cur.execute(template, params)``) + rather than f-string-baking the version/description into the SQL + string. Tests assert on the bound ``params`` tuple rather than + on substrings of the template — manual escaping in the template + is exactly the pattern the security review rejected. See + ``backend.pg_executor.run_trusted_sql`` for the full contract. + """ + out: list[tuple[str, tuple[object, ...]]] = [] + for c in cursor_mock.execute.call_args_list: + sql = c.args[0] + if "INSERT INTO" in sql and "dq_migrations" in sql: + params = tuple(c.args[1]) if len(c.args) >= 2 else () + out.append((sql, params)) + return out + + +def _insert_meta_versions(cursor_mock: MagicMock) -> list[object]: + """Just the migration-version values bound to dq_migrations INSERTs. + + The first positional ``%s`` in the template is ``version``; this + helper hands the test the *bound* values so assertions don't have + to know about the placeholder/params split. + """ + return [params[0] for _, params in _insert_meta_calls(cursor_mock) if params] + + +# --------------------------------------------------------------------------- +# PG_MIGRATIONS catalogue invariants +# --------------------------------------------------------------------------- + + +class TestPgMigrationsCatalogue: + def test_versions_are_unique(self): + versions = [m.version for m in PG_MIGRATIONS] + assert len(versions) == len(set(versions)), "PG_MIGRATIONS version numbers must be unique" + + def test_versions_are_monotonically_increasing(self): + versions = [m.version for m in PG_MIGRATIONS] + assert versions == sorted(versions), "PG_MIGRATIONS entries must be listed in ascending version order" + + def test_every_migration_has_a_description(self): + # Empty descriptions break the dq_migrations INSERT (NOT NULL). + for m in PG_MIGRATIONS: + assert m.description.strip(), f"v{m.version} migration has an empty description" + + +# --------------------------------------------------------------------------- +# PgMigration dataclass behaviour +# --------------------------------------------------------------------------- + + +class TestPgMigrationDataclass: + def test_is_frozen(self): + m = PgMigration(version=99, description="test", sql="SELECT 1") + with pytest.raises(dataclasses.FrozenInstanceError): + m.version = 100 # type: ignore[misc] + + def test_equality_is_by_value(self): + a = PgMigration(version=1, description="d", sql="s") + b = PgMigration(version=1, description="d", sql="s") + c = PgMigration(version=2, description="d", sql="s") + assert a == b + assert a != c + + +# --------------------------------------------------------------------------- +# run_all — pending/applied dispatch +# --------------------------------------------------------------------------- + + +class TestRunAllDispatch: + def test_applies_all_when_none_applied(self): + exec_mock = _make_executor(applied_versions=()) + runner = PgMigrationRunner(exec_mock) + + applied = runner.run_all() + + assert applied == len(PG_MIGRATIONS) + # Each migration records its version in dq_migrations. + assert len(_insert_meta_sqls(_cursor_of(exec_mock))) == len(PG_MIGRATIONS) + + def test_skips_all_when_already_applied(self): + all_versions = tuple(m.version for m in PG_MIGRATIONS) + exec_mock = _make_executor(applied_versions=all_versions) + runner = PgMigrationRunner(exec_mock) + + applied = runner.run_all() + + assert applied == 0 + # _apply is never entered, so connection() must not be touched. + exec_mock.connection.assert_not_called() + + def test_runs_only_pending_when_some_already_applied(self): + """With v1 already applied, only the new v2 should run.""" + fake_migrations = [ + PgMigration(version=1, description="v1", sql="CREATE TABLE {schema}.t1 (id int);"), + PgMigration(version=2, description="v2", sql="CREATE TABLE {schema}.t2 (id int);"), + ] + exec_mock = _make_executor(applied_versions=(1,)) + runner = PgMigrationRunner(exec_mock, migrations=fake_migrations) + + applied = runner.run_all() + + assert applied == 1 + cur = _cursor_of(exec_mock) + inserts = _insert_meta_calls(cur) + assert len(inserts) == 1 + # Version 2 is bound through psycopg's parameter binder, not + # baked into the SQL string. + assert _insert_meta_versions(cur) == [2] + # The template uses ``%s`` placeholders for both value + # columns — confirms we didn't regress to f-string interpolation. + template = inserts[0][0] + assert "VALUES (%s, %s," in template + # v1's DDL must not be re-applied. + sqls = _executed_sqls(cur) + assert not any(".t1 " in s for s in sqls), "v1 DDL ran a second time" + assert any(".t2 " in s for s in sqls), "v2 DDL did not run" + + def test_bootstraps_schema_and_meta_table_before_any_migration(self): + exec_mock = _make_executor(applied_versions=()) + runner = PgMigrationRunner(exec_mock) + runner.run_all() + + # _ensure_schema and _ensure_meta_table go through executor.execute + # (NOT cursor.execute) — they are one-shot statements applied + # before the transactional _apply loop. + bootstrap = [c.args[0] for c in exec_mock.execute.call_args_list] + assert any("CREATE SCHEMA IF NOT EXISTS" in s for s in bootstrap) + assert any("CREATE TABLE IF NOT EXISTS" in s and "dq_migrations" in s for s in bootstrap) + + def test_schema_placeholder_is_substituted(self): + """Every {schema} placeholder must be replaced before SQL leaves the runner.""" + exec_mock = _make_executor(applied_versions=(), schema="custom_app_schema") + runner = PgMigrationRunner(exec_mock) + runner.run_all() + + cur = _cursor_of(exec_mock) + for sql in _executed_sqls(cur): + assert "{schema}" not in sql, f"Unsubstituted placeholder in: {sql!r}" + # And the configured schema name actually appears somewhere. + assert any("custom_app_schema" in s for s in _executed_sqls(cur)) + + +# --------------------------------------------------------------------------- +# Constructor injection — migrations= kwarg behaviour +# --------------------------------------------------------------------------- + + +class TestMigrationsInjection: + """The ``migrations=`` constructor seam exists so tests don't have + to monkey-patch the module-level constant. These tests pin down the + contract so a future refactor doesn't silently drop the seam.""" + + def test_default_is_the_module_level_catalogue(self): + """``migrations=None`` (the default) must wire to ``PG_MIGRATIONS``.""" + exec_mock = _make_executor(applied_versions=tuple(m.version for m in PG_MIGRATIONS)) + runner = PgMigrationRunner(exec_mock) + # If the default wiring is correct then "every version already + # applied" causes zero new applications. + assert runner.run_all() == 0 + + def test_injected_list_is_used_in_place_of_default(self): + """An explicit ``migrations=`` list must shadow ``PG_MIGRATIONS``.""" + fake = [ + PgMigration(version=10, description="injected v10", sql="CREATE TABLE {schema}.t10 (id int);"), + PgMigration(version=11, description="injected v11", sql="CREATE TABLE {schema}.t11 (id int);"), + ] + exec_mock = _make_executor(applied_versions=()) + runner = PgMigrationRunner(exec_mock, migrations=fake) + + assert runner.run_all() == len(fake) + cur = _cursor_of(exec_mock) + inserts = _insert_meta_calls(cur) + # Only the injected versions land in dq_migrations — no leak + # from the production catalogue. The versions are bound via + # psycopg's parameter binder, so we read them from the + # ``params`` tuple rather than substring-searching the SQL. + assert len(inserts) == len(fake) + bound_versions = _insert_meta_versions(cur) + assert 10 in bound_versions + assert 11 in bound_versions + production_versions = {m.version for m in PG_MIGRATIONS} + assert not (set(bound_versions) & production_versions), ( + "Injected runner leaked production-catalogue versions: " f"{set(bound_versions) & production_versions}" + ) + + def test_injected_list_is_snapshotted_defensively(self): + """Mutating the caller's list after construction must not move the runner's view. + + The runner stores a tuple snapshot to prevent the + action-at-a-distance footgun where a test fixture mutates its + shared list and quietly changes runner behaviour mid-suite. + """ + fake = [PgMigration(version=1, description="v1", sql="CREATE TABLE {schema}.t (id int);")] + runner = PgMigrationRunner(_make_executor(), migrations=fake) + + fake.append(PgMigration(version=2, description="v2", sql="CREATE TABLE {schema}.t2 (id int);")) + # Snapshot length is what was passed at construction time. + assert len(runner._migrations) == 1 + + def test_empty_injected_list_runs_zero_migrations_but_still_bootstraps(self): + """An empty migration list still produces schema + meta-table bootstrap.""" + exec_mock = _make_executor(applied_versions=()) + runner = PgMigrationRunner(exec_mock, migrations=[]) + + assert runner.run_all() == 0 + bootstrap = [c.args[0] for c in exec_mock.execute.call_args_list] + assert any("CREATE SCHEMA IF NOT EXISTS" in s for s in bootstrap) + assert any("dq_migrations" in s and "CREATE TABLE IF NOT EXISTS" in s for s in bootstrap) + # No _apply ever runs → connection() must not be touched. + exec_mock.connection.assert_not_called() + + +# --------------------------------------------------------------------------- +# _apply — atomicity contract +# --------------------------------------------------------------------------- + + +class TestApplyAtomicity: + """The contract Laurence's review surfaced: DDL + version row are atomic.""" + + def test_one_commit_per_migration(self): + """Each migration ends with exactly one ``conn.commit()`` — the + DDL and the dq_migrations INSERT share that single transaction.""" + exec_mock = _make_executor(applied_versions=()) + runner = PgMigrationRunner(exec_mock) + runner.run_all() + + conn = _connection_of(exec_mock) + assert conn.commit.call_count == len(PG_MIGRATIONS) + + def test_per_statement_split_emits_multiple_cursor_executes(self): + """The runner still splits on ``;`` so cursor errors pinpoint the + exact failing DDL statement rather than a position inside a + multi-kilobyte compound string.""" + exec_mock = _make_executor(applied_versions=()) + runner = PgMigrationRunner(exec_mock) + runner.run_all() + + cur = _cursor_of(exec_mock) + # Floor: every migration produces at least one DDL execute plus + # the bookkeeping INSERT. + assert cur.execute.call_count >= 2 * len(PG_MIGRATIONS) + + def test_version_row_insert_runs_through_the_same_cursor(self): + """The dq_migrations INSERT must reach the same transactional + cursor as the DDL — never the executor's auto-committing + ``execute`` path.""" + exec_mock = _make_executor(applied_versions=()) + runner = PgMigrationRunner(exec_mock) + runner.run_all() + + cur_inserts = _insert_meta_sqls(_cursor_of(exec_mock)) + exec_inserts = [c.args[0] for c in exec_mock.execute.call_args_list if "INSERT INTO" in c.args[0]] + assert cur_inserts, "dq_migrations INSERT must go through the transactional cursor" + assert not any("dq_migrations" in s for s in exec_inserts), ( + "dq_migrations INSERT leaked onto the auto-committing executor.execute path — " + "would break migration atomicity" + ) + + def test_failure_during_ddl_aborts_transaction_and_propagates(self): + """If a DDL statement raises, the transaction rolls back (no + commit), the version row is never written, and the exception + propagates so the next start retries cleanly.""" + # Inject a deterministic single migration via the constructor + # seam so we know exactly which DDL statement should fail. + fake_migration = PgMigration( + version=42, + description="failing test", + sql=( + "CREATE TABLE {schema}.t1 (id int);" + "CREATE TABLE {schema}.t2 (id int);" + "CREATE TABLE {schema}.t3 (id int);" + ), + ) + + exec_mock = _make_executor(applied_versions=()) + cur = _cursor_of(exec_mock) + + # Content-keyed failure instead of a positional ``[None, None, + # RuntimeError, ...]`` list. The earlier form was brittle: if + # the runner ever changed how it splits or batches statements + # (e.g. adds an explicit ``BEGIN``, switches from per-statement + # ``execute`` to ``executemany``, or skips empty trailing + # tokens differently), the failure index would drift and the + # test would pass for the wrong reason — green even though the + # rollback path was never actually exercised. Keying the + # failure to a substring of the offending DDL means the test + # fails on the same statement regardless of how the runner + # batches around it. + boom_marker = ".t2 " # space avoids matching the un-related ``t2x`` future-proofing + + def _fail_on_t2(sql: str, *_args: object, **_kwargs: object) -> None: + if boom_marker in sql: + raise RuntimeError("boom") + + cur.execute.side_effect = _fail_on_t2 + + runner = PgMigrationRunner(exec_mock, migrations=[fake_migration]) + with pytest.raises(RuntimeError, match="boom"): + runner.run_all() + + conn = _connection_of(exec_mock) + # commit() must NOT have run — Postgres would roll the + # transaction back when the connection returns to the pool. + conn.commit.assert_not_called() + # And the version row must never have been requested either. + assert not _insert_meta_sqls(cur) + + # The failing statement WAS the one we targeted (the assertion + # is what makes the content-keyed predicate load-bearing — if + # the runner stops emitting ``.t2`` entirely the test fails + # loudly here rather than silently). Equally, anything emitted + # AFTER ``.t2`` proves the abort path is broken: a single + # failed statement must stop the whole migration in its + # tracks, not "let it skip and move on". + sqls = _executed_sqls(cur) + assert any(boom_marker in s for s in sqls), ( + "The targeted DDL statement was never issued — the runner " + "may have stopped emitting it, in which case the rollback " + "path is no longer exercised by this test." + ) + idx_of_failure = next(i for i, s in enumerate(sqls) if boom_marker in s) + assert not any(".t3 " in s for s in sqls[idx_of_failure + 1 :]), ( + "Runner kept executing DDL after a statement raised — the " + "transaction's all-or-nothing semantics rely on the loop " + "short-circuiting on first failure." + ) + + def test_failure_in_first_migration_skips_subsequent_migrations(self): + """A failed migration must short-circuit the rest of the run — + otherwise a later migration could land on top of a partially- + rolled-back schema.""" + fake_migrations = [ + PgMigration(version=1, description="v1", sql="BROKEN;"), + PgMigration(version=2, description="v2", sql="CREATE TABLE {schema}.t2 (id int);"), + ] + + exec_mock = _make_executor(applied_versions=()) + cur = _cursor_of(exec_mock) + cur.execute.side_effect = RuntimeError("boom") # every execute fails + + runner = PgMigrationRunner(exec_mock, migrations=fake_migrations) + with pytest.raises(RuntimeError, match="boom"): + runner.run_all() + + # v2's DDL must not have been attempted. + sqls = _executed_sqls(cur) + assert not any(".t2 " in s for s in sqls), "v2 ran despite v1 failing — short-circuit broken" + + +# --------------------------------------------------------------------------- +# Executor-protocol consolidation +# --------------------------------------------------------------------------- + + +class TestExecutorProtocolConsolidation: + """The migration runner's ``_Executor`` Protocol MUST extend + :class:`OltpExecutorProtocol`. + + Why this matters: pre-consolidation the module re-declared its own + ``schema`` / ``execute`` / ``query`` / ``q`` Protocol members + independently of the shared OLTP contract. That drift is silent — + a future addition to ``OltpExecutorProtocol`` (a new upsert + variant, a new dialect helper) would NOT propagate here, and the + structural-typing surface across the two callers would slowly + diverge until someone tripped over a runtime AttributeError on a + code path that only fires in production. + + Locking the inheritance relationship as a structural assertion + means a regression to the independent-declaration form fails this + test loudly instead of waiting for an integration smoke run. + """ + + def test_executor_protocol_extends_oltp_executor_protocol(self): + from databricks_labs_dqx_app.backend.migrations.postgres import _Executor + from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol + + # Direct MRO check — any future refactor that drops the + # inheritance (e.g. someone "simplifies" by re-listing the + # members inline) trips this immediately. + assert OltpExecutorProtocol in _Executor.__mro__, ( + "_Executor must extend OltpExecutorProtocol so the executor " + "surface stays single-sourced. If you removed the inheritance, " + "you've re-introduced the duplication the consolidation fixed." + ) + + def test_executor_protocol_adds_only_connection_beyond_oltp_surface(self): + """The runner-specific surface is exactly one method: ``connection()``. + + Anything else added here is either (a) redundant with the OLTP + contract — push it up to :class:`OltpExecutorProtocol` — or + (b) leaking psycopg-specific concerns into the shared layer. + Either way it deserves explicit review, which a failing assert + forces. + """ + from databricks_labs_dqx_app.backend.migrations.postgres import _Executor + from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol + + # The Protocol member set lives on ``__protocol_attrs__`` on + # Python 3.12+; on 3.11 we approximate via __annotations__ + + # dir() filtering. Either way: subtract the OLTP surface to + # see what's unique to the migration runner. + oltp_members = set(dir(OltpExecutorProtocol)) - set(dir(object)) + runner_members = set(dir(_Executor)) - set(dir(object)) + runner_only = runner_members - oltp_members + + assert runner_only == {"connection"}, ( + f"_Executor must add exactly 'connection' beyond OltpExecutorProtocol; " + f"got runner-only members: {sorted(runner_only)}. " + "Anything else either belongs on OltpExecutorProtocol " + "(if the OLTP services would use it) or signals new " + "psycopg-specific surface leaking into the shared layer." + ) + + def test_pg_executor_structurally_satisfies_extended_protocol(self): + """PgExecutor remains the canonical implementation post-extension. + + Because ``OltpExecutorProtocol`` is ``@runtime_checkable``, an + ``isinstance`` check passes iff PgExecutor exposes every member + the Protocol requires. If a future PgExecutor refactor drops a + Protocol-required method (or rename without an alias), this + test fails BEFORE production hits the AttributeError. + """ + from databricks_labs_dqx_app.backend.sql_executor import OltpExecutorProtocol + + exec_mock = _make_executor() + assert isinstance(exec_mock, OltpExecutorProtocol), ( + "create_autospec(PgExecutor) must satisfy OltpExecutorProtocol — " + "if this fails, PgExecutor no longer structurally satisfies the " + "shared Protocol and every OLTP service annotation is now lying." + ) + # And the runner-only surface is on the mock too. + assert hasattr(exec_mock, "connection") diff --git a/app/tests/test_retention.py b/app/tests/test_retention.py new file mode 100644 index 000000000..a9240d9d5 --- /dev/null +++ b/app/tests/test_retention.py @@ -0,0 +1,441 @@ +"""Tests for the retention sweep + admin settings surface. + +Three layers exercised: +* ``AppSettingsService.get_retention_days`` / + ``get_quarantine_retention_days`` and their setters — round-trip the + integer-day setting through ``dq_app_settings``. +* ``backend.routes.v1.config._validate_retention_days`` — request-side + bounds checking for the admin PUT endpoint. +* ``SchedulerService._resolve_retention_days`` and + ``_resolve_quarantine_retention_days`` — fall-back behaviour when the + setting is missing / unparseable / read fails, plus the safety floor. + +The retention defaults differ on purpose — global is 90 days (trend +dashboards) and quarantine is 30 days (tighter PII window) — so the +tests exercise both resolvers separately. +""" + +from __future__ import annotations + +import pytest + + +# --------------------------------------------------------------------------- +# AppSettingsService — storage layer +# --------------------------------------------------------------------------- + + +class TestAppSettingsRetention: + @pytest.fixture + def svc(self, sql_executor_mock): + from databricks_labs_dqx_app.backend.services.app_settings_service import AppSettingsService + + return AppSettingsService(sql_executor_mock), sql_executor_mock + + def test_get_returns_none_when_unset(self, svc): + s, sql = svc + sql.query.return_value = [] + assert s.get_retention_days() is None + assert s.get_quarantine_retention_days() is None + + def test_get_returns_none_when_blank(self, svc): + s, sql = svc + sql.query.return_value = [(None,)] + assert s.get_retention_days() is None + sql.query.return_value = [("",)] + assert s.get_retention_days() is None + + def test_get_parses_integer(self, svc): + s, sql = svc + sql.query.return_value = [("45",)] + assert s.get_retention_days() == 45 + sql.query.return_value = [("14",)] + assert s.get_quarantine_retention_days() == 14 + + def test_get_returns_none_on_garbage_value(self, svc): + s, sql = svc + sql.query.return_value = [("not-a-number",)] + assert s.get_retention_days() is None + + def test_save_persists_integer_string(self, svc): + s, sql = svc + assert s.save_retention_days(60) == 60 + assert sql.upsert.called + kwargs = sql.upsert.call_args.kwargs + if "value_cols" in kwargs: + payload = kwargs["value_cols"]["setting_value"] + else: + payload = sql.upsert.call_args.args[2]["setting_value"] + assert payload == "60" + + def test_save_quarantine_persists_integer_string(self, svc): + s, sql = svc + assert s.save_quarantine_retention_days(21) == 21 + # The most recent upsert should carry the new key. + kwargs = sql.upsert.call_args.kwargs + if "key_cols" in kwargs: + keys = kwargs["key_cols"] + else: + keys = sql.upsert.call_args.args[1] + assert keys["setting_key"] == "quarantine_retention_days" + + +# --------------------------------------------------------------------------- +# Route-level validator +# --------------------------------------------------------------------------- + + +class TestRetentionDaysValidator: + @pytest.fixture + def validate(self): + from databricks_labs_dqx_app.backend.routes.v1.config import _validate_retention_days + + return _validate_retention_days + + def test_min_floor_enforced(self, validate): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + validate(1, field="retention_days") + assert exc.value.status_code == 400 + assert "at least" in exc.value.detail.lower() + + def test_max_ceiling_enforced(self, validate): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + validate(10_000, field="retention_days") + assert exc.value.status_code == 400 + assert "at most" in exc.value.detail.lower() + + def test_valid_value_passes_through(self, validate): + assert validate(30, field="quarantine_retention_days") == 30 + assert validate(90, field="retention_days") == 90 + + def test_field_label_appears_in_message(self, validate): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + validate(2, field="quarantine_retention_days") + assert "quarantine_retention_days" in exc.value.detail + + +# --------------------------------------------------------------------------- +# Scheduler service — both resolvers honour the floor and fall back safely +# --------------------------------------------------------------------------- + + +class TestSchedulerResolveRetention: + """Scheduler must never crash on a missing / corrupt retention setting. + + Built via the real ``SchedulerService(...)`` constructor through + the shared :func:`make_scheduler` factory so a rename of the + public ``oltp_sql=`` constructor parameter fails at construction + time rather than silently producing AttributeErrors later. + """ + + @pytest.fixture + def scheduler(self, make_scheduler): + return make_scheduler() + + def test_global_falls_back_to_default_when_unset(self, scheduler): + from databricks_labs_dqx_app.backend.services.scheduler_service import _RETENTION_DAYS_DEFAULT + + svc, mocks = scheduler + mocks.oltp.query.return_value = [] + assert svc._resolve_retention_days() == _RETENTION_DAYS_DEFAULT + + def test_quarantine_falls_back_to_its_own_default(self, scheduler): + from databricks_labs_dqx_app.backend.services.scheduler_service import ( + _QUARANTINE_RETENTION_DAYS_DEFAULT, + ) + + svc, mocks = scheduler + mocks.oltp.query.return_value = [] + assert svc._resolve_quarantine_retention_days() == _QUARANTINE_RETENTION_DAYS_DEFAULT + + def test_quarantine_default_is_tighter_than_global(self): + """Sanity check on the constants: quarantine (PII) ages out faster.""" + from databricks_labs_dqx_app.backend.services.scheduler_service import ( + _QUARANTINE_RETENTION_DAYS_DEFAULT, + _RETENTION_DAYS_DEFAULT, + ) + + assert _QUARANTINE_RETENTION_DAYS_DEFAULT < _RETENTION_DAYS_DEFAULT + + def test_returns_persisted_value(self, scheduler): + svc, mocks = scheduler + mocks.oltp.query.return_value = [("45",)] + assert svc._resolve_retention_days() == 45 + + def test_floor_protects_against_too_small_value(self, scheduler): + from databricks_labs_dqx_app.backend.services.scheduler_service import _RETENTION_DAYS_MIN + + svc, mocks = scheduler + mocks.oltp.query.return_value = [("1",)] + # Stored value of 1 must be raised to the floor — never wipe inside the safety window. + assert svc._resolve_retention_days() == _RETENTION_DAYS_MIN + assert svc._resolve_quarantine_retention_days() == _RETENTION_DAYS_MIN + + def test_swallows_query_exception(self, scheduler): + from databricks_labs_dqx_app.backend.services.scheduler_service import ( + _QUARANTINE_RETENTION_DAYS_DEFAULT, + _RETENTION_DAYS_DEFAULT, + ) + + svc, mocks = scheduler + mocks.oltp.query.side_effect = RuntimeError("warehouse offline") + # Must not propagate — sweep retries on the next tick. + assert svc._resolve_retention_days() == _RETENTION_DAYS_DEFAULT + assert svc._resolve_quarantine_retention_days() == _QUARANTINE_RETENTION_DAYS_DEFAULT + + def test_unparseable_value_returns_default(self, scheduler): + from databricks_labs_dqx_app.backend.services.scheduler_service import _RETENTION_DAYS_DEFAULT + + svc, mocks = scheduler + mocks.oltp.query.return_value = [("not-a-number",)] + assert svc._resolve_retention_days() == _RETENTION_DAYS_DEFAULT + + +# --------------------------------------------------------------------------- +# Scheduler service — _run_retention applies the right cutoff per table +# --------------------------------------------------------------------------- + + +class TestRunRetentionUsesQuarantineCutoff: + """The Delta sweep must use the quarantine-specific cutoff for that one table.""" + + @pytest.fixture + def scheduler(self, make_scheduler): + # ``distinct_sql=True`` so we can introspect Delta-side + # ``execute`` calls independently of OLTP ones below. + return make_scheduler(distinct_sql=True) + + def test_quarantine_table_uses_its_own_default(self, scheduler): + from databricks_labs_dqx_app.backend.services.scheduler_service import ( + _QUARANTINE_RETENTION_DAYS_DEFAULT, + _RETENTION_DAYS_DEFAULT, + ) + + svc, mocks = scheduler + svc._run_retention() + + # Collect every DELETE statement that actually fired. + delta_stmts = [c.args[0] for c in mocks.sql.execute.call_args_list] + quarantine_stmts = [s for s in delta_stmts if "dq_quarantine_records" in s] + non_quarantine_delta = [s for s in delta_stmts if "dq_quarantine_records" not in s] + + assert quarantine_stmts, "Expected at least one DELETE for the quarantine table" + # Every quarantine DELETE should use the *quarantine* default, + # not the global one. + for stmt in quarantine_stmts: + assert f"INTERVAL {_QUARANTINE_RETENTION_DAYS_DEFAULT} DAY" in stmt + assert f"INTERVAL {_RETENTION_DAYS_DEFAULT} DAY" not in stmt + + # And the other Delta tables should still use the global default. + for stmt in non_quarantine_delta: + assert f"INTERVAL {_RETENTION_DAYS_DEFAULT} DAY" in stmt + + +# --------------------------------------------------------------------------- +# Scheduler service — _run_retention uses dialect-specific INTERVAL syntax +# --------------------------------------------------------------------------- + + +class TestRunRetentionPostgresDialect: + """OLTP retention must emit Postgres-flavoured ``INTERVAL`` syntax. + + Delta speaks ``INTERVAL 90 DAY`` (no quotes, uppercase singular); + Postgres speaks ``INTERVAL '90 days'`` (single-quoted literal, + lowercase plural). The two syntaxes are mutually incompatible — + Postgres throws ``ERROR: syntax error at or near "DAY"`` on the + Delta form, and Delta throws on the quoted form — so the + scheduler branches on ``self._oltp_sql.dialect``. This class + parallels :class:`TestRunRetentionUsesQuarantineCutoff` (which + covers the Delta dialect) for the Postgres branch. + + The reviewer's specific contract to lock down: + + 1. OLTP DELETEs use ``INTERVAL ' days'`` (and NEVER + ``INTERVAL DAY``) when the executor reports + ``dialect == "postgres"``. + 2. Cutoff routing is correct: OLTP tables today use the *global* + ``days`` (no OLTP quarantine table), while the Delta loop still + routes ``dq_quarantine_records`` to ``quarantine_days``. + 3. The Delta loop is untouched by the OLTP dialect flip — its + statements keep using the Delta ``INTERVAL N DAY`` form. + """ + + @pytest.fixture + def scheduler(self, make_scheduler): + # ``distinct_sql=True`` keeps ``mocks.sql`` (analytical Delta) + # and ``mocks.oltp`` (Postgres-flavoured) addressable as + # separate mocks so the test can assert which loop emitted + # which DELETE. + return make_scheduler(oltp_dialect="postgres", catalog="dqx", schema="public", distinct_sql=True) + + @staticmethod + def _delta_stmts(mocks) -> list[str]: + return [c.args[0] for c in mocks.sql.execute.call_args_list] + + @staticmethod + def _oltp_stmts(mocks) -> list[str]: + return [c.args[0] for c in mocks.oltp.execute.call_args_list] + + def test_oltp_loop_uses_postgres_interval_syntax(self, scheduler): + from databricks_labs_dqx_app.backend.services.scheduler_service import ( + _OLTP_RETENTION_TABLES, + _RETENTION_DAYS_DEFAULT, + ) + + svc, mocks = scheduler + svc._run_retention() + + oltp_stmts = self._oltp_stmts(mocks) + # One DELETE per OLTP retention table. + assert len(oltp_stmts) == len(_OLTP_RETENTION_TABLES) + + for stmt in oltp_stmts: + # Postgres form (with single quotes, lowercase plural "days"). + assert ( + f"INTERVAL '{_RETENTION_DAYS_DEFAULT} days'" in stmt + ), f"OLTP DELETE missing Postgres INTERVAL literal: {stmt!r}" + # Delta form must NOT appear — Postgres would reject it. + assert ( + f"INTERVAL {_RETENTION_DAYS_DEFAULT} DAY" not in stmt + ), f"OLTP DELETE leaked Delta INTERVAL syntax into Postgres: {stmt!r}" + + def test_oltp_loop_uses_global_cutoff_not_quarantine(self, scheduler): + """No OLTP table is in the quarantine list today — every OLTP DELETE uses ``days``.""" + from databricks_labs_dqx_app.backend.services.scheduler_service import ( + _QUARANTINE_RETENTION_DAYS_DEFAULT, + _QUARANTINE_TABLE_NAME, + _OLTP_RETENTION_TABLES, + _RETENTION_DAYS_DEFAULT, + ) + + # Guard the test's premise: if anyone adds the quarantine + # table to the OLTP list, this assertion forces them to + # extend the cutoff-routing logic for the OLTP loop too. + assert _QUARANTINE_TABLE_NAME not in {t for t, _ in _OLTP_RETENTION_TABLES}, ( + "Quarantine table moved to OLTP — extend _run_retention to route " + "quarantine_days for the OLTP branch and update this test." + ) + + svc, mocks = scheduler + svc._run_retention() + + oltp_stmts = self._oltp_stmts(mocks) + for stmt in oltp_stmts: + assert f"INTERVAL '{_RETENTION_DAYS_DEFAULT} days'" in stmt + assert f"INTERVAL '{_QUARANTINE_RETENTION_DAYS_DEFAULT} days'" not in stmt + + def test_delta_loop_untouched_by_postgres_oltp_dialect(self, scheduler): + """Flipping ``_oltp_sql.dialect`` MUST NOT change the Delta loop's syntax.""" + from databricks_labs_dqx_app.backend.services.scheduler_service import ( + _QUARANTINE_RETENTION_DAYS_DEFAULT, + _RETENTION_DAYS_DEFAULT, + ) + + svc, mocks = scheduler + svc._run_retention() + + delta_stmts = self._delta_stmts(mocks) + for stmt in delta_stmts: + # Delta keeps its no-quotes, uppercase-DAY form regardless + # of what the OLTP executor speaks. + if "dq_quarantine_records" in stmt: + assert f"INTERVAL {_QUARANTINE_RETENTION_DAYS_DEFAULT} DAY" in stmt + else: + assert f"INTERVAL {_RETENTION_DAYS_DEFAULT} DAY" in stmt + # No Postgres syntax should ever appear in a Delta DELETE. + assert "days'" not in stmt + assert "'" not in stmt.split("INTERVAL", 1)[1] if "INTERVAL" in stmt else True + + def test_oltp_delete_targets_use_executor_fqn(self, scheduler): + """OLTP loop must build the table path via ``self._oltp_sql.fqn(...)``. + + Hard-coding ``catalog.schema.table`` would break Postgres (no + catalog) and bypass the executor's dialect-aware quoting. + """ + from databricks_labs_dqx_app.backend.services.scheduler_service import _OLTP_RETENTION_TABLES + + svc, mocks = scheduler + svc._run_retention() + + oltp_stmts = self._oltp_stmts(mocks) + # ``make_scheduler`` configures ``oltp.fqn`` to return + # ``catalog.schema.table`` — the same shape :class:`SqlExecutor` + # produces. Any code change that bypasses ``fqn()`` and + # hardcodes its own FQN shape will fail this assertion. + for table_name, _ in _OLTP_RETENTION_TABLES: + assert any( + f"DELETE FROM dqx.public.{table_name} " in s for s in oltp_stmts + ), f"No OLTP DELETE used fqn()'d path for {table_name}" + # A hardcoded Delta-quoted three-part form must not leak in. + assert not any(f"`dqx`.`public`.{table_name}" in s for s in oltp_stmts) + + def test_oltp_failure_does_not_abort_remaining_tables(self, scheduler): + """Parity with the Delta loop — one bad DELETE must not skip the others.""" + from databricks_labs_dqx_app.backend.services.scheduler_service import _OLTP_RETENTION_TABLES + + svc, mocks = scheduler + # First OLTP DELETE blows up; the rest must still fire. + call_count = {"n": 0} + + def flaky_execute(*_, **__): + call_count["n"] += 1 + if call_count["n"] == 1: + raise RuntimeError("transient Lakebase failure") + + mocks.oltp.execute.side_effect = flaky_execute + # Must not propagate — _run_retention catches per-table. + svc._run_retention() + # All OLTP tables were attempted despite the early failure. + assert call_count["n"] == len(_OLTP_RETENTION_TABLES) + + +class TestRunRetentionExecutorContract: + """``_run_retention`` is dialect-agnostic — it just calls + :meth:`OltpExecutorProtocol.interval_days_expr` and trusts the + executor to render the right form. + + The pre-Protocol design used a defensive + ``getattr(oltp, "dialect", "delta") == "postgres"`` branch inside + ``_run_retention``. That fallback no longer exists — every OLTP + executor must implement the Protocol — so the contract worth + locking down now is *the service must NOT inspect ``dialect`` + directly* and *must call ``interval_days_expr`` exactly once per + retention sweep* (DRY: don't re-derive the literal per row). + """ + + def test_interval_literal_is_sourced_from_executor_not_a_dialect_branch(self, make_scheduler): + from databricks_labs_dqx_app.backend.services.scheduler_service import ( + _OLTP_RETENTION_TABLES, + _RETENTION_DAYS_DEFAULT, + ) + + # Configure the OLTP mock to return a clearly identifiable + # sentinel interval — if any DELETE statement is missing the + # sentinel, the service is hand-building the literal instead + # of delegating to the executor. + svc, mocks = make_scheduler(oltp_dialect="postgres", distinct_sql=True) + mocks.oltp.interval_days_expr.side_effect = None + mocks.oltp.interval_days_expr.return_value = "INTERVAL_SENTINEL_42" + + svc._run_retention() + + # interval_days_expr is called exactly once with the resolved + # retention-day count (DRY: not re-derived per OLTP table). + mocks.oltp.interval_days_expr.assert_called_once_with(_RETENTION_DAYS_DEFAULT) + + # Every OLTP DELETE statement embedded the sentinel literally. + oltp_stmts = [c.args[0] for c in mocks.oltp.execute.call_args_list] + assert len(oltp_stmts) == len(_OLTP_RETENTION_TABLES) + for stmt in oltp_stmts: + assert "INTERVAL_SENTINEL_42" in stmt, f"OLTP DELETE bypassed executor.interval_days_expr: {stmt!r}" + # And the service must NOT have synthesised either dialect's + # literal on its own — both would be a regression to the + # branch-on-dialect pattern. + assert f"INTERVAL '{_RETENTION_DAYS_DEFAULT} days'" not in stmt + assert f"INTERVAL {_RETENTION_DAYS_DEFAULT} DAY" not in stmt diff --git a/app/tests/test_role_service.py b/app/tests/test_role_service.py index d3940fdd7..cb700cdfa 100644 --- a/app/tests/test_role_service.py +++ b/app/tests/test_role_service.py @@ -164,3 +164,74 @@ def test_invalidate_forces_next_lookup(self, role_service, sql_executor_mock): sql_executor_mock.query.return_value = [] role_service.list_mappings(use_cache=True) sql_executor_mock.query.assert_called_once() + + +# --------------------------------------------------------------------------- +# create_mapping — dialect-agnostic upsert via the executor Protocol +# --------------------------------------------------------------------------- + + +class TestCreateMappingDelegatesToProtocol: + """``create_mapping`` must NOT branch on ``self._sql.dialect``. + + The pre-refactor code hand-built two SQL strings (MERGE for Delta, + INSERT ... ON CONFLICT for Postgres). The Protocol-based design + delegates that choice to :meth:`OltpExecutorProtocol.upsert_with_audit` + so adding a new backend (e.g. SQLite for local dev) only requires + implementing the Protocol — not touching every service. + + These tests lock down the contract from the service side: + + 1. ``create_mapping`` calls ``upsert_with_audit`` exactly once. + 2. It never calls ``execute`` directly with a hand-built SQL string + (a regression would mean the dialect branch came back). + 3. The kwargs forwarded to the Protocol carry the right audit + semantics: ``preserve_created=True`` and no ``increment_on_update`` + (role mappings don't have a version column). + 4. ``created_*`` and ``updated_*`` are passed as ``RawSql("now()")`` + so the executor renders the dialect-correct timestamp function. + """ + + def test_calls_upsert_with_audit_once_with_correct_audit_kwargs(self, role_service, sql_executor_mock): + from databricks_labs_dqx_app.backend.sql_executor import RawSql + + role_service.create_mapping( + role=UserRole.RULE_AUTHOR.value, + group_name="writers", + user_email="alice@example.com", + ) + + # The service MUST funnel everything through the Protocol's + # upsert_with_audit. Going around it (e.g. via .execute()) would + # mean the dialect branch came back. + sql_executor_mock.upsert_with_audit.assert_called_once() + sql_executor_mock.execute.assert_not_called() + + call = sql_executor_mock.upsert_with_audit.call_args + # Audit semantics: preserve created_* on UPDATE; no version column. + assert call.kwargs.get("preserve_created") is True + assert call.kwargs.get("increment_on_update") is None + + # Keys are the natural identity of the row. + assert call.kwargs["key_cols"] == { + "role": UserRole.RULE_AUTHOR.value, + "group_name": "writers", + } + + # Audit cols carry the actor + dialect-portable "now" sentinel. + value_cols = call.kwargs["value_cols"] + assert value_cols["created_by"] == "alice@example.com" + assert value_cols["updated_by"] == "alice@example.com" + # Timestamps go through RawSql so the executor renders the + # right function call per dialect (now() on both, but the + # plumbing must be Protocol-mediated, not stringified). + assert isinstance(value_cols["created_at"], RawSql) + assert isinstance(value_cols["updated_at"], RawSql) + + def test_invalid_role_short_circuits_before_touching_executor(self, role_service, sql_executor_mock): + # Validation runs before the executor — a bad role must not + # leak a partial write attempt. + with pytest.raises(ValueError, match="Invalid role"): + role_service.create_mapping(role="not-a-real-role", group_name="g", user_email="u@x") + sql_executor_mock.upsert_with_audit.assert_not_called() + sql_executor_mock.execute.assert_not_called() diff --git a/app/tests/test_rules_catalog_service.py b/app/tests/test_rules_catalog_service.py index e8dbb3b90..398375562 100644 --- a/app/tests/test_rules_catalog_service.py +++ b/app/tests/test_rules_catalog_service.py @@ -220,10 +220,11 @@ def test_returns_empty_when_table_has_no_rules(self, svc, sql_executor_mock): def test_detects_existing_match(self, svc, sql_executor_mock): existing_check = {"check": {"function": "is_not_null", "arguments": {"column": "x"}}} # _row_to_entry expects 10 columns from _SELECT_COLS — emulate that. + # Column 1 is to_json(`check`): a bare check object, not an array. sql_executor_mock.query.return_value = [ ( "a.b.c", - json.dumps([existing_check]), + json.dumps(existing_check), 1, "draft", "u@x", @@ -243,7 +244,7 @@ def test_rejected_rules_do_not_count(self, svc, sql_executor_mock): sql_executor_mock.query.return_value = [ ( "a.b.c", - json.dumps([existing_check]), + json.dumps(existing_check), 1, "rejected", # rejected → not a duplicate "u@x", @@ -261,7 +262,7 @@ def test_exclude_rule_id_skips_self(self, svc, sql_executor_mock): sql_executor_mock.query.return_value = [ ( "a.b.c", - json.dumps([existing_check]), + json.dumps(existing_check), 1, "approved", "u@x", diff --git a/app/tests/test_scheduler_service.py b/app/tests/test_scheduler_service.py index 79a30447a..8c3190ff3 100644 --- a/app/tests/test_scheduler_service.py +++ b/app/tests/test_scheduler_service.py @@ -13,7 +13,6 @@ from __future__ import annotations from datetime import datetime, timedelta, timezone -from unittest.mock import MagicMock import pytest @@ -27,22 +26,45 @@ ) -def _make_service() -> SchedulerService: - """Construct a SchedulerService without invoking ``__init__``. +# --------------------------------------------------------------------------- +# Fixture +# +# The shared :func:`make_scheduler` factory (in ``conftest.py``) routes +# through the real ``SchedulerService(...)`` constructor and injects mocks +# via the public ``oltp_sql=`` parameter. We use ``distinct_sql=True`` and +# ``distinct_tmp_sql=True`` because the view-GC tests need to inspect +# ``_sql`` (in-use cross-check) and ``_tmp_sql`` (DROP VIEW execution) as +# independently addressable mocks. The previous ``object.__new__`` + +# private-attribute pattern coupled the tests to internal field names; the +# helper makes the rename surface (``oltp_sql=`` constructor parameter) +# the test contract instead. +# --------------------------------------------------------------------------- - The real constructor builds two ``SqlExecutor`` instances (which - requires a ``WorkspaceClient`` and a warehouse id), neither of which - is needed for these tests. We bypass it via ``object.__new__`` and - set the handful of attributes the methods we test actually read. + +@pytest.fixture +def gc_scheduler(make_scheduler): + """Construct a scheduler with distinct ``_sql`` and ``_tmp_sql`` mocks. + + Returns ``(svc, mocks)``. ``mocks.sql`` is the analytical Delta + executor (for the in-use cross-check); ``mocks.tmp`` is the tmp- + schema executor (for SHOW VIEWS / DROP VIEW). ``_next_view_gc_at`` + starts at year 2099 (the make_scheduler factory schedules a real + next-Saturday value; tests that exercise the GC schedule + explicitly override this). """ - svc = object.__new__(SchedulerService) - svc._catalog = "main" - svc._schema = "dqx" - svc._tmp_schema = "dqx_tmp" - svc._sql = MagicMock(name="SqlExecutor (main schema)") - svc._tmp_sql = MagicMock(name="SqlExecutor (tmp schema)") + svc, mocks = make_scheduler( + catalog="main", + schema="dqx", + tmp_schema="dqx_tmp", + distinct_sql=True, + distinct_tmp_sql=True, + ) + # Park the next-fire well into the future so tests opt into firing + # by overriding ``_next_view_gc_at`` explicitly. (This one field + # has no constructor seam because in production it's computed from + # the current wall clock — the override is an inherent test seam.) svc._next_view_gc_at = datetime(2099, 1, 1, tzinfo=timezone.utc) - return svc + return svc, mocks # --------------------------------------------------------------------------- @@ -116,75 +138,75 @@ def test_microseconds_normalised(self): class TestGcOrphanViews: - def test_no_candidates_returns_quickly(self): - svc = _make_service() - svc._tmp_sql.query.return_value = [] + def test_no_candidates_returns_quickly(self, gc_scheduler): + svc, mocks = gc_scheduler + mocks.tmp.query.return_value = [] svc._gc_orphan_views() # No DROP VIEW issued. - svc._tmp_sql.execute.assert_not_called() + mocks.tmp.execute.assert_not_called() - def test_drops_eligible_views_skipping_in_use(self): - svc = _make_service() + def test_drops_eligible_views_skipping_in_use(self, gc_scheduler): + svc, mocks = gc_scheduler # Three candidates: tmp_view_aaaa1111, tmp_view_bbbb2222, tmp_view_cccc3333 - svc._tmp_sql.query.return_value = [ + mocks.tmp.query.return_value = [ ("tmp_view_aaaa1111",), ("tmp_view_bbbb2222",), ("tmp_view_cccc3333",), ] # tmp_view_bbbb2222 is currently in use by a RUNNING dryrun - svc._sql.query.return_value = [("main.dqx_tmp.tmp_view_bbbb2222",)] + mocks.sql.query.return_value = [("main.dqx_tmp.tmp_view_bbbb2222",)] svc._gc_orphan_views() - executed = [c.args[0] for c in svc._tmp_sql.execute.call_args_list] + executed = [c.args[0] for c in mocks.tmp.execute.call_args_list] assert any("tmp_view_aaaa1111" in s for s in executed) assert any("tmp_view_cccc3333" in s for s in executed) assert not any("tmp_view_bbbb2222" in s for s in executed) - def test_invalid_view_names_are_filtered_out(self): - svc = _make_service() + def test_invalid_view_names_are_filtered_out(self, gc_scheduler): + svc, mocks = gc_scheduler # The regex requires tmp_view_<8-32 lowercase hex>; everything else # is skipped to keep the GC laser-focused. - svc._tmp_sql.query.return_value = [ + mocks.tmp.query.return_value = [ ("tmp_view_DEADBEEF",), # uppercase hex → not matched ("tmp_view_short",), # too short / non-hex ("tmp_view_aaaa1111",), # valid ("real_user_view",), # not the right prefix ("tmp_view_zzzzzzzzzz",), # non-hex characters ] - svc._sql.query.return_value = [] + mocks.sql.query.return_value = [] svc._gc_orphan_views() - executed = [c.args[0] for c in svc._tmp_sql.execute.call_args_list] + executed = [c.args[0] for c in mocks.tmp.execute.call_args_list] assert len(executed) == 1 assert "tmp_view_aaaa1111" in executed[0] - def test_in_use_query_failure_does_not_abort_cleanup(self): - svc = _make_service() - svc._tmp_sql.query.return_value = [("tmp_view_aaaa1111",)] - svc._sql.query.side_effect = RuntimeError("warehouse cold") + def test_in_use_query_failure_does_not_abort_cleanup(self, gc_scheduler): + svc, mocks = gc_scheduler + mocks.tmp.query.return_value = [("tmp_view_aaaa1111",)] + mocks.sql.query.side_effect = RuntimeError("warehouse cold") # Should still drop the candidate using age-only criteria. svc._gc_orphan_views() - svc._tmp_sql.execute.assert_called_once() + mocks.tmp.execute.assert_called_once() - def test_individual_drop_failure_continues_loop(self): - svc = _make_service() - svc._tmp_sql.query.return_value = [ + def test_individual_drop_failure_continues_loop(self, gc_scheduler): + svc, mocks = gc_scheduler + mocks.tmp.query.return_value = [ ("tmp_view_aaaa1111",), ("tmp_view_bbbb2222",), ] - svc._sql.query.return_value = [] - svc._tmp_sql.execute.side_effect = [RuntimeError("locked"), None] + mocks.sql.query.return_value = [] + mocks.tmp.execute.side_effect = [RuntimeError("locked"), None] # Both DROPs are attempted even when the first one blows up. svc._gc_orphan_views() - assert svc._tmp_sql.execute.call_count == 2 + assert mocks.tmp.execute.call_count == 2 - def test_list_query_failure_returns_silently(self): - svc = _make_service() - svc._tmp_sql.query.side_effect = RuntimeError("warehouse down") + def test_list_query_failure_returns_silently(self, gc_scheduler): + svc, mocks = gc_scheduler + mocks.tmp.query.side_effect = RuntimeError("warehouse down") svc._gc_orphan_views() - svc._tmp_sql.execute.assert_not_called() + mocks.tmp.execute.assert_not_called() # Must not have attempted the in-use cross-check either. - svc._sql.query.assert_not_called() + mocks.sql.query.assert_not_called() # --------------------------------------------------------------------------- @@ -194,8 +216,8 @@ def test_list_query_failure_returns_silently(self): class TestMaybeGcOrphanViews: @pytest.mark.asyncio - async def test_skips_when_not_yet_due(self): - svc = _make_service() + async def test_skips_when_not_yet_due(self, gc_scheduler): + svc, mocks = gc_scheduler future = datetime(2099, 1, 1, tzinfo=timezone.utc) svc._next_view_gc_at = future @@ -203,29 +225,29 @@ async def test_skips_when_not_yet_due(self): # Schedule untouched, no SQL run. assert svc._next_view_gc_at == future - svc._tmp_sql.query.assert_not_called() + mocks.tmp.query.assert_not_called() @pytest.mark.asyncio - async def test_advances_schedule_before_running(self): - svc = _make_service() + async def test_advances_schedule_before_running(self, gc_scheduler): + svc, mocks = gc_scheduler # Schedule a fire on 2026-05-02 01:00 UTC. due = datetime(2026, 5, 2, 1, 0, tzinfo=timezone.utc) svc._next_view_gc_at = due # No candidates → no work to do, but the schedule MUST advance. - svc._tmp_sql.query.return_value = [] + mocks.tmp.query.return_value = [] await svc._maybe_gc_orphan_views(due + timedelta(seconds=1)) # Next fire is the following Saturday at 01:00. assert svc._next_view_gc_at == datetime(2026, 5, 9, 1, 0, tzinfo=timezone.utc) @pytest.mark.asyncio - async def test_failure_in_gc_does_not_propagate(self): - svc = _make_service() + async def test_failure_in_gc_does_not_propagate(self, gc_scheduler): + svc, mocks = gc_scheduler due = datetime(2026, 5, 2, 1, 0, tzinfo=timezone.utc) svc._next_view_gc_at = due - svc._tmp_sql.query.side_effect = RuntimeError("nope") + mocks.tmp.query.side_effect = RuntimeError("nope") # Must not raise; just log and reschedule. await svc._maybe_gc_orphan_views(due + timedelta(minutes=1)) assert svc._next_view_gc_at > due diff --git a/app/uv.lock b/app/uv.lock index b4f399763..553feb804 100644 --- a/app/uv.lock +++ b/app/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-04-27T20:27:01.708704Z" +exclude-newer = "2026-05-04T20:36:45.499629Z" exclude-newer-span = "P7D" [options.exclude-newer-package] @@ -900,6 +900,7 @@ dependencies = [ { name = "databricks-sql-connector", extra = ["pyarrow"] }, { name = "fastapi" }, { name = "openpyxl" }, + { name = "psycopg", extra = ["binary", "pool"] }, { name = "pydantic-settings" }, { name = "uvicorn" }, ] @@ -924,6 +925,7 @@ requires-dist = [ { name = "databricks-sql-connector", extras = ["pyarrow"], specifier = "==4.2.5" }, { name = "fastapi", specifier = "~=0.119" }, { name = "openpyxl", specifier = ">=3.1" }, + { name = "psycopg", extras = ["binary", "pool"], specifier = ">=3.2" }, { name = "pydantic-settings", specifier = "~=2.11" }, { name = "uvicorn", specifier = "~=0.37" }, ] @@ -3270,6 +3272,90 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c9/ad/33b2ccec09bf96c2b2ef3f9a6f66baac8253d7565d8839e024a6b905d45d/psutil-7.1.3-cp37-abi3-win_arm64.whl", hash = "sha256:bd0d69cee829226a761e92f28140bec9a5ee9d5b4fb4b0cc589068dbfff559b1", size = 244608, upload-time = "2025-11-02T12:26:36.136Z" }, ] +[[package]] +name = "psycopg" +version = "3.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "tzdata", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d3/b6/379d0a960f8f435ec78720462fd94c4863e7a31237cf81bf76d0af5883bf/psycopg-3.3.3.tar.gz", hash = "sha256:5e9a47458b3c1583326513b2556a2a9473a1001a56c9efe9e587245b43148dd9", size = 165624, upload-time = "2026-02-18T16:52:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/5b/181e2e3becb7672b502f0ed7f16ed7352aca7c109cfb94cf3878a9186db9/psycopg-3.3.3-py3-none-any.whl", hash = "sha256:f96525a72bcfade6584ab17e89de415ff360748c766f0106959144dcbb38c698", size = 212768, upload-time = "2026-02-18T16:46:27.365Z" }, +] + +[package.optional-dependencies] +binary = [ + { name = "psycopg-binary", marker = "implementation_name != 'pypy'" }, +] +pool = [ + { name = "psycopg-pool" }, +] + +[[package]] +name = "psycopg-binary" +version = "3.3.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/be/c0/b389119dd754483d316805260f3e73cdcad97925839107cc7a296f6132b1/psycopg_binary-3.3.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a89bb9ee11177b2995d87186b1d9fa892d8ea725e85eab28c6525e4cc14ee048", size = 4609740, upload-time = "2026-02-18T16:47:51.093Z" }, + { url = "https://files.pythonhosted.org/packages/cf/e3/9976eef20f61840285174d360da4c820a311ab39d6b82fa09fbb545be825/psycopg_binary-3.3.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9f7d0cf072c6fbac3795b08c98ef9ea013f11db609659dcfc6b1f6cc31f9e181", size = 4676837, upload-time = "2026-02-18T16:47:55.523Z" }, + { url = "https://files.pythonhosted.org/packages/9f/f2/d28ba2f7404fd7f68d41e8a11df86313bd646258244cb12a8dd83b868a97/psycopg_binary-3.3.3-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:90eecd93073922f085967f3ed3a98ba8c325cbbc8c1a204e300282abd2369e13", size = 5497070, upload-time = "2026-02-18T16:47:59.929Z" }, + { url = "https://files.pythonhosted.org/packages/de/2f/6c5c54b815edeb30a281cfcea96dc93b3bb6be939aea022f00cab7aa1420/psycopg_binary-3.3.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:dac7ee2f88b4d7bb12837989ca354c38d400eeb21bce3b73dac02622f0a3c8d6", size = 5172410, upload-time = "2026-02-18T16:48:05.665Z" }, + { url = "https://files.pythonhosted.org/packages/51/75/8206c7008b57de03c1ada46bd3110cc3743f3fd9ed52031c4601401d766d/psycopg_binary-3.3.3-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b62cf8784eb6d35beaee1056d54caf94ec6ecf2b7552395e305518ab61eb8fd2", size = 6763408, upload-time = "2026-02-18T16:48:13.541Z" }, + { url = "https://files.pythonhosted.org/packages/d4/5a/ea1641a1e6c8c8b3454b0fcb43c3045133a8b703e6e824fae134088e63bd/psycopg_binary-3.3.3-cp311-cp311-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a39f34c9b18e8f6794cca17bfbcd64572ca2482318db644268049f8c738f35a6", size = 5006255, upload-time = "2026-02-18T16:48:22.176Z" }, + { url = "https://files.pythonhosted.org/packages/aa/fb/538df099bf55ae1637d52d7ccb6b9620b535a40f4c733897ac2b7bb9e14c/psycopg_binary-3.3.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:883d68d48ca9ff3cb3d10c5fdebea02c79b48eecacdddbf7cce6e7cdbdc216b8", size = 4532694, upload-time = "2026-02-18T16:48:27.338Z" }, + { url = "https://files.pythonhosted.org/packages/a1/d1/00780c0e187ea3c13dfc53bd7060654b2232cd30df562aac91a5f1c545ac/psycopg_binary-3.3.3-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:cab7bc3d288d37a80aa8c0820033250c95e40b1c2b5c57cf59827b19c2a8b69d", size = 4222833, upload-time = "2026-02-18T16:48:31.221Z" }, + { url = "https://files.pythonhosted.org/packages/7a/34/a07f1ff713c51d64dc9f19f2c32be80299a2055d5d109d5853662b922cb4/psycopg_binary-3.3.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:56c767007ca959ca32f796b42379fc7e1ae2ed085d29f20b05b3fc394f3715cc", size = 3952818, upload-time = "2026-02-18T16:48:35.869Z" }, + { url = "https://files.pythonhosted.org/packages/d3/67/d33f268a7759b4445f3c9b5a181039b01af8c8263c865c1be7a6444d4749/psycopg_binary-3.3.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:da2f331a01af232259a21573a01338530c6016dcfad74626c01330535bcd8628", size = 4258061, upload-time = "2026-02-18T16:48:41.365Z" }, + { url = "https://files.pythonhosted.org/packages/b4/3b/0d8d2c5e8e29ccc07d28c8af38445d9d9abcd238d590186cac82ee71fc84/psycopg_binary-3.3.3-cp311-cp311-win_amd64.whl", hash = "sha256:19f93235ece6dbfc4036b5e4f6d8b13f0b8f2b3eeb8b0bd2936d406991bcdd40", size = 3558915, upload-time = "2026-02-18T16:48:46.679Z" }, + { url = "https://files.pythonhosted.org/packages/90/15/021be5c0cbc5b7c1ab46e91cc3434eb42569f79a0592e67b8d25e66d844d/psycopg_binary-3.3.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6698dbab5bcef8fdb570fc9d35fd9ac52041771bfcfe6fd0fc5f5c4e36f1e99d", size = 4591170, upload-time = "2026-02-18T16:48:55.594Z" }, + { url = "https://files.pythonhosted.org/packages/f1/54/a60211c346c9a2f8c6b272b5f2bbe21f6e11800ce7f61e99ba75cf8b63e1/psycopg_binary-3.3.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:329ff393441e75f10b673ae99ab45276887993d49e65f141da20d915c05aafd8", size = 4670009, upload-time = "2026-02-18T16:49:03.608Z" }, + { url = "https://files.pythonhosted.org/packages/c1/53/ac7c18671347c553362aadbf65f92786eef9540676ca24114cc02f5be405/psycopg_binary-3.3.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:eb072949b8ebf4082ae24289a2b0fd724da9adc8f22743409d6fd718ddb379df", size = 5469735, upload-time = "2026-02-18T16:49:10.128Z" }, + { url = "https://files.pythonhosted.org/packages/7f/c3/4f4e040902b82a344eff1c736cde2f2720f127fe939c7e7565706f96dd44/psycopg_binary-3.3.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:263a24f39f26e19ed7fc982d7859a36f17841b05bebad3eb47bb9cd2dd785351", size = 5152919, upload-time = "2026-02-18T16:49:16.335Z" }, + { url = "https://files.pythonhosted.org/packages/0c/e7/d929679c6a5c212bcf738806c7c89f5b3d0919f2e1685a0e08d6ff877945/psycopg_binary-3.3.3-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5152d50798c2fa5bd9b68ec68eb68a1b71b95126c1d70adaa1a08cd5eefdc23d", size = 6738785, upload-time = "2026-02-18T16:49:22.687Z" }, + { url = "https://files.pythonhosted.org/packages/69/b0/09703aeb69a9443d232d7b5318d58742e8ca51ff79f90ffe6b88f1db45e7/psycopg_binary-3.3.3-cp312-cp312-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9d6a1e56dd267848edb824dbeb08cf5bac649e02ee0b03ba883ba3f4f0bd54f2", size = 4979008, upload-time = "2026-02-18T16:49:27.313Z" }, + { url = "https://files.pythonhosted.org/packages/cc/a6/e662558b793c6e13a7473b970fee327d635270e41eded3090ef14045a6a5/psycopg_binary-3.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:73eaaf4bb04709f545606c1db2f65f4000e8a04cdbf3e00d165a23004692093e", size = 4508255, upload-time = "2026-02-18T16:49:31.575Z" }, + { url = "https://files.pythonhosted.org/packages/5f/7f/0f8b2e1d5e0093921b6f324a948a5c740c1447fbb45e97acaf50241d0f39/psycopg_binary-3.3.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:162e5675efb4704192411eaf8e00d07f7960b679cd3306e7efb120bb8d9456cc", size = 4189166, upload-time = "2026-02-18T16:49:35.801Z" }, + { url = "https://files.pythonhosted.org/packages/92/ec/ce2e91c33bc8d10b00c87e2f6b0fb570641a6a60042d6a9ae35658a3a797/psycopg_binary-3.3.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:fab6b5e37715885c69f5d091f6ff229be71e235f272ebaa35158d5a46fd548a0", size = 3924544, upload-time = "2026-02-18T16:49:41.129Z" }, + { url = "https://files.pythonhosted.org/packages/c5/2f/7718141485f73a924205af60041c392938852aa447a94c8cbd222ff389a1/psycopg_binary-3.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a4aab31bd6d1057f287c96c0effca3a25584eb9cc702f282ecb96ded7814e830", size = 4235297, upload-time = "2026-02-18T16:49:46.726Z" }, + { url = "https://files.pythonhosted.org/packages/57/f9/1add717e2643a003bbde31b1b220172e64fbc0cb09f06429820c9173f7fc/psycopg_binary-3.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:59aa31fe11a0e1d1bcc2ce37ed35fe2ac84cd65bb9036d049b1a1c39064d0f14", size = 3547659, upload-time = "2026-02-18T16:49:52.999Z" }, + { url = "https://files.pythonhosted.org/packages/03/0a/cac9fdf1df16a269ba0e5f0f06cac61f826c94cadb39df028cdfe19d3a33/psycopg_binary-3.3.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:05f32239aec25c5fb15f7948cffdc2dc0dac098e48b80a140e4ba32b572a2e7d", size = 4590414, upload-time = "2026-02-18T16:50:01.441Z" }, + { url = "https://files.pythonhosted.org/packages/9c/c0/d8f8508fbf440edbc0099b1abff33003cd80c9e66eb3a1e78834e3fb4fb9/psycopg_binary-3.3.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7c84f9d214f2d1de2fafebc17fa68ac3f6561a59e291553dfc45ad299f4898c1", size = 4669021, upload-time = "2026-02-18T16:50:08.803Z" }, + { url = "https://files.pythonhosted.org/packages/04/05/097016b77e343b4568feddf12c72171fc513acef9a4214d21b9478569068/psycopg_binary-3.3.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:e77957d2ba17cada11be09a5066d93026cdb61ada7c8893101d7fe1c6e1f3925", size = 5467453, upload-time = "2026-02-18T16:50:14.985Z" }, + { url = "https://files.pythonhosted.org/packages/91/23/73244e5feb55b5ca109cede6e97f32ef45189f0fdac4c80d75c99862729d/psycopg_binary-3.3.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:42961609ac07c232a427da7c87a468d3c82fee6762c220f38e37cfdacb2b178d", size = 5151135, upload-time = "2026-02-18T16:50:24.82Z" }, + { url = "https://files.pythonhosted.org/packages/11/49/5309473b9803b207682095201d8708bbc7842ddf3f192488a69204e36455/psycopg_binary-3.3.3-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ae07a3114313dd91fce686cab2f4c44af094398519af0e0f854bc707e1aeedf1", size = 6737315, upload-time = "2026-02-18T16:50:35.106Z" }, + { url = "https://files.pythonhosted.org/packages/d4/5d/03abe74ef34d460b33c4d9662bf6ec1dd38888324323c1a1752133c10377/psycopg_binary-3.3.3-cp313-cp313-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d257c58d7b36a621dcce1d01476ad8b60f12d80eb1406aee4cf796f88b2ae482", size = 4979783, upload-time = "2026-02-18T16:50:42.067Z" }, + { url = "https://files.pythonhosted.org/packages/f0/6c/3fbf8e604e15f2f3752900434046c00c90bb8764305a1b81112bff30ba24/psycopg_binary-3.3.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:07c7211f9327d522c9c47560cae00a4ecf6687f4e02d779d035dd3177b41cb12", size = 4509023, upload-time = "2026-02-18T16:50:50.116Z" }, + { url = "https://files.pythonhosted.org/packages/9c/6b/1a06b43b7c7af756c80b67eac8bfaa51d77e68635a8a8d246e4f0bb7604a/psycopg_binary-3.3.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:8e7e9eca9b363dbedeceeadd8be97149d2499081f3c52d141d7cd1f395a91f83", size = 4185874, upload-time = "2026-02-18T16:50:55.97Z" }, + { url = "https://files.pythonhosted.org/packages/2b/d3/bf49e3dcaadba510170c8d111e5e69e5ae3f981c1554c5bb71c75ce354bb/psycopg_binary-3.3.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:cb85b1d5702877c16f28d7b92ba030c1f49ebcc9b87d03d8c10bf45a2f1c7508", size = 3925668, upload-time = "2026-02-18T16:51:03.299Z" }, + { url = "https://files.pythonhosted.org/packages/f8/92/0aac830ed6a944fe334404e1687a074e4215630725753f0e3e9a9a595b62/psycopg_binary-3.3.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4d4606c84d04b80f9138d72f1e28c6c02dc5ae0c7b8f3f8aaf89c681ce1cd1b1", size = 4234973, upload-time = "2026-02-18T16:51:09.097Z" }, + { url = "https://files.pythonhosted.org/packages/2e/96/102244653ee5a143ece5afe33f00f52fe64e389dfce8dbc87580c6d70d3d/psycopg_binary-3.3.3-cp313-cp313-win_amd64.whl", hash = "sha256:74eae563166ebf74e8d950ff359be037b85723d99ca83f57d9b244a871d6c13b", size = 3551342, upload-time = "2026-02-18T16:51:13.892Z" }, + { url = "https://files.pythonhosted.org/packages/a2/71/7a57e5b12275fe7e7d84d54113f0226080423a869118419c9106c083a21c/psycopg_binary-3.3.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:497852c5eaf1f0c2d88ab74a64a8097c099deac0c71de1cbcf18659a8a04a4b2", size = 4607368, upload-time = "2026-02-18T16:51:19.295Z" }, + { url = "https://files.pythonhosted.org/packages/c7/04/cb834f120f2b2c10d4003515ef9ca9d688115b9431735e3936ae48549af8/psycopg_binary-3.3.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:258d1ea53464d29768bf25930f43291949f4c7becc706f6e220c515a63a24edd", size = 4687047, upload-time = "2026-02-18T16:51:23.84Z" }, + { url = "https://files.pythonhosted.org/packages/40/e9/47a69692d3da9704468041aa5ed3ad6fc7f6bb1a5ae788d261a26bbca6c7/psycopg_binary-3.3.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:111c59897a452196116db12e7f608da472fbff000693a21040e35fc978b23430", size = 5487096, upload-time = "2026-02-18T16:51:29.645Z" }, + { url = "https://files.pythonhosted.org/packages/0b/b6/0e0dd6a2f802864a4ae3dbadf4ec620f05e3904c7842b326aafc43e5f464/psycopg_binary-3.3.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:17bb6600e2455993946385249a3c3d0af52cd70c1c1cdbf712e9d696d0b0bf1b", size = 5168720, upload-time = "2026-02-18T16:51:36.499Z" }, + { url = "https://files.pythonhosted.org/packages/6f/0d/977af38ac19a6b55d22dff508bd743fd7c1901e1b73657e7937c7cccb0a3/psycopg_binary-3.3.3-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:642050398583d61c9856210568eb09a8e4f2fe8224bf3be21b67a370e677eead", size = 6762076, upload-time = "2026-02-18T16:51:43.167Z" }, + { url = "https://files.pythonhosted.org/packages/34/40/912a39d48322cf86895c0eaf2d5b95cb899402443faefd4b09abbba6b6e1/psycopg_binary-3.3.3-cp314-cp314-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:533efe6dc3a7cba5e2a84e38970786bb966306863e45f3db152007e9f48638a6", size = 4997623, upload-time = "2026-02-18T16:51:47.707Z" }, + { url = "https://files.pythonhosted.org/packages/98/0c/c14d0e259c65dc7be854d926993f151077887391d5a081118907a9d89603/psycopg_binary-3.3.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:5958dbf28b77ce2033482f6cb9ef04d43f5d8f4b7636e6963d5626f000efb23e", size = 4532096, upload-time = "2026-02-18T16:51:51.421Z" }, + { url = "https://files.pythonhosted.org/packages/39/21/8b7c50a194cfca6ea0fd4d1f276158307785775426e90700ab2eba5cd623/psycopg_binary-3.3.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:a6af77b6626ce92b5817bf294b4d45ec1a6161dba80fc2d82cdffdd6814fd023", size = 4208884, upload-time = "2026-02-18T16:51:57.336Z" }, + { url = "https://files.pythonhosted.org/packages/c7/2c/a4981bf42cf30ebba0424971d7ce70a222ae9b82594c42fc3f2105d7b525/psycopg_binary-3.3.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:47f06fcbe8542b4d96d7392c476a74ada521c5aebdb41c3c0155f6595fc14c8d", size = 3944542, upload-time = "2026-02-18T16:52:04.266Z" }, + { url = "https://files.pythonhosted.org/packages/60/e9/b7c29b56aa0b85a4e0c4d89db691c1ceef08f46a356369144430c155a2f5/psycopg_binary-3.3.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e7800e6c6b5dc4b0ca7cc7370f770f53ac83886b76afda0848065a674231e856", size = 4254339, upload-time = "2026-02-18T16:52:10.444Z" }, + { url = "https://files.pythonhosted.org/packages/98/5a/291d89f44d3820fffb7a04ebc8f3ef5dda4f542f44a5daea0c55a84abf45/psycopg_binary-3.3.3-cp314-cp314-win_amd64.whl", hash = "sha256:165f22ab5a9513a3d7425ffb7fcc7955ed8ccaeef6d37e369d6cc1dff1582383", size = 3652796, upload-time = "2026-02-18T16:52:14.02Z" }, +] + +[[package]] +name = "psycopg-pool" +version = "3.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/56/9a/9470d013d0d50af0da9c4251614aeb3c1823635cab3edc211e3839db0bcf/psycopg_pool-3.3.0.tar.gz", hash = "sha256:fa115eb2860bd88fce1717d75611f41490dec6135efb619611142b24da3f6db5", size = 31606, upload-time = "2025-12-01T11:34:33.11Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/c3/26b8a0908a9db249de3b4169692e1c7c19048a9bc41a4d3209cee7dbb758/psycopg_pool-3.3.0-py3-none-any.whl", hash = "sha256:2e44329155c410b5e8666372db44276a8b1ebd8c90f1c3026ebba40d4bc81063", size = 39995, upload-time = "2025-12-01T11:34:29.761Z" }, +] + [[package]] name = "py4j" version = "0.10.9.7" diff --git a/app/yarn.lock b/app/yarn.lock index ac1bf5216..c1753d1ed 100644 --- a/app/yarn.lock +++ b/app/yarn.lock @@ -43,8 +43,8 @@ picocolors "^1.1.1" "@babel/compat-data@^7.28.6": - version "7.29.0" - integrity sha512-T1NCJqT/j9+cn8fvkt7jtwbLBfLC/1y1c7NtCeXFRgzGTsafi68MRv8yzkYSapBnFA6L3U2VSc02ciDzoAJhJg== + version "7.29.3" + integrity sha512-LIVqM46zQWZhj17qA8wb4nW/ixr2y1Nw+r1etiAWgRM6U1IqP+LNhL1yg440jYZR72jCWcWbLWzIosH+uP1fqg== "@babel/core@^7.23.7", "@babel/core@^7.28.5", "@babel/core@^7.29.0": version "7.29.0" @@ -129,8 +129,8 @@ "@babel/types" "^7.29.0" "@babel/parser@^7.1.0", "@babel/parser@^7.20.7", "@babel/parser@^7.23.6", "@babel/parser@^7.28.5", "@babel/parser@^7.28.6", "@babel/parser@^7.29.0": - version "7.29.2" - integrity sha512-4GgRzy/+fsBa72/RZVJmGKPmZu9Byn8o4MoLpmNe1m8ZfYnz5emHLQz3U4gLud6Zwl0RZIcgiLD7Uq7ySFuDLA== + version "7.29.3" + integrity sha512-b3ctpQwp+PROvU/cttc4OYl4MzfJUWy6FZg+PMXfzmt/+39iHVF0sDfqay8TQM3JA2EUOyKcFZt75jWriQijsA== dependencies: "@babel/types" "^7.29.0" @@ -969,105 +969,105 @@ version "1.0.0-rc.3" integrity sha512-eybk3TjzzzV97Dlj5c+XrBFW57eTNhzod66y9HrBlzJ6NsCrWCp/2kaPS3K9wJmurBC0Tdw4yPjXKZqlznim3Q== -"@rollup/rollup-android-arm-eabi@4.60.2": - version "4.60.2" - integrity sha512-dnlp69efPPg6Uaw2dVqzWRfAWRnYVb1XJ8CyyhIbZeaq4CA5/mLeZ1IEt9QqQxmbdvagjLIm2ZL8BxXv5lH4Yw== +"@rollup/rollup-android-arm-eabi@4.60.3": + version "4.60.3" + integrity sha512-x35CNW/ANXG3hE/EZpRU8MXX1JDN86hBb2wMGAtltkz7pc6cxgjpy1OMMfDosOQ+2hWqIkag/fGok1Yady9nGw== -"@rollup/rollup-android-arm64@4.60.2": - version "4.60.2" - integrity sha512-OqZTwDRDchGRHHm/hwLOL7uVPB9aUvI0am/eQuWMNyFHf5PSEQmyEeYYheA0EPPKUO/l0uigCp+iaTjoLjVoHg== +"@rollup/rollup-android-arm64@4.60.3": + version "4.60.3" + integrity sha512-xw3xtkDApIOGayehp2+Rz4zimfkaX65r4t47iy+ymQB2G4iJCBBfj0ogVg5jpvjpn8UWn/+q9tprxleYeNp3Hw== -"@rollup/rollup-darwin-arm64@4.60.2": - version "4.60.2" - integrity sha512-UwRE7CGpvSVEQS8gUMBe1uADWjNnVgP3Iusyda1nSRwNDCsRjnGc7w6El6WLQsXmZTbLZx9cecegumcitNfpmA== +"@rollup/rollup-darwin-arm64@4.60.3": + version "4.60.3" + integrity sha512-vo6Y5Qfpx7/5EaamIwi0WqW2+zfiusVihKatLvtN1VFVy3D13uERk/6gZLU1UiHRL6fDXqj/ELIeVRGnvcTE1g== -"@rollup/rollup-darwin-x64@4.60.2": - version "4.60.2" - integrity sha512-gjEtURKLCC5VXm1I+2i1u9OhxFsKAQJKTVB8WvDAHF+oZlq0GTVFOlTlO1q3AlCTE/DF32c16ESvfgqR7343/g== +"@rollup/rollup-darwin-x64@4.60.3": + version "4.60.3" + integrity sha512-D+0QGcZhBzTN82weOnsSlY7V7+RMmPuF1CkbxyMAGE8+ZHeUjyb76ZiWmBlCu//AQQONvxcqRbwZTajZKqjuOw== -"@rollup/rollup-freebsd-arm64@4.60.2": - version "4.60.2" - integrity sha512-Bcl6CYDeAgE70cqZaMojOi/eK63h5Me97ZqAQoh77VPjMysA/4ORQBRGo3rRy45x4MzVlU9uZxs8Uwy7ZaKnBw== +"@rollup/rollup-freebsd-arm64@4.60.3": + version "4.60.3" + integrity sha512-6HnvHCT7fDyj6R0Ph7A6x8dQS/S38MClRWeDLqc0MdfWkxjiu1HSDYrdPhqSILzjTIC/pnXbbJbo+ft+gy/9hQ== -"@rollup/rollup-freebsd-x64@4.60.2": - version "4.60.2" - integrity sha512-LU+TPda3mAE2QB0/Hp5VyeKJivpC6+tlOXd1VMoXV/YFMvk/MNk5iXeBfB4MQGRWyOYVJ01625vjkr0Az98OJQ== +"@rollup/rollup-freebsd-x64@4.60.3": + version "4.60.3" + integrity sha512-KHLgC3WKlUYW3ShFKnnosZDOJ0xjg9zp7au3sIm2bs/tGBeC2ipmvRh/N7JKi0t9Ue20C0dpEshi8WUubg+cnA== -"@rollup/rollup-linux-arm-gnueabihf@4.60.2": - version "4.60.2" - integrity sha512-2QxQrM+KQ7DAW4o22j+XZ6RKdxjLD7BOWTP0Bv0tmjdyhXSsr2Ul1oJDQqh9Zf5qOwTuTc7Ek83mOFaKnodPjg== +"@rollup/rollup-linux-arm-gnueabihf@4.60.3": + version "4.60.3" + integrity sha512-DV6fJoxEYWJOvaZIsok7KrYl0tPvga5OZ2yvKHNNYyk/2roMLqQAbGhr78EQ5YhHpnhLKJD3S1WFusAkmUuV5g== -"@rollup/rollup-linux-arm-musleabihf@4.60.2": - version "4.60.2" - integrity sha512-TbziEu2DVsTEOPif2mKWkMeDMLoYjx95oESa9fkQQK7r/Orta0gnkcDpzwufEcAO2BLBsD7mZkXGFqEdMRRwfw== +"@rollup/rollup-linux-arm-musleabihf@4.60.3": + version "4.60.3" + integrity sha512-mQKoJAzvuOs6F+TZybQO4GOTSMUu7v0WdxEk24krQ/uUxXoPTtHjuaUuPmFhtBcM4K0ons8nrE3JyhTuCFtT/w== -"@rollup/rollup-linux-arm64-gnu@4.60.2": - version "4.60.2" - integrity sha512-bO/rVDiDUuM2YfuCUwZ1t1cP+/yqjqz+Xf2VtkdppefuOFS2OSeAfgafaHNkFn0t02hEyXngZkxtGqXcXwO8Rg== +"@rollup/rollup-linux-arm64-gnu@4.60.3": + version "4.60.3" + integrity sha512-Whjj2qoiJ6+OOJMGptTYazaJvjOJm+iKHpXQM1P3LzGjt7Ff++Tp7nH4N8J/BUA7R9IHfDyx4DJIflifwnbmIA== -"@rollup/rollup-linux-arm64-musl@4.60.2": - version "4.60.2" - integrity sha512-hr26p7e93Rl0Za+JwW7EAnwAvKkehh12BU1Llm9Ykiibg4uIr2rbpxG9WCf56GuvidlTG9KiiQT/TXT1yAWxTA== +"@rollup/rollup-linux-arm64-musl@4.60.3": + version "4.60.3" + integrity sha512-4YTNHKqGng5+yiZt3mg77nmyuCfmNfX4fPmyUapBcIk+BdwSwmCWGXOUxhXbBEkFHtoN5boLj/5NON+u5QC9tg== -"@rollup/rollup-linux-loong64-gnu@4.60.2": - version "4.60.2" - integrity sha512-pOjB/uSIyDt+ow3k/RcLvUAOGpysT2phDn7TTUB3n75SlIgZzM6NKAqlErPhoFU+npgY3/n+2HYIQVbF70P9/A== +"@rollup/rollup-linux-loong64-gnu@4.60.3": + version "4.60.3" + integrity sha512-SU3kNlhkpI4UqlUc2VXPGK9o886ZsSeGfMAX2ba2b8DKmMXq4AL7KUrkSWVbb7koVqx41Yczx6dx5PNargIrEA== -"@rollup/rollup-linux-loong64-musl@4.60.2": - version "4.60.2" - integrity sha512-2/w+q8jszv9Ww1c+6uJT3OwqhdmGP2/4T17cu8WuwyUuuaCDDJ2ojdyYwZzCxx0GcsZBhzi3HmH+J5pZNXnd+Q== +"@rollup/rollup-linux-loong64-musl@4.60.3": + version "4.60.3" + integrity sha512-6lDLl5h4TXpB1mTf2rQWnAk/LcXrx9vBfu/DT5TIPhvMhRWaZ5MxkIc8u4lJAmBo6klTe1ywXIUHFjylW505sg== -"@rollup/rollup-linux-ppc64-gnu@4.60.2": - version "4.60.2" - integrity sha512-11+aL5vKheYgczxtPVVRhdptAM2H7fcDR5Gw4/bTcteuZBlH4oP9f5s9zYO9aGZvoGeBpqXI/9TZZihZ609wKw== +"@rollup/rollup-linux-ppc64-gnu@4.60.3": + version "4.60.3" + integrity sha512-BMo8bOw8evlup/8G+cj5xWtPyp93xPdyoSN16Zy90Q2QZ0ZYRhCt6ZJSwbrRzG9HApFabjwj2p25TUPDWrhzqQ== -"@rollup/rollup-linux-ppc64-musl@4.60.2": - version "4.60.2" - integrity sha512-i16fokAGK46IVZuV8LIIwMdtqhin9hfYkCh8pf8iC3QU3LpwL+1FSFGej+O7l3E/AoknL6Dclh2oTdnRMpTzFQ== +"@rollup/rollup-linux-ppc64-musl@4.60.3": + version "4.60.3" + integrity sha512-E0L8X1dZN1/Rph+5VPF6Xj2G7JJvMACVXtamTJIDrVI44Y3K+G8gQaMEAavbqCGTa16InptiVrX6eM6pmJ+7qA== -"@rollup/rollup-linux-riscv64-gnu@4.60.2": - version "4.60.2" - integrity sha512-49FkKS6RGQoriDSK/6E2GkAsAuU5kETFCh7pG4yD/ylj9rKhTmO3elsnmBvRD4PgJPds5W2PkhC82aVwmUcJ7A== +"@rollup/rollup-linux-riscv64-gnu@4.60.3": + version "4.60.3" + integrity sha512-oZJ/WHaVfHUiRAtmTAeo3DcevNsVvH8mbvodjZy7D5QKvCefO371SiKRpxoDcCxB3PTRTLayWBkvmDQKTcX/sw== -"@rollup/rollup-linux-riscv64-musl@4.60.2": - version "4.60.2" - integrity sha512-mjYNkHPfGpUR00DuM1ZZIgs64Hpf4bWcz9Z41+4Q+pgDx73UwWdAYyf6EG/lRFldmdHHzgrYyge5akFUW0D3mQ== +"@rollup/rollup-linux-riscv64-musl@4.60.3": + version "4.60.3" + integrity sha512-Dhbyh7j9FybM3YaTgaHmVALwA8AkUwTPccyCQ79TG9AJUsMQqgN1DDEZNr4+QUfwiWvLDumW5vdwzoeUF+TNxQ== -"@rollup/rollup-linux-s390x-gnu@4.60.2": - version "4.60.2" - integrity sha512-ALyvJz965BQk8E9Al/JDKKDLH2kfKFLTGMlgkAbbYtZuJt9LU8DW3ZoDMCtQpXAltZxwBHevXz5u+gf0yA0YoA== +"@rollup/rollup-linux-s390x-gnu@4.60.3": + version "4.60.3" + integrity sha512-cJd1X5XhHHlltkaypz1UcWLA8AcoIi1aWhsvaWDskD1oz2eKCypnqvTQ8ykMNI0RSmm7NkTdSqSSD7zM0xa6Ig== -"@rollup/rollup-linux-x64-gnu@4.60.2": - version "4.60.2" - integrity sha512-UQjrkIdWrKI626Du8lCQ6MJp/6V1LAo2bOK9OTu4mSn8GGXIkPXk/Vsp4bLHCd9Z9Iz2OTEaokUE90VweJgIYQ== +"@rollup/rollup-linux-x64-gnu@4.60.3": + version "4.60.3" + integrity sha512-DAZDBHQfG2oQuhY7mc6I3/qB4LU2fQCjRvxbDwd/Jdvb9fypP4IJ4qmtu6lNjes6B531AI8cg1aKC2di97bUxA== -"@rollup/rollup-linux-x64-musl@4.60.2": - version "4.60.2" - integrity sha512-bTsRGj6VlSdn/XD4CGyzMnzaBs9bsRxy79eTqTCBsA8TMIEky7qg48aPkvJvFe1HyzQ5oMZdg7AnVlWQSKLTnw== +"@rollup/rollup-linux-x64-musl@4.60.3": + version "4.60.3" + integrity sha512-cRxsE8c13mZOh3vP+wLDxpQBRrOHDIGOWyDL93Sy0Ga8y515fBcC2pjUfFwUe5T7tqvTvWbCpg1URM/AXdWIXA== -"@rollup/rollup-openbsd-x64@4.60.2": - version "4.60.2" - integrity sha512-6d4Z3534xitaA1FcMWP7mQPq5zGwBmGbhphh2DwaA1aNIXUu3KTOfwrWpbwI4/Gr0uANo7NTtaykFyO2hPuFLg== +"@rollup/rollup-openbsd-x64@4.60.3": + version "4.60.3" + integrity sha512-QaWcIgRxqEdQdhJqW4DJctsH6HCmo5vHxY0krHSX4jMtOqfzC+dqDGuHM87bu4H8JBeibWx7jFz+h6/4C8wA5Q== -"@rollup/rollup-openharmony-arm64@4.60.2": - version "4.60.2" - integrity sha512-NetAg5iO2uN7eB8zE5qrZ3CSil+7IJt4WDFLcC75Ymywq1VZVD6qJ6EvNLjZ3rEm6gB7XW5JdT60c6MN35Z85Q== +"@rollup/rollup-openharmony-arm64@4.60.3": + version "4.60.3" + integrity sha512-AaXwSvUi3QIPtroAUw1t5yHGIyqKEXwH54WUocFolZhpGDruJcs8c+xPNDRn4XiQsS7MEwnYsHW2l0MBLDMkWg== -"@rollup/rollup-win32-arm64-msvc@4.60.2": - version "4.60.2" - integrity sha512-NCYhOotpgWZ5kdxCZsv6Iudx0wX8980Q/oW4pNFNihpBKsDbEA1zpkfxJGC0yugsUuyDZ7gL37dbzwhR0VI7pQ== +"@rollup/rollup-win32-arm64-msvc@4.60.3": + version "4.60.3" + integrity sha512-65LAKM/bAWDqKNEelHlcHvm2V+Vfb8C6INFxQXRHCvaVN1rJfwr4NvdP4FyzUaLqWfaCGaadf6UbTm8xJeYfEg== -"@rollup/rollup-win32-ia32-msvc@4.60.2": - version "4.60.2" - integrity sha512-RXsaOqXxfoUBQoOgvmmijVxJnW2IGB0eoMO7F8FAjaj0UTywUO/luSqimWBJn04WNgUkeNhh7fs7pESXajWmkg== +"@rollup/rollup-win32-ia32-msvc@4.60.3": + version "4.60.3" + integrity sha512-EEM2gyhBF5MFnI6vMKdX1LAosE627RGBzIoGMdLloPZkXrUN0Ckqgr2Qi8+J3zip/8NVVro3/FjB+tjhZUgUHA== -"@rollup/rollup-win32-x64-gnu@4.60.2": - version "4.60.2" - integrity sha512-qdAzEULD+/hzObedtmV6iBpdL5TIbKVztGiK7O3/KYSf+HIzU257+MX1EXJcyIiDbMAqmbwaufcYPvyRryeZtA== +"@rollup/rollup-win32-x64-gnu@4.60.3": + version "4.60.3" + integrity sha512-E5Eb5H/DpxaoXH++Qkv28RcUJboMopmdDUALBczvHMf7hNIxaDZqwY5lK12UK1BHacSmvupoEWGu+n993Z0y1A== -"@rollup/rollup-win32-x64-msvc@4.60.2": - version "4.60.2" - integrity sha512-Nd/SgG27WoA9e+/TdK74KnHz852TLa94ovOYySo/yMPuTmpckK/jIF2jSwS3g7ELSKXK13/cVdmg1Z/DaCWKxA== +"@rollup/rollup-win32-x64-msvc@4.60.3": + version "4.60.3" + integrity sha512-hPt/bgL5cE+Qp+/TPHBqptcAgPzgj46mPcg/16zNUmbQk0j+mOEQV/+Lqu8QRtDV3Ek95Q6FeFITpuhl6OTsAA== "@shikijs/engine-oniguruma@^3.23.0": version "3.23.0" @@ -1379,15 +1379,15 @@ version "1.161.6" integrity sha512-NaOGLRrddszbQj9upGat6HG/4TKvXLvu+osAIgfxPYA+eIvYKv8GKDJOrY2D3/U9MRnKfMWD7bU4jeD4xmqyIg== -"@tanstack/query-core@5.100.5": - version "5.100.5" - integrity sha512-t20KrhKkf0HXzqQkPbJ5erhFesup68BAbwFgYmTrS7bxMF7O5MdmL8jUkik4thsG7Hg00fblz30h6yF1d5TxGg== +"@tanstack/query-core@5.100.9": + version "5.100.9" + integrity sha512-SJSFw1S8+kQ0+knv/XGfrbocWoAlT7vDKsSImtLx3ZPQmEcR46hkDjLSvynSy25N8Ms4tIEini1FuBd5k7IscQ== "@tanstack/react-query@^5.90.5": - version "5.100.5" - integrity sha512-aNwj1mi2v2bQ9IxkyR1grLOUkv3BYWoykHy9KDyLNbjC3tsahbOHJibK+Wjtr1wRhG59/AvJhiJG5OlthaCgJA== + version "5.100.9" + integrity sha512-Oa44XkaI3kCNN6ME0KByU3xT3SEUNOMfZpHxL6+wFoTm+OeUFYHKdeYVe0aOXlRDm/f15sgLwEt2HDorIdW8+A== dependencies: - "@tanstack/query-core" "5.100.5" + "@tanstack/query-core" "5.100.9" "@tanstack/react-router-devtools@^1.133.36": version "1.166.13" @@ -1396,12 +1396,12 @@ "@tanstack/router-devtools-core" "1.167.3" "@tanstack/react-router@^1.133.36": - version "1.168.25" - integrity sha512-4U/E76dc+fYuLixjV1RLNfqrkQoexSL8MqGNpIHOodtvY3fMPGaALrvDVtBDQYBEU4z5r5fHaV6+kclWAVFP9A== + version "1.169.1" + integrity sha512-MBtQKSvac3OCcsSa6oBpDrrN90IV47I6Gtv05NxhbFVh+gVjtqvs6HSU4XM9+y5sHZPgS+35eArflX4vM8GEnQ== dependencies: "@tanstack/history" "1.161.6" "@tanstack/react-store" "^0.9.3" - "@tanstack/router-core" "1.168.17" + "@tanstack/router-core" "1.169.1" isbot "^5.1.22" "@tanstack/react-store@^0.9.3": @@ -1411,9 +1411,9 @@ "@tanstack/store" "0.9.3" use-sync-external-store "^1.6.0" -"@tanstack/router-core@1.168.17": - version "1.168.17" - integrity sha512-VDq7HCqRK3sdpxoETwYoTXTaYi+OVQC197g1fdzaiZBUmhntfjn+PQc15OzTqNNhf8Menk6r6ftmuphybMKdig== +"@tanstack/router-core@1.169.1": + version "1.169.1" + integrity sha512-x+2gIGKTTE1qAn7tLieGfrB5ciOviDmmi2ox9fAWUubRV+yTU5ruGFXocoCIWF+lB+SOtnHjo2E9BLSWyYoEmA== dependencies: "@tanstack/history" "1.161.6" cookie-es "^3.0.0" @@ -1427,12 +1427,12 @@ clsx "^2.1.1" goober "^2.1.16" -"@tanstack/router-generator@1.166.36": - version "1.166.36" - integrity sha512-ce8Sg+ONwdd483kXJBYhTcdIAjEwSlWUOkoLsgPdNUIfA05hdnd9JkNnM4X1OnzpFL8/+TBSMo4WYQp9CHhDPg== +"@tanstack/router-generator@1.166.39": + version "1.166.39" + integrity sha512-j2OW/UvpjM/DT9tHVmuhWW1k6UOezTRrPqBPZFFmIth0fY7iTPqK+Erqpo8r5yGTRGCbMvOS4sL3H2MldnIZew== dependencies: "@babel/types" "^7.28.5" - "@tanstack/router-core" "1.168.17" + "@tanstack/router-core" "1.169.1" "@tanstack/router-utils" "1.161.7" "@tanstack/virtual-file-routes" "1.161.7" jiti "^2.6.1" @@ -1441,8 +1441,8 @@ zod "^3.24.2" "@tanstack/router-plugin@^1.133.36": - version "1.167.28" - integrity sha512-O23ba7JaKvx5Eu0l6iTpknu79QcdkMmoW1VtZdsZe5NoQ6dHHru6caoapDc/uOxmz7h7VYfSuLjs/UYg7EA1cA== + version "1.167.32" + integrity sha512-i9BA6GzUCoM20UYZ77orXzHwD5zM0OQTtLuPNbqTTSG38CvR6viRFP/d+QFo2aRNyCvun8PR7zSa49bslSggEQ== dependencies: "@babel/core" "^7.28.5" "@babel/plugin-syntax-jsx" "^7.27.1" @@ -1450,8 +1450,8 @@ "@babel/template" "^7.27.2" "@babel/traverse" "^7.28.5" "@babel/types" "^7.28.5" - "@tanstack/router-core" "1.168.17" - "@tanstack/router-generator" "1.166.36" + "@tanstack/router-core" "1.169.1" + "@tanstack/router-generator" "1.166.39" "@tanstack/router-utils" "1.161.7" "@tanstack/virtual-file-routes" "1.161.7" chokidar "^3.6.0" @@ -1481,8 +1481,8 @@ integrity sha512-olW33+Cn+bsCsZKPwEGhlkqS6w3M2slFv11JIobdnCFKMLG97oAI2kWKdx5/zsywTL8flpnoIgaZZPlQTFYhdQ== "@tybys/wasm-util@^0.10.1": - version "0.10.1" - integrity sha512-9tTaPJLSiejZKx+Bmog4uSubteqTvFrVrURwkmHixBo0G4seD0zUxp98E1DzUBJxLQ3NPwXrGKDiVjwx/DpPsg== + version "0.10.2" + integrity sha512-RoBvJ2X0wuKlWFIjrwffGw1IqZHKQqzIchKaadZZfnNpsAYp2mM0h36JtPCjNDAHGgYez/15uMBpfGwchhiMgg== dependencies: tslib "^2.4.0" @@ -1690,10 +1690,10 @@ available-typed-arrays@^1.0.7: possible-typed-array-names "^1.0.0" axios@^1.13.1: - version "1.15.2" - integrity sha512-wLrXxPtcrPTsNlJmKjkPnNPK2Ihe0hn0wGSaTEiHRPxwjvJwT3hKmXF4dpqxmPO9SoNb2FsYXj/xEo0gHN+D5A== + version "1.16.0" + integrity sha512-6hp5CwvTPlN2A31g5dxnwAX0orzM7pmCRDLnZSX772mv8WDqICwFjowHuPs04Mc8deIld1+ejhtaMn5vp6b+1w== dependencies: - follow-redirects "^1.15.11" + follow-redirects "^1.16.0" form-data "^4.0.5" proxy-from-env "^2.1.0" @@ -1715,8 +1715,8 @@ balanced-match@^4.0.2: integrity sha512-BLrgEcRTwX2o6gGxGOCNyMvGSp35YofuYzw9h1IMTRmKqttAZZVU67bdb9Pr2vUHA8+j3i2tJfjO6C6+4myGTA== baseline-browser-mapping@^2.10.12: - version "2.10.23" - integrity sha512-xwVXGqevyKPsiuQdLj+dZMVjidjJV508TBqexND5HrF89cGdCYCJFB3qhcxRHSeMctdCfbR1jrxBajhDy7o29g== + version "2.10.27" + integrity sha512-zEs/ufmZoUd7WftKpKyXaT6RFxpQ5Qm9xytKRHvJfxFV9DFJkZph9RvJ1LcOUi0Z1ZVijMte65JbILeV+8QQEA== binary-extensions@^2.0.0: version "2.3.0" @@ -1956,8 +1956,8 @@ dunder-proto@^1.0.0, dunder-proto@^1.0.1: gopd "^1.2.0" electron-to-chromium@^1.5.328: - version "1.5.344" - integrity sha512-4MxfbmNDm+KPh066EZy+eUnkcDPcZ35wNmOWzFuh/ijvHsve6kbLTLURy88uCNK5FbpN+yk2nQY6BYh1GEt+wg== + version "1.5.349" + integrity sha512-QsWVGyRuY07Aqb234QytTfwd5d9AJlfNIQ5wIOl1L+PZDzI9d9+Fn0FRale/QYlFxt/bUnB0/nLd1jFPGxGK1A== emoji-regex@^8.0.0: version "8.0.0" @@ -2203,8 +2203,8 @@ fast-safe-stringify@^2.0.7: integrity sha512-W+KJc2dmILlPplD/H4K9l9LcAHAfPtP6BY84uVLXQ6Evcz9Lcg33Y2z1IVblT6xdY54PXYVHEv+0Wpq8Io6zkA== fast-uri@^3.0.1: - version "3.1.0" - integrity sha512-iPeeDKJSWf4IEOasVVrknXpaBV0IApz/gp7S2bb7Z4Lljbl2MGJRqInZiUrQwV16cpzw/D3S5j5Julj/gT52AA== + version "3.1.1" + integrity sha512-h2r7rcm6Ee/J8o0LD5djLuFVcfbZxhvho4vvsbeV0aMvXjUgqv4YpxpkEx0d68l6+IleVfLAdVEfhR7QNMkGHQ== fastq@^1.6.0: version "1.20.1" @@ -2229,7 +2229,7 @@ find-up@5.0.0: locate-path "^6.0.0" path-exists "^4.0.0" -follow-redirects@^1.15.11: +follow-redirects@^1.16.0: version "1.16.0" integrity sha512-y5rN/uOsadFT/JfYwhxRS5R7Qce+g3zG97+JrtFZlC9klX/W5hD7iiLzScI4nZqUS7DNUdhPgw4xI8W2LuXlUw== @@ -2890,8 +2890,8 @@ ms@^2.1.3: integrity sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA== nanoid@^3.3.11: - version "3.3.11" - integrity sha512-N8SpfPUnUp1bK+PMYW8qSWdl9U+wwNWI4QKxOYDy9JAro3WMX7p2OeVRF9v+347pnakNevPmiHhNmZ2HbFA76w== + version "3.3.12" + integrity sha512-ZB9RH/39qpq5Vu6Y+NmUaFhQR6pp+M2Xt76XBnEwDaGcVAqhlvxrl3B2bKS5D3NH3QR76v3aSrKaF/Kiy7lEtQ== nimma@0.2.3: version "0.2.3" @@ -3099,8 +3099,8 @@ possible-typed-array-names@^1.0.0: integrity sha512-/+5VFTchJDoVj3bhoqi6UeymcD00DAwb1nJwamzPvHEszJ4FpF6SNNbUbOS8yI56qHzdV8eK0qEfOSiodkTdxg== postcss@^8.5.6: - version "8.5.12" - integrity sha512-W62t/Se6rA0Az3DfCL0AqJwXuKwBeYg6nOaIgzP+xZ7N5BFCI7DYi1qs6ygUYT6rvfi6t9k65UMLJC+PHZpDAA== + version "8.5.14" + integrity sha512-SoSL4+OSEtR99LHFZQiJLkT59C5B1amGO1NzTwj7TT1qCUgUO6hxOvzkOYxD+vMrXBM3XJIKzokoERdqQq/Zmg== dependencies: nanoid "^3.3.11" picocolors "^1.1.1" @@ -3215,36 +3215,36 @@ reusify@^1.0.4: integrity sha512-g6QUff04oZpHs0eG5p83rFLhHeV00ug/Yf9nZM6fLeUrPguBTkTQOdpAWWspMh55TZfVQDPaN3NQJfbVRAxdIw== rollup@^4.43.0: - version "4.60.2" - integrity sha512-J9qZyW++QK/09NyN/zeO0dG/1GdGfyp9lV8ajHnRVLfo/uFsbji5mHnDgn/qYdUHyCkM2N+8VyspgZclfAh0eQ== + version "4.60.3" + integrity sha512-pAQK9HalE84QSm4Po3EmWIZPd3FnjkShVkiMlz1iligWYkWQ7wHYd1PF/T7QZ5TVSD6uSTon5gBVMSM4JfBV+A== dependencies: "@types/estree" "1.0.8" optionalDependencies: - "@rollup/rollup-android-arm-eabi" "4.60.2" - "@rollup/rollup-android-arm64" "4.60.2" - "@rollup/rollup-darwin-arm64" "4.60.2" - "@rollup/rollup-darwin-x64" "4.60.2" - "@rollup/rollup-freebsd-arm64" "4.60.2" - "@rollup/rollup-freebsd-x64" "4.60.2" - "@rollup/rollup-linux-arm-gnueabihf" "4.60.2" - "@rollup/rollup-linux-arm-musleabihf" "4.60.2" - "@rollup/rollup-linux-arm64-gnu" "4.60.2" - "@rollup/rollup-linux-arm64-musl" "4.60.2" - "@rollup/rollup-linux-loong64-gnu" "4.60.2" - "@rollup/rollup-linux-loong64-musl" "4.60.2" - "@rollup/rollup-linux-ppc64-gnu" "4.60.2" - "@rollup/rollup-linux-ppc64-musl" "4.60.2" - "@rollup/rollup-linux-riscv64-gnu" "4.60.2" - "@rollup/rollup-linux-riscv64-musl" "4.60.2" - "@rollup/rollup-linux-s390x-gnu" "4.60.2" - "@rollup/rollup-linux-x64-gnu" "4.60.2" - "@rollup/rollup-linux-x64-musl" "4.60.2" - "@rollup/rollup-openbsd-x64" "4.60.2" - "@rollup/rollup-openharmony-arm64" "4.60.2" - "@rollup/rollup-win32-arm64-msvc" "4.60.2" - "@rollup/rollup-win32-ia32-msvc" "4.60.2" - "@rollup/rollup-win32-x64-gnu" "4.60.2" - "@rollup/rollup-win32-x64-msvc" "4.60.2" + "@rollup/rollup-android-arm-eabi" "4.60.3" + "@rollup/rollup-android-arm64" "4.60.3" + "@rollup/rollup-darwin-arm64" "4.60.3" + "@rollup/rollup-darwin-x64" "4.60.3" + "@rollup/rollup-freebsd-arm64" "4.60.3" + "@rollup/rollup-freebsd-x64" "4.60.3" + "@rollup/rollup-linux-arm-gnueabihf" "4.60.3" + "@rollup/rollup-linux-arm-musleabihf" "4.60.3" + "@rollup/rollup-linux-arm64-gnu" "4.60.3" + "@rollup/rollup-linux-arm64-musl" "4.60.3" + "@rollup/rollup-linux-loong64-gnu" "4.60.3" + "@rollup/rollup-linux-loong64-musl" "4.60.3" + "@rollup/rollup-linux-ppc64-gnu" "4.60.3" + "@rollup/rollup-linux-ppc64-musl" "4.60.3" + "@rollup/rollup-linux-riscv64-gnu" "4.60.3" + "@rollup/rollup-linux-riscv64-musl" "4.60.3" + "@rollup/rollup-linux-s390x-gnu" "4.60.3" + "@rollup/rollup-linux-x64-gnu" "4.60.3" + "@rollup/rollup-linux-x64-musl" "4.60.3" + "@rollup/rollup-openbsd-x64" "4.60.3" + "@rollup/rollup-openharmony-arm64" "4.60.3" + "@rollup/rollup-win32-arm64-msvc" "4.60.3" + "@rollup/rollup-win32-ia32-msvc" "4.60.3" + "@rollup/rollup-win32-x64-gnu" "4.60.3" + "@rollup/rollup-win32-x64-msvc" "4.60.3" fsevents "~2.3.2" run-parallel@^1.1.9: @@ -3291,12 +3291,12 @@ semver@^6.3.1: integrity sha512-BR7VvDCVHO+q2xBEWskxS6DJE1qRnb7DxzUrogb71CWoSficBxYsiAGd+Kl0mmq/MprG9yArRkyrQxTO6XjMzA== seroval-plugins@^1.5.0: - version "1.5.2" - integrity sha512-qpY0Cl+fKYFn4GOf3cMiq6l72CpuVaawb6ILjubOQ+diJ54LfOWaSSPsaswN8DRPIPW4Yq+tE1k5aKd7ILyaFg== + version "1.5.3" + integrity sha512-LhVh4KjjkKmCxOUjoaUwtqbDjyMfnA535yEmmGDuwZcIYtw8ns6tZmeszNTECeUg/3sJpnEjsz/KhQrcPXPw1Q== seroval@^1.5.0: - version "1.5.2" - integrity sha512-xcRN39BdsnO9Tf+VzsE7b3JyTJASItIV1FVFewJKCFcW4s4haIKS3e6vj8PGB9qBwC7tnuOywQMdv5N4qkzi7Q== + version "1.5.3" + integrity sha512-BXe0x4buEeYiIKaRUnth1WqCILQ3k4O67KP/B4pC3pVz0Mv2c96ngA9QDREUYxWY1sb2RZVRqwI9RcpVMyHCVw== set-function-length@^1.2.2: version "1.2.2" @@ -3791,8 +3791,8 @@ yaml@^1.10.0: integrity sha512-vIYeF1u3CjlhAFekPPAk2h/Kv4T3mAkMox5OymRiJQB0spDP10LHvt+K7G9Ny6NuuMAb25/6n1qyUjAcGNf/AA== yaml@^2.8.0, yaml@^2.8.3: - version "2.8.3" - integrity sha512-AvbaCLOO2Otw/lW5bmh9d/WEdcDFdQp2Z2ZUH3pX9U2ihyUY0nvLv7J6TrWowklRGPYbB/IuIMfYgxaCPg5Bpg== + version "2.8.4" + integrity sha512-ml/JPOj9fOQK8RNnWojA67GbZ0ApXAUlN2UQclwv2eVgTgn7O9gg9o7paZWKMp4g0H3nTLtS9LVzhkpOFIKzog== yargs-parser@^21.1.1: version "21.1.1" diff --git a/docs/dqx/docs/dev/contributing.mdx b/docs/dqx/docs/dev/contributing.mdx index f3883dc23..4fc7b883b 100644 --- a/docs/dqx/docs/dev/contributing.mdx +++ b/docs/dqx/docs/dev/contributing.mdx @@ -200,8 +200,11 @@ Before running `make app-start-dev`, configure authentication to a Databricks wo DATABRICKS_CONFIG_PROFILE= # from ~/.databrickscfg DATABRICKS_WAREHOUSE_ID= # SQL Warehouses → connection details DQX_JOB_ID= # optional locally; required for profiler / dry-run +DQX_LAKEBASE_INSTANCE_NAME= # optional locally; empty = OLTP tables run on Delta ``` +Leave `DQX_LAKEBASE_INSTANCE_NAME` empty for most local dev — the app falls back to Delta for the OLTP tables (rules catalog, app settings, RBAC, comments, schedules) so you don't need a Lakebase instance to iterate. To exercise the Lakebase path locally, deploy the bundle once and point the variable at the resulting instance. See [app/DEVELOPMENT.md](https://github.com/databrickslabs/dqx/blob/v0.14.0/app/DEVELOPMENT.md) for the full set of `DQX_LAKEBASE_*` variables. + If you don't have a profile yet, run `databricks auth login --host -p ` first. See the [Development Mode](https://github.com/databrickslabs/dqx/blob/v0.14.0/app/README.md#development-mode) section of the app README for more detail. @@ -227,7 +230,7 @@ Deploying DQX Studio to a workspace is required when you want to: - verify a change behaves correctly under the production identity model (service principal + on-behalf-of), or - run a review pass against a deployed app before merging. -For the full step-by-step (service principal creation, asset-bundle deploy, schema/volume permission grants, app start, troubleshooting) follow [app/DEPLOYMENT.md](https://github.com/databrickslabs/dqx/blob/v0.14.0/app/DEPLOYMENT.md). +For the full step-by-step (service principal creation, one-time storage bootstrap, asset-bundle deploy, Lakebase opt-out, post-deploy grants, app start, troubleshooting) follow [app/DEPLOYMENT.md](https://github.com/databrickslabs/dqx/blob/v0.14.0/app/DEPLOYMENT.md). ### Running integration tests and code coverage diff --git a/docs/dqx/docs/installation.mdx b/docs/dqx/docs/installation.mdx index 5a24ef9fe..06f136203 100644 --- a/docs/dqx/docs/installation.mdx +++ b/docs/dqx/docs/installation.mdx @@ -438,7 +438,8 @@ Databricks CLI will confirm a few options: - **Databricks Apps** feature enabled on the workspace. - **User token passthrough** enabled for Databricks Apps — DQX Studio uses On-Behalf-Of (OBO) tokens to access Unity Catalog with the end user's identity. - **Serverless compute** enabled on the workspace — the `dqx-studio-task-runner` job is serverless-only. -- An **existing Unity Catalog catalog** where the studio's schemas and volumes will be created. The bundle does not create the catalog itself. +- An **existing Unity Catalog catalog** where the studio's schemas and volumes will be provisioned. The bundle does not create the catalog itself; `make app-deploy` provisions the schemas and wheels volume *inside* the catalog automatically. +- **Lakebase Postgres** enabled on the workspace if you keep the default backend layout. DQX Studio stores its OLTP state (rules catalog, app settings, RBAC, comments, schedule configs, scheduler bookkeeping) in a Lakebase instance for sub-millisecond reads. The bundle creates the instance for you. If your workspace doesn't have Lakebase, you can opt out and run everything on Delta — see the **Hybrid storage backend** admonition under the install steps below. ### Install DQX Studio using a Declarative Automation Bundle @@ -464,26 +465,45 @@ Databricks CLI will confirm a few options: # admin_group: # default: admins # app_name: # default: dqx-studio # sql_warehouse_name: # default: dqx-studio-sql-warehouse - # schema_name: # default: dqx_app + # schema_name: # default: dqx_studio + + # Lakebase backend (default: enabled). Set lakebase_instance_name + # to "" to disable Lakebase and run all OLTP tables on Delta. + # lakebase_instance_name: dqx-studio-lakebase + # lakebase_database_name: dqx_studio + # lakebase_capacity: CU_1 # CU_1 / CU_2 / CU_4 / CU_8 presets: trigger_pause_status: PAUSED ``` - See the [DQX Studio deployment guide](https://github.com/databrickslabs/dqx/blob/v0.14.0/app/DEPLOYMENT.md#step-3-configure-databricksyml) for the full reference of each variable, including security implications of `admin_group` and when to override per target. + See the [DQX Studio deployment guide](https://github.com/databrickslabs/dqx/blob/v0.14.0/app/DEPLOYMENT.md#step-4-configure-databricksyml) for the full reference of each variable, including security implications of `admin_group` and when to override per target. + +4. (One-time, only on a workspace whose schemas / volume / Lakebase instance were created out-of-band before adopting this layout) adopt them into bundle management: + ```commandline + make app-bind PROFILE= TARGET= + ``` -4. Build, deploy, grant permissions, and start the studio in a single command: +5. Build, deploy, grant permissions, and start the studio in a single command: ```commandline make app-deploy PROFILE= TARGET= ``` - This runs `make app-build`, `databricks bundle deploy`, `app/scripts/post_deploy_grants.sh`, and `databricks bundle run` in sequence. + This runs `make app-build`, `databricks bundle deploy` (provisions the schemas, wheels volume, Lakebase Postgres instance, SQL warehouse, task-runner job, and Databricks App in dependency order; the stateful resources — schemas, volume, and instance — carry `lifecycle.prevent_destroy: true`), `app/scripts/post_deploy_grants.sh`, and `databricks bundle run` in sequence. The app connects to the always-present `databricks_postgres` admin database on the instance and creates its own `dqx_studio` Postgres schema there on first start; no logical-database provisioning step is needed. -5. Open the deployed app from the **Apps** page in your Databricks workspace. +6. Open the deployed app from the **Apps** page in your Databricks workspace. -For the full walkthrough — including step-by-step commands, manual `GRANT` statements, troubleshooting, and target-specific configuration — see the [DQX Studio deployment guide](https://github.com/databrickslabs/dqx/blob/v0.14.0/app/DEPLOYMENT.md). + +The studio's schemas (`dqx_studio`, `dqx_studio_tmp`), wheels volume, and Lakebase Postgres instance are declared as bundle resources with `lifecycle.prevent_destroy: true` (Databricks CLI 0.268+). `databricks bundle destroy` is blocked from dropping them, so production data survives accidental destroy/replace operations. The app's `dqx_studio` Postgres schema (inside the `databricks_postgres` admin database on the Lakebase instance) is created at startup and is not itself a bundle resource, but it is protected transitively by the instance-level guard — as long as the instance survives, the schema and its tables survive. To intentionally tear something down, remove the flag, `databricks bundle deployment unbind `, then destroy manually. + + +For the full walkthrough — including step-by-step commands, manual `GRANT` statements, troubleshooting, target-specific configuration, and the bind workflow for adopting existing resources — see the [DQX Studio deployment guide](https://github.com/databrickslabs/dqx/blob/v0.14.0/app/DEPLOYMENT.md). + + +DQX Studio splits its data across two physical backends: high-volume append-mostly tables (`dq_validation_runs`, `dq_profiling_results`, `dq_quarantine_records`, `dq_metrics`) live in **Delta Lake** because they're written by Spark; transactional tables (rules catalog, app settings, RBAC, comments, schedule configs) live in **Lakebase Postgres** for fast row-level reads/writes from the FastAPI request handlers. + -On its first start, DQX Studio runs database migrations and uploads its DQX wheel files to the Unity Catalog volume. If the task-runner job is triggered before the app has fully started at least once, it will fail to find its wheels. Wait for `Uploaded databricks_labs_dqx-...` in the app logs before triggering any profiler or dry-run jobs. +On its first start, DQX Studio runs database migrations (Delta and, if enabled, Lakebase) and uploads its DQX wheel files to the Unity Catalog volume. If the task-runner job is triggered before the app has fully started at least once, it will fail to find its wheels. Wait for `Uploaded databricks_labs_dqx-...` in the app logs before triggering any profiler or dry-run jobs. If Lakebase is enabled, also wait for `Lakebase OLTP routing enabled` before opening the UI — when Lakebase is configured (`DQX_LAKEBASE_INSTANCE_NAME` non-empty) and init fails, the app refuses to start and the Databricks Apps platform will restart it; silent fallback to Delta would split-brain OLTP writes across two physical stores and is not safe. To run on Delta only, unset `DQX_LAKEBASE_INSTANCE_NAME`. ### Upgrade DQX Studio @@ -495,17 +515,41 @@ git pull make app-deploy PROFILE= TARGET= ``` -Database migrations run automatically on app startup and preserve existing rules, runs, schedules, and configuration. +Database migrations run automatically on app startup and preserve existing rules, runs, schedules, and configuration. All stateful resources are declared with `lifecycle.prevent_destroy: true`, so even an accidental `databricks bundle destroy` won't drop them. ### Uninstall DQX Studio -Destroy the bundle to remove the app, the task-runner job, and the SQL warehouse: - -```commandline -cd app && databricks bundle destroy -p -t -``` +`databricks bundle destroy` removes the app, the task-runner job, and the SQL warehouse, but is **blocked** from dropping the bundle-managed stateful resources (schemas, volume, Lakebase instance) by their `prevent_destroy` flag. The DQX `dqx_studio` Postgres schema inside `databricks_postgres` lives below the resource layer DABs models and is therefore unaffected by `bundle destroy` — drop it manually if you want a fully clean wipe. To fully uninstall: + +1. Edit `app/databricks.yml` and remove `lifecycle.prevent_destroy: true` from each of the stateful resources you want to drop. +2. Unbind the resources so the bundle no longer tracks them, then destroy: + ```bash + cd app + databricks bundle deployment unbind main_schema -t + databricks bundle deployment unbind tmp_schema -t + databricks bundle deployment unbind wheels -t + databricks bundle deployment unbind lakebase -t + databricks bundle destroy -p -t + ``` +3. Drop the now-unbound UC resources manually if desired: + ```sql + DROP VOLUME IF EXISTS .dqx_studio.wheels; + DROP SCHEMA IF EXISTS .dqx_studio CASCADE; + DROP SCHEMA IF EXISTS .dqx_studio_tmp CASCADE; + ``` +4. Drop the DQX Postgres schema and the Lakebase instance if desired: + ```bash + # DQX schema inside ``databricks_postgres`` (connect to the + # Lakebase instance with psql via ``databricks_postgres``): + # DROP SCHEMA IF EXISTS dqx_studio CASCADE; + + # Lakebase instance: + databricks database delete-database-instance dqx-studio-lakebase -p + ``` -The Unity Catalog catalog and the studio's schemas (which contain historical rules and run results) are not removed by `bundle destroy` — drop them manually if you no longer need the data. + +`prevent_destroy` is what stops a `databricks bundle destroy` (or a forced replace from a deploy) from wiping your rules, schedules, role mappings, and comments. The multi-step uninstall above is by design — silent data loss is the bigger risk to protect against. + ## Installing Dashboard diff --git a/tests/integration/test_app_backend.py b/tests/integration/test_app_backend.py index 5a5fb4c61..621206618 100644 --- a/tests/integration/test_app_backend.py +++ b/tests/integration/test_app_backend.py @@ -5,7 +5,7 @@ import pytest from fastapi.testclient import TestClient -from databricks_labs_dqx_app.backend.common.authorization import UserRole +from databricks_labs_dqx_app.backend.common.authorization import UserRole, get_user_email from databricks_labs_dqx_app.backend.dependencies import get_obo_ws, get_user_role, get_app_settings_service from databricks_labs_dqx_app.backend.models import InstallationSettings from databricks_labs_dqx_app.backend.services.app_settings_service import AppSettingsService @@ -107,15 +107,30 @@ async def override_get_user_role() -> UserRole: warehouse_id = os.environ.get("DATABRICKS_WAREHOUSE_ID", "") test_schema = make_schema(catalog_name=TEST_CATALOG) _sql = SqlExecutor(ws=ws, warehouse_id=warehouse_id, catalog=TEST_CATALOG, schema=test_schema.name) + + # Create the OLTP fallback tables (dq_app_settings, dq_quality_rules, ...) + # in the test schema. ``AppSettingsService.ensure_table()`` is a no-op + # since the Lakebase refactor — DDL now lives in MigrationRunner. + from databricks_labs_dqx_app.backend.migrations import MigrationRunner + + MigrationRunner(_sql).run_all(include_oltp_fallback=True) + _settings_svc = AppSettingsService(sql=_sql) - _settings_svc.ensure_table() async def override_get_app_settings_service() -> AppSettingsService: return _settings_svc + # The production proxy injects ``X-Forwarded-Email``; ``TestClient`` does + # not, so without this override every write endpoint returns 401. + test_user_email = ws.current_user.me().user_name or "test@example.com" + + def override_get_user_email() -> str: + return test_user_email + app.dependency_overrides[get_obo_ws] = override_get_obo_ws app.dependency_overrides[get_user_role] = override_get_user_role app.dependency_overrides[get_app_settings_service] = override_get_app_settings_service + app.dependency_overrides[get_user_email] = override_get_user_email client = TestClient(app) yield client diff --git a/tests/unit/test_app_backend.py b/tests/unit/test_app_backend.py index f66a3f182..69dea0c9f 100644 --- a/tests/unit/test_app_backend.py +++ b/tests/unit/test_app_backend.py @@ -116,9 +116,15 @@ def _failed_response(message: str = "boom") -> StatementResponse: _SAMPLE_CHECKS = [{"criticality": "error", "check": {"function": "is_not_null", "arguments": {"column": "id"}}}] +# Row[1] must JSON-decode to a *bare* check object, not a list. After +# the v1 baseline split each catalog row stores one check in the +# VARIANT/JSONB ``check`` column rather than an array of checks +# (see RulesCatalogService._row_to_entry — anything that decodes to a +# non-dict is dropped). ``_SAMPLE_CHECKS`` keeps its list shape because +# it is the *input* to ``save`` (which writes one row per element). _SAMPLE_ROW = [ "catalog.schema.table", - json.dumps(_SAMPLE_CHECKS), + json.dumps(_SAMPLE_CHECKS[0]), "3", "draft", "alice@example.com",