databricks4py 0.2.0__py3-none-any.whl

databricks4py/observability/query_listener.py

@@ -0,0 +1,236 @@
+"""Streaming query progress observer via PySpark StreamingQueryListener.
+
+Wraps the PySpark 3.4+ ``StreamingQueryListener`` API into a simple observer
+that collects progress snapshots, optionally emits them to a
+:class:`~databricks4py.metrics.base.MetricsSink`, and exposes the latest
+state for health checks.
+
+Requires PySpark >= 3.4. On older versions, :meth:`QueryProgressObserver.attach`
+raises ``ImportError`` with a clear message.
+
+Example::
+
+    observer = QueryProgressObserver(spark=spark, metrics_sink=my_sink)
+    observer.attach()
+
+    query = df.writeStream.start()
+    # ... wait for progress ...
+
+    latest = observer.latest_progress()
+    if latest:
+        print(f"Batch {latest.batch_id}: {latest.num_input_rows} rows")
+
+    observer.detach()
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections import deque
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from databricks4py.observability._utils import parse_duration_ms
+from databricks4py.spark_session import active_fallback
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from pyspark.sql import SparkSession
+
+    from databricks4py.metrics.base import MetricsSink
+
+__all__ = ["QueryProgressObserver", "QueryProgressSnapshot"]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class QueryProgressSnapshot:
+    """Immutable snapshot of a single streaming query progress event.
+
+    Attributes:
+        query_id: Unique query identifier (UUID string).
+        query_name: Optional query name (set via ``.queryName()``).
+        batch_id: Monotonically increasing batch counter.
+        batch_duration_ms: Wall-clock time for the batch in milliseconds.
+        input_rows_per_second: Rate of data arriving from the source.
+        processed_rows_per_second: Rate of data being processed.
+        num_input_rows: Total rows read in this batch.
+        timestamp: UTC time when the progress was recorded.
+        sources: Per-source offset and rate info.
+        sink: Sink commit progress info.
+        raw_json: Full progress JSON for custom parsing.
+    """
+
+    query_id: str
+    query_name: str | None
+    batch_id: int
+    batch_duration_ms: int
+    input_rows_per_second: float
+    processed_rows_per_second: float
+    num_input_rows: int
+    timestamp: datetime
+    sources: list[dict[str, Any]] = field(default_factory=list)
+    sink: dict[str, Any] = field(default_factory=dict)
+    raw_json: str = ""
+
+    @classmethod
+    def from_progress(cls, progress: Any) -> QueryProgressSnapshot:
+        """Build from a PySpark ``StreamingQueryProgress`` object.
+
+        Expects an object with a ``.json`` property returning a JSON string
+        (standard ``StreamingQueryProgress`` in PySpark 3.4+).
+        """
+        import json as _json
+
+        raw = progress.json
+        data = _json.loads(raw)
+
+        return cls(
+            query_id=data.get("id", ""),
+            query_name=data.get("name"),
+            batch_id=data.get("batchId", -1),
+            batch_duration_ms=parse_duration_ms(data.get("batchDuration", "0 ms")),
+            input_rows_per_second=data.get("inputRowsPerSecond", 0.0),
+            processed_rows_per_second=data.get("processedRowsPerSecond", 0.0),
+            num_input_rows=data.get("numInputRows", 0),
+            timestamp=datetime.now(tz=timezone.utc),
+            sources=data.get("sources", []),
+            sink=data.get("sink", {}),
+            raw_json=raw,
+        )
+
+
+class QueryProgressObserver:
+    """Collects streaming query progress events and routes them to a metrics sink.
+
+    Attaches a ``StreamingQueryListener`` to the SparkSession. Each progress
+    event is captured as a :class:`QueryProgressSnapshot`, stored in a bounded
+    history, and optionally emitted as a ``MetricEvent``.
+
+    Args:
+        spark: SparkSession to attach the listener to.
+        metrics_sink: Optional sink for ``query_progress`` metric events.
+        on_progress: Optional callback invoked on each progress event.
+        history_size: Maximum number of snapshots to retain. Oldest are evicted.
+        query_name_filter: If set, only track queries with this name.
+    """
+
+    def __init__(
+        self,
+        *,
+        spark: SparkSession | None = None,
+        metrics_sink: MetricsSink | None = None,
+        on_progress: Callable[[QueryProgressSnapshot], None] | None = None,
+        history_size: int = 100,
+        query_name_filter: str | None = None,
+    ) -> None:
+        self._spark = active_fallback(spark)
+        self._metrics_sink = metrics_sink
+        self._on_progress = on_progress
+        self._history: deque[QueryProgressSnapshot] = deque(maxlen=history_size)
+        self._query_name_filter = query_name_filter
+        self._listener: Any = None
+        self._attached = False
+
+    def attach(self) -> None:
+        """Register the listener with the SparkSession.
+
+        Raises:
+            ImportError: If PySpark < 3.4 (no Python listener support).
+        """
+        try:
+            from pyspark.sql.streaming.listener import StreamingQueryListener
+        except ImportError as exc:
+            raise ImportError(
+                "StreamingQueryListener requires PySpark >= 3.4. "
+                "Upgrade PySpark or use polling via query.lastProgress instead."
+            ) from exc
+
+        observer = self
+
+        class _Listener(StreamingQueryListener):
+            def onQueryStarted(self, event: Any) -> None:
+                logger.debug("Query started: id=%s name=%s", event.id, event.name)
+
+            def onQueryProgress(self, event: Any) -> None:
+                observer._handle_progress(event.progress)
+
+            def onQueryTerminated(self, event: Any) -> None:
+                exc_str = str(event.exception)[:500] if event.exception else None
+                logger.info("Query terminated: id=%s exception=%s", event.id, exc_str)
+
+        self._listener = _Listener()
+        self._spark.streams.addListener(self._listener)
+        self._attached = True
+        logger.info("QueryProgressObserver attached")
+
+    def detach(self) -> None:
+        """Remove the listener from the SparkSession."""
+        if self._listener is not None and self._attached:
+            self._spark.streams.removeListener(self._listener)
+            self._attached = False
+            logger.info("QueryProgressObserver detached")
+
+    def _handle_progress(self, progress: Any) -> None:
+        try:
+            snapshot = QueryProgressSnapshot.from_progress(progress)
+        except (json.JSONDecodeError, KeyError, TypeError, ValueError, AttributeError):
+            logger.exception("Failed to process query progress event")
+            return
+
+        if self._query_name_filter is not None and snapshot.query_name != self._query_name_filter:
+            return
+
+        self._history.append(snapshot)
+
+        if self._on_progress is not None:
+            self._on_progress(snapshot)
+
+        if self._metrics_sink is not None:
+            from databricks4py.metrics.base import MetricEvent
+
+            self._metrics_sink.emit(
+                MetricEvent(
+                    job_name=snapshot.query_name or snapshot.query_id,
+                    event_type="query_progress",
+                    timestamp=snapshot.timestamp,
+                    duration_ms=snapshot.batch_duration_ms,
+                    row_count=snapshot.num_input_rows,
+                    batch_id=snapshot.batch_id,
+                    metadata={
+                        "input_rows_per_second": snapshot.input_rows_per_second,
+                        "processed_rows_per_second": snapshot.processed_rows_per_second,
+                    },
+                )
+            )
+
+        logger.debug(
+            "Progress: query=%s batch=%d rows=%d rate=%.1f rows/s",
+            snapshot.query_name or snapshot.query_id,
+            snapshot.batch_id,
+            snapshot.num_input_rows,
+            snapshot.processed_rows_per_second,
+        )
+
+    def latest_progress(self) -> QueryProgressSnapshot | None:
+        """Most recent progress snapshot, or None if no events received yet."""
+        return self._history[-1] if self._history else None
+
+    def history(self, limit: int | None = None) -> list[QueryProgressSnapshot]:
+        """Recent progress snapshots, newest last.
+
+        Args:
+            limit: Maximum number of snapshots to return. Returns all if None.
+        """
+        items = list(self._history)
+        if limit is not None:
+            items = items[-limit:]
+        return items
+
+    @property
+    def is_attached(self) -> bool:
+        return self._attached

databricks4py/quality/__init__.py

@@ -0,0 +1,26 @@
+"""Data quality expectations and enforcement."""
+
+from databricks4py.quality.base import Expectation, ExpectationResult, QualityReport
+from databricks4py.quality.expectations import (
+    ColumnExists,
+    InRange,
+    MatchesRegex,
+    NotNull,
+    RowCount,
+    Unique,
+)
+from databricks4py.quality.gate import QualityError, QualityGate
+
+__all__ = [
+    "ColumnExists",
+    "Expectation",
+    "ExpectationResult",
+    "InRange",
+    "MatchesRegex",
+    "NotNull",
+    "QualityError",
+    "QualityGate",
+    "QualityReport",
+    "RowCount",
+    "Unique",
+]

databricks4py/quality/base.py

@@ -0,0 +1,54 @@
+"""Core quality abstractions: expectations, results, and reports."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pyspark.sql import Column, DataFrame, Row
+
+__all__ = ["Expectation", "ExpectationResult", "QualityReport"]
+
+
+@dataclass(frozen=True)
+class ExpectationResult:
+    """Result of a single expectation check."""
+
+    expectation: str
+    passed: bool
+    total_rows: int
+    failing_rows: int = 0
+    sample: list[Row] | None = None
+
+
+@dataclass(frozen=True)
+class QualityReport:
+    """Aggregated results from running multiple expectations."""
+
+    results: list[ExpectationResult]
+    passed: bool
+
+    def summary(self) -> str:
+        lines = []
+        for r in self.results:
+            status = "PASS" if r.passed else "FAIL"
+            lines.append(f"[{status}] {r.expectation} ({r.failing_rows}/{r.total_rows} failing)")
+        overall = "PASSED" if self.passed else "FAILED"
+        lines.append(f"Overall: {overall}")
+        return "\n".join(lines)
+
+
+class Expectation(ABC):
+    """Abstract base for DataFrame quality expectations."""
+
+    @abstractmethod
+    def validate(self, df: DataFrame) -> ExpectationResult: ...
+
+    def failing_condition(self) -> Column | None:
+        """Return a Column expression that is True for failing rows.
+
+        Returns None for aggregate checks that can't filter individual rows.
+        """
+        return None

databricks4py/quality/expectations.py

@@ -0,0 +1,184 @@
+"""Built-in expectation implementations."""
+
+from __future__ import annotations
+
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from databricks4py.quality.base import Expectation, ExpectationResult
+
+if TYPE_CHECKING:
+    from pyspark.sql import Column, DataFrame
+
+__all__ = ["ColumnExists", "InRange", "MatchesRegex", "NotNull", "RowCount", "Unique"]
+
+
+class NotNull(Expectation):
+    """Validates that specified columns contain no null values."""
+
+    def __init__(self, *columns: str) -> None:
+        self._columns = columns
+
+    def validate(self, df: DataFrame) -> ExpectationResult:
+        total = df.count()
+        condition = self.failing_condition()
+        failing = df.where(condition).count() if condition is not None else 0
+        return ExpectationResult(
+            expectation=repr(self),
+            passed=failing == 0,
+            total_rows=total,
+            failing_rows=failing,
+        )
+
+    def failing_condition(self) -> Column | None:
+        from pyspark.sql import functions as F
+
+        conditions = [F.col(c).isNull() for c in self._columns]
+        return reduce(lambda a, b: a | b, conditions)
+
+    def __repr__(self) -> str:
+        cols = ", ".join(repr(c) for c in self._columns)
+        return f"NotNull({cols})"
+
+
+class InRange(Expectation):
+    """Validates that a column's values fall within bounds."""
+
+    def __init__(
+        self, column: str, *, min_val: float | None = None, max_val: float | None = None
+    ) -> None:
+        self._column = column
+        self._min_val = min_val
+        self._max_val = max_val
+
+    def validate(self, df: DataFrame) -> ExpectationResult:
+        total = df.count()
+        condition = self.failing_condition()
+        failing = df.where(condition).count() if condition is not None else 0
+        return ExpectationResult(
+            expectation=repr(self),
+            passed=failing == 0,
+            total_rows=total,
+            failing_rows=failing,
+        )
+
+    def failing_condition(self) -> Column | None:
+        from pyspark.sql import functions as F
+
+        conditions = []
+        col = F.col(self._column)
+        if self._min_val is not None:
+            conditions.append(col < self._min_val)
+        if self._max_val is not None:
+            conditions.append(col > self._max_val)
+        if not conditions:
+            return None
+        return reduce(lambda a, b: a | b, conditions)
+
+    def __repr__(self) -> str:
+        return f"InRange({self._column!r}, min_val={self._min_val!r}, max_val={self._max_val!r})"
+
+
+class Unique(Expectation):
+    """Validates no duplicate rows exist for the specified columns."""
+
+    def __init__(self, *columns: str) -> None:
+        self._columns = columns
+
+    def validate(self, df: DataFrame) -> ExpectationResult:
+        total = df.count()
+        distinct = df.select(*self._columns).distinct().count()
+        failing = total - distinct
+        return ExpectationResult(
+            expectation=repr(self),
+            passed=failing == 0,
+            total_rows=total,
+            failing_rows=failing,
+        )
+
+    def __repr__(self) -> str:
+        cols = ", ".join(repr(c) for c in self._columns)
+        return f"Unique({cols})"
+
+
+class RowCount(Expectation):
+    """Validates that the DataFrame row count falls within bounds."""
+
+    def __init__(self, *, min_count: int | None = None, max_count: int | None = None) -> None:
+        self._min_count = min_count
+        self._max_count = max_count
+
+    def validate(self, df: DataFrame) -> ExpectationResult:
+        total = df.count()
+        passed = True
+        if self._min_count is not None and total < self._min_count:
+            passed = False
+        if self._max_count is not None and total > self._max_count:
+            passed = False
+        return ExpectationResult(
+            expectation=repr(self),
+            passed=passed,
+            total_rows=total,
+        )
+
+    def __repr__(self) -> str:
+        return f"RowCount(min_count={self._min_count!r}, max_count={self._max_count!r})"
+
+
+class MatchesRegex(Expectation):
+    """Validates that a column's values match a regex pattern."""
+
+    def __init__(self, column: str, pattern: str) -> None:
+        self._column = column
+        self._pattern = pattern
+
+    def validate(self, df: DataFrame) -> ExpectationResult:
+        total = df.count()
+        condition = self.failing_condition()
+        failing = df.where(condition).count() if condition is not None else 0
+        return ExpectationResult(
+            expectation=repr(self),
+            passed=failing == 0,
+            total_rows=total,
+            failing_rows=failing,
+        )
+
+    def failing_condition(self) -> Column | None:
+        from pyspark.sql import functions as F
+
+        return ~F.col(self._column).rlike(self._pattern)
+
+    def __repr__(self) -> str:
+        return f"MatchesRegex({self._column!r}, {self._pattern!r})"
+
+
+class ColumnExists(Expectation):
+    """Validates that specified columns exist in the DataFrame schema."""
+
+    def __init__(self, *columns: str, dtype: str | None = None) -> None:
+        self._columns = columns
+        self._dtype = dtype
+
+    def validate(self, df: DataFrame) -> ExpectationResult:
+        total = df.count()
+        actual_cols = set(df.columns)
+        missing = [c for c in self._columns if c not in actual_cols]
+
+        if not missing and self._dtype:
+            schema_fields = {f.name: f.dataType.simpleString() for f in df.schema.fields}
+            for c in self._columns:
+                if schema_fields.get(c) != self._dtype:
+                    missing.append(c)
+
+        return ExpectationResult(
+            expectation=repr(self),
+            passed=len(missing) == 0,
+            total_rows=total,
+            failing_rows=len(missing),
+        )
+
+    def __repr__(self) -> str:
+        cols = ", ".join(repr(c) for c in self._columns)
+        if self._dtype:
+            return f"ColumnExists({cols}, dtype={self._dtype!r})"
+        return f"ColumnExists({cols})"

databricks4py/quality/gate.py

@@ -0,0 +1,90 @@
+"""Quality gate: orchestrate expectations and enforce data quality."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Literal
+
+from databricks4py.quality.base import Expectation, QualityReport
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame
+
+__all__ = ["QualityError", "QualityGate"]
+
+logger = logging.getLogger(__name__)
+
+
+class QualityError(Exception):
+    """Raised when a quality gate fails in 'raise' mode."""
+
+    def __init__(self, report: QualityReport) -> None:
+        self.report = report
+        super().__init__(f"Quality check failed:\n{report.summary()}")
+
+
+class QualityGate:
+    """Runs expectations against a DataFrame and enforces quality policy.
+
+    Args:
+        *expectations: Expectation instances to evaluate.
+        on_fail: Action on failure: "raise", "warn", or "quarantine".
+        quarantine_handler: Callable receiving the bad-rows DataFrame.
+            Required when on_fail="quarantine".
+    """
+
+    def __init__(
+        self,
+        *expectations: Expectation,
+        on_fail: Literal["raise", "warn", "quarantine"] = "raise",
+        quarantine_handler: Callable[[DataFrame], None] | None = None,
+    ) -> None:
+        if on_fail == "quarantine" and quarantine_handler is None:
+            raise ValueError("quarantine_handler is required when on_fail='quarantine'")
+        self._expectations = expectations
+        self._on_fail = on_fail
+        self._quarantine_handler = quarantine_handler
+
+    def check(self, df: DataFrame) -> QualityReport:
+        """Run all expectations and return a report."""
+        results = [exp.validate(df) for exp in self._expectations]
+        passed = all(r.passed for r in results)
+        return QualityReport(results=results, passed=passed)
+
+    def enforce(self, df: DataFrame) -> DataFrame:
+        """Run expectations and enforce the configured failure policy.
+
+        Returns the clean DataFrame (original if all pass, or filtered
+        if quarantine mode splits out bad rows).
+        """
+        report = self.check(df)
+        if report.passed:
+            return df
+
+        if self._on_fail == "raise":
+            raise QualityError(report)
+
+        if self._on_fail == "warn":
+            logger.warning("Quality check failed:\n%s", report.summary())
+            return df
+
+        # quarantine: split bad rows using failing_condition()
+        from functools import reduce
+
+        conditions = []
+        for exp in self._expectations:
+            cond = exp.failing_condition()
+            if cond is not None:
+                conditions.append(cond)
+
+        if conditions:
+            bad_condition = reduce(lambda a, b: a | b, conditions)
+            bad_rows = df.where(bad_condition)
+            clean_rows = df.where(~bad_condition)
+            self._quarantine_handler(bad_rows)  # type: ignore[misc]
+            return clean_rows
+
+        # no row-level conditions available — quarantine nothing
+        logger.warning("Quarantine requested but no row-level failing conditions available")
+        return df

databricks4py/retry.py

@@ -0,0 +1,102 @@
+"""Retry decorator with exponential backoff for transient failures."""
+
+from __future__ import annotations
+
+import functools
+import logging
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+__all__ = ["RetryConfig", "retry"]
+
+logger = logging.getLogger(__name__)
+
+
+def _default_retryable_exceptions() -> tuple[type[BaseException], ...]:
+    exceptions: list[type[BaseException]] = [ConnectionError, TimeoutError, OSError]
+    try:
+        from py4j.protocol import Py4JNetworkError
+
+        exceptions.append(Py4JNetworkError)
+    except ImportError:
+        pass
+    return tuple(exceptions)
+
+
+@dataclass(frozen=True)
+class RetryConfig:
+    """Configuration for :func:`retry` behaviour.
+
+    Args:
+        max_attempts: Total number of tries (including the first).
+        base_delay_seconds: Initial delay before the first retry.
+        max_delay_seconds: Upper cap on the exponentially increasing delay.
+        backoff_factor: Multiplier applied to the delay after each failure.
+        retryable_exceptions: Exception types that trigger a retry.
+            Defaults to ``ConnectionError``, ``TimeoutError``, ``OSError``,
+            and ``Py4JNetworkError`` (if py4j is installed).
+    """
+
+    max_attempts: int = 3
+    base_delay_seconds: float = 1.0
+    max_delay_seconds: float = 60.0
+    backoff_factor: float = 2.0
+    retryable_exceptions: tuple[type[BaseException], ...] = field(default_factory=tuple)
+
+    def __post_init__(self) -> None:
+        if not self.retryable_exceptions:
+            object.__setattr__(self, "retryable_exceptions", _default_retryable_exceptions())
+
+
+def retry(config: RetryConfig | None = None) -> Callable:
+    """Decorator factory for retrying functions with exponential backoff.
+
+    Example::
+
+        @retry(RetryConfig(max_attempts=5, base_delay_seconds=2.0))
+        def fetch_data():
+            return requests.get(url).json()
+
+    Args:
+        config: Retry parameters. Uses defaults if ``None``.
+    """
+    if config is None:
+        config = RetryConfig()
+
+    def decorator(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            last_exception: BaseException | None = None
+            for attempt in range(1, config.max_attempts + 1):
+                try:
+                    return fn(*args, **kwargs)
+                except config.retryable_exceptions as exc:
+                    last_exception = exc
+                    if attempt == config.max_attempts:
+                        logger.error(
+                            "%s failed after %d attempts: %s",
+                            fn.__name__,
+                            config.max_attempts,
+                            exc,
+                        )
+                        raise
+                    delay = min(
+                        config.base_delay_seconds * (config.backoff_factor ** (attempt - 1)),
+                        config.max_delay_seconds,
+                    )
+                    logger.warning(
+                        "%s attempt %d/%d failed: %s. Retrying in %.1fs",
+                        fn.__name__,
+                        attempt,
+                        config.max_attempts,
+                        exc,
+                        delay,
+                    )
+                    time.sleep(delay)
+            raise last_exception  # type: ignore[misc]
+
+        return wrapper
+
+    return decorator