streamcase 0.1.0a1__py3-none-any.whl

streamcase/__init__.py

@@ -0,0 +1,23 @@
+"""Deterministic testing for Apache Spark Structured Streaming."""
+
+from streamcase._version import __version__
+from streamcase.actions import Batch, Restart, batch, restart
+from streamcase.assertions import assert_batch_count, assert_rows_equal, assert_unique_keys
+from streamcase.results import CapturedBatch, ScenarioResult
+from streamcase.scenario import Action, Scenario, scenario
+
+__all__ = [
+    "Action",
+    "Batch",
+    "CapturedBatch",
+    "Restart",
+    "Scenario",
+    "ScenarioResult",
+    "__version__",
+    "assert_batch_count",
+    "assert_rows_equal",
+    "assert_unique_keys",
+    "batch",
+    "restart",
+    "scenario",
+]

streamcase/_cleanup.py

@@ -0,0 +1,33 @@
+"""Preserve execution failures while attempting runner-owned cleanup."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+
+
+def _record_cleanup_failure(primary: BaseException, resource: str, error: BaseException) -> None:
+    failures = primary.__dict__.get("_streamcase_cleanup_failures", ())
+    primary.__dict__["_streamcase_cleanup_failures"] = (*failures, (resource, error))
+
+    note = f"Streamcase cleanup failed for {resource}: {type(error).__name__}: {error}"
+    add_note = getattr(primary, "add_note", None)
+    if callable(add_note):
+        add_note(note)
+    else:
+        primary.args = (*primary.args, note)
+
+
+@contextmanager
+def _cleanup_on_exit(resource: str, cleanup: Callable[[], None]) -> Iterator[None]:
+    """Run cleanup and keep a primary failure when cleanup also fails."""
+    try:
+        yield
+    except BaseException as primary:
+        try:
+            cleanup()
+        except BaseException as error:
+            _record_cleanup_failure(primary, resource, error)
+        raise
+    else:
+        cleanup()

streamcase/_diagnostics.py

@@ -0,0 +1,108 @@
+"""Deterministic internal formatting for assertion failures."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import TypeAlias, cast
+
+_MAX_DIAGNOSTIC_ITEMS = 5
+
+_Location: TypeAlias = tuple[int, int]
+_DuplicateKey: TypeAlias = tuple[
+    tuple[str, ...],
+    tuple[object, ...],
+    Sequence[_Location],
+]
+
+
+def _render_value(value: object) -> str:
+    if isinstance(value, Mapping):
+        mapping = cast(Mapping[str, object], value)
+        rendered_mapping_items = (
+            f"{key!r}: {_render_value(mapping[key])}" for key in sorted(mapping)
+        )
+        return "{" + ", ".join(rendered_mapping_items) + "}"
+
+    if isinstance(value, tuple):
+        rendered_tuple_items = ", ".join(_render_value(item) for item in value)
+        trailing_comma = "," if len(value) == 1 else ""
+        return f"({rendered_tuple_items}{trailing_comma})"
+
+    return repr(value)
+
+
+def _format_item_section(title: str, items: Sequence[str]) -> list[str]:
+    lines = [f"{title} ({len(items)}):"]
+    visible_items = items[:_MAX_DIAGNOSTIC_ITEMS]
+    lines.extend(f"  - {item}" for item in visible_items)
+
+    omitted_count = len(items) - len(visible_items)
+    if omitted_count:
+        lines.append(f"  ... {omitted_count} more item(s) omitted")
+    return lines
+
+
+def _format_batch_count_mismatch(expected_count: int, actual_count: int) -> str:
+    return "\n".join(
+        (
+            "Captured batch count differs.",
+            f"Expected: {expected_count} batch(es).",
+            f"Actual: {actual_count} batch(es).",
+        ),
+    )
+
+
+def _format_rows_mismatch(
+    expected_count: int,
+    actual_count: int,
+    missing: Sequence[Mapping[str, object]],
+    unexpected: Sequence[Mapping[str, object]],
+) -> str:
+    lines = [
+        "Rows differ.",
+        f"Expected: {expected_count} row(s).",
+        f"Actual: {actual_count} row(s).",
+    ]
+    if missing:
+        rendered_missing = [_render_value(row) for row in missing]
+        lines.extend(_format_item_section("Missing rows", rendered_missing))
+    if unexpected:
+        rendered_unexpected = [_render_value(row) for row in unexpected]
+        lines.extend(_format_item_section("Unexpected rows", rendered_unexpected))
+    return "\n".join(lines)
+
+
+def _format_missing_key_fields(
+    batch_id: int,
+    row_index: int,
+    missing_fields: Sequence[str],
+) -> str:
+    items = [f"batch {batch_id} row {row_index}: {field!r}" for field in missing_fields]
+    lines = ["Unique-key assertion failed."]
+    lines.extend(_format_item_section("Missing key fields", items))
+    return "\n".join(lines)
+
+
+def _format_duplicate_keys(duplicates: Sequence[_DuplicateKey]) -> str:
+    items: list[str] = []
+    for fields, key, locations in duplicates:
+        rendered_key = ", ".join(
+            f"{field}={_render_value(value)}" for field, value in zip(fields, key, strict=True)
+        )
+
+        visible_locations = locations[:_MAX_DIAGNOSTIC_ITEMS]
+        rendered_locations = ", ".join(
+            f"batch {batch_id} row {row_index}" for batch_id, row_index in visible_locations
+        )
+        omitted_count = len(locations) - len(visible_locations)
+        if omitted_count:
+            rendered_locations += f", ... {omitted_count} more occurrence(s) omitted"
+
+        items.append(
+            f"Duplicate key ({rendered_key}) occurred {len(locations)} times at "
+            f"{rendered_locations}."
+        )
+
+    lines = ["Unique-key assertion failed."]
+    lines.extend(_format_item_section("Duplicate keys", items))
+    return "\n".join(lines)

streamcase/_directories.py

@@ -0,0 +1,107 @@
+"""Private isolated filesystem layouts for scenario runs."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import tempfile
+from dataclasses import dataclass, field
+from pathlib import Path
+
+_RUN_DIRECTORY_PREFIX = "streamcase-run-"
+
+
+def _resolve_base_directory(base_dir: str | os.PathLike[str]) -> Path:
+    try:
+        raw_path = os.fspath(base_dir)
+    except TypeError as error:
+        raise TypeError("Runner base directory must be a string or path-like object.") from error
+
+    if not isinstance(raw_path, str):
+        raise TypeError("Runner base directory must resolve to a string path.")
+    if not raw_path:
+        raise ValueError("Runner base directory must not be empty.")
+
+    path = Path(raw_path)
+    if not path.exists():
+        raise FileNotFoundError(f"Runner base directory does not exist: {path}")
+    if not path.is_dir():
+        raise NotADirectoryError(f"Runner base path is not a directory: {path}")
+    return path.resolve(strict=True)
+
+
+def _require_direct_child(path: Path, parent: Path, description: str) -> None:
+    if path.parent != parent:
+        raise RuntimeError(f"Generated {description} must be a direct child of {parent}: {path}")
+
+
+@dataclass(frozen=True, slots=True)
+class RunDirectories:
+    """One private, immutable set of runner-owned directories."""
+
+    root: Path
+    input_dir: Path
+    checkpoint_dir: Path
+    temporary_dir: Path
+    retain_artifacts: bool
+    _base_dir: Path = field(repr=False, compare=False)
+
+    def cleanup(self) -> None:
+        """Remove the generated run root unless artifact retention is enabled."""
+        if self.retain_artifacts:
+            return
+        if not self.root.exists():
+            return
+
+        resolved_root = self.root.resolve(strict=True)
+        _require_direct_child(resolved_root, self._base_dir, "run root")
+        shutil.rmtree(resolved_root)
+
+
+def create_run_directories(
+    *,
+    base_dir: str | os.PathLike[str] | None = None,
+    retain_artifacts: bool = False,
+) -> RunDirectories:
+    """Create a unique, isolated directory layout for one scenario run."""
+    if not isinstance(retain_artifacts, bool):
+        raise TypeError("retain_artifacts must be a bool.")
+    if retain_artifacts and base_dir is None:
+        raise ValueError("retain_artifacts=True requires an explicit base directory.")
+
+    if base_dir is None:
+        root = Path(tempfile.mkdtemp(prefix=_RUN_DIRECTORY_PREFIX)).resolve(strict=True)
+        resolved_base = root.parent.resolve(strict=True)
+    else:
+        resolved_base = _resolve_base_directory(base_dir)
+        root = Path(tempfile.mkdtemp(prefix=_RUN_DIRECTORY_PREFIX, dir=resolved_base)).resolve(
+            strict=True
+        )
+
+    _require_direct_child(root, resolved_base, "run root")
+
+    input_dir = root / "input"
+    checkpoint_dir = root / "checkpoint"
+    temporary_dir = root / "temporary"
+    try:
+        input_dir.mkdir()
+        input_dir = input_dir.resolve(strict=True)
+        _require_direct_child(input_dir, root, "input directory")
+        checkpoint_dir.mkdir()
+        checkpoint_dir = checkpoint_dir.resolve(strict=True)
+        _require_direct_child(checkpoint_dir, root, "checkpoint directory")
+        temporary_dir.mkdir()
+        temporary_dir = temporary_dir.resolve(strict=True)
+        _require_direct_child(temporary_dir, root, "temporary directory")
+    except BaseException:
+        shutil.rmtree(root)
+        raise
+
+    return RunDirectories(
+        root=root,
+        input_dir=input_dir,
+        checkpoint_dir=checkpoint_dir,
+        temporary_dir=temporary_dir,
+        retain_artifacts=retain_artifacts,
+        _base_dir=resolved_base,
+    )

streamcase/_input_files.py

@@ -0,0 +1,59 @@
+"""Private atomic input-file publication for scenario batches."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from streamcase._directories import RunDirectories, _require_direct_child
+from streamcase._serialization import _encode_batch_json_lines
+from streamcase.actions import Batch
+
+_BATCH_INDEX_WIDTH = 20
+
+
+def _batch_file_name(index: int) -> str:
+    return f"batch-{index:0{_BATCH_INDEX_WIDTH}d}.json"
+
+
+@dataclass(slots=True)
+class AtomicBatchWriter:
+    """Publish encoded batches sequentially into one isolated input directory."""
+
+    directories: RunDirectories
+    _next_index: int = field(default=0, init=False, repr=False)
+
+    def publish(self, action: Batch) -> Path:
+        """Atomically publish one batch and return its final input path."""
+        encoded = _encode_batch_json_lines(action)
+        file_name = _batch_file_name(self._next_index)
+        destination = self.directories.input_dir / file_name
+        temporary = self.directories.input_dir / f".{file_name}.tmp"
+
+        _require_direct_child(destination, self.directories.input_dir, "batch input file")
+        _require_direct_child(temporary, self.directories.input_dir, "temporary batch file")
+
+        if destination.exists():
+            raise FileExistsError(f"Batch input destination already exists: {destination}")
+        if temporary.exists():
+            raise FileExistsError(f"Temporary batch input file already exists: {temporary}")
+
+        temporary_created = False
+        try:
+            with temporary.open("x", encoding="utf-8", newline="\n") as handle:
+                temporary_created = True
+                handle.write(encoded)
+                handle.flush()
+                os.fsync(handle.fileno())
+
+            if destination.exists():
+                raise FileExistsError(f"Batch input destination already exists: {destination}")
+            temporary.rename(destination)
+        except BaseException:
+            if temporary_created:
+                temporary.unlink(missing_ok=True)
+            raise
+
+        self._next_index += 1
+        return destination

streamcase/_serialization.py

@@ -0,0 +1,31 @@
+"""Internal serialization helpers for scenario actions."""
+
+from __future__ import annotations
+
+import json
+
+from streamcase.actions import Batch
+
+
+def _encode_batch_json_lines(action: Batch) -> str:
+    """Encode a batch as deterministic, newline-terminated JSON Lines text."""
+    lines: list[str] = []
+    for row_index, row in enumerate(action.rows):
+        try:
+            encoded = json.dumps(
+                dict(row),
+                allow_nan=False,
+                ensure_ascii=False,
+                separators=(",", ":"),
+                sort_keys=True,
+            )
+        except TypeError as error:
+            message = f"Batch row at index {row_index} could not be encoded as JSON: {error}"
+            raise TypeError(message) from error
+        except ValueError as error:
+            message = f"Batch row at index {row_index} could not be encoded as JSON: {error}"
+            raise ValueError(message) from error
+
+        lines.append(encoded)
+
+    return "\n".join(lines) + "\n"

streamcase/_spark_capture.py

@@ -0,0 +1,33 @@
+"""Private driver-side capture for Spark ``foreachBatch`` callbacks."""
+
+from __future__ import annotations
+
+from _thread import LockType
+from dataclasses import dataclass, field
+from threading import Lock
+from typing import TYPE_CHECKING
+
+from streamcase.results import CapturedBatch
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame
+
+
+@dataclass(slots=True)
+class _BatchCapture:
+    """Collect complete immutable output batches in callback order."""
+
+    _batches: list[CapturedBatch] = field(default_factory=list, init=False, repr=False)
+    _lock: LockType = field(default_factory=Lock, init=False, repr=False)
+
+    def callback(self, dataframe: DataFrame, batch_id: int) -> None:
+        """Collect and freeze one Spark output micro-batch on the driver."""
+        with self._lock:
+            rows = (row.asDict(recursive=True) for row in dataframe.collect())
+            captured = CapturedBatch(batch_id, rows)
+            self._batches.append(captured)
+
+    def snapshot(self) -> tuple[CapturedBatch, ...]:
+        """Return an immutable snapshot after any active callback completes."""
+        with self._lock:
+            return tuple(self._batches)

streamcase/_spark_execution.py

@@ -0,0 +1,95 @@
+"""Private execution of Spark streaming scenario actions."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import TYPE_CHECKING
+
+from streamcase._cleanup import _cleanup_on_exit, _record_cleanup_failure
+from streamcase._directories import RunDirectories
+from streamcase._input_files import AtomicBatchWriter
+from streamcase._spark_capture import _BatchCapture
+from streamcase._spark_lifecycle import _QueryLifecycle
+from streamcase.actions import Batch, Restart
+from streamcase.results import ScenarioResult
+from streamcase.scenario import Scenario
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame
+
+
+def _process_batch_action(
+    index: int,
+    action: Batch,
+    input_writer: AtomicBatchWriter,
+    lifecycle: _QueryLifecycle,
+) -> None:
+    try:
+        query = lifecycle.require_active()
+        input_writer.publish(action)
+        query.processAllAvailable()
+        lifecycle.require_active()
+    except Exception as error:
+        raise RuntimeError(f"Batch action at index {index} failed: {error}") from error
+
+
+def _process_restart_action(
+    index: int,
+    rebuild_stream: Callable[[], DataFrame],
+    lifecycle: _QueryLifecycle,
+) -> None:
+    try:
+        lifecycle.require_active()
+        lifecycle.stop()
+    except Exception as error:
+        raise _restart_failure(index, "stopping the active query", error) from error
+
+    try:
+        replacement_stream = rebuild_stream()
+    except Exception as error:
+        raise _restart_failure(index, "rebuilding the stream", error) from error
+
+    try:
+        lifecycle.start(replacement_stream, action_index=index + 1)
+        lifecycle.require_active()
+    except Exception as error:
+        raise _restart_failure(index, "starting the replacement query", error) from error
+
+
+def _restart_failure(index: int, transition: str, error: Exception) -> RuntimeError:
+    failure = RuntimeError(f"Restart action at index {index} failed while {transition}: {error}")
+    for resource, cleanup_error in getattr(error, "_streamcase_cleanup_failures", ()):
+        _record_cleanup_failure(failure, resource, cleanup_error)
+    return failure
+
+
+def _execute_batches(
+    stream: DataFrame,
+    scenario: Scenario,
+    directories: RunDirectories,
+    capture: _BatchCapture,
+    *,
+    output_mode: str = "append",
+    query_options: Mapping[str, str] | None = None,
+    rebuild_stream: Callable[[], DataFrame] | None = None,
+) -> ScenarioResult:
+    """Execute each action and return an immutable output snapshot."""
+    if rebuild_stream is None and any(isinstance(action, Restart) for action in scenario.actions):
+        raise ValueError("Restart actions require a stream rebuild callback.")
+
+    lifecycle = _QueryLifecycle(
+        directories,
+        capture,
+        output_mode=output_mode,
+        query_options=query_options,
+    )
+    with _cleanup_on_exit("streaming query", lifecycle.stop):
+        lifecycle.start(stream, action_index=0)
+        input_writer = AtomicBatchWriter(directories)
+        for index, action in enumerate(scenario.actions):
+            if isinstance(action, Batch):
+                _process_batch_action(index, action, input_writer, lifecycle)
+            else:
+                assert rebuild_stream is not None
+                _process_restart_action(index, rebuild_stream, lifecycle)
+    return ScenarioResult(capture.snapshot())

streamcase/_spark_lifecycle.py

@@ -0,0 +1,81 @@
+"""Private ownership and transitions for one Spark streaming query at a time."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING
+
+from streamcase._cleanup import _cleanup_on_exit
+from streamcase._directories import RunDirectories
+from streamcase._spark_capture import _BatchCapture
+from streamcase._spark_query_options import _prepare_query_configuration
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame
+    from pyspark.sql.streaming import StreamingQuery
+
+
+def _stop_named_query(stream: DataFrame, name: str) -> None:
+    for active_query in stream.sparkSession.streams.active:
+        if active_query.name == name:
+            active_query.stop()
+
+
+class _QueryLifecycle:
+    """Keep runner-owned query transitions and writer settings in one place."""
+
+    def __init__(
+        self,
+        directories: RunDirectories,
+        capture: _BatchCapture,
+        *,
+        output_mode: str = "append",
+        query_options: Mapping[str, str] | None = None,
+    ) -> None:
+        self._directories = directories
+        self._capture = capture
+        self._output_mode, self._query_options = _prepare_query_configuration(
+            output_mode, query_options
+        )
+        self._query: StreamingQuery | None = None
+        self._pending_start_stream: DataFrame | None = None
+
+    def require_active(self) -> StreamingQuery:
+        """Return the owned active query, rejecting missing or terminated state."""
+        if self._query is None:
+            raise RuntimeError("No streaming query is owned by this run.")
+        if not self._query.isActive:
+            raise RuntimeError("The streaming query terminated.")
+        return self._query
+
+    def start(self, stream: DataFrame, *, action_index: int) -> None:
+        """Start a query only when the previous one was successfully stopped."""
+        if self._query is not None or self._pending_start_stream is not None:
+            raise RuntimeError("Stop the owned streaming query before starting another.")
+
+        writer = stream.writeStream
+        for name, value in self._query_options.items():
+            writer = writer.option(name, value)
+        writer = (
+            writer.foreachBatch(self._capture.callback)
+            .outputMode(self._output_mode)
+            .option("checkpointLocation", str(self._directories.checkpoint_dir))
+            .queryName(self._directories.root.name)
+        )
+        try:
+            self._query = writer.start()
+        except Exception as error:
+            self._pending_start_stream = stream
+            with _cleanup_on_exit("streaming query", self.stop):
+                raise RuntimeError(
+                    f"Could not start the query before Batch action at index {action_index}."
+                ) from error
+
+    def stop(self) -> None:
+        """Stop once; repeat calls are safe after success or before first start."""
+        if self._query is not None:
+            self._query.stop()
+            self._query = None
+        if self._pending_start_stream is not None:
+            _stop_named_query(self._pending_start_stream, self._directories.root.name)
+            self._pending_start_stream = None

streamcase/_spark_managed.py

@@ -0,0 +1,43 @@
+"""Private ownership boundary for one Spark scenario run."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable, Mapping
+from typing import TYPE_CHECKING
+
+from streamcase._cleanup import _cleanup_on_exit
+from streamcase._directories import RunDirectories, create_run_directories
+from streamcase._spark_capture import _BatchCapture
+from streamcase._spark_execution import _execute_batches
+from streamcase._spark_query_options import _prepare_query_configuration
+from streamcase.results import ScenarioResult
+from streamcase.scenario import Scenario
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame
+
+
+def _run_managed_batches(
+    scenario: Scenario,
+    build_stream: Callable[[RunDirectories], DataFrame],
+    *,
+    base_dir: str | os.PathLike[str] | None = None,
+    retain_artifacts: bool = False,
+    output_mode: str = "append",
+    query_options: Mapping[str, str] | None = None,
+) -> ScenarioResult:
+    """Own run directories while a caller-owned session builds the stream."""
+    approved_mode, approved_options = _prepare_query_configuration(output_mode, query_options)
+    directories = create_run_directories(base_dir=base_dir, retain_artifacts=retain_artifacts)
+    with _cleanup_on_exit("run directory", directories.cleanup):
+        stream = build_stream(directories)
+        return _execute_batches(
+            stream,
+            scenario,
+            directories,
+            _BatchCapture(),
+            output_mode=approved_mode,
+            query_options=approved_options,
+            rebuild_stream=lambda: build_stream(directories),
+        )

streamcase/_spark_query_options.py

@@ -0,0 +1,67 @@
+"""Validate the narrow, runner-safe Spark streaming writer configuration."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+_OUTPUT_MODES = frozenset({"append", "complete", "update"})
+_RESERVED_QUERY_OPTIONS = frozenset(
+    {
+        "availablenow",
+        "checkpointlocation",
+        "clusterby",
+        "continuous",
+        "foreach",
+        "foreachbatch",
+        "format",
+        "once",
+        "outputmode",
+        "partitionby",
+        "path",
+        "paths",
+        "processingtime",
+        "queryname",
+        "sink",
+        "table",
+        "totable",
+        "trigger",
+    }
+)
+
+
+def _prepare_query_configuration(
+    output_mode: str,
+    query_options: Mapping[str, str] | None,
+) -> tuple[str, dict[str, str]]:
+    """Return a validated mode and detached copy of caller writer options."""
+    if not isinstance(output_mode, str):
+        raise TypeError("Spark output mode must be a string.")
+    if output_mode not in _OUTPUT_MODES:
+        raise ValueError("Spark output mode must be 'append', 'complete', or 'update'.")
+
+    if query_options is None:
+        return output_mode, {}
+    if not isinstance(query_options, Mapping):
+        raise TypeError("Spark query options must be a mapping of strings to strings.")
+
+    copied = dict(query_options)
+    seen: set[str] = set()
+    conflicts: list[str] = []
+    for name, value in copied.items():
+        if not isinstance(name, str):
+            raise TypeError("Spark query option names must be strings.")
+        if not isinstance(value, str):
+            raise TypeError(f"Spark query option {name!r} must have a string value.")
+
+        normalized = name.casefold()
+        if normalized in seen:
+            raise ValueError(f"Spark query option {name!r} duplicates a case-insensitive name.")
+        seen.add(normalized)
+        if normalized in _RESERVED_QUERY_OPTIONS:
+            conflicts.append(name)
+
+    if conflicts:
+        rendered = ", ".join(repr(name) for name in sorted(conflicts, key=str.casefold))
+        raise ValueError(f"Runner-owned query options cannot be overridden: {rendered}.")
+
+    return output_mode, copied