databricks4py 0.2.0__py3-none-any.whl

databricks4py/__init__.py

@@ -0,0 +1,56 @@
+"""databricks4py: Spark, Delta Lake, and Databricks utility library.
+
+A collection of reusable abstractions for building PySpark applications
+on Databricks and locally.
+"""
+
+__version__ = "0.2.0"
+
+from databricks4py.catalog import CatalogSchema
+from databricks4py.config import Environment, JobConfig, UnityConfig
+from databricks4py.logging import configure_logging, get_logger
+from databricks4py.metrics import CompositeMetricsSink, LoggingMetricsSink, MetricEvent, MetricsSink
+from databricks4py.retry import RetryConfig, retry
+from databricks4py.secrets import SecretFetcher
+from databricks4py.spark_session import active_fallback, get_active, get_or_create_local_session
+from databricks4py.workflow import Workflow
+
+
+def inject_dbutils(dbutils_module):
+    """Unified dbutils injection for secrets and DBFS operations."""
+    from databricks4py.io.dbfs import _set_dbutils_module
+
+    SecretFetcher.dbutils = dbutils_module
+    _set_dbutils_module(dbutils_module)
+
+
+__all__ = [
+    "__version__",
+    # SparkSession
+    "get_active",
+    "active_fallback",
+    "get_or_create_local_session",
+    # Catalog
+    "CatalogSchema",
+    # Config
+    "Environment",
+    "JobConfig",
+    "UnityConfig",
+    # Logging
+    "configure_logging",
+    "get_logger",
+    # Metrics
+    "CompositeMetricsSink",
+    "LoggingMetricsSink",
+    "MetricEvent",
+    "MetricsSink",
+    # Retry
+    "RetryConfig",
+    "retry",
+    # Secrets
+    "SecretFetcher",
+    # Workflow
+    "Workflow",
+    # Utilities
+    "inject_dbutils",
+]

databricks4py/catalog.py

@@ -0,0 +1,65 @@
+"""Catalog and schema-qualified table name management."""
+
+from __future__ import annotations
+
+__all__ = ["CatalogSchema"]
+
+
+class CatalogSchema:
+    """Schema-qualified table name registry.
+
+    Provides attribute-based access to fully qualified table names
+    within a schema. Supports versioned table aliases.
+
+    Example::
+
+        sales = CatalogSchema(
+            "sales",
+            tables=["orders", "customers"],
+            versioned_tables={"metrics": "metrics_v3"},
+        )
+
+        # Access table names
+        sales.orders      # "sales.orders"
+        sales.customers   # "sales.customers"
+        sales.metrics     # "sales.metrics_v3"
+
+    Args:
+        name: The schema name.
+        tables: List of table names in this schema.
+        versioned_tables: Dict mapping logical names to versioned names.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        tables: list[str] | None = None,
+        versioned_tables: dict[str, str] | None = None,
+    ) -> None:
+        self._name = name
+        self._tables: dict[str, str] = {}
+
+        for table in tables or []:
+            self._tables[table] = f"{name}.{table}"
+
+        for logical_name, versioned_name in (versioned_tables or {}).items():
+            self._tables[logical_name] = f"{name}.{versioned_name}"
+
+    @property
+    def schema_name(self) -> str:
+        """The schema name."""
+        return self._name
+
+    def __getattr__(self, name: str) -> str:
+        if name.startswith("_"):
+            raise AttributeError(name)
+        try:
+            return self._tables[name]
+        except KeyError:
+            raise AttributeError(
+                f"'{type(self).__name__}' has no table '{name}'. "
+                f"Available: {sorted(self._tables.keys())}"
+            ) from None
+
+    def __repr__(self) -> str:
+        return f"CatalogSchema({self._name!r}, tables={sorted(self._tables.keys())})"

databricks4py/config/__init__.py

@@ -0,0 +1,6 @@
+"""Configuration classes for Databricks job entry points."""
+
+from databricks4py.config.base import Environment, JobConfig
+from databricks4py.config.unity import UnityConfig
+
+__all__ = ["Environment", "JobConfig", "UnityConfig"]

databricks4py/config/base.py

@@ -0,0 +1,119 @@
+"""Base configuration for Databricks jobs."""
+
+from __future__ import annotations
+
+import logging
+import os
+from enum import Enum
+
+__all__ = ["Environment", "JobConfig"]
+
+logger = logging.getLogger(__name__)
+
+
+class Environment(Enum):
+    """Deployment environment. Resolved automatically from Databricks widgets or env vars."""
+
+    DEV = "dev"
+    STAGING = "staging"
+    PROD = "prod"
+
+
+class JobConfig:
+    """Configuration container for Databricks job parameters.
+
+    Resolves the deployment environment from (in priority order):
+    1. ``spark.databricks.widget.env`` conf (Databricks widget)
+    2. ``ENV`` or ``ENVIRONMENT`` environment variable
+    3. Defaults to ``DEV``
+
+    Example::
+
+        config = JobConfig(
+            tables={"events": "catalog.bronze.events", "users": "catalog.silver.users"},
+            secret_scope="my-scope",
+            spark_configs={"spark.sql.shuffle.partitions": "8"},
+        )
+        table_name = config.table("events")  # "catalog.bronze.events"
+
+    Args:
+        tables: Mapping of logical names to fully qualified table names.
+        secret_scope: Databricks secret scope for :meth:`secret` lookups.
+        storage_root: Optional root path for storage operations.
+        spark_configs: Spark configuration overrides applied by
+            :meth:`~databricks4py.workflow.Workflow.execute`.
+    """
+
+    def __init__(
+        self,
+        tables: dict[str, str],
+        *,
+        secret_scope: str | None = None,
+        storage_root: str | None = None,
+        spark_configs: dict[str, str] | None = None,
+    ) -> None:
+        self.tables = tables
+        self.secret_scope = secret_scope
+        self.storage_root = storage_root
+        self.spark_configs = spark_configs or {}
+        self.env = self._resolve_env()
+
+    def _resolve_env(self) -> Environment:
+        raw: str | None = None
+
+        # Try Databricks widget parameter first
+        try:
+            from pyspark.sql import SparkSession
+
+            spark = SparkSession.getActiveSession()
+            if spark is not None:
+                raw = spark.conf.get("spark.databricks.widget.env", None)
+        except Exception:
+            pass
+
+        if raw is None:
+            raw = os.getenv("ENV") or os.getenv("ENVIRONMENT")
+
+        if raw is None:
+            return Environment.DEV
+
+        try:
+            return Environment(raw.lower())
+        except ValueError:
+            logger.warning("Unknown environment '%s', defaulting to DEV", raw)
+            return Environment.DEV
+
+    def table(self, name: str) -> str:
+        """Look up a fully qualified table name by logical key.
+
+        Raises:
+            KeyError: If *name* is not in the configured tables.
+        """
+        try:
+            return self.tables[name]
+        except KeyError:
+            available = sorted(self.tables.keys())
+            raise KeyError(f"Table '{name}' not configured. Available: {available}") from None
+
+    def secret(self, key: str) -> str:
+        """Fetch a secret from Databricks using the configured scope.
+
+        Raises:
+            ValueError: If no ``secret_scope`` was configured.
+        """
+        if self.secret_scope is None:
+            raise ValueError("No secret_scope configured on this JobConfig")
+
+        from databricks4py.secrets import SecretFetcher
+
+        return SecretFetcher.fetch_secret(self.secret_scope, key)
+
+    @classmethod
+    def from_env(cls, **kwargs) -> JobConfig:
+        return cls(**kwargs)
+
+    def __repr__(self) -> str:
+        return (
+            f"JobConfig(env={self.env.value!r}, tables={len(self.tables)}, "
+            f"secret_scope={self.secret_scope!r})"
+        )

databricks4py/config/unity.py

@@ -0,0 +1,72 @@
+"""Unity Catalog-aware configuration."""
+
+from __future__ import annotations
+
+from databricks4py.config.base import JobConfig
+
+__all__ = ["UnityConfig"]
+
+
+class UnityConfig(JobConfig):
+    """Environment-aware Unity Catalog configuration.
+
+    Builds a catalog name from ``{catalog_prefix}_{env}`` (e.g. ``myapp_prod``)
+    and resolves table references in ``schema.table`` format to fully qualified
+    three-part names.
+
+    Example::
+
+        config = UnityConfig(
+            catalog_prefix="myapp",
+            schemas=["bronze", "silver"],
+        )
+        # In production: config.catalog == "myapp_prod"
+        config.table("bronze.events")  # "myapp_prod.bronze.events"
+
+    Args:
+        catalog_prefix: Base name prepended to the resolved environment.
+        schemas: Allowed schema names. :meth:`table` rejects unknown schemas.
+        secret_scope: Forwarded to :class:`JobConfig`.
+        storage_root: Forwarded to :class:`JobConfig`.
+        spark_configs: Forwarded to :class:`JobConfig`.
+    """
+
+    def __init__(
+        self,
+        catalog_prefix: str,
+        schemas: list[str],
+        *,
+        secret_scope: str | None = None,
+        storage_root: str | None = None,
+        spark_configs: dict[str, str] | None = None,
+    ) -> None:
+        super().__init__(
+            tables={},
+            secret_scope=secret_scope,
+            storage_root=storage_root,
+            spark_configs=spark_configs,
+        )
+        self.catalog_prefix = catalog_prefix
+        self.schemas = schemas
+        self.catalog = f"{catalog_prefix}_{self.env.value}"
+
+    def table(self, name: str) -> str:
+        """Resolve ``'schema.table'`` to ``'catalog.schema.table'``.
+
+        Raises:
+            ValueError: If *name* is not in ``schema.table`` format.
+            KeyError: If the schema is not in the configured list.
+        """
+        parts = name.split(".")
+        if len(parts) != 2:
+            raise ValueError(f"Expected 'schema.table' format, got '{name}'")
+        schema, table_name = parts
+        if schema not in self.schemas:
+            raise KeyError(f"Schema '{schema}' not in configured schemas: {sorted(self.schemas)}")
+        return f"{self.catalog}.{schema}.{table_name}"
+
+    def __repr__(self) -> str:
+        return (
+            f"UnityConfig(catalog={self.catalog!r}, schemas={self.schemas}, "
+            f"secret_scope={self.secret_scope!r})"
+        )

databricks4py/filters/__init__.py

@@ -0,0 +1,17 @@
+"""DataFrame filter pipeline."""
+
+from databricks4py.filters.base import (
+    ColumnFilter,
+    DropDuplicates,
+    Filter,
+    FilterPipeline,
+    WhereFilter,
+)
+
+__all__ = [
+    "ColumnFilter",
+    "DropDuplicates",
+    "Filter",
+    "FilterPipeline",
+    "WhereFilter",
+]

databricks4py/filters/base.py

@@ -0,0 +1,154 @@
+"""DataFrame filter abstractions."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+
+from pyspark.sql import DataFrame
+
+__all__ = [
+    "ColumnFilter",
+    "DropDuplicates",
+    "Filter",
+    "FilterPipeline",
+    "WhereFilter",
+]
+
+
+class Filter(ABC):
+    """Abstract base class for DataFrame filters.
+
+    Subclasses implement :meth:`apply` to transform a DataFrame.
+    Filters are callable — ``filter(df)`` is equivalent to ``filter.apply(df)``.
+
+    Example::
+
+        class ActiveOnly(Filter):
+            def apply(self, df: DataFrame) -> DataFrame:
+                return df.where("is_active = true")
+
+        df = ActiveOnly()(raw_df)
+    """
+
+    @abstractmethod
+    def apply(self, df: DataFrame) -> DataFrame:
+        """Apply the filter to a DataFrame.
+
+        Args:
+            df: Input DataFrame.
+
+        Returns:
+            Filtered DataFrame (schema should be unchanged or a subset).
+        """
+        ...
+
+    def __call__(self, df: DataFrame) -> DataFrame:
+        """Apply the filter (callable interface).
+
+        Equivalent to calling :meth:`apply`.
+        """
+        return self.apply(df)
+
+
+class FilterPipeline(Filter):
+    """Chain of filters applied sequentially.
+
+    Example::
+
+        pipeline = FilterPipeline([
+            DropDuplicates(),
+            WhereFilter("score > 50"),
+            ColumnFilter(["id", "name"]),
+        ])
+        result = pipeline(df)
+    """
+
+    def __init__(self, filters: Sequence[Filter] | None = None) -> None:
+        self._filters: list[Filter] = list(filters or [])
+
+    def add(self, filter_: Filter) -> None:
+        """Append a filter to the end of the pipeline.
+
+        Args:
+            filter_: The filter to add.
+        """
+        self._filters.append(filter_)
+
+    def apply(self, df: DataFrame) -> DataFrame:
+        """Apply all filters in sequence.
+
+        Args:
+            df: Input DataFrame.
+
+        Returns:
+            DataFrame after all filters have been applied.
+        """
+        for f in self._filters:
+            df = f(df)
+        return df
+
+    def __len__(self) -> int:
+        return len(self._filters)
+
+    def __repr__(self) -> str:
+        return f"FilterPipeline({self._filters!r})"
+
+
+class DropDuplicates(Filter):
+    """Remove duplicate rows.
+
+    Args:
+        subset: Optional column names to consider for deduplication.
+            If None, all columns are used.
+    """
+
+    def __init__(self, subset: Sequence[str] | None = None) -> None:
+        self._subset = list(subset) if subset else None
+
+    def apply(self, df: DataFrame) -> DataFrame:
+        if self._subset:
+            return df.dropDuplicates(self._subset)
+        return df.dropDuplicates()
+
+    def __repr__(self) -> str:
+        return f"DropDuplicates(subset={self._subset!r})"
+
+
+class ColumnFilter(Filter):
+    """Select specific columns from a DataFrame.
+
+    Args:
+        columns: Column names to keep.
+
+    Raises:
+        ValueError: If columns is empty.
+    """
+
+    def __init__(self, columns: Sequence[str]) -> None:
+        if not columns:
+            raise ValueError("ColumnFilter requires at least one column")
+        self._columns = list(columns)
+
+    def apply(self, df: DataFrame) -> DataFrame:
+        return df.select(*self._columns)
+
+    def __repr__(self) -> str:
+        return f"ColumnFilter(columns={self._columns!r})"
+
+
+class WhereFilter(Filter):
+    """Apply a SQL WHERE condition.
+
+    Args:
+        condition: SQL condition string (e.g. ``"age > 18"``).
+    """
+
+    def __init__(self, condition: str) -> None:
+        self._condition = condition
+
+    def apply(self, df: DataFrame) -> DataFrame:
+        return df.where(self._condition)
+
+    def __repr__(self) -> str:
+        return f"WhereFilter({self._condition!r})"

databricks4py/io/__init__.py

@@ -0,0 +1,40 @@
+"""I/O utilities for Delta Lake, DBFS, and streaming."""
+
+from databricks4py.io.checkpoint import CheckpointInfo, CheckpointManager
+from databricks4py.io.dbfs import copy_from_remote, inject_dbutils_module, ls, mkdirs, mv, rm
+from databricks4py.io.delta import (
+    DeltaTable,
+    DeltaTableAppender,
+    DeltaTableOverwriter,
+    GeneratedColumn,
+    optimize_table,
+    vacuum_table,
+)
+from databricks4py.io.merge import MergeBuilder, MergeResult
+from databricks4py.io.streaming import StreamingTableReader, StreamingTriggerOptions
+
+__all__ = [
+    # Checkpoint
+    "CheckpointInfo",
+    "CheckpointManager",
+    # DBFS
+    "copy_from_remote",
+    "inject_dbutils_module",
+    "ls",
+    "mkdirs",
+    "mv",
+    "rm",
+    # Delta
+    "DeltaTable",
+    "DeltaTableAppender",
+    "DeltaTableOverwriter",
+    "GeneratedColumn",
+    "optimize_table",
+    "vacuum_table",
+    # Merge
+    "MergeBuilder",
+    "MergeResult",
+    # Streaming
+    "StreamingTableReader",
+    "StreamingTriggerOptions",
+]

databricks4py/io/checkpoint.py

@@ -0,0 +1,98 @@
+"""Streaming checkpoint lifecycle management."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+import shutil
+from dataclasses import dataclass
+from typing import Any
+
+from pyspark.sql import SparkSession
+
+__all__ = [
+    "CheckpointInfo",
+    "CheckpointManager",
+]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class CheckpointInfo:
+    """Immutable snapshot of a streaming checkpoint's state."""
+
+    path: str
+    last_batch_id: int | None
+    offsets: dict[str, Any] | None
+    size_bytes: int
+
+
+def _dir_size(path: str) -> int:
+    total = 0
+    for dirpath, _dirnames, filenames in os.walk(path):
+        for fname in filenames:
+            total += os.path.getsize(os.path.join(dirpath, fname))
+    return total
+
+
+def _sanitize(name: str) -> str:
+    return re.sub(r"[^a-zA-Z0-9]", "_", name)
+
+
+class CheckpointManager:
+    """Manages streaming checkpoint directories.
+
+    Args:
+        base_path: Root directory under which checkpoints are stored.
+        spark: Optional SparkSession (reserved for future DBFS/cloud support).
+    """
+
+    def __init__(
+        self,
+        base_path: str,
+        *,
+        spark: SparkSession | None = None,
+    ) -> None:
+        self._base_path = base_path
+        self._spark = spark
+
+    def path_for(self, source: str, sink: str) -> str:
+        """Generate a deterministic checkpoint path for a source/sink pair."""
+        return f"{self._base_path}/{_sanitize(source)}__{_sanitize(sink)}"
+
+    def exists(self, path: str) -> bool:
+        return os.path.isdir(path)
+
+    def reset(self, path: str) -> None:
+        """Delete a checkpoint directory. No-op if it doesn't exist."""
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            logger.info("Deleted checkpoint at %s", path)
+
+    def info(self, path: str) -> CheckpointInfo:
+        """Read checkpoint metadata from the offset log."""
+        offsets_dir = os.path.join(path, "offsets")
+        last_batch_id: int | None = None
+        offsets: dict | None = None
+
+        if os.path.isdir(offsets_dir):
+            batch_ids = []
+            for entry in os.listdir(offsets_dir):
+                if entry.isdigit():
+                    batch_ids.append(int(entry))
+
+            if batch_ids:
+                last_batch_id = max(batch_ids)
+                offset_file = os.path.join(offsets_dir, str(last_batch_id))
+                with open(offset_file) as f:
+                    offsets = json.load(f)
+
+        return CheckpointInfo(
+            path=path,
+            last_batch_id=last_batch_id,
+            offsets=offsets,
+            size_bytes=_dir_size(path),
+        )