contractforge-databricks 0.1.0__py3-none-any.whl

contractforge_databricks/runtime/sources.py

@@ -0,0 +1,208 @@
+"""Databricks runtime source resolution.
+
+The functions in this module intentionally receive ``spark`` as an argument so
+the adapter stays importable outside Databricks and PySpark remains optional.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.connectors import (
+    catalog_source_query,
+    catalog_source_table_or_path,
+    delta_share_options,
+    eventhubs_bounded_options,
+    file_reader_options,
+    file_source_format,
+    is_available_now_stream_source,
+    is_bounded_stream_source,
+    is_catalog_source,
+    is_delta_share_source,
+    is_file_source,
+    is_http_file_source,
+    is_kafka_stream_source,
+    is_rest_api_connector,
+    jdbc_common_options,
+    kafka_bounded_options,
+    stream_source_format,
+)
+from contractforge_core.runtime import PreparedInput
+from contractforge_core.runtime import QueryOne
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.preparation import apply_contract_preparation, apply_write_staging
+from contractforge_databricks.runtime.file_selection import selected_file_load_path
+from contractforge_databricks.runtime.http_file import resolve_http_file_dataframe
+from contractforge_databricks.runtime.rest_api import resolve_rest_api_dataframe
+from contractforge_databricks.runtime.source_schema import apply_declared_schema
+from contractforge_databricks.runtime.source_metadata import (
+    schema_types,
+    source_metadata,
+    source_metadata_with_watermark,
+    source_name,
+)
+from contractforge_databricks.runtime.source_registry import get_source_resolver
+from contractforge_databricks.runtime.watermark import collect_previous_watermark
+from contractforge_databricks.runtime.storage_auth import configure_object_storage_access
+from contractforge_databricks.security import resolve_databricks_secret_placeholders, validate_source_security
+from contractforge_databricks.sources.interpret import interpret_incremental_files_source, is_incremental_file_source
+from contractforge_databricks.sources.rds_iam_runtime import materialize_rds_iam_options
+from contractforge_databricks.sources.table_refs import (
+    contract_with_databricks_source_refs,
+    databricks_table_ref_resolver,
+)
+
+_JDBC_SOURCE_ALIASES = {"jdbc", "postgres", "mysql", "sqlserver", "oracle", "redshift", "db2", "mariadb", "snowflake_jdbc", "bigquery_jdbc"}
+
+
+def resolve_source_dataframe(spark: Any, source: dict[str, Any], *, contract: SemanticContract | None = None) -> Any:
+    """Resolve a core source contract into a Databricks DataFrame."""
+
+    validate_source_security(source)
+    source = resolve_databricks_secret_placeholders(source)
+    source_type = source.get("type")
+    custom_resolver = get_source_resolver(str(source.get("connector") or source_type or ""))
+    if custom_resolver is not None:
+        return custom_resolver.resolve(spark, source)
+    if is_catalog_source(source):
+        return _resolve_catalog_source(spark, source, contract=contract)
+    if is_incremental_file_source(source):
+        return _resolve_autoloader_source(spark, source)
+    if is_file_source(source):
+        options = file_reader_options(source)
+        path, options = configure_object_storage_access(spark, source, options)
+        source = {**source, "path": path} if path is not None else source
+        return _read_source_with_options(
+            spark.read,
+            file_source_format(source),
+            options,
+            selected_file_load_path(spark, source, options),
+            source,
+        )
+    if is_http_file_source(source):
+        return resolve_http_file_dataframe(spark, source)
+    if is_rest_api_connector(source):
+        return resolve_rest_api_dataframe(spark, source)
+    if source_type in {"jdbc", "connector"} or source.get("connector") in _JDBC_SOURCE_ALIASES:
+        jdbc_options = materialize_rds_iam_options(
+            jdbc_common_options(source),
+            auth=source.get("auth"),
+        )
+        return _read_with_options(spark.read, "jdbc", jdbc_options, None)
+    if is_bounded_stream_source(source) or is_available_now_stream_source(source):
+        options = kafka_bounded_options(source) if is_kafka_stream_source(source) else eventhubs_bounded_options(source)
+        source_format = stream_source_format(source)
+        reader = spark.readStream if is_available_now_stream_source(source) else spark.read
+        return _read_with_options(reader, source_format, options, None)
+    if is_delta_share_source(source):
+        return _read_with_options(spark.read, "deltaSharing", delta_share_options(source), None)
+    raise ValueError(f"source.type={source_type!r} cannot be resolved by the Databricks runtime source resolver")
+
+
+def prepare_source_view(
+    spark: Any,
+    source: dict[str, Any],
+    *,
+    view_name: str,
+    collect_metrics: bool = False,
+) -> PreparedInput:
+    """Resolve a source and register it as a temporary view for ingestion."""
+
+    df = resolve_source_dataframe(spark, source)
+    df.createOrReplaceTempView(view_name)
+    columns = tuple(str(column) for column in getattr(df, "columns", ()) or ())
+    rows_read = int(df.count()) if collect_metrics else 0
+    return PreparedInput(
+        source_view=view_name,
+        source_columns=columns,
+        source_schema=schema_types(df),
+        rows_read=rows_read,
+        source_name=source_name(source),
+        source_metadata=source_metadata(source),
+    )
+
+
+def prepare_contract_source_view(
+    spark: Any,
+    contract: SemanticContract,
+    *,
+    view_name: str,
+    collect_metrics: bool = False,
+    query_one: QueryOne | None = None,
+    evidence_catalog: str = "main",
+    evidence_schema: str = "ops",
+) -> PreparedInput:
+    """Resolve, prepare and register the contract source as a temporary view."""
+
+    if not contract.source.raw:
+        raise ValueError("prepare_contract_source_view requires a structured source contract")
+    runtime_contract = contract_with_databricks_source_refs(contract)
+    df = resolve_source_dataframe(spark, runtime_contract.source.raw or {}, contract=runtime_contract)
+    watermark_column, watermark_previous = collect_previous_watermark(
+        contract=contract,
+        query_one=query_one,
+        catalog=evidence_catalog,
+        schema=evidence_schema,
+    )
+    df = apply_contract_preparation(
+        df,
+        contract,
+        watermark_column=watermark_column,
+        watermark_previous=watermark_previous,
+    )
+    df = apply_write_staging(df, contract)
+    df.createOrReplaceTempView(view_name)
+    columns = tuple(str(column) for column in getattr(df, "columns", ()) or ())
+    rows_read = int(df.count()) if collect_metrics else 0
+    return PreparedInput(
+        source_view=view_name,
+        source_columns=columns,
+        source_schema=schema_types(df),
+        rows_read=rows_read,
+        source_name=runtime_contract.source.name,
+        source_metadata=source_metadata_with_watermark(runtime_contract.source.raw or {}, watermark_previous),
+    )
+
+
+def _resolve_catalog_source(spark: Any, source: dict[str, Any], *, contract: SemanticContract | None = None) -> Any:
+    resolver = databricks_table_ref_resolver(contract) if contract is not None else None
+    if source.get("type") == "sql" or source.get("connector") == "sql":
+        return spark.sql(catalog_source_query(source, table_ref_resolver=resolver))
+    table_or_path = catalog_source_table_or_path(source, table_ref_resolver=resolver)
+    if source.get("path") and not source.get("table"):
+        source_type = str(source.get("type") or "delta")
+        source_format = "delta" if source_type == "delta_table" else source_type.replace("_table", "")
+        return _read_with_options(spark.read, source_format, {}, table_or_path)
+    return spark.table(str(table_or_path))
+
+
+def _resolve_autoloader_source(spark: Any, source: dict[str, Any]) -> Any:
+    interpreted = interpret_incremental_files_source(source)
+    options = {"cloudFiles.format": str(interpreted.get("format") or "json")}
+    options.update({str(key): str(value) for key, value in interpreted.get("options", {}).items()})
+    if interpreted.get("schema_tracking_location"):
+        options["cloudFiles.schemaLocation"] = str(interpreted["schema_tracking_location"])
+    if interpreted.get("schema_hints"):
+        options["cloudFiles.schemaHints"] = str(interpreted["schema_hints"])
+    return _read_with_options(spark.readStream, "cloudFiles", options, interpreted.get("path"))
+
+
+def _read_with_options(reader: Any, source_format: str, options: dict[str, str], path: object | None) -> Any:
+    builder = reader.format(source_format)
+    for key, value in sorted(options.items()):
+        builder = builder.option(key, value)
+    return builder.load(path if isinstance(path, list) else str(path)) if path is not None else builder.load()
+
+
+def _read_source_with_options(
+    reader: Any,
+    source_format: str,
+    options: dict[str, str],
+    path: object | None,
+    source: dict[str, Any],
+) -> Any:
+    builder = reader.format(source_format)
+    for key, value in sorted(options.items()):
+        builder = builder.option(key, value)
+    builder = apply_declared_schema(builder, source)
+    return builder.load(path if isinstance(path, list) else str(path)) if path is not None else builder.load()

contractforge_databricks/runtime/spark.py

@@ -0,0 +1,183 @@
+"""Databricks/Spark runtime convenience helpers with lazy imports."""
+
+from __future__ import annotations
+
+import json
+import logging
+import platform
+from typing import Any
+
+from contractforge_databricks.capabilities.runtime import is_serverless_conf
+from contractforge_databricks.runtime.detection import _collect_spark_conf
+from contractforge_databricks.sql import quote_identifier, quote_table_name
+
+logger = logging.getLogger(__name__)
+_SERVERLESS_CACHE: dict[int, bool] = {}
+
+
+def get_active_spark() -> Any:
+    """Resolve the active Databricks or PySpark session at call time."""
+    try:
+        from databricks.sdk.runtime import spark as dbx_spark  # type: ignore
+
+        if dbx_spark is not None:
+            return dbx_spark
+    except Exception as exc:
+        logger.debug("Databricks runtime spark session was not available; falling back to PySpark.", exc_info=exc)
+    try:
+        from pyspark.sql import SparkSession
+    except Exception as exc:
+        raise RuntimeError("PySpark is required to resolve an active Spark session.") from exc
+    session = SparkSession.getActiveSession() or getattr(SparkSession, "_instantiatedSession", None)
+    if session is None:
+        raise RuntimeError("No active SparkSession was found.")
+    return session
+
+
+def runtime_info(spark: Any | None = None) -> dict[str, str | None]:
+    session = spark or _maybe_active_spark()
+    version = getattr(session, "version", None) if session is not None else None
+    return {
+        "runtime_type": "serverless" if detect_serverless(session) else "classic",
+        "spark_version": version,
+        "python_version": platform.python_version(),
+    }
+
+
+def detect_serverless(spark: Any | None = None) -> bool:
+    session = spark or _maybe_active_spark()
+    if session is None:
+        return False
+    cache_key = id(session)
+    if cache_key in _SERVERLESS_CACHE:
+        return _SERVERLESS_CACHE[cache_key]
+    conf = _collect_spark_conf(session) if session is not None else {}
+    detected = is_serverless_conf(conf) or conf.get("spark.databricks.clusterUsageTags.clusterSource") == "JOB_SERVERLESS"
+    _SERVERLESS_CACHE[cache_key] = detected
+    return detected
+
+
+def safe_cache(df: Any, *, enabled: bool = True, serverless: bool | None = None) -> Any:
+    if not enabled or (detect_serverless() if serverless is None else serverless):
+        return df
+    try:
+        return df.cache()
+    except Exception as exc:
+        if _is_unsupported_cache_error(exc):
+            return df
+        raise
+
+
+def safe_unpersist(df: Any, *, enabled: bool = True, serverless: bool | None = None) -> None:
+    if not enabled or (detect_serverless() if serverless is None else serverless):
+        return
+    try:
+        df.unpersist()
+    except Exception as exc:
+        if not _is_unsupported_cache_error(exc):
+            raise
+
+
+def safe_cache_table(spark: Any, table_name: str, *, enabled: bool = True, serverless: bool | None = None) -> bool:
+    if not enabled or (detect_serverless(spark) if serverless is None else serverless):
+        return False
+    try:
+        catalog = getattr(spark, "catalog", None)
+        cache_table = getattr(catalog, "cacheTable", None)
+        if callable(cache_table):
+            cache_table(table_name)
+        else:
+            spark.sql(f"CACHE TABLE {quote_identifier(table_name)}")
+        return True
+    except Exception as exc:
+        if _is_unsupported_cache_error(exc):
+            return False
+        raise
+
+
+def safe_uncache_table(spark: Any, table_name: str, *, enabled: bool = True, serverless: bool | None = None) -> None:
+    if not enabled or (detect_serverless(spark) if serverless is None else serverless):
+        return
+    try:
+        catalog = getattr(spark, "catalog", None)
+        uncache_table = getattr(catalog, "uncacheTable", None)
+        if callable(uncache_table):
+            uncache_table(table_name)
+        else:
+            spark.sql(f"UNCACHE TABLE {quote_identifier(table_name)}")
+    except Exception as exc:
+        if not _is_unsupported_cache_error(exc):
+            raise
+
+
+def table_exists(full_name: str, *, spark: Any | None = None) -> bool:
+    session = spark or get_active_spark()
+    try:
+        if session.catalog.tableExists(full_name):
+            return True
+    except Exception as exc:
+        logger.debug("Spark catalog tableExists failed for %s; falling back to DESCRIBE TABLE.", full_name, exc_info=exc)
+    try:
+        session.sql(f"DESCRIBE TABLE {quote_table_name(full_name)}")
+        return True
+    except Exception as exc:
+        logger.debug("Spark DESCRIBE TABLE failed for %s.", full_name, exc_info=exc)
+        return False
+
+
+def schema_signature(df: Any) -> str:
+    return json.dumps(
+        [(field.name, field.dataType.simpleString(), field.nullable) for field in df.schema.fields],
+        ensure_ascii=False,
+    )
+
+
+def fix_encoding(df: Any, *, enabled: bool, encoding: str, columns: tuple[str, ...] = ()) -> Any:
+    if not enabled:
+        return df
+    from pyspark.sql import functions as functions  # type: ignore
+
+    string_cols = [field.name for field in df.schema.fields if field.dataType.typeName() == "string"]
+    cols_to_fix = columns or tuple(string_cols)
+    for column in cols_to_fix:
+        if column in string_cols:
+            df = df.withColumn(column, functions.decode(functions.col(column).cast("binary"), encoding))
+    return df
+
+
+def sync_delta_schema(
+    *,
+    df: Any,
+    target_table: str,
+    schema_changes: dict[str, Any],
+    policy: str,
+    spark: Any | None = None,
+) -> None:
+    session = spark or get_active_spark()
+    if policy not in {"permissive", "additive_only"} or not table_exists(target_table, spark=session):
+        return
+    fields = {field.name: field.dataType.simpleString() for field in df.schema.fields}
+    added = [column for column in schema_changes.get("added_columns", ()) if column in fields]
+    if added:
+        cols_sql = ", ".join(f"{quote_identifier(column)} {fields[column]}" for column in added)
+        session.sql(f"ALTER TABLE {quote_table_name(target_table)} ADD COLUMNS ({cols_sql})")
+    for change in schema_changes.get("type_changes", ()):
+        if not change.get("allowed"):
+            continue
+        column = str(change["column"])
+        source_type = str(change["source"])
+        session.sql(f"ALTER TABLE {quote_table_name(target_table)} ALTER COLUMN {quote_identifier(column)} TYPE {source_type}")
+        change["applied"] = True
+
+
+def _maybe_active_spark() -> Any | None:
+    try:
+        return get_active_spark()
+    except Exception as exc:
+        logger.debug("No active Spark session could be resolved.", exc_info=exc)
+        return None
+
+
+def _is_unsupported_cache_error(exc: Exception) -> bool:
+    text = str(exc).upper()
+    return "NOT_SUPPORTED" in text or "SERVERLESS" in text

contractforge_databricks/runtime/spark_defaults.py

@@ -0,0 +1,35 @@
+"""Spark-backed runtime defaults for Databricks bundle execution."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import Any
+
+from contractforge_core.runtime import QueryOne
+from contractforge_databricks.runtime.models import DatabricksIngestOptions
+
+
+def with_spark_runtime_defaults(spark: Any, opts: DatabricksIngestOptions, target: str) -> DatabricksIngestOptions:
+    if opts.target_schema is not None or not opts.ensure_table:
+        return opts
+    target_schema = spark_target_schema(spark, target)
+    return replace(opts, target_schema=target_schema) if target_schema is not None else opts
+
+
+def spark_target_schema(spark: Any, target: str) -> dict[str, str] | None:
+    try:
+        schema = spark.table(target).schema
+    except Exception:
+        return None
+    return {str(field.name): str(field.dataType.simpleString()).lower() for field in schema.fields}
+
+
+def spark_query_one(spark: Any) -> QueryOne | None:
+    if not callable(getattr(spark, "sql", None)):
+        return None
+
+    def query_one(statement: str) -> dict[str, Any] | None:
+        rows = spark.sql(statement).limit(1).collect()
+        return rows[0].asDict() if rows else None
+
+    return query_one

contractforge_databricks/runtime/storage_auth.py

@@ -0,0 +1,132 @@
+"""Databricks runtime object-storage credential helpers."""
+
+from __future__ import annotations
+
+import urllib.parse
+from typing import Any
+
+from contractforge_core.connectors import object_storage_provider
+
+
+def configure_object_storage_access(
+    spark: Any,
+    source: dict[str, Any],
+    options: dict[str, str],
+) -> tuple[object | None, dict[str, str]]:
+    """Configure adapter-owned Spark storage credentials and return read path/options."""
+
+    provider = object_storage_provider(source)
+    path = source.get("path")
+    if provider == "s3":
+        return path, _configure_s3(spark, source, options)
+    if provider == "azure_blob":
+        return _configure_azure_blob(spark, source, path), options
+    return path, options
+
+
+def _configure_s3(spark: Any, source: dict[str, Any], options: dict[str, str]) -> dict[str, str]:
+    reader_options: dict[str, str] = {}
+    for key, value in options.items():
+        if key.startswith("fs.s3a.") or key.startswith("spark.hadoop.fs.s3a."):
+            _set_conf(spark, key, value)
+        else:
+            reader_options[key] = value
+    auth = _dict(source.get("auth"))
+    access_key = auth.get("access_key_id") or auth.get("access_key") or auth.get("aws_access_key_id")
+    secret_key = auth.get("secret_access_key") or auth.get("secret_key") or auth.get("aws_secret_access_key")
+    session_token = auth.get("session_token") or auth.get("token") or auth.get("aws_session_token")
+    if bool(access_key) != bool(secret_key):
+        raise ValueError("source.auth for connector=s3 requires access_key_id and secret_access_key together")
+    if access_key and secret_key:
+        _set_conf(spark, "fs.s3a.access.key", str(access_key))
+        _set_conf(spark, "fs.s3a.secret.key", str(secret_key))
+        if session_token:
+            _set_conf(spark, "fs.s3a.session.token", str(session_token))
+            _set_conf(spark, "fs.s3a.aws.credentials.provider", "org.apache.hadoop.fs.s3a.TemporaryAWSCredentialsProvider")
+        else:
+            _set_conf(spark, "fs.s3a.aws.credentials.provider", "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider")
+    return reader_options
+
+
+def _configure_azure_blob(spark: Any, source: dict[str, Any], path: object | None) -> object | None:
+    auth = _dict(source.get("auth"))
+    sas_token = auth.get("sas_token") or auth.get("token")
+    if not sas_token:
+        return path
+    account_url = str(source.get("account_url") or "").strip()
+    container = str(source.get("container") or "").strip()
+    if account_url or container:
+        account = _azure_account_from_url(account_url)
+        if not account:
+            raise ValueError("source.account_url is required for connector=azure_blob when source.container is used")
+        if not container:
+            raise ValueError("source.container is required for connector=azure_blob when source.account_url is used")
+        _configure_azure_blob_sas(spark, account, container, str(sas_token))
+        if path and "://" not in str(path):
+            return f"wasbs://{container}@{account}.blob.core.windows.net/{str(path).lstrip('/')}"
+        return path
+    account, inferred_container = _azure_account_container_from_uri(str(path or ""))
+    if not account or not inferred_container:
+        raise ValueError(
+            "auth.sas_token in connector=azure_blob requires source.account_url/source.container "
+            "or path wasbs://container@account.blob.core.windows.net/..."
+        )
+    _configure_azure_blob_sas(spark, account, inferred_container, str(sas_token))
+    return path
+
+
+def _configure_azure_blob_sas(spark: Any, account: str, container: str, sas_token: str) -> None:
+    token = sas_token.strip()
+    if token.startswith("?"):
+        token = token[1:]
+    if not token:
+        raise ValueError("auth.sas_token cannot be empty for connector=azure_blob")
+    _set_conf(spark, f"fs.azure.sas.{container}.{account}.blob.core.windows.net", token)
+
+
+def _azure_account_from_url(account_url: str) -> str:
+    if not account_url:
+        return ""
+    parsed = urllib.parse.urlparse(account_url if "://" in account_url else f"https://{account_url}")
+    host = parsed.netloc or parsed.path
+    return host.split(".", 1)[0].strip()
+
+
+def _azure_account_container_from_uri(path: str) -> tuple[str, str]:
+    parsed = urllib.parse.urlparse(path)
+    if parsed.scheme not in {"wasbs", "wasb", "abfss", "abfs"} or "@" not in parsed.netloc:
+        return "", ""
+    container, host = parsed.netloc.split("@", 1)
+    return host.split(".", 1)[0].strip(), container.strip()
+
+
+def _set_conf(spark: Any, key: str, value: str) -> None:
+    conf = getattr(spark, "conf", None)
+    if conf is None or not hasattr(conf, "set"):
+        raise RuntimeError("Object-storage source auth requires a Spark session with spark.conf.set")
+    try:
+        conf.set(key, value)
+    except Exception as exc:
+        if _is_spark_config_blocked(exc):
+            if key.startswith("fs.azure.sas."):
+                raise RuntimeError(
+                    "Databricks serverless/Spark Connect blocked Spark SAS configuration. "
+                    "Use a Unity Catalog External Location or Volume, or configure direct SAS only in a runtime "
+                    "where Hadoop config fs.azure.sas.* is allowed."
+                ) from exc
+            if key.startswith("fs.s3a.") or key.startswith("spark.hadoop.fs.s3a."):
+                raise RuntimeError(
+                    "Databricks serverless/Spark Connect blocked Spark S3 credential configuration. "
+                    "Use a Unity Catalog External Location or Volume, or configure source.auth for S3 only in a "
+                    "runtime where Hadoop config fs.s3a.* is allowed."
+                ) from exc
+        raise
+
+
+def _is_spark_config_blocked(exc: Exception) -> bool:
+    message = str(exc)
+    return "CONFIG_NOT_AVAILABLE" in message or "Configuration fs.azure.sas" in message
+
+
+def _dict(value: object) -> dict[str, Any]:
+    return dict(value) if isinstance(value, dict) else {}