From 421cd393fb005f8fab89febdff1066b74c0c81e2 Mon Sep 17 00:00:00 2001 From: Manuel Konrad <84141230+manuelkonrad@users.noreply.github.com> Date: Sun, 27 Jul 2025 16:50:34 +0000 Subject: [PATCH] refactored engine and removed feature store Signed-off-by: Manuel Konrad <84141230+manuelkonrad@users.noreply.github.com> --- .github/workflows/bandit.yml | 2 + .gitignore | 8 +- README.md | 140 ++- pyproject.toml | 18 +- src/batchwise/__init__.py | 4 +- src/batchwise/cli.py | 62 - src/batchwise/config.py | 44 - src/batchwise/constants.py | 101 -- src/batchwise/dataset.py | 823 ------------- src/batchwise/engine.py | 270 ++--- src/batchwise/processor.py | 464 ++++++-- src/batchwise/store.py | 57 - .../__main__.py => tests/__init__.py | 5 - tests/conftest.py | 38 + tests/example_engine.py | 19 - tests/test_cron_parsing.py | 135 +++ tests/test_engine.py | 515 ++++++--- tests/test_imports.py | 33 + tests/test_integration.py | 211 ++++ tests/test_processor.py | 1016 +++++++++++++++++ tests/test_window.py | 58 + 21 files changed, 2491 insertions(+), 1532 deletions(-) delete mode 100644 src/batchwise/cli.py delete mode 100644 src/batchwise/config.py delete mode 100644 src/batchwise/constants.py delete mode 100644 src/batchwise/dataset.py delete mode 100644 src/batchwise/store.py rename src/batchwise/__main__.py => tests/__init__.py (53%) create mode 100644 tests/conftest.py delete mode 100644 tests/example_engine.py create mode 100644 tests/test_cron_parsing.py create mode 100644 tests/test_imports.py create mode 100644 tests/test_integration.py create mode 100644 tests/test_processor.py create mode 100644 tests/test_window.py diff --git a/.github/workflows/bandit.yml b/.github/workflows/bandit.yml index c42a066..95a0a18 100644 --- a/.github/workflows/bandit.yml +++ b/.github/workflows/bandit.yml @@ -19,3 +19,5 @@ jobs: steps: - name: Perform Bandit Analysis uses: PyCQA/bandit-action@v1 + with: + targets: "src" diff --git a/.gitignore b/.gitignore index c56c353..ab26589 100644 --- a/.gitignore +++ b/.gitignore @@ -34,9 +34,13 @@ _version.py .mypy_cache/ .ruff_cache/ reports/ +.vscode/* +!.vscode/settings.json +.idea/ # Misc .DS_Store ~* -.vscode/* -!.vscode/settings.json + +# Batchwise +batchwise_checkpoints/ diff --git a/README.md b/README.md index 88f6cc3..d1425ec 100644 --- a/README.md +++ b/README.md @@ -13,19 +13,157 @@ [![Security - Bandit](https://img.shields.io/badge/security-Bandit-yellow.svg)](https://github.com/PyCQA/bandit) -Lightweight batch processing engine and feature store. +Lightweight batch processing engine. ## Table of Contents - [Getting Started](#getting_started) +- [Examples](#examples) - [License](#license) ## Getting Started +### Installation + ```console pip install batchwise ``` +### Basic engine and processor + +The core of `batchwise` consists of an `Engine` that manages one or more `Processor` functions. Each processor operates on time-based windows determined by cron expressions. + +```python +import datetime +from batchwise import Engine, Window + +# Initialize the engine (defaults are shown here) +engine = Engine( + checkpoint_path="./batchwise_checkpoints", # path or uri to checkpoint directory + logger_name="batchwise", # name of logger instance + timezone=datetime.timezone.utc, # timezone (datetime.tzinfo object) +) + +# Register a processor using the decorator +@engine.processor( + interval="*/30 * * * *", # Run every 30 minutes + delay="1h", # Wait 1 hour before considering window complete + lookback="1d", # Look back 1 day for processing windows + include_incomplete=False, # Only process complete windows +) +def my_processor(window: Window): + print(f"Processing window: {window.start} to {window.end}") + print(f"Window complete: {window.complete}") + # Your processing logic here + pass + +# Run the engine +engine() +``` + +The `Window` object provides: +- `window_id`: Unique identifier for the window +- `start`: Start time of the window (datetime.datetime) +- `end`: End time of the window (datetime.datetime) +- `complete`: Boolean indicating if the window is complete + +### Using an fsspec-compatible filesystem + +`batchwise` supports any `fsspec`-compatible filesystem for checkpoint storage, enabling cloud storage integration: + +```python +from pathlib import Path +from fsspec.implementations.dirfs import DirFileSystem +from fsspec.implementations.local import LocalFileSystem +from batchwise import Engine, Window + +# For example, construct a DirFileSystem to wrap a base filesystem +fs = LocalFileSystem() +dir_fs = DirFileSystem(path=str(Path.cwd()), fs=fs) + +engine = Engine( + checkpoint_path="batchwise_checkpoints", # sub-path on filesystem + fs=dir_fs, +) + +@engine.processor( + interval="0 * * * *", + delay="2h", + lookback="1d", +) +def my_processor(window: Window): + # Your processing logic here + pass + +engine() +``` + +For cloud storage, use the appropriate fsspec implementation (e.g., `s3fs`, `adlfs`, `gcsfs`). + +### Additional context + +You can pass additional context to your processors: + +```python +from batchwise import Engine, Window + +engine = Engine() + +# Define a dictionary to be passed to the processor +config = { + "database_url": "postgresql://...", + "threshold": 100, +} + +@engine.processor( + interval="0 0 * * *", + delay="1d", + lookback="7d", + context=config +) +def processor_with_context(window: Window, context: dict): + # Access context parameters + database_url = context["database_url"] + threshold = context["threshold"] + # Your processing logic here + pass + +engine() +``` + +**Important:** When using `context`, your processor function must accept a `context` parameter. + +### Parallel and/or continuous processing + +Run multiple processors in parallel using multiprocessing: + +```python +from batchwise import Engine, Window + +engine = Engine() + +@engine.processor(interval="*/15 * * * *", delay="30m", lookback="2h") +def processor_1(window: Window): + pass + +@engine.processor(interval="*/30 * * * *", delay="1h", lookback="4h") +def processor_2(window: Window): + pass + +@engine.processor(interval="0 * * * *", delay="2h", lookback="12h") +def processor_3(window: Window): + pass + +# Run sequentially (default) +engine() + +# Run with 3 parallel processes +engine(num_processes=3) + +# Or run continuously with a minimum time between full cycles +engine(num_processes=3, every_seconds=60) # Check and run every 60 seconds +``` + ## License `batchwise` is distributed under the terms of the [MIT](https://spdx.org/licenses/MIT.html) license. diff --git a/pyproject.toml b/pyproject.toml index 1b391de..03d1ff0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "hatchling.build" [project] name = "batchwise" dynamic = ["version"] -description = "Lightweight batch processing engine and feature store." +description = "Lightweight batch processing engine." readme = "README.md" requires-python = ">=3.10" license = "MIT" @@ -25,22 +25,14 @@ classifiers = [ "Programming Language :: Python :: 3.13", "Programming Language :: Python :: Implementation :: CPython", ] -dependencies = [ - "pandas>=1,<3", - "pyyaml~=6.0", - "pydantic~=2.3", - "pydantic-settings~=2.3", - "pyarrow>=11.0", - "pillow>=11.0", - "numpy>=1,<3", - "fsspec>=2022", -] +dependencies = [] [project.optional-dependencies] test = [ "pytest~=8.3.3", "pytest-cov~=6.0.0", "pytest-html~=4.1.1", + "fsspec>=2024", ] dev = [ "pre-commit~=4.0.1", @@ -57,9 +49,6 @@ Documentation = "https://github.com/manuelkonrad/batchwise#readme" Issues = "https://github.com/manuelkonrad/batchwise/issues" Source = "https://github.com/manuelkonrad/batchwise" -[project.scripts] -batchwise = "batchwise.cli:cli" - # ~~~~~ # # Hatch # # ~~~~~ # @@ -198,4 +187,5 @@ ignore_missing_imports = true [tool.bandit] exclude_dirs = [ "scripts", + "tests", ] diff --git a/src/batchwise/__init__.py b/src/batchwise/__init__.py index ccc9977..8119860 100644 --- a/src/batchwise/__init__.py +++ b/src/batchwise/__init__.py @@ -3,6 +3,6 @@ # SPDX-License-Identifier: MIT from batchwise.engine import Engine -from batchwise.store import FeatureStore +from batchwise.processor import PreventCompletion, Window -__all__ = ["Engine", "FeatureStore"] +__all__ = ["Engine", "PreventCompletion", "Window"] diff --git a/src/batchwise/cli.py b/src/batchwise/cli.py deleted file mode 100644 index bf22315..0000000 --- a/src/batchwise/cli.py +++ /dev/null @@ -1,62 +0,0 @@ -# SPDX-FileCopyrightText: 2025 Manuel Konrad -# -# SPDX-License-Identifier: MIT - -import importlib.util -import logging -import time -from pathlib import Path - -from batchwise.config import config - -logger = logging.getLogger("batchwise") - - -def cli() -> None: - """Run the command line interface loop.""" - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - ) - processing = True - while processing: - start_time = time.time() - logger.info("Discovering and running engines.") - for path in config.engine_paths: - if Path(path).is_dir(): - file_paths = sorted(Path(path).rglob("*.py")) - else: - file_paths = [Path(path)] - for file_path in file_paths: - logger.info(f"Inspecting file {file_path}.") - try: - spec = importlib.util.spec_from_file_location( - file_path.name.replace(".py", ""), file_path - ) - if spec is None: - logger.warning( - f"Could not load spec from engine file path {file_path}." - ) - continue - module = importlib.util.module_from_spec(spec) - if spec.loader is None: - logger.warning( - f"Could not load module from spec for {file_path}." - ) - continue - spec.loader.exec_module(module) - if hasattr(module, "engine") and callable(module.engine): - logger.info(f"Running engine in {file_path}.") - module.engine() - else: - logger.info(f"No engine found in file {file_path}. Skipping.") - except Exception as e: - logger.error(f"Error while running engine file {file_path}: {e}") - if config.interval: - logger.info("Waiting for next iteration.") - time.sleep(max(0, config.interval - (time.time() - start_time))) - else: - logger.info( - "Tried to run all engines once. Exiting since no interval is set." - ) - processing = False diff --git a/src/batchwise/config.py b/src/batchwise/config.py deleted file mode 100644 index 7733338..0000000 --- a/src/batchwise/config.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-FileCopyrightText: 2025 Manuel Konrad -# -# SPDX-License-Identifier: MIT - -import os -import sys -from typing import Any - -import yaml -from pydantic import model_validator -from pydantic_settings import BaseSettings - -if os.path.basename(sys.argv[0]) == "pytest": - cli_parse_args = False -else: - cli_parse_args = True - - -class BatchwiseConfig(BaseSettings, cli_parse_args=cli_parse_args): # type: ignore - """Configuration for batch processing.""" - - config_path: str | None = None - catalog_paths: list[str] = [] - engine_paths: list[str] = [] - interval: int | None = None - processor_default_config: dict[str, Any] = {} - processor_configs: dict[str, dict[str, Any]] = {} - - @model_validator(mode="before") - @classmethod - def load_yaml(cls, values: dict) -> dict: - """Load configuration from a YAML file if a config_path is present.""" - if values.get("config_path"): - with open(values["config_path"], "r", encoding="utf-8") as config_file: - file_config = yaml.safe_load(config_file) - - # init, commandline and environment args have precendence over file config - file_config.update(values) - return file_config - else: - return values - - -config = BatchwiseConfig() diff --git a/src/batchwise/constants.py b/src/batchwise/constants.py deleted file mode 100644 index f762fe3..0000000 --- a/src/batchwise/constants.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-FileCopyrightText: 2025 Manuel Konrad -# -# SPDX-License-Identifier: MIT - -from enum import Enum - - -class ColumnType(str, Enum): - """Types of columns in a DataFrame.""" - - # String types - IDENTIFIER = "identifier" - CATEGORICAL = "categorical" - TEXTUAL = "textual" - DATE = "date" - - # Int8 types - MONTH = "month" - DAY = "day" - HOUR = "hour" - MINUTE = "minute" - SECOND = "second" - - # Int16 types - YEAR = "year" - - # Int32 types - MICROSECOND = "microsecond" - NANOSECOND = "nanosecond" - - # Timestamp types - TIMESTAMP = "timestamp" - - # Float64 types - NUMERICAL = "numerical" - - # Array of struct types - CURVE = "curve" - - # Object types - OBJECT = "object" - - -class TypeGroups: - """Mapping of column types to groups of data types.""" - - STRING_TYPES = { - ColumnType.IDENTIFIER, - ColumnType.CATEGORICAL, - ColumnType.TEXTUAL, - ColumnType.DATE, - } - - INT8_TYPES = { - ColumnType.MONTH, - ColumnType.DAY, - ColumnType.HOUR, - ColumnType.MINUTE, - ColumnType.SECOND, - } - - INT16_TYPES = {ColumnType.YEAR} - INT32_TYPES = {ColumnType.MICROSECOND, ColumnType.NANOSECOND} - TIMESTAMP_TYPES = {ColumnType.TIMESTAMP} - FLOAT64_TYPES = {ColumnType.NUMERICAL} - ARRAY_OF_STRUCT_TYPES = {ColumnType.CURVE} - OBJECT_TYPES = {ColumnType.OBJECT} - - DATETIME_VARS = [ - ColumnType.YEAR, - ColumnType.MONTH, - ColumnType.DAY, - ColumnType.HOUR, - ColumnType.MINUTE, - ColumnType.SECOND, - ColumnType.MICROSECOND, - ] - - -class StringFormat(str, Enum): - """Text-based formats that can be loaded using yaml.safe_load""" - - JSON = "json" - YAML = "yaml" - - -class ImageFormat(str, Enum): - """Image formats that can be loaded using PIL.Image""" - - IMG = "img" - BMP = "bmp" - PNG = "png" - JPG = "jpg" - JPEG = "jpeg" - - -class RawFormat(str, Enum): - """Raw formats requiring explicit encoding""" - - UTF8 = "utf-8" - BINARY = "binary" diff --git a/src/batchwise/dataset.py b/src/batchwise/dataset.py deleted file mode 100644 index ea5a464..0000000 --- a/src/batchwise/dataset.py +++ /dev/null @@ -1,823 +0,0 @@ -# SPDX-FileCopyrightText: 2025 Manuel Konrad -# -# SPDX-License-Identifier: MIT - -import datetime -import functools -import json -import operator -from abc import abstractmethod -from pathlib import Path, PurePosixPath -from typing import Any, Literal - -import fsspec -import numpy as np -import pandas as pd -import pyarrow as pa -import pyarrow.compute as pc -import pyarrow.parquet as pq -import yaml -from fsspec.implementations.arrow import ArrowFSWrapper -from fsspec.implementations.dirfs import DirFileSystem -from PIL import Image -from pyarrow import dataset, fs -from pydantic import BaseModel, model_validator - -from batchwise.constants import ( - ColumnType, - ImageFormat, - RawFormat, - StringFormat, - TypeGroups, -) - - -class FsConnection(BaseModel): - """Manages filesystem connections via PyArrow or fsspec.""" - - backend_fs_type: Literal["pyarrow", "fsspec"] = "pyarrow" - protocol: str | None = "local" - base_path: str | None = None - kwargs: dict = {} - - def __getitem__(self, key: str) -> bytes: - """Return file content as bytes from a given key.""" - fs = self.get_fsspec_fs() - with fs.open(key, "rb") as f: - return f.read() - - def __setitem__(self, key: str, value: bytes) -> None: - """Write bytes content to a given key.""" - fs = self.get_fsspec_fs() - with fs.open(key, "wb") as f: - f.write(value) - - @model_validator(mode="after") - def validate_base_path(self) -> "FsConnection": - """Validate local base path if specified, adjusting it to absolute.""" - if self.protocol == "local" and self.base_path: - self.base_path = str(Path(self.base_path).resolve()) - return self - - def get_fsspec_fs( - self, - ) -> fsspec.AbstractFileSystem | DirFileSystem | ArrowFSWrapper: - """Get an fsspec-based filesystem.""" - if self.backend_fs_type == "fsspec": - kwargs = {"auto_mkdir": True} - kwargs.update(self.kwargs) - base_fs = fsspec.filesystem( - self.protocol, - **kwargs, - ) - if self.base_path: - return DirFileSystem(self.base_path, base_fs) - else: - return base_fs - elif self.backend_fs_type == "pyarrow": - return ArrowFSWrapper(self.get_pyarrow_fs()) - - def get_pyarrow_fs(self) -> fs.FileSystem: - """Get a PyArrow filesystem from the current configuration.""" - if self.backend_fs_type == "fsspec": - fsspec_fs = fsspec.filesystem(self.protocol, **self.kwargs) - base_fs = fs.PyFileSystem(fs.FSSpecHandler(fsspec_fs)) - elif self.backend_fs_type == "pyarrow": - if self.protocol == "local": - base_fs = fs.LocalFileSystem(**self.kwargs) - elif self.protocol == "s3": - base_fs = fs.S3FileSystem(**self.kwargs) - elif self.protocol == "gcs": - base_fs = fs.GcsFileSystem(**self.kwargs) - elif self.protocol == "hdfs": - base_fs = fs.HadoopFileSystem(**self.kwargs) - else: - raise ValueError( - "Unknown pyarrow fs protocol: {}".format(self.protocol) - ) - else: - raise ValueError("Unknown backend fs type: {}".format(self.backend_fs_type)) - if self.base_path: - return fs.SubTreeFileSystem(self.base_path, base_fs) - else: - return base_fs - - -class Column(BaseModel): - """Represents a dataset column and its storage details.""" - - column_type: ColumnType - group: str | None = None - uri_connection: FsConnection | None = None - object_format: StringFormat | ImageFormat | RawFormat | None = None - type_string: str | None = None - description: str | None = None - - @model_validator(mode="before") - @classmethod - def extract_pa_type(cls, values: dict[str, Any]) -> dict[str, Any]: - """Extract an optional PyArrow type string from the column_type field.""" - if len(column_type_split := values["column_type"].split(":")) == 2: - values["column_type"], values["type_string"] = column_type_split - return values - - def read_uri_object(self, uri: str) -> Any: - """Read object from URI respecting the column's specified format.""" - if self.uri_connection is None: - raise ValueError("Column URI connection not specified.") - fs = self.uri_connection.get_fsspec_fs() - with fs.open(uri, "rb") as uri_file: - if self.object_format in StringFormat: - decoded = yaml.safe_load(uri_file) - elif self.object_format in ImageFormat: - decoded = np.array(Image.open(uri_file)) - elif self.object_format in RawFormat: - if self.object_format == "binary": - decoded = uri_file.read() - else: - decoded = uri_file.read().decode(self.object_format) - elif self.object_format is None: - decoded = uri_file.read() - else: - raise ValueError("Unknown object format: {}".format(self.object_format)) - return decoded - - def write_uri_object(self, uri: str, binary_object: Any) -> None: - """Write object to URI in the column's specified format.""" - if self.uri_connection is None: - raise ValueError("Column URI connection not specified.") - fs = self.uri_connection.get_fsspec_fs() - with fs.open(uri, "wb", auto_mkdir=True) as uri_file: - if self.object_format in StringFormat: - yaml.safe_dump(binary_object, uri_file) - elif self.object_format in ImageFormat: - Image.fromarray(binary_object).save(uri_file, self.object_format) - elif self.object_format in RawFormat: - if self.object_format == "binary": - uri_file.write(binary_object) - else: - uri_file.write(binary_object.encode(self.object_format)) - elif self.object_format is None: - uri_file.write(binary_object.encode(self.object_format)) - else: - raise ValueError("Unknown object format: {}".format(self.object_format)) - - -class AnnotationConfig(BaseModel): - """Holds configuration for annotation tasks and file paths.""" - - annotation_type: Literal["label", "value", "bbox", "polygon", "rle"] - classes: dict[int, str] | None = None - annotations_uri: str | None = None - annotations_connection: FsConnection | None = FsConnection() - task_id_column: str | None = None - target_columns: list[str] | None = None - - def _get_full_path(self, path: str) -> str: - """Build the full path to an annotation file.""" - if self.annotations_uri is None: - raise ValueError("Annotations URI not specified.") - full_path = str(PurePosixPath(self.annotations_uri, path)).replace(":/", "://") - return full_path - - def set_annotation(self, path: str, annotation: dict[str, Any]) -> None: - """Save annotation to the specified path.""" - if self.annotations_connection is None: - raise ValueError("Annotations connection not specified.") - full_path = self._get_full_path(path) - self.annotations_connection[full_path] = json.dumps(annotation).encode("utf-8") - - def get_annotation(self, path: str) -> dict[str, Any]: - """Load annotation from the specified path.""" - if self.annotations_connection is None: - raise ValueError("Annotations connection not specified.") - full_path = self._get_full_path(path) - return json.loads(self.annotations_connection[full_path]) - - -class Dataset(BaseModel): - """Base representation of a structured dataset with columns and connections.""" - - name: str - catalog: str | None = None - tags: list[str] | None = None - description: str | None = None - connections: dict[str, FsConnection] | None = None - columns: dict[str, Column] = {} - datetime_columns: list[str] = [] - timestamp_column: str | None = None - primary_id_column: str | None = None - annotation_configs: dict[str, AnnotationConfig] | None = None - group_dag: dict[str, list] | None = None - - @model_validator(mode="before") - @classmethod - def inject_connections(cls, values: dict[str, Any]) -> dict[str, Any]: - """Add default connections into FsConnection objects.""" - for column in values["columns"].values(): - if isinstance(column.get("uri_connection"), str): - column["uri_connection"] = values["connections"][ - column["uri_connection"] - ] - for annotation_config in values.get("annotation_configs", {}).values(): - if isinstance(annotation_config.get("annotations_connection"), str): - annotation_config["annotations_connection"] = values["connections"][ - annotation_config["annotations_connection"] - ] - return values - - def get_column_names_by_type( - self, - ) -> tuple[str | None, str | None, dict[str, list[str]]]: - """Return primary key column, timestamp column, and columns grouped by type.""" - columns_by_type: dict[str, list[str]] = {} - for column_type in ColumnType: - columns_by_type[column_type.value] = [] - for column_name, column in self.columns.items(): - if ( - column_name == self.primary_id_column - and not column.column_type == "identifier" - ): - raise ValueError("Primary key must be of type identifier.") - if ( - column_name == self.timestamp_column - and not column.column_type == "timestamp" - ): - raise ValueError("Timestamp column must be of type timestamp.") - if column.column_type in columns_by_type: - columns_by_type[column.column_type].append(column_name) - else: - raise ValueError("Unknown column type: {}".format(column.column_type)) - return self.primary_id_column, self.timestamp_column, columns_by_type - - def _check_every( - self, timestamp: datetime.datetime, every: str, window_unit: str - ) -> bool: - """Check if the 'every' schedule matches the given timestamp.""" - every_list = every.split() - minute = timestamp.minute - hour = timestamp.hour - day = timestamp.day - month = timestamp.month - day_of_week = (timestamp.weekday() + 1) % 7 - - if len(every_list) != 5: - raise ValueError( - "Every must have 5 components: minute hour day month day_of_week" - ) - - if window_unit == "hour" and every_list[0] != "0": - raise ValueError("String must start with 0 for hourly windows.") - elif window_unit == "day" and (every_list[0] != "0" or every_list[1] != "0"): - raise ValueError("String must start with 0 0 for daily windows.") - elif window_unit == "month" and ( - every_list[0] != "0" - or every_list[1] != "0" - or every_list[2] != "1" - or every_list[4] != "*" - ): - raise ValueError( - "String must start with 0 0 1 and end with * for monthly windows." - ) - elif window_unit == "year" and ( - every_list[0] != "0" - or every_list[1] != "0" - or every_list[2] != "1" - or every_list[3] != "1" - or every_list[4] != "*" - ): - raise ValueError("String must be 0 0 1 1 * for yearly windows.") - - def check_field(field_value, pattern): - # Wildcard means match everything - if pattern == "*": - return True - - # Handle multiple patterns separated by commas - if "," in str(pattern): - for p in str(pattern).split(","): - if check_field(field_value, p): - return True - return False - - # Handle step values (with or without ranges) - if "/" in str(pattern): - range_part, step_part = str(pattern).split("/") - step = int(step_part) - - # Handle */step syntax - if range_part == "*": - return field_value % step == 0 - - # Handle range/step syntax - if "-" in range_part: - start, end = map(int, range_part.split("-")) - if start <= field_value <= end: - return (field_value - start) % step == 0 - return False - - # Handle single value/step - return int(range_part) == field_value - - # Handle ranges - if "-" in str(pattern): - start, end = map(int, str(pattern).split("-")) - return start <= field_value <= end - - # Handle specific values - try: - return int(pattern) == field_value - except ValueError: - return False - - minute_match = check_field(minute, every_list[0]) - hour_match = check_field(hour, every_list[1]) - day_match = check_field(day, every_list[2]) - month_match = check_field(month, every_list[3]) - day_of_week_match = check_field(day_of_week, every_list[4]) - - return ( - minute_match - and hour_match - and day_match - and month_match - and day_of_week_match - ) - - def get_windows( - self, - max_lookback: int, - extend_before: int, - extend_after: int, - completion_delay: str, - output_mode: str, - every: str | None, - ) -> list[list[Any]]: - """Generate time-based windows according to dataset settings.""" - datetime_vars = TypeGroups.DATETIME_VARS - - if len(self.datetime_columns) > 0: - if self.columns[self.datetime_columns[-1]].column_type == "date": - window_unit = ColumnType.DAY.value - else: - window_unit = self.columns[self.datetime_columns[-1]].column_type.value - else: - raise ValueError("No datetime columns specified.") - - replacement_values = { - key.value: 0 - for key in datetime_vars[ - datetime_vars.index(ColumnType[window_unit.upper()]) + 1 : - ] - } - - now = datetime.datetime.now(datetime.timezone.utc) - now_rounded = now.replace(**replacement_values) # type: ignore - windows = [] - unit_delta = datetime.timedelta(**{window_unit + "s": 1}) - for i in range(max_lookback): - segments: dict[str, dict[str, Any]] = {} - for segment_name in ["inner", "before", "after"]: - segments[segment_name] = { - "parts": [], - "start": None, - "end": None, - "unit": window_unit, - } - delta = datetime.timedelta(**{window_unit + "s": i}) - inner_timestamp = now_rounded - delta - if every and not self._check_every(inner_timestamp, every, window_unit): - continue - segments["inner"]["parts"].append(inner_timestamp) - segments["inner"]["start"] = inner_timestamp - segments["inner"]["end"] = latest_end = inner_timestamp + unit_delta - for j in range(extend_before, 0, -1): - segments["before"]["parts"].append( - inner_timestamp - datetime.timedelta(**{window_unit + "s": j}) - ) - if len(segments["before"]["parts"]) > 0: - segments["before"]["start"] = segments["before"]["parts"][0] - segments["before"]["end"] = segments["before"]["parts"][-1] + delta - for j in range(1, extend_after + 1): - segments["after"]["parts"].append( - inner_timestamp + datetime.timedelta(**{window_unit + "s": j}) - ) - if len(segments["after"]["parts"]) > 0: - segments["after"]["start"] = segments["after"]["parts"][0] - segments["after"]["end"] = latest_end = ( - segments["after"]["parts"][-1] + unit_delta - ) - completed = self.check_completion(segments) - can_be_completed = now - latest_end > datetime.timedelta( - **{completion_delay.split()[1]: int(completion_delay.split()[0])} - ) - if not completed and ( - can_be_completed or output_mode in ["overwrite", "append"] - ): - windows.append([segments, can_be_completed]) - return windows - - @abstractmethod - def get_dataframes( - self, segments: dict[str, dict[str, Any]] - ) -> dict[str, pd.DataFrame]: - """Retrieve dataframes for each segment.""" - pass - - -class RwDataset(Dataset): - """Abstract base for datasets that support reading and writing.""" - - @abstractmethod - def write_dataframe( - self, dataframe: pd.DataFrame, segments: dict[str, Any] - ) -> None: - """Write the given dataframe, partitioned by segments.""" - pass - - @abstractmethod - def check_completion(self, segments: dict[str, Any]) -> bool: - """Check if the dataset is marked complete for the given segments.""" - pass - - @abstractmethod - def set_completion(self, segments: dict[str, Any]) -> None: - """Mark the dataset as complete for the given segments.""" - pass - - -class ArrowDataset(RwDataset): - """An Arrow-based dataset for partitioned reads and writes.""" - - dataset_uri: str - dataset_connection: FsConnection = FsConnection() - file_format: Literal["parquet", "csv", "json", "feather", "binary"] = "parquet" - partitioning_columns: list[str] = [] - partitioning_scheme: Literal["hive", "directory"] = "hive" - infer_schema: bool = False - - @model_validator(mode="before") - @classmethod - def inject_dataset_connection(cls, values: dict[str, Any]) -> dict[str, Any]: - """Inject dataset_connections from default connections.""" - if isinstance(values.get("dataset_connection"), str): - values["dataset_connection"] = values["connections"][ - values["dataset_connection"] - ] - return values - - @model_validator(mode="after") - def _validate_partitions(self) -> "ArrowDataset": - """Ensure partition columns match datetime columns where needed.""" - if len(self.partitioning_columns) == 0: - self.partitioning_columns = self.datetime_columns - if ( - not self.datetime_columns - == self.partitioning_columns[: len(self.datetime_columns)] - ): - raise ValueError("Datetime columns must match the first partition columns.") - - if self.file_format == "binary": - if ( - self.partitioning_columns + ["path"] == list(self.columns.keys()) - and self.columns["path"].column_type != "object" - ): - raise ValueError("Path column must be of type object.") - elif self.partitioning_columns == list(self.columns.keys()): - self.columns["path"] = Column( - column_type=ColumnType.OBJECT, - uri_connection=self.dataset_connection, - object_format=RawFormat.BINARY, - ) - else: - raise ValueError("Invalid columns for binary format.") - self._get_datetime_partitioning_string() - - return self - - def get_dataframes( - self, segments: dict[str, dict[str, Any]] - ) -> dict[str, pd.DataFrame]: - """Return dataframes for each requested segment from Arrow dataset.""" - filter_expressions = self._get_filter_expression(segments) - dataframes = {} - for segment_name, filter_expression in filter_expressions.items(): - if self.file_format == "binary": - pa_ds = self._get_arrow_dataset() - dataframes[segment_name] = pd.DataFrame( - { - "path": [ - i.path - for i in pa_ds.get_fragments(filter=filter_expression) - ] - } - ) - else: - pa_ds = self._get_arrow_dataset(filter_expression=filter_expression) - dataframes[segment_name] = pa_ds.to_table().to_pandas() - return dataframes - - def get_dataframes_date_range( - self, start: datetime.date, end: datetime.date - ) -> pd.DataFrame: - """Return a dataframe covering the specified date range.""" - unit_delta = datetime.timedelta(days=1) - parts = [] - for i in range((end - start).days + 1): - parts.append(start + datetime.timedelta(days=i)) - segments = { - "range": { - "parts": parts, - "start": parts[0], - "end": parts[-1] + unit_delta, - "unit": "day", - } - } - return self.get_dataframes(segments)["range"] - - def check_completion(self, segments: dict[str, Any]) -> bool: - """Check if the given segments are marked complete.""" - time_partition = segments["inner"]["parts"][0].strftime( - self._get_datetime_partitioning_string() - ) - fs = self.dataset_connection.get_fsspec_fs() - return fs.exists(time_partition + ".COMPLETED") - - def set_completion(self, segments: dict[str, Any]) -> None: - """Mark the specified segments as complete.""" - time_partition = segments["inner"]["parts"][0].strftime( - self._get_datetime_partitioning_string() - ) - fs = self.dataset_connection.get_fsspec_fs() - partition_path = str(PurePosixPath(self.dataset_uri, time_partition)).replace( - ":/", "://" - ) - fs.mkdir(partition_path, create_parents=True) - fs.touch(str(PurePosixPath(partition_path, ".COMPLETED")).replace(":/", "://")) - - def write_dataframe( - self, - dataframe: pd.DataFrame, - segments: dict[str, Any], - overwrite: bool = False, - mismatch: str = "error", - ) -> None: - """Write dataframe into partitioned files according to segments.""" - if len(self.datetime_columns) == 0: - raise ValueError("No datetime columns specified.") - - if len(dataframe) == 0: - return - - window_datetime = segments["inner"]["parts"][0] - datetime_partitions = dataframe[self.datetime_columns].drop_duplicates() - - if mismatch == "error": - if len(datetime_partitions) > 1: - raise ValueError("Invalid window time range.") - for column_name in self.datetime_columns: - column = self.columns[column_name] - if column.column_type == "date": - if datetime_partitions.iloc[0][ - column_name - ] != window_datetime.strftime("%Y-%m-%d"): - raise ValueError("Invalid dataframe partition.") - elif column.column_type in TypeGroups.DATETIME_VARS[:5]: - if datetime_partitions.iloc[0][column_name] != getattr( - window_datetime, column.column_type - ): - raise ValueError("Invalid dataframe partition.") - else: - raise ValueError("Invalid dataframe partition.") - elif mismatch == "ignore": - pass - else: - raise ValueError("Invalid mismatch parameter.") - datetime_partitioning_string = window_datetime.strftime( - self._get_datetime_partitioning_string() - ) - partitions = dataframe[self.partitioning_columns].drop_duplicates() - fs = self.dataset_connection.get_pyarrow_fs() - arrow_schema = None if self.infer_schema else self._get_arrow_schema() - if len(self.partitioning_columns) > len(self.datetime_columns): - if self.partitioning_scheme == "hive": - partitioning_string = ( - "/".join( - [ - column_name + "={}" - for column_name in self.partitioning_columns[ - len(self.datetime_columns) : - ] - ] - ) - + "/" - ) - elif self.partitioning_scheme == "directory": - partitioning_string = "{}/" * len( - self.partitioning_columns[len(self.datetime_columns) :] - ) - else: - raise ValueError("Invalid partitioning scheme.") - else: - partitioning_string = "" - for partition in partitions.itertuples(index=False): - full_partitioning_string = datetime_partitioning_string - filtered_dataframe = dataframe - if len(self.partitioning_columns) > len(self.datetime_columns): - full_partitioning_string += partitioning_string.format( - *partition[len(self.datetime_columns) :] - ) - filtered_dataframe = dataframe[ - functools.reduce( - operator.and_, - [ - dataframe[column_name] == getattr(partition, column_name) - for column_name in self.partitioning_columns[ - len(self.datetime_columns) : - ] - ], - ) - ] - - arrow_table = pa.table(filtered_dataframe, schema=arrow_schema) - # fs.create_dir(full_partitioning_string, recursive=True) - partition_path = str( - PurePosixPath(self.dataset_uri, full_partitioning_string) - ).replace(":/", "://") - fs.create_dir(partition_path, recursive=True) - table_path = str( - PurePosixPath(partition_path, "data." + self.file_format) - ).replace(":/", "://") - if not fs.get_file_info(table_path).is_file or overwrite: - if self.file_format == "parquet": - pq.write_table( - arrow_table, - table_path, - filesystem=fs, - ) - else: - raise ValueError("Invalid output file format.") - - def _get_datetime_partitioning_string(self) -> str: - """Build directory structure format based on partition scheme and datetime columns.""" - if self.partitioning_scheme == "hive": - format_strings = { - ("date",): "{}=%Y-%m-%d/", - ("date", "hour"): "{}=%Y-%m-%d/{}=%H/", - ("date", "hour", "minute"): "{}=%Y-%m-%d/{}=%H/{}=%M/", - ("year", "month"): "{}=%Y/{}=%m/", - ("year", "month", "day"): "{}=%Y/{}=%m/{}=%d/", - ("year", "month", "day", "hour"): "{}=%Y/{}=%m/{}=%d/{}=%H/", - ( - "year", - "month", - "day", - "hour", - "minute", - ): "{}=%Y/{}=%m/{}=%d/{}=%H/{}=%M/", - } - elif self.partitioning_scheme == "directory": - format_strings = { - ("date",): "%Y-%m-%d/", - ("date", "hour"): "%Y-%m-%d/%H/", - ("date", "hour", "minute"): "%Y-%m-%d/%H/%M/", - ("year", "month"): "%Y/%m/", - ("year", "month", "day"): "%Y/%m/%d/", - ("year", "month", "day", "hour"): "%Y/%m/%d/%H/", - ("year", "month", "day", "hour", "minute"): "%Y/%m/%d/%H/%M/", - } - else: - raise ValueError("Invalid partitioning scheme.") - - key = tuple(self._get_datetime_column_types()) - if key in format_strings: - return format_strings[key].format(*self.datetime_columns) - else: - raise ValueError("Invalid partition columns.") - - def _get_datetime_column_types(self) -> list[str]: - """Return the column types for datetime columns in the dataset.""" - return [ - self.columns[column_name].column_type - for column_name in self.datetime_columns - ] - - def _get_filter_expression( - self, segments: dict[str, dict[str, Any]] - ) -> dict[str, pc.Expression]: - """Generate filter expressions for each segment.""" - datetime_column_types = self._get_datetime_column_types() - window_unit = None - if len(self.datetime_columns) > 0: - if datetime_column_types[-1] == "date": - window_unit = "day" - else: - window_unit = datetime_column_types[-1] - - expressions = { - segments_name: pc.scalar(False) for segments_name in segments.keys() - } - expressions["undefined"] = pc.scalar(False) - - def build_expression(segment, part): - filters = [] - for column_name in self.datetime_columns: - column = self.columns[column_name] - if column.column_type == "date": - filters.append(pc.field(column_name) == part.strftime("%Y-%m-%d")) - elif column.column_type in TypeGroups.DATETIME_VARS[:5]: - filters.append( - pc.field(column_name) == getattr(part, column.column_type) - ) - if column.column_type == segment["unit"]: - break - return functools.reduce(operator.and_, filters) - - for segment_name, segment in segments.items(): - if window_unit: - for part in segment["parts"]: - expressions[segment_name] = expressions[ - segment_name - ] | build_expression(segment, part) - - if ( - segment["unit"] is None or segment["unit"] != window_unit - ) and self.timestamp_column: - expressions[segment_name] = expressions[segment_name] | ( - (pc.field(self.timestamp_column) >= segment["start"]) - & (pc.field(self.timestamp_column) < segment["end"]) - ) - elif segment["unit"] != window_unit: - raise ValueError( - "No suitable datetime or timestamp columns for filtering." - ) - - else: - expressions["undefined"] = pc.scalar(True) - - return expressions - - def _get_arrow_dataset( - self, filter_expression: pc.Expression | None = None - ) -> dataset.Dataset: - """Build and optionally filter a PyArrow dataset from the dataset URI.""" - pa_fs = self.dataset_connection.get_pyarrow_fs() - pa_fs.create_dir(self.dataset_uri, recursive=True) - arrow_schema = None if self.infer_schema else self._get_arrow_schema() - arrow_dataset = dataset.dataset( - self.dataset_uri, - format=self.file_format if self.file_format != "binary" else "parquet", - filesystem=pa_fs, - schema=arrow_schema, - partitioning=self.partitioning_scheme - if self.partitioning_scheme != "directory" - else None, - ) - if filter_expression is not None: - arrow_dataset = arrow_dataset.filter(filter_expression) - return arrow_dataset - - def _get_arrow_schema(self) -> pa.Schema: - """Generate a PyArrow schema based on the dataset's column definitions.""" - fields = [] - for column_name, column in self.columns.items(): - fields.append(pa.field(column_name, self._get_column_schema(column))) - return pa.schema(fields) - - def _get_column_schema(self, column: Column) -> pa.DataType: - """Determine the PyArrow data type from the column's definition.""" - if column.type_string: - return pa.type_for_alias(column.type_string) - elif column.column_type in TypeGroups.STRING_TYPES: - return pa.string() - elif column.column_type in TypeGroups.FLOAT64_TYPES: - return pa.float64() - elif column.column_type in TypeGroups.INT8_TYPES: - return pa.uint8() - elif column.column_type in TypeGroups.INT16_TYPES: - return pa.uint32() - elif column.column_type in TypeGroups.INT32_TYPES: - return pa.uint16() - elif column.column_type in TypeGroups.TIMESTAMP_TYPES: - return pa.timestamp("ns") - elif column.column_type in TypeGroups.ARRAY_OF_STRUCT_TYPES: - return pa.list_(pa.struct([("x", pa.float64()), ("y", pa.float64())])) - elif column.column_type in TypeGroups.OBJECT_TYPES: - if column.uri_connection: - return pa.string() - elif ( - column.object_format in StringFormat - or column.object_format == RawFormat.UTF8 - ): - return pa.string() - elif ( - column.object_format in ImageFormat - or column.object_format == RawFormat.BINARY - ): - return pa.binary() - else: - raise ValueError( - "Unknown object format: {}".format(column.object_format) - ) - else: - raise ValueError("Unknown column type: {}".format(column.column_type)) diff --git a/src/batchwise/engine.py b/src/batchwise/engine.py index 8037ccf..d8fa5c4 100644 --- a/src/batchwise/engine.py +++ b/src/batchwise/engine.py @@ -2,182 +2,134 @@ # # SPDX-License-Identifier: MIT +import datetime import logging -import uuid +import time from collections import OrderedDict -from typing import Any, Callable, Literal +from multiprocessing import Pool +from pathlib import Path +from typing import Callable -from batchwise.config import config -from batchwise.dataset import ArrowDataset from batchwise.processor import Processor -from batchwise.store import FeatureStore -logger = logging.getLogger("batchwise") + +def run_processor(config: dict) -> None: + """Run a single processor by name.""" + processor_name = config["name"] + logger = logging.getLogger(config["logger_name"]) + try: + processor = Processor(**config) + logger.info(f"Running processor '{processor_name}'.") + processor() + except Exception as e: + if config["abort_on_exception"]: + raise e + else: + logger.error(f"Processor {processor_name} failed with exception: {e}") class Engine: - """Manages processors and orchestrates data processing.""" + """Batch processing engine to manage and run processors.""" def __init__( self, - feature_store: FeatureStore | None = None, - processor_default_config: dict[str, Any] | None = None, + checkpoint_path: Path | str = "./batchwise_checkpoints", + logger_name: str = "batchwise", + timezone: datetime.tzinfo | None = datetime.timezone.utc, + fs=None, ) -> None: - """Initialize with an optional feature store.""" - self._feature_store = feature_store or FeatureStore( - catalog_paths=config.catalog_paths - ) - self._processor_default_config = ( - processor_default_config or config.processor_default_config or {} - ) - self._processor_configs = config.processor_configs or {} - self._processors: OrderedDict[str, Processor] = OrderedDict() - self._sinks: list[str] = [] + """Initialize the batch processing engine. + + Args: + checkpoint_path (Path | str): Path to store processor checkpoints. + logger_name (str): Name of the logger to use. + timezone (datetime.tzinfo | None): Timezone for processing windows. + fs: Optional filesystem abstraction. + """ + self._checkpoint_path = Path(checkpoint_path) + self._processor_configs: dict[str, dict] = OrderedDict() + self._abort_on_exception = False + self._timezone = timezone + self._logger_name = logger_name + self._fs = fs def processor( self, - name: str, - sink: str, - source: str | list[str] | None = None, - extend_before: int = 0, - extend_after: int = 0, - max_lookback: int = 10, - completion_delay: str = "2 minutes", - output_mode: Literal["overwrite", "complete", "append"] = "overwrite", - post_processing_callback: Callable[..., None] | None = None, - every: str | None = None, - config: dict[str, Any] | None = None, - ) -> Callable[..., None]: - """Decorator to register a processor with given parameters.""" - - def processor_inner(function): - self.add_processor( - name=name, - source=source, - sink=sink, - function=function, - extend_before=extend_before, - extend_after=extend_after, - max_lookback=max_lookback, - completion_delay=completion_delay, - output_mode=output_mode, - post_processing_callback=post_processing_callback, - every=every, - config=config, + interval: str, + delay: str, + lookback: str, + include_incomplete: bool = False, + context: dict | None = None, + ) -> Callable: + """Decorator to register a processor function. + + Args: + interval (str): Cron expression for processing intervals. + delay (str): Delay before considering a window complete. + lookback (str): Lookback period for processing windows. + include_incomplete (bool): Whether to include incomplete windows. + context (dict | None): Additional context for the processor. + + Returns: + Callable: Decorator function. + """ + + def _processor_inner(func: Callable) -> Callable: + if func.__name__ in self._processor_configs: + raise ValueError("Processor name must be unique.") + self._processor_configs[func.__name__] = dict( + name=func.__name__, + func=func, + interval=interval, + delay=delay, + lookback=lookback, + include_incomplete=include_incomplete, + context=context, + abort_on_exception=self._abort_on_exception, + checkpoint_path=self._checkpoint_path, + timezone=self._timezone, + logger_name=self._logger_name, + fs=self._fs, ) + return func - return processor_inner - - def _create_simple_dataset_handler(self, dataset_string: str) -> ArrowDataset: - """Create a minimal Dataset handler from a dataset string.""" - file_format, dataset_path = dataset_string.split("@") - columns = dict() - partitioning_columns = [] - datetime_columns = [] - if "[" in dataset_path: - dataset_path, partitioning_string = dataset_path.split("[") - partitioning_string = partitioning_string.strip("]") - partitioning_scheme, partition_columns = partitioning_string.split(":") - partitioning_column_strings = partition_columns.split(",") - - for partitioning_column in partitioning_column_strings: - column_name, column_type = partitioning_column.split("=") - partitioning_columns.append(column_name) - if column_type in ["date", "year", "month", "day", "hour", "minute"]: - datetime_columns.append(column_name) - columns[column_name] = { - "column_type": column_type, - } - return ArrowDataset( - name=str(uuid.uuid4()), - file_format=file_format, - dataset_uri=dataset_path, - columns=columns, - partitioning_columns=partitioning_columns, - datetime_columns=datetime_columns, - ) - - def _get_dataset_handler(self, dataset_string: str) -> ArrowDataset: - """Retrieve or build a dataset handler from registry or path.""" - if "@" in dataset_string: - return self._create_simple_dataset_handler(dataset_string) - else: - catalog_name, dataset_name = dataset_string.split(":") - if self._feature_store is not None: - return self._feature_store.catalogs[catalog_name][dataset_name] - else: - raise ValueError("Feature store must be specified.") - - def _get_dataset_handlers( - self, source: str | list[str] | None, sink: str - ) -> dict[str, Any]: - """Get dataset handlers for source and sink.""" - dataset_handlers: dict[str, Any] = {"source": None, "sink": None} - if isinstance(source, list): - dataset_handlers["source"] = [] - for dataset_string in source: - dataset_handlers["source"].append( - self._get_dataset_handler(dataset_string) - ) - elif source: - dataset_handlers["source"] = self._get_dataset_handler(source) - else: - dataset_handlers["source"] = None - if sink: - dataset_handlers["sink"] = self._get_dataset_handler(sink) - else: - raise ValueError("Sink must be specified.") - return dataset_handlers + return _processor_inner - def add_processor( + def _run_in_loop(self) -> None: + """Run registered processors sequentially.""" + for processor_config in self._processor_configs.values(): + run_processor(processor_config) + + def _run_in_pool(self, num_processes: int) -> None: + """Run registered processors in parallel.""" + with Pool(processes=num_processes) as pool: + list(pool.imap_unordered(run_processor, self._processor_configs.values())) + + def __call__( self, - function: Callable[..., Any], - name: str, - source: str | list[str] | None, - sink: str, - extend_before: int, - extend_after: int, - max_lookback: int, - completion_delay: str, - output_mode: str, - post_processing_callback: Callable[..., None] | None, - every: str | None, - config: dict[str, Any] | None, + num_processes: int = 1, + every_seconds: int | float | None = None, ) -> None: - """Add a new processor to the Engine.""" - if sink not in self._sinks: - self._sinks.append(sink) - else: - raise ValueError("Sink must be unique.") - if name in self._processors: - raise ValueError("Processor name must be unique.") - processor_config = dict(self._processor_default_config) - if name in self._processor_configs: - processor_config.update(self._processor_configs[name]) - if config: - processor_config.update(config) - self._processors[name] = Processor( - function=function, - dataset_handlers=self._get_dataset_handlers(source, sink), - extend_before=extend_before, - extend_after=extend_after, - max_lookback=max_lookback, - completion_delay=completion_delay, - output_mode=output_mode, - post_processing_callback=post_processing_callback, - every=every, - config=processor_config, - ) - - def run_sequentially(self) -> None: - """Run all registered processors in sequence.""" - for processor_name, processor in self._processors.items(): - logger.info(f"Running processor {processor_name}.") - try: - processor() - except Exception as e: - logger.error(f"Error while running processor {processor_name}: {e}") - - def __call__(self) -> None: - """Executes run_sequentially when called.""" - self.run_sequentially() + """Execute all registered processors. + + Args: + num_processes (int): Number of processes to use for parallel execution. If 1, runs sequentially. + every_seconds (int | float | None): Minimum time in seconds between full cycles. If None, runs only once. + """ + logger = logging.getLogger(self._logger_name) + while True: + start_time = time.time() + logger.info("Running engine iteration.") + num_processes = int(num_processes) + if num_processes > 1: + self._run_in_pool(num_processes) + elif num_processes == 1: + self._run_in_loop() + else: + raise ValueError("num_processes must cast to a positive integer.") + if every_seconds is not None: + sleep_time = start_time + every_seconds - time.time() + if sleep_time > 0: + time.sleep(sleep_time) + else: + break diff --git a/src/batchwise/processor.py b/src/batchwise/processor.py index 06e6e08..73eea3d 100644 --- a/src/batchwise/processor.py +++ b/src/batchwise/processor.py @@ -2,107 +2,385 @@ # # SPDX-License-Identifier: MIT -import inspect -from typing import Any, Callable +import datetime +import logging +import re +from dataclasses import dataclass +from inspect import signature +from pathlib import Path, PurePosixPath +from typing import Callable + + +@dataclass +class Window: + """Data class representing a processing window. + + Attributes: + window_id (str): Unique identifier for the window. + start (datetime.datetime): Start time of the window. + end (datetime.datetime): End time of the window. + complete (bool): Whether the window is complete. + """ + + window_id: str + start: datetime.datetime + end: datetime.datetime + complete: bool + + +class PreventCompletion(Exception): + """Exception which can be raised to prevent window completion.""" + + pass class Processor: - """Processes data windows using a user-defined function.""" + """Batch processor handling time windows based on cron expressions. + + Attributes: + name (str): Name of the processor. + interval (str): Cron expression for processing intervals. + delay (str): Delay before considering a window complete. + lookback (str): Lookback period for processing windows. + """ def __init__( self, - function: Callable[..., Any], - dataset_handlers: dict[str, Any], - extend_before: int, - extend_after: int, - max_lookback: int, - completion_delay: str, - output_mode: str, - post_processing_callback: Callable[..., None] | None, - every: str | None, - config: dict[str, Any], - ) -> None: - """Initialize with function, dataset handlers, and scheduling parameters.""" - self._function = function - self._dataset_handlers = dataset_handlers - self._extend_before = extend_before - self._extend_after = extend_after - self._max_lookback = max_lookback - self._completion_delay = completion_delay - self._output_mode = output_mode - self._post_processing_callback = post_processing_callback - self._every = every - self._config = config - self._signature = set(inspect.signature(self._function).parameters.keys()) - valid_parameters = ( - "source_data", - "sink_data", - "source_fs", - "sink_fs", - "context", - ) - if not self._signature.issubset(valid_parameters): - raise ValueError("Invalid function signature.") + name: str, + func: Callable, + interval: str, + delay: str, + lookback: str, + include_incomplete: bool, + context: dict | None, + abort_on_exception: bool, + checkpoint_path: Path, + timezone: datetime.tzinfo | None, + logger_name: str, + fs=None, + ): + """Initialize the processor. - def __call__(self) -> None: - """Execute the user-defined function over configured data windows.""" - windows = self._dataset_handlers["sink"].get_windows( - self._max_lookback, - self._extend_before, - self._extend_after, - self._completion_delay, - self._output_mode, - self._every, + Args: + name (str): Name of the processor. + func (Callable): Processor function to be called for each window. + interval (str): Cron expression for processing intervals. + delay (str): Delay before considering a window complete. + lookback (str): Lookback period for processing windows. + include_incomplete (bool): Whether to include incomplete windows. + context (dict | None): Additional context for the processor. + abort_on_exception (bool): Whether to abort on exceptions. + checkpoint_path (Path): Path to store processor checkpoints. + timezone (datetime.tzinfo | None): Timezone for processing windows. + logger_name (str): Name of the logger to use. + fs: Optional filesystem abstraction. + """ + self._name = name + self._interval = interval + self._delay = delay + self._lookback = lookback + self._include_incomplete = include_incomplete + self._abort_on_exception = abort_on_exception + self._func = func + self._timezone = timezone + self._logger_name = logger_name + self._fs = fs + + # check function signature and prepare kwargs + _expected_args = signature(self._func).parameters + if "window" not in _expected_args: + raise ValueError("Processor function must accept 'window' argument.") + self._additional_kwargs = {} + if context is not None: + if "context" not in _expected_args: + raise ValueError( + "Processor function does not accept 'context' argument." + ) + self._additional_kwargs["context"] = context + + # prepare checkpoint path + if self._fs: + self.processor_checkpoint_uri = str( + PurePosixPath(checkpoint_path / self._name) + ) + if not self._fs.exists(self.processor_checkpoint_uri): + self._fs.mkdir(self.processor_checkpoint_uri, create_parents=True) + else: + self.processor_checkpoint_path = Path(checkpoint_path) / self._name + self.processor_checkpoint_path.mkdir(parents=True, exist_ok=True) + + # set or check cron interval file + self._set_or_check_interval() + + @property + def name(self): + return self._name + + @property + def interval(self): + return self._interval + + @property + def delay(self): + return self._delay + + @property + def lookback(self): + return self._lookback + + def _raise_changed_interval(self, existing: str) -> None: + """Raise error for changed cron interval.""" + raise ValueError( + f"Cron interval mismatch for processor {self._name}: " + f"existing '{existing}', new '{self._interval}'" ) - for segments, can_be_completed in windows: - input_dict: dict[str, Any] = {} - if "source_data" in self._signature: - if isinstance(self._dataset_handlers.get("source"), list): - input_dict["source_data"] = [ - handler.get_dataframes(segments) - for handler in self._dataset_handlers["source"] - ] - elif self._dataset_handlers.get("source"): - input_dict["source_data"] = self._dataset_handlers[ - "source" - ].get_dataframes(segments) + + def _set_or_check_interval(self) -> None: + """Set or check the cron interval for the processor.""" + + CRON_EXPRESSION_FILENAME = "cron_expression" + if self._fs: + cron_uri = f"{self.processor_checkpoint_uri}/{CRON_EXPRESSION_FILENAME}" + if not self._fs.exists(cron_uri): + with self._fs.open(cron_uri, "w") as f: + f.write(self._interval) + else: + with self._fs.open(cron_uri, "r") as f: + existing_interval = f.read().strip() + if existing_interval != self._interval: + self._raise_changed_interval(existing_interval) + else: + cron_path = self.processor_checkpoint_path / CRON_EXPRESSION_FILENAME + if not cron_path.exists(): + with open(cron_path, "w") as f: + f.write(self._interval) + else: + with open(cron_path, "r") as f: + existing_interval = f.read().strip() + if existing_interval != self._interval: + self._raise_changed_interval(existing_interval) + + def _set_checkpoint(self, window_id: str) -> None: + """Set checkpoint for the given processor and window ID.""" + if self._fs: + self._fs.touch(f"{self.processor_checkpoint_uri}/{window_id}") + else: + (Path(self.processor_checkpoint_path) / f"{window_id}").touch() + + def _has_checkpoint(self, window_id: str) -> bool: + """Check if checkpoint exists for the given processor and window ID.""" + if self._fs: + return self._fs.exists(f"{self.processor_checkpoint_uri}/{window_id}") + else: + return (Path(self.processor_checkpoint_path) / f"{window_id}").exists() + + def _parse_timedelta(self, td_str: str) -> datetime.timedelta: + """Parse timedelta string like '1d', '12h', '30m' into timedelta object.""" + abbreviations = { + "w": "weeks", + "d": "days", + "h": "hours", + "m": "minutes", + "s": "seconds", + } + match = re.match(r"^(\d+)([wdhms])$", td_str) + if ( + not match + or match.group(2) not in abbreviations + or not match.group(1).isdigit() + ): + raise ValueError(f"Invalid timedelta string: {td_str}") + else: + return datetime.timedelta( + **{abbreviations[match.group(2)]: int(match.group(1))} + ) + + def _check_in_range(self, value: int, min_val: int, max_val: int) -> None: + """Check if a value is within a specified range.""" + if not min_val <= value <= max_val: + raise ValueError(f"Cron {value} out of range ({min_val}-{max_val})") + + def _parse_cron_field(self, field: str, min_val: int, max_val: int) -> list[int]: + """Parse a single cron field and return list of matching values.""" + if field == "*": + return list(range(min_val, max_val + 1)) + + elif "," in field: + values = [] + for part in field.split(","): + values.extend(self._parse_cron_field(part, min_val, max_val)) + return sorted(set(values)) + + elif "/" in field: + range_part, step_str = field.split("/") + step = int(step_str) + self._check_in_range(step, 1, max_val - min_val + 1) + if range_part == "*": + return list(range(min_val, max_val + 1, step)) + elif "-" in range_part: + start_str, end_str = range_part.split("-") + start = int(start_str.strip()) + end = int(end_str.strip()) + self._check_in_range(start, min_val, max_val) + self._check_in_range(end, min_val, max_val) + if start >= end: + raise ValueError(f"Invalid cron range field: {field}") + return list(range(start, end + 1, step)) + elif range_part.strip().isdigit(): + single_value = int(range_part.strip()) + self._check_in_range(single_value, min_val, max_val) + return list(range(single_value, max_val + 1, step)) + + elif "-" in field: + start_str, end_str = field.split("-") + start = int(start_str.strip()) + end = int(end_str.strip()) + self._check_in_range(start, min_val, max_val) + self._check_in_range(end, min_val, max_val) + if start >= end: + raise ValueError(f"Invalid cron range field: {field}") + return list(range(start, end + 1)) + + elif field.strip().isdigit(): + single_value = int(field.strip()) + self._check_in_range(single_value, min_val, max_val) + return [single_value] + + raise ValueError(f"Invalid cron field: {field}") + + def _parse_cron(self, cron_expr: str) -> dict[str, list[int]]: + """Parse cron expression into component fields.""" + parts = cron_expr.split() + if len(parts) != 5: + raise ValueError( + f"Invalid cron expression: {cron_expr}. Expected 5 fields." + ) + + min_day, max_day = (1, 7) if "7" in parts[4] else (0, 6) + + return { + "minute": self._parse_cron_field(parts[0], 0, 59), + "hour": self._parse_cron_field(parts[1], 0, 23), + "day": self._parse_cron_field(parts[2], 1, 31), + "month": self._parse_cron_field(parts[3], 1, 12), + "dow": self._parse_cron_field(parts[4], min_day, max_day), + } + + def _get_trigger_times( + self, + start: datetime.datetime, + end: datetime.datetime, + cron_fields: dict[str, list[int]], + ) -> list[datetime.datetime]: + """Get all trigger times between start and end that match the cron expression.""" + trigger_times = [] + + current = start.replace(second=0, microsecond=0) + + while current < end + datetime.timedelta( + days=366 + ): # safety limit to avoid infinite loop + if ( + current.month in cron_fields["month"] + and ( + current.day in cron_fields["day"] + or current.weekday() in [(d + 6) % 7 for d in cron_fields["dow"]] + ) # Convert Sunday=0 to Monday=0 + ): + if current.hour in cron_fields["hour"]: + if current.minute in cron_fields["minute"]: + trigger_times.append(current) + if current >= end: + break + current_minute_index = cron_fields["minute"].index( + current.minute + ) + if current_minute_index + 1 < len(cron_fields["minute"]): + current = current.replace( + minute=cron_fields["minute"][current_minute_index + 1] + ) + else: + current_hour_index = cron_fields["hour"].index(current.hour) + if current_hour_index + 1 < len(cron_fields["hour"]): + current = current.replace( + hour=cron_fields["hour"][current_hour_index + 1], + minute=cron_fields["minute"][0], + ) + else: + current += datetime.timedelta(days=1) + current = current.replace( + hour=cron_fields["hour"][0], + minute=cron_fields["minute"][0], + ) + else: + current += datetime.timedelta(minutes=1) else: - input_dict["source_data"] = None - if "sink_data" in self._signature: - input_dict["sink_data"] = self._dataset_handlers["sink"].get_dataframes( - segments + current += datetime.timedelta(hours=1) + current = current.replace(minute=cron_fields["minute"][0]) + else: + current += datetime.timedelta(days=1) + current = current.replace( + hour=cron_fields["hour"][0], minute=cron_fields["minute"][0] ) - if "source_fs" in self._signature: - if isinstance(self._dataset_handlers["source"], list): - input_dict["source_fs"] = [] - for handler in self._dataset_handlers["source"]: - source_fs_dict = {} - for column_name, column in handler.columns.items(): - if column.column_type == "object" and column.uri_connection: - source_fs_dict[column_name] = column.uri_connection - input_dict["source_fs"].append(source_fs_dict) - elif self._dataset_handlers["source"]: - input_dict["source_fs"] = {} - for column_name, column in self._dataset_handlers[ - "source" - ].columns.items(): - if column.column_type == "object" and column.uri_connection: - input_dict["source_fs"][column_name] = column.uri_connection + + return trigger_times + + def _get_windows(self) -> list[Window]: + """Convert cron expression to list of windows.""" + now = datetime.datetime.now(tz=self._timezone) + + # Parse timedelta strings + delay = self._parse_timedelta(self._delay) + lookback = self._parse_timedelta(self._lookback) + + # Calculate the earliest time to look back + earliest_time = now - lookback + + # Calculate the latest time when a window can be considered complete + latest_complete_time = now - delay + + # Parse cron expression + cron_fields = self._parse_cron(self._interval) + trigger_times = self._get_trigger_times(earliest_time, now, cron_fields) + + # Convert trigger times to windows + windows = [] + for i in range(len(trigger_times) - 1): + start_time = trigger_times[i] + end_time = trigger_times[i + 1] + + # Determine if this window is complete + complete = end_time <= latest_complete_time + window_id = f"{start_time.isoformat()}_{end_time.isoformat()}" + + if (complete or self._include_incomplete) and not self._has_checkpoint( + window_id=window_id + ): + # Create window + window = Window( + window_id=window_id, + start=start_time, + end=end_time, + complete=complete, + ) + windows.append(window) + + return windows + + def __call__(self) -> None: + """Process all applicable windows.""" + logger = logging.getLogger(self._logger_name) + for window in self._get_windows(): + try: + self._func(window=window, **self._additional_kwargs) + if window.complete: + self._set_checkpoint(window.window_id) + except PreventCompletion as e: + logger.info(f"Window {window.window_id} completion prevented: {e}") + except Exception as e: + if self._abort_on_exception: + raise e else: - input_dict["source_fs"] = None - if "sink_fs" in self._signature: - input_dict["sink_fs"] = {} - for column in self._dataset_handlers["sink"].columns: - if column.column_type == "object" and column.uri_connection: - input_dict["sink_fs"][column.name] = column.uri_connection - if "context" in self._signature: - input_dict["context"] = {"segments": segments, "config": self._config} - result = self._function(**input_dict) - overwrite = self._output_mode == "overwrite" - self._dataset_handlers["sink"].write_dataframe( - result, segments, overwrite=overwrite - ) - if can_be_completed: - self._dataset_handlers["sink"].set_completion(segments) - if self._post_processing_callback: - self._post_processing_callback(result) + logger.error( + f"Processing window {window.window_id} failed with exception: {e}" + ) diff --git a/src/batchwise/store.py b/src/batchwise/store.py deleted file mode 100644 index 4dc4dea..0000000 --- a/src/batchwise/store.py +++ /dev/null @@ -1,57 +0,0 @@ -# SPDX-FileCopyrightText: 2025 Manuel Konrad -# -# SPDX-License-Identifier: MIT - - -from pathlib import Path - -import yaml -from pydantic import BaseModel, model_validator - -from batchwise.dataset import ArrowDataset - - -class FeatureStore(BaseModel): - """Feature store model for managing dataset catalogs.""" - - catalogs: dict[str, dict[str, ArrowDataset]] = {} - catalog_paths: list[str] = [] - - @model_validator(mode="before") - @classmethod - def load_feature_store(cls, values: dict) -> dict: - """Load all datasets from catalog paths.""" - catalog_paths = values["catalog_paths"] - catalogs: dict[str, dict[str, ArrowDataset]] = dict() - for catalog_path in catalog_paths: - catalog_path = Path(catalog_path) - if ( - Path(catalog_path, "defaults.yaml").exists() - and Path(catalog_path, "defaults.json").exists() - ): - raise ValueError( - "Only one of defaults.yaml and defaults.json is allowed." - ) - elif Path(catalog_path, "defaults.yaml").exists(): - defaults_path = Path(catalog_path, "defaults.yaml") - elif Path(catalog_path, "defaults.json").exists(): - defaults_path = Path(catalog_path, "defaults.json") - else: - defaults_path = None - defaults = dict() - if defaults_path: - with open(defaults_path, "r", encoding="utf-8") as defaults_file: - defaults = yaml.safe_load(defaults_file) - for file_path in list(catalog_path.rglob("*.yaml")) + list( - catalog_path.rglob("*.json") - ): - with open(file_path, "r", encoding="utf-8") as config_file: - dataset_dict = yaml.safe_load(config_file) - dataset_dict.setdefault("connections", dict()) - for key, val in defaults.get("connections", {}).items(): - dataset_dict["connections"].setdefault(key, val) - dataset_dict.setdefault("catalog", defaults.get("catalog")) - catalog = catalogs.setdefault(dataset_dict["catalog"], dict()) - catalog[dataset_dict["name"]] = dataset_dict - values["catalogs"] = catalogs - return values diff --git a/src/batchwise/__main__.py b/tests/__init__.py similarity index 53% rename from src/batchwise/__main__.py rename to tests/__init__.py index 3cd7302..ea529ed 100644 --- a/src/batchwise/__main__.py +++ b/tests/__init__.py @@ -1,8 +1,3 @@ # SPDX-FileCopyrightText: 2025 Manuel Konrad # # SPDX-License-Identifier: MIT - -from batchwise.cli import cli - -if __name__ == "__main__": - cli() diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..7ffd3c9 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: 2025 Manuel Konrad +# +# SPDX-License-Identifier: MIT + +import datetime +import logging +import tempfile +from pathlib import Path + +import pytest + + +@pytest.fixture +def temp_checkpoint_dir(): + """Fixture providing a temporary checkpoint directory.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + + +@pytest.fixture +def test_logger_name(): + """Fixture providing a test logger name.""" + logger_name = "test_batchwise" + logger = logging.getLogger(logger_name) + logger.setLevel(logging.DEBUG) + return logger_name + + +@pytest.fixture +def utc_timezone(): + """Fixture providing UTC timezone.""" + return datetime.timezone.utc + + +@pytest.fixture +def sample_context(): + """Fixture providing sample context data.""" + return {"key": "value", "number": 42} diff --git a/tests/example_engine.py b/tests/example_engine.py deleted file mode 100644 index bb76689..0000000 --- a/tests/example_engine.py +++ /dev/null @@ -1,19 +0,0 @@ -# SPDX-FileCopyrightText: 2025 Manuel Konrad -# -# SPDX-License-Identifier: MIT - -from batchwise import Engine - -engine = Engine() - - -@engine.processor( - name="test", - source="test:test_1", - sink="test:test_2", -) -def test_processor(source_data, sink_data, source_fs, context): - return source_data["inner"] - - -engine() diff --git a/tests/test_cron_parsing.py b/tests/test_cron_parsing.py new file mode 100644 index 0000000..97e91ae --- /dev/null +++ b/tests/test_cron_parsing.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: 2025 Manuel Konrad +# +# SPDX-License-Identifier: MIT + +import pytest + +from batchwise.processor import Processor, Window + + +class TestCronParsing: + """Detailed tests for cron expression parsing.""" + + def create_processor( + self, interval, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Helper to create a processor with given interval.""" + + def sample_func(window: Window): + pass + + return Processor( + name="test_processor", + func=sample_func, + interval=interval, + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + def test_every_minute(self, temp_checkpoint_dir, test_logger_name, utc_timezone): + """Test every minute cron expression.""" + processor = self.create_processor( + "* * * * *", temp_checkpoint_dir, test_logger_name, utc_timezone + ) + result = processor._parse_cron("* * * * *") + assert len(result["minute"]) == 60 + + def test_every_five_minutes( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test every 5 minutes cron expression.""" + processor = self.create_processor( + "*/5 * * * *", temp_checkpoint_dir, test_logger_name, utc_timezone + ) + result = processor._parse_cron("*/5 * * * *") + assert result["minute"] == [0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55] + + def test_specific_times(self, temp_checkpoint_dir, test_logger_name, utc_timezone): + """Test specific times cron expression.""" + processor = self.create_processor( + "0 9,12,18 * * *", temp_checkpoint_dir, test_logger_name, utc_timezone + ) + result = processor._parse_cron("0 9,12,18 * * *") + assert result["minute"] == [0] + assert result["hour"] == [9, 12, 18] + + def test_weekday_only(self, temp_checkpoint_dir, test_logger_name, utc_timezone): + """Test weekday only cron expression.""" + processor = self.create_processor( + "0 9 * * 1-5", temp_checkpoint_dir, test_logger_name, utc_timezone + ) + result = processor._parse_cron("0 9 * * 1-5") + assert result["dow"] == [1, 2, 3, 4, 5] + + def test_first_day_of_month( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test first day of month cron expression.""" + processor = self.create_processor( + "0 0 1 * *", temp_checkpoint_dir, test_logger_name, utc_timezone + ) + result = processor._parse_cron("0 0 1 * *") + assert result["day"] == [1] + assert result["hour"] == [0] + assert result["minute"] == [0] + + def test_complex_expression( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test complex cron expression.""" + processor = self.create_processor( + "15,45 9-17/2 * 1,6,12 *", + temp_checkpoint_dir, + test_logger_name, + utc_timezone, + ) + result = processor._parse_cron("15,45 9-17/2 * 1,6,12 *") + assert result["minute"] == [15, 45] + assert result["hour"] == [9, 11, 13, 15, 17] + assert result["month"] == [1, 6, 12] + + def test_invalid_range_order( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test invalid range order raises error.""" + processor = self.create_processor( + "0 0 * * *", temp_checkpoint_dir, test_logger_name, utc_timezone + ) + with pytest.raises(ValueError, match="Invalid cron range field"): + processor._parse_cron("30-10 * * * *") + + def test_step_with_range(self, temp_checkpoint_dir, test_logger_name, utc_timezone): + """Test step with range.""" + processor = self.create_processor( + "10-50/10 * * * *", temp_checkpoint_dir, test_logger_name, utc_timezone + ) + result = processor._parse_cron("10-50/10 * * * *") + assert result["minute"] == [10, 20, 30, 40, 50] + + def test_sunday_variations( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test Sunday can be represented as 0 or 7.""" + # Use different checkpoint dirs to avoid interval mismatch + import tempfile + from pathlib import Path + + with tempfile.TemporaryDirectory() as tmpdir1: + processor = self.create_processor( + "0 0 * * 0", Path(tmpdir1), test_logger_name, utc_timezone + ) + result = processor._parse_cron("0 0 * * 0") + assert 0 in result["dow"] + + with tempfile.TemporaryDirectory() as tmpdir2: + processor2 = self.create_processor( + "0 0 * * 7", Path(tmpdir2), test_logger_name, utc_timezone + ) + result2 = processor2._parse_cron("0 0 * * 7") + assert 7 in result2["dow"] diff --git a/tests/test_engine.py b/tests/test_engine.py index 4e4c37f..6ad890a 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -3,153 +3,368 @@ # SPDX-License-Identifier: MIT import datetime -import shutil -import subprocess # nosec -from pathlib import Path - -import pandas as pd -import yaml -from pytest import fixture - -from batchwise import Engine, FeatureStore - - -@fixture -def prepared_tmp_paths(tmp_path): - """ - Prepare temporary paths and datasets for testing. - """ - now = datetime.datetime.now() - - # Create parquet files with three partitions - dataset_1_path = tmp_path / "test_1_dataset" - for i in range(3): - current = now - datetime.timedelta(hours=i) - df = pd.DataFrame( - { - "id": [str(i)], - "value": [i], - "date": [current.strftime("%Y-%m-%d")], - "hour": [current.hour], - } - ) - output_path = ( - dataset_1_path / f"date={current.strftime('%Y-%m-%d')}/hour={current.hour}" - ) - output_path.mkdir(parents=True, exist_ok=True) - df.to_parquet(output_path / "data.parquet") - - # Prepare dataset catalog definitions - catalog_paths = tmp_path / "catalogs" - catalog_paths.mkdir(parents=True, exist_ok=True) - - dataset_1_dict = { - "name": "test_1", - "catalog": "test", - "dataset_uri": "test_1_dataset", - "primary_id_column": "id", - "dataset_connection": { - "base_path": str(tmp_path), - }, - "format": "parquet", - "datetime_columns": ["date", "hour"], - "columns": { - "id": {"column_type": "identifier"}, - "value": {"column_type": "numerical"}, - "date": {"column_type": "date"}, - "hour": {"column_type": "hour"}, - }, - } - - with open(catalog_paths / "test_1.yaml", "w") as file: - file.write(yaml.dump(dataset_1_dict)) - - dataset_2_dict = { - "name": "test_2", - "catalog": "test", - "dataset_uri": "test_2_dataset", - "primary_id_column": "id", - "dataset_connection": { - "base_path": str(tmp_path), - }, - "format": "parquet", - "datetime_columns": ["date", "hour"], - "columns": { - "id": {"column_type": "identifier"}, - "value": {"column_type": "numerical"}, - "date": {"column_type": "date"}, - "hour": {"column_type": "hour"}, - }, - } - - with open(catalog_paths / "test_2.yaml", "w") as file: - file.write(yaml.dump(dataset_2_dict)) - - # Copy example engine - engine_paths = tmp_path / "engines" - engine_paths.mkdir(parents=True, exist_ok=True) - example_engine_path = Path(__file__).parent / "example_engine.py" - shutil.copy(example_engine_path, engine_paths / "example_engine.py") - - return tmp_path, catalog_paths, engine_paths, dataset_1_path - - -def test_simple(prepared_tmp_paths): - """ - Test a simple engine processor. - """ - base_path, _, _, dataset_1_path = prepared_tmp_paths - engine = Engine() - - output_path = base_path / "output_simple" - output_path.mkdir(parents=True, exist_ok=True) - - @engine.processor( - name="test", - source=f"parquet@{str(dataset_1_path)}[hive:date=date,hour=hour]", - sink=f"parquet@{str(output_path)}[hive:date=date,hour=hour]", - ) - def test_processor(source_data, sink_data, source_fs, context): - return source_data["inner"] - - engine() - - -def test_feature_store(prepared_tmp_paths): - """ - Test the feature store integration with the engine. - """ - base_path, catalog_paths, engine_paths, _ = prepared_tmp_paths - - output_path = base_path / "output_feature_store" - output_path.mkdir(parents=True, exist_ok=True) - feature_store = FeatureStore(catalog_paths=[str(catalog_paths)]) - engine = Engine(feature_store=feature_store) - - @engine.processor( - name="test", - source="test:test_1", - sink=f"parquet@{str(output_path)}[hive:date=date,hour=hour]", - ) - def test_processor(source_data, sink_data, source_fs, context): - return source_data["inner"] - - engine() - - -def test_cli(prepared_tmp_paths): - """ - Test the CLI functionality. - """ - _, catalog_paths, engine_paths, _ = prepared_tmp_paths - - cli_run = subprocess.run( # nosec - [ - "batchwise", - "--engine_paths", - str(engine_paths), - "--catalog_paths", - str(catalog_paths), - ] - ) - cli_run.check_returncode() +import logging +import pickle +from pathlib import Path, PurePosixPath +from unittest.mock import Mock, patch + +import pytest +from fsspec.implementations.dirfs import DirFileSystem + +from batchwise import Engine, Window +from batchwise.engine import run_processor +from batchwise.processor import Processor + + +class TestEngine: + """Test suite for Engine class.""" + + def test_engine_initialization( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test Engine initialization.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=utc_timezone, + ) + + assert engine._checkpoint_path == temp_checkpoint_dir + assert engine._logger_name == test_logger_name + assert engine._timezone == utc_timezone + assert len(engine._processor_configs) == 0 + + def test_engine_default_initialization(self): + """Test Engine with default parameters.""" + engine = Engine() + + assert engine._checkpoint_path == Path("./batchwise_checkpoints") + assert isinstance(engine._logger_name, str) + assert engine._timezone == datetime.timezone.utc + + def test_engine_initialization_with_dirfilesystem( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Ensure Engine can be initialized with an fsspec DirFileSystem.""" + dir_fs = DirFileSystem(path=str(temp_checkpoint_dir)) + engine = Engine( + checkpoint_path=Path("checkpoints"), + logger_name=test_logger_name, + timezone=utc_timezone, + fs=dir_fs, + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def dirfs_processor(window: Window): + pass + + Processor(**engine._processor_configs["dirfs_processor"]) + processor_uri = PurePosixPath(Path("checkpoints") / "dirfs_processor") + cron_uri = processor_uri / "cron_expression" + + assert engine._fs is dir_fs + assert dir_fs.exists(str(processor_uri)) + with dir_fs.open(str(cron_uri), "r") as cron_file: + assert cron_file.read().strip() == "*/30 * * * *" + + def test_engine_with_dirfilesystem_is_picklable( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Engine with fsspec handler remains picklable.""" + dir_fs = DirFileSystem(path=str(temp_checkpoint_dir)) + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=utc_timezone, + fs=dir_fs, + ) + + pickled = pickle.dumps(engine) + restored = pickle.loads(pickled) + + assert isinstance(restored, Engine) + assert isinstance(restored._fs, DirFileSystem) + assert restored._checkpoint_path == temp_checkpoint_dir + assert restored._logger_name == test_logger_name + assert restored._timezone == utc_timezone + + def test_register_processor(self, temp_checkpoint_dir, test_logger_name): + """Test registering a processor using decorator.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="0 0 * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + assert "test_processor" in engine._processor_configs + assert engine._processor_configs["test_processor"]["name"] == "test_processor" + + def test_register_processor_with_context( + self, temp_checkpoint_dir, test_logger_name, sample_context + ): + """Test registering a processor with context.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor( + interval="0 0 * * *", + delay="1h", + lookback="1d", + context=sample_context, + ) + def test_processor_with_context(window: Window, context: dict): + pass + + assert "test_processor_with_context" in engine._processor_configs + + def test_duplicate_processor_name_raises_error( + self, temp_checkpoint_dir, test_logger_name + ): + """Test that duplicate processor names raise ValueError.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="0 0 * * *", delay="1h", lookback="1d") + def duplicate_processor(window: Window): + pass + + with pytest.raises(ValueError, match="Processor name must be unique"): + + @engine.processor(interval="0 0 * * *", delay="1h", lookback="1d") + def duplicate_processor(window: Window): + pass + + def test_run_processor_invokes_processor( + self, temp_checkpoint_dir, test_logger_name + ): + """run_processor instantiates Processor and calls it.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + config = engine._processor_configs["test_processor"] + + with patch("batchwise.engine.Processor") as mock_processor_cls: + instance = Mock() + mock_processor_cls.return_value = instance + + run_processor(config) + + mock_processor_cls.assert_called_once_with(**config) + instance.assert_called_once_with() + + def test_run_processor_handles_exception_no_abort( + self, temp_checkpoint_dir, test_logger_name, caplog + ): + """run_processor logs and swallows when abort flag is False.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def failing_processor(window: Window): + pass + + config = engine._processor_configs["failing_processor"] + + with ( + patch("batchwise.engine.Processor") as mock_processor_cls, + caplog.at_level(logging.ERROR, logger=test_logger_name), + ): + instance = Mock() + instance.side_effect = RuntimeError("Test error") + mock_processor_cls.return_value = instance + + run_processor(config) + + assert "failed with exception" in caplog.text + + def test_run_processor_handles_exception_with_abort( + self, temp_checkpoint_dir, test_logger_name + ): + """run_processor re-raises when abort flag is True.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + engine._abort_on_exception = True + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def failing_processor(window: Window): + pass + + config = engine._processor_configs["failing_processor"] + + with patch("batchwise.engine.Processor") as mock_processor_cls: + instance = Mock() + instance.side_effect = RuntimeError("Test error") + mock_processor_cls.return_value = instance + + with pytest.raises(RuntimeError, match="Test error"): + run_processor(config) + + def test_run_in_loop(self, temp_checkpoint_dir, test_logger_name): + """Test running processors sequentially in loop.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + calls = [] + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def processor_1(window: Window): + calls.append("processor_1") + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def processor_2(window: Window): + calls.append("processor_2") + + with patch("batchwise.engine.run_processor") as mock_run: + engine._run_in_loop() + + assert mock_run.call_count == 2 + assert list(engine._processor_configs.keys()) == ["processor_1", "processor_2"] + + def test_engine_call_single_process(self, temp_checkpoint_dir, test_logger_name): + """Test calling engine with single process.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + with patch.object(engine, "_run_in_loop") as mock_run: + engine(num_processes=1) + mock_run.assert_called_once() + + def test_engine_call_multi_process(self, temp_checkpoint_dir, test_logger_name): + """Test calling engine with multiple processes.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + with patch.object(engine, "_run_in_pool") as mock_run: + engine(num_processes=2) + mock_run.assert_called_once_with(2) + + def test_engine_call_invalid_num_processes( + self, temp_checkpoint_dir, test_logger_name + ): + """Test calling engine with invalid num_processes raises ValueError.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + with pytest.raises( + ValueError, match="num_processes must cast to a positive integer" + ): + engine(num_processes=0) + + def test_engine_with_every_seconds(self, temp_checkpoint_dir, test_logger_name): + """Test engine respects minimum interval.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + # Mock _run_in_loop to avoid infinite loop, then test every_seconds logic + with patch.object(engine, "_run_in_loop", side_effect=StopIteration): + try: + engine(num_processes=1, every_seconds=2) + except StopIteration: + pass # Expected from mocked _run_in_loop + + def test_engine_every_seconds_triggers_sleep( + self, temp_checkpoint_dir, test_logger_name + ): + """Engine waits when every_seconds is provided.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + with ( + patch.object(engine, "_run_in_loop") as mock_loop, + patch("batchwise.engine.time.time", return_value=100.0), + patch("batchwise.engine.time.sleep") as mock_sleep, + ): + mock_loop.return_value = None + + def _sleep(duration): + assert duration == pytest.approx(1.0) + raise StopIteration() + + mock_sleep.side_effect = _sleep + + with pytest.raises(StopIteration): + engine(num_processes=1, every_seconds=1) + + mock_loop.assert_called_once() + mock_sleep.assert_called_once() + + def test_engine_every_seconds_non_positive_sleep_time( + self, temp_checkpoint_dir, test_logger_name + ): + """Engine skips sleeping when computed interval is non-positive.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + with ( + patch.object( + engine, "_run_in_loop", side_effect=[None, StopIteration] + ) as mock_loop, + patch("batchwise.engine.time.time", side_effect=[100.0, 101.0, 200.0]), + patch("batchwise.engine.time.sleep") as mock_sleep, + ): + with pytest.raises(StopIteration): + engine(num_processes=1, every_seconds=0) + + assert mock_loop.call_count >= 1 + mock_sleep.assert_not_called() + + def test_run_in_pool_uses_pool(self, temp_checkpoint_dir, test_logger_name): + """Ensure _run_in_pool leverages multiprocessing Pool.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, logger_name=test_logger_name + ) + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="1d") + def test_processor(window: Window): + pass + + with patch("batchwise.engine.Pool") as mock_pool: + pool_instance = mock_pool.return_value.__enter__.return_value + engine._run_in_pool(3) + + mock_pool.assert_called_once_with(processes=3) + pool_instance.imap_unordered.assert_called_once() + called_func, called_values = pool_instance.imap_unordered.call_args.args + assert called_func is run_processor + assert list(called_values) == list(engine._processor_configs.values()) diff --git a/tests/test_imports.py b/tests/test_imports.py new file mode 100644 index 0000000..a25c1b5 --- /dev/null +++ b/tests/test_imports.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: 2025 Manuel Konrad +# +# SPDX-License-Identifier: MIT + +"""Test that imports work correctly.""" + + +def test_import_engine(): + """Test importing Engine.""" + from batchwise import Engine + + assert Engine is not None + + +def test_import_window(): + """Test importing Window.""" + from batchwise import Window + + assert Window is not None + + +def test_import_prevent_completion(): + """Test importing PreventCompletion.""" + from batchwise import PreventCompletion + + assert PreventCompletion is not None + + +def test_import_all(): + """Test importing all public API.""" + from batchwise import Engine, PreventCompletion, Window + + assert all([Engine, Window, PreventCompletion]) diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 0000000..db1bf2b --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: 2025 Manuel Konrad +# +# SPDX-License-Identifier: MIT + +import datetime + +from batchwise import Engine, PreventCompletion, Window + + +class TestIntegration: + """Integration tests for the complete batch processing workflow.""" + + def test_end_to_end_single_processor(self, temp_checkpoint_dir, test_logger_name): + """Test end-to-end workflow with a single processor.""" + processed_windows = [] + + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=datetime.timezone.utc, + ) + + @engine.processor(interval="0 * * * *", delay="2h", lookback="6h") + def hourly_processor(window: Window): + processed_windows.append(window) + + # Run engine once + engine(num_processes=1) + + # Verify processor was registered and executed + assert "hourly_processor" in engine._processor_configs + assert len(processed_windows) >= 0 + + def test_end_to_end_multiple_processors( + self, temp_checkpoint_dir, test_logger_name + ): + """Test end-to-end workflow with multiple processors.""" + processor1_windows = [] + processor2_windows = [] + + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=datetime.timezone.utc, + ) + + @engine.processor(interval="0 * * * *", delay="2h", lookback="6h") + def hourly_processor(window: Window): + processor1_windows.append(window) + + @engine.processor(interval="0 0 * * *", delay="1d", lookback="7d") + def daily_processor(window: Window): + processor2_windows.append(window) + + # Run engine once + engine(num_processes=1) + + # Verify both processors were registered and executed + assert "hourly_processor" in engine._processor_configs + assert "daily_processor" in engine._processor_configs + + def test_checkpoint_prevents_reprocessing( + self, temp_checkpoint_dir, test_logger_name + ): + """Test that checkpoints prevent reprocessing of completed windows.""" + call_count = {"count": 0} + + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=datetime.timezone.utc, + ) + + @engine.processor(interval="0 * * * *", delay="2h", lookback="6h") + def counting_processor(window: Window): + call_count["count"] += 1 + + # Run engine twice + engine(num_processes=1) + first_count = call_count["count"] + + engine(num_processes=1) + second_count = call_count["count"] + + # Second run should not process any new windows (if all were complete) + assert second_count == first_count + + def test_processor_with_context(self, temp_checkpoint_dir, test_logger_name): + """Test processor using context data.""" + context_data = {"api_key": "test_key", "threshold": 100} + received_context = [] + + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=datetime.timezone.utc, + ) + + @engine.processor( + interval="0 * * * *", + delay="2h", + lookback="6h", + context=context_data, + ) + def context_processor(window: Window, context: dict): + received_context.append(context) + + engine(num_processes=1) + + # Verify context was passed correctly + for ctx in received_context: + assert ctx == context_data + + def test_prevent_completion_integration( + self, temp_checkpoint_dir, test_logger_name + ): + """Test PreventCompletion exception in integrated workflow.""" + call_count = {"count": 0} + first_window_id = {"id": None} + + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=datetime.timezone.utc, + ) + + @engine.processor(interval="0 * * * *", delay="2h", lookback="6h") + def conditional_processor(window: Window): + call_count["count"] += 1 + # Store first window ID + if first_window_id["id"] is None: + first_window_id["id"] = window.window_id + # Prevent completion only on first call + if call_count["count"] == 1: + raise PreventCompletion("Conditions not met") + + # Run engine twice + engine(num_processes=1) + first_count = call_count["count"] + + engine(num_processes=1) + second_count = call_count["count"] + + # Second run should reprocess first window since completion was prevented + assert second_count > first_count or first_count == 0 + + def test_parallel_processing(self, temp_checkpoint_dir, test_logger_name): + """Test parallel processing with multiple workers.""" + # Note: Multiprocessing doesn't work well with local functions in tests + # This test just verifies the engine accepts num_processes > 1 + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=datetime.timezone.utc, + ) + + @engine.processor(interval="0 * * * *", delay="2h", lookback="6h") + def processor1(window: Window): + pass + + @engine.processor(interval="0 0 * * *", delay="1d", lookback="7d") + def processor2(window: Window): + pass + + # Run with single process (multiprocessing doesn't work with local functions) + engine(num_processes=1) + + # Just verify it completes without error + assert len(engine._processor_configs) == 2 + + def test_different_intervals(self, temp_checkpoint_dir, test_logger_name): + """Test processors with different cron intervals.""" + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=datetime.timezone.utc, + ) + + @engine.processor(interval="*/15 * * * *", delay="30m", lookback="2h") + def fifteen_min_processor(window: Window): + pass + + @engine.processor(interval="*/30 * * * *", delay="1h", lookback="4h") + def thirty_min_processor(window: Window): + pass + + @engine.processor(interval="0 * * * *", delay="2h", lookback="12h") + def hourly_processor(window: Window): + pass + + engine(num_processes=1) + + assert len(engine._processor_configs) == 3 + + def test_timezone_handling(self, temp_checkpoint_dir, test_logger_name): + """Test that timezone is properly handled.""" + est = datetime.timezone(datetime.timedelta(hours=-5)) + + engine = Engine( + checkpoint_path=temp_checkpoint_dir, + logger_name=test_logger_name, + timezone=est, + ) + + @engine.processor(interval="0 0 * * *", delay="1d", lookback="7d") + def timezone_processor(window: Window): + # Verify window times have correct timezone + assert window.start.tzinfo == est + assert window.end.tzinfo == est + + engine(num_processes=1) diff --git a/tests/test_processor.py b/tests/test_processor.py new file mode 100644 index 0000000..bdc64b5 --- /dev/null +++ b/tests/test_processor.py @@ -0,0 +1,1016 @@ +# SPDX-FileCopyrightText: 2025 Manuel Konrad +# +# SPDX-License-Identifier: MIT + +import datetime +from pathlib import Path +from unittest.mock import patch + +import pytest +from fsspec.implementations.dirfs import DirFileSystem + +from batchwise.processor import PreventCompletion, Processor, Window + + +class TestProcessor: + """Test suite for Processor class.""" + + def test_processor_initialization( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test Processor initialization.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + assert processor.name == "test_processor" + assert processor.interval == "0 0 * * *" + assert processor.delay == "1h" + assert processor.lookback == "1d" + + def test_processor_missing_window_argument( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test that processor function must accept 'window' argument.""" + + def invalid_func(): + pass + + with pytest.raises(ValueError, match="must accept 'window' argument"): + Processor( + name="invalid_processor", + func=invalid_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + def test_processor_with_context_missing_argument( + self, temp_checkpoint_dir, test_logger_name, utc_timezone, sample_context + ): + """Test that processor with context must accept 'context' argument.""" + + def func_without_context(window: Window): + pass + + with pytest.raises(ValueError, match="does not accept 'context' argument"): + Processor( + name="test_processor", + func=func_without_context, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=sample_context, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + def test_fs_interval_mismatch_raises_error( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Ensure cron mismatch is detected when using filesystem abstraction.""" + + def sample_func(window: Window): + pass + + dir_fs = DirFileSystem(path=str(temp_checkpoint_dir)) + checkpoint_path = Path("fs_checkpoints") + + Processor( + name="fs_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=checkpoint_path, + timezone=utc_timezone, + logger_name=test_logger_name, + fs=dir_fs, + ) + + with pytest.raises(ValueError, match="Cron interval mismatch"): + Processor( + name="fs_processor", + func=sample_func, + interval="30 1 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=checkpoint_path, + timezone=utc_timezone, + logger_name=test_logger_name, + fs=dir_fs, + ) + + def test_fs_interval_same_value_is_accepted( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Re-registering with same interval should be allowed on filesystem.""" + + def sample_func(window: Window): + pass + + dir_fs = DirFileSystem(path=str(temp_checkpoint_dir)) + checkpoint_path = Path("fs_checkpoints") + + Processor( + name="fs_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=checkpoint_path, + timezone=utc_timezone, + logger_name=test_logger_name, + fs=dir_fs, + ) + + # Should not raise because interval remains unchanged + Processor( + name="fs_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=checkpoint_path, + timezone=utc_timezone, + logger_name=test_logger_name, + fs=dir_fs, + ) + + def test_checkpoint_helpers_with_fs( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """_set_checkpoint and _has_checkpoint operate with filesystem abstraction.""" + + def sample_func(window: Window): + pass + + dir_fs = DirFileSystem(path=str(temp_checkpoint_dir)) + checkpoint_path = Path("fs_checkpoints") + + processor = Processor( + name="fs_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=checkpoint_path, + timezone=utc_timezone, + logger_name=test_logger_name, + fs=dir_fs, + ) + + window_id = "2025-01-01T00:00:00_2025-01-01T01:00:00" + processor._set_checkpoint(window_id) + checkpoint_uri = processor.processor_checkpoint_uri + + assert dir_fs.exists(f"{checkpoint_uri}/{window_id}") + assert processor._has_checkpoint(window_id) + + def test_parse_timedelta_days( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing timedelta with days.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1d", + lookback="7d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_timedelta("5d") + assert result == datetime.timedelta(days=5) + + def test_parse_timedelta_hours( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing timedelta with hours.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 * * * *", + delay="2h", + lookback="24h", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_timedelta("12h") + assert result == datetime.timedelta(hours=12) + + def test_parse_timedelta_minutes( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing timedelta with minutes.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="*/15 * * * *", + delay="30m", + lookback="120m", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_timedelta("45m") + assert result == datetime.timedelta(minutes=45) + + def test_parse_timedelta_weeks( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing timedelta with weeks.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * 0", + delay="1w", + lookback="4w", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_timedelta("2w") + assert result == datetime.timedelta(weeks=2) + + def test_parse_timedelta_seconds( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing timedelta with seconds.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="* * * * *", + delay="30s", + lookback="300s", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_timedelta("90s") + assert result == datetime.timedelta(seconds=90) + + def test_parse_timedelta_invalid( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing invalid timedelta raises ValueError.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + with pytest.raises(ValueError, match="Invalid timedelta string"): + processor._parse_timedelta("invalid") + + with pytest.raises(ValueError, match="Invalid timedelta string"): + processor._parse_timedelta("10x") + + with pytest.raises(ValueError, match="Invalid timedelta string"): + processor._parse_timedelta("d10") + + def test_parse_cron_all_wildcards( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing cron expression with all wildcards.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="* * * * *", + delay="1m", + lookback="1h", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_cron("* * * * *") + assert len(result["minute"]) == 60 + assert len(result["hour"]) == 24 + assert len(result["day"]) == 31 + assert len(result["month"]) == 12 + assert len(result["dow"]) == 7 + + def test_parse_cron_specific_values( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing cron expression with specific values.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 12 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_cron("0 12 * * *") + assert result["minute"] == [0] + assert result["hour"] == [12] + assert len(result["day"]) == 31 + assert len(result["month"]) == 12 + + def test_parse_cron_ranges( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing cron expression with ranges.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0-30 9-17 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_cron("0-30 9-17 * * *") + assert result["minute"] == list(range(0, 31)) + assert result["hour"] == list(range(9, 18)) + + def test_parse_cron_field_single_value_step( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Parse cron field with numeric start and explicit step.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_cron_field("5/20", 0, 59) + assert result == [5, 25, 45] + + def test_parse_cron_field_invalid_range( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Invalid ranges should raise ValueError.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + with pytest.raises(ValueError, match="Invalid cron range field"): + processor._parse_cron_field("10-5", 0, 59) + + def test_parse_cron_field_invalid_range_with_step( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Invalid range inside stepped field should raise ValueError.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + with pytest.raises(ValueError, match="Invalid cron range field"): + processor._parse_cron_field("10-5/2", 0, 59) + + def test_parse_cron_field_invalid_value( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Completely invalid cron field should raise ValueError.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + with pytest.raises(ValueError, match="Invalid cron field"): + processor._parse_cron_field("abc", 0, 59) + + def test_parse_cron_field_invalid_step_expression( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Invalid step expressions should fall back to cron field error.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + with pytest.raises(ValueError, match="Invalid cron field"): + processor._parse_cron_field("foo/10", 0, 59) + + def test_get_trigger_times_advances_to_matching_day( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Ensure _get_trigger_times advances dates when no fields match.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + start = datetime.datetime(2025, 1, 1, tzinfo=utc_timezone) + end = start + datetime.timedelta(days=3) + cron_fields = { + "minute": [0], + "hour": [0], + "day": [start.day + 1], + "month": [start.month], + # choose a cron weekday that maps to a different python weekday so only DOM matches + "dow": [(((start.weekday() + 1) % 7) - 6) % 7], + } + + result = processor._get_trigger_times(start, end, cron_fields) + + assert result + assert result[0].day == start.day + 1 + + def test_get_trigger_times_empty_when_outside_range( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """No trigger times when start is far beyond end horizon.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + start = datetime.datetime(2025, 1, 1, tzinfo=utc_timezone) + end = start - datetime.timedelta(days=400) + cron_fields = { + "minute": [0], + "hour": [0], + "day": [start.day], + "month": [start.month], + "dow": list(range(0, 7)), + } + + result = processor._get_trigger_times(start, end, cron_fields) + + assert result == [] + + def test_get_trigger_times_dom_or_dow( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """DOM and DOW constraints should behave like logical OR.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 1 * 1", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + start = datetime.datetime(2025, 1, 1, tzinfo=utc_timezone) + end = start + datetime.timedelta(days=8) + + cron_fields = processor._parse_cron(processor.interval) + trigger_times = processor._get_trigger_times(start, end, cron_fields) + trigger_dates = {ts.date() for ts in trigger_times} + + # January 1st, 2025 is a Wednesday (matches DOM=1) and January 6th is a Monday (matches DOW=1) + assert datetime.date(2025, 1, 1) in trigger_dates + assert datetime.date(2025, 1, 6) in trigger_dates + + def test_parse_cron_steps( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing cron expression with steps.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="*/15 * * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_cron("*/15 * * * *") + assert result["minute"] == [0, 15, 30, 45] + + def test_parse_cron_comma_separated( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing cron expression with comma-separated values.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0,15,30,45 * * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + result = processor._parse_cron("0,15,30,45 * * * *") + assert result["minute"] == [0, 15, 30, 45] + + def test_parse_cron_invalid_field_count( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing cron expression with invalid field count.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + with pytest.raises(ValueError, match="Expected 5 fields"): + processor._parse_cron("0 0 * *") + + def test_parse_cron_out_of_range( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test parsing cron with out-of-range values.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + with pytest.raises(ValueError, match="out of range"): + processor._parse_cron("60 * * * *") + + def test_checkpoint_operations( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test checkpoint set and check operations.""" + + def sample_func(window: Window): + pass + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + window_id = "test_window_id" + assert not processor._has_checkpoint(window_id) + + processor._set_checkpoint(window_id) + assert processor._has_checkpoint(window_id) + + def test_cron_interval_persistence( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test that cron interval is persisted and checked.""" + + def sample_func(window: Window): + pass + + # Create first processor + Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + # Create second processor with same name and interval - should succeed + Processor( + name="test_processor", + func=sample_func, + interval="0 0 * * *", + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + # Try to create processor with different interval - should fail + with pytest.raises(ValueError, match="Cron interval mismatch"): + Processor( + name="test_processor", + func=sample_func, + interval="0 12 * * *", # Different interval + delay="1h", + lookback="1d", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + def test_processor_call_with_complete_window( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test processor call creates checkpoint for complete windows.""" + called_windows = [] + + def sample_func(window: Window): + called_windows.append(window) + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 * * * *", # Every hour + delay="2h", # 2 hour delay + lookback="6h", # Look back 6 hours + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + # Just call the processor - it will use real time + processor() + + # Should process some windows (depends on current time) + assert len(called_windows) >= 0 + + def test_processor_prevents_completion( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test PreventCompletion exception prevents checkpoint creation.""" + + def sample_func(window: Window): + raise PreventCompletion("Testing prevention") + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 * * * *", + delay="2h", + lookback="6h", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + # This should not raise, just log + processor() + + def test_processor_exception_handling_abort( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test processor exception handling with abort enabled.""" + + def failing_func(window: Window): + raise RuntimeError("Test error") + + processor = Processor( + name="test_processor", + func=failing_func, + interval="0 * * * *", + delay="2h", + lookback="6h", + include_incomplete=False, + context=None, + abort_on_exception=True, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + # Should raise the exception + with pytest.raises(RuntimeError, match="Test error"): + with patch.object(processor, "_get_windows") as mock_windows: + start = datetime.datetime(2025, 1, 1, 0, 0, 0, tzinfo=utc_timezone) + end = datetime.datetime(2025, 1, 1, 1, 0, 0, tzinfo=utc_timezone) + mock_windows.return_value = [ + Window( + window_id=f"{start.isoformat()}_{end.isoformat()}", + start=start, + end=end, + complete=True, + ) + ] + processor() + + def test_processor_exception_handling_no_abort( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test processor exception handling without abort.""" + + def failing_func(window: Window): + raise RuntimeError("Test error") + + processor = Processor( + name="test_processor", + func=failing_func, + interval="0 * * * *", + delay="2h", + lookback="6h", + include_incomplete=False, + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + # Should not raise the exception + with patch.object(processor, "_get_windows") as mock_windows: + start = datetime.datetime(2025, 1, 1, 0, 0, 0, tzinfo=utc_timezone) + end = datetime.datetime(2025, 1, 1, 1, 0, 0, tzinfo=utc_timezone) + mock_windows.return_value = [ + Window( + window_id=f"{start.isoformat()}_{end.isoformat()}", + start=start, + end=end, + complete=True, + ) + ] + processor() # Should not raise + + def test_include_incomplete_windows( + self, temp_checkpoint_dir, test_logger_name, utc_timezone + ): + """Test that include_incomplete flag allows processing incomplete windows.""" + called_windows = [] + + def sample_func(window: Window): + called_windows.append(window) + + processor = Processor( + name="test_processor", + func=sample_func, + interval="0 * * * *", + delay="2h", + lookback="6h", + include_incomplete=True, # Include incomplete windows + context=None, + abort_on_exception=False, + checkpoint_path=temp_checkpoint_dir, + timezone=utc_timezone, + logger_name=test_logger_name, + ) + + processor() + + # Check if any incomplete windows were processed + incomplete_count = sum(1 for w in called_windows if not w.complete) + # Depending on timing, there might be incomplete windows + assert incomplete_count >= 0 diff --git a/tests/test_window.py b/tests/test_window.py new file mode 100644 index 0000000..4ccccc1 --- /dev/null +++ b/tests/test_window.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: 2025 Manuel Konrad +# +# SPDX-License-Identifier: MIT + +import datetime + +from batchwise.processor import Window + + +class TestWindow: + """Test suite for Window dataclass.""" + + def test_window_creation(self): + """Test creating a Window instance.""" + start = datetime.datetime(2025, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc) + end = datetime.datetime(2025, 1, 1, 1, 0, 0, tzinfo=datetime.timezone.utc) + window_id = f"{start.isoformat()}_{end.isoformat()}" + + window = Window( + window_id=window_id, + start=start, + end=end, + complete=True, + ) + + assert window.window_id == window_id + assert window.start == start + assert window.end == end + assert window.complete is True + + def test_window_incomplete(self): + """Test creating an incomplete window.""" + start = datetime.datetime(2025, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc) + end = datetime.datetime(2025, 1, 1, 1, 0, 0, tzinfo=datetime.timezone.utc) + window_id = f"{start.isoformat()}_{end.isoformat()}" + + window = Window( + window_id=window_id, + start=start, + end=end, + complete=False, + ) + + assert window.complete is False + + def test_window_equality(self): + """Test window equality comparison.""" + start = datetime.datetime(2025, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc) + end = datetime.datetime(2025, 1, 1, 1, 0, 0, tzinfo=datetime.timezone.utc) + window_id = f"{start.isoformat()}_{end.isoformat()}" + + window1 = Window(window_id=window_id, start=start, end=end, complete=True) + window2 = Window(window_id=window_id, start=start, end=end, complete=True) + + assert window1.window_id == window2.window_id + assert window1.start == window2.start + assert window1.end == window2.end + assert window1.complete == window2.complete