databricks4py 0.2.0__py3-none-any.whl

databricks4py/io/merge.py

@@ -0,0 +1,176 @@
+"""Fluent MERGE INTO builder for Delta Lake tables."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from pyspark.sql import DataFrame, SparkSession
+
+from databricks4py.spark_session import active_fallback
+
+if TYPE_CHECKING:
+    from databricks4py.metrics.base import MetricsSink
+
+__all__ = ["MergeBuilder", "MergeResult"]
+
+logger = logging.getLogger(__name__)
+
+_SOURCE_ALIAS = "source"
+_TARGET_ALIAS = "target"
+
+
+@dataclass(frozen=True)
+class MergeResult:
+    """Outcome metrics from a MERGE operation."""
+
+    rows_inserted: int
+    rows_updated: int
+    rows_deleted: int
+
+
+class MergeBuilder:
+    """Fluent builder for Delta Lake MERGE INTO operations.
+
+    Example::
+
+        result = (
+            MergeBuilder("catalog.schema.target", source_df, spark)
+            .on("id")
+            .when_matched_update()
+            .when_not_matched_insert()
+            .execute()
+        )
+    """
+
+    def __init__(
+        self,
+        target_table_name: str,
+        source: DataFrame,
+        spark: SparkSession | None = None,
+        *,
+        metrics_sink: MetricsSink | None = None,
+    ) -> None:
+        self._target_table_name = target_table_name
+        self._source = source
+        self._spark = active_fallback(spark)
+        self._metrics_sink = metrics_sink
+
+        self._join_keys: list[str] = []
+        self._join_condition: str | None = None
+        self._actions: list[dict[str, Any]] = []
+
+    def on(self, *keys: str) -> MergeBuilder:
+        """Set merge join keys (ANDed equality conditions)."""
+        self._join_keys = list(keys)
+        return self
+
+    def on_condition(self, condition: str) -> MergeBuilder:
+        """Set a custom merge condition expression instead of key-based equality."""
+        self._join_condition = condition
+        return self
+
+    def when_matched_update(self, columns: list[str] | None = None) -> MergeBuilder:
+        """Update matched rows. If columns is None, updates all columns."""
+        self._actions.append({"type": "matched_update", "columns": columns})
+        return self
+
+    def when_matched_delete(self, condition: str | None = None) -> MergeBuilder:
+        """Delete matched rows, optionally filtered by condition."""
+        self._actions.append({"type": "matched_delete", "condition": condition})
+        return self
+
+    def when_not_matched_insert(self, columns: list[str] | None = None) -> MergeBuilder:
+        """Insert non-matched source rows. If columns is None, inserts all."""
+        self._actions.append({"type": "not_matched_insert", "columns": columns})
+        return self
+
+    def when_not_matched_by_source_delete(self, condition: str | None = None) -> MergeBuilder:
+        """Delete target rows not present in source."""
+        self._actions.append({"type": "not_matched_by_source_delete", "condition": condition})
+        return self
+
+    def _build_condition(self) -> str:
+        if self._join_condition:
+            return self._join_condition
+        parts = [f"{_TARGET_ALIAS}.{k} = {_SOURCE_ALIAS}.{k}" for k in self._join_keys]
+        return " AND ".join(parts)
+
+    def execute(self) -> MergeResult:
+        """Execute the merge and return metrics."""
+        from delta.tables import DeltaTable
+
+        target_dt = DeltaTable.forName(self._spark, self._target_table_name)
+        condition = self._build_condition()
+
+        merger = target_dt.alias(_TARGET_ALIAS).merge(self._source.alias(_SOURCE_ALIAS), condition)
+
+        for action in self._actions:
+            merger = self._apply_action(merger, action)
+
+        merger.execute()
+
+        result = self._read_metrics()
+        if self._metrics_sink:
+            self._emit_metrics(result)
+        return result
+
+    def _apply_action(self, merger: Any, action: dict[str, Any]) -> Any:
+        action_type = action["type"]
+
+        if action_type == "matched_update":
+            columns = action["columns"]
+            if columns:
+                update_map = {col: f"{_SOURCE_ALIAS}.{col}" for col in columns}
+                return merger.whenMatchedUpdate(set=update_map)
+            return merger.whenMatchedUpdateAll()
+
+        if action_type == "matched_delete":
+            cond = action.get("condition")
+            return merger.whenMatchedDelete(condition=cond) if cond else merger.whenMatchedDelete()
+
+        if action_type == "not_matched_insert":
+            columns = action["columns"]
+            if columns:
+                insert_map = {col: f"{_SOURCE_ALIAS}.{col}" for col in columns}
+                return merger.whenNotMatchedInsert(values=insert_map)
+            return merger.whenNotMatchedInsertAll()
+
+        if action_type == "not_matched_by_source_delete":
+            cond = action.get("condition")
+            if cond:
+                return merger.whenNotMatchedBySourceDelete(condition=cond)
+            return merger.whenNotMatchedBySourceDelete()
+
+        msg = f"Unknown merge action: {action_type}"
+        raise ValueError(msg)
+
+    def _read_metrics(self) -> MergeResult:
+        history = self._spark.sql(f"DESCRIBE HISTORY {self._target_table_name} LIMIT 1")
+        row = history.collect()[0]
+        metrics: dict[str, str] = row["operationMetrics"] or {}
+
+        return MergeResult(
+            rows_inserted=int(metrics.get("numTargetRowsInserted", 0)),
+            rows_updated=int(metrics.get("numTargetRowsUpdated", 0)),
+            rows_deleted=int(metrics.get("numTargetRowsDeleted", 0)),
+        )
+
+    def _emit_metrics(self, result: MergeResult) -> None:
+        from databricks4py.metrics.base import MetricEvent
+
+        event = MetricEvent(
+            job_name="merge",
+            event_type="merge_complete",
+            timestamp=datetime.now(tz=timezone.utc),
+            row_count=result.rows_inserted + result.rows_updated,
+            table_name=self._target_table_name,
+            metadata={
+                "rows_inserted": result.rows_inserted,
+                "rows_updated": result.rows_updated,
+                "rows_deleted": result.rows_deleted,
+            },
+        )
+        self._metrics_sink.emit(event)  # type: ignore[union-attr]

databricks4py/io/streaming.py

@@ -0,0 +1,281 @@
+"""Structured Streaming utilities."""
+
+from __future__ import annotations
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.streaming import StreamingQuery
+
+from databricks4py.filters.base import Filter
+from databricks4py.spark_session import active_fallback
+
+if TYPE_CHECKING:
+    from databricks4py.io.checkpoint import CheckpointManager
+    from databricks4py.metrics.base import MetricsSink
+
+__all__ = [
+    "StreamingTableReader",
+    "StreamingTriggerOptions",
+]
+
+logger = logging.getLogger(__name__)
+
+
+class StreamingTriggerOptions(Enum):
+    """Common streaming trigger configurations.
+
+    Values are dicts suitable for passing to ``writeStream.trigger(**value)``.
+
+    Example::
+
+        trigger = StreamingTriggerOptions.PROCESSING_TIME_1M
+        stream.writeStream.trigger(**trigger.value).start()
+    """
+
+    PROCESSING_TIME_10S = {"processingTime": "10 seconds"}
+    PROCESSING_TIME_30S = {"processingTime": "30 seconds"}
+    PROCESSING_TIME_1M = {"processingTime": "1 minute"}
+    PROCESSING_TIME_5M = {"processingTime": "5 minutes"}
+    PROCESSING_TIME_10M = {"processingTime": "10 minutes"}
+    AVAILABLE_NOW = {"availableNow": True}
+
+
+class StreamingTableReader(ABC):
+    """Abstract base for streaming micro-batch processors.
+
+    Subclasses implement :meth:`process_batch` to handle each micro-batch.
+    The :meth:`start` method wires up ``foreachBatch`` and returns a
+    ``StreamingQuery``.
+
+    Example::
+
+        class MyProcessor(StreamingTableReader):
+            def process_batch(self, df, batch_id):
+                df.write.format("delta").mode("append").saveAsTable("output")
+
+        reader = MyProcessor(
+            source_table="catalog.schema.input",
+            trigger=StreamingTriggerOptions.PROCESSING_TIME_1M,
+            checkpoint_location="/checkpoints/my_reader",
+        )
+        query = reader.start()
+        query.awaitTermination()
+
+    Args:
+        source_table: Table name or path to read as a stream.
+        trigger: Trigger configuration, a raw dict, or None for the default.
+        checkpoint_location: Path for streaming checkpoints. Auto-generated
+            when a ``checkpoint_manager`` is provided and this is None.
+        source_format: Source format (default ``"delta"``).
+        row_filter: Optional Filter to apply before processing each batch.
+        skip_empty_batches: Skip batches with 0 rows (default True).
+        read_options: Additional read options as key-value pairs.
+        checkpoint_manager: Optional CheckpointManager for auto-generating
+            checkpoint paths.
+        metrics_sink: Optional MetricsSink for emitting batch metrics.
+        dead_letter_table: Fully qualified table name to write failed batches
+            to. When set and ``process_batch`` raises, the offending DataFrame
+            is written here with ``_dlq_error_message``, ``_dlq_error_timestamp``,
+            and ``_dlq_batch_id`` columns appended. Uses ``mergeSchema=true`` so
+            the table is auto-created on first failure.
+        spark: Optional SparkSession.
+    """
+
+    _DEFAULT_TRIGGER: dict[str, str] = {"processingTime": "10 seconds"}
+
+    def __init__(
+        self,
+        source_table: str,
+        trigger: StreamingTriggerOptions | dict | None = None,
+        checkpoint_location: str | None = None,
+        *,
+        source_format: str = "delta",
+        row_filter: Filter | None = None,
+        skip_empty_batches: bool = True,
+        read_options: dict[str, str] | None = None,
+        checkpoint_manager: CheckpointManager | None = None,
+        metrics_sink: MetricsSink | None = None,
+        dead_letter_table: str | None = None,
+        spark: SparkSession | None = None,
+    ) -> None:
+        self._spark = active_fallback(spark)
+        self._source_table = source_table
+        self._metrics_sink = metrics_sink
+        self._dead_letter_table = dead_letter_table
+        self._query: StreamingQuery | None = None
+
+        # Resolve trigger to a plain dict
+        if trigger is None:
+            self._trigger_dict = self._DEFAULT_TRIGGER
+        elif isinstance(trigger, StreamingTriggerOptions):
+            self._trigger_dict = trigger.value
+        else:
+            self._trigger_dict = trigger
+
+        # Auto-generate checkpoint path when manager is provided
+        if checkpoint_location is None and checkpoint_manager is not None:
+            checkpoint_location = checkpoint_manager.path_for(source_table, self.__class__.__name__)
+        if checkpoint_location is None:
+            raise ValueError(
+                "checkpoint_location is required when no checkpoint_manager is provided"
+            )
+
+        self._checkpoint_location = checkpoint_location
+        self._source_format = source_format
+        self._filter = row_filter
+        self._skip_empty_batches = skip_empty_batches
+        self._read_options = read_options or {}
+
+    @abstractmethod
+    def process_batch(self, df: DataFrame, batch_id: int) -> None:
+        """Process a single micro-batch.
+
+        Args:
+            df: The micro-batch DataFrame.
+            batch_id: The batch identifier.
+        """
+        ...
+
+    def _write_to_dlq(self, df: DataFrame, batch_id: int, error_msg: str) -> None:
+        """Append a failed batch to the dead-letter table with error metadata."""
+        assert self._dead_letter_table is not None  # only called when dlq is configured
+        from datetime import datetime, timezone
+
+        from pyspark.sql import functions as F
+
+        error_df = (
+            df.withColumn("_dlq_error_message", F.lit(error_msg))
+            .withColumn(
+                "_dlq_error_timestamp",
+                F.lit(datetime.now(tz=timezone.utc).isoformat()).cast("timestamp"),
+            )
+            .withColumn("_dlq_batch_id", F.lit(batch_id))
+        )
+        (
+            error_df.write.format("delta")
+            .mode("append")
+            .option("mergeSchema", "true")
+            .saveAsTable(self._dead_letter_table)
+        )
+        logger.warning(
+            "Wrote batch %d to DLQ %s: %s",
+            batch_id,
+            self._dead_letter_table,
+            error_msg[:200],
+        )
+
+    def _foreach_batch_wrapper(self, df: DataFrame, batch_id: int) -> None:
+        """Internal wrapper handling empty batch detection, filtering, and metrics."""
+        if self._skip_empty_batches and df.isEmpty():
+            logger.debug("Skipping empty batch %d", batch_id)
+            return
+
+        if self._filter is not None:
+            df = self._filter(df)
+            if self._skip_empty_batches and df.isEmpty():
+                logger.debug("Skipping batch %d (empty after filtering)", batch_id)
+                return
+
+        count = df.count()
+        logger.info("Processing batch %d (%d rows)", batch_id, count)
+
+        start = time.monotonic()
+        try:
+            self.process_batch(df, batch_id)
+        except Exception:  # noqa: BLE001 — broad catch intentional: DLQ must capture all failures
+            if self._dead_letter_table is not None:
+                import traceback
+
+                original_tb = traceback.format_exc()
+                try:
+                    self._write_to_dlq(df, batch_id, original_tb)
+                except Exception:
+                    logger.error(
+                        "DLQ write failed for batch %d; original error: %s",
+                        batch_id,
+                        original_tb[:500],
+                        exc_info=True,
+                    )
+                    raise
+                return
+            raise
+        duration_ms = (time.monotonic() - start) * 1000
+
+        if self._metrics_sink is not None:
+            from datetime import datetime, timezone
+
+            from databricks4py.metrics.base import MetricEvent
+
+            self._metrics_sink.emit(
+                MetricEvent(
+                    job_name=self.__class__.__name__,
+                    event_type="batch_complete",
+                    timestamp=datetime.now(tz=timezone.utc),
+                    duration_ms=duration_ms,
+                    row_count=count,
+                    batch_id=batch_id,
+                    table_name=self._source_table,
+                )
+            )
+
+    def _build_read_stream(self) -> DataFrame:
+        """Build the readStream DataFrame."""
+        reader = self._spark.readStream.format(self._source_format)
+        for key, value in self._read_options.items():
+            reader = reader.option(key, value)
+
+        if self._source_format == "delta":
+            return reader.table(self._source_table)
+        return reader.load(self._source_table)
+
+    def start(self) -> StreamingQuery:
+        """Start the streaming query.
+
+        Returns:
+            The active StreamingQuery.
+        """
+        stream_df = self._build_read_stream()
+
+        self._query = (
+            stream_df.writeStream.foreachBatch(self._foreach_batch_wrapper)
+            .trigger(**self._trigger_dict)
+            .option("checkpointLocation", self._checkpoint_location)
+            .start()
+        )
+
+        logger.info(
+            "Started streaming query from %s with trigger %s",
+            self._source_table,
+            self._trigger_dict,
+        )
+        return self._query
+
+    def stop(self, timeout_seconds: int = 30) -> None:
+        """Stop the streaming query and wait for graceful termination.
+
+        Args:
+            timeout_seconds: Maximum seconds to wait after calling stop.
+
+        Raises:
+            ValueError: If the query has not been started yet.
+        """
+        if self._query is None:
+            raise ValueError("No active query. Call start() first.")
+        self._query.stop()
+        self._query.awaitTermination(timeout=timeout_seconds)
+        logger.info("Streaming query stopped")
+
+    @property
+    def query(self) -> StreamingQuery | None:
+        """The active StreamingQuery, or None if start() has not been called."""
+        return self._query
+
+    @property
+    def is_active(self) -> bool:
+        """True if the streaming query is currently running."""
+        return self._query is not None and self._query.isActive

databricks4py/logging.py

@@ -0,0 +1,39 @@
+"""Logging configuration for Spark applications."""
+
+from __future__ import annotations
+
+import logging
+import os
+
+__all__ = ["configure_logging", "get_logger"]
+
+
+def configure_logging(level: int | str | None = None) -> None:
+    """Configure root logger with standard formatting.
+
+    Reads ``LOG_LEVEL`` environment variable if no level is provided.
+    Silences noisy ``py4j`` logger.
+
+    Args:
+        level: Log level (e.g. ``logging.INFO``, ``"DEBUG"``).
+            Defaults to ``LOG_LEVEL`` env var or ``INFO``.
+    """
+    if level is None:
+        level = os.getenv("LOG_LEVEL", "INFO")
+
+    formatting = "[%(levelname)s] [%(asctime)s] %(name)s - %(message)s"
+    logging.basicConfig(level=level, format=formatting, force=True)
+
+    # Silence py4j noise
+    logging.getLogger("py4j").setLevel(logging.ERROR)
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Get a named logger.
+
+    Convenience wrapper around ``logging.getLogger``.
+
+    Args:
+        name: Logger name (typically ``__name__``).
+    """
+    return logging.getLogger(name)

databricks4py/metrics/__init__.py

@@ -0,0 +1,22 @@
+"""Metrics collection with pluggable sinks."""
+
+from __future__ import annotations
+
+from databricks4py.metrics.base import CompositeMetricsSink, MetricEvent, MetricsSink
+from databricks4py.metrics.logging_sink import LoggingMetricsSink
+
+__all__ = [
+    "CompositeMetricsSink",
+    "DeltaMetricsSink",
+    "LoggingMetricsSink",
+    "MetricEvent",
+    "MetricsSink",
+]
+
+
+def __getattr__(name: str):
+    if name == "DeltaMetricsSink":
+        from databricks4py.metrics.delta_sink import DeltaMetricsSink
+
+        return DeltaMetricsSink
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

databricks4py/metrics/base.py

@@ -0,0 +1,66 @@
+"""Core metric types and sink abstractions."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+__all__ = ["CompositeMetricsSink", "MetricEvent", "MetricsSink"]
+
+
+@dataclass(frozen=True)
+class MetricEvent:
+    """A single metrics observation."""
+
+    job_name: str
+    event_type: str
+    timestamp: datetime
+    duration_ms: int | None = None
+    row_count: int | None = None
+    table_name: str | None = None
+    batch_id: int | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class MetricsSink(ABC):
+    """Abstract base for metrics destinations.
+
+    Subclasses must implement :meth:`emit`. Override :meth:`flush` if the
+    sink buffers events (e.g. :class:`~databricks4py.metrics.delta_sink.DeltaMetricsSink`).
+    """
+
+    @abstractmethod
+    def emit(self, event: MetricEvent) -> None:
+        """Send a single metric event to the destination."""
+        ...
+
+    def flush(self) -> None:  # noqa: B027
+        """Flush any buffered events. Default is a no-op."""
+
+
+class CompositeMetricsSink(MetricsSink):
+    """Fans out events to multiple sinks.
+
+    Example::
+
+        sink = CompositeMetricsSink(
+            DeltaMetricsSink("catalog.schema.metrics"),
+            LoggingMetricsSink(),
+        )
+
+    Args:
+        sinks: One or more MetricsSink instances to delegate to.
+    """
+
+    def __init__(self, *sinks: MetricsSink) -> None:
+        self._sinks = sinks
+
+    def emit(self, event: MetricEvent) -> None:
+        for sink in self._sinks:
+            sink.emit(event)
+
+    def flush(self) -> None:
+        for sink in self._sinks:
+            sink.flush()

databricks4py/metrics/delta_sink.py

@@ -0,0 +1,75 @@
+"""Delta Lake metrics sink."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import asdict
+from typing import TYPE_CHECKING
+
+from databricks4py.metrics.base import MetricEvent, MetricsSink
+from databricks4py.spark_session import active_fallback
+
+if TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+
+__all__ = ["DeltaMetricsSink"]
+
+logger = logging.getLogger(__name__)
+
+
+class DeltaMetricsSink(MetricsSink):
+    """Buffers metric events and writes them to a Delta table on flush or threshold."""
+
+    def __init__(
+        self,
+        table_name: str,
+        *,
+        spark: SparkSession | None = None,
+        buffer_size: int = 100,
+    ) -> None:
+        self._table_name = table_name
+        self._spark = active_fallback(spark)
+        self._buffer_size = buffer_size
+        self._buffer: list[MetricEvent] = []
+
+    def emit(self, event: MetricEvent) -> None:
+        self._buffer.append(event)
+        if len(self._buffer) >= self._buffer_size:
+            self.flush()
+
+    def flush(self) -> None:
+        if not self._buffer:
+            return
+
+        from pyspark.sql.types import (
+            DoubleType,
+            IntegerType,
+            StringType,
+            StructField,
+            StructType,
+        )
+
+        schema = StructType(
+            [
+                StructField("job_name", StringType()),
+                StructField("event_type", StringType()),
+                StructField("timestamp", StringType()),
+                StructField("duration_ms", DoubleType()),
+                StructField("row_count", IntegerType()),
+                StructField("batch_id", IntegerType()),
+                StructField("table_name", StringType()),
+                StructField("metadata", StringType()),
+            ]
+        )
+
+        rows = [
+            {
+                k: str(v) if not isinstance(v, (int, float, str, type(None))) else v
+                for k, v in asdict(event).items()
+            }
+            for event in self._buffer
+        ]
+        df = self._spark.createDataFrame(rows, schema=schema)
+        df.write.format("delta").mode("append").saveAsTable(self._table_name)
+        logger.info("Flushed %d metric events to %s", len(self._buffer), self._table_name)
+        self._buffer.clear()

databricks4py/metrics/logging_sink.py

@@ -0,0 +1,20 @@
+"""Logging-based metrics sink."""
+
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import asdict
+
+from databricks4py.metrics.base import MetricEvent, MetricsSink
+
+__all__ = ["LoggingMetricsSink"]
+
+logger = logging.getLogger(__name__)
+
+
+class LoggingMetricsSink(MetricsSink):
+    """Emits metric events as JSON via the standard logger."""
+
+    def emit(self, event: MetricEvent) -> None:
+        logger.info(json.dumps(asdict(event), default=str))