From e9ee4237063ba995200a5a56006c2c984cb110b9 Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:19:56 +0100 Subject: [PATCH 1/8] refactor(iac): one bounded_ident for names a service caps A capped name kept its first characters and ended in an 8-digit digest of the whole in gcp_governance._bounded; the Databricks key constraints need the same with an underscore. naming.bounded_ident is that helper, and _bounded delegates to it (same output). --- fluid_build/iac/naming.py | 8 ++++++++ fluid_build/iac/providers/gcp_governance.py | 7 ++----- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/fluid_build/iac/naming.py b/fluid_build/iac/naming.py index cbe1df48..f16a8028 100644 --- a/fluid_build/iac/naming.py +++ b/fluid_build/iac/naming.py @@ -48,6 +48,14 @@ def digest_ident(value: str) -> str: return f"{safe_ident(value)}_{_digest(value, 32)}" +def bounded_ident(text: str, limit: int, sep: str = "_") -> str: + """``text``, or, when it is longer than ``limit``, its first characters, ``sep`` and an + 8-digit digest of the whole: a name a service caps, whose cut forms stay distinct.""" + if len(text) <= limit: + return text + return f"{text[: limit - 9]}{sep}{_digest(text, 8)}" + + def _digest(value: str, length: int) -> str: return hashlib.sha256(value.encode("utf-8")).hexdigest()[:length] diff --git a/fluid_build/iac/providers/gcp_governance.py b/fluid_build/iac/providers/gcp_governance.py index 2b8864ac..7d226be4 100644 --- a/fluid_build/iac/providers/gcp_governance.py +++ b/fluid_build/iac/providers/gcp_governance.py @@ -76,7 +76,7 @@ from ..base import UnsupportedBindingError from ..column_access import authz_readers, column_masks, column_readers, restrictions_for from ..governance_labels import governance_labels -from ..naming import safe_ident +from ..naming import bounded_ident, safe_ident from ..principals import GCP, gcp_grants, principal_map, resolve_principal from ..provider_match import canonical_cloud from ..row_access import row_filters_for @@ -249,10 +249,7 @@ def taxonomy_region(bq_location: str) -> str: def _bounded(text: str, limit: int) -> str: """``text`` cut to ``limit`` characters, a short hash keeping cut names distinct.""" - if len(text) <= limit: - return text - digest = hashlib.sha256(text.encode("utf-8")).hexdigest()[:8] - return f"{text[: limit - 9]}-{digest}" + return bounded_ident(text, limit, "-") def product_key_ring(cid: str, dataset: str) -> str: From 8de4ac46a2e7fd316d70b38dc50a242f785cc526 Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:19:56 +0100 Subject: [PATCH 2/8] feat(databricks): Unity Catalog primary and foreign keys from the contract A column tagged primary-key becomes the table's PRIMARY KEY (composite in schema order; each column must be required), and one tagged foreign-key with labels.references: . a FOREIGN KEY to that product's table, so Catalog Explorer draws the ERD and Genie can join the products. FLUID 0.7.6/0.7.7 has no key field ($defs.column is closed), so these are the tag conventions the ODCS exporter and the changelog rules read. - databricks_keys.py: compile, resolve, plan, reconcile. A parent in the same contract is one of its exposes; any other is found live by the fluid_product_id tag every forge table carries (system.information_schema.table_tags), its live PRIMARY KEY must match, the child's catalog wins a tie, an ambiguous one is skipped and said. - fluid apply declares them in post_apply on the SQL warehouse (terraform-provider-databricks has no constraint argument: #4157, PR #4205 unmerged), re-declares what is missing, replaces or drops only keys it named (fluid_pk_, fluid_fk_
_), and never fails the run on a key: each skip says why and what to do. --mode dry-run prints the statements. - Measured on a warehouse (dbsql 2026.45): CREATE OR REPLACE drops a table's keys and every foreign key referencing it; INSERT OVERWRITE, tags and masks survive. So a Databricks dbt job that writes a keyed table gets a second task, fluid_keys, a SQL file task whose CONTINUE handler passes over a failing statement (re-adding keys dbt dropped), and a dbt build that amend-and-build runs locally is followed by the same re-declaration (post_build hook). - The pattern is dbt-databricks' own: ALTER TABLE ... ADD CONSTRAINT after every create or replace (macros/relations/constraints.sql). - _fake_databricks serves the keys with the live rules and error texts, and the table_tags lookup. --- CHANGELOG.md | 20 + fluid_build/cli/_apply_opentofu_engine.py | 37 + fluid_build/cli/apply.py | 4 +- fluid_build/iac/providers/databricks.py | 173 +++- fluid_build/iac/providers/databricks_jobs.py | 64 +- fluid_build/iac/providers/databricks_keys.py | 750 ++++++++++++++++++ tests/iac/_fake_databricks.py | 238 +++++- tests/iac/test_iac_databricks_fake_e2e.py | 44 + .../iac/test_iac_databricks_jobs_fake_e2e.py | 26 + tests/iac/test_iac_databricks_keys.py | 663 ++++++++++++++++ 10 files changed, 2001 insertions(+), 18 deletions(-) create mode 100644 fluid_build/iac/providers/databricks_keys.py create mode 100644 tests/iac/test_iac_databricks_keys.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 166198fd..7fb4ba20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- **Databricks: Unity Catalog primary and foreign keys from the contract.** A column tagged + `primary-key` (all of them, in schema order, for a composite key; each must be + `required: true`) becomes the table's PRIMARY KEY, and one tagged `foreign-key` with + `labels.references: .` becomes a FOREIGN KEY to that product's table, + so Catalog Explorer draws the entity-relationship diagram and Genie can join the products. + The parent may live in another contract: it is found by the `fluid_product_id` tag every + forge table carries (`system.information_schema.table_tags`), wherever an overlay put it. + `fluid apply` declares the keys on the SQL warehouse after `tofu apply` (the provider has no + constraint argument), re-declares any that are missing, replaces or drops the ones it + declared (named `fluid_pk_
` / `fluid_fk_
_`) when the contract changes, + and leaves any other key in place. A key that cannot be declared yet (its parent not + deployed, say) is skipped with the reason and the fix, never failing the run; + `--mode dry-run` prints the statements. Because a dbt `table` build's `CREATE OR REPLACE` + drops a table's keys and every foreign key that references it (measured on a warehouse), + each Databricks dbt job that writes a keyed table gets a second task, `fluid_keys`, that + declares them again after dbt, in a script that never fails the run; a dbt build that + `--mode amend-and-build` runs locally is followed by the same re-declaration. + ### Fixed - **`fluid ai models --provider` works, and takes the names `fluid forge --llm-provider` diff --git a/fluid_build/cli/_apply_opentofu_engine.py b/fluid_build/cli/_apply_opentofu_engine.py index 6a78babc..8d7794df 100644 --- a/fluid_build/cli/_apply_opentofu_engine.py +++ b/fluid_build/cli/_apply_opentofu_engine.py @@ -363,6 +363,43 @@ def _enforce_access( ) from exc +def build_after_apply( + args, logger: logging.Logger, contract: Mapping[str, Any], build_kwargs: Mapping[str, Any] +) -> int: + """The build phase of a build-augmented mode, after a ``tofu apply`` that succeeded, then + the plugin's ``post_build`` hook once every build succeeded (:func:`after_builds`).""" + from fluid_build.build_runners import run_builds_from_args + + rc = run_builds_from_args(args, logger, force_run=True, **build_kwargs) + if rc == 0: + after_builds(args, contract) + return rc + + +def after_builds(args, contract: Mapping[str, Any]) -> None: + """The plugin's ``post_build`` hook, if it has one, after a build phase that succeeded. + + Optional plugin capability, discovered like ``plan_notices``, and advisory: a dbt build run + here may have rebuilt a table with CREATE OR REPLACE, which drops what the apply declared on + it that dbt does not (the Databricks plugin's Unity Catalog keys), and the plugin declares + it again. It never changes the run's exit code. + """ + try: + plugin = get_iac_plugin( + _resolve_provider(contract, getattr(args, "provider", None) or "auto") + ) + except Exception: # noqa: BLE001 — the apply resolved it already; nothing to add here + return + run = getattr(plugin, "post_build", None) + if not callable(run): + return + try: + with contract_source(source_contract_file(getattr(args, "contract", None))): + run(contract, say=cprint) + except Exception as exc: # noqa: BLE001 — advisory: the build's own result stands + cprint(f" after the build: {type(exc).__name__}: {redact_secret_text(str(exc))}") + + def _enforce_after_failed_apply(plugin: Any, contract: Mapping[str, Any]) -> None: """Best effort, bind-only ``post_apply`` for a ``tofu apply`` that failed part way. diff --git a/fluid_build/cli/apply.py b/fluid_build/cli/apply.py index 0c8aeca8..9f5f49f8 100644 --- a/fluid_build/cli/apply.py +++ b/fluid_build/cli/apply.py @@ -2086,9 +2086,9 @@ def run(args, logger: logging.Logger) -> int: # not silently dropped on the cloud-provider path. if rc == 0 and needs_build(resolved_mode): args.build_id = resolved_build_id - from fluid_build.build_runners import run_builds_from_args + from fluid_build.cli._apply_opentofu_engine import build_after_apply - return run_builds_from_args(args, logger, force_run=True, **build_kwargs) + return build_after_apply(args, logger, source_contract, build_kwargs) return rc try: diff --git a/fluid_build/iac/providers/databricks.py b/fluid_build/iac/providers/databricks.py index 3a4297de..19fc5ddf 100644 --- a/fluid_build/iac/providers/databricks.py +++ b/fluid_build/iac/providers/databricks.py @@ -150,7 +150,7 @@ import logging import os import re -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple from ...providers._sql_safety import validate_ident, validate_sql_type_param_payload @@ -176,7 +176,16 @@ from ..versions import required_providers from .databricks_access import TableEnforcement, compile_enforcement from .databricks_auth import tofu_env_overlay -from .databricks_enforce import Target, post_apply, pre_apply, preview, unprotected +from .databricks_enforce import ( + SqlRunner, + Target, + default_client, + post_apply, + pre_apply, + preview, + scrub, + unprotected, +) from .databricks_jobs import ( CompiledJobs, DbtJob, @@ -184,6 +193,17 @@ dbt_written_exposes, run_as_allowlist, ) +from .databricks_keys import ( + Candidate, + KeyTarget, + TableKeys, + compile_keys, + foreign_keys, + own_candidates, + rebuild_script, + reconcile, +) +from .databricks_keys import lookup as lookup_key_tables from .databricks_retention import Retention, check_expire _LOG = logging.getLogger(__name__) @@ -296,6 +316,8 @@ class _Table: column_comments: Tuple[Tuple[str, str], ...] = () #: ``(column, sensitivity)`` per column the contract classifies. classified: Tuple[Tuple[str, str], ...] = () + #: Its primary key and foreign key references (``databricks_keys.py``), if it has any. + keys: Optional[TableKeys] = None @property def full_name(self) -> str: @@ -933,6 +955,7 @@ def _resolve_table( contract: Mapping[str, Any], validating: bool, builds: Sequence[Mapping[str, Any]] = (), + warnings: Optional[List[str]] = None, ) -> Optional[_Table]: """One Databricks expose → :class:`_Table`, or ``None`` with its refusals.""" binding = exposure.get("binding") or {} @@ -970,6 +993,8 @@ def _resolve_table( refusals.append(exc) if len(refusals) > before: return None + full_name = f"{names['database']}.{names['schema']}.{names['table']}" + keys = compile_keys(where, full_name, exposure, columns, [] if warnings is None else warnings) return _Table( index=index, expose_id=expose_id, @@ -986,6 +1011,7 @@ def _resolve_table( comment=comment, column_comments=tuple(column_comments), classified=tuple(classified), + keys=keys, ) @@ -1304,9 +1330,14 @@ def _refuse_build_hosts( def _compile( - contract: Mapping[str, Any], env: Mapping[str, str], *, validating: bool = False + contract: Mapping[str, Any], + env: Mapping[str, str], + *, + validating: bool = False, + key_lookup: Optional[KeyLookup] = None, ) -> _Compiled: - """The whole module for ``contract``, and every refusal, in one pass.""" + """The whole module for ``contract``, and every refusal, in one pass (``key_lookup``: + :func:`_rebuild_scripts`).""" compiled = _Compiled() refusals = compiled.refusals _refuse_build_hosts(contract, refusals) @@ -1349,6 +1380,7 @@ def _compile( ) _emit_grants(compiled, table, cid, privileges.get(table.full_name, {})) _emit_tags(compiled, table, cid, contract) + _check_own_keys(compiled, contract) _merge_jobs( compiled, compile_dbt_jobs( @@ -1357,6 +1389,7 @@ def _compile( warehouse, validating=validating, run_as_allowed=run_as_allowlist(env), + rebuild_script=_rebuild_scripts(compiled, contract, builds, key_lookup), ), ) @@ -1369,6 +1402,110 @@ def _compile( return compiled +#: ``products → {product id: its tables}``, the live half of a foreign key's resolution. +KeyLookup = Callable[[Sequence[str]], Mapping[str, Sequence[Candidate]]] + + +def _check_own_keys(compiled: _Compiled, contract: Mapping[str, Any]) -> None: + """A ``fluid validate`` warning for each foreign key to this contract's own product that no + expose's declared primary key takes. One to another product is resolved at apply time.""" + contract_id = _text(contract.get("id")) + keyed = [(t, t.keys) for t in compiled.tables if t.keys is not None] + own = own_candidates(contract_id, [keys for _t, keys in keyed]) + for table, keys in keyed: + mine = tuple(r for r in keys.references if r.product == contract_id) + if not mine: + continue + notes: List[str] = [] + foreign_keys(replace(keys, references=mine), own, notes, own_id=contract_id) + compiled.warnings.extend(f"{table.where}: {note}" for note in notes) + + +def _rebuild_scripts( + compiled: _Compiled, + contract: Mapping[str, Any], + builds: Sequence[Mapping[str, Any]], + key_lookup: Optional[KeyLookup], +) -> Callable[[Mapping[str, Any]], Optional[str]]: + """The ``fluid_keys`` script of each Databricks dbt job, by its build; ``None`` for a job + that writes no table with a key (``databricks_keys.rebuild_script``). + + Its foreign keys to other products need ``key_lookup``, which only ``emit`` passes; without + it (``fluid validate``, the native planner) a script holds the primary keys and the foreign + keys to this contract's own tables. + """ + contract_id = _text(contract.get("id")) + keyed = [(t, t.keys) for t in compiled.tables if t.keys is not None] + written: Dict[int, List[TableKeys]] = {} + for build in builds: + if not is_databricks_job_build(build): + continue + tables = [ + keys + for table, keys in keyed + if any( + w is build for w in builds_writing_expose(table.exposure.get("exposeId"), builds) + ) + ] + if tables: + written[id(build)] = tables + if not written: + return lambda build: None + candidates: Dict[str, Sequence[Candidate]] = dict( + own_candidates(contract_id, [keys for _t, keys in keyed]) + ) + others = sorted( + {r.product for tables in written.values() for k in tables for r in k.references} + - {contract_id} + ) + if others and key_lookup is not None: + candidates.update(key_lookup(others)) + scripts = { + build_id: rebuild_script( + [(keys, foreign_keys(keys, candidates, own_id=contract_id)) for keys in tables] + ) + for build_id, tables in written.items() + } + return lambda build: scripts.get(id(build)) + + +def _declare_keys( + contract: Mapping[str, Any], + say: Callable[[str], None], + client_factory: Optional[Callable[[], Any]] = None, + *, + dry_run: bool = False, +) -> None: + """Declare the contract's primary and foreign keys on its live tables + (``databricks_keys.reconcile``). A key is metadata: this never raises, and says what it + could not do.""" + try: + compiled = _compile(contract, os.environ) + if compiled.refusals: + return # the hooks that gate have refused already + targets = [KeyTarget(t.full_name, t.keys) for t in compiled.tables] + reconcile(targets, _text(contract.get("id")), say, client_factory, dry_run=dry_run) + except Exception as exc: # noqa: BLE001 — a key never fails the run + say(f"\n Unity Catalog keys: not declared ({type(exc).__name__}: {scrub(exc)}).") + + +def _live_key_lookup(products: Sequence[str]) -> Mapping[str, Sequence[Candidate]]: + """Other products' live tables, for the ``fluid_keys`` scripts ``emit`` writes. + + Best effort, empty when the workspace cannot be read: a module still emits without a + credential (``fluid generate iac``), and ``fluid apply`` resolves, declares and reports the + keys itself after ``tofu apply`` (``databricks_keys.reconcile``). + """ + try: + warehouse = str(os.environ.get(WAREHOUSE_ENV) or "").strip() + if not warehouse: + return {} + client = default_client() + return lookup_key_tables(SqlRunner(client, warehouse), client, products) + except Exception: # noqa: BLE001 — no live answer is no foreign key in the script + return {} + + def _resolve_tables( compiled: _Compiled, exposures: Sequence[Tuple[int, Mapping[str, Any]]], @@ -1387,6 +1524,7 @@ def _resolve_tables( contract=contract, validating=validating, builds=builds, + warnings=compiled.warnings, ) if table is not None: compiled.tables.append(table) @@ -2044,7 +2182,7 @@ class DatabricksIacPlugin: def emit( self, contract: Mapping[str, Any], actions: Iterable[Mapping[str, Any]] = () ) -> Dict[str, Any]: - compiled = _compile(contract, os.environ) + compiled = _compile(contract, os.environ, key_lookup=_live_key_lookup) if compiled.refusals: raise compiled.refusals[0] _enforce_sovereignty(contract, compiled.tables) @@ -2186,6 +2324,7 @@ def pre_apply( say = say or print if dry_run: preview(targets, say, client_factory) + _declare_keys(contract, say, client_factory, dry_run=True) else: pre_apply(targets, say, client_factory) @@ -2200,13 +2339,35 @@ def post_apply( longer restricts, and revoke ``SELECT`` from the restricted principals of any table it could not bind (``databricks_enforce.post_apply``). ``after_failed_apply`` is the bind-only mode for an apply that failed part way: it drops nothing and skips a table - that was never created.""" + that was never created. Then, after an apply that succeeded, the contract's Unity + Catalog keys (:func:`_declare_keys`).""" post_apply( self._enforcement_targets(contract), say or print, client_factory, after_failed_apply=after_failed_apply, ) + if not after_failed_apply: + _declare_keys(contract, say or print, client_factory) + + def post_build( + self, + contract: Mapping[str, Any], + say: Optional[Callable[[str], None]] = None, + client_factory: Optional[Callable[[], Any]] = None, + ) -> None: + """After ``fluid apply --mode amend-and-build`` ran the contract's builds here. + + A dbt build run on this machine against a Databricks table rebuilt it with CREATE OR + REPLACE, which drops the keys the apply declared, so they are declared again. A build + that runs as a Databricks job does that in its job's ``fluid_keys`` task, and a landing + (``INSERT OVERWRITE``) keeps them: neither needs this. + """ + from ...build_runners.base import is_dbt_build + + builds = [b for b in contract.get("builds") or [] if isinstance(b, Mapping)] + if any(is_dbt_build(b) and not is_databricks_job_build(b) for b in builds): + _declare_keys(contract, say or print, client_factory) def plan_notices(self, contract: Mapping[str, Any]) -> List[str]: """What the apply cannot do itself and someone must: ``USE CATALOG``; and the dbt jobs diff --git a/fluid_build/iac/providers/databricks_jobs.py b/fluid_build/iac/providers/databricks_jobs.py index 8fa0a46f..91ec4a79 100644 --- a/fluid_build/iac/providers/databricks_jobs.py +++ b/fluid_build/iac/providers/databricks_jobs.py @@ -46,7 +46,18 @@ import shlex from dataclasses import dataclass, field from pathlib import Path, PurePosixPath -from typing import Any, Dict, FrozenSet, List, Mapping, Optional, Sequence, Set, Tuple +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + List, + Mapping, + Optional, + Sequence, + Set, + Tuple, +) from urllib.parse import parse_qs, urlsplit, urlunsplit from ...util.binding_paths import ENV_PLACEHOLDER_RE @@ -55,6 +66,7 @@ from ..base import UnsupportedBindingError from ..naming import TofuExpr, digest_ident, safe_ident, tofu_ref, unique_ident from ..source import current_contract_source +from .databricks_keys import REBUILD_FILE, REBUILD_TASK, REBUILD_TIMEOUT_SECONDS #: The serverless environment the dbt task runs in: client version 2 and the #: dbt adapter, installed per run (Phase 0b ran exactly this against the provider). @@ -737,6 +749,7 @@ def compile_dbt_jobs( *, validating: bool = False, run_as_allowed: FrozenSet[str] = frozenset(), + rebuild_script: Optional[Callable[[Mapping[str, Any]], Optional[str]]] = None, ) -> CompiledJobs: """Every Databricks dbt job of ``contract``, with its workspace files and refusals. @@ -744,6 +757,10 @@ def compile_dbt_jobs( job runs dbt against the catalog and schema of the first one its build writes. ``fluid validate`` compiles with ``validating`` and no contract location may be known; the project files are then not read. + + ``rebuild_script`` gives, for a build, the SQL that declares again the Unity Catalog + keys its dbt run drops (``databricks_keys.rebuild_script``), or ``None``: a job with one + runs it in a second task, ``fluid_keys``, after dbt. """ compiled = CompiledJobs() builds = [b for b in contract.get("builds") or [] if isinstance(b, Mapping)] @@ -760,7 +777,8 @@ def compile_dbt_jobs( table = _target_table(item, builds, tables, compiled.refusals) source = _project_source(item, compiled.refusals, validating=validating) if len(compiled.refusals) == before and table is not None and source is not None: - _emit_job(compiled, item, table, settings, source) + script = rebuild_script(build) if rebuild_script is not None else None + _emit_job(compiled, item, table, settings, source, script) return compiled @@ -866,8 +884,17 @@ def _project_source( def _emit_job( - compiled: CompiledJobs, item: _Build, table: Any, settings: _Settings, source: _Source + compiled: CompiledJobs, + item: _Build, + table: Any, + settings: _Settings, + source: _Source, + keys_script: Optional[str] = None, ) -> None: + """The job, its uploaded project, and, with ``keys_script``, the ``fluid_keys`` task that + runs it after dbt (``databricks_keys.py``): uploaded next to the project, whose own files + dbt reads never include a top-level ``.sql`` file.""" + keys_file = [(REBUILD_FILE, keys_script.encode("utf-8"))] if keys_script is not None else [] dbt_task: Dict[str, Any] = { "commands": settings.commands, "warehouse_id": item.warehouse, @@ -894,16 +921,20 @@ def _emit_job( dbt_task["source"] = "GIT" if source.directory: dbt_task["project_directory"] = source.directory + uploads = keys_file else: dbt_task["source"] = "WORKSPACE" # A workspace path as the jobs service reads it (Asset Bundles' file_path form). dbt_task["project_directory"] = _under_home("/Workspace", item.root) - depends = _emit_files(compiled, item, source.files) - if depends: - job["depends_on"] = depends + uploads = [*source.files, *keys_file] + compiled.data.setdefault(CURRENT_USER, {})["me"] = {} + depends = _emit_files(compiled, item, uploads) + if depends: + job["depends_on"] = depends compiled.data.setdefault(CURRENT_USER, {})["me"] = {} job["task"] = [ - {"task_key": "dbt", "environment_key": "dbt", "dbt_task": dbt_task, **settings.task} + {"task_key": "dbt", "environment_key": "dbt", "dbt_task": dbt_task, **settings.task}, + *([_keys_task(item)] if keys_file else []), ] compiled.resources.append(("databricks_job", item.key, f"{item.where} job", job)) compiled.jobs.append( @@ -913,12 +944,29 @@ def _emit_job( address=f"databricks_job.{item.key}", source=dbt_task["source"], trigger=settings.trigger, - files=tuple(rel for rel, _ in source.files), + files=tuple(rel for rel, _ in uploads), root=item.root, ) ) +def _keys_task(item: _Build) -> Dict[str, Any]: + """The ``fluid_keys`` task: once dbt has succeeded, the uploaded script on the same + warehouse.""" + return { + "task_key": REBUILD_TASK, + "depends_on": [{"task_key": "dbt"}], + "timeout_seconds": REBUILD_TIMEOUT_SECONDS, + "sql_task": { + "warehouse_id": item.warehouse, + "file": { + "path": _under_home("/Workspace", f"{item.root}/{REBUILD_FILE}"), + "source": "WORKSPACE", + }, + }, + } + + def _emit_files( compiled: CompiledJobs, item: _Build, files: Sequence[Tuple[str, bytes]] ) -> List[str]: diff --git a/fluid_build/iac/providers/databricks_keys.py b/fluid_build/iac/providers/databricks_keys.py new file mode 100644 index 00000000..c15692b6 --- /dev/null +++ b/fluid_build/iac/providers/databricks_keys.py @@ -0,0 +1,750 @@ +# Copyright 2025 Agentics Transformation Limited +# SPDX-License-Identifier: Apache-2.0 + +"""Unity Catalog primary and foreign keys from the contract: read, resolve, declare, re-declare. + +A FLUID 0.7.x column has no key field (``$defs.column`` is closed), so the keys come from the +conventions forge already reads elsewhere (``api/changelog_rules.py``, the ODCS exporter): + +* ``tags: [primary-key]`` (or ``primaryKey``) on each key column: the table's PRIMARY KEY, + its columns in schema order. Each must be ``required: true``: Unity Catalog declares a + key only on NOT NULL columns, and ``required`` is what makes the column NOT NULL. +* ``tags: [foreign-key]`` (or ``foreignKey``) and ``labels.references: .``: + a FOREIGN KEY to the table of that product whose primary key is that column. Two or more + columns that reference one product together form one key when they are exactly the + columns of a composite primary key there. + +The keys are informational: Unity Catalog does not enforce them, and the default ``NORELY`` is +kept, so no query is rewritten on their word. Catalog Explorer draws its entity-relationship +diagram from them, and Genie joins on them. The ``databricks/databricks`` provider (1.140.0) +has no constraint argument on ``databricks_sql_table`` (terraform-provider-databricks#4157; +PR #4205, which would add one, is unmerged), so, like the masks (``databricks_enforce.py``), +the statements run on the SQL warehouse ``DATABRICKS_WAREHOUSE_ID`` names, outside OpenTofu, +after ``tofu apply``. + +**Where a parent is.** A product in the same contract is found among its own exposes, by the +primary keys they declare. Any other product is found live: every table this provider creates +carries the tag ``fluid_product_id`` (``databricks._PRODUCT_TAGS``), so +``system.information_schema.table_tags`` names the tables a product id deployed, wherever an +overlay put them and whichever contract or repository applied them. A candidate must have the +referenced columns as its live PRIMARY KEY; of several, one in the child's catalog wins, and +more than one left is ambiguous, so that key is skipped and said. + +**Names.** ``fluid_pk_
`` and ``fluid_fk_
_`` (a key's name is unique in +its schema; one longer than 255 characters is cut and given a digest). The ``fluid_`` prefix +marks what this module owns: it replaces or drops only a key so named, and leaves any other in +place, saying so. A key that already holds what the contract asks, whatever its name, is kept. + +**A key never fails the run**: it is metadata. One that cannot be declared yet (its parent +not deployed, a parent without its primary key, a nullable column) is skipped with the reason +and the fix, and the next ``fluid apply`` declares it. + +Measured on Databricks SQL (serverless warehouse, dbsql 2026.45, 2026-11-12): + +* ``CREATE OR REPLACE TABLE`` (as ``SELECT``, or with a column list) drops every key the table + had and every FOREIGN KEY on another table that references it. ``INSERT OVERWRITE`` keeps + them; table tags and column masks survive both. ``DROP COLUMN`` takes a key on that column + with it. +* ``ADD CONSTRAINT`` has no ``IF NOT EXISTS``: adding a key that exists fails + (``DELTA_CONSTRAINT_ALREADY_EXISTS``), and so does a second key on the same columns. + Dropping a primary key that a foreign key references fails without ``CASCADE``. +* A SQL script whose ``DECLARE CONTINUE HANDLER FOR SQLEXCEPTION`` passes over each statement + that fails runs as a job's SQL file task, and the run succeeds. + +So a dbt ``table`` build, a ``CREATE OR REPLACE``, would leave its table with no keys after +every run. Each Databricks dbt job that writes a keyed table therefore gets a second task, +``fluid_keys`` (``databricks_jobs.py``), which runs :func:`rebuild_script` on the same +warehouse once dbt has succeeded: the keys the apply resolved, each added again, and each that +fails (a key still in place, a parent whose own task has not run yet) passed over, so no run +fails on one. A parent's rebuild also drops the foreign keys that reference it; each child's +next build, or ``fluid apply``, declares them again. + +The pattern is dbt-databricks' own (``macros/relations/constraints.sql``: ``ALTER TABLE ... ADD +CONSTRAINT`` after every create or replace), for keys the contract declares rather than a model. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union + +from ...providers._sql_safety import reject_control_chars, validate_ident +from ..naming import bounded_ident +from .databricks_access import quote, quote_table, string_literal +from .databricks_enforce import ( + DATABRICKS_EXTRA_INSTALL, + WAREHOUSE_ENV, + SqlRunner, + StatementFailed, + default_client, + read_table, + scrub, +) + +#: The column tags that mark a key column, spelt as the ODCS exporter reads them. +PRIMARY_KEY_TAGS = frozenset({"primary-key", "primaryKey"}) +FOREIGN_KEY_TAGS = frozenset({"foreign-key", "foreignKey"}) +#: The column label that says what a foreign key references: ``.``. +REFERENCES_LABEL = "references" +#: The table tag that names the product a table belongs to (``databricks._PRODUCT_TAGS``). +PRODUCT_TAG = "fluid_product_id" +#: The prefixes of the keys this module declares, and so owns. +PK_PREFIX = "fluid_pk_" +FK_PREFIX = "fluid_fk_" +#: The job task, and the workspace file next to the dbt project, that re-declare the keys. +REBUILD_TASK = "fluid_keys" +REBUILD_FILE = "fluid_keys.sql" +#: That task's cap: a few ``ALTER TABLE`` statements on a warehouse dbt has just used. +REBUILD_TIMEOUT_SECONDS = 600 +#: A Unity Catalog name holds 255 characters. +_MAX_NAME = 255 + +Say = Callable[[str], None] + + +# ── what the contract declares ───────────────────────────────────────── + + +@dataclass(frozen=True) +class Reference: + """One ``labels.references`` on a ``foreign-key`` column.""" + + column: str + product: str + parent_column: str + + +@dataclass(frozen=True) +class TableKeys: + """The keys one table's ``contract.schema`` declares.""" + + full_name: str + primary: Tuple[str, ...] = () + references: Tuple[Reference, ...] = () + + @property + def table(self) -> str: + return self.full_name.split(".")[-1] + + @property + def catalog(self) -> str: + return self.full_name.split(".")[0] + + @property + def primary_name(self) -> str: + return constraint_name(PK_PREFIX, self.table) + + +@dataclass(frozen=True) +class KeyTarget: + """A table of the contract, and its keys (``None``: it declares none).""" + + full_name: str + keys: Optional[TableKeys] + + +@dataclass(frozen=True) +class Candidate: + """A table a referenced product has, and its primary key.""" + + full_name: str + primary: Tuple[str, ...] + + @property + def catalog(self) -> str: + return self.full_name.split(".")[0] + + +@dataclass(frozen=True) +class ForeignKey: + name: str + columns: Tuple[str, ...] + parent: str + parent_columns: Tuple[str, ...] + + def describe(self) -> str: + return ( + f"foreign key {self.name} ({', '.join(self.columns)}) references {self.parent} " + f"({', '.join(self.parent_columns)})" + ) + + +def constraint_name(prefix: str, *parts: str) -> str: + """``prefix`` and ``parts`` joined by ``_``; cut, with a digest, past 255 characters.""" + return bounded_ident(prefix + "_".join(parts), _MAX_NAME) + + +def _tags(raw: Any) -> frozenset: + if not isinstance(raw, (list, tuple)): + return frozenset() + return frozenset(str(tag).strip() for tag in raw if isinstance(tag, str)) + + +def _reference(labels: Any, column: str, at: str, warnings: List[str]) -> Optional[Reference]: + raw = labels.get(REFERENCES_LABEL) if isinstance(labels, Mapping) else None + text = str(raw or "").strip() + product, dot, parent_column = text.rpartition(".") + product, parent_column = product.strip(), parent_column.strip().lower() + if not dot or not product or not parent_column: + warnings.append( + f"{at} is tagged foreign-key, but its labels.references is {text!r}, not " + "., so Unity Catalog gets no foreign key for it. Set " + "labels.references to the referenced product's id and its key column." + ) + return None + try: + validate_ident(parent_column) + reject_control_chars(product, "labels.references") + except ValueError: + warnings.append( + f"{at}: labels.references {text!r} names no plain product id and column, so " + "Unity Catalog gets no foreign key for it." + ) + return None + return Reference(column, product, parent_column) + + +def compile_keys( + where: str, + full_name: str, + exposure: Mapping[str, Any], + columns: Sequence[Mapping[str, Any]], + warnings: List[str], +) -> Optional[TableKeys]: + """The keys ``exposure``'s ``contract.schema`` declares, or ``None`` when it declares none. + + ``columns`` are the table's compiled columns (lower-cased names and ``nullable``). A problem + is a warning, not a refusal: the table applies, without that key. + """ + nullable = {str(c.get("name")): bool(c.get("nullable")) for c in columns} + primary: List[str] = [] + references: List[Reference] = [] + schema = (exposure.get("contract") or {}).get("schema") or [] + for column in schema if isinstance(schema, list) else []: + if not isinstance(column, Mapping): + continue + name = str(column.get("name") or "").strip().lower() + if name not in nullable: + continue # refused by the column compile, which says why + tags = _tags(column.get("tags")) + if tags & PRIMARY_KEY_TAGS: + primary.append(name) + if tags & FOREIGN_KEY_TAGS: + ref = _reference(column.get("labels"), name, f"{where} column {name}", warnings) + if ref is not None: + references.append(ref) + loose = [c for c in primary if nullable[c]] + if loose: + warnings.append( + f"{where} tags {', '.join(loose)} primary-key without required: true, and Unity " + "Catalog declares a primary key only on NOT NULL columns, so the table gets none " + f"(and no foreign key can reference it). Set required: true on {', '.join(loose)}." + ) + primary = [] + if not primary and not references: + return None + return TableKeys(full_name, tuple(primary), tuple(references)) + + +def own_candidates(contract_id: str, keyed: Iterable[TableKeys]) -> Dict[str, List[Candidate]]: + """The contract's own tables with a primary key, as candidates for its own id.""" + own = [Candidate(k.full_name, k.primary) for k in keyed if k.primary] + return {contract_id: own} if contract_id and own else {} + + +# ── where each foreign key points ────────────────────────────────────── + + +def _pick(matches: List[Candidate], catalog: str) -> Tuple[Optional[Candidate], List[Candidate]]: + """The one match (one in ``catalog`` wins a tie), or ``None`` and every match left.""" + if len(matches) > 1: + local = [m for m in matches if m.catalog == catalog] + matches = local or matches + return (matches[0], []) if len(matches) == 1 else (None, matches) + + +def foreign_keys( + keys: TableKeys, + candidates: Mapping[str, Sequence[Candidate]], + notes: Optional[List[str]] = None, + *, + own_id: str = "", +) -> List[ForeignKey]: + """The foreign keys ``keys`` resolves to among ``candidates`` (product id → its tables). + + A reference that resolves to no table, or to more than one, is left out, and ``notes`` + says why and what to do. ``own_id`` is the contract's own product id, whose candidates + are its own exposes. + """ + said = notes if notes is not None else [] + out: List[ForeignKey] = [] + by_product: Dict[str, List[Reference]] = {} + for ref in keys.references: + by_product.setdefault(ref.product, []).append(ref) + for product, refs in by_product.items(): + children = ", ".join(r.column for r in refs) + found = list(candidates.get(product) or ()) + if not found: + if product == own_id: + said.append( + f"foreign key ({children}) to {product}: no expose of this contract " + "declares the primary key it references. Tag that column primary-key." + ) + else: + said.append( + f"foreign key ({children}) to {product}: no table tagged {PRODUCT_TAG} = " + f"{product!r} is visible to this identity. Apply {product} first, then " + "run fluid apply here again." + ) + continue + wanted = [r.parent_column for r in refs] + if len(refs) > 1 and len(set(wanted)) == len(wanted): + whole = [c for c in found if len(c.primary) > 1 and set(c.primary) == set(wanted)] + if whole: + pick, tied = _pick(whole, keys.catalog) + if pick is None: + said.append(_ambiguous(children, product, wanted, tied)) + continue + child = {r.parent_column: r.column for r in refs} + columns = tuple(child[c] for c in pick.primary) + name = constraint_name(FK_PREFIX, keys.table, *columns) + out.append(ForeignKey(name, columns, pick.full_name, pick.primary)) + continue + for ref in refs: + single = [c for c in found if c.primary == (ref.parent_column,)] + pick, tied = _pick(single, keys.catalog) + if pick is not None: + name = constraint_name(FK_PREFIX, keys.table, ref.column) + out.append(ForeignKey(name, (ref.column,), pick.full_name, pick.primary)) + elif tied: + said.append(_ambiguous(ref.column, product, [ref.parent_column], tied)) + else: + tables = ", ".join(sorted(c.full_name for c in found)) + said.append( + f"foreign key ({ref.column}) to {product}.{ref.parent_column}: {product} " + f"has {tables}, and none has PRIMARY KEY ({ref.parent_column}). Tag " + f"{ref.parent_column} primary-key in {product}'s contract and apply it, then " + "run fluid apply here again." + ) + return out + + +def _ambiguous( + children: str, product: str, parents: Sequence[str], tied: Sequence[Candidate] +) -> str: + return ( + f"foreign key ({children}) to {product}: more than one of its tables has PRIMARY KEY " + f"({', '.join(parents)}) ({', '.join(sorted(c.full_name for c in tied))}), and not " + "exactly one is in this table's catalog. Keep one of them visible to this identity." + ) + + +def live_keys( + live: Mapping[str, Any], +) -> Tuple[Optional[Tuple[str, Tuple[str, ...]]], List[ForeignKey]]: + """``((name, columns) or None, [foreign keys])`` from a Tables API read + (``table_constraints[]``: ``primary_key_constraint`` / ``foreign_key_constraint``).""" + primary: Optional[Tuple[str, Tuple[str, ...]]] = None + foreign: List[ForeignKey] = [] + + def names(raw: Any) -> Tuple[str, ...]: + return tuple(str(c).lower() for c in raw or [] if isinstance(c, str)) + + for entry in live.get("table_constraints") or []: + if not isinstance(entry, Mapping): + continue + pk = entry.get("primary_key_constraint") + if isinstance(pk, Mapping): + primary = (str(pk.get("name") or ""), names(pk.get("child_columns"))) + fk = entry.get("foreign_key_constraint") + if isinstance(fk, Mapping): + foreign.append( + ForeignKey( + str(fk.get("name") or ""), + names(fk.get("child_columns")), + str(fk.get("parent_table") or "").lower(), + names(fk.get("parent_columns")), + ) + ) + return primary, foreign + + +def lookup(runner: SqlRunner, client: Any, products: Iterable[str]) -> Dict[str, List[Candidate]]: + """The tables tagged with each product id, with their live primary keys. + + One query on ``system.information_schema.table_tags``, which lists what this identity can + see, then one Tables API read per table found. A table whose name this module could not + quote is no candidate. Raises when the query or a read fails. + """ + wanted = sorted({p for p in products if p}) + if not wanted: + return {} + statement = ( + "SELECT catalog_name, schema_name, table_name, tag_value " + "FROM system.information_schema.table_tags " + f"WHERE tag_name = {string_literal(PRODUCT_TAG)} " + f"AND tag_value IN ({', '.join(string_literal(p) for p in wanted)})" + ) + doc = runner.execute(statement) + rows = (doc.get("result") or {}).get("data_array") or [] + out: Dict[str, List[Candidate]] = {} + for row in rows: + if not isinstance(row, (list, tuple)) or len(row) < 4: + continue + catalog, schema, table, product = (str(value or "") for value in row[:4]) + full_name = f"{catalog}.{schema}.{table}".lower() + try: + quote_table(full_name) + except ValueError: + continue + live = read_table(client, full_name) + if live is None: + continue + primary, _ = live_keys(live) + out.setdefault(product, []).append(Candidate(full_name, primary[1] if primary else ())) + return out + + +# ── the statements ───────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class Step: + what: str + statement: str + + +@dataclass +class TablePlan: + """What brings one table's keys to the contract, by phase.""" + + full_name: str + exists: bool + fk_drops: List[Step] = field(default_factory=list) + pk_drops: List[Step] = field(default_factory=list) + pk_adds: List[Step] = field(default_factory=list) + fk_adds: List[Step] = field(default_factory=list) + #: Why a key the contract declares is not declared, and what to do. + notes: List[str] = field(default_factory=list) + + @property + def steps(self) -> List[Step]: + return [*self.fk_drops, *self.pk_drops, *self.pk_adds, *self.fk_adds] + + +def _cols(columns: Sequence[str]) -> str: + return ", ".join(quote(c) for c in columns) + + +def primary_key_statement(keys: TableKeys) -> str: + return ( + f"ALTER TABLE {quote_table(keys.full_name)} ADD CONSTRAINT {quote(keys.primary_name)} " + f"PRIMARY KEY ({_cols(keys.primary)})" + ) + + +def foreign_key_statement(full_name: str, fk: ForeignKey) -> str: + return ( + f"ALTER TABLE {quote_table(full_name)} ADD CONSTRAINT {quote(fk.name)} " + f"FOREIGN KEY ({_cols(fk.columns)}) REFERENCES {quote_table(fk.parent)} " + f"({_cols(fk.parent_columns)})" + ) + + +def _drop(full_name: str, name: str, what: str) -> Step: + return Step( + what, f"ALTER TABLE {quote_table(full_name)} DROP CONSTRAINT IF EXISTS {quote(name)}" + ) + + +def plan_table( + full_name: str, + keys: Optional[TableKeys], + fks: Sequence[ForeignKey], + live: Optional[Mapping[str, Any]], +) -> TablePlan: + """The steps from ``live`` (``None``: not created yet) to what the contract declares. + + A key already in place is kept whatever its name. One this module declared (its prefix) + that the contract changed or dropped is dropped; any other is left, and a contract key it + is in the way of is noted. A foreign key of this module whose parent cannot be resolved now + is kept while the contract still references its columns. + """ + plan = TablePlan(full_name, exists=live is not None) + live_pk, live_fks = live_keys(live or {}) + want_pk = keys.primary if keys is not None else () + if live_pk is not None and live_pk[1] != want_pk: + name, cols = live_pk + if name.startswith(PK_PREFIX): + why = "its columns changed" if want_pk else "the contract no longer declares it" + plan.pk_drops.append( + _drop(full_name, name, f"drop primary key {name} ({', '.join(cols)}): {why}") + ) + live_pk = None + elif want_pk: + plan.notes.append( + f"primary key ({', '.join(want_pk)}): the table has PRIMARY KEY {name} " + f"({', '.join(cols)}), which fluid did not declare, and a table holds one. Drop " + f"{name} to let fluid apply declare the contract's." + ) + if keys is not None and want_pk and live_pk is None: + plan.pk_adds.append( + Step( + f"primary key {keys.primary_name} ({', '.join(want_pk)})", + primary_key_statement(keys), + ) + ) + wanted = {fk.columns: fk for fk in fks} + referenced = {r.column for r in keys.references} if keys is not None else set() + for fk in live_fks: + want = wanted.get(fk.columns) + if want is not None and (fk.parent, fk.parent_columns) == ( + want.parent, + want.parent_columns, + ): + del wanted[fk.columns] # in place, whatever its name + continue + if not fk.name.startswith(FK_PREFIX): + if want is not None: + plan.notes.append( + f"foreign key ({', '.join(fk.columns)}) to {want.parent}: the columns " + f"already carry foreign key {fk.name}, which fluid did not declare, " + f"referencing {fk.parent}. Drop {fk.name} to let fluid apply declare the " + "contract's." + ) + del wanted[fk.columns] + continue + if want is None and set(fk.columns) <= referenced: + continue # still referenced, its parent unresolved now: kept + why = "its parent changed" if want is not None else "the contract no longer declares it" + plan.fk_drops.append(_drop(full_name, fk.name, f"drop {fk.describe()}: {why}")) + for fk in wanted.values(): + plan.fk_adds.append(Step(fk.describe(), foreign_key_statement(full_name, fk))) + return plan + + +def _hint(error: str) -> str: + """What to do about a key statement that failed, when the cause is a known one.""" + text = error.lower() + if "is nullable" in text: + return ( + " A key column is nullable on the live table; a dbt model sets its own, so " + "declare not_null on it there (with the model contract enforced)." + ) + if "does not have either a primary key" in text or "do not match the referenced" in text: + return ( + " The parent has no primary key on those columns now; its own apply or build " + "declares it." + ) + if "non empty child constraints" in text: + return ( + " Foreign keys on other tables reference it: remove their foreign-key tags and " + "apply those products first." + ) + return "" + + +# ── the hook ─────────────────────────────────────────────────────────── + + +def _session(factory: Optional[Callable[[], Any]]) -> Union[Tuple[Any, SqlRunner], str]: + """A client and the warehouse runner, or why there are none.""" + try: + client = (factory or default_client)() + except ImportError: + return f"databricks-sdk is not installed ({DATABRICKS_EXTRA_INSTALL})" + except Exception as exc: # noqa: BLE001 — a bad credential must not become a traceback + return f"no Databricks client ({type(exc).__name__}: {scrub(exc)})" + warehouse = str(os.environ.get(WAREHOUSE_ENV) or "").strip() + if not warehouse: + return f"{WAREHOUSE_ENV} is not set, and the keys are declared on that warehouse" + return client, SqlRunner(client, warehouse) + + +def resolve( + keyed: Sequence[TableKeys], + contract_id: str, + runner: SqlRunner, + client: Any, + notes: List[str], +) -> Dict[str, List[Candidate]]: + """Every table the contract's references can resolve to: its own exposes' declared + primary keys, and the live tables of each other product it references.""" + candidates: Dict[str, List[Candidate]] = dict(own_candidates(contract_id, keyed)) + others = sorted({r.product for k in keyed for r in k.references if r.product != contract_id}) + if not others: + return candidates + try: + candidates.update(lookup(runner, client, others)) + except Exception as exc: # noqa: BLE001 — a lookup that fails declares no such key + notes.append( + f"foreign keys to {', '.join(others)}: their tables could not be looked up " + f"({type(exc).__name__}: {scrub(exc)})." + ) + return candidates + + +@dataclass(frozen=True) +class _Read: + target: KeyTarget + live: Optional[Mapping[str, Any]] + + +def reconcile( + targets: Sequence[KeyTarget], + contract_id: str, + say: Say, + client_factory: Optional[Callable[[], Any]] = None, + *, + dry_run: bool = False, +) -> None: + """Declare the contract's keys on its live tables, and replace or drop the keys this + module declared that the contract changed or no longer asks for. With ``dry_run``, print + the statements and run none. + + Never raises on a key: one that is not declared is said, with what to do. A contract that + declares no key prints nothing unless a key this module declared earlier is dropped. + """ + if not targets: + return + declared = any(t.keys is not None for t in targets) + session = _session(client_factory) + if isinstance(session, str): + if declared: + say(f"\n Unity Catalog keys: none declared, {session}.") + return + client, runner = session + general: List[str] = [] + reads: List[_Read] = [] + for target in targets: + try: + live = read_table(client, target.full_name) + except Exception as exc: # noqa: BLE001 — one unreadable table must not stop the rest + if target.keys is not None: + general.append( + f"keys of {target.full_name}: it could not be read ({type(exc).__name__}: " + f"{scrub(exc)})." + ) + continue + if live is None and not dry_run: + if target.keys is not None: + general.append(f"keys of {target.full_name}: the table does not exist.") + continue + reads.append(_Read(target, live)) + keyed = [r.target.keys for r in reads if r.target.keys is not None] + candidates = resolve(keyed, contract_id, runner, client, general) if keyed else {} + plans: List[TablePlan] = [] + for read in reads: + keys = read.target.keys + notes: List[str] = [] + fks = foreign_keys(keys, candidates, notes, own_id=contract_id) if keys else [] + plan = plan_table(read.target.full_name, keys, fks, read.live) + plan.notes[:0] = notes + if plan.steps or plan.notes: + plans.append(plan) + if not plans and not general: + return + say( + "\n Unity Catalog keys (run on the SQL warehouse" + + (" after the apply" if dry_run else "") + + ", not by OpenTofu):" + ) + for line in general: + say(f" not declared: {line}") + if dry_run: + for plan in plans: + say(f" {plan.full_name}" + ("" if plan.exists else " (not created yet)")) + for step in plan.steps: + say(f" {step.what}") + say(f" {step.statement}") + for note in plan.notes: + say(f" not declared: {note}") + if any(plan.steps for plan in plans): + say(" dry-run: none of these statements ran.") + return + _run(plans, runner, say) + + +def _run(plans: Sequence[TablePlan], runner: SqlRunner, say: Say) -> None: + """Every plan's steps, phase by phase across the tables: a foreign key is dropped before + the primary key it references, and a primary key added before a foreign key to it.""" + lines: Dict[str, List[str]] = {p.full_name: [] for p in plans} + stuck: set = set() # tables whose old primary key could not be dropped + for phase in ("fk_drops", "pk_drops", "pk_adds", "fk_adds"): + for plan in plans: + for step in getattr(plan, phase): + if phase == "pk_adds" and plan.full_name in stuck: + lines[plan.full_name].append( + f" not declared: {step.what}: the old primary key is still there." + ) + continue + try: + runner.run(step.statement) + except StatementFailed as exc: + if phase == "pk_drops": + stuck.add(plan.full_name) + lines[plan.full_name].append( + f" failed: {step.what}: {exc}.{_hint(str(exc))} The statement was: " + f"{step.statement}" + ) + continue + lines[plan.full_name].append(f" done: {step.what}") + for plan in plans: + say(f" {plan.full_name}") + for line in lines[plan.full_name]: + say(line) + for note in plan.notes: + say(f" not declared: {note}") + + +# ── after a dbt build ────────────────────────────────────────────────── + + +def rebuild_script(entries: Sequence[Tuple[TableKeys, Sequence[ForeignKey]]]) -> str: + """The SQL file a dbt job's ``fluid_keys`` task runs once the build has succeeded. + + Every primary key first, then every foreign key, each inside one script whose + ``CONTINUE`` handler passes over a statement that fails, so the task, and the run, never + fail on a key. A job whose tables resolve to no statement still gets a valid file. + """ + body = [primary_key_statement(keys) for keys, _ in entries if keys.primary] + body += [foreign_key_statement(keys.full_name, fk) for keys, fks in entries for fk in fks] + header = [ + "-- Written by fluid apply (data-product-forge); the next apply rewrites it.", + "-- A dbt table build replaces its table with CREATE OR REPLACE, which drops the table's", + "-- Unity Catalog keys and the foreign keys that reference it. This task, run after a", + "-- build that succeeded, declares this job's keys again. A statement that fails (a key", + "-- still in place, a parent whose own key is not back yet) is passed over: no run fails", + "-- on a key.", + ] + if not body: + return "\n".join([*header, "SELECT 'no key to declare' AS fluid_keys;", ""]) + return "\n".join( + [ + *header, + "BEGIN", + " DECLARE CONTINUE HANDLER FOR SQLEXCEPTION BEGIN END;", + *(f" {statement};" for statement in body), + "END;", + "", + ] + ) + + +__all__ = [ + "Candidate", + "ForeignKey", + "KeyTarget", + "REBUILD_FILE", + "REBUILD_TASK", + "TableKeys", + "compile_keys", + "foreign_keys", + "lookup", + "own_candidates", + "plan_table", + "rebuild_script", + "reconcile", + "resolve", +] diff --git a/tests/iac/_fake_databricks.py b/tests/iac/_fake_databricks.py index 70b51155..a8139499 100644 --- a/tests/iac/_fake_databricks.py +++ b/tests/iac/_fake_databricks.py @@ -17,7 +17,13 @@ refuses it), ALTER TABLE SET / DROP ROW FILTER, ALTER COLUMN SET / DROP MASK, and REVOKE SELECT. The bindings are served on the table (``row_filter``, ``columns[].mask``) as Unity Catalog serves them, and a CREATE OR REPLACE TABLE keeps them (Databricks docs: REPLACE preserves row filters -and column masks). ``fail_on(substring, message)`` makes every statement containing it FAIL, +and column masks). The key constraints forge-cli declares (``databricks_keys.py``): ALTER TABLE +ADD CONSTRAINT ... PRIMARY KEY / FOREIGN KEY ... REFERENCES and DROP CONSTRAINT [IF EXISTS], +served as ``table_constraints``, with the rules and error texts a warehouse showed (a key needs +NOT NULL columns, a foreign key its parent's primary key, a referenced primary key does not +drop), and the ``fluid_product_id`` lookup on ``system.information_schema.table_tags``. A CREATE +OR REPLACE TABLE drops the table's keys and the foreign keys that reference it, and a DROP +COLUMN the keys on that column, as they do live. ``fail_on(substring, message)`` makes every statement containing it FAIL, for the fail-safe tests; ``fail_route(method, path_part)`` answers matching requests with an error, to fail a ``tofu apply`` part way; ``timeline`` orders statements against permission changes. Any other statement FAILS (state ``FAILED``) and stays visible in ``sql_log``, so a new @@ -69,7 +75,7 @@ import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path -from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Set, Tuple from urllib.parse import parse_qs, unquote, urlparse NOT_FOUND = "NOT_FOUND" @@ -258,6 +264,169 @@ def _new_table(cols: List[Dict[str, Any]], tail: Dict[str, Any], marker: int) -> "table_type": "MANAGED", "created_at": int(time.time() * 1000), "marker": marker, + "constraints": [], + } + + +# ---------------------------------------------------------------------- key constraints +# As measured on a Databricks SQL warehouse (dbsql 2026.45, 2026-11-12; see +# ``fluid_build/iac/providers/databricks_keys.py``): the error texts below are the live ones. +def _constraint_action( + fake: "FakeDatabricks", full: str, t: Dict[str, Any], action: str +) -> Optional[Tuple[bool, str]]: + """``ALTER TABLE `` ADD CONSTRAINT ... PRIMARY KEY / FOREIGN KEY and DROP CONSTRAINT; + ``None`` for any other action.""" + m = re.match( + r"ADD\s+CONSTRAINT\s+`?([^`\s]+)`?\s+PRIMARY\s+KEY\s*\(([^)]*)\)\s*$", action, re.I + ) + if m: + return _add_primary_key(fake, full, t, m.group(1), _names(m.group(2))) + m = re.match( + r"ADD\s+CONSTRAINT\s+`?([^`\s]+)`?\s+FOREIGN\s+KEY\s*\(([^)]*)\)\s+REFERENCES\s+" + rf"{_FULL}\s*\(([^)]*)\)\s*$", + action, + re.I, + ) + if m: + cname, inner, pcat, psch, pname, pinner = m.groups() + parent = f"{pcat}.{psch}.{pname}" + return _add_foreign_key(fake, full, t, cname, _names(inner), parent, _names(pinner)) + m = re.match(r"DROP\s+CONSTRAINT\s+(IF\s+EXISTS\s+)?`?([^`\s]+)`?\s*$", action, re.I) + if m: + return _drop_constraint(fake, full, t, m.group(2), bool(m.group(1))) + return None + + +def _schema_constraint_names(fake: "FakeDatabricks", full: str) -> Set[str]: + schema = full.rsplit(".", 1)[0] + "." + return { + c["name"] + for name, table in fake.tables.items() + if name.startswith(schema) + for c in table.get("constraints") or [] + } + + +def _add_primary_key( + fake: "FakeDatabricks", full: str, t: Dict[str, Any], name: str, cols: List[str] +) -> Tuple[bool, str]: + have = {c["name"]: c for c in t["columns"]} + missing = [c for c in cols if c not in have] + if missing: + return False, f"[UNRESOLVED_COLUMN] {missing[0]}" + if name in _schema_constraint_names(fake, full): + return False, f"[DELTA_CONSTRAINT_ALREADY_EXISTS] Constraint '{name}' already exists." + held = next((c for c in t["constraints"] if c["type"] == "PRIMARY KEY"), None) + if held is not None: + return ( + False, + "[RESOURCE_ALREADY_EXISTS.UC_CHILD_COLUMNS_ALREADY_EXIST_IN_PRIMARY_KEY_OR_UNIQUE_" + f"CONSTRAINT] Failed to create PRIMARY KEY constraint {name}: table already has a " + f"PRIMARY KEY constraint: {held['name']}.", + ) + loose = [c for c in cols if have[c]["nullable"]] + if loose: + return ( + False, + f"Cannot create the primary key `{name}` because its child column(s) `{loose[0]}` " + "is nullable. Please change the column nullability and retry.", + ) + t["constraints"].append({"type": "PRIMARY KEY", "name": name, "columns": cols}) + return True, "" + + +def _add_foreign_key( + fake: "FakeDatabricks", + full: str, + t: Dict[str, Any], + name: str, + cols: List[str], + parent: str, + parent_cols: List[str], +) -> Tuple[bool, str]: + missing = [c for c in cols if c not in {col["name"] for col in t["columns"]}] + if missing: + return False, f"[UNRESOLVED_COLUMN] {missing[0]}" + target = fake.tables.get(parent) + if target is None: + return False, f"[TABLE_OR_VIEW_NOT_FOUND] The table or view {parent} cannot be found." + pk = next((c for c in target["constraints"] if c["type"] == "PRIMARY KEY"), None) + if pk is None: + return ( + False, + "[NOT_FOUND.UC_REFERENTIAL_CONSTRAINT_DOES_NOT_EXIST] Table " + f"`{parent}` does not have either a primary key constraint or a UNIQUE constraint.", + ) + if list(pk["columns"]) != list(parent_cols): + return ( + False, + "The foreign key parent columns do not match the referenced primary key child " + f"columns or any unique constraint columns. Foreign key parent columns are " + f"({', '.join(parent_cols)}).", + ) + if name in _schema_constraint_names(fake, full): + return False, f"[DELTA_CONSTRAINT_ALREADY_EXISTS] Constraint '{name}' already exists." + same = next( + (c for c in t["constraints"] if c["type"] == "FOREIGN KEY" and c["columns"] == cols), None + ) + if same is not None: + return ( + False, + f"[CONSTRAINT_ALREADY_EXISTS_WITH_SERVER_MESSAGE] Failed to create foreign key " + f"constraint `{name}`: table `{full}` already has a foreign key constraint: " + f"`{same['name']}`, that has the same set of child columns. SQLSTATE: 42710", + ) + t["constraints"].append( + { + "type": "FOREIGN KEY", + "name": name, + "columns": cols, + "parent": parent, + "parent_columns": parent_cols, + } + ) + return True, "" + + +def _drop_constraint( + fake: "FakeDatabricks", full: str, t: Dict[str, Any], name: str, if_exists: bool +) -> Tuple[bool, str]: + held = next((c for c in t["constraints"] if c["name"] == name), None) + if held is None: + if if_exists: + return True, "" + return False, f"[DELTA_CONSTRAINT_DOES_NOT_EXIST] Cannot drop nonexistent constraint {name}" + if held["type"] == "PRIMARY KEY" and _referencing(fake, full): + return ( + False, + f"Cannot delete constraint `{name}` with non empty child constraints. Set the " + "argument CASCADE to true to delete both this constraint and all child constraints.", + ) + t["constraints"].remove(held) + return True, "" + + +def _referencing(fake: "FakeDatabricks", full: str) -> List[Tuple[str, Dict[str, Any]]]: + """``(table, foreign key)`` for every foreign key that references ``full``.""" + return [ + (name, c) + for name, table in fake.tables.items() + for c in table.get("constraints") or [] + if c["type"] == "FOREIGN KEY" and c["parent"] == full + ] + + +def _constraint_doc(c: Dict[str, Any]) -> Dict[str, Any]: + """One ``table_constraints[]`` entry, as the Tables API serves it.""" + if c["type"] == "PRIMARY KEY": + return {"primary_key_constraint": {"name": c["name"], "child_columns": list(c["columns"])}} + return { + "foreign_key_constraint": { + "name": c["name"], + "child_columns": list(c["columns"]), + "parent_table": c["parent"], + "parent_columns": list(c["parent_columns"]), + } } @@ -302,6 +471,10 @@ def _execute(fake: "FakeDatabricks", sql: str) -> Tuple[bool, str]: for col in new["columns"]: if col["name"] in kept: col["mask"] = kept[col["name"]] + # Its keys are not kept, and every foreign key elsewhere that referenced it goes + # too (measured on a warehouse, 2026-11-12). + for _child, fk in _referencing(fake, full): + fake.tables[_child]["constraints"].remove(fk) fake.tables[full] = new fake.table_rows.pop(full, None) # a new table holds no rows return True, "" @@ -459,6 +632,10 @@ def find(col: str) -> int: cols[idx].pop("mask") return True, "" + keyed = _constraint_action(fake, full, t, action) + if keyed is not None: + return keyed + m = re.match(r"ADD\s+COLUMN\s+(.*)$", action, re.S | re.I) if m: body = m.group(1).strip() @@ -498,6 +675,8 @@ def find(col: str) -> int: if missing and not ie: return False, f"[UNRESOLVED_COLUMN] {missing[0]}" t["columns"] = [c for c in cols if c["name"] not in drop] + # A key on a dropped column goes with it (measured, 2026-11-12). + t["constraints"] = [c for c in t["constraints"] if not set(c["columns"]) & set(drop)] return True, "" m = re.match(r"RENAME\s+COLUMN\s+`?([^`\s]+)`?\s+TO\s+`?([^`\s]+)`?\s*$", action, re.I) @@ -574,6 +753,45 @@ def _count_doc(n: int) -> Dict[str, Any]: } +def _table_tags_query(fake: "FakeDatabricks", s: str) -> Optional[Tuple[bool, str, Dict[str, Any]]]: + """The tables carrying a tag with one of the given values, from + ``system.information_schema.table_tags`` (the parent lookup of ``databricks_keys.py``).""" + m = re.match( + r"SELECT\s+catalog_name\s*,\s*schema_name\s*,\s*table_name\s*,\s*tag_value\s+FROM\s+" + rf"system\.information_schema\.table_tags\s+WHERE\s+tag_name\s*=\s*{_QUOTED}\s+AND\s+" + r"tag_value\s+IN\s*\((.*)\)\s*$", + s, + re.S | re.I, + ) + if m is None: + return None + key = _unescape(m.group(1)) + values = {_unescape(v) for v in re.findall(_QUOTED, m.group(2))} + rows = [ + [*name.split("."), value] + for (kind, name, tag), value in sorted(fake.tags.items()) + if kind == "tables" and tag == key and value in values and name in fake.tables + ] + return True, "", _rows_doc(("catalog_name", "schema_name", "table_name", "tag_value"), rows) + + +def _rows_doc(names: Sequence[str], rows: List[List[str]]) -> Dict[str, Any]: + """The result of a query of string columns, as the Statement Execution API sends it.""" + return { + "manifest": { + "format": "JSON_ARRAY", + "schema": { + "column_count": len(names), + "columns": [ + {"name": n, "type_name": "STRING", "position": i} for i, n in enumerate(names) + ], + }, + "total_row_count": len(rows), + }, + "result": {"row_count": len(rows), "data_array": rows}, + } + + def _insert_doc(n: int, report: bool) -> Dict[str, Any]: """What ``INSERT`` answers: ``num_affected_rows`` and ``num_inserted_rows`` (Databricks SQL).""" if not report: @@ -671,6 +889,9 @@ def _data_statement(fake: "FakeDatabricks", sql: str) -> Optional[Tuple[bool, st for row in fetched ] return True, "", _insert_doc(len(fetched), fake.insert_reports_count) + tagged = _table_tags_query(fake, s) + if tagged is not None: + return tagged m = re.match(rf"SELECT\s+COUNT\(\s*(?:\*|1)\s*\)\s+FROM\s+{_FULL}\s*$", s, re.I) if m: full = ".".join(m.groups()) @@ -974,6 +1195,17 @@ def tags_on(self, entity_type: str, full_name: str) -> Dict[str, str]: """The tags an entity holds, ``{key: value}``.""" return {k: v for (t, n, k), v in self.tags.items() if t == entity_type and n == full_name} + def keys_on(self, full_name: str) -> List[Tuple[str, ...]]: + """A table's keys: ``("PRIMARY KEY", name, cols)`` / ``("FOREIGN KEY", name, cols, + parent, parent_cols)``, columns comma-joined.""" + out: List[Tuple[str, ...]] = [] + for c in self.tables[full_name]["constraints"]: + row: Tuple[str, ...] = (c["type"], c["name"], ",".join(c["columns"])) + if c["type"] == "FOREIGN KEY": + row += (c["parent"], ",".join(c["parent_columns"])) + out.append(row) + return out + def grant_catalog(self, catalog: str, principal: str, *privileges: str) -> None: """Seed catalog-level privileges (e.g. ``USE_CATALOG``); schemas and tables inherit them.""" self.grant("catalog", catalog, principal, *privileges) @@ -1284,6 +1516,8 @@ def _table_doc(self, full: str) -> Dict[str, Any]: col["partition_index"] = t["partitions"].index(c["name"]) cols.append(col) extra = {"row_filter": t["row_filter"]} if t.get("row_filter") else {} + if t.get("constraints"): + extra["table_constraints"] = [_constraint_doc(c) for c in t["constraints"]] return { **extra, "name": name, diff --git a/tests/iac/test_iac_databricks_fake_e2e.py b/tests/iac/test_iac_databricks_fake_e2e.py index d750b77e..4cc59b97 100644 --- a/tests/iac/test_iac_databricks_fake_e2e.py +++ b/tests/iac/test_iac_databricks_fake_e2e.py @@ -966,3 +966,47 @@ def test_a_dbt_job_created_with_its_restricted_table_is_paused_until_the_next_ap assert "is created paused" not in out assert "tofu plan: +0 ~0 -0" in apply_contract(capfd, path, "--mode", "dry-run")[1] + + +# ── primary and foreign keys ──────────────────────────────────────────────── +def _keyed(product_id: str = "dbx.fake.orders", table: str = "orders", ref: str = "") -> dict: + """``orders`` with ``id`` its primary key; with ``ref``, ``region`` a foreign key to it.""" + contract = make_contract(product_id=product_id) + expose = contract["exposes"][0] + expose["exposeId"] = table + expose["binding"]["location"]["table"] = table + schema = expose["contract"]["schema"] + schema[0]["tags"] = ["primary-key"] + if ref: + schema[2].update(tags=["foreign-key"], labels={"references": ref}) + return contract + + +@requires_sdk +def test_apply_declares_the_keys_and_a_child_finds_its_parent_by_its_product_tag( + workspace, fake, capfd +): + parent = write_contract(workspace, _keyed()) + rc, out = apply_contract(capfd, parent) + assert rc == 0, out + assert fake.keys_on(TABLE) == [("PRIMARY KEY", "fluid_pk_orders", "id")] + assert "done: primary key fluid_pk_orders (id)" in out + + # The child is another contract: its parent is found by the fluid_product_id tag + # tofu put on the parent's table, not by any assumption about where it lives. + (workspace / "child").mkdir() + child = _keyed("dbx.fake.returns", "returns", ref="dbx.fake.orders.id") + child["exposes"][0]["contract"]["schema"][2]["type"] = "bigint" + child_path = write_contract(workspace / "child", child) + rc, out = apply_contract(capfd, child_path) + assert rc == 0, out + assert fake.keys_on("main.sales.returns") == [ + ("PRIMARY KEY", "fluid_pk_returns", "id"), + ("FOREIGN KEY", "fluid_fk_returns_region", "region", TABLE, "id"), + ] + + statements = len(fake.sql_log) + rc, out = apply_contract(capfd, child_path) + assert rc == 0 and "tofu plan: +0 ~0 -0" in out, out + assert not [s for s in fake.sql_log[statements:] if "CONSTRAINT" in s] + assert fake.unknown == [] diff --git a/tests/iac/test_iac_databricks_jobs_fake_e2e.py b/tests/iac/test_iac_databricks_jobs_fake_e2e.py index 9c3f22e6..d7050592 100644 --- a/tests/iac/test_iac_databricks_jobs_fake_e2e.py +++ b/tests/iac/test_iac_databricks_jobs_fake_e2e.py @@ -241,3 +241,29 @@ def test_the_table_a_dbt_build_writes_keeps_what_dbt_made_and_the_drift_is_said( assert "amount: contract=double live=decimal(38,6)" in out assert "loaded_at: present in the live table but not declared" in out assert fake.sql_log[statements:] == [] and fake.tables[TABLE]["marker"] == marker + + +def test_a_job_that_writes_a_keyed_table_declares_its_keys_again_after_dbt(workspace, fake, capfd): + document = job_contract() + document["exposes"][0]["contract"]["schema"][0]["tags"] = ["primary-key"] + path = write_contract(workspace, document) + + rc, out = apply_contract(capfd, path) + assert rc == 0, out + settings = _only_job(fake)["settings"] + dbt_task, keys_task = settings["tasks"] + assert dbt_task["task_key"] == "dbt" + assert keys_task["task_key"] == "fluid_keys" + assert keys_task["depends_on"] == [{"task_key": "dbt"}] + assert keys_task["sql_task"] == { + "warehouse_id": fake.warehouse_id, + "file": {"path": f"/Workspace{ROOT}/fluid_keys.sql", "source": "WORKSPACE"}, + } + script = fake.file_content(f"{ROOT}/fluid_keys.sql").decode("utf-8") + assert "DECLARE CONTINUE HANDLER FOR SQLEXCEPTION" in script + assert ( + "ALTER TABLE `main`.`sales`.`orders` ADD CONSTRAINT `fluid_pk_orders` PRIMARY KEY (`id`)" + in script + ) + # The provider reads back the task and the file as it sent them. + assert "tofu plan: +0 ~0 -0" in apply_contract(capfd, path, "--mode", "dry-run")[1] diff --git a/tests/iac/test_iac_databricks_keys.py b/tests/iac/test_iac_databricks_keys.py new file mode 100644 index 00000000..2083ecab --- /dev/null +++ b/tests/iac/test_iac_databricks_keys.py @@ -0,0 +1,663 @@ +# Copyright 2025 Agentics Transformation Limited +# SPDX-License-Identifier: Apache-2.0 + +"""Unity Catalog primary and foreign keys from the contract (``databricks_keys.py``). + +Pure parts first (what a contract declares, where a foreign key points, the statements from a +live table, the script a dbt job runs after its build), then the jobs' ``fluid_keys`` task, then +the plugin's hooks against ``_fake_databricks.FakeDatabricks``, whose key rules and error texts +are the ones a warehouse showed. No ``tofu``, no network. +""" + +from __future__ import annotations + +import base64 +import json +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional + +import pytest + +from fluid_build.iac import IAC_PLUGINS +from fluid_build.iac.providers import databricks as dbx +from fluid_build.iac.providers import databricks_keys as keys +from fluid_build.iac.providers.databricks_enforce import SqlRunner +from fluid_build.iac.source import contract_source + +from ._fake_databricks import FakeDatabricks + +pytestmark = [pytest.mark.unit, pytest.mark.databricks] + +PLUGIN = IAC_PLUGINS["databricks"] +WAREHOUSE = "wh-0123456789abcdef" +SITES = "geo.sites_v1" +TICKETS = "ops.tickets_v1" + + +def col(name: str, kind: str = "string", **kw: Any) -> Dict[str, Any]: + """A column; ``pk=True`` tags it primary-key (and required), ``ref=`` foreign-key.""" + out: Dict[str, Any] = {"name": name, "type": kind} + tags: List[str] = [] + if kw.pop("pk", False): + tags.append("primary-key") + out["required"] = True + ref = kw.pop("ref", None) + if ref is not None: + tags.append("foreign-key") + out["labels"] = {"references": ref} + if tags: + out["tags"] = tags + out.update(kw) + return out + + +def expose(table: str, columns: List[Dict[str, Any]], schema: str, catalog: str = "main"): + return { + "exposeId": table, + "kind": "table", + "binding": { + "platform": "databricks", + "format": "delta_table", + "location": {"database": catalog, "schema": schema, "table": table}, + }, + "contract": {"schema": columns}, + } + + +def product(pid: str, *exposes: Dict[str, Any], **extra: Any) -> Dict[str, Any]: + return { + "fluidVersion": "0.7.6", + "kind": "DataProduct", + "id": pid, + "name": pid, + "exposes": list(exposes), + **extra, + } + + +def sites(**kw: Any) -> Dict[str, Any]: + return product(SITES, expose("sites", [col("site_id", pk=True), col("name")], "geo"), **kw) + + +def tickets(ref: str = f"{SITES}.site_id", **kw: Any) -> Dict[str, Any]: + columns = [col("ticket_id", pk=True), col("site_id", ref=ref), col("note")] + return product(TICKETS, expose("tickets", columns, "ops"), **kw) + + +def compiled(contract: Dict[str, Any], **kw: Any) -> dbx._Compiled: + return dbx._compile(contract, {dbx.WAREHOUSE_ENV: WAREHOUSE}, **kw) + + +def table_keys(contract: Dict[str, Any]) -> keys.TableKeys: + (table,) = compiled(contract).tables + assert table.keys is not None + return table.keys + + +# ── what a contract declares ─────────────────────────────────────────────── +class TestCompile: + def test_tags_and_labels_become_a_primary_key_and_a_reference(self): + k = table_keys(tickets()) + assert k.full_name == "main.ops.tickets" + assert k.primary == ("ticket_id",) + assert k.references == (keys.Reference("site_id", SITES, "site_id"),) + assert k.primary_name == "fluid_pk_tickets" + + def test_a_composite_key_keeps_schema_order_and_either_tag_spelling(self): + columns = [ + col("day", "date", required=True, tags=["primaryKey"]), + col("v"), + col("site_id", pk=True), + ] + k = table_keys(product("x.y_v1", expose("t", columns, "s"))) + assert k.primary == ("day", "site_id") + + def test_a_key_column_that_is_not_required_gives_no_key_and_says_why(self): + columns = [col("id", tags=["primary-key"]), col("v")] + out = compiled(product("x.y_v1", expose("t", columns, "s"))) + assert out.tables[0].keys is None + (warning,) = [w for w in out.warnings if "primary-key" in w] + assert "required: true" in warning and "Set required: true on id" in warning + + @pytest.mark.parametrize("ref", ["", "site_id", "geo.sites_v1.", ".site_id", "x.bad-col"]) + def test_a_foreign_key_tag_without_a_usable_reference_is_a_warning(self, ref): + columns = [col("id", pk=True), col("site_id", tags=["foreign-key"], labels={})] + if ref: + columns[1]["labels"] = {"references": ref} + out = compiled(product("x.y_v1", expose("t", columns, "s"))) + assert out.tables[0].keys.references == () + assert any("labels.references" in w for w in out.warnings) + + def test_a_table_without_key_tags_has_no_keys(self): + out = compiled(product("x.y_v1", expose("t", [col("a"), col("b")], "s"))) + assert out.tables[0].keys is None + assert not [w for w in out.warnings if "key" in w] + + def test_a_reference_to_its_own_product_that_no_expose_takes_is_a_validate_warning(self): + columns = [col("id", pk=True), col("parent_id", ref="x.y_v1.other")] + out = compiled(product("x.y_v1", expose("t", columns, "s"))) + (warning,) = [w for w in out.warnings if "foreign key" in w] + assert warning.startswith("exposes[0] (t): foreign key (parent_id) to x.y_v1") + + def test_the_product_tag_is_the_one_the_tables_carry(self): + assert keys.PRODUCT_TAG in dict(dbx._PRODUCT_TAGS) + assert dict(dbx._PRODUCT_TAGS)[keys.PRODUCT_TAG] == "id" + + def test_a_long_name_is_cut_with_a_digest(self): + name = keys.constraint_name(keys.FK_PREFIX, "t" * 200, "c" * 100) + assert len(name) == 255 and name.startswith("fluid_fk_ttt") + assert name != keys.constraint_name(keys.FK_PREFIX, "t" * 200, "c" * 101) + + +# ── where a foreign key points ───────────────────────────────────────────── +def _keys(*refs: keys.Reference, table: str = "main.ops.tickets") -> keys.TableKeys: + return keys.TableKeys(table, ("ticket_id",), tuple(refs)) + + +class TestResolve: + REF = keys.Reference("site_id", SITES, "site_id") + + def test_one_table_with_the_referenced_primary_key(self): + cands = {SITES: [keys.Candidate("main.geo.sites", ("site_id",))]} + (fk,) = keys.foreign_keys(_keys(self.REF), cands) + assert fk == keys.ForeignKey( + "fluid_fk_tickets_site_id", ("site_id",), "main.geo.sites", ("site_id",) + ) + + def test_columns_that_cover_a_composite_key_form_one_in_its_order(self): + refs = ( + keys.Reference("s", "p.q_v1", "site_id"), + keys.Reference("d", "p.q_v1", "day"), + ) + cands = {"p.q_v1": [keys.Candidate("main.p.q", ("day", "site_id"))]} + (fk,) = keys.foreign_keys(_keys(*refs), cands) + assert fk.columns == ("d", "s") and fk.parent_columns == ("day", "site_id") + assert fk.name == "fluid_fk_tickets_d_s" + + def test_two_columns_to_one_key_are_two_foreign_keys(self): + refs = ( + keys.Reference("from_site", SITES, "site_id"), + keys.Reference("to_site", SITES, "site_id"), + ) + cands = {SITES: [keys.Candidate("main.geo.sites", ("site_id",))]} + out = keys.foreign_keys(_keys(*refs), cands) + assert [fk.columns for fk in out] == [("from_site",), ("to_site",)] + + def test_a_tie_goes_to_the_table_in_the_childs_catalog(self): + cands = { + SITES: [ + keys.Candidate("dev.geo.sites", ("site_id",)), + keys.Candidate("main.geo.sites", ("site_id",)), + ] + } + (fk,) = keys.foreign_keys(_keys(self.REF), cands) + assert fk.parent == "main.geo.sites" + + def test_a_tie_no_catalog_breaks_is_skipped_and_said(self): + cands = { + SITES: [ + keys.Candidate("a.geo.sites", ("site_id",)), + keys.Candidate("b.geo.sites", ("site_id",)), + ] + } + notes: List[str] = [] + assert keys.foreign_keys(_keys(self.REF), cands, notes) == [] + assert "more than one" in notes[0] and "a.geo.sites, b.geo.sites" in notes[0] + + def test_a_product_not_deployed_is_said_with_what_to_do(self): + notes: List[str] = [] + assert keys.foreign_keys(_keys(self.REF), {}, notes) == [] + assert notes == [ + f"foreign key (site_id) to {SITES}: no table tagged fluid_product_id = " + f"'{SITES}' is visible to this identity. Apply {SITES} first, then run fluid " + "apply here again." + ] + + def test_a_parent_without_that_primary_key_is_said(self): + notes: List[str] = [] + cands = {SITES: [keys.Candidate("main.geo.sites", ())]} + assert keys.foreign_keys(_keys(self.REF), cands, notes) == [] + assert "none has PRIMARY KEY (site_id)" in notes[0] + + +# ── the statements from a live table ─────────────────────────────────────── +FK = keys.ForeignKey("fluid_fk_tickets_site_id", ("site_id",), "main.geo.sites", ("site_id",)) + + +def _live(*constraints: Dict[str, Any]) -> Dict[str, Any]: + return {"name": "tickets", "table_constraints": list(constraints)} + + +def _pk(name: str, *cols: str) -> Dict[str, Any]: + return {"primary_key_constraint": {"name": name, "child_columns": list(cols)}} + + +def _fk(name: str, col_: str, parent: str = "main.geo.sites") -> Dict[str, Any]: + return { + "foreign_key_constraint": { + "name": name, + "child_columns": [col_], + "parent_table": parent, + "parent_columns": ["site_id"], + } + } + + +class TestPlan: + K = _keys(TestResolve.REF) + + def test_a_new_table_gets_its_primary_key_then_its_foreign_keys(self): + plan = keys.plan_table("main.ops.tickets", self.K, [FK], None) + assert [s.statement for s in plan.steps] == [ + "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT `fluid_pk_tickets` PRIMARY KEY " + "(`ticket_id`)", + "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT `fluid_fk_tickets_site_id` " + "FOREIGN KEY (`site_id`) REFERENCES `main`.`geo`.`sites` (`site_id`)", + ] + assert plan.exists is False + + def test_keys_in_place_under_any_name_are_kept(self): + live = _live(_pk("their_pk", "ticket_id"), _fk("their_fk", "site_id")) + plan = keys.plan_table("main.ops.tickets", self.K, [FK], live) + assert plan.steps == [] and plan.notes == [] + + def test_a_changed_key_of_ours_is_dropped_and_added(self): + live = _live( + _pk("fluid_pk_tickets", "note"), + _fk("fluid_fk_tickets_site_id", "site_id", "main.old.sites"), + ) + plan = keys.plan_table("main.ops.tickets", self.K, [FK], live) + assert [s.what.split(" ")[0] for s in plan.steps] == ["drop", "drop", "primary", "foreign"] + assert "its columns changed" in plan.pk_drops[0].what + assert "its parent changed" in plan.fk_drops[0].what + + def test_a_key_someone_else_declared_is_left_and_said(self): + live = _live(_pk("admin_pk", "note"), _fk("admin_fk", "site_id", "main.x.y")) + plan = keys.plan_table("main.ops.tickets", self.K, [FK], live) + assert plan.steps == [] + assert "admin_pk" in plan.notes[0] and "fluid did not declare" in plan.notes[0] + assert "admin_fk" in plan.notes[1] + + def test_keys_of_ours_the_contract_dropped_are_dropped_and_no_other(self): + live = _live( + _pk("fluid_pk_tickets", "ticket_id"), + _fk("fluid_fk_tickets_old", "note"), + _fk("admin_fk", "x"), + ) + plan = keys.plan_table("main.ops.tickets", None, [], live) + assert [s.statement for s in plan.steps] == [ + "ALTER TABLE `main`.`ops`.`tickets` DROP CONSTRAINT IF EXISTS `fluid_fk_tickets_old`", + "ALTER TABLE `main`.`ops`.`tickets` DROP CONSTRAINT IF EXISTS `fluid_pk_tickets`", + ] + + def test_a_foreign_key_of_ours_whose_parent_is_unresolved_now_is_kept(self): + live = _live( + _pk("fluid_pk_tickets", "ticket_id"), _fk("fluid_fk_tickets_site_id", "site_id") + ) + plan = keys.plan_table("main.ops.tickets", self.K, [], live) + assert plan.steps == [] + + +# ── after a dbt build ────────────────────────────────────────────────────── +class TestRebuildScript: + def test_every_primary_key_then_every_foreign_key_inside_one_handler(self): + child = _keys(TestResolve.REF) + parent = keys.TableKeys("main.geo.sites", ("site_id",)) + script = keys.rebuild_script([(child, [FK]), (parent, [])]) + body = script.split("BEGIN\n", 1)[1] + assert body.splitlines() == [ + " DECLARE CONTINUE HANDLER FOR SQLEXCEPTION BEGIN END;", + " " + keys.primary_key_statement(child) + ";", + " " + keys.primary_key_statement(parent) + ";", + " " + keys.foreign_key_statement("main.ops.tickets", FK) + ";", + "END;", + ] + assert script.startswith("-- Written by fluid apply") + + def test_nothing_to_declare_is_still_a_valid_file(self): + only_refs = keys.TableKeys("main.ops.tickets", (), (TestResolve.REF,)) + script = keys.rebuild_script([(only_refs, [])]) + assert "BEGIN" not in script and script.rstrip().endswith("AS fluid_keys;") + + +# ── the jobs' fluid_keys task ────────────────────────────────────────────── +def _dbt(contract: Dict[str, Any], repository: str = "./dbt") -> Dict[str, Any]: + contract["builds"] = [ + { + "id": "transform", + "engine": "dbt", + "repository": repository, + "outputs": [contract["exposes"][0]["exposeId"]], + "properties": {"model": "m"}, + "execution": {"runtime": {"platform": "databricks"}}, + } + ] + return contract + + +def _project(tmp_path: Path) -> None: + root = tmp_path / "dbt" + (root / "models").mkdir(parents=True) + (root / "dbt_project.yml").write_text("name: p\n", encoding="utf-8") + (root / "models" / "m.sql").write_text("select 1\n", encoding="utf-8") + + +def _emitted(tmp_path: Path, contract: Dict[str, Any], monkeypatch, found=None) -> Dict[str, Any]: + """``emit``'s resources, its live lookup answered by ``found`` (product → candidates).""" + asked: List[List[str]] = [] + + def lookup(products): + asked.append(list(products)) + return found or {} + + monkeypatch.setattr(dbx, "_live_key_lookup", lookup) + monkeypatch.setenv(dbx.WAREHOUSE_ENV, WAREHOUSE) + with contract_source(tmp_path / "contract.fluid.yaml"): + out = PLUGIN.emit(contract) + out["_asked"] = asked + return out + + +def _script(resources: Dict[str, Any]) -> str: + body = next( + v + for v in resources["databricks_workspace_file"].values() + if v["path"].endswith("/fluid_keys.sql") + ) + return base64.b64decode(body["content_base64"]).decode("utf-8") + + +class TestJobTask: + def test_a_job_that_writes_a_keyed_table_declares_them_again_after_dbt( + self, tmp_path, monkeypatch + ): + _project(tmp_path) + cands = {SITES: [keys.Candidate("main.geo.sites", ("site_id",))]} + out = _emitted(tmp_path, _dbt(tickets()), monkeypatch, cands) + assert out["_asked"] == [[SITES]] + (job,) = out["databricks_job"].values() + dbt_task, keys_task = job["task"] + assert dbt_task["task_key"] == "dbt" + assert keys_task["task_key"] == "fluid_keys" + assert keys_task["depends_on"] == [{"task_key": "dbt"}] + assert keys_task["sql_task"]["warehouse_id"] == WAREHOUSE + assert keys_task["sql_task"]["file"]["source"] == "WORKSPACE" + assert str(keys_task["sql_task"]["file"]["path"]).startswith("/Workspace${") + assert str(keys_task["sql_task"]["file"]["path"]).endswith("/files/fluid_keys.sql") + script = _script(out) + assert "ADD CONSTRAINT `fluid_pk_tickets` PRIMARY KEY (`ticket_id`)" in script + assert "REFERENCES `main`.`geo`.`sites` (`site_id`)" in script + # The file exists before the job that runs it. + address = next( + f"databricks_workspace_file.{k}" + for k, v in out["databricks_workspace_file"].items() + if v["path"].endswith("/fluid_keys.sql") + ) + assert address in job["depends_on"] + + def test_an_unresolved_foreign_key_leaves_only_the_primary_key(self, tmp_path, monkeypatch): + _project(tmp_path) + script = _script(_emitted(tmp_path, _dbt(tickets()), monkeypatch)) + assert "PRIMARY KEY" in script and "FOREIGN KEY" not in script + + def test_a_job_whose_tables_have_no_key_is_unchanged(self, tmp_path, monkeypatch): + _project(tmp_path) + contract = _dbt(product("x.y_v1", expose("t", [col("a")], "s"))) + out = _emitted(tmp_path, contract, monkeypatch) + (job,) = out["databricks_job"].values() + assert [t["task_key"] for t in job["task"]] == ["dbt"] + assert not any( + "fluid_keys" in json.dumps(v) for v in out["databricks_workspace_file"].values() + ) + assert out["_asked"] == [] + + def test_a_git_job_gets_the_file_in_its_own_folder(self, tmp_path, monkeypatch): + contract = _dbt(tickets(), repository="https://github.com/o/r?tag=v1.2.0") + out = _emitted(tmp_path, contract, monkeypatch) + (job,) = out["databricks_job"].values() + assert job["task"][0]["dbt_task"]["source"] == "GIT" + assert job["task"][1]["sql_task"]["file"]["source"] == "WORKSPACE" + assert len(out["databricks_directory"]) == 1 + assert "PRIMARY KEY" in _script(out) + # The file's path is under the deploying identity's home, read by a data source. + assert "me" in compiled(contract).data["databricks_current_user"] + + def test_the_plan_lists_the_file(self, tmp_path, monkeypatch): + _project(tmp_path) + monkeypatch.setenv(dbx.WAREHOUSE_ENV, WAREHOUSE) + with contract_source(tmp_path / "contract.fluid.yaml"): + actions = dbx.planned_objects(_dbt(tickets())) + files = [a["resource_id"] for a in actions if a["op"] == "databricks.workspace_file.ensure"] + assert any(f.endswith("/fluid_keys.sql") for f in files) + + +# ── the hooks, against the stand-in ──────────────────────────────────────── +@pytest.fixture +def fake(monkeypatch) -> Iterator[FakeDatabricks]: + server = FakeDatabricks().start() + monkeypatch.setenv(dbx.WAREHOUSE_ENV, server.warehouse_id) + try: + yield server + finally: + server.stop() + + +class Hooks: + def __init__(self, fake: FakeDatabricks): + self.fake = fake + self.said: List[str] = [] + + def post(self, doc: Dict[str, Any]) -> None: + PLUGIN.post_apply(doc, say=self.said.append, client_factory=self.fake.client) + + def dry(self, doc: Dict[str, Any]) -> None: + PLUGIN.pre_apply(doc, say=self.said.append, dry_run=True, client_factory=self.fake.client) + + @property + def text(self) -> str: + return "\n".join(self.said) + + def sql(self, statement: str) -> None: + SqlRunner(self.fake.client(), self.fake.warehouse_id).run(statement) + + +def deploy_sites(fake: FakeDatabricks) -> None: + """What ``tofu apply`` of :func:`sites` leaves: the table, tagged with its product id.""" + fake.create_table("main.geo.sites", [("site_id", "string", False), ("name", "string", True)]) + fake.tag("tables", "main.geo.sites", keys.PRODUCT_TAG, SITES) + + +def deploy_tickets(fake: FakeDatabricks) -> None: + fake.create_table( + "main.ops.tickets", + [("ticket_id", "string", False), ("site_id", "string", True), ("note", "string", True)], + ) + fake.tag("tables", "main.ops.tickets", keys.PRODUCT_TAG, TICKETS) + + +PK_SITES = ("PRIMARY KEY", "fluid_pk_sites", "site_id") +PK_TICKETS = ("PRIMARY KEY", "fluid_pk_tickets", "ticket_id") +FK_TICKETS = ("FOREIGN KEY", "fluid_fk_tickets_site_id", "site_id", "main.geo.sites", "site_id") + + +class TestHooks: + def test_the_parent_then_the_child_across_two_contracts(self, fake): + run = Hooks(fake) + deploy_sites(fake) + run.post(sites()) + assert fake.keys_on("main.geo.sites") == [PK_SITES] + deploy_tickets(fake) + run.post(tickets()) + assert fake.keys_on("main.ops.tickets") == [PK_TICKETS, FK_TICKETS] + assert ( + "done: foreign key fluid_fk_tickets_site_id (site_id) references main.geo.sites (site_id)" + in run.text + ) + # One lookup, by tag, for the other product: the parent's location is never assumed. + lookups = [s for s in fake.sql_log if "information_schema.table_tags" in s] + assert len(lookups) == 1 and f"'{SITES}'" in lookups[0] + + def test_a_second_apply_runs_no_key_statement(self, fake): + run = Hooks(fake) + deploy_sites(fake) + deploy_tickets(fake) + run.post(sites()) + run.post(tickets()) + before = len(fake.sql_log) + run.post(sites()) + run.post(tickets()) + assert not [s for s in fake.sql_log[before:] if "CONSTRAINT" in s] + + def test_a_parent_not_deployed_yet_is_said_and_declared_by_the_next_apply(self, fake): + run = Hooks(fake) + deploy_tickets(fake) + run.post(tickets()) # never raises + assert fake.keys_on("main.ops.tickets") == [PK_TICKETS] + assert f"not declared: foreign key (site_id) to {SITES}: no table tagged" in run.text + deploy_sites(fake) + run.post(sites()) + run.post(tickets()) + assert fake.keys_on("main.ops.tickets") == [PK_TICKETS, FK_TICKETS] + + def test_a_rebuild_drops_them_and_the_next_apply_declares_them_again(self, fake): + run = Hooks(fake) + deploy_sites(fake) + deploy_tickets(fake) + run.post(sites()) + run.post(tickets()) + # dbt rebuilds the parent: its key, and the child's foreign key to it, are gone. + run.sql("CREATE OR REPLACE TABLE main.geo.sites (site_id string NOT NULL, name string)") + assert fake.keys_on("main.geo.sites") == [] + assert fake.keys_on("main.ops.tickets") == [PK_TICKETS] + run.post(sites()) + run.post(tickets()) + assert fake.keys_on("main.geo.sites") == [PK_SITES] + assert fake.keys_on("main.ops.tickets") == [PK_TICKETS, FK_TICKETS] + + def test_a_tag_removed_drops_our_key_and_leaves_an_admins(self, fake): + run = Hooks(fake) + deploy_sites(fake) + deploy_tickets(fake) + run.post(sites()) + run.post(tickets()) + run.sql( + "ALTER TABLE main.ops.tickets ADD CONSTRAINT admin_fk FOREIGN KEY (note) REFERENCES main.geo.sites (site_id)" + ) + untagged = tickets() + untagged["exposes"][0]["contract"]["schema"][1].pop("tags") + run.post(untagged) + assert fake.keys_on("main.ops.tickets") == [ + PK_TICKETS, + ("FOREIGN KEY", "admin_fk", "note", "main.geo.sites", "site_id"), + ] + assert "done: drop foreign key fluid_fk_tickets_site_id" in run.text + + def test_a_dry_run_prints_the_statements_and_runs_none(self, fake): + run = Hooks(fake) + deploy_sites(fake) + run.post(sites()) + before = [s for s in fake.sql_log if "CONSTRAINT" in s] + run.dry(tickets()) # the child is not created yet + assert [s for s in fake.sql_log if "CONSTRAINT" in s] == before + assert "main.ops.tickets (not created yet)" in run.text + assert "REFERENCES `main`.`geo`.`sites` (`site_id`)" in run.text + assert "dry-run: none of these statements ran." in run.text + + def test_a_nullable_live_key_column_is_said_with_the_fix_and_fails_nothing(self, fake): + run = Hooks(fake) + # A dbt model that does not declare not_null on its key column. + fake.create_table("main.geo.sites", [("site_id", "string", True), ("name", "string", True)]) + run.post(sites()) + assert fake.keys_on("main.geo.sites") == [] + assert "failed: primary key fluid_pk_sites (site_id):" in run.text + assert "declare not_null on it there" in run.text + + def test_a_contract_with_no_key_says_nothing(self, fake): + run = Hooks(fake) + fake.create_table("main.s.t", [("a", "string", True)]) + run.post(product("x.y_v1", expose("t", [col("a")], "s"))) + assert "Unity Catalog keys" not in run.text + + def test_no_client_is_one_line_for_a_contract_with_keys_and_silence_without(self, monkeypatch): + monkeypatch.setenv(dbx.WAREHOUSE_ENV, WAREHOUSE) + said: List[str] = [] + + def broken() -> Any: + raise RuntimeError("no credential") + + dbx._declare_keys(tickets(), said.append, broken) + assert ( + len(said) == 1 and "Unity Catalog keys: none declared, no Databricks client" in said[0] + ) + said.clear() + dbx._declare_keys(product("x.y_v1", expose("t", [col("a")], "s")), said.append, broken) + assert said == [] + + def test_parent_and_child_in_one_contract_on_the_first_apply(self, fake): + run = Hooks(fake) + both = product( + "geo.all_v1", + expose("sites", [col("site_id", pk=True)], "geo"), + expose( + "visits", + [col("visit_id", pk=True), col("site_id", ref="geo.all_v1.site_id")], + "geo", + ), + ) + fake.create_table("main.geo.sites", [("site_id", "string", False)]) + fake.create_table( + "main.geo.visits", [("visit_id", "string", False), ("site_id", "string", True)] + ) + run.post(both) + assert fake.keys_on("main.geo.visits") == [ + ("PRIMARY KEY", "fluid_pk_visits", "visit_id"), + ("FOREIGN KEY", "fluid_fk_visits_site_id", "site_id", "main.geo.sites", "site_id"), + ] + assert not [s for s in fake.sql_log if "information_schema" in s] + + def test_a_failed_apply_declares_no_key(self, fake): + run = Hooks(fake) + deploy_sites(fake) + PLUGIN.post_apply( + sites(), say=run.said.append, client_factory=fake.client, after_failed_apply=True + ) + assert fake.keys_on("main.geo.sites") == [] + + +# ── after builds run here ────────────────────────────────────────────────── +class TestAfterBuilds: + def _calls(self, monkeypatch, contract: Dict[str, Any]) -> List[List[str]]: + seen: List[List[str]] = [] + monkeypatch.setenv(dbx.WAREHOUSE_ENV, WAREHOUSE) + monkeypatch.setattr( + dbx, "reconcile", lambda targets, *a, **k: seen.append([t.full_name for t in targets]) + ) + PLUGIN.post_build(contract, say=lambda line: None) + return seen + + def test_a_dbt_build_run_here_has_its_keys_declared_again(self, monkeypatch): + contract = _dbt(tickets()) + contract["builds"][0]["execution"] = {} + assert self._calls(monkeypatch, contract) == [["main.ops.tickets"]] + + def test_a_job_or_a_landing_needs_nothing(self, monkeypatch): + assert self._calls(monkeypatch, _dbt(tickets())) == [] # its job's fluid_keys task does + landing = tickets(builds=[{"id": "ingest", "engine": "duckdb", "outputs": ["tickets"]}]) + assert self._calls(monkeypatch, landing) == [] # INSERT OVERWRITE keeps them + + def test_the_engine_runs_the_hook_and_nothing_it_raises_escapes(self, monkeypatch, capsys): + from fluid_build.cli import _apply_opentofu_engine as engine + + class Plugin: + calls: List[Dict[str, Any]] = [] + + def post_build(self, contract, say): + self.calls.append(contract) + raise RuntimeError("boom") + + monkeypatch.setattr(engine, "_resolve_provider", lambda contract, requested: "databricks") + monkeypatch.setattr(engine, "get_iac_plugin", lambda provider: Plugin()) + args = type("Args", (), {"contract": None, "provider": None})() + engine.after_builds(args, {"id": "x"}) + assert Plugin.calls == [{"id": "x"}] + assert "after the build: RuntimeError: boom" in capsys.readouterr().out From 45d46064595261bbdff265661ff9a7c72de9df9c Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:28:09 +0100 Subject: [PATCH 3/8] fix(databricks): a key's parent the deploying identity owns wins a tie Any table owner can tag a table fluid_product_id, so a lookalike table in the child's own catalog outranked the parent this pipeline deployed in another catalog. A candidate the deploying identity owns (read-back owner vs SCIM Me, as the mask functions are checked) now wins a tie before the catalog does. Nothing is refused: a parent another identity deployed is still used when it is the only match, so a cross-team foreign key keeps working. --- fluid_build/iac/providers/databricks.py | 4 +- fluid_build/iac/providers/databricks_keys.py | 42 +++++++++++++++----- tests/iac/test_iac_databricks_keys.py | 31 +++++++++++++++ 3 files changed, 67 insertions(+), 10 deletions(-) diff --git a/fluid_build/iac/providers/databricks.py b/fluid_build/iac/providers/databricks.py index 19fc5ddf..dcc4d69a 100644 --- a/fluid_build/iac/providers/databricks.py +++ b/fluid_build/iac/providers/databricks.py @@ -198,6 +198,7 @@ KeyTarget, TableKeys, compile_keys, + deploying_identity, foreign_keys, own_candidates, rebuild_script, @@ -1501,7 +1502,8 @@ def _live_key_lookup(products: Sequence[str]) -> Mapping[str, Sequence[Candidate if not warehouse: return {} client = default_client() - return lookup_key_tables(SqlRunner(client, warehouse), client, products) + runner = SqlRunner(client, warehouse) + return lookup_key_tables(runner, client, products, deploying_identity(client)) except Exception: # noqa: BLE001 — no live answer is no foreign key in the script return {} diff --git a/fluid_build/iac/providers/databricks_keys.py b/fluid_build/iac/providers/databricks_keys.py index c15692b6..cb44fe9f 100644 --- a/fluid_build/iac/providers/databricks_keys.py +++ b/fluid_build/iac/providers/databricks_keys.py @@ -78,6 +78,8 @@ SqlRunner, StatementFailed, default_client, + deploy_identity, + owner_matches, read_table, scrub, ) @@ -146,10 +148,16 @@ class KeyTarget: @dataclass(frozen=True) class Candidate: - """A table a referenced product has, and its primary key.""" + """A table a referenced product has, and its primary key. + + ``trusted``: it is one of the contract's own exposes, or the identity this run deploys as + owns it. Anyone who owns a table can tag it ``fluid_product_id``, so a table this pipeline + deployed wins a tie against one it did not. + """ full_name: str primary: Tuple[str, ...] + trusted: bool = True @property def catalog(self) -> str: @@ -257,10 +265,11 @@ def own_candidates(contract_id: str, keyed: Iterable[TableKeys]) -> Dict[str, Li def _pick(matches: List[Candidate], catalog: str) -> Tuple[Optional[Candidate], List[Candidate]]: - """The one match (one in ``catalog`` wins a tie), or ``None`` and every match left.""" - if len(matches) > 1: - local = [m for m in matches if m.catalog == catalog] - matches = local or matches + """The one match, or ``None`` and every match left. A tie goes to the trusted ones, then + to the ones in ``catalog``.""" + for prefer in (lambda m: m.trusted, lambda m: m.catalog == catalog): + if len(matches) > 1: + matches = [m for m in matches if prefer(m)] or matches return (matches[0], []) if len(matches) == 1 else (None, matches) @@ -370,8 +379,11 @@ def names(raw: Any) -> Tuple[str, ...]: return primary, foreign -def lookup(runner: SqlRunner, client: Any, products: Iterable[str]) -> Dict[str, List[Candidate]]: - """The tables tagged with each product id, with their live primary keys. +def lookup( + runner: SqlRunner, client: Any, products: Iterable[str], owner: str = "" +) -> Dict[str, List[Candidate]]: + """The tables tagged with each product id, with their live primary keys; one ``owner`` + owns (the deploying identity) is ``trusted``. One query on ``system.information_schema.table_tags``, which lists what this identity can see, then one Tables API read per table found. A table whose name this module could not @@ -402,7 +414,10 @@ def lookup(runner: SqlRunner, client: Any, products: Iterable[str]) -> Dict[str, if live is None: continue primary, _ = live_keys(live) - out.setdefault(product, []).append(Candidate(full_name, primary[1] if primary else ())) + trusted = bool(owner) and owner_matches(live.get("owner"), {owner}) + out.setdefault(product, []).append( + Candidate(full_name, primary[1] if primary else (), trusted) + ) return out @@ -576,7 +591,7 @@ def resolve( if not others: return candidates try: - candidates.update(lookup(runner, client, others)) + candidates.update(lookup(runner, client, others, deploying_identity(client))) except Exception as exc: # noqa: BLE001 — a lookup that fails declares no such key notes.append( f"foreign keys to {', '.join(others)}: their tables could not be looked up " @@ -585,6 +600,15 @@ def resolve( return candidates +def deploying_identity(client: Any) -> str: + """Who this run deploys as, or ``""`` when that cannot be read (then no table is trusted + over another).""" + try: + return deploy_identity(client) + except Exception: # noqa: BLE001 — a tie is then broken by catalog alone + return "" + + @dataclass(frozen=True) class _Read: target: KeyTarget diff --git a/tests/iac/test_iac_databricks_keys.py b/tests/iac/test_iac_databricks_keys.py index 2083ecab..f472f07a 100644 --- a/tests/iac/test_iac_databricks_keys.py +++ b/tests/iac/test_iac_databricks_keys.py @@ -193,6 +193,18 @@ def test_a_tie_goes_to_the_table_in_the_childs_catalog(self): (fk,) = keys.foreign_keys(_keys(self.REF), cands) assert fk.parent == "main.geo.sites" + def test_a_table_this_pipeline_owns_beats_a_tagged_one_it_does_not(self): + # Anyone who owns a table can tag it fluid_product_id: a decoy in the child's own + # catalog must not outrank the parent the deploying identity created elsewhere. + cands = { + SITES: [ + keys.Candidate("main.scratch.sites", ("site_id",), trusted=False), + keys.Candidate("prod.geo.sites", ("site_id",), trusted=True), + ] + } + (fk,) = keys.foreign_keys(_keys(self.REF), cands) + assert fk.parent == "prod.geo.sites" + def test_a_tie_no_catalog_breaks_is_skipped_and_said(self): cands = { SITES: [ @@ -497,6 +509,25 @@ def test_the_parent_then_the_child_across_two_contracts(self, fake): lookups = [s for s in fake.sql_log if "information_schema.table_tags" in s] assert len(lookups) == 1 and f"'{SITES}'" in lookups[0] + def test_a_decoy_someone_else_owns_does_not_become_the_parent(self, fake): + run = Hooks(fake) + fake.create_table("prod.geo.sites", [("site_id", "string", False)]) + fake.tag("tables", "prod.geo.sites", keys.PRODUCT_TAG, SITES) + run.sql("ALTER TABLE prod.geo.sites ADD CONSTRAINT fluid_pk_sites PRIMARY KEY (site_id)") + fake.create_table("main.scratch.sites", [("site_id", "string", False)]) + fake.tables["main.scratch.sites"]["owner"] = "mallory@example.test" + fake.tag("tables", "main.scratch.sites", keys.PRODUCT_TAG, SITES) + run.sql("ALTER TABLE main.scratch.sites ADD CONSTRAINT decoy_pk PRIMARY KEY (site_id)") + deploy_tickets(fake) + run.post(tickets()) + assert ( + "FOREIGN KEY", + "fluid_fk_tickets_site_id", + "site_id", + "prod.geo.sites", + "site_id", + ) in fake.keys_on("main.ops.tickets") + def test_a_second_apply_runs_no_key_statement(self, fake): run = Hooks(fake) deploy_sites(fake) From 1d57b8545fc8a694481e578a956b61ca93c9abb5 Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:48:11 +0100 Subject: [PATCH 4/8] fix(databricks): keys review round: dry runs, templates, lookups, names, coverage From the diff's security, delivery and defect reviews: - A dry run of --mode amend-and-build no longer declares keys: the post-build hook runs only when the apply is real, and runs after a build that failed part way too (it may already have rebuilt a table). - The post-build hook gets the contract as the apply did, {{ env.* }} templates resolved; before, a templated binding made it skip the keys without a word. - fluid validate (and the verify inventory) of a templated location no longer raises: a key statement for a name that is no plain identifier yet is left out of the job script. - emit looks up other products' tables only after the refusal and sovereignty checks, waits at most 120 s, and a lookup that fails is a plan notice instead of silence; one unreadable or oddly named tagged table no longer discards the others. - A composite key and a single-column key to two tables of one product are both declared; a column that is half of a composite key is said as such. - A foreign key this module declared is kept only while it still says what the contract does (or its product cannot be looked up); one a new key covers, or whose reference moved, is dropped. - A job's script also re-declares the contract's own foreign keys into the tables it rebuilds. - Key names can no longer collide in a schema: fluid_fk_
__[__], with a digest when a part could read two ways. - Own exposes' declared primary keys count even when a parent's table could not be read. - The hooks' output (server error classes, the install hint) prints literally. - Docs: the job task passes over a key it cannot add (a nullable key column) and needs a warehouse that runs SQL scripting. --- CHANGELOG.md | 7 +- fluid_build/cli/_apply_opentofu_engine.py | 31 ++- fluid_build/iac/providers/databricks.py | 76 +++++- fluid_build/iac/providers/databricks_keys.py | 247 ++++++++++++----- tests/iac/test_iac_databricks_fake_e2e.py | 2 +- tests/iac/test_iac_databricks_keys.py | 263 +++++++++++++++++-- 6 files changed, 520 insertions(+), 106 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7fb4ba20..aba2a06f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,13 +18,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 forge table carries (`system.information_schema.table_tags`), wherever an overlay put it. `fluid apply` declares the keys on the SQL warehouse after `tofu apply` (the provider has no constraint argument), re-declares any that are missing, replaces or drops the ones it - declared (named `fluid_pk_
` / `fluid_fk_
_`) when the contract changes, + declared (named `fluid_pk_
` / `fluid_fk_
__`) when the contract changes, and leaves any other key in place. A key that cannot be declared yet (its parent not deployed, say) is skipped with the reason and the fix, never failing the run; `--mode dry-run` prints the statements. Because a dbt `table` build's `CREATE OR REPLACE` drops a table's keys and every foreign key that references it (measured on a warehouse), each Databricks dbt job that writes a keyed table gets a second task, `fluid_keys`, that - declares them again after dbt, in a script that never fails the run; a dbt build that + declares them again after dbt, with the contract's own foreign keys into those tables, in a + script that passes over a key it cannot add instead of failing the run. Such a key (a key + column the dbt model leaves nullable, say) is reported by the next `fluid apply`, and the + script needs a SQL warehouse that runs SQL scripting. A dbt build that `--mode amend-and-build` runs locally is followed by the same re-declaration. ### Fixed diff --git a/fluid_build/cli/_apply_opentofu_engine.py b/fluid_build/cli/_apply_opentofu_engine.py index 8d7794df..22f94aee 100644 --- a/fluid_build/cli/_apply_opentofu_engine.py +++ b/fluid_build/cli/_apply_opentofu_engine.py @@ -345,16 +345,17 @@ def _enforce_access( binds Unity Catalog column masks and row filters, which the provider cannot manage (``iac/providers/databricks_enforce.py``). Unlike the advisory hooks this one gates: a refusal or failure is a typed error and the run exits non-zero. With ``dry_run`` the - plugin prints the statements and runs none. + plugin prints the statements and runs none. What a hook prints is data (table names, + statements, a server's error text with its ``[ERROR_CLASS]``), so it goes out literally. """ run = getattr(plugin, hook, None) if not callable(run): return try: if dry_run: - run(contract, say=cprint, dry_run=True) + run(contract, say=cprint_literal, dry_run=True) else: - run(contract, say=cprint) + run(contract, say=cprint_literal) except UnsupportedBindingError as exc: raise CLIError( 1, @@ -367,26 +368,30 @@ def build_after_apply( args, logger: logging.Logger, contract: Mapping[str, Any], build_kwargs: Mapping[str, Any] ) -> int: """The build phase of a build-augmented mode, after a ``tofu apply`` that succeeded, then - the plugin's ``post_build`` hook once every build succeeded (:func:`after_builds`).""" + the plugin's ``post_build`` hook (:func:`after_builds`), whatever the builds returned: one + that failed part way may already have rebuilt a table. A dry run changes nothing, so it + runs no hook (its plan already printed what the apply would declare).""" from fluid_build.build_runners import run_builds_from_args rc = run_builds_from_args(args, logger, force_run=True, **build_kwargs) - if rc == 0: - after_builds(args, contract) + if not getattr(args, "dry_run", False): + after_builds(args, logger, contract) return rc -def after_builds(args, contract: Mapping[str, Any]) -> None: - """The plugin's ``post_build`` hook, if it has one, after a build phase that succeeded. +def after_builds(args, logger: logging.Logger, contract: Mapping[str, Any]) -> None: + """The plugin's ``post_build`` hook, if it has one, after the build phase. Optional plugin capability, discovered like ``plan_notices``, and advisory: a dbt build run here may have rebuilt a table with CREATE OR REPLACE, which drops what the apply declared on it that dbt does not (the Databricks plugin's Unity Catalog keys), and the plugin declares - it again. It never changes the run's exit code. + it again. It gets the contract as the apply did, its ``{{ env.* }}`` templates resolved + (:func:`_load_contract`), and never changes the run's exit code. """ try: + resolved = _load_contract(args, logger, contract) plugin = get_iac_plugin( - _resolve_provider(contract, getattr(args, "provider", None) or "auto") + _resolve_provider(resolved, getattr(args, "provider", None) or "auto") ) except Exception: # noqa: BLE001 — the apply resolved it already; nothing to add here return @@ -395,9 +400,9 @@ def after_builds(args, contract: Mapping[str, Any]) -> None: return try: with contract_source(source_contract_file(getattr(args, "contract", None))): - run(contract, say=cprint) + run(resolved, say=cprint_literal) except Exception as exc: # noqa: BLE001 — advisory: the build's own result stands - cprint(f" after the build: {type(exc).__name__}: {redact_secret_text(str(exc))}") + cprint_literal(f" after the build: {type(exc).__name__}: {exc}") def _enforce_after_failed_apply(plugin: Any, contract: Mapping[str, Any]) -> None: @@ -415,7 +420,7 @@ def _enforce_after_failed_apply(plugin: Any, contract: Mapping[str, Any]) -> Non "exist before the run stops (no mask or filter is dropped)." ) try: - run(contract, say=cprint, after_failed_apply=True) + run(contract, say=cprint_literal, after_failed_apply=True) except UnsupportedBindingError as exc: cprint_literal( f" access enforcement after the failed apply: {exc} {' '.join(exc.remediation)}" diff --git a/fluid_build/iac/providers/databricks.py b/fluid_build/iac/providers/databricks.py index dcc4d69a..d1c17fea 100644 --- a/fluid_build/iac/providers/databricks.py +++ b/fluid_build/iac/providers/databricks.py @@ -1431,6 +1431,8 @@ def _rebuild_scripts( """The ``fluid_keys`` script of each Databricks dbt job, by its build; ``None`` for a job that writes no table with a key (``databricks_keys.rebuild_script``). + A job's script declares again what its build's ``CREATE OR REPLACE`` drops: the keys of the + tables it writes, and the foreign keys of this contract's other tables that reference them. Its foreign keys to other products need ``key_lookup``, which only ``emit`` passes; without it (``fluid validate``, the native planner) a script holds the primary keys and the foreign keys to this contract's own tables. @@ -1461,12 +1463,16 @@ def _rebuild_scripts( ) if others and key_lookup is not None: candidates.update(key_lookup(others)) - scripts = { - build_id: rebuild_script( - [(keys, foreign_keys(keys, candidates, own_id=contract_id)) for keys in tables] - ) - for build_id, tables in written.items() - } + resolved = [ + (keys.full_name, fk) + for _t, keys in keyed + for fk in foreign_keys(keys, candidates, own_id=contract_id) + ] + scripts: Dict[int, str] = {} + for build_id, tables in written.items(): + names = {keys.full_name for keys in tables} + drops = [(child, fk) for child, fk in resolved if child in names or fk.parent in names] + scripts[build_id] = rebuild_script(tables, drops) return lambda build: scripts.get(id(build)) @@ -1490,22 +1496,43 @@ def _declare_keys( say(f"\n Unity Catalog keys: not declared ({type(exc).__name__}: {scrub(exc)}).") -def _live_key_lookup(products: Sequence[str]) -> Mapping[str, Sequence[Candidate]]: +#: How long ``emit`` waits for the live key lookup: a plan should not hang on a warehouse. +_KEY_LOOKUP_TIMEOUT_SECONDS = 120 + + +def _live_key_lookup( + products: Sequence[str], failures: Optional[List[str]] = None +) -> Mapping[str, Sequence[Candidate]]: """Other products' live tables, for the ``fluid_keys`` scripts ``emit`` writes. - Best effort, empty when the workspace cannot be read: a module still emits without a - credential (``fluid generate iac``), and ``fluid apply`` resolves, declares and reports the - keys itself after ``tofu apply`` (``databricks_keys.reconcile``). + Best effort, and what it could not read is added to ``failures``: a module still emits + without a credential (``fluid generate iac``), and ``fluid apply`` resolves, declares and + reports the keys itself after ``tofu apply`` (``databricks_keys.reconcile``). """ + said: List[str] = [] try: warehouse = str(os.environ.get(WAREHOUSE_ENV) or "").strip() if not warehouse: + said.append(f"{WAREHOUSE_ENV} is not set") return {} client = default_client() - runner = SqlRunner(client, warehouse) - return lookup_key_tables(runner, client, products, deploying_identity(client)) - except Exception: # noqa: BLE001 — no live answer is no foreign key in the script + runner = SqlRunner(client, warehouse, timeout=_KEY_LOOKUP_TIMEOUT_SECONDS) + return lookup_key_tables(runner, client, products, deploying_identity(client), said) + except Exception as exc: # noqa: BLE001 — no live answer is no foreign key in the script + said.append(f"{type(exc).__name__}: {scrub(exc)}") return {} + finally: + if failures is not None: + failures.extend(said) + + +def _references_another_product(compiled: _Compiled, contract: Mapping[str, Any]) -> bool: + """Whether a dbt job's script may need the live lookup: the contract has a job, and a + table that references a product other than its own.""" + own = _text(contract.get("id")) + return bool(compiled.jobs) and any( + r.product != own for t in compiled.tables if t.keys for r in t.keys.references + ) def _resolve_tables( @@ -2184,10 +2211,21 @@ class DatabricksIacPlugin: def emit( self, contract: Mapping[str, Any], actions: Iterable[Mapping[str, Any]] = () ) -> Dict[str, Any]: - compiled = _compile(contract, os.environ, key_lookup=_live_key_lookup) + compiled = _compile(contract, os.environ) if compiled.refusals: raise compiled.refusals[0] _enforce_sovereignty(contract, compiled.tables) + # Only now, with the contract compiled and its sovereignty upheld, may the dbt jobs' + # fluid_keys scripts look up other products' tables on the warehouse; ``plan_notices`` + # says when it could not. + failures: List[str] = [] + if _references_another_product(compiled, contract): + compiled = _compile( + contract, + os.environ, + key_lookup=lambda products: _live_key_lookup(products, failures), + ) + self._key_lookup_failures = (_contract_key(contract), failures) # Decided once, here: ``plan_notices`` of the same run says what this paused, so a # transient read error cannot pause a job without the notice (or the reverse). to_pause = _jobs_to_pause(compiled, contract) @@ -2389,6 +2427,16 @@ def plan_notices(self, contract: Mapping[str, Any]) -> List[str]: "so this identity needs the ASSIGN permission on it, which an account or " "workspace admin grants, or the tag assignment fails." ) + failed = getattr(self, "_key_lookup_failures", None) + self._key_lookup_failures = None # used once, as the pause decision is + if failed is not None and failed[1] and failed[0] == _contract_key(contract): + notices.append( + "the fluid_keys task of this contract's dbt jobs declares no foreign key to " + "another product's table it could not look up now (" + + "; ".join(failed[1]) + + "). fluid apply still declares them on the tables after the apply, and the " + "next apply that can look them up writes them into the task's script." + ) decided = getattr(self, "_pause_decision", None) self._pause_decision = None # used once: the next run decides again if decided is not None and decided[0] is not None and decided[0] == _contract_key(contract): diff --git a/fluid_build/iac/providers/databricks_keys.py b/fluid_build/iac/providers/databricks_keys.py index cb44fe9f..51e8bcae 100644 --- a/fluid_build/iac/providers/databricks_keys.py +++ b/fluid_build/iac/providers/databricks_keys.py @@ -30,10 +30,12 @@ referenced columns as its live PRIMARY KEY; of several, one in the child's catalog wins, and more than one left is ambiguous, so that key is skipped and said. -**Names.** ``fluid_pk_
`` and ``fluid_fk_
_`` (a key's name is unique in -its schema; one longer than 255 characters is cut and given a digest). The ``fluid_`` prefix -marks what this module owns: it replaces or drops only a key so named, and leaves any other in -place, saying so. A key that already holds what the contract asks, whatever its name, is kept. +**Names.** ``fluid_pk_
`` and ``fluid_fk_
__[__]``: a key's name +must be unique in its schema, and a table's name is, so are a foreign key's columns on it +(:func:`constraint_name`; a part that could read two ways adds a digest, and a name past 255 +characters is cut and given one). The ``fluid_`` prefix marks what this module owns: it +replaces or drops only a key so named, and leaves any other in place, saying so. A key that +already holds what the contract asks, whatever its name, is kept. **A key never fails the run**: it is metadata. One that cannot be declared yet (its parent not deployed, a parent without its primary key, a nullable column) is skipped with the reason @@ -54,10 +56,16 @@ So a dbt ``table`` build, a ``CREATE OR REPLACE``, would leave its table with no keys after every run. Each Databricks dbt job that writes a keyed table therefore gets a second task, ``fluid_keys`` (``databricks_jobs.py``), which runs :func:`rebuild_script` on the same -warehouse once dbt has succeeded: the keys the apply resolved, each added again, and each that -fails (a key still in place, a parent whose own task has not run yet) passed over, so no run -fails on one. A parent's rebuild also drops the foreign keys that reference it; each child's -next build, or ``fluid apply``, declares them again. +warehouse once dbt has succeeded: the keys of the tables it rebuilt and the contract's own +foreign keys into them, each added again, and each that fails passed over, so no run fails +on one. What it cannot add it passes over without a word: a key column the dbt model leaves +nullable (declare ``not_null`` on it in the model, with the model contract enforced) is +reported only by the next ``fluid apply``. The script is SQL scripting (a ``BEGIN ... END`` +block), which a current Databricks SQL warehouse runs; on one that cannot, the task, and the +job's run, fail. A parent's rebuild also drops the foreign keys of other contracts' tables +that reference it: such a child gets them back from its own job's task after its next build, +or from its next ``fluid apply``. A dbt build that ``fluid apply --mode amend-and-build`` runs +on this machine is followed by the same re-declaration (``post_build``). The pattern is dbt-databricks' own (``macros/relations/constraints.sql``: ``ALTER TABLE ... ADD CONSTRAINT`` after every create or replace), for keys the contract declares rather than a model. @@ -65,9 +73,24 @@ from __future__ import annotations +import hashlib import os +from contextlib import suppress from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union +from typing import ( + Any, + Callable, + Collection, + Dict, + FrozenSet, + Iterable, + List, + Mapping, + Optional, + Sequence, + Tuple, + Union, +) from ...providers._sql_safety import reject_control_chars, validate_ident from ..naming import bounded_ident @@ -179,8 +202,14 @@ def describe(self) -> str: def constraint_name(prefix: str, *parts: str) -> str: - """``prefix`` and ``parts`` joined by ``_``; cut, with a digest, past 255 characters.""" - return bounded_ident(prefix + "_".join(parts), _MAX_NAME) + """``prefix`` and ``parts`` joined by ``__``, which no part a schema can mistake holds: + ``fluid_pk_
``, ``fluid_fk_
__[__]``. A part that holds ``__`` + or starts or ends with ``_`` could read two ways, so a digest of the parts then follows. + Cut, with a digest, past 255 characters.""" + name = prefix + "__".join(parts) + if any("__" in p or p.startswith("_") or p.endswith("_") for p in parts): + name += "_" + hashlib.sha256("\0".join(parts).encode("utf-8")).hexdigest()[:8] + return bounded_ident(name, _MAX_NAME) def _tags(raw: Any) -> frozenset: @@ -307,20 +336,8 @@ def foreign_keys( "run fluid apply here again." ) continue - wanted = [r.parent_column for r in refs] - if len(refs) > 1 and len(set(wanted)) == len(wanted): - whole = [c for c in found if len(c.primary) > 1 and set(c.primary) == set(wanted)] - if whole: - pick, tied = _pick(whole, keys.catalog) - if pick is None: - said.append(_ambiguous(children, product, wanted, tied)) - continue - child = {r.parent_column: r.column for r in refs} - columns = tuple(child[c] for c in pick.primary) - name = constraint_name(FK_PREFIX, keys.table, *columns) - out.append(ForeignKey(name, columns, pick.full_name, pick.primary)) - continue - for ref in refs: + left = _composites(keys, product, refs, found, out, said) + for ref in left: single = [c for c in found if c.primary == (ref.parent_column,)] pick, tied = _pick(single, keys.catalog) if pick is not None: @@ -329,16 +346,67 @@ def foreign_keys( elif tied: said.append(_ambiguous(ref.column, product, [ref.parent_column], tied)) else: - tables = ", ".join(sorted(c.full_name for c in found)) - said.append( - f"foreign key ({ref.column}) to {product}.{ref.parent_column}: {product} " - f"has {tables}, and none has PRIMARY KEY ({ref.parent_column}). Tag " - f"{ref.parent_column} primary-key in {product}'s contract and apply it, then " - "run fluid apply here again." - ) + said.append(_unmatched(ref, product, found)) return out +def _composites( + keys: TableKeys, + product: str, + refs: Sequence[Reference], + found: Sequence[Candidate], + out: List[ForeignKey], + said: List[str], +) -> List[Reference]: + """The composite foreign keys among ``refs``, appended to ``out``; the references left. + + For each composite primary key a table of ``product`` has, widest first: the references + that name each of its columns exactly once are one foreign key, in that key's order. What + is left is matched column by column, so one product can take a composite key to one of its + tables and a single-column key to another. + """ + by_key: Dict[frozenset, List[Candidate]] = {} + for cand in found: + if len(cand.primary) > 1: + by_key.setdefault(frozenset(cand.primary), []).append(cand) + left = list(refs) + for columns in sorted(by_key, key=lambda k: (-len(k), sorted(k))): + named: Dict[str, List[Reference]] = {} + for ref in left: + if ref.parent_column in columns: + named.setdefault(ref.parent_column, []).append(ref) + if set(named) != columns or any(len(v) != 1 for v in named.values()): + continue + used = [v[0] for v in named.values()] + left = [r for r in left if r not in used] + pick, tied = _pick(by_key[columns], keys.catalog) + if pick is None: + children = ", ".join(r.column for r in used) + said.append(_ambiguous(children, product, sorted(columns), tied)) + continue + child = tuple(named[c][0].column for c in pick.primary) + name = constraint_name(FK_PREFIX, keys.table, *child) + out.append(ForeignKey(name, child, pick.full_name, pick.primary)) + return left + + +def _unmatched(ref: Reference, product: str, found: Sequence[Candidate]) -> str: + """Why ``ref`` matches no table of ``product``, and what to do.""" + part = next((c for c in found if ref.parent_column in c.primary and len(c.primary) > 1), None) + if part is not None: + return ( + f"foreign key ({ref.column}) to {product}.{ref.parent_column}: that column is part " + f"of the composite PRIMARY KEY ({', '.join(part.primary)}) of {part.full_name}. " + "Reference each of its columns, once each, from this table's columns." + ) + tables = ", ".join(sorted(c.full_name for c in found)) + return ( + f"foreign key ({ref.column}) to {product}.{ref.parent_column}: {product} has {tables}, " + f"and none has PRIMARY KEY ({ref.parent_column}). Tag {ref.parent_column} primary-key " + f"in {product}'s contract and apply it, then run fluid apply here again." + ) + + def _ambiguous( children: str, product: str, parents: Sequence[str], tied: Sequence[Candidate] ) -> str: @@ -380,15 +448,21 @@ def names(raw: Any) -> Tuple[str, ...]: def lookup( - runner: SqlRunner, client: Any, products: Iterable[str], owner: str = "" + runner: SqlRunner, + client: Any, + products: Iterable[str], + owner: str = "", + notes: Optional[List[str]] = None, ) -> Dict[str, List[Candidate]]: """The tables tagged with each product id, with their live primary keys; one ``owner`` owns (the deploying identity) is ``trusted``. One query on ``system.information_schema.table_tags``, which lists what this identity can - see, then one Tables API read per table found. A table whose name this module could not - quote is no candidate. Raises when the query or a read fails. + see, then one Tables API read per table found. A table that cannot be read, or whose name + is no plain identifier, is no candidate, and ``notes`` says which. Raises when the query + fails. """ + said = notes if notes is not None else [] wanted = sorted({p for p in products if p}) if not wanted: return {} @@ -408,9 +482,19 @@ def lookup( full_name = f"{catalog}.{schema}.{table}".lower() try: quote_table(full_name) + live = read_table(client, full_name) except ValueError: + said.append( + f"{full_name} is tagged {PRODUCT_TAG} = {product!r}, but its name is no plain " + "identifier, so no foreign key references it." + ) + continue + except Exception as exc: # noqa: BLE001 — one unreadable table is not every table + said.append( + f"{full_name}, tagged {PRODUCT_TAG} = {product!r}, could not be read " + f"({type(exc).__name__}: {scrub(exc)}), so no foreign key references it now." + ) continue - live = read_table(client, full_name) if live is None: continue primary, _ = live_keys(live) @@ -478,13 +562,17 @@ def plan_table( keys: Optional[TableKeys], fks: Sequence[ForeignKey], live: Optional[Mapping[str, Any]], + *, + candidates: Optional[Mapping[str, Sequence[Candidate]]] = None, + unknown: Collection[str] = frozenset(), ) -> TablePlan: """The steps from ``live`` (``None``: not created yet) to what the contract declares. A key already in place is kept whatever its name. One this module declared (its prefix) that the contract changed or dropped is dropped; any other is left, and a contract key it - is in the way of is noted. A foreign key of this module whose parent cannot be resolved now - is kept while the contract still references its columns. + is in the way of is noted. A foreign key of this module that no wanted key replaces is kept + while its columns still reference the one product whose tables could not be looked up, or + whose tables still include its parent (:func:`_still_meant`): a parent unresolved for now. """ plan = TablePlan(full_name, exists=live is not None) live_pk, live_fks = live_keys(live or {}) @@ -511,7 +599,6 @@ def plan_table( ) ) wanted = {fk.columns: fk for fk in fks} - referenced = {r.column for r in keys.references} if keys is not None else set() for fk in live_fks: want = wanted.get(fk.columns) if want is not None and (fk.parent, fk.parent_columns) == ( @@ -530,8 +617,8 @@ def plan_table( ) del wanted[fk.columns] continue - if want is None and set(fk.columns) <= referenced: - continue # still referenced, its parent unresolved now: kept + if want is None and _still_meant(fk, keys, fks, candidates, unknown): + continue # its parent unresolved for now: kept why = "its parent changed" if want is not None else "the contract no longer declares it" plan.fk_drops.append(_drop(full_name, fk.name, f"drop {fk.describe()}: {why}")) for fk in wanted.values(): @@ -539,6 +626,29 @@ def plan_table( return plan +def _still_meant( + fk: ForeignKey, + keys: Optional[TableKeys], + wanted: Sequence[ForeignKey], + candidates: Optional[Mapping[str, Sequence[Candidate]]], + unknown: Collection[str], +) -> bool: + """Whether a foreign key this module declared, that the contract's resolved keys do not + include, still says what the contract does: no wanted key covers any of its columns, its + columns all reference one product, and that product's tables could not be looked up (or + were not, ``candidates`` ``None``) or still include its parent.""" + if any(set(w.columns) & set(fk.columns) for w in wanted): + return False + refs = {r.column: r.product for r in keys.references} if keys is not None else {} + products = {refs.get(c) for c in fk.columns} + if len(products) != 1 or None in products: + return False + (product,) = products + if candidates is None or product in unknown: + return True + return any(c.full_name == fk.parent for c in candidates.get(str(product)) or ()) + + def _hint(error: str) -> str: """What to do about a key statement that failed, when the cause is a known one.""" text = error.lower() @@ -583,21 +693,24 @@ def resolve( runner: SqlRunner, client: Any, notes: List[str], -) -> Dict[str, List[Candidate]]: - """Every table the contract's references can resolve to: its own exposes' declared - primary keys, and the live tables of each other product it references.""" +) -> Tuple[Dict[str, List[Candidate]], FrozenSet[str]]: + """Every table the contract's references can resolve to (its own exposes' declared + primary keys, and the live tables of each other product it references), and the products + whose tables could not be looked up at all.""" candidates: Dict[str, List[Candidate]] = dict(own_candidates(contract_id, keyed)) others = sorted({r.product for k in keyed for r in k.references if r.product != contract_id}) if not others: - return candidates + return candidates, frozenset() try: - candidates.update(lookup(runner, client, others, deploying_identity(client))) + owner = deploying_identity(client) + candidates.update(lookup(runner, client, others, owner, notes)) except Exception as exc: # noqa: BLE001 — a lookup that fails declares no such key notes.append( f"foreign keys to {', '.join(others)}: their tables could not be looked up " f"({type(exc).__name__}: {scrub(exc)})." ) - return candidates + return candidates, frozenset(others) + return candidates, frozenset() def deploying_identity(client: Any) -> str: @@ -656,14 +769,20 @@ def reconcile( general.append(f"keys of {target.full_name}: the table does not exist.") continue reads.append(_Read(target, live)) - keyed = [r.target.keys for r in reads if r.target.keys is not None] - candidates = resolve(keyed, contract_id, runner, client, general) if keyed else {} + # Own candidates come from what every expose declares, read or not: a parent this run could + # not read still declares its primary key. + keyed = [t.keys for t in targets if t.keys is not None] + candidates, unknown = ( + resolve(keyed, contract_id, runner, client, general) if keyed else ({}, frozenset()) + ) plans: List[TablePlan] = [] for read in reads: keys = read.target.keys notes: List[str] = [] fks = foreign_keys(keys, candidates, notes, own_id=contract_id) if keys else [] - plan = plan_table(read.target.full_name, keys, fks, read.live) + plan = plan_table( + read.target.full_name, keys, fks, read.live, candidates=candidates, unknown=unknown + ) plan.notes[:0] = notes if plan.steps or plan.notes: plans.append(plan) @@ -725,22 +844,32 @@ def _run(plans: Sequence[TablePlan], runner: SqlRunner, say: Say) -> None: # ── after a dbt build ────────────────────────────────────────────────── -def rebuild_script(entries: Sequence[Tuple[TableKeys, Sequence[ForeignKey]]]) -> str: +def rebuild_script(rebuilt: Sequence[TableKeys], foreign: Sequence[Tuple[str, ForeignKey]]) -> str: """The SQL file a dbt job's ``fluid_keys`` task runs once the build has succeeded. - Every primary key first, then every foreign key, each inside one script whose - ``CONTINUE`` handler passes over a statement that fails, so the task, and the run, never - fail on a key. A job whose tables resolve to no statement still gets a valid file. + The primary keys of the tables the build rebuilds, then ``foreign`` (``(child table, key)``), + each inside one script whose ``CONTINUE`` handler passes over a statement that fails, so + the task, and the run, never fail on a key: a key still in place, a parent whose own key is + not back yet, a key column the dbt model leaves nullable. A statement for a name that is no + plain identifier yet (``fluid validate`` compiles before ``{{ env.* }}`` templates are + resolved) is left out. A job whose tables resolve to no statement still gets a valid file. """ - body = [primary_key_statement(keys) for keys, _ in entries if keys.primary] - body += [foreign_key_statement(keys.full_name, fk) for keys, fks in entries for fk in fks] + body: List[str] = [] + for keys in rebuilt: + if keys.primary: + with suppress(ValueError): + body.append(primary_key_statement(keys)) + for child, fk in foreign: + with suppress(ValueError): + body.append(foreign_key_statement(child, fk)) header = [ "-- Written by fluid apply (data-product-forge); the next apply rewrites it.", "-- A dbt table build replaces its table with CREATE OR REPLACE, which drops the table's", "-- Unity Catalog keys and the foreign keys that reference it. This task, run after a", - "-- build that succeeded, declares this job's keys again. A statement that fails (a key", - "-- still in place, a parent whose own key is not back yet) is passed over: no run fails", - "-- on a key.", + "-- build that succeeded, declares them again. A statement that fails (a key still in", + "-- place, a parent whose own key is not back yet, a key column the model leaves", + "-- nullable) is passed over: no run fails on a key, and fluid apply says why one is", + "-- missing.", ] if not body: return "\n".join([*header, "SELECT 'no key to declare' AS fluid_keys;", ""]) diff --git a/tests/iac/test_iac_databricks_fake_e2e.py b/tests/iac/test_iac_databricks_fake_e2e.py index 4cc59b97..750dfe03 100644 --- a/tests/iac/test_iac_databricks_fake_e2e.py +++ b/tests/iac/test_iac_databricks_fake_e2e.py @@ -1002,7 +1002,7 @@ def test_apply_declares_the_keys_and_a_child_finds_its_parent_by_its_product_tag assert rc == 0, out assert fake.keys_on("main.sales.returns") == [ ("PRIMARY KEY", "fluid_pk_returns", "id"), - ("FOREIGN KEY", "fluid_fk_returns_region", "region", TABLE, "id"), + ("FOREIGN KEY", "fluid_fk_returns__region", "region", TABLE, "id"), ] statements = len(fake.sql_log) diff --git a/tests/iac/test_iac_databricks_keys.py b/tests/iac/test_iac_databricks_keys.py index f472f07a..322bf398 100644 --- a/tests/iac/test_iac_databricks_keys.py +++ b/tests/iac/test_iac_databricks_keys.py @@ -13,6 +13,8 @@ import base64 import json +import logging +import re from pathlib import Path from typing import Any, Dict, Iterator, List, Optional @@ -161,7 +163,7 @@ def test_one_table_with_the_referenced_primary_key(self): cands = {SITES: [keys.Candidate("main.geo.sites", ("site_id",))]} (fk,) = keys.foreign_keys(_keys(self.REF), cands) assert fk == keys.ForeignKey( - "fluid_fk_tickets_site_id", ("site_id",), "main.geo.sites", ("site_id",) + "fluid_fk_tickets__site_id", ("site_id",), "main.geo.sites", ("site_id",) ) def test_columns_that_cover_a_composite_key_form_one_in_its_order(self): @@ -172,7 +174,7 @@ def test_columns_that_cover_a_composite_key_form_one_in_its_order(self): cands = {"p.q_v1": [keys.Candidate("main.p.q", ("day", "site_id"))]} (fk,) = keys.foreign_keys(_keys(*refs), cands) assert fk.columns == ("d", "s") and fk.parent_columns == ("day", "site_id") - assert fk.name == "fluid_fk_tickets_d_s" + assert fk.name == "fluid_fk_tickets__d__s" def test_two_columns_to_one_key_are_two_foreign_keys(self): refs = ( @@ -233,7 +235,7 @@ def test_a_parent_without_that_primary_key_is_said(self): # ── the statements from a live table ─────────────────────────────────────── -FK = keys.ForeignKey("fluid_fk_tickets_site_id", ("site_id",), "main.geo.sites", ("site_id",)) +FK = keys.ForeignKey("fluid_fk_tickets__site_id", ("site_id",), "main.geo.sites", ("site_id",)) def _live(*constraints: Dict[str, Any]) -> Dict[str, Any]: @@ -263,7 +265,7 @@ def test_a_new_table_gets_its_primary_key_then_its_foreign_keys(self): assert [s.statement for s in plan.steps] == [ "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT `fluid_pk_tickets` PRIMARY KEY " "(`ticket_id`)", - "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT `fluid_fk_tickets_site_id` " + "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT `fluid_fk_tickets__site_id` " "FOREIGN KEY (`site_id`) REFERENCES `main`.`geo`.`sites` (`site_id`)", ] assert plan.exists is False @@ -276,7 +278,7 @@ def test_keys_in_place_under_any_name_are_kept(self): def test_a_changed_key_of_ours_is_dropped_and_added(self): live = _live( _pk("fluid_pk_tickets", "note"), - _fk("fluid_fk_tickets_site_id", "site_id", "main.old.sites"), + _fk("fluid_fk_tickets__site_id", "site_id", "main.old.sites"), ) plan = keys.plan_table("main.ops.tickets", self.K, [FK], live) assert [s.what.split(" ")[0] for s in plan.steps] == ["drop", "drop", "primary", "foreign"] @@ -293,18 +295,18 @@ def test_a_key_someone_else_declared_is_left_and_said(self): def test_keys_of_ours_the_contract_dropped_are_dropped_and_no_other(self): live = _live( _pk("fluid_pk_tickets", "ticket_id"), - _fk("fluid_fk_tickets_old", "note"), + _fk("fluid_fk_tickets__old", "note"), _fk("admin_fk", "x"), ) plan = keys.plan_table("main.ops.tickets", None, [], live) assert [s.statement for s in plan.steps] == [ - "ALTER TABLE `main`.`ops`.`tickets` DROP CONSTRAINT IF EXISTS `fluid_fk_tickets_old`", + "ALTER TABLE `main`.`ops`.`tickets` DROP CONSTRAINT IF EXISTS `fluid_fk_tickets__old`", "ALTER TABLE `main`.`ops`.`tickets` DROP CONSTRAINT IF EXISTS `fluid_pk_tickets`", ] def test_a_foreign_key_of_ours_whose_parent_is_unresolved_now_is_kept(self): live = _live( - _pk("fluid_pk_tickets", "ticket_id"), _fk("fluid_fk_tickets_site_id", "site_id") + _pk("fluid_pk_tickets", "ticket_id"), _fk("fluid_fk_tickets__site_id", "site_id") ) plan = keys.plan_table("main.ops.tickets", self.K, [], live) assert plan.steps == [] @@ -315,7 +317,7 @@ class TestRebuildScript: def test_every_primary_key_then_every_foreign_key_inside_one_handler(self): child = _keys(TestResolve.REF) parent = keys.TableKeys("main.geo.sites", ("site_id",)) - script = keys.rebuild_script([(child, [FK]), (parent, [])]) + script = keys.rebuild_script([child, parent], [("main.ops.tickets", FK)]) body = script.split("BEGIN\n", 1)[1] assert body.splitlines() == [ " DECLARE CONTINUE HANDLER FOR SQLEXCEPTION BEGIN END;", @@ -328,7 +330,7 @@ def test_every_primary_key_then_every_foreign_key_inside_one_handler(self): def test_nothing_to_declare_is_still_a_valid_file(self): only_refs = keys.TableKeys("main.ops.tickets", (), (TestResolve.REF,)) - script = keys.rebuild_script([(only_refs, [])]) + script = keys.rebuild_script([only_refs], []) assert "BEGIN" not in script and script.rstrip().endswith("AS fluid_keys;") @@ -358,7 +360,7 @@ def _emitted(tmp_path: Path, contract: Dict[str, Any], monkeypatch, found=None) """``emit``'s resources, its live lookup answered by ``found`` (product → candidates).""" asked: List[List[str]] = [] - def lookup(products): + def lookup(products, failures=None): asked.append(list(products)) return found or {} @@ -489,7 +491,7 @@ def deploy_tickets(fake: FakeDatabricks) -> None: PK_SITES = ("PRIMARY KEY", "fluid_pk_sites", "site_id") PK_TICKETS = ("PRIMARY KEY", "fluid_pk_tickets", "ticket_id") -FK_TICKETS = ("FOREIGN KEY", "fluid_fk_tickets_site_id", "site_id", "main.geo.sites", "site_id") +FK_TICKETS = ("FOREIGN KEY", "fluid_fk_tickets__site_id", "site_id", "main.geo.sites", "site_id") class TestHooks: @@ -502,7 +504,7 @@ def test_the_parent_then_the_child_across_two_contracts(self, fake): run.post(tickets()) assert fake.keys_on("main.ops.tickets") == [PK_TICKETS, FK_TICKETS] assert ( - "done: foreign key fluid_fk_tickets_site_id (site_id) references main.geo.sites (site_id)" + "done: foreign key fluid_fk_tickets__site_id (site_id) references main.geo.sites (site_id)" in run.text ) # One lookup, by tag, for the other product: the parent's location is never assumed. @@ -522,7 +524,7 @@ def test_a_decoy_someone_else_owns_does_not_become_the_parent(self, fake): run.post(tickets()) assert ( "FOREIGN KEY", - "fluid_fk_tickets_site_id", + "fluid_fk_tickets__site_id", "site_id", "prod.geo.sites", "site_id", @@ -581,7 +583,7 @@ def test_a_tag_removed_drops_our_key_and_leaves_an_admins(self, fake): PK_TICKETS, ("FOREIGN KEY", "admin_fk", "note", "main.geo.sites", "site_id"), ] - assert "done: drop foreign key fluid_fk_tickets_site_id" in run.text + assert "done: drop foreign key fluid_fk_tickets__site_id" in run.text def test_a_dry_run_prints_the_statements_and_runs_none(self, fake): run = Hooks(fake) @@ -642,7 +644,7 @@ def test_parent_and_child_in_one_contract_on_the_first_apply(self, fake): run.post(both) assert fake.keys_on("main.geo.visits") == [ ("PRIMARY KEY", "fluid_pk_visits", "visit_id"), - ("FOREIGN KEY", "fluid_fk_visits_site_id", "site_id", "main.geo.sites", "site_id"), + ("FOREIGN KEY", "fluid_fk_visits__site_id", "site_id", "main.geo.sites", "site_id"), ] assert not [s for s in fake.sql_log if "information_schema" in s] @@ -689,6 +691,233 @@ def post_build(self, contract, say): monkeypatch.setattr(engine, "_resolve_provider", lambda contract, requested: "databricks") monkeypatch.setattr(engine, "get_iac_plugin", lambda provider: Plugin()) args = type("Args", (), {"contract": None, "provider": None})() - engine.after_builds(args, {"id": "x"}) + engine.after_builds(args, logging.getLogger("t"), {"id": "x"}) assert Plugin.calls == [{"id": "x"}] assert "after the build: RuntimeError: boom" in capsys.readouterr().out + + +# ── the review round: each finding pinned ────────────────────────────────── +R = keys.Reference + + +class TestCompositeAmongOthers: + CANDS = { + "p.q_v1": [ + keys.Candidate("main.p.orders", ("order_id", "line_no")), + keys.Candidate("main.p.cust", ("cust_id",)), + ] + } + + def test_a_composite_key_and_a_single_one_to_two_tables_of_one_product(self): + refs = ( + R("o", "p.q_v1", "order_id"), + R("l", "p.q_v1", "line_no"), + R("c", "p.q_v1", "cust_id"), + ) + out = keys.foreign_keys(_keys(*refs), self.CANDS) + assert {(fk.columns, fk.parent) for fk in out} == { + (("o", "l"), "main.p.orders"), + (("c",), "main.p.cust"), + } + + def test_two_composite_keys_to_two_tables_of_one_product(self): + cands = { + "p.q_v1": [ + keys.Candidate("main.p.a", ("x", "y")), + keys.Candidate("main.p.b", ("u", "v")), + ] + } + refs = ( + R("cx", "p.q_v1", "x"), + R("cy", "p.q_v1", "y"), + R("cu", "p.q_v1", "u"), + R("cv", "p.q_v1", "v"), + ) + out = keys.foreign_keys(_keys(*refs), cands) + assert sorted((fk.columns, fk.parent) for fk in out) == [ + (("cu", "cv"), "main.p.b"), + (("cx", "cy"), "main.p.a"), + ] + + def test_half_a_composite_key_is_said_as_such(self): + notes: List[str] = [] + assert keys.foreign_keys(_keys(R("o", "p.q_v1", "order_id")), self.CANDS, notes) == [] + assert "part of the composite PRIMARY KEY (order_id, line_no) of main.p.orders" in notes[0] + + +class TestStaleForeignKeys: + K = keys.TableKeys( + "main.ops.tickets", ("ticket_id",), (R("a", "p.q_v1", "x"), R("b", "p.q_v1", "y")) + ) + LIVE = _live( + _pk("fluid_pk_tickets", "ticket_id"), + { + "foreign_key_constraint": { + "name": "fluid_fk_tickets__a", + "child_columns": ["a"], + "parent_table": "main.c.old", + "parent_columns": ["x"], + } + }, + ) + + def test_one_a_wanted_composite_key_covers_is_dropped(self): + want = keys.ForeignKey("fluid_fk_tickets__a__b", ("a", "b"), "main.c.new", ("x", "y")) + plan = keys.plan_table("main.ops.tickets", self.K, [want], self.LIVE, candidates={}) + assert [s.statement.split(" CONSTRAINT ")[-1].split(" ")[-1] for s in plan.fk_drops] == [ + "`fluid_fk_tickets__a`" + ] + assert [s.what.split(" (")[0] for s in plan.fk_adds] == [ + "foreign key fluid_fk_tickets__a__b" + ] + + def test_one_whose_product_now_resolves_elsewhere_or_nowhere_is_dropped(self): + plan = keys.plan_table("main.ops.tickets", self.K, [], self.LIVE, candidates={}) + assert len(plan.fk_drops) == 1 + + def test_one_whose_product_cannot_be_looked_up_now_is_kept(self): + plan = keys.plan_table( + "main.ops.tickets", self.K, [], self.LIVE, candidates={}, unknown={"p.q_v1"} + ) + assert plan.fk_drops == [] + + +def test_names_cannot_collide_in_a_schema(): + a = keys.constraint_name(keys.FK_PREFIX, "a", "b", "c") + b = keys.constraint_name(keys.FK_PREFIX, "a_b", "c") + c = keys.constraint_name(keys.FK_PREFIX, "a", "b_c") + assert len({a, b, c}) == 3 + assert a == "fluid_fk_a__b__c" and b == "fluid_fk_a_b__c" + ambiguous = keys.constraint_name(keys.FK_PREFIX, "a__b", "c") + assert re.fullmatch(r"fluid_fk_a__b__c_[0-9a-f]{8}", ambiguous) + + +def test_lookup_skips_a_table_it_cannot_read_or_name_and_says_so(): + class Runner: + def execute(self, statement: str) -> Dict[str, Any]: + rows = [ + ["main", "geo", "sites", SITES], + ["main", "geo", "busy", SITES], + ["main", "geo", "bad-name", SITES], + ] + return {"result": {"data_array": rows}} + + class Api: + def do(self, method: str, path: str, **_kw: Any) -> Dict[str, Any]: + if path.endswith("main.geo.busy"): + raise RuntimeError("429 too many") + pk = {"primary_key_constraint": {"name": "p", "child_columns": ["site_id"]}} + return {"owner": "me", "table_constraints": [pk]} + + notes: List[str] = [] + client = type("Client", (), {"api_client": Api()})() + found = keys.lookup(Runner(), client, [SITES], "me", notes) + assert found == {SITES: [keys.Candidate("main.geo.sites", ("site_id",), True)]} + assert any("main.geo.busy" in n and "429 too many" in n for n in notes) + assert any("main.geo.bad-name" in n and "no plain identifier" in n for n in notes) + + +class TestJobScriptCoverage: + def test_a_rebuilt_parent_gets_its_same_contract_childs_foreign_key_back( + self, tmp_path, monkeypatch + ): + _project(tmp_path) + contract = product( + "sales.star_v1", + expose("dim", [col("id", pk=True)], "sales"), + expose( + "fact", [col("fact_id", pk=True), col("dim_id", ref="sales.star_v1.id")], "sales" + ), + ) + contract = _dbt(contract) # the job writes dim only; fact is no build's + script = _script(_emitted(tmp_path, contract, monkeypatch)) + assert "ADD CONSTRAINT `fluid_pk_dim` PRIMARY KEY (`id`)" in script + assert ( + "`fluid_fk_fact__dim_id` FOREIGN KEY (`dim_id`) REFERENCES `main`.`sales`.`dim` (`id`)" + in script + ) + assert "fluid_pk_fact" not in script # fact is not rebuilt + + def test_validate_of_a_templated_location_does_not_raise(self, monkeypatch): + monkeypatch.setenv(dbx.WAREHOUSE_ENV, WAREHOUSE) + contract = _dbt(tickets(), repository="https://github.com/o/r?tag=v1.2.0") + contract["exposes"][0]["binding"]["location"]["database"] = "{{ env.CAT }}" + out = dbx._compile(contract, {dbx.WAREHOUSE_ENV: WAREHOUSE}, validating=True) + assert not out.refusals, [r.kind for r in out.refusals] + + +class TestEmitLookup: + def test_a_lookup_that_fails_is_a_plan_notice(self, tmp_path, monkeypatch): + _project(tmp_path) + monkeypatch.setenv(dbx.WAREHOUSE_ENV, WAREHOUSE) + + def failing(products, failures=None): + failures.append("RuntimeError: warehouse asleep") + return {} + + monkeypatch.setattr(dbx, "_live_key_lookup", failing) + contract = _dbt(tickets()) + with contract_source(tmp_path / "contract.fluid.yaml"): + PLUGIN.emit(contract) + notices = PLUGIN.plan_notices(contract) + assert any("could not look up now (RuntimeError: warehouse asleep)" in n for n in notices) + + def test_a_refused_contract_never_reaches_the_warehouse(self, tmp_path, monkeypatch): + _project(tmp_path) + monkeypatch.setenv(dbx.WAREHOUSE_ENV, WAREHOUSE) + asked: List[Any] = [] + monkeypatch.setattr(dbx, "_live_key_lookup", lambda *a, **k: asked.append(a) or {}) + contract = _dbt(tickets()) + contract["exposes"][0]["contract"]["schema"][2]["type"] = "not_a_type" + with contract_source(tmp_path / "contract.fluid.yaml"), pytest.raises(Exception): + PLUGIN.emit(contract) + assert asked == [] + + +class TestBuildPhaseHook: + def _engine(self, monkeypatch, plugin: Any): + from fluid_build.cli import _apply_opentofu_engine as engine + + monkeypatch.setattr(engine, "_resolve_provider", lambda contract, requested: "databricks") + monkeypatch.setattr(engine, "get_iac_plugin", lambda provider: plugin) + return engine + + def test_a_dry_run_declares_nothing(self, monkeypatch): + import fluid_build.build_runners as runners + + engine = self._engine(monkeypatch, object()) + monkeypatch.setattr(runners, "run_builds_from_args", lambda *a, **k: 0) + called: List[Any] = [] + monkeypatch.setattr(engine, "after_builds", lambda *a: called.append(a)) + args = type("Args", (), {"dry_run": True, "contract": None, "provider": None})() + assert engine.build_after_apply(args, logging.getLogger("t"), {"id": "x"}, {}) == 0 + assert called == [] + args.dry_run = False + assert engine.build_after_apply(args, logging.getLogger("t"), {"id": "x"}, {}) == 0 + assert len(called) == 1 + + def test_a_build_that_failed_part_way_is_still_followed_by_the_hook(self, monkeypatch): + import fluid_build.build_runners as runners + + engine = self._engine(monkeypatch, object()) + monkeypatch.setattr(runners, "run_builds_from_args", lambda *a, **k: 1) + called: List[Any] = [] + monkeypatch.setattr(engine, "after_builds", lambda *a: called.append(a)) + args = type("Args", (), {"dry_run": False, "contract": None, "provider": None})() + assert engine.build_after_apply(args, logging.getLogger("t"), {"id": "x"}, {}) == 1 + assert len(called) == 1 + + def test_the_hook_gets_the_contract_with_its_templates_resolved(self, monkeypatch): + seen: List[Dict[str, Any]] = [] + + class Plugin: + def post_build(self, contract, say): + seen.append(contract) + + engine = self._engine(monkeypatch, Plugin()) + monkeypatch.setenv("CAT", "main") + contract = tickets() + contract["exposes"][0]["binding"]["location"]["database"] = "{{ env.CAT }}" + args = type("Args", (), {"contract": None, "provider": None})() + engine.after_builds(args, logging.getLogger("t"), contract) + assert seen[0]["exposes"][0]["binding"]["location"]["database"] == "main" From 7da17bca6315e64769eaf28057f5160a098738cd Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:54:19 +0100 Subject: [PATCH 5/8] fix(databricks): a parent that cannot be read now keeps its foreign key The review round made one unreadable tagged table no longer discard the others, but it then left its product "resolved" without that table, so the next fluid apply dropped a correct foreign key to it and said the contract no longer declared it. A product with a table that could not be read is now unknown, like one whose lookup failed, and its keys stay. A key that is dropped says what changed: a new key replaces it, its columns reference another product now, or the contract no longer declares it. --- fluid_build/iac/providers/databricks_keys.py | 53 ++++++++++++-------- tests/iac/test_iac_databricks_keys.py | 34 +++++++++++++ 2 files changed, 67 insertions(+), 20 deletions(-) diff --git a/fluid_build/iac/providers/databricks_keys.py b/fluid_build/iac/providers/databricks_keys.py index 51e8bcae..cdea956a 100644 --- a/fluid_build/iac/providers/databricks_keys.py +++ b/fluid_build/iac/providers/databricks_keys.py @@ -88,6 +88,7 @@ Mapping, Optional, Sequence, + Set, Tuple, Union, ) @@ -453,14 +454,16 @@ def lookup( products: Iterable[str], owner: str = "", notes: Optional[List[str]] = None, + unreadable: Optional[Set[str]] = None, ) -> Dict[str, List[Candidate]]: """The tables tagged with each product id, with their live primary keys; one ``owner`` owns (the deploying identity) is ``trusted``. One query on ``system.information_schema.table_tags``, which lists what this identity can see, then one Tables API read per table found. A table that cannot be read, or whose name - is no plain identifier, is no candidate, and ``notes`` says which. Raises when the query - fails. + is no plain identifier, is no candidate, and ``notes`` says which; the product of one that + could not be read goes in ``unreadable``, since its tables are then not all known. Raises + when the query fails. """ said = notes if notes is not None else [] wanted = sorted({p for p in products if p}) @@ -494,6 +497,8 @@ def lookup( f"{full_name}, tagged {PRODUCT_TAG} = {product!r}, could not be read " f"({type(exc).__name__}: {scrub(exc)}), so no foreign key references it now." ) + if unreadable is not None: + unreadable.add(product) continue if live is None: continue @@ -572,7 +577,7 @@ def plan_table( that the contract changed or dropped is dropped; any other is left, and a contract key it is in the way of is noted. A foreign key of this module that no wanted key replaces is kept while its columns still reference the one product whose tables could not be looked up, or - whose tables still include its parent (:func:`_still_meant`): a parent unresolved for now. + whose tables still include its parent (:func:`_stale`): a parent unresolved for now. """ plan = TablePlan(full_name, exists=live is not None) live_pk, live_fks = live_keys(live or {}) @@ -617,36 +622,43 @@ def plan_table( ) del wanted[fk.columns] continue - if want is None and _still_meant(fk, keys, fks, candidates, unknown): + why = ( + "its parent changed" if want is not None else _stale(fk, keys, fks, candidates, unknown) + ) + if why is None: continue # its parent unresolved for now: kept - why = "its parent changed" if want is not None else "the contract no longer declares it" plan.fk_drops.append(_drop(full_name, fk.name, f"drop {fk.describe()}: {why}")) for fk in wanted.values(): plan.fk_adds.append(Step(fk.describe(), foreign_key_statement(full_name, fk))) return plan -def _still_meant( +def _stale( fk: ForeignKey, keys: Optional[TableKeys], wanted: Sequence[ForeignKey], candidates: Optional[Mapping[str, Sequence[Candidate]]], unknown: Collection[str], -) -> bool: - """Whether a foreign key this module declared, that the contract's resolved keys do not - include, still says what the contract does: no wanted key covers any of its columns, its - columns all reference one product, and that product's tables could not be looked up (or - were not, ``candidates`` ``None``) or still include its parent.""" - if any(set(w.columns) & set(fk.columns) for w in wanted): - return False +) -> Optional[str]: + """Why a foreign key this module declared, that the contract's resolved keys do not + include, no longer says what the contract does; ``None`` keeps it. It is kept while its + columns all reference one product whose tables could not all be looked up (or were not, + ``candidates`` ``None``), or that still include its parent: a parent unresolved for now.""" + covering = next((w for w in wanted if set(w.columns) & set(fk.columns)), None) + if covering is not None: + return f"the contract's foreign key ({', '.join(covering.columns)}) replaces it" refs = {r.column: r.product for r in keys.references} if keys is not None else {} products = {refs.get(c) for c in fk.columns} - if len(products) != 1 or None in products: - return False + if None in products: + return "the contract no longer declares it" + if len(products) != 1: + return "its columns now reference more than one product" (product,) = products if candidates is None or product in unknown: - return True - return any(c.full_name == fk.parent for c in candidates.get(str(product)) or ()) + return None + if any(c.full_name == fk.parent for c in candidates.get(str(product)) or ()): + return None + return f"the contract now references {product}, and {fk.parent} is not one of its tables" def _hint(error: str) -> str: @@ -696,21 +708,22 @@ def resolve( ) -> Tuple[Dict[str, List[Candidate]], FrozenSet[str]]: """Every table the contract's references can resolve to (its own exposes' declared primary keys, and the live tables of each other product it references), and the products - whose tables could not be looked up at all.""" + whose tables could not all be looked up.""" candidates: Dict[str, List[Candidate]] = dict(own_candidates(contract_id, keyed)) others = sorted({r.product for k in keyed for r in k.references if r.product != contract_id}) if not others: return candidates, frozenset() try: owner = deploying_identity(client) - candidates.update(lookup(runner, client, others, owner, notes)) + unreadable: Set[str] = set() + candidates.update(lookup(runner, client, others, owner, notes, unreadable)) except Exception as exc: # noqa: BLE001 — a lookup that fails declares no such key notes.append( f"foreign keys to {', '.join(others)}: their tables could not be looked up " f"({type(exc).__name__}: {scrub(exc)})." ) return candidates, frozenset(others) - return candidates, frozenset() + return candidates, frozenset(unreadable) def deploying_identity(client: Any) -> str: diff --git a/tests/iac/test_iac_databricks_keys.py b/tests/iac/test_iac_databricks_keys.py index 322bf398..c087e4a1 100644 --- a/tests/iac/test_iac_databricks_keys.py +++ b/tests/iac/test_iac_databricks_keys.py @@ -921,3 +921,37 @@ def post_build(self, contract, say): args = type("Args", (), {"contract": None, "provider": None})() engine.after_builds(args, logging.getLogger("t"), contract) assert seen[0]["exposes"][0]["binding"]["location"]["database"] == "main" + + +def test_a_parent_that_cannot_be_read_now_keeps_its_foreign_key(): + # A transient read failure of the parent must not drop a correct key, nor call it undeclared. + class Runner: + def execute(self, statement: str) -> Dict[str, Any]: + return {"result": {"data_array": [["main", "geo", "sites", SITES]]}} + + class Api: + def do(self, method: str, path: str, **_kw: Any) -> Dict[str, Any]: + if "scim" in path: + return {"userName": "me"} + raise RuntimeError("429 too many") + + client = type("Client", (), {"api_client": Api()})() + k = _keys(TestResolve.REF) + notes: List[str] = [] + cands, unknown = keys.resolve([k], TICKETS, Runner(), client, notes) + assert unknown == {SITES} and any("429 too many" in n for n in notes) + live = _live(_pk("fluid_pk_tickets", "ticket_id"), _fk("fluid_fk_tickets__site_id", "site_id")) + fks = keys.foreign_keys(k, cands) + plan = keys.plan_table("main.ops.tickets", k, fks, live, candidates=cands, unknown=unknown) + assert plan.fk_drops == [] and plan.fk_adds == [] + + +def test_a_dropped_key_says_what_changed(): + live = _live(_pk("fluid_pk_tickets", "ticket_id"), _fk("fluid_fk_tickets__site_id", "site_id")) + moved = _keys(keys.Reference("site_id", "geo.other_v1", "site_id")) + plan = keys.plan_table("main.ops.tickets", moved, [], live, candidates={}) + (drop,) = plan.fk_drops + assert ( + "the contract now references geo.other_v1, and main.geo.sites is not one of its tables" + in drop.what + ) From 499c3941db9771b402c0658728a6965c06866a14 Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:57:34 +0100 Subject: [PATCH 6/8] fix(databricks): a key stays put on a partial view of its parent's product When one table of the referenced product could not be read and a sibling with the same primary key could, the apply moved a correct foreign key to the sibling, where the next apply then left it. While not every table of that product can be read, a key that is in place stays there, and the apply says why it was not moved. --- fluid_build/iac/providers/databricks_keys.py | 16 ++++++++++++++++ tests/iac/test_iac_databricks_keys.py | 15 +++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/fluid_build/iac/providers/databricks_keys.py b/fluid_build/iac/providers/databricks_keys.py index cdea956a..70e66a08 100644 --- a/fluid_build/iac/providers/databricks_keys.py +++ b/fluid_build/iac/providers/databricks_keys.py @@ -622,6 +622,15 @@ def plan_table( ) del wanted[fk.columns] continue + if want is not None and _product(fk, keys) in unknown: + # Not every table of its product could be read, so ``want`` is a guess from part of + # them: the key stays where it is rather than move to a sibling of its parent. + del wanted[fk.columns] + plan.notes.append( + f"foreign key ({', '.join(fk.columns)}) is kept on {fk.parent}: not every table " + f"of {_product(fk, keys)} could be read, so it is not moved to {want.parent}." + ) + continue why = ( "its parent changed" if want is not None else _stale(fk, keys, fks, candidates, unknown) ) @@ -633,6 +642,13 @@ def plan_table( return plan +def _product(fk: ForeignKey, keys: Optional[TableKeys]) -> Optional[str]: + """The one product a foreign key's columns reference now, or ``None``.""" + refs = {r.column: r.product for r in keys.references} if keys is not None else {} + products = {refs.get(c) for c in fk.columns} + return next(iter(products)) if len(products) == 1 else None + + def _stale( fk: ForeignKey, keys: Optional[TableKeys], diff --git a/tests/iac/test_iac_databricks_keys.py b/tests/iac/test_iac_databricks_keys.py index c087e4a1..a5111795 100644 --- a/tests/iac/test_iac_databricks_keys.py +++ b/tests/iac/test_iac_databricks_keys.py @@ -955,3 +955,18 @@ def test_a_dropped_key_says_what_changed(): "the contract now references geo.other_v1, and main.geo.sites is not one of its tables" in drop.what ) + + +def test_a_key_is_not_moved_to_a_sibling_on_a_partial_view_of_its_product(): + # main.geo.sites (the live parent) could not be read; main.geo.sites2 could and has the + # same primary key. The key stays on sites rather than move to sites2. + k = _keys(TestResolve.REF) + live = _live(_pk("fluid_pk_tickets", "ticket_id"), _fk("fluid_fk_tickets__site_id", "site_id")) + cands = {SITES: [keys.Candidate("main.geo.sites2", ("site_id",))]} + fks = keys.foreign_keys(k, cands) + assert fks[0].parent == "main.geo.sites2" + plan = keys.plan_table("main.ops.tickets", k, fks, live, candidates=cands, unknown={SITES}) + assert plan.fk_drops == [] and plan.fk_adds == [] + assert ( + "is kept on main.geo.sites: not every table of geo.sites_v1 could be read" in plan.notes[0] + ) From 8be192a06430a74a691d149bd54c50402f0a72af Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:59:26 +0100 Subject: [PATCH 7/8] fix(databricks): a key whose parent's key columns changed says so The partial-view guard kept a key only to stop a move to another table; a key whose parent is the same and whose parent columns changed is replaced, and the drop says that instead of "its parent changed" (fix re-check, LOW). --- fluid_build/iac/providers/databricks_keys.py | 13 +++++++++---- tests/iac/test_iac_databricks_keys.py | 10 ++++++++++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/fluid_build/iac/providers/databricks_keys.py b/fluid_build/iac/providers/databricks_keys.py index 70e66a08..ca0b862b 100644 --- a/fluid_build/iac/providers/databricks_keys.py +++ b/fluid_build/iac/providers/databricks_keys.py @@ -622,7 +622,7 @@ def plan_table( ) del wanted[fk.columns] continue - if want is not None and _product(fk, keys) in unknown: + if want is not None and want.parent != fk.parent and _product(fk, keys) in unknown: # Not every table of its product could be read, so ``want`` is a guess from part of # them: the key stays where it is rather than move to a sibling of its parent. del wanted[fk.columns] @@ -631,9 +631,14 @@ def plan_table( f"of {_product(fk, keys)} could be read, so it is not moved to {want.parent}." ) continue - why = ( - "its parent changed" if want is not None else _stale(fk, keys, fks, candidates, unknown) - ) + if want is not None: + why: Optional[str] = ( + "its parent changed" + if want.parent != fk.parent + else "its parent's key columns changed" + ) + else: + why = _stale(fk, keys, fks, candidates, unknown) if why is None: continue # its parent unresolved for now: kept plan.fk_drops.append(_drop(full_name, fk.name, f"drop {fk.describe()}: {why}")) diff --git a/tests/iac/test_iac_databricks_keys.py b/tests/iac/test_iac_databricks_keys.py index a5111795..961fd688 100644 --- a/tests/iac/test_iac_databricks_keys.py +++ b/tests/iac/test_iac_databricks_keys.py @@ -970,3 +970,13 @@ def test_a_key_is_not_moved_to_a_sibling_on_a_partial_view_of_its_product(): assert ( "is kept on main.geo.sites: not every table of geo.sites_v1 could be read" in plan.notes[0] ) + + +def test_a_key_whose_parent_key_columns_changed_is_replaced_and_says_so(): + live = _live(_pk("fluid_pk_tickets", "ticket_id"), _fk("fluid_fk_tickets__site_id", "site_id")) + want = keys.ForeignKey("fluid_fk_tickets__site_id", ("site_id",), "main.geo.sites", ("code",)) + plan = keys.plan_table( + "main.ops.tickets", _keys(TestResolve.REF), [want], live, candidates={}, unknown={SITES} + ) + assert [s.what.split(": ")[-1] for s in plan.fk_drops] == ["its parent's key columns changed"] + assert len(plan.fk_adds) == 1 and plan.notes == [] From 45c8a73f3e78449c32a0453e418b33a7a47ff0b6 Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Thu, 12 Nov 2026 22:17:18 +0100 Subject: [PATCH 8/8] test(databricks): address the code-quality review on #860 The expected statements in a key-plan test are joined with an explicit +, not adjacent literals; the unused Optional import goes; REBUILD_TIMEOUT_SECONDS is in __all__, since databricks_jobs imports it (the per-file analysis saw no use). --- fluid_build/iac/providers/databricks_keys.py | 1 + tests/iac/test_iac_databricks_keys.py | 11 ++++++----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/fluid_build/iac/providers/databricks_keys.py b/fluid_build/iac/providers/databricks_keys.py index ca0b862b..205949b8 100644 --- a/fluid_build/iac/providers/databricks_keys.py +++ b/fluid_build/iac/providers/databricks_keys.py @@ -925,6 +925,7 @@ def rebuild_script(rebuilt: Sequence[TableKeys], foreign: Sequence[Tuple[str, Fo "KeyTarget", "REBUILD_FILE", "REBUILD_TASK", + "REBUILD_TIMEOUT_SECONDS", "TableKeys", "compile_keys", "foreign_keys", diff --git a/tests/iac/test_iac_databricks_keys.py b/tests/iac/test_iac_databricks_keys.py index 961fd688..c9d63e9c 100644 --- a/tests/iac/test_iac_databricks_keys.py +++ b/tests/iac/test_iac_databricks_keys.py @@ -16,7 +16,7 @@ import logging import re from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List import pytest @@ -262,11 +262,12 @@ class TestPlan: def test_a_new_table_gets_its_primary_key_then_its_foreign_keys(self): plan = keys.plan_table("main.ops.tickets", self.K, [FK], None) + table = "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT " assert [s.statement for s in plan.steps] == [ - "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT `fluid_pk_tickets` PRIMARY KEY " - "(`ticket_id`)", - "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT `fluid_fk_tickets__site_id` " - "FOREIGN KEY (`site_id`) REFERENCES `main`.`geo`.`sites` (`site_id`)", + table + "`fluid_pk_tickets` PRIMARY KEY (`ticket_id`)", + table + + "`fluid_fk_tickets__site_id` FOREIGN KEY (`site_id`) REFERENCES " + + "`main`.`geo`.`sites` (`site_id`)", ] assert plan.exists is False