contractforge-databricks 0.1.0__py3-none-any.whl

contractforge_databricks/runtime/streaming.py

@@ -0,0 +1,131 @@
+"""Runtime payload helpers for Databricks available-now streaming."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any
+
+from contractforge_core.config import CTRL_SCHEMA_VERSION, FRAMEWORK_VERSION
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.rendering.names import target_full_name
+
+
+def stream_metrics_from_batches(batch_results: list[dict[str, Any]]) -> dict[str, int]:
+    """Aggregate metrics returned by child batch ingestions."""
+
+    return {
+        "batches_processed": len(batch_results),
+        "total_rows_read": sum(_int_metric(result, "rows_read") for result in batch_results),
+        "total_rows_written": sum(_int_metric(result, "rows_written") for result in batch_results),
+        "total_rows_quarantined": sum(_int_metric(result, "rows_quarantined") for result in batch_results),
+    }
+
+
+def prefer_child_stream_metrics(local: dict[str, int], child: dict[str, int]) -> bool:
+    """Return true when persisted child-run metrics are more complete."""
+
+    if child.get("batches_processed", 0) <= 0:
+        return False
+    local_rows = local.get("total_rows_read", 0) + local.get("total_rows_written", 0) + local.get("total_rows_quarantined", 0)
+    child_rows = child.get("total_rows_read", 0) + child.get("total_rows_written", 0) + child.get("total_rows_quarantined", 0)
+    return local.get("batches_processed", 0) == 0 or child.get("batches_processed", 0) > local.get("batches_processed", 0) or child_rows > local_rows
+
+
+def stream_start_payload(
+    contract: SemanticContract,
+    *,
+    stream_run_id: str,
+    status: str = "RUNNING",
+    started_at_utc: datetime | None = None,
+    idempotency_key: str | None = None,
+    idempotency_policy: str = "always_run",
+    runtime_metadata: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    source = contract.source.raw or {}
+    operations = _operations_metadata(contract)
+    runtime = runtime_metadata or {}
+    return {
+        "stream_run_id": stream_run_id,
+        "idempotency_key": idempotency_key if idempotency_key is not None else operations.get("idempotency_key"),
+        "idempotency_policy": operations.get("idempotency_policy") or idempotency_policy,
+        "target_table": target_full_name(contract),
+        "target_catalog": _target_catalog(contract),
+        "target_layer": contract.target.layer,
+        "runtime_entrypoint": runtime.get("notebook_name") or operations.get("notebook_name"),
+        "source_type": source.get("type") or contract.source.kind,
+        "source_path": source.get("path") or source.get("url") or source.get("table") or contract.source.location,
+        "trigger": source.get("trigger") or "available_now",
+        "checkpoint_location": source.get("progress_location") or source.get("checkpoint_location"),
+        "status": status,
+        "started_at_utc": _timestamp(started_at_utc),
+        "batches_processed": 0,
+        "total_rows_read": 0,
+        "total_rows_written": 0,
+        "total_rows_quarantined": 0,
+        "framework_version": FRAMEWORK_VERSION,
+        "ctrl_schema_version": CTRL_SCHEMA_VERSION,
+        "master_job_id": operations.get("master_job_id"),
+        "master_run_id": operations.get("master_run_id"),
+        "parent_run_id": operations.get("parent_run_id"),
+        "run_group_id": operations.get("run_group_id"),
+        **runtime,
+    }
+
+
+def stream_result_payload(
+    contract: SemanticContract,
+    *,
+    stream_run_id: str,
+    status: str,
+    started_at_utc: datetime,
+    batch_results: list[dict[str, Any]],
+    stage_durations: dict[str, float] | None = None,
+    error_message: str | None = None,
+    skip_reason: str | None = None,
+    skipped_by_stream_run_id: str | None = None,
+    stream_metrics: dict[str, int] | None = None,
+    runtime_metadata: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    finished = datetime.now(timezone.utc)
+    metrics = stream_metrics or stream_metrics_from_batches(batch_results)
+    return {
+        **stream_start_payload(
+            contract,
+            stream_run_id=stream_run_id,
+            status=status,
+            started_at_utc=started_at_utc,
+            runtime_metadata=runtime_metadata,
+        ),
+        "ended_at_utc": _timestamp(finished),
+        "duration_seconds": (finished - started_at_utc).total_seconds(),
+        "batches_processed": metrics["batches_processed"],
+        "total_rows_read": metrics["total_rows_read"],
+        "total_rows_written": metrics["total_rows_written"],
+        "total_rows_quarantined": metrics["total_rows_quarantined"],
+        "batch_results": batch_results,
+        "stage_durations": stage_durations or {},
+        "error_message": error_message,
+        "skip_reason": skip_reason,
+        "skipped_by_stream_run_id": skipped_by_stream_run_id,
+    }
+
+
+def _int_metric(payload: dict[str, Any], key: str) -> int:
+    return int(payload.get(key) or 0)
+
+
+def _timestamp(value: datetime | None) -> str:
+    value = value or datetime.now(timezone.utc)
+    if value.tzinfo is None:
+        value = value.replace(tzinfo=timezone.utc)
+    return value.astimezone(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
+
+
+def _operations_metadata(contract: SemanticContract) -> dict[str, Any]:
+    return dict(contract.operations.metadata or {}) if contract.operations and contract.operations.metadata else {}
+
+
+def _target_catalog(contract: SemanticContract) -> str | None:
+    if not contract.target.namespace:
+        return None
+    return contract.target.namespace.split(".", 1)[0]

contractforge_databricks/runtime/success.py

@@ -0,0 +1,104 @@
+"""Successful Databricks runtime finalization."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.execution import ExecutionOutcome
+from contractforge_core.quality import QualityRuleResult
+from contractforge_core.runtime import PreparedInput, QueryOne
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.evidence import EvidenceWriter
+from contractforge_databricks.runtime.finalization import finalize_ingest
+from contractforge_databricks.runtime.lineage import write_runtime_diagnostics
+from contractforge_databricks.runtime.metrics import collect_write_metrics
+from contractforge_databricks.runtime.models import DatabricksIngestOptions
+from contractforge_databricks.runtime.utils import utc_now_str
+from contractforge_databricks.runtime.watermark import collect_previous_watermark, collect_watermark_candidate
+from contractforge_databricks.state import StateWriter
+
+
+def finalize_success(
+    *,
+    evidence: EvidenceWriter,
+    state: StateWriter,
+    contract: SemanticContract,
+    prepared: PreparedInput,
+    opts: DatabricksIngestOptions,
+    run_id: str,
+    target: str,
+    started: str,
+    outcome: ExecutionOutcome | None,
+    logical_rows_written: int,
+    quality_status_value: str,
+    schema_changes: dict[str, Any],
+    governance_results: dict[str, Any],
+    query_one: QueryOne | None,
+    quality_results: tuple[QualityRuleResult, ...] = (),
+    write_started_at: str | None = None,
+    write_finished_at: str | None = None,
+    stage_durations: dict[str, float] | None = None,
+) -> dict[str, Any]:
+    rows_written, operation_metrics = collect_write_metrics(
+        contract=contract,
+        target_table=target,
+        rows_written=logical_rows_written,
+        query_one=query_one,
+    )
+    watermark_column, watermark_current = collect_watermark_candidate(
+        contract=contract,
+        prepared=prepared,
+        query_one=query_one,
+    )
+    _, watermark_previous = collect_previous_watermark(
+        contract=contract,
+        query_one=query_one,
+        catalog=opts.catalog,
+        schema=opts.schema,
+    )
+    source_metadata = prepared.source_metadata or {}
+    watermark_previous = source_metadata.get("watermark_previous") or watermark_previous
+    diagnostics = write_runtime_diagnostics(
+        runner=evidence.runner,
+        contract=contract,
+        prepared=prepared,
+        run_id=run_id,
+        target=target,
+        status="SUCCESS",
+        started=started,
+        finished=_utc_now(),
+        rows_written=rows_written,
+        operation_metrics=operation_metrics,
+        catalog=opts.catalog,
+        schema=opts.schema,
+        query_one=query_one,
+        runtime_metadata=opts.runtime_metadata,
+    )
+    return finalize_ingest(
+        evidence,
+        state,
+        contract,
+        prepared,
+        opts,
+        run_id,
+        target,
+        "SUCCESS",
+        started,
+        rows_written=rows_written,
+        quality_status_value=quality_status_value,
+        quality_results=quality_results,
+        operation_metrics=operation_metrics,
+        schema_changes=schema_changes,
+        governance_results=governance_results,
+        write_started_at=write_started_at,
+        write_finished_at=write_finished_at,
+        stage_durations=stage_durations,
+        watermark_column=watermark_column,
+        watermark_previous=watermark_previous,
+        watermark_current=watermark_current,
+        diagnostics=diagnostics,
+    )
+
+
+def _utc_now() -> str:
+    return utc_now_str()

contractforge_databricks/runtime/utils.py

@@ -0,0 +1,52 @@
+"""Small Databricks runtime utility helpers without Spark import requirements."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Callable, Iterable
+
+
+def utc_now_ts() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def utc_now_str() -> str:
+    return utc_now_ts().strftime("%Y-%m-%d %H:%M:%S")
+
+
+def today_str() -> str:
+    return utc_now_ts().strftime("%Y-%m-%d")
+
+
+def new_run_id() -> str:
+    return str(uuid.uuid4())
+
+
+def resolve_run_id(run_id: str | None, run_id_factory: Callable[[], str] | None = None) -> str:
+    if run_id:
+        return run_id
+    if run_id_factory:
+        return str(run_id_factory())
+    return f"run-{uuid.uuid4()}"
+
+
+def safe_truncate(text: str | None, max_len: int = 100_000) -> str | None:
+    if text is None or len(text) <= max_len:
+        return text
+    return text[:max_len] + "\n...TRUNCATED..."
+
+
+def as_list(value: str | Iterable[Any] | None, sep: str = "|") -> list[str]:
+    if not value:
+        return []
+    if isinstance(value, str):
+        return [item.strip() for item in value.split(sep) if item.strip()]
+    return [str(item).strip() for item in value if str(item).strip()]
+
+
+def validate_columns(df: Any, columns: Iterable[str], context: str = "columns") -> None:
+    available = set(getattr(df, "columns", ()) or ())
+    missing = [column for column in columns if column not in available]
+    if missing:
+        raise ValueError(f"{context} not found: {missing}")

contractforge_databricks/runtime/watermark.py

@@ -0,0 +1,71 @@
+"""Runtime watermark collection for Databricks prepared views."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.runtime import PreparedInput, QueryOne
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.rendering.names import target_full_name
+from contractforge_databricks.state.queries import render_select_previous_watermark_sql
+from contractforge_databricks.watermark import render_select_watermark_candidate_sql
+
+
+def collect_watermark_candidate(
+    *,
+    contract: SemanticContract,
+    prepared: PreparedInput,
+    query_one: QueryOne | None,
+) -> tuple[str | None, str | None]:
+    columns = _watermark_columns(contract)
+    if not columns or query_one is None:
+        return (None, None)
+    row = query_one(
+        render_select_watermark_candidate_sql(
+            table_name=prepared.source_view,
+            columns=columns,
+            types=prepared.source_schema,
+        )
+    )
+    value = _row_value(row, "watermark_value")
+    return ("|".join(columns), None if value is None else str(value))
+
+
+def collect_previous_watermark(
+    *,
+    contract: SemanticContract,
+    query_one: QueryOne | None,
+    catalog: str = "main",
+    schema: str = "ops",
+) -> tuple[str | None, str | None]:
+    columns = _watermark_columns(contract)
+    if not columns or query_one is None:
+        return (None, None)
+    row = query_one(
+        render_select_previous_watermark_sql(
+            target_table=target_full_name(contract),
+            state_table=f"{catalog}.{schema}.ctrl_ingestion_state",
+        )
+    )
+    value = _row_value(row, "watermark_value")
+    return ("|".join(columns), None if value is None else str(value))
+
+
+def _watermark_columns(contract: SemanticContract) -> tuple[str, ...]:
+    metadata = contract.operations.metadata if contract.operations and contract.operations.metadata else {}
+    value = metadata.get("watermark_columns")
+    if isinstance(value, str):
+        return tuple(part.strip() for part in value.split(",") if part.strip())
+    if isinstance(value, (list, tuple)):
+        return tuple(str(part).strip() for part in value if str(part).strip())
+    return ()
+
+
+def _row_value(row: Any, key: str) -> Any:
+    if row is None:
+        return None
+    if isinstance(row, dict):
+        return row.get(key)
+    if hasattr(row, "asDict"):
+        return row.asDict().get(key)
+    return getattr(row, key, None)

contractforge_databricks/runtime/windows.py

@@ -0,0 +1,184 @@
+"""Databricks runtime orchestration for execution windows."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import Any
+from uuid import uuid4
+
+from contractforge_core.contracts import semantic_contract_from_mapping
+from contractforge_core.execution import ExecutionWindow, build_time_windows, summarize_window_results
+from contractforge_core.quality import QualityRuleResult
+from contractforge_core.runtime import QueryOne
+from contractforge_core.watermark import extract_watermark_field_value
+from contractforge_databricks.execution import build_child_window_plan
+from contractforge_databricks.rendering.names import target_full_name
+from contractforge_databricks.runtime.models import DatabricksIngestOptions
+from contractforge_databricks.runtime.orchestrator import ingest_databricks_contract
+from contractforge_databricks.runtime.sources import prepare_contract_source_view
+from contractforge_databricks.state.queries import render_select_previous_watermark_sql
+
+
+def has_windowed_execution(contract_mapping: dict[str, Any]) -> bool:
+    execution = contract_mapping.get("execution")
+    if not isinstance(execution, dict):
+        return False
+    catchup = execution.get("catchup")
+    return isinstance(execution.get("window"), dict) or (isinstance(catchup, dict) and bool(catchup.get("enabled")))
+
+
+def ingest_windowed_databricks_contract(
+    contract_mapping: dict[str, Any],
+    *,
+    spark: Any,
+    runner: Any,
+    options: DatabricksIngestOptions,
+    query_one: QueryOne | None = None,
+    quality_results: tuple[QualityRuleResult, ...] = (),
+    view_name: str | None = None,
+    collect_metrics: bool = False,
+) -> dict[str, Any]:
+    window_config = _window_config(contract_mapping, options=options, query_one=query_one)
+    windows = _windows(window_config)
+    parent_run_id = options.run_id or f"run-{uuid4()}"
+    results: list[dict[str, Any]] = []
+    for index, window in enumerate(windows, start=1):
+        child_plan = build_child_window_plan(
+            parent_run_id=parent_run_id,
+            column=str(window_config["column"]),
+            window=window,
+            index=index,
+            existing_filter=contract_mapping.get("filter_expression"),
+            base_idempotency_key=contract_mapping.get("idempotency_key"),
+        )
+        child_mapping = _child_contract_mapping(contract_mapping, child_plan)
+        child_contract = semantic_contract_from_mapping(child_mapping)
+        child_opts = replace(
+            options,
+            run_id=f"{parent_run_id}:window:{index:04d}",
+            idempotency_key=child_plan.idempotency_key or options.idempotency_key,
+        )
+        prepared = prepare_contract_source_view(
+            spark,
+            child_contract,
+            view_name=_child_view_name(child_contract, view_name, index),
+            collect_metrics=collect_metrics,
+            query_one=query_one,
+            evidence_catalog=child_opts.catalog,
+            evidence_schema=child_opts.schema,
+        )
+        result = ingest_databricks_contract(
+            child_contract,
+            runner=runner,
+            prepared=prepared,
+            options=child_opts,
+            query_one=query_one,
+            quality_results=quality_results,
+        )
+        result["execution_window"] = _window_payload(child_plan.window, str(window_config["column"]))
+        results.append(result)
+        if result.get("status") == "FAILED" and window_config.get("stop_on_failure", True):
+            break
+    return _summary(parent_run_id, windows, results)
+
+
+def _window_config(
+    contract_mapping: dict[str, Any],
+    *,
+    options: DatabricksIngestOptions,
+    query_one: QueryOne | None,
+) -> dict[str, Any]:
+    execution = dict(contract_mapping.get("execution") or {})
+    if isinstance(execution.get("window"), dict):
+        return dict(execution["window"])
+    catchup = dict(execution.get("catchup") or {})
+    if not catchup.get("enabled"):
+        raise ValueError("windowed execution requires execution.window or enabled execution.catchup")
+    start = catchup.get("start") or _previous_watermark_start(contract_mapping, options, query_one, catchup)
+    return {
+        "column": catchup.get("column") or _single_watermark_column(contract_mapping),
+        "start": start,
+        "end": catchup.get("end"),
+        "every": catchup.get("every"),
+        "stop_on_failure": catchup.get("stop_on_failure", True),
+    }
+
+
+def _windows(config: dict[str, Any]) -> tuple[ExecutionWindow, ...]:
+    explicit = config.get("windows")
+    if explicit:
+        return tuple(
+            ExecutionWindow(start=str(item["start"]), end=str(item["end"]), label=str(item.get("label") or ""))
+            for item in explicit
+        )
+    missing = [key for key in ("column", "start", "end", "every") if not config.get(key)]
+    if missing:
+        raise ValueError(f"execution window requires: {', '.join(missing)}")
+    return build_time_windows(str(config["start"]), str(config["end"]), str(config["every"]))
+
+
+def _child_contract_mapping(contract_mapping: dict[str, Any], child_plan: Any) -> dict[str, Any]:
+    child = dict(contract_mapping)
+    child["filter_expression"] = child_plan.filter_expression
+    child["parent_run_id"] = child_plan.parent_run_id
+    if child_plan.idempotency_key:
+        child["idempotency_key"] = child_plan.idempotency_key
+    runtime = dict(child.get("runtime_parameters") or {})
+    runtime.update(child_plan.runtime_parameters)
+    child["runtime_parameters"] = runtime
+    child.pop("execution", None)
+    return child
+
+
+def _summary(parent_run_id: str, windows: tuple[ExecutionWindow, ...], results: list[dict[str, Any]]) -> dict[str, Any]:
+    summary = dict(summarize_window_results(results))
+    summary.update(
+        {
+            "run_id": parent_run_id,
+            "parent_run_id": parent_run_id,
+            "windows_total": len(windows),
+            "windows_processed": len(results),
+            "window_results": results,
+        }
+    )
+    return summary
+
+
+def _previous_watermark_start(
+    contract_mapping: dict[str, Any],
+    options: DatabricksIngestOptions,
+    query_one: QueryOne | None,
+    catchup: dict[str, Any],
+) -> str:
+    if query_one is None:
+        raise ValueError("execution.catchup.start is required when query_one is not provided")
+    contract = semantic_contract_from_mapping(contract_mapping)
+    row = query_one(
+        render_select_previous_watermark_sql(
+            target_table=target_full_name(contract),
+            state_table=f"{options.catalog}.{options.schema}.ctrl_ingestion_state",
+        )
+    )
+    raw = row.get("watermark_value") if isinstance(row, dict) else None
+    start = extract_watermark_field_value(raw, catchup.get("column") or _single_watermark_column(contract_mapping))
+    if not start:
+        raise ValueError("execution.catchup.start is required when no previous watermark exists")
+    return start
+
+
+def _single_watermark_column(contract_mapping: dict[str, Any]) -> str:
+    value = contract_mapping.get("watermark_columns")
+    columns = [value] if isinstance(value, str) else list(value or ())
+    columns = [str(item).strip() for item in columns if str(item).strip()]
+    if len(columns) != 1:
+        raise ValueError("execution.catchup.column is required unless exactly one watermark column is configured")
+    return columns[0]
+
+
+def _child_view_name(contract: Any, view_name: str | None, index: int) -> str:
+    base = view_name or f"cf_source_{target_full_name(contract).replace('`', '').replace('.', '_')}"
+    return f"{base}_{index:04d}"
+
+
+def _window_payload(window: ExecutionWindow, column: str) -> dict[str, str]:
+    return {"label": window.label, "column": column, "start": window.start, "end": window.end}

contractforge_databricks/runtime/write.py

@@ -0,0 +1,66 @@
+"""Dispatch prepared Databricks views to write-mode executors."""
+
+from __future__ import annotations
+
+from contractforge_core.runtime import PreparedInput, QueryOne
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.execution import (
+    ExecutionOutcome,
+    SqlRunner,
+    execute_append,
+    execute_hash_diff_insert,
+    execute_overwrite,
+    execute_replace_partitions,
+    execute_scd1_merge,
+    execute_scd2_merge,
+    execute_snapshot_soft_delete,
+)
+from contractforge_databricks.write_modes.registry import execute_registered_write_mode
+
+
+def execute_prepared_write(
+    *,
+    runner: SqlRunner,
+    contract: SemanticContract,
+    prepared: PreparedInput,
+    target_partition_predicate: str | None = None,
+    replace_partition_predicate: str | None = None,
+    target_schema: dict[str, str] | None = None,
+    query_one: QueryOne | None = None,
+) -> ExecutionOutcome:
+    kwargs = {"runner": runner, "contract": contract, "source_view": prepared.source_view}
+    if contract.write.mode == "scd0_append":
+        return execute_append(**kwargs)
+    if contract.write.mode == "scd0_overwrite":
+        return execute_overwrite(**kwargs)
+    if contract.write.mode == "scd1_upsert":
+        if replace_partition_predicate:
+            return execute_replace_partitions(**kwargs, predicate=replace_partition_predicate)
+        return execute_scd1_merge(
+            **kwargs,
+            source_columns=prepared.source_columns,
+            target_partition_predicate=target_partition_predicate,
+        )
+    if contract.write.mode == "scd1_hash_diff":
+        return execute_hash_diff_insert(
+            **kwargs,
+            source_columns=prepared.source_columns,
+            target_schema=target_schema,
+            query_one=query_one,
+        )
+    if contract.write.mode == "scd2_historical":
+        return execute_scd2_merge(**kwargs, insert_columns=prepared.source_columns)
+    if contract.write.mode == "snapshot_soft_delete":
+        return execute_snapshot_soft_delete(**kwargs, source_columns=prepared.source_columns)
+    if contract.write.mode.startswith("custom:"):
+        return execute_registered_write_mode(
+            contract.write.mode,
+            runner=runner,
+            contract=contract,
+            prepared=prepared,
+            target_partition_predicate=target_partition_predicate,
+            replace_partition_predicate=replace_partition_predicate,
+            target_schema=target_schema,
+            query_one=query_one,
+        )
+    raise ValueError(f"Unsupported Databricks write mode: {contract.write.mode}")