databricks4py 0.2.0__py3-none-any.whl

databricks4py/migrations/__init__.py

@@ -0,0 +1,27 @@
+"""Migration framework for Delta Lake table evolution."""
+
+from databricks4py.migrations.alter import TableAlter
+from databricks4py.migrations.runner import MigrationRunner, MigrationRunResult, MigrationStep
+from databricks4py.migrations.schema_diff import (
+    ColumnChange,
+    SchemaDiff,
+    SchemaEvolutionError,
+)
+from databricks4py.migrations.validators import (
+    MigrationError,
+    TableValidator,
+    ValidationResult,
+)
+
+__all__ = [
+    "ColumnChange",
+    "MigrationError",
+    "MigrationRunResult",
+    "MigrationRunner",
+    "MigrationStep",
+    "SchemaDiff",
+    "SchemaEvolutionError",
+    "TableAlter",
+    "TableValidator",
+    "ValidationResult",
+]

databricks4py/migrations/alter.py

@@ -0,0 +1,114 @@
+"""Fluent DDL builder for Delta table schema changes."""
+
+from __future__ import annotations
+
+import logging
+
+from pyspark.sql import SparkSession
+
+from databricks4py.spark_session import active_fallback
+
+__all__ = ["TableAlter"]
+
+logger = logging.getLogger(__name__)
+
+
+def _sql_str(value: str) -> str:
+    """Escape single quotes for embedding in SQL string literals."""
+    return value.replace("'", "''")
+
+
+class TableAlter:
+    """Fluent interface for batching Delta table DDL operations.
+
+    Operations are queued and executed one-by-one via ``ALTER TABLE`` when
+    :meth:`apply` is called. Each operation is logged at INFO level.
+
+    Note:
+        ``rename_column`` and ``drop_column`` require Delta column mapping to be
+        enabled on the target table
+        (``delta.columnMapping.mode = 'name'``). Use :meth:`set_property` to
+        enable it before renaming or dropping.
+
+    Example::
+
+        TableAlter("catalog.schema.events", spark=spark) \\
+            .add_column("region", "STRING", comment="ISO-3166 region code") \\
+            .set_property("delta.enableChangeDataFeed", "true") \\
+            .apply()
+
+    Args:
+        table_name: Fully qualified table name.
+        spark: Optional SparkSession.
+    """
+
+    def __init__(self, table_name: str, *, spark: SparkSession | None = None) -> None:
+        self._table_name = table_name
+        self._spark = active_fallback(spark)
+        self._ops: list[str] = []
+
+    def add_column(
+        self,
+        name: str,
+        data_type: str,
+        *,
+        after: str | None = None,
+        nullable: bool = True,
+        comment: str | None = None,
+    ) -> TableAlter:
+        """Queue an ADD COLUMN operation.
+
+        Args:
+            name: Column name.
+            data_type: Spark SQL type string (e.g. ``"STRING"``, ``"DOUBLE"``).
+            after: If set, place the new column after this existing column.
+            nullable: Whether the column allows nulls (default True).
+            comment: Optional column comment.
+        """
+        not_null = "" if nullable else " NOT NULL"
+        comment_clause = f" COMMENT '{_sql_str(comment)}'" if comment else ""
+        after_clause = f" AFTER {after}" if after else ""
+        self._ops.append(f"ADD COLUMN ({name} {data_type}{not_null}{comment_clause}{after_clause})")
+        return self
+
+    def rename_column(self, old_name: str, new_name: str) -> TableAlter:
+        """Queue a RENAME COLUMN operation.
+
+        Requires ``delta.columnMapping.mode = 'name'`` on the target table.
+        """
+        self._ops.append(f"RENAME COLUMN {old_name} TO {new_name}")
+        return self
+
+    def drop_column(self, name: str) -> TableAlter:
+        """Queue a DROP COLUMN operation.
+
+        Requires ``delta.columnMapping.mode = 'name'`` on the target table.
+        """
+        self._ops.append(f"DROP COLUMN {name}")
+        return self
+
+    def set_property(self, key: str, value: str) -> TableAlter:
+        """Queue a SET TBLPROPERTIES operation.
+
+        Args:
+            key: Property key (e.g. ``"delta.enableChangeDataFeed"``).
+            value: Property value string.
+        """
+        self._ops.append(f"SET TBLPROPERTIES ('{_sql_str(key)}' = '{_sql_str(value)}')")
+        return self
+
+    def apply(self) -> None:
+        """Execute all queued DDL operations against the target table.
+
+        Each operation runs as a separate ``ALTER TABLE`` statement. The queue
+        is cleared after a successful apply. If any statement raises, the
+        remaining operations are not run and the queue is not cleared — a
+        subsequent ``apply()`` call will retry from the beginning.
+        """
+        if not self._ops:
+            return
+        for op in self._ops:
+            sql = f"ALTER TABLE {self._table_name} {op}"
+            logger.info("Executing: %s", sql)
+            self._spark.sql(sql)
+        self._ops.clear()

databricks4py/migrations/runner.py

@@ -0,0 +1,241 @@
+"""Ordered, idempotent migration runner for Delta Lake tables."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+
+from pyspark.sql import SparkSession
+from pyspark.sql.types import BooleanType, StringType, StructField, StructType, TimestampType
+
+from databricks4py.spark_session import active_fallback
+
+__all__ = ["MigrationRunResult", "MigrationRunner", "MigrationStep"]
+
+logger = logging.getLogger(__name__)
+
+_HISTORY_SCHEMA = StructType(
+    [
+        StructField("version", StringType(), False),
+        StructField("description", StringType(), True),
+        StructField("applied_at", TimestampType(), False),
+        StructField("success", BooleanType(), False),
+        StructField("error_message", StringType(), True),
+    ]
+)
+
+
+@dataclass(frozen=True)
+class MigrationStep:
+    """A single versioned migration step.
+
+    Steps are sorted by ``version`` before execution, so lexicographic ordering
+    determines the run order. Use a fixed-width prefix (``"V001"``, ``"V002"``, etc.)
+    to keep ordering stable as the number of steps grows.
+
+    Args:
+        version: Unique version string. Determines execution order.
+        description: Human-readable description stored in the history table.
+        up: Callable that receives a SparkSession and applies the migration.
+            May execute any Spark SQL, Delta, or Python logic.
+        pre_validate: Optional guard — called before ``up``. Return ``False``
+            to abort the step with a ``MigrationError``.
+        post_validate: Optional check — called after ``up``. Return ``False``
+            to mark the step as failed and halt the run.
+    """
+
+    version: str
+    description: str
+    up: Callable[[SparkSession], None]
+    pre_validate: Callable[[SparkSession], bool] | None = None
+    post_validate: Callable[[SparkSession], bool] | None = None
+
+
+@dataclass
+class MigrationRunResult:
+    """Summary of a :meth:`MigrationRunner.run` execution.
+
+    Attributes:
+        applied: Versions applied in this run, in execution order.
+        skipped: Versions already recorded as applied (idempotent skips).
+        failed: The version that caused a failure, if any.
+        dry_run: Whether this was a dry-run (no changes written).
+    """
+
+    applied: list[str] = field(default_factory=list)
+    skipped: list[str] = field(default_factory=list)
+    failed: str | None = None
+    dry_run: bool = False
+
+
+class MigrationRunner:
+    """Ordered, idempotent migration runner for Delta Lake.
+
+    Tracks applied versions in a Delta history table so each step runs exactly
+    once across all environments. Steps execute in lexicographic version order.
+
+    Example::
+
+        def add_audit_columns(spark: SparkSession) -> None:
+            spark.sql(
+                "ALTER TABLE catalog.schema.events "
+                "ADD COLUMNS (created_at TIMESTAMP, updated_at TIMESTAMP)"
+            )
+
+        runner = MigrationRunner(
+            history_table="catalog.schema._migration_history",
+        )
+        runner.register(
+            MigrationStep(
+                version="V001",
+                description="Add audit columns to events",
+                up=add_audit_columns,
+            )
+        )
+        result = runner.run()
+
+    Args:
+        history_table: Fully qualified Delta table name used to record applied steps.
+            Created automatically on first use.
+        spark: Optional SparkSession.
+    """
+
+    def __init__(
+        self,
+        history_table: str,
+        *,
+        spark: SparkSession | None = None,
+    ) -> None:
+        self._spark = active_fallback(spark)
+        self._history_table = history_table
+        self._steps: list[MigrationStep] = []
+        self._ensure_history_table()
+
+    def _ensure_history_table(self) -> None:
+        self._spark.sql(
+            f"""
+            CREATE TABLE IF NOT EXISTS {self._history_table} (
+                version     STRING    NOT NULL,
+                description STRING,
+                applied_at  TIMESTAMP NOT NULL,
+                success     BOOLEAN   NOT NULL,
+                error_message STRING
+            )
+            USING DELTA
+            """
+        )
+
+    def _applied_versions(self) -> set[str]:
+        rows = self._spark.sql(
+            f"SELECT version FROM {self._history_table} WHERE success = true"
+        ).collect()
+        return {row["version"] for row in rows}
+
+    def _record(
+        self,
+        step: MigrationStep,
+        *,
+        success: bool,
+        error_message: str | None = None,
+    ) -> None:
+        from pyspark.sql import Row
+
+        row = Row(
+            version=step.version,
+            description=step.description,
+            applied_at=datetime.now(tz=timezone.utc),
+            success=success,
+            error_message=error_message,
+        )
+        (
+            self._spark.createDataFrame([row], schema=_HISTORY_SCHEMA)
+            .write.format("delta")
+            .mode("append")
+            .saveAsTable(self._history_table)
+        )
+
+    def register(self, *steps: MigrationStep) -> MigrationRunner:
+        """Add one or more migration steps. Returns self for chaining.
+
+        Steps can be registered in any call order — execution order is always
+        determined by ``version``, not registration order.
+        """
+        self._steps.extend(steps)
+        return self
+
+    def pending(self) -> list[MigrationStep]:
+        """Return steps not yet applied, sorted by version."""
+        applied = self._applied_versions()
+        return sorted(
+            (s for s in self._steps if s.version not in applied),
+            key=lambda s: s.version,
+        )
+
+    def applied(self) -> list[str]:
+        """Return successfully applied versions in sorted order."""
+        return sorted(self._applied_versions())
+
+    def run(self, *, dry_run: bool = False) -> MigrationRunResult:
+        """Run all pending migration steps in version order.
+
+        Steps already in the history table are skipped (idempotent). Execution
+        halts on the first failure — the failed version is recorded with
+        ``success=False`` and returned in :attr:`MigrationRunResult.failed`.
+
+        Args:
+            dry_run: Log what would run without executing or writing history.
+
+        Returns:
+            MigrationRunResult summarising applied, skipped, and failed steps.
+
+        Raises:
+            MigrationError: If a step's pre-validation returns False.
+        """
+        from databricks4py.migrations.validators import MigrationError
+
+        result = MigrationRunResult(dry_run=dry_run)
+        applied_set = self._applied_versions()
+        all_sorted = sorted(self._steps, key=lambda s: s.version)
+
+        for step in all_sorted:
+            if step.version in applied_set:
+                result.skipped.append(step.version)
+                logger.debug("Skipping already-applied step %s", step.version)
+                continue
+
+            logger.info("Running migration %s: %s", step.version, step.description)
+
+            if step.pre_validate is not None and not step.pre_validate(self._spark):
+                raise MigrationError(
+                    step.version,
+                    [f"Pre-validation failed for step {step.version}: {step.description}"],
+                )
+
+            if dry_run:
+                logger.info("[dry-run] Would apply %s: %s", step.version, step.description)
+                result.applied.append(step.version)
+                continue
+
+            try:
+                step.up(self._spark)
+            except Exception as exc:
+                error_msg = str(exc)
+                logger.error("Migration %s failed: %s", step.version, error_msg)
+                self._record(step, success=False, error_message=error_msg)
+                result.failed = step.version
+                return result
+
+            if step.post_validate is not None and not step.post_validate(self._spark):
+                error_msg = f"Post-validation failed for step {step.version}"
+                logger.error(error_msg)
+                self._record(step, success=False, error_message=error_msg)
+                result.failed = step.version
+                return result
+
+            self._record(step, success=True)
+            result.applied.append(step.version)
+            logger.info("Applied migration %s", step.version)
+
+        return result

databricks4py/migrations/schema_diff.py

@@ -0,0 +1,136 @@
+"""Schema diff detection for Delta Lake table evolution."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.types import StructType
+
+from databricks4py.spark_session import active_fallback
+
+__all__ = ["ColumnChange", "SchemaDiff", "SchemaEvolutionError"]
+
+
+class SchemaEvolutionError(Exception):
+    """Raised on breaking schema changes."""
+
+
+@dataclass(frozen=True)
+class ColumnChange:
+    column: str
+    change_type: Literal["added", "removed", "type_changed", "nullable_changed"]
+    old_value: str | None = None
+    new_value: str | None = None
+    severity: Literal["info", "warning", "breaking"] = "info"
+
+
+class SchemaDiff:
+    """Compares two StructType schemas and reports column-level changes."""
+
+    def __init__(self, current: StructType, incoming: StructType) -> None:
+        self._current = current
+        self._incoming = incoming
+        self._changes: list[ColumnChange] | None = None
+
+    @classmethod
+    def from_tables(
+        cls,
+        table_name: str,
+        incoming_df: DataFrame,
+        *,
+        spark: SparkSession | None = None,
+    ) -> SchemaDiff:
+        """Create a diff between an existing table's schema and an incoming DataFrame.
+
+        Args:
+            table_name: Fully qualified table to read the current schema from.
+            incoming_df: DataFrame whose schema represents the proposed change.
+            spark: Optional SparkSession.
+        """
+        spark = active_fallback(spark)
+        current_schema = spark.read.table(table_name).schema
+        return cls(current=current_schema, incoming=incoming_df.schema)
+
+    def changes(self) -> list[ColumnChange]:
+        """Compute and return the list of column-level changes between schemas."""
+        if self._changes is not None:
+            return self._changes
+
+        result: list[ColumnChange] = []
+        current_fields = {f.name: f for f in self._current.fields}
+        incoming_fields = {f.name: f for f in self._incoming.fields}
+
+        for name in sorted(incoming_fields.keys() - current_fields.keys()):
+            field = incoming_fields[name]
+            result.append(
+                ColumnChange(
+                    column=name,
+                    change_type="added",
+                    new_value=str(field.dataType),
+                    severity="info",
+                )
+            )
+
+        for name in sorted(current_fields.keys() - incoming_fields.keys()):
+            field = current_fields[name]
+            result.append(
+                ColumnChange(
+                    column=name,
+                    change_type="removed",
+                    old_value=str(field.dataType),
+                    severity="breaking",
+                )
+            )
+
+        for name in sorted(current_fields.keys() & incoming_fields.keys()):
+            cur = current_fields[name]
+            inc = incoming_fields[name]
+
+            if cur.dataType != inc.dataType:
+                result.append(
+                    ColumnChange(
+                        column=name,
+                        change_type="type_changed",
+                        old_value=str(cur.dataType),
+                        new_value=str(inc.dataType),
+                        severity="breaking",
+                    )
+                )
+            elif cur.nullable != inc.nullable:
+                result.append(
+                    ColumnChange(
+                        column=name,
+                        change_type="nullable_changed",
+                        old_value=str(cur.nullable),
+                        new_value=str(inc.nullable),
+                        severity="warning",
+                    )
+                )
+
+        self._changes = result
+        return result
+
+    def has_breaking_changes(self) -> bool:
+        """True if any change has ``severity='breaking'`` (column removal or type change)."""
+        return any(c.severity == "breaking" for c in self.changes())
+
+    def summary(self) -> str:
+        """Return a human-readable table of all detected changes."""
+        changes = self.changes()
+        if not changes:
+            return "No schema changes detected."
+
+        lines = [f"{'Column':<30} {'Change':<20} {'Severity':<10} {'Details'}"]
+        lines.append("-" * 80)
+        for c in changes:
+            details = ""
+            if c.old_value and c.new_value:
+                details = f"{c.old_value} -> {c.new_value}"
+            elif c.new_value:
+                details = c.new_value
+            elif c.old_value:
+                details = c.old_value
+            lines.append(f"{c.column:<30} {c.change_type:<20} {c.severity:<10} {details}")
+        return "\n".join(lines)