databricks4py 0.2.0__py3-none-any.whl

databricks4py/secrets.py

@@ -0,0 +1,69 @@
+"""Databricks secret management via dbutils."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import pyspark.sql
+
+from databricks4py.spark_session import active_fallback
+
+__all__ = ["SecretFetcher", "inject_dbutils"]
+
+logger = logging.getLogger(__name__)
+
+
+class SecretFetcher:
+    """Fetch secrets from Databricks secret vaults via dbutils.
+
+    The ``dbutils`` module must be injected before use, either by calling
+    :func:`inject_dbutils` or by setting :attr:`dbutils` directly.
+
+    Example::
+
+        from databricks4py.secrets import SecretFetcher, inject_dbutils
+        import pyspark.dbutils  # only available on Databricks
+
+        inject_dbutils(pyspark.dbutils)
+        value = SecretFetcher.fetch_secret("my-scope", "api-key")
+    """
+
+    dbutils: Any = None
+
+    @staticmethod
+    def fetch_secret(
+        secret_scope: str,
+        secret_key: str,
+        *,
+        spark: pyspark.sql.SparkSession | None = None,
+    ) -> str:
+        """Fetch a secret from the Databricks secret vault.
+
+        Args:
+            secret_scope: The secret scope name.
+            secret_key: The secret key name.
+            spark: Optional SparkSession (used to create DBUtils instance).
+
+        Raises:
+            RuntimeError: If dbutils has not been injected.
+        """
+        if SecretFetcher.dbutils is None:
+            raise RuntimeError(
+                "SecretFetcher.dbutils has not been set. "
+                "Call inject_dbutils(pyspark.dbutils) first."
+            )
+
+        spark = active_fallback(spark)
+        _dbutils = SecretFetcher.dbutils.DBUtils(spark)
+        return _dbutils.secrets.get(scope=secret_scope, key=secret_key)
+
+
+def inject_dbutils(dbutils_module: Any) -> None:
+    """Inject the dbutils module for secret fetching.
+
+    Args:
+        dbutils_module: The ``pyspark.dbutils`` module (only available on Databricks).
+    """
+    SecretFetcher.dbutils = dbutils_module
+    logger.debug("Injected dbutils module: %s", dbutils_module)

databricks4py/spark_session.py

@@ -0,0 +1,68 @@
+"""SparkSession management utilities."""
+
+from __future__ import annotations
+
+import pyspark.sql
+
+__all__ = ["get_active", "active_fallback", "get_or_create_local_session"]
+
+
+def get_active() -> pyspark.sql.SparkSession:
+    """Get the currently active SparkSession.
+
+    Raises:
+        RuntimeError: If no active SparkSession exists.
+    """
+    spark = pyspark.sql.SparkSession.getActiveSession()
+    if spark is None:
+        raise RuntimeError(
+            "No active SparkSession found. "
+            "Create one with SparkSession.builder.getOrCreate() or "
+            "use get_or_create_local_session()."
+        )
+    return spark
+
+
+def active_fallback(spark: pyspark.sql.SparkSession | None = None) -> pyspark.sql.SparkSession:
+    """Return the given SparkSession, or fall back to the active one.
+
+    This is the standard pattern for library functions that accept an
+    optional ``spark`` parameter::
+
+        def my_function(data, *, spark=None):
+            spark = active_fallback(spark)
+            ...
+
+    Args:
+        spark: An explicit SparkSession, or None to use the active session.
+
+    Raises:
+        RuntimeError: If spark is None and no active session exists.
+    """
+    return spark if spark is not None else get_active()
+
+
+def get_or_create_local_session() -> pyspark.sql.SparkSession:
+    """Create a local SparkSession configured for Delta Lake.
+
+    Suitable for local development and testing. Configures:
+    - ``local[*]`` master
+    - Delta Lake SQL extensions and catalog
+    - Local Derby metastore
+
+    Returns:
+        A SparkSession with Delta Lake support.
+    """
+    from delta import configure_spark_with_delta_pip
+
+    builder = (
+        pyspark.sql.SparkSession.builder.master("local[*]")
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config(
+            "spark.sql.catalog.spark_catalog",
+            "org.apache.spark.sql.delta.catalog.DeltaCatalog",
+        )
+        .config("spark.sql.warehouse.dir", "spark-warehouse")
+    )
+
+    return configure_spark_with_delta_pip(builder).getOrCreate()

databricks4py/testing/__init__.py

@@ -0,0 +1,35 @@
+"""Test utilities and fixtures for databricks4py.
+
+Provides pytest fixtures and mock objects for testing Spark and
+Databricks applications locally.
+
+Usage in your ``conftest.py``::
+
+    from databricks4py.testing.fixtures import *  # noqa: F401,F403
+"""
+
+from databricks4py.testing.assertions import assert_frame_equal, assert_schema_equal
+from databricks4py.testing.builders import DataFrameBuilder
+from databricks4py.testing.fixtures import (
+    clear_env,
+    df_builder,
+    spark_session,
+    spark_session_function,
+    temp_delta,
+)
+from databricks4py.testing.mocks import MockDBUtils, MockDBUtilsModule
+from databricks4py.testing.temp_table import TempDeltaTable
+
+__all__ = [
+    "DataFrameBuilder",
+    "MockDBUtils",
+    "MockDBUtilsModule",
+    "TempDeltaTable",
+    "assert_frame_equal",
+    "assert_schema_equal",
+    "clear_env",
+    "df_builder",
+    "spark_session",
+    "spark_session_function",
+    "temp_delta",
+]

databricks4py/testing/assertions.py

@@ -0,0 +1,111 @@
+"""DataFrame and schema assertion helpers for testing."""
+
+from __future__ import annotations
+
+from pyspark.sql import DataFrame
+from pyspark.sql.types import StructType
+
+__all__ = ["assert_frame_equal", "assert_schema_equal"]
+
+
+def assert_schema_equal(
+    actual: StructType,
+    expected: StructType,
+    *,
+    check_nullable: bool = False,
+) -> None:
+    """Compare two StructTypes field-by-field with a clear diff on mismatch.
+
+    Args:
+        actual: The actual schema.
+        expected: The expected schema.
+        check_nullable: Whether to compare nullable flags (default False).
+
+    Raises:
+        AssertionError: If the schemas differ.
+    """
+    actual_fields = actual.fields
+    expected_fields = expected.fields
+
+    if len(actual_fields) != len(expected_fields):
+        raise AssertionError(
+            f"Schema length mismatch: actual has {len(actual_fields)} fields, "
+            f"expected {len(expected_fields)}"
+        )
+
+    for i, (a, e) in enumerate(zip(actual_fields, expected_fields, strict=True)):
+        if a.name != e.name:
+            raise AssertionError(
+                f"Field {i}: name mismatch — actual={a.name!r}, expected={e.name!r}"
+            )
+        if a.dataType != e.dataType:
+            raise AssertionError(
+                f"Field {a.name!r}: type mismatch — actual={a.dataType}, expected={e.dataType}"
+            )
+        if check_nullable and a.nullable != e.nullable:
+            raise AssertionError(
+                f"Field {a.name!r}: nullable mismatch — actual={a.nullable}, expected={e.nullable}"
+            )
+
+
+def assert_frame_equal(
+    actual: DataFrame,
+    expected: DataFrame,
+    *,
+    check_order: bool = False,
+    check_schema: bool = True,
+    check_nullable: bool = False,
+) -> None:
+    """Assert two DataFrames are equal.
+
+    Uses Spark 3.5's ``assertDataFrameEqual`` when available, otherwise
+    falls back to a manual row-by-row comparison.
+
+    Args:
+        actual: The actual DataFrame.
+        expected: The expected DataFrame.
+        check_order: Whether row order matters (default False).
+        check_schema: Whether to compare schemas (default True).
+        check_nullable: Whether to compare nullable flags (default False).
+
+    Raises:
+        AssertionError: If the DataFrames differ.
+    """
+    if check_schema:
+        assert_schema_equal(actual.schema, expected.schema, check_nullable=check_nullable)
+
+    try:
+        from pyspark.testing.utils import assertDataFrameEqual
+
+        assertDataFrameEqual(actual, expected, checkRowOrder=check_order)
+    except ImportError:
+        _manual_frame_compare(actual, expected, check_order=check_order)
+
+
+def _manual_frame_compare(
+    actual: DataFrame,
+    expected: DataFrame,
+    *,
+    check_order: bool,
+) -> None:
+    actual_rows = actual.collect()
+    expected_rows = expected.collect()
+
+    if len(actual_rows) != len(expected_rows):
+        raise AssertionError(
+            f"Row count mismatch: actual={len(actual_rows)}, expected={len(expected_rows)}"
+        )
+
+    if check_order:
+        for i, (a, e) in enumerate(zip(actual_rows, expected_rows, strict=True)):
+            if a != e:
+                raise AssertionError(f"Row {i} mismatch:\n  actual:   {a}\n  expected: {e}")
+    else:
+        actual_sorted = sorted(actual_rows, key=str)
+        expected_sorted = sorted(expected_rows, key=str)
+        if actual_sorted != expected_sorted:
+            raise AssertionError(
+                f"DataFrames differ (ignoring order):\n"
+                f"  actual:   {actual_sorted}\n"
+                f"  expected: {expected_sorted}"
+            )

databricks4py/testing/builders.py

@@ -0,0 +1,127 @@
+"""Fluent builder for test DataFrames."""
+
+from __future__ import annotations
+
+import random
+from typing import Any
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.types import (
+    BooleanType,
+    DataType,
+    DateType,
+    DoubleType,
+    FloatType,
+    IntegerType,
+    LongType,
+    ShortType,
+    StringType,
+    StructField,
+    StructType,
+    TimestampType,
+)
+
+__all__ = ["DataFrameBuilder"]
+
+_TYPE_MAP: dict[str, DataType] = {
+    "string": StringType(),
+    "int": IntegerType(),
+    "integer": IntegerType(),
+    "long": LongType(),
+    "bigint": LongType(),
+    "short": ShortType(),
+    "smallint": ShortType(),
+    "float": FloatType(),
+    "double": DoubleType(),
+    "boolean": BooleanType(),
+    "date": DateType(),
+    "timestamp": TimestampType(),
+}
+
+
+def _resolve_type(type_str: str) -> DataType:
+    resolved = _TYPE_MAP.get(type_str.lower())
+    if resolved is not None:
+        return resolved
+    from pyspark.sql.types import _parse_datatype_string
+
+    return _parse_datatype_string(type_str)
+
+
+class DataFrameBuilder:
+    """Fluent builder for constructing test DataFrames.
+
+    Example::
+
+        df = (
+            DataFrameBuilder(spark)
+            .with_columns({"id": "int", "name": "string"})
+            .with_rows((1, "alice"), (2, "bob"))
+            .build()
+        )
+    """
+
+    def __init__(self, spark: SparkSession) -> None:
+        self._spark = spark
+        self._schema: StructType | None = None
+        self._rows: list[tuple] = []
+
+    def with_columns(self, schema: dict[str, str]) -> DataFrameBuilder:
+        """Define columns from a dict of ``{name: type_string}``."""
+        fields = [StructField(name, _resolve_type(t)) for name, t in schema.items()]
+        self._schema = StructType(fields)
+        return self
+
+    def with_schema(self, schema: StructType) -> DataFrameBuilder:
+        """Define columns from a StructType directly."""
+        self._schema = schema
+        return self
+
+    def with_rows(self, *rows: tuple) -> DataFrameBuilder:
+        """Add explicit data rows."""
+        self._rows.extend(rows)
+        return self
+
+    def with_sequential(self, column: str, start: int = 1, count: int = 10) -> DataFrameBuilder:
+        """Generate sequential integer rows for a single column.
+
+        If rows already exist, this replaces them.
+        """
+        if self._schema is None:
+            self._schema = StructType([StructField(column, IntegerType())])
+        self._rows = [(i,) for i in range(start, start + count)]
+        return self
+
+    def with_nulls(
+        self, column: str, frequency: float = 0.1, *, seed: int | None = None
+    ) -> DataFrameBuilder:
+        """Inject nulls into a column at the given frequency.
+
+        Must be called after rows are populated.
+        """
+        if self._schema is None:
+            raise ValueError("Schema must be defined before injecting nulls")
+
+        col_idx = next((i for i, f in enumerate(self._schema.fields) if f.name == column), None)
+        if col_idx is None:
+            raise ValueError(f"Column {column!r} not found in schema")
+
+        rng = random.Random(seed)
+        new_rows: list[tuple[Any, ...]] = []
+        for row in self._rows:
+            row_list = list(row)
+            if rng.random() < frequency:
+                row_list[col_idx] = None
+            new_rows.append(tuple(row_list))
+        self._rows = new_rows
+        return self
+
+    def build(self) -> DataFrame:
+        """Construct the DataFrame.
+
+        Raises:
+            ValueError: If no schema has been defined.
+        """
+        if self._schema is None:
+            raise ValueError("No schema defined — call with_columns() or with_schema() first")
+        return self._spark.createDataFrame(self._rows, schema=self._schema)

databricks4py/testing/fixtures.py

@@ -0,0 +1,134 @@
+"""Pytest fixtures for Spark and Delta Lake testing."""
+
+from __future__ import annotations
+
+import os
+import shutil
+from collections.abc import Callable, Generator
+from typing import Any
+
+import pyspark.sql
+import pytest
+
+from databricks4py.testing.builders import DataFrameBuilder
+from databricks4py.testing.temp_table import TempDeltaTable
+
+__all__ = [
+    "clear_env",
+    "df_builder",
+    "spark_session",
+    "spark_session_function",
+    "temp_delta",
+]
+
+
+@pytest.fixture(scope="session")
+def spark_session(
+    tmp_path_factory: pytest.TempPathFactory,
+) -> Generator[pyspark.sql.SparkSession, None, None]:
+    """Session-scoped SparkSession with Delta Lake support.
+
+    Creates a single SparkSession for the entire test session to avoid
+    the overhead of starting/stopping the JVM repeatedly. Uses a temporary
+    directory for the Derby metastore and Spark warehouse.
+    """
+    from delta import configure_spark_with_delta_pip
+
+    warehouse_dir = str(tmp_path_factory.mktemp("spark-warehouse"))
+    derby_dir = str(tmp_path_factory.mktemp("derby"))
+
+    builder = (
+        pyspark.sql.SparkSession.builder.master("local[*]")
+        .appName("databricks4py-tests")
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config(
+            "spark.sql.catalog.spark_catalog",
+            "org.apache.spark.sql.delta.catalog.DeltaCatalog",
+        )
+        .config("spark.sql.warehouse.dir", warehouse_dir)
+        .config("javax.jdo.option.ConnectionURL", f"jdbc:derby:{derby_dir}/metastore;create=true")
+        .config("spark.driver.extraJavaOptions", f"-Dderby.system.home={derby_dir}")
+        .config("spark.ui.enabled", "false")
+        .config("spark.sql.shuffle.partitions", "2")
+    )
+
+    spark = configure_spark_with_delta_pip(builder).getOrCreate()
+    yield spark
+    spark.stop()
+
+
+@pytest.fixture()
+def spark_session_function(
+    spark_session: pyspark.sql.SparkSession,
+) -> Generator[pyspark.sql.SparkSession, None, None]:
+    """Function-scoped SparkSession that cleans up between tests.
+
+    Reuses the session-scoped SparkSession but clears the catalog
+    and cache after each test to ensure isolation.
+    """
+    yield spark_session
+
+    # Clean up tables
+    for db in spark_session.catalog.listDatabases():
+        for table in spark_session.catalog.listTables(db.name):
+            spark_session.sql(f"DROP TABLE IF EXISTS {db.name}.{table.name}")
+
+    spark_session.catalog.clearCache()
+
+    # Clean up any leftover warehouse files
+    warehouse = spark_session.conf.get("spark.sql.warehouse.dir")
+    if warehouse and os.path.exists(warehouse):
+        for item in os.listdir(warehouse):
+            item_path = os.path.join(warehouse, item)
+            if os.path.isdir(item_path):
+                shutil.rmtree(item_path, ignore_errors=True)
+
+
+@pytest.fixture(autouse=True)
+def clear_env() -> Generator[None, None, None]:
+    """Auto-use fixture that restores environment variables after each test."""
+    original = os.environ.copy()
+    yield
+    os.environ.clear()
+    os.environ.update(original)
+
+
+@pytest.fixture()
+def df_builder(spark_session: pyspark.sql.SparkSession) -> DataFrameBuilder:
+    """Return a DataFrameBuilder bound to the session-scoped SparkSession."""
+    return DataFrameBuilder(spark_session)
+
+
+@pytest.fixture()
+def temp_delta(
+    spark_session_function: pyspark.sql.SparkSession,
+) -> Generator[Callable[..., TempDeltaTable]]:
+    """Factory fixture that creates TempDeltaTables and cleans up on exit.
+
+    Usage::
+
+        def test_something(temp_delta):
+            with temp_delta(schema={"id": "int"}, data=[(1,)]) as table:
+                assert table.dataframe().count() == 1
+    """
+    tables: list[TempDeltaTable] = []
+
+    def _factory(
+        *,
+        table_name: str | None = None,
+        schema: dict[str, str] | None = None,
+        data: list[tuple[Any, ...]] | None = None,
+    ) -> TempDeltaTable:
+        t = TempDeltaTable(
+            spark_session_function,
+            table_name=table_name,
+            schema=schema,
+            data=data,
+        )
+        tables.append(t)
+        return t
+
+    yield _factory
+
+    for t in tables:
+        spark_session_function.sql(f"DROP TABLE IF EXISTS {t.table_name}")

databricks4py/testing/mocks.py

@@ -0,0 +1,106 @@
+"""Mock objects for Databricks dbutils in local testing."""
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = ["MockDBUtils", "MockDBUtilsModule"]
+
+
+class _MockSecrets:
+    """Mock for dbutils.secrets."""
+
+    def __init__(self) -> None:
+        self._secrets: dict[tuple[str, str], str] = {}
+
+    def get(self, scope: str, key: str) -> str:
+        """Get a secret value.
+
+        Args:
+            scope: The secret scope.
+            key: The secret key.
+
+        Raises:
+            KeyError: If the secret is not found.
+        """
+        try:
+            return self._secrets[(scope, key)]
+        except KeyError:
+            raise KeyError(f"Secret not found: scope={scope!r}, key={key!r}") from None
+
+    def put(self, scope: str, key: str, value: str) -> None:
+        """Store a secret value (for test setup)."""
+        self._secrets[(scope, key)] = value
+
+
+class _MockFS:
+    """Mock for dbutils.fs."""
+
+    def __init__(self) -> None:
+        self._copies: list[tuple[str, str, bool]] = []
+        self._moves: list[tuple[str, str, bool]] = []
+        self._removes: list[tuple[str, bool]] = []
+        self._mkdirs: list[str] = []
+        self._ls_results: dict[str, list[Any]] = {}
+
+    def cp(self, source: str, dest: str, recurse: bool = False) -> bool:
+        """Record a copy operation."""
+        self._copies.append((source, dest, recurse))
+        return True
+
+    def mv(self, source: str, dest: str, recurse: bool = False) -> bool:
+        """Record a move operation."""
+        self._moves.append((source, dest, recurse))
+        return True
+
+    def rm(self, path: str, recurse: bool = False) -> bool:
+        """Record a remove operation."""
+        self._removes.append((path, recurse))
+        return True
+
+    def mkdirs(self, path: str) -> bool:
+        """Record a mkdirs operation."""
+        self._mkdirs.append(path)
+        return True
+
+    def ls(self, path: str) -> list[Any]:
+        """List files at path."""
+        return self._ls_results.get(path, [])
+
+
+class MockDBUtils:
+    """Mock Databricks DBUtils for local testing.
+
+    Provides mock implementations of ``secrets`` and ``fs`` modules.
+
+    Example::
+
+        mock = MockDBUtils()
+        mock.secrets.put("my-scope", "api-key", "secret-value")
+        assert mock.secrets.get("my-scope", "api-key") == "secret-value"
+    """
+
+    def __init__(self) -> None:
+        self.secrets = _MockSecrets()
+        self.fs = _MockFS()
+
+
+class MockDBUtilsModule:
+    """Mock for the ``pyspark.dbutils`` module.
+
+    Mimics the Databricks runtime's ``pyspark.dbutils`` module which
+    provides a ``DBUtils`` class that accepts a SparkSession.
+
+    Example::
+
+        mock_module = MockDBUtilsModule()
+        dbutils = mock_module.DBUtils(spark)
+        dbutils.secrets.get("scope", "key")
+    """
+
+    def __init__(self, mock_dbutils: MockDBUtils | None = None) -> None:
+        self._mock_dbutils = mock_dbutils or MockDBUtils()
+
+    def DBUtils(self, spark: Any = None) -> MockDBUtils:  # noqa: N802
+        """Create a DBUtils instance (ignores spark argument)."""
+        return self._mock_dbutils