contractforge-databricks 0.1.0__py3-none-any.whl

contractforge_databricks/runtime/maintenance.py

@@ -0,0 +1,43 @@
+"""Post-write Databricks maintenance hooks."""
+
+from __future__ import annotations
+
+from contractforge_core.execution import ExecutionOutcome
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.contract_extensions import databricks_extensions
+from contractforge_databricks.execution import SqlRunner
+from contractforge_databricks.maintenance import MaintenancePlan, execute_maintenance_plan
+
+
+def run_post_write_maintenance(
+    *,
+    runner: SqlRunner,
+    contract: SemanticContract,
+    target_table: str,
+    outcome: ExecutionOutcome | None,
+    rows_written: int,
+) -> tuple[str, ...]:
+    extensions = databricks_extensions(contract)
+    if not extensions.get("optimize_after_write"):
+        return ()
+    rows_written = int((outcome.metrics if outcome else {}).get("rows_written", rows_written) or 0)
+    if rows_written <= 0:
+        return ()
+    return execute_maintenance_plan(
+        runner,
+        MaintenancePlan(
+            target_table=target_table,
+            optimize=True,
+            zorder_columns=_tuple(extensions.get("zorder_columns")),
+        ),
+    )
+
+
+def _tuple(value: object) -> tuple[str, ...]:
+    if value is None:
+        return ()
+    if isinstance(value, str):
+        return (value,)
+    if isinstance(value, (list, tuple, set)):
+        return tuple(str(item) for item in value)
+    return (str(value),)

contractforge_databricks/runtime/merge_validation.py

@@ -0,0 +1,98 @@
+"""Prepared-source safety checks for Databricks MERGE writes."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.quality import QualityRuleResult, quality_status
+from contractforge_core.runtime import PreparedInput, QueryOne
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.sql import quote_identifier, quote_table_name
+
+MERGE_WRITE_MODES = {"scd1_upsert", "scd2_historical", "snapshot_soft_delete"}
+
+
+def validate_merge_source_safety(
+    *,
+    contract: SemanticContract,
+    prepared: PreparedInput,
+    query_one: QueryOne | None,
+    quality_results: tuple[QualityRuleResult, ...] = (),
+) -> dict[str, Any]:
+    """Validate source key safety before executing Databricks MERGE patterns."""
+    if contract.write.mode not in MERGE_WRITE_MODES:
+        return {"status": "SKIPPED", "reason": "not_merge_mode"}
+    if query_one is None:
+        return {"status": "SKIPPED", "reason": "query_one_not_configured"}
+    if not contract.write.merge_keys or prepared.rows_read <= 0:
+        return {"status": "SKIPPED", "reason": "no_merge_keys_or_rows"}
+
+    _validate_columns(contract.write.merge_keys, prepared.source_columns)
+    null_row = query_one(render_merge_key_nulls_sql(prepared.source_view, contract.write.merge_keys))
+    all_null_count = _int_row_value(null_row, "all_keys_null_rows")
+    if all_null_count == prepared.rows_read:
+        raise ValueError(
+            f"mode={contract.write.mode} received {prepared.rows_read} rows with fully null merge_keys. "
+            f"keys={list(contract.write.merge_keys)}. Fix the source or add quality_rules.not_null."
+        )
+
+    if _skip_duplicate_check(contract, quality_results):
+        return {"status": "PASSED", "all_null_key_rows": all_null_count, "duplicate_check": "SKIPPED"}
+
+    duplicate_row = query_one(render_merge_key_duplicates_sql(prepared.source_view, contract.write.merge_keys))
+    duplicate_groups = _int_row_value(duplicate_row, "duplicate_key_groups")
+    duplicate_rows = _int_row_value(duplicate_row, "duplicate_rows")
+    if duplicate_groups:
+        raise ValueError(
+            f"mode={contract.write.mode} received {duplicate_rows} duplicate source rows across "
+            f"{duplicate_groups} merge_key groups. keys={list(contract.write.merge_keys)}. "
+            "Fix the composite key, declare quality_rules.unique_key, or apply transform.deduplicate."
+        )
+    return {"status": "PASSED", "all_null_key_rows": all_null_count, "duplicate_key_groups": duplicate_groups}
+
+
+def render_merge_key_nulls_sql(source_view: str, merge_keys: tuple[str, ...]) -> str:
+    all_keys_null = " AND ".join(f"{quote_identifier(key)} IS NULL" for key in merge_keys)
+    return (
+        f"SELECT count(*) AS all_keys_null_rows "
+        f"FROM {quote_table_name(source_view)} WHERE {all_keys_null}"
+    )
+
+
+def render_merge_key_duplicates_sql(source_view: str, merge_keys: tuple[str, ...]) -> str:
+    key_list = ", ".join(quote_identifier(key) for key in merge_keys)
+    return (
+        "SELECT count(*) AS duplicate_key_groups, coalesce(sum(row_count), 0) AS duplicate_rows "
+        f"FROM (SELECT {key_list}, count(*) AS row_count FROM {quote_table_name(source_view)} "
+        f"GROUP BY {key_list} HAVING count(*) > 1)"
+    )
+
+
+def _validate_columns(keys: tuple[str, ...], source_columns: tuple[str, ...]) -> None:
+    if not source_columns:
+        return
+    missing = [key for key in keys if key not in source_columns]
+    if missing:
+        raise ValueError(f"merge_keys missing from prepared source columns: {missing}")
+
+
+def _skip_duplicate_check(contract: SemanticContract, quality_results: tuple[QualityRuleResult, ...]) -> bool:
+    if quality_status(quality_results) != "PASSED":
+        return False
+    unique_rules = tuple(rule for rule in contract.quality if rule.rule == "unique_key")
+    return any(set(rule.columns) == set(contract.write.merge_keys) for rule in unique_rules)
+
+
+def _int_row_value(row: Any, key: str) -> int:
+    if row is None:
+        return 0
+    if isinstance(row, dict):
+        value = row.get(key)
+    elif hasattr(row, "asDict"):
+        value = row.asDict().get(key)
+    else:
+        value = getattr(row, key, None)
+    try:
+        return int(value or 0)
+    except (TypeError, ValueError):
+        return 0

contractforge_databricks/runtime/metadata.py

@@ -0,0 +1,21 @@
+"""Runtime contract metadata payload helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.semantic import SemanticContract
+
+
+def contract_metadata(contract: SemanticContract, operations: dict[str, Any]) -> dict[str, Any]:
+    return {
+        "description": operations.get("description"),
+        "owner": contract.governance.owner if contract.governance else None,
+        "domain": contract.target.domain,
+        "tags": operations.get("tags"),
+        "sla": operations.get("sla"),
+        "runtime_parameters": operations.get("runtime_parameters"),
+        "operations": contract.operations.metadata if contract.operations else None,
+        "applied_presets": operations.get("applied_presets"),
+        "target_schema": contract.target.namespace,
+    }

contractforge_databricks/runtime/metrics.py

@@ -0,0 +1,34 @@
+"""Runtime write metric collection for Databricks."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.metrics import normalize_rows_written
+from contractforge_core.runtime import QueryOne
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.metrics import (
+    latest_operation_metrics_from_history_row,
+    render_delta_history_query,
+    resolve_write_metrics,
+)
+
+
+def collect_write_metrics(
+    *,
+    contract: SemanticContract,
+    target_table: str,
+    rows_written: int,
+    query_one: QueryOne | None,
+) -> tuple[int, dict[str, Any]]:
+    delta_metrics = {}
+    if query_one is not None:
+        delta_metrics = latest_operation_metrics_from_history_row(
+            query_one(render_delta_history_query(target_table=target_table))
+        )
+    row_metrics, operation_metrics, metrics_source = resolve_write_metrics(contract, rows_written, delta_metrics)
+    operation_metrics["metrics_source"] = metrics_source
+    normalized = normalize_rows_written(rows_written, row_metrics)
+    row_metrics["rows_affected"] = normalized
+    operation_metrics["normalizedRowMetrics"] = row_metrics
+    return normalized, operation_metrics

contractforge_databricks/runtime/models.py

@@ -0,0 +1,32 @@
+"""Runtime input models for Databricks ingestion orchestration."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from contractforge_core.runtime import PreparedInput
+from contractforge_databricks.runtime.hooks import DatabricksIngestionHooks
+
+
+PreparedViewInput = PreparedInput
+
+
+@dataclass(frozen=True)
+class DatabricksIngestOptions:
+    catalog: str = "main"
+    schema: str = "ops"
+    dry_run: bool = False
+    ensure_table: bool = True
+    lock_enabled: bool = False
+    lock_owner: str | None = None
+    idempotency_key: str | None = None
+    idempotency_policy: str = "always_run"
+    quality_action: str = "fail"
+    run_id: str | None = None
+    run_id_factory: Callable[[], str] | None = None
+    runtime_metadata: dict[str, Any] | None = None
+    target_schema: dict[str, str] | None = None
+    allow_review_required: bool = False
+    raise_on_failure: bool = True
+    hooks: DatabricksIngestionHooks | None = None

contractforge_databricks/runtime/options.py

@@ -0,0 +1,33 @@
+"""Databricks runtime option resolution from core contract semantics."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import Any
+
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.contract_extensions import databricks_extensions
+from contractforge_databricks.runtime.hooks import DatabricksIngestionHooks
+from contractforge_databricks.runtime.models import DatabricksIngestOptions
+
+
+def effective_ingest_options(contract: SemanticContract, options: DatabricksIngestOptions) -> DatabricksIngestOptions:
+    metadata = contract.operations.metadata if contract.operations and contract.operations.metadata else {}
+    extensions = databricks_extensions(contract)
+    updates: dict[str, Any] = {}
+    if options.idempotency_key is None and metadata.get("idempotency_key"):
+        updates["idempotency_key"] = str(metadata["idempotency_key"])
+    if options.idempotency_policy == "always_run" and metadata.get("idempotency_policy"):
+        updates["idempotency_policy"] = str(metadata["idempotency_policy"])
+    if options.quality_action == "fail" and metadata.get("on_quality_fail"):
+        updates["quality_action"] = str(metadata["on_quality_fail"])
+    if options.hooks is None and extensions.get("hooks") is not None:
+        hooks = extensions["hooks"]
+        if isinstance(hooks, dict):
+            hooks = DatabricksIngestionHooks(**hooks)
+        if not isinstance(hooks, DatabricksIngestionHooks):
+            raise ValueError("extensions.databricks.hooks must be DatabricksIngestionHooks")
+        updates["hooks"] = hooks
+    if not options.lock_enabled and extensions.get("lock_enabled"):
+        updates["lock_enabled"] = True
+    return replace(options, **updates) if updates else options

contractforge_databricks/runtime/orchestration_context.py

@@ -0,0 +1,185 @@
+"""Private helpers for Databricks runtime orchestration."""
+
+from __future__ import annotations
+
+from typing import Any
+from typing import NamedTuple
+
+from contractforge_core.contracts import semantic_contract_from_mapping
+from contractforge_core.errors import raise_for_failure_result
+from contractforge_core.quality import QualityRuleResult, quality_policy_status
+from contractforge_core.runtime import PreparedInput, QueryOne
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.contract_extensions import normalize_databricks_contract
+from contractforge_databricks.evidence import EvidenceWriter
+from contractforge_databricks.execution import SqlRunner
+from contractforge_databricks.rendering.names import target_full_name
+from contractforge_databricks.runtime.dry_run import finalize_dry_run
+from contractforge_databricks.runtime.errors import error_log_payload
+from contractforge_databricks.runtime.finalization import finalize_ingest
+from contractforge_databricks.runtime.models import DatabricksIngestOptions
+from contractforge_databricks.runtime.options import effective_ingest_options
+from contractforge_databricks.runtime.success import finalize_success
+from contractforge_databricks.runtime.utils import resolve_run_id, utc_now_str
+from contractforge_databricks.runtime.write_flow import WriteFlowResult
+from contractforge_databricks.security import exception_message
+from contractforge_databricks.state import StateWriter
+
+
+class RuntimeContext(NamedTuple):
+    semantic: SemanticContract
+    opts: DatabricksIngestOptions
+    target: str
+    run_id: str
+    started: str
+    evidence: EvidenceWriter
+    state: StateWriter
+    quality_status_value: str
+
+
+class RuntimeProgress:
+    __slots__ = ("prepared", "schema_changes", "governance_results")
+
+    def __init__(self, prepared: PreparedInput) -> None:
+        self.prepared = prepared
+        self.schema_changes: dict[str, Any] = {}
+        self.governance_results: dict[str, Any] = {}
+
+
+def build_runtime_context(
+    contract: dict[str, Any] | SemanticContract,
+    *,
+    runner: SqlRunner,
+    options: DatabricksIngestOptions | None,
+    query_one: QueryOne | None,
+    quality_results: tuple[QualityRuleResult, ...],
+) -> RuntimeContext:
+    base_opts = options or DatabricksIngestOptions()
+    semantic = contract if isinstance(contract, SemanticContract) else semantic_contract_from_mapping(normalize_databricks_contract(contract))
+    opts = effective_ingest_options(semantic, base_opts)
+    target = target_full_name(semantic)
+    return RuntimeContext(
+        semantic=semantic,
+        opts=opts,
+        target=target,
+        run_id=resolve_run_id(opts.run_id, opts.run_id_factory),
+        started=utc_now_str(),
+        evidence=EvidenceWriter(runner, catalog=opts.catalog, schema=opts.schema),
+        state=StateWriter(runner, catalog=opts.catalog, schema=opts.schema, query_one=query_one),
+        quality_status_value=quality_policy_status(quality_results, on_quality_fail=opts.quality_action),
+    )
+
+
+def complete_result(ctx: RuntimeContext, result: dict[str, Any]) -> dict[str, Any]:
+    if ctx.opts.hooks and ctx.opts.hooks.after_finalize:
+        ctx.opts.hooks.after_finalize(ctx.semantic, result)
+    if ctx.opts.raise_on_failure:
+        raise_for_failure_result(result)
+    return result
+
+
+def finalize_skipped_result(
+    ctx: RuntimeContext,
+    progress: RuntimeProgress,
+    *,
+    quality_results: tuple[QualityRuleResult, ...],
+    skipped_by_run_id: object,
+) -> dict[str, Any]:
+    return finalize_ingest(
+        ctx.evidence,
+        ctx.state,
+        ctx.semantic,
+        progress.prepared,
+        ctx.opts,
+        ctx.run_id,
+        ctx.target,
+        "SKIPPED",
+        ctx.started,
+        rows_written=0,
+        quality_status_value="SKIPPED",
+        quality_results=quality_results,
+        skip_reason="idempotency_key_already_succeeded",
+        skipped_by_run_id=str(skipped_by_run_id) if skipped_by_run_id else None,
+    )
+
+
+def finalize_dry_run_result(ctx: RuntimeContext, progress: RuntimeProgress) -> dict[str, Any]:
+    return finalize_dry_run(
+        evidence=ctx.evidence,
+        state=ctx.state,
+        contract=ctx.semantic,
+        prepared=progress.prepared,
+        opts=ctx.opts,
+        run_id=ctx.run_id,
+        target=ctx.target,
+        started=ctx.started,
+        quality_status_value=ctx.quality_status_value,
+    )
+
+
+def finalize_success_result(
+    ctx: RuntimeContext,
+    progress: RuntimeProgress,
+    *,
+    write_flow: WriteFlowResult,
+    quality_results: tuple[QualityRuleResult, ...],
+    query_one: QueryOne | None,
+) -> dict[str, Any]:
+    return finalize_success(
+        evidence=ctx.evidence,
+        state=ctx.state,
+        contract=ctx.semantic,
+        prepared=progress.prepared,
+        opts=ctx.opts,
+        run_id=ctx.run_id,
+        target=ctx.target,
+        started=ctx.started,
+        outcome=write_flow.outcome,
+        logical_rows_written=write_flow.logical_rows_written,
+        quality_status_value=ctx.quality_status_value,
+        quality_results=quality_results,
+        schema_changes=progress.schema_changes,
+        governance_results=progress.governance_results,
+        write_started_at=write_flow.write_started_at,
+        write_finished_at=write_flow.write_finished_at,
+        stage_durations=write_flow.stage_durations,
+        query_one=query_one,
+    )
+
+
+def finalize_failure_result(
+    ctx: RuntimeContext,
+    progress: RuntimeProgress,
+    exc: Exception,
+    *,
+    quality_results: tuple[QualityRuleResult, ...],
+) -> dict[str, Any]:
+    error_message = exception_message(exc)
+    if not ctx.opts.dry_run:
+        ctx.evidence.write_error_log(
+            error_log_payload(
+                exc,
+                run_id=ctx.run_id,
+                target=ctx.target,
+                source_table=progress.prepared.source_name or progress.prepared.source_view,
+                mode=ctx.semantic.write.mode,
+                runtime_metadata=ctx.opts.runtime_metadata,
+            )
+        )
+    return finalize_ingest(
+        ctx.evidence,
+        ctx.state,
+        ctx.semantic,
+        progress.prepared,
+        ctx.opts,
+        ctx.run_id,
+        ctx.target,
+        "FAILED",
+        ctx.started,
+        rows_written=0,
+        quality_status_value=ctx.quality_status_value,
+        quality_results=quality_results,
+        error_message=error_message,
+        schema_changes=progress.schema_changes,
+        governance_results=progress.governance_results,
+    )

contractforge_databricks/runtime/orchestrator.py

@@ -0,0 +1,147 @@
+"""Databricks runtime orchestration over prepared source views."""
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.quality import QualityRuleResult, quality_status
+from contractforge_core.runtime import PreparedInput, QueryOne
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.adapter import DatabricksAdapter
+from contractforge_databricks.execution import SqlRunner
+from contractforge_databricks.rendering.names import target_full_name
+from contractforge_databricks.runtime.hooks import apply_prepared_hook
+from contractforge_databricks.runtime.models import DatabricksIngestOptions
+from contractforge_databricks.runtime.orchestration_context import (
+    RuntimeContext,
+    RuntimeProgress,
+    build_runtime_context,
+    complete_result,
+    finalize_dry_run_result,
+    finalize_failure_result,
+    finalize_skipped_result,
+    finalize_success_result,
+)
+from contractforge_databricks.runtime.write_flow import execute_runtime_write_flow
+from contractforge_databricks.state import render_find_idempotent_run_sql
+
+
+def ingest_databricks_contract(
+    contract: dict[str, Any] | SemanticContract,
+    *,
+    runner: SqlRunner,
+    prepared: PreparedInput,
+    options: DatabricksIngestOptions | None = None,
+    query_one: QueryOne | None = None,
+    quality_results: tuple[QualityRuleResult, ...] = (),
+) -> dict[str, Any]:
+    """Execute one Databricks contract over an already prepared source view."""
+    ctx = build_runtime_context(
+        contract,
+        runner=runner,
+        options=options,
+        query_one=query_one,
+        quality_results=quality_results,
+    )
+    progress = RuntimeProgress(prepared=prepared)
+    raw_quality_status = quality_status(quality_results)
+
+    try:
+        _validate_planning(ctx.semantic, ctx.opts)
+        progress.prepared = _apply_after_prepare(ctx, progress.prepared)
+        skipped = _idempotency_skip(ctx.target, ctx.opts, query_one)
+        if skipped:
+            return complete_result(
+                ctx,
+                finalize_skipped_result(
+                    ctx,
+                    progress,
+                    quality_results=quality_results,
+                    skipped_by_run_id=skipped.get("run_id"),
+                ),
+            )
+        _raise_for_quality_failure(raw_quality_status, ctx.opts)
+        _acquire_lock(ctx)
+        progress.prepared = _apply_before_write(ctx, progress.prepared)
+        if ctx.opts.dry_run:
+            return complete_result(ctx, finalize_dry_run_result(ctx, progress))
+        write_flow = execute_runtime_write_flow(
+            runner=runner,
+            evidence=ctx.evidence,
+            contract=ctx.semantic,
+            prepared=progress.prepared,
+            opts=ctx.opts,
+            run_id=ctx.run_id,
+            target=ctx.target,
+            query_one=query_one,
+            quality_results=quality_results,
+        )
+        progress.schema_changes = write_flow.schema_changes
+        progress.governance_results = write_flow.governance_results
+        if ctx.opts.hooks and ctx.opts.hooks.after_write:
+            ctx.opts.hooks.after_write(ctx.semantic, progress.prepared, write_flow.outcome)
+        return complete_result(
+            ctx,
+            finalize_success_result(
+                ctx,
+                progress,
+                write_flow=write_flow,
+                quality_results=quality_results,
+                query_one=query_one,
+            ),
+        )
+    except Exception as exc:
+        return complete_result(ctx, finalize_failure_result(ctx, progress, exc, quality_results=quality_results))
+    finally:
+        _release_lock(ctx)
+
+
+def _apply_after_prepare(ctx: RuntimeContext, prepared: PreparedInput) -> PreparedInput:
+    return apply_prepared_hook(ctx.opts.hooks.after_prepare if ctx.opts.hooks else None, ctx.semantic, prepared)
+
+
+def _apply_before_write(ctx: RuntimeContext, prepared: PreparedInput) -> PreparedInput:
+    return apply_prepared_hook(ctx.opts.hooks.before_write if ctx.opts.hooks else None, ctx.semantic, prepared)
+
+
+def _raise_for_quality_failure(raw_quality_status: str, opts: DatabricksIngestOptions) -> None:
+    if raw_quality_status == "FAILED" and opts.quality_action == "fail":
+        raise ValueError("Quality gates failed before Databricks write")
+
+
+def _acquire_lock(ctx: RuntimeContext) -> None:
+    if ctx.opts.lock_enabled and not ctx.opts.dry_run:
+        ctx.state.acquire_lock(target_table=ctx.target, run_id=ctx.run_id, owner=ctx.opts.lock_owner)
+
+
+def _release_lock(ctx: RuntimeContext) -> None:
+    if ctx.opts.lock_enabled and not ctx.opts.dry_run:
+        ctx.state.release_lock(target_table=ctx.target, run_id=ctx.run_id)
+
+
+def _validate_planning(contract: SemanticContract, opts: DatabricksIngestOptions) -> None:
+    runtime = dict(opts.runtime_metadata or {})
+    result = DatabricksAdapter.from_evidence(
+        target_table=target_full_name(contract),
+        runtime_type=str(runtime.get("runtime_type") or "serverless"),
+        spark_version=str(runtime["spark_version"]) if runtime.get("spark_version") else None,
+    ).plan(contract)
+    if result.status == "UNSUPPORTED" or (result.status == "REVIEW_REQUIRED" and not opts.allow_review_required):
+        blockers = "; ".join(blocker.message for blocker in result.blockers)
+        raise ValueError(f"Databricks planning status {result.status}: {blockers}")
+
+
+def _idempotency_skip(target: str, opts: DatabricksIngestOptions, query_one: QueryOne | None) -> dict[str, Any] | None:
+    if not opts.idempotency_key or opts.idempotency_policy not in {"skip_if_success", "rerun_if_failed", "fail_if_success"}:
+        return None
+    statement = render_find_idempotent_run_sql(
+        target_table=target,
+        idempotency_key=opts.idempotency_key,
+        status="SUCCESS",
+        runs_table=f"{opts.catalog}.{opts.schema}.ctrl_ingestion_runs",
+    )
+    previous = query_one(statement) if query_one else None
+    if not previous:
+        return None
+    if opts.idempotency_policy == "fail_if_success":
+        raise ValueError(f"idempotency_key={opts.idempotency_key!r} already succeeded")
+    return previous