diff --git a/CHANGELOG.md b/CHANGELOG.md index 166198fd..aba2a06f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,29 @@ 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, 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 - **`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..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, @@ -363,6 +364,47 @@ 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 (: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 not getattr(args, "dry_run", False): + after_builds(args, logger, contract) + return rc + + +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 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(resolved, 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(resolved, say=cprint_literal) + except Exception as exc: # noqa: BLE001 — advisory: the build's own result stands + cprint_literal(f" after the build: {type(exc).__name__}: {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. @@ -378,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/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/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/databricks.py b/fluid_build/iac/providers/databricks.py index 3a4297de..d1c17fea 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,18 @@ dbt_written_exposes, run_as_allowlist, ) +from .databricks_keys import ( + Candidate, + KeyTarget, + TableKeys, + compile_keys, + deploying_identity, + 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 +317,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 +956,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 +994,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 +1012,7 @@ def _resolve_table( comment=comment, column_comments=tuple(column_comments), classified=tuple(classified), + keys=keys, ) @@ -1304,9 +1331,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 +1381,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 +1390,7 @@ def _compile( warehouse, validating=validating, run_as_allowed=run_as_allowlist(env), + rebuild_script=_rebuild_scripts(compiled, contract, builds, key_lookup), ), ) @@ -1369,6 +1403,138 @@ 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``). + + 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. + """ + 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)) + 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)) + + +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)}).") + + +#: 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, 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, 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( compiled: _Compiled, exposures: Sequence[Tuple[int, Mapping[str, Any]]], @@ -1387,6 +1553,7 @@ def _resolve_tables( contract=contract, validating=validating, builds=builds, + warnings=compiled.warnings, ) if table is not None: compiled.tables.append(table) @@ -2048,6 +2215,17 @@ def emit( 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) @@ -2186,6 +2364,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 +2379,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 @@ -2226,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_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..205949b8 --- /dev/null +++ b/fluid_build/iac/providers/databricks_keys.py @@ -0,0 +1,938 @@ +# 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 +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 +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 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. +""" + +from __future__ import annotations + +import hashlib +import os +from contextlib import suppress +from dataclasses import dataclass, field +from typing import ( + Any, + Callable, + Collection, + Dict, + FrozenSet, + Iterable, + List, + Mapping, + Optional, + Sequence, + Set, + 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, + deploy_identity, + owner_matches, + 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. + + ``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: + 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 ``__``, 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: + 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, 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) + + +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 + 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: + 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: + 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: + 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], + 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; 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}) + 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) + 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." + ) + if unreadable is not None: + unreadable.add(product) + continue + if live is None: + continue + primary, _ = live_keys(live) + 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 + + +# ── 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]], + *, + 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 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:`_stale`): a parent unresolved for now. + """ + 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} + 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 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] + 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 + 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}")) + for fk in wanted.values(): + plan.fk_adds.append(Step(fk.describe(), foreign_key_statement(full_name, fk))) + 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], + wanted: Sequence[ForeignKey], + candidates: Optional[Mapping[str, Sequence[Candidate]]], + unknown: Collection[str], +) -> 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 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 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: + """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], +) -> 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 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) + 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(unreadable) + + +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 + 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)) + # 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, candidates=candidates, unknown=unknown + ) + 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(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. + + 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: 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 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;", ""]) + 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", + "REBUILD_TIMEOUT_SECONDS", + "TableKeys", + "compile_keys", + "foreign_keys", + "lookup", + "own_candidates", + "plan_table", + "rebuild_script", + "reconcile", + "resolve", +] 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: 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..750dfe03 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..c9d63e9c --- /dev/null +++ b/tests/iac/test_iac_databricks_keys.py @@ -0,0 +1,983 @@ +# 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 +import logging +import re +from pathlib import Path +from typing import Any, Dict, Iterator, List + +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_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: [ + 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) + table = "ALTER TABLE `main`.`ops`.`tickets` ADD CONSTRAINT " + assert [s.statement for s in plan.steps] == [ + 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 + + 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, parent], [("main.ops.tickets", FK)]) + 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, failures=None): + 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_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) + 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, 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" + + +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 + ) + + +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] + ) + + +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 == []