contractforge-databricks 0.1.0__py3-none-any.whl

contractforge_databricks/sources/rds_iam_runtime.py

@@ -0,0 +1,191 @@
+"""Materialize RDS IAM JDBC tokens at adapter runtime.
+
+The core's ``rds_iam_review_options`` only writes a placeholder password
+(``{{rds_iam_token}}``) plus three ``contractforge.rdsIamHost / Port /
+Region`` metadata options. At runtime the Databricks adapter is the
+layer that has the AWS credentials and the network reach to actually
+mint the IAM auth token, so the materialization belongs here.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import time
+from typing import Any, Mapping
+
+from contractforge_core.connectors import generate_rds_iam_auth_token
+
+RDS_IAM_TOKEN_PLACEHOLDER = "{{rds_iam_token}}"
+_HOST_OPTION = "contractforge.rdsIamHost"
+_PORT_OPTION = "contractforge.rdsIamPort"
+_REGION_OPTION = "contractforge.rdsIamRegion"
+_METADATA_OPTIONS = (_HOST_OPTION, _PORT_OPTION, _REGION_OPTION)
+_TOKEN_TTL_SECONDS = 14 * 60
+_TOKEN_CACHE: dict[str, tuple[float, str]] = {}
+
+
+def materialize_rds_iam_options(
+    options: Mapping[str, str],
+    *,
+    auth: Mapping[str, Any] | None,
+) -> dict[str, str]:
+    """Replace the RDS IAM placeholder with a freshly minted IAM auth token.
+
+    Returns a new options dict with the contractforge.rdsIam* metadata
+    keys removed. If the placeholder is not present, the options are
+    returned untouched so non-IAM JDBC sources are unaffected.
+    """
+
+    options_dict = dict(options)
+    if options_dict.get("password") != RDS_IAM_TOKEN_PLACEHOLDER:
+        for key in _METADATA_OPTIONS:
+            options_dict.pop(key, None)
+        return options_dict
+    host = options_dict.pop(_HOST_OPTION, None)
+    port = options_dict.pop(_PORT_OPTION, None)
+    region = options_dict.pop(_REGION_OPTION, None)
+    if not host or not port or not region:
+        raise ValueError(
+            "RDS IAM materialization requires contractforge.rdsIamHost,"
+            " contractforge.rdsIamPort and contractforge.rdsIamRegion"
+            " (produced by the core for source.auth.type='rds_iam')"
+        )
+    auth_dict = dict(auth or {})
+    username = str(options_dict.get("user") or auth_dict.get("username") or "")
+    if not username:
+        raise ValueError("RDS IAM auth requires a JDBC user or source.auth.username")
+    host_text = str(host)
+    port_int = int(port)
+    region_text = str(region)
+    access_key = auth_dict.get("access_key_id")
+    secret_key = auth_dict.get("secret_access_key")
+    session_token = auth_dict.get("session_token")
+    token = _cached_token(
+        host=host_text,
+        port=port_int,
+        region=region_text,
+        username=username,
+        access_key=str(access_key or ""),
+        secret_key=str(secret_key or ""),
+        session_token=str(session_token or ""),
+    )
+    if not token:
+        if access_key and secret_key:
+            token = generate_rds_iam_auth_token(
+                host=host_text,
+                port=port_int,
+                region=region_text,
+                username=username,
+                access_key=str(access_key),
+                secret_key=str(secret_key),
+                session_token=str(session_token) if session_token else None,
+            )
+        else:
+            token = _boto3_rds_iam_token(host=host_text, port=port_int, region=region_text, username=username)
+        _store_token(
+            host=host_text,
+            port=port_int,
+            region=region_text,
+            username=username,
+            access_key=str(access_key or ""),
+            secret_key=str(secret_key or ""),
+            session_token=str(session_token or ""),
+            token=token,
+        )
+    options_dict["password"] = token
+    return options_dict
+
+
+def _cache_key(
+    *,
+    host: str,
+    port: int,
+    region: str,
+    username: str,
+    access_key: str,
+    secret_key: str,
+    session_token: str,
+) -> str:
+    digest = hashlib.sha256()
+    for component in (host, str(port), region, username, access_key, secret_key, session_token):
+        digest.update(component.encode("utf-8"))
+        digest.update(b"\x1f")
+    return digest.hexdigest()
+
+
+def _cached_token(
+    *,
+    host: str,
+    port: int,
+    region: str,
+    username: str,
+    access_key: str,
+    secret_key: str,
+    session_token: str,
+) -> str | None:
+    key = _cache_key(
+        host=host,
+        port=port,
+        region=region,
+        username=username,
+        access_key=access_key,
+        secret_key=secret_key,
+        session_token=session_token,
+    )
+    cached = _TOKEN_CACHE.get(key)
+    if not cached:
+        return None
+    created_at, token = cached
+    if time.time() - created_at >= _TOKEN_TTL_SECONDS:
+        _TOKEN_CACHE.pop(key, None)
+        return None
+    return token
+
+
+def _store_token(
+    *,
+    host: str,
+    port: int,
+    region: str,
+    username: str,
+    access_key: str,
+    secret_key: str,
+    session_token: str,
+    token: str,
+) -> None:
+    key = _cache_key(
+        host=host,
+        port=port,
+        region=region,
+        username=username,
+        access_key=access_key,
+        secret_key=secret_key,
+        session_token=session_token,
+    )
+    _TOKEN_CACHE[key] = (time.time(), token)
+
+
+def _boto3_rds_iam_token(*, host: str, port: int, region: str, username: str) -> str:
+    try:
+        import boto3  # type: ignore[import-not-found]
+    except Exception as exc:  # pragma: no cover - depends on runtime image
+        raise ValueError(
+            "RDS IAM auth requires either source.auth.access_key_id/source.auth.secret_access_key "
+            "or Databricks cluster AWS credentials with boto3 available"
+        ) from exc
+    try:
+        client = boto3.Session(region_name=region).client("rds", region_name=region)
+        return str(
+            client.generate_db_auth_token(
+                DBHostname=host,
+                Port=port,
+                DBUsername=username,
+                Region=region,
+            )
+        )
+    except Exception as exc:
+        raise ValueError(
+            "RDS IAM auth could not generate a token from Databricks cluster AWS credentials; "
+            "configure an instance profile/default AWS credential chain or declare "
+            "source.auth.access_key_id/source.auth.secret_access_key as secret placeholders"
+        ) from exc

contractforge_databricks/sources/rest_api.py

@@ -0,0 +1,33 @@
+"""Databricks review planning for generic REST API connector contracts."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from contractforge_core.connectors import is_rest_api_connector as is_rest_api_connector, rest_api_descriptor
+
+
+def render_rest_api_review_plan(source: dict[str, Any]) -> str:
+    descriptor = rest_api_descriptor(source, redaction="***REDACTED***")
+    payload = {
+        "kind": "databricks_rest_api_review_plan",
+        **descriptor,
+        "recommended_databricks_targets": _recommended_targets(source, descriptor["pagination"]),
+        "notes": [
+            "Generic paginated REST API ingestion is not rendered as a custom Databricks Python API client by default.",
+            "For bounded public files, prefer portable http_file/http_json/http_csv source types.",
+            "For proprietary SaaS APIs, prefer native_passthrough so the adapter can target Lakeflow Connect or Databricks Connections.",
+            "If no native connector exists, execute a reviewed landing step outside the semantic core and ingest landed files with incremental_files.",
+        ],
+    }
+    return json.dumps(payload, indent=2, sort_keys=True) + "\n"
+
+
+def _recommended_targets(source: dict[str, Any], pagination: dict[str, Any]) -> list[str]:
+    if not pagination or pagination.get("type") in {None, "none"}:
+        return ["http_file_if_bounded", "databricks_connection"]
+    system = str(source.get("provider") or source.get("name") or "").lower()
+    if system in {"salesforce", "workday", "servicenow"}:
+        return ["lakeflow_connect", "native_passthrough"]
+    return ["native_passthrough", "land_to_object_storage_then_incremental_files"]

contractforge_databricks/sources/support.py

@@ -0,0 +1,50 @@
+"""Databricks source connector support declarations."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_databricks.sources.classification import (
+    DatabricksSourceClassification,
+    classify_databricks_source,
+)
+
+
+def databricks_source_support(source: dict[str, Any] | str) -> dict[str, Any]:
+    """Return Databricks support metadata for a source connector.
+
+    This is adapter-owned documentation data. It does not affect core planning
+    and does not make Databricks names portable.
+    """
+
+    return _entry(classify_databricks_source(source))
+
+
+def list_databricks_source_support() -> tuple[dict[str, Any], ...]:
+    sources = (
+        "table",
+        "sql",
+        "csv",
+        "s3",
+        "incremental_files",
+        "http_json",
+        "rest_api",
+        "jdbc",
+        "kafka_bounded",
+        "eventhubs_bounded",
+        "delta_share",
+        "native_passthrough",
+    )
+    return tuple(databricks_source_support(source) for source in sources)
+
+
+def _entry(classification: DatabricksSourceClassification) -> dict[str, Any]:
+    entry: dict[str, Any] = {
+        "adapter": "databricks",
+        "source_type": classification.source_type,
+        "status": classification.status,
+        "note": classification.note,
+    }
+    if classification.native_mapping:
+        entry["native_mapping"] = classification.native_mapping
+    return entry

contractforge_databricks/sources/table_refs.py

@@ -0,0 +1,65 @@
+"""Databricks resolution for core logical table references."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+from contractforge_core.connectors import catalog_source_query, source_logical_table_reference
+from contractforge_core.connectors import LogicalTableReference
+from contractforge_core.semantic import SemanticContract
+
+
+def databricks_table_ref_resolver(contract: SemanticContract):
+    """Return a resolver that maps logical refs to Databricks table names."""
+
+    def _resolve(ref: LogicalTableReference) -> str:
+        if ref.catalog and ref.schema:
+            return f"{ref.catalog}.{ref.schema}.{ref.table}"
+        if ref.schema:
+            catalog = _target_catalog(contract)
+            return f"{catalog}.{ref.schema}.{ref.table}" if catalog else f"{ref.schema}.{ref.table}"
+        namespace = _namespace_for_layer(contract, ref.layer)
+        return f"{namespace}.{ref.table}" if namespace else ref.table
+
+    return _resolve
+
+
+def resolve_databricks_source_table_refs(contract: SemanticContract) -> dict | None:
+    """Return source.raw with logical refs resolved to Databricks table names."""
+
+    if not contract.source.raw:
+        return None
+    source = dict(contract.source.raw)
+    resolver = databricks_table_ref_resolver(contract)
+    table_ref = source_logical_table_reference(source)
+    if table_ref is not None:
+        source["table"] = resolver(table_ref)
+        source.pop("ref", None)
+        source.pop("table_ref", None)
+    if source.get("query"):
+        source["query"] = catalog_source_query(source, table_ref_resolver=resolver)
+    return source
+
+
+def contract_with_databricks_source_refs(contract: SemanticContract) -> SemanticContract:
+    source = resolve_databricks_source_table_refs(contract)
+    if source is None:
+        return contract
+    return replace(contract, source=replace(contract.source, raw=source))
+
+
+def _namespace_for_layer(contract: SemanticContract, layer: str) -> str | None:
+    namespace = contract.target.namespace
+    if not namespace:
+        return None
+    current_layer = contract.target.layer
+    if namespace.endswith(f"_{current_layer}"):
+        return f"{namespace[: -(len(current_layer) + 1)]}_{layer}"
+    return namespace
+
+
+def _target_catalog(contract: SemanticContract) -> str | None:
+    namespace = contract.target.namespace
+    if not namespace or "." not in namespace:
+        return None
+    return namespace.split(".", 1)[0]

contractforge_databricks/sql/__init__.py

@@ -0,0 +1,4 @@
+from contractforge_databricks.sql.identifiers import quote_identifier, quote_table_name
+from contractforge_databricks.sql.literals import sql_int, sql_json, sql_literal, sql_string
+
+__all__ = ["quote_identifier", "quote_table_name", "sql_int", "sql_json", "sql_literal", "sql_string"]

contractforge_databricks/sql/identifiers.py

@@ -0,0 +1,17 @@
+"""Databricks SQL identifier helpers."""
+
+from __future__ import annotations
+
+
+def quote_identifier(identifier: str) -> str:
+    if not identifier or not identifier.strip():
+        raise ValueError("identifier must not be empty")
+    return f"`{identifier.replace('`', '``')}`"
+
+
+def quote_table_name(table_name: str) -> str:
+    parts = [part.strip() for part in table_name.split(".") if part.strip()]
+    if not parts:
+        raise ValueError("table name must not be empty")
+    return ".".join(quote_identifier(part) for part in parts)
+

contractforge_databricks/sql/literals.py

@@ -0,0 +1,36 @@
+"""Databricks SQL literal helpers."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+
+def sql_string(value: object) -> str:
+    if value is None:
+        return "NULL"
+    return "'" + str(value).replace("'", "''") + "'"
+
+
+def sql_literal(value: Any) -> str:
+    if value is None:
+        return "NULL"
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, int) and not isinstance(value, bool):
+        return str(value)
+    if isinstance(value, float):
+        return str(value)
+    return sql_string(value)
+
+
+def sql_int(value: int | None) -> str:
+    return "NULL" if value is None else str(int(value))
+
+
+def sql_json(value: Any) -> str:
+    try:
+        payload = json.dumps(value, default=str, ensure_ascii=False, sort_keys=True)
+    except Exception:
+        payload = json.dumps(str(value), ensure_ascii=False)
+    return sql_string(payload)

contractforge_databricks/state/__init__.py

@@ -0,0 +1,39 @@
+from contractforge_databricks.state.ddl import render_create_state_tables_sql
+from contractforge_databricks.state.migrations import (
+    control_table_additive_migrations,
+    render_control_table_migrations_sql,
+)
+from contractforge_databricks.state.queries import (
+    render_control_metadata_current_sql,
+    render_find_idempotent_run_sql,
+    render_find_idempotent_stream_sql,
+    render_has_successful_run_sql,
+    render_lock_status_sql,
+    render_record_control_metadata_sql,
+    render_select_previous_watermark_sql,
+)
+from contractforge_databricks.state.sql import (
+    render_acquire_lock_sql,
+    render_release_lock_sql,
+    render_upsert_state_sql,
+)
+from contractforge_databricks.state.tables import state_table_names
+from contractforge_databricks.state.writer import StateWriter
+
+__all__ = [
+    "StateWriter",
+    "control_table_additive_migrations",
+    "render_acquire_lock_sql",
+    "render_control_metadata_current_sql",
+    "render_control_table_migrations_sql",
+    "render_create_state_tables_sql",
+    "render_find_idempotent_run_sql",
+    "render_find_idempotent_stream_sql",
+    "render_has_successful_run_sql",
+    "render_lock_status_sql",
+    "render_record_control_metadata_sql",
+    "render_release_lock_sql",
+    "render_select_previous_watermark_sql",
+    "render_upsert_state_sql",
+    "state_table_names",
+]

contractforge_databricks/state/ddl.py

@@ -0,0 +1,24 @@
+"""Databricks Delta DDL for operational state tables."""
+
+from __future__ import annotations
+
+from contractforge_core.evidence import STATE_TABLE_SCHEMAS
+from contractforge_databricks.sql import quote_table_name
+from contractforge_databricks.state.tables import state_table_names
+
+
+def render_create_state_tables_sql(*, catalog: str = "main", schema: str = "ops") -> str:
+    names = state_table_names(catalog, schema)
+    statements = [f"CREATE SCHEMA IF NOT EXISTS {quote_table_name(f'{catalog}.{schema}')};"]
+    for name, table in names.items():
+        statements.append(
+            "\n".join(
+                [
+                    f"CREATE TABLE IF NOT EXISTS {quote_table_name(table)} (",
+                    f"  {STATE_TABLE_SCHEMAS[name]}",
+                    ")",
+                    "USING DELTA;",
+                ]
+            )
+        )
+    return "\n\n".join(statements) + "\n"

contractforge_databricks/state/migrations.py

@@ -0,0 +1,146 @@
+"""Additive Databricks control-table migration planning."""
+
+from __future__ import annotations
+
+from contractforge_databricks.evidence.tables import evidence_table_names
+from contractforge_databricks.sql import quote_identifier, quote_table_name
+from contractforge_databricks.state.tables import state_table_names
+
+
+EVIDENCE_ADDITIVE_COLUMNS: dict[str, dict[str, str]] = {
+    "runs": {
+        "idempotency_key": "STRING",
+        "idempotency_policy": "STRING",
+        "skip_reason": "STRING",
+        "skipped_by_run_id": "STRING",
+        "metrics_source": "STRING",
+        "framework_version": "STRING",
+        "ctrl_schema_version": "BIGINT",
+        "runtime_type": "STRING",
+        "engine_version": "STRING",
+        "python_version": "STRING",
+        "stage_durations_json": "STRING",
+        "contract_description": "STRING",
+        "contract_owner": "STRING",
+        "contract_domain": "STRING",
+        "contract_tags_json": "STRING",
+        "contract_sla": "STRING",
+        "runtime_parameters_json": "STRING",
+        "annotations_status": "STRING",
+        "annotations_result_json": "STRING",
+        "ownership_json": "STRING",
+        "operations_json": "STRING",
+        "source_type": "STRING",
+        "source_connector": "STRING",
+        "source_name": "STRING",
+        "source_provider": "STRING",
+        "source_format": "STRING",
+        "source_path": "STRING",
+        "source_options_json": "STRING",
+        "source_read_json": "STRING",
+        "source_request_json": "STRING",
+        "source_auth_json": "STRING",
+        "source_pagination_json": "STRING",
+        "source_response_json": "STRING",
+        "source_incremental_json": "STRING",
+        "source_limits_json": "STRING",
+        "source_capabilities_json": "STRING",
+        "source_metrics_json": "STRING",
+        "source_system": "STRING",
+        "write_engine_requested": "STRING",
+        "write_engine_selected": "STRING",
+        "write_engine_status": "STRING",
+        "write_engine_reason": "STRING",
+        "write_engine_fallback_policy": "STRING",
+        "write_started_at_utc": "TIMESTAMP",
+        "write_finished_at_utc": "TIMESTAMP",
+        "table_version_before": "STRING",
+        "table_version_after": "STRING",
+        "write_committed": "BOOLEAN",
+        "parent_run_id": "STRING",
+        "run_group_id": "STRING",
+        "master_job_id": "STRING",
+        "master_run_id": "STRING",
+        "rows_expired": "BIGINT",
+        "metrics_json": "STRING",
+    },
+    "errors": {"error_class": "STRING", "occurred_at_utc": "TIMESTAMP"},
+    "quality": {"severity": "STRING", "observed_value": "STRING", "message": "STRING"},
+    "quarantine": {"record_ref": "STRING", "reason": "STRING"},
+    "schema_changes": {"payload_json": "STRING", "changed_at_utc": "TIMESTAMP"},
+    "lineage": {"source_name": "STRING"},
+    "metadata": {
+        "run_id": "STRING",
+        "target_table": "STRING",
+        "source_metadata_json": "STRING",
+        "captured_at_utc": "TIMESTAMP",
+    },
+    "streams": {
+        "run_id": "STRING",
+        "batch_id": "STRING",
+        "batch_metrics_json": "STRING",
+        "captured_at_utc": "TIMESTAMP",
+    },
+    "annotations": {"previous_value": "STRING", "annotation_date": "DATE"},
+    "operations": {"ownership_json": "STRING"},
+    "access": {
+        "access_run_id": "STRING",
+        "action": "STRING",
+        "column_name": "STRING",
+        "function_name": "STRING",
+        "new_value": "STRING",
+        "mode": "STRING",
+        "drift_policy": "STRING",
+        "revoke_unmanaged": "BOOLEAN",
+        "access_date": "DATE",
+        "payload_json": "STRING",
+        "applied_at_utc": "TIMESTAMP",
+    },
+}
+
+STATE_ADDITIVE_COLUMNS: dict[str, dict[str, str]] = {
+    "state": {
+        "parent_run_id": "STRING",
+        "run_group_id": "STRING",
+        "master_job_id": "STRING",
+        "master_run_id": "STRING",
+        "last_table_version": "STRING",
+        "last_write_completed_at_utc": "TIMESTAMP",
+        "last_watermark_candidate": "STRING",
+    },
+    "locks": {
+        "owner": "STRING",
+        "ttl_minutes": "BIGINT",
+        "released_at_utc": "TIMESTAMP",
+    },
+}
+
+
+def control_table_additive_migrations(*, catalog: str = "main", schema: str = "ops") -> dict[str, dict[str, str]]:
+    tables = {**evidence_table_names(catalog, schema), **state_table_names(catalog, schema)}
+    migrations: dict[str, dict[str, str]] = {}
+    for key, columns in {**EVIDENCE_ADDITIVE_COLUMNS, **STATE_ADDITIVE_COLUMNS}.items():
+        migrations[tables[key]] = dict(columns)
+    return migrations
+
+
+def render_control_table_migrations_sql(*, catalog: str = "main", schema: str = "ops") -> str:
+    lines = [
+        "-- Databricks control-table additive migrations.",
+        "-- Review existing schemas before execution; apply only columns that are missing.",
+        "",
+    ]
+    for table, columns in control_table_additive_migrations(catalog=catalog, schema=schema).items():
+        rendered_columns = ",\n  ".join(
+            f"{quote_identifier(column)} {column_type}" for column, column_type in columns.items()
+        )
+        lines.extend(
+            [
+                f"-- {table}",
+                f"ALTER TABLE {quote_table_name(table)} ADD COLUMNS (",
+                f"  {rendered_columns}",
+                ");",
+                "",
+            ]
+        )
+    return "\n".join(lines).rstrip() + "\n"