diff --git a/CHANGELOG.md b/CHANGELOG.md index 31365ea..95231e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- **Stream reset for incremental syncs**: re-fetch an incremental stream in full and *replace* the destination table for one run, then resume incremental from it. Previously the only way to rebuild a drifted table was to delete backend rows by hand — the watermark (`last_run`, taken from the last succeeded job) had no escape hatch. Request one with `bizon run config.yml --reset`, with `source.reset: true` in the config, or with `bizon stream reset ` (`--cancel` to withdraw, `--stream` to target a stream other than the config's for templated configs). Requests are scoped to a single stream — keyed on `(name, source_name, stream_name)`, the same triple as the watermark they override — so resetting one stream never affects another under the same pipeline name. The last form records the request in the backend and the next run consumes it, so pipelines whose command line is fixed by a scheduler need no change. During a reset the producer skips the watermark and calls `get()` instead of `get_records_after()`, and the run reaches destinations as `sync_mode: full_refresh` so they replace their table through their existing full-refresh path (for `bigquery`: staging into `{table}_temp`, then a `WRITE_TRUNCATE` copy job). The job row stays `incremental`, so the reset run becomes the next run's watermark. The request stays bound to the job running it, so a crashed reset is retried as a reset rather than silently degrading into an append. Only meaningful for `sync_mode: incremental` (ignored with a warning otherwise) and supported by every destination with a working full-refresh path — the exception is `bigquery_streaming`, which has no staging table and appends even on a full refresh, so a reset there is rejected at validation instead of duplicating data. Adds a `stream_resets` table, created automatically alongside the existing ones (no migration needed). + ## [0.4.1] - 2026-06-29 ### Fixed diff --git a/CLAUDE.md b/CLAUDE.md index 429afbe..f53f907 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -29,6 +29,7 @@ uv run ruff check --fix . # Lint and auto-fix uv run bizon run config.yml # Run a pipeline from YAML config uv run bizon source list # List available sources uv run bizon stream list # List streams for a source +uv run bizon stream reset config.yml # Queue a stream reset for the next run ``` ## Releasing @@ -213,6 +214,39 @@ Then register in: - `INCREMENTAL` - Only new/updated records since last run (append-only) - `STREAM` - Continuous streaming mode +### Stream Reset + +One incremental run that ignores the watermark, re-fetches everything, and **replaces** the +destination table — then incremental resumes from that run. See `README.md#stream-reset` for the +user-facing docs. + +The whole feature hangs off a single config field, `source.reset`, so it needs no new plumbing: +`init_job()` (`bizon/engine/runner/runner.py`) runs in the parent before the producer and consumer +are submitted, and both are handed the same `bizon_config` / `config` objects. + +- **Trigger** — `bizon run --reset`, `source.reset: true`, or a pending row in `stream_resets` + written by `bizon stream reset ` (the only form that reaches a run whose command line a + scheduler owns; `--stream` overrides the config's stream). `AbstractRunner.resolve_reset()` + collapses all three into one bool. +- **Granularity** — keyed on `(name, source_name, stream_name)`, the same triple as + `get_last_successful_stream_job`, so a reset is exactly as scoped as the watermark it overrides. + Multi-stream configs (the `streams:` block) can never be reset: they require `sync_mode: stream`. +- **Producer** (`pipeline/producer.py`) — skips the `get_last_successful_stream_job` lookup and falls + through to `source.get()`. +- **Destination** — needs no reset-specific code. `SyncMetadata.from_bizon_config()` maps a reset onto + `sync_mode: full_refresh`, so the existing full-refresh path is reused with no new finalize branch + and every destination that can replace its table supports reset for free. Note this is the sync mode + of the *materialization*, not of the job. Destinations that append even on a full refresh (only + `bigquery_streaming`, which has no `finalize()`) are listed in `RESET_UNSUPPORTED_DESTINATIONS` in + `bizon/common/models.py` and rejected at config validation. +- **Job row** — stays `incremental`, so `get_last_successful_stream_job` picks the reset run up as the + next watermark automatically. +- **Crash safety** — every reset job has a consumed `stream_resets` row pointing at it + (`bind_stream_reset_to_job`). That is how a retry knows the in-flight job is a reset instead of + degrading into an append. `BigQueryDestination._ensure_clean_temp_table()` only drops the stale temp + table when the job has written no destination cursor yet — otherwise it would discard iterations the + resuming producer will not re-fetch. + ### Implementing Incremental Sync Incremental sync requires implementation in both **sources** and **destinations**. diff --git a/README.md b/README.md index b32c11c..a4a8523 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ staying small enough to read end to end. - [Destinations](#destinations) - [Sync Modes](#sync-modes) - [Incremental Sync](#incremental-sync) + - [Stream Reset](#stream-reset) - [Engine Configuration](#engine-configuration) - [Backends](#backends-state-storage) - [Queues](#queues) @@ -166,6 +167,7 @@ The CLI entry point is `bizon` (`bizon.cli.main:cli`). | `bizon run ` | Run a pipeline from a YAML config | | `bizon source list` | List available sources and their streams | | `bizon stream list ` | List a source's streams, flagged `[Supports incremental]` / `[Full refresh only]` | +| `bizon stream reset ` | Queue a [stream reset](#stream-reset) for the next run of that pipeline | | `bizon secrets check ` | Dry-run every `gsm://` / `env://` reference and report (masked) results | | `bizon destination` | Subcommand group (no subcommands yet) | @@ -220,6 +222,7 @@ Common `SourceConfig` fields (`bizon/source/config.py`); each connector adds its | `cursor_field` | `None` | Timestamp field for incremental filtering (e.g. `updated_at`) | | `authentication` | `None` | Auth block (`type` + `params`); connector-specific | | `force_ignore_checkpoint` | `false` | Ignore existing checkpoints and restart from iteration 0 | +| `reset` | `false` | Re-fetch the whole stream and replace the destination table, then resume incremental ([details](#stream-reset)) | | `max_iterations` | `None` | Cap iterations per run (default: run until source is exhausted) | | `api_config.retry_limit` | `10` | Retries before giving up on an API call | | `source_file_path` | `None` | Path to a custom source file (same as `--custom-source`) | @@ -417,6 +420,52 @@ On the first incremental run (no previous successful job): - The job is marked successful - Subsequent runs use `get_records_after()` with the `last_run` timestamp +#### Stream Reset + +A reset re-fetches the whole stream once and **replaces** the destination table, then resumes +incremental from that run. Use it when the table has drifted — a backfill, a bug in a transform, or +records the source changed without bumping its cursor field. + +There are three ways to ask for one; all do the same thing: + +```bash +# One-shot, for a run you launch yourself +bizon run config.yml --reset + +# Queued in the backend, for a pipeline whose command line is fixed by a scheduler. +# The next `bizon run config.yml` picks it up — no change to the cron/Airflow job. +bizon stream reset config.yml +bizon stream reset config.yml --cancel # changed your mind + +# One config templated across streams? Pick which stream to reset. +bizon stream reset config.yml --stream deals +``` + +```yaml +source: + sync_mode: incremental + reset: true # same effect, set in the config +``` + +What the run does differently: +- **Producer** ignores the watermark and calls `get()` instead of `get_records_after()` +- **Destination** materializes the run as a full refresh, replacing the table rather than appending to + it (for `bigquery`: staging into `{table}_temp`, then a `WRITE_TRUNCATE` copy job) +- The job row stays `incremental`, so the next run uses *this* run as its new watermark + +Notes: +- **Scoped to a single stream.** The request is keyed on `(name, source_name, stream_name)` — the same + triple as the watermark it overrides — so resetting one stream never affects another, even under the + same pipeline name. `bizon stream reset` takes the stream from the config; `--stream` overrides it. +- Only meaningful with `sync_mode: incremental`; ignored (with a warning) for the other modes. +- Supported by every destination with a working full-refresh path, since that is what a reset + materializes as. The exception is `bigquery_streaming`, which has no staging table and appends even + on a full refresh; a reset there is rejected at config validation rather than duplicating your data. +- A reset that crashes is retried as a reset — the request stays bound to its job, so a rerun cannot + silently degrade into an append. +- `bizon stream reset` writes to the backend, so it must reach the same one the pipeline uses. With + the default file-based `sqlite` backend that means the same machine and file. + ## Engine Configuration The `engine` block configures three pluggable subsystems. All have defaults, so `engine` is diff --git a/bizon/cli/main.py b/bizon/cli/main.py index 662eb4d..8665925 100644 --- a/bizon/cli/main.py +++ b/bizon/cli/main.py @@ -1,19 +1,28 @@ import click from dotenv import find_dotenv, load_dotenv -from bizon.engine.engine import RunnerFactory, replace_env_variables_in_config +from bizon.common.models import BizonConfig +from bizon.engine.backend.backend import BackendFactory +from bizon.engine.backend.config import BackendTypes +from bizon.engine.engine import ( + RunnerFactory, + replace_env_variables_in_config, + resolve_config, +) from bizon.engine.resolvers import ( ReferenceResolutionError, ResolverRegistry, collect_references_in_config, ) from bizon.engine.runner.config import LoggerLevel +from bizon.source.config import SourceSyncModes from bizon.source.discover import discover_all_sources from .utils import ( parse_from_yaml, set_custom_source_path_in_config, set_log_level, + set_reset_in_config, set_runner_in_config, ) @@ -71,6 +80,105 @@ def list(source_name: str): # noqa click.echo(f"{stream_mode} - {stream.name}") +@stream.command() +@click.argument("filename", type=click.Path(exists=True)) +@click.option( + "--env-file", + required=False, + type=click.Path(exists=True), + help="Path to .env file to load environment variables from.", +) +@click.option("--cancel", is_flag=True, default=False, help="Cancel pending reset requests instead of adding one.") +@click.option( + "--stream", + "stream_name", + required=False, + help="Reset this stream instead of the one named in the config, for configs templated across streams.", +) +def reset(filename: str, env_file: str, cancel: bool, stream_name: str): + """Request a reset of the incremental stream defined by a config file. + + The request is stored in the backend and consumed by the next run of that pipeline, so scheduled + pipelines pick it up without any change to their command line. It is scoped to a single stream: + resetting one stream never affects another, even under the same pipeline name. + """ + + if env_file: + load_dotenv(env_file) + else: + load_dotenv(find_dotenv(".env")) + + config = resolve_config(parse_from_yaml(filename)) + + if stream_name: + config["source"]["stream"] = stream_name + + bizon_config = BizonConfig.model_validate(obj=config) + + if bizon_config.source.sync_mode != SourceSyncModes.INCREMENTAL: + raise click.exceptions.ClickException( + f"Only incremental streams can be reset, but sync_mode is '{bizon_config.source.sync_mode.value}'. " + f"A '{bizon_config.source.sync_mode.value}' stream already rebuilds its destination on every run." + ) + + backend = BackendFactory.get_backend(config=bizon_config.engine.backend) + backend.check_prerequisites() + backend.create_all_tables() + + if bizon_config.engine.backend.type == BackendTypes.SQLITE: + click.secho( + "Warning: the sqlite backend is a local file, so this request is only visible to runs using the same file.", + fg="yellow", + ) + + stream_label = f"{bizon_config.source.name} - {bizon_config.source.stream}" + + if cancel: + cancelled = backend.cancel_pending_stream_resets( + name=bizon_config.name, + source_name=bizon_config.source.name, + stream_name=bizon_config.source.stream, + ) + + if cancelled: + click.secho(f"Cancelled {cancelled} pending reset request(s) for {stream_label}.", fg="green") + else: + click.echo(f"No pending reset request for {stream_label}.") + return + + if backend.get_pending_stream_reset( + name=bizon_config.name, + source_name=bizon_config.source.name, + stream_name=bizon_config.source.stream, + ): + click.echo(f"A reset is already pending for {stream_label}, nothing to do.") + return + + # Nothing validates the stream name here (the source is never instantiated), so a typo — most + # likely via --stream — would otherwise queue a reset that silently never fires. + if not backend.get_last_successful_stream_job( + name=bizon_config.name, + source_name=bizon_config.source.name, + stream_name=bizon_config.source.stream, + ): + click.secho( + f"Warning: no previous successful run found for {stream_label}. Check the stream name — " + f"a stream that has never run already fetches everything on its next run.", + fg="yellow", + ) + + backend.create_stream_reset( + name=bizon_config.name, + source_name=bizon_config.source.name, + stream_name=bizon_config.source.stream, + ) + click.secho( + f"Reset requested for {stream_label}. The next run will re-fetch the full stream, replace the " + f"destination table, and resume incremental from there.", + fg="green", + ) + + # Create a 'destination' group under 'bizon' @cli.group() def destination(): @@ -160,12 +268,20 @@ def check(filename: str, env_file: str): type=click.Path(exists=True), help="Path to .env file to load environment variables from.", ) +@click.option( + "--reset", + is_flag=True, + default=False, + help="Reset the incremental stream: re-fetch it in full and replace the destination table, " + "then resume incremental from this run.", +) def run( filename: str, custom_source: str, runner: str, log_level: LoggerLevel, env_file: str, + reset: bool, help="Run a bizon pipeline from a YAML file.", ): """Run a bizon pipeline from a YAML file.""" @@ -188,6 +304,9 @@ def run( # Override runner param in config set_runner_in_config(config=config, runner=runner) + # Override reset param in config + set_reset_in_config(config=config, reset=reset) + runner = RunnerFactory.create_from_config_dict(config=config) result = runner.run() diff --git a/bizon/cli/utils.py b/bizon/cli/utils.py index 852d40e..226e09f 100644 --- a/bizon/cli/utils.py +++ b/bizon/cli/utils.py @@ -25,6 +25,12 @@ def set_custom_source_path_in_config(config: dict, custom_source: str): config["source"]["source_file_path"] = custom_source +def set_reset_in_config(config: dict, reset: bool): + # Only ever written when the flag is passed: writing False would clobber `reset: true` set in YAML. + if reset: + config["source"]["reset"] = True + + # TODO: Refacto def set_runner_in_config(config: dict, runner: str): if runner: diff --git a/bizon/common/models.py b/bizon/common/models.py index 712280f..4e648bd 100644 --- a/bizon/common/models.py +++ b/bizon/common/models.py @@ -15,12 +15,19 @@ ) from bizon.connectors.destinations.file.src.config import FileDestinationConfig from bizon.connectors.destinations.logger.src.config import LoggerConfig +from bizon.destination.config import DestinationTypes from bizon.engine.config import EngineConfig from bizon.engine.resolvers.config import SecretsConfig from bizon.monitoring.config import MonitoringConfig from bizon.source.config import SourceConfig, SourceSyncModes from bizon.transform.config import TransformModel +# A reset reaches destinations as `sync_mode: full_refresh` (see SyncMetadata.from_bizon_config), so +# any destination with a working full-refresh path supports it for free. `bigquery_streaming` is the +# exception: it has no `finalize()` and no staging table, so even a plain full refresh appends to the +# final table instead of replacing it. +RESET_UNSUPPORTED_DESTINATIONS = {DestinationTypes.BIGQUERY_STREAMING} + class StreamSourceConfig(BaseModel): """Source-specific stream routing configuration. @@ -176,6 +183,26 @@ def validate_streams_config(cls, v: Optional[list[StreamConfig]], info) -> Optio return v + @model_validator(mode="after") + def validate_reset_is_supported_by_destination(self) -> "BizonConfig": + """Reject a reset on destinations that cannot replace their table. + + A reset re-fetches the whole stream, so a destination that appends instead of replacing would + silently duplicate the data. Fail loudly instead. + """ + # Only incremental runs act on the flag (see AbstractRunner.resolve_reset), so a reset that is + # already going to be ignored must not be rejected here. + if not self.source.reset or self.source.sync_mode != SourceSyncModes.INCREMENTAL: + return self + + if self.destination.name in RESET_UNSUPPORTED_DESTINATIONS: + raise ValueError( + f"Configuration Error: source.reset is not supported by destination " + f"'{self.destination.name}', which appends to its table instead of replacing it." + ) + + return self + @model_validator(mode="before") @classmethod def inject_config_from_streams(cls, data: Any) -> Any: @@ -260,15 +287,27 @@ class SyncMetadata(BaseModel): sync_mode: SourceSyncModes destination_name: str destination_alias: str + reset: bool = False @classmethod def from_bizon_config(cls, job_id: str, config: BizonConfig) -> "SyncMetadata": + sync_mode = config.source.sync_mode + + # A reset materializes as a full refresh: it re-fetches the whole stream and replaces the + # table. Mapping it here means every destination with a working full-refresh path supports + # reset with no changes of its own. Note this is the sync mode of the *materialization*, not + # of the job: the job row is written from bizon_config and stays `incremental`, so + # get_last_successful_stream_job keeps using the reset run as the next watermark. + if config.source.reset and sync_mode == SourceSyncModes.INCREMENTAL: + sync_mode = SourceSyncModes.FULL_REFRESH + return cls( name=config.name, job_id=job_id, source_name=config.source.name, stream_name=config.source.stream, - sync_mode=config.source.sync_mode, + sync_mode=sync_mode, destination_name=config.destination.name, destination_alias=config.destination.alias, + reset=config.source.reset, ) diff --git a/bizon/connectors/destinations/bigquery/src/destination.py b/bizon/connectors/destinations/bigquery/src/destination.py index 756bce1..c57b47c 100644 --- a/bizon/connectors/destinations/bigquery/src/destination.py +++ b/bizon/connectors/destinations/bigquery/src/destination.py @@ -58,6 +58,7 @@ def __init__( self._any_load_failed = False self._dataset_ensured = False + self._temp_table_ensured = False @property def table_id(self) -> str: @@ -138,6 +139,28 @@ def _ensure_dataset(self): self._dataset_ensured = True + def _ensure_clean_temp_table(self): + """Drop a stale temp table once per reset run, before the first load. + + Loads always WRITE_APPEND into the temp table, so rows left behind by an earlier crashed run + would be published by finalize()'s WRITE_TRUNCATE copy and end up in a table the user asked to + be replaced. Only resets need this: a reset shares the `_temp` staging table with full refresh + and is the one incremental case that publishes with WRITE_TRUNCATE. + """ + if self._temp_table_ensured or not self.sync_metadata.reset: + return + + self._temp_table_ensured = True + + # A reset that already wrote cursors is being resumed after a crash: the producer restarts from + # the last destination cursor, so the temp table holds iterations it will not re-fetch. + if self.backend.get_last_cursor_by_job_id(job_id=self.sync_metadata.job_id) is not None: + logger.info(f"Resuming stream reset, keeping temp table {self.temp_table_id} ...") + return + + logger.info(f"Stream reset: dropping stale temp table {self.temp_table_id} ...") + self.bq_client.delete_table(self.temp_table_id, not_found_ok=True) + def check_connection(self) -> bool: self._ensure_dataset() return True @@ -232,6 +255,7 @@ def load_to_bigquery(self, gcs_file: str, df_destination_records: pl.DataFrame = def write_records(self, df_destination_records: pl.DataFrame) -> Tuple[bool, str]: self._ensure_dataset() + self._ensure_clean_temp_table() gs_file_name = self.convert_and_upload_to_buffer( df_destination_records=self._rename_for_bq(df_destination_records) ) @@ -260,6 +284,7 @@ def buffer_flush_handler(self, session=None) -> DestinationIteration: return super().buffer_flush_handler(session=session) self._ensure_dataset() + self._ensure_clean_temp_table() # Snapshot iteration metadata before the buffer is flushed by the caller. destination_iteration = DestinationIteration( diff --git a/bizon/engine/backend/adapters/sqlalchemy/backend.py b/bizon/engine/backend/adapters/sqlalchemy/backend.py index 604ff2c..c3884fd 100644 --- a/bizon/engine/backend/adapters/sqlalchemy/backend.py +++ b/bizon/engine/backend/adapters/sqlalchemy/backend.py @@ -14,12 +14,14 @@ TABLE_DESTINATION_CURSOR, TABLE_SOURCE_CURSOR, TABLE_STREAM_INFO, + TABLE_STREAM_RESET, Base, CursorStatus, DestinationCursor, JobStatus, SourceCursor, StreamJob, + StreamReset, ) from .config import BigQueryConfigDetails, PostgresConfigDetails, SQLiteConfigDetails @@ -160,6 +162,10 @@ def check_prerequisites(self) -> bool: all_entities_exist = False logger.info(f"Table {TABLE_DESTINATION_CURSOR} does not exist in the database, we will create it") + if not inspect(engine).has_table(TABLE_STREAM_RESET): + all_entities_exist = False + logger.info(f"Table {TABLE_STREAM_RESET} does not exist in the database, we will create it") + return all_entities_exist def _add_and_commit(self, obj, session: Optional[Session] = None): @@ -270,6 +276,74 @@ def get_last_successful_stream_job(self, name: str, source_name: str, stream_nam logger.info(f"No last successful job found for source={source_name} stream={stream_name}") return None + #### STREAM RESET #### + + def create_stream_reset( + self, name: str, source_name: str, stream_name: str, session: Optional[Session] = None + ) -> StreamReset: + """Record a pending reset request for the given stream and return it""" + + new_stream_reset = StreamReset(name=name, source_name=source_name, stream_name=stream_name) + new_stream_reset = self._add_and_commit(new_stream_reset, session=session) + logger.debug(f"New stream reset has been requested: {new_stream_reset}") + return new_stream_reset + + def get_pending_stream_reset( + self, name: str, source_name: str, stream_name: str, session: Optional[Session] = None + ) -> Optional[StreamReset]: + """Get the most recent reset request for the given stream that no run has consumed yet""" + + query = ( + select(StreamReset) + .filter( + StreamReset.name == name, + StreamReset.source_name == source_name, + StreamReset.stream_name == stream_name, + StreamReset.consumed_at.is_(None), + ) + .order_by(StreamReset.requested_at.desc()) + .limit(1) + ) + + return self._execute(query, session=session).scalar_one_or_none() + + def get_stream_reset_by_job_id(self, job_id: str, session: Optional[Session] = None) -> Optional[StreamReset]: + """Get the reset request consumed by the given job, if that job is a reset job""" + + query = select(StreamReset).filter(StreamReset.consumed_by_job_id == job_id).limit(1) + return self._execute(query, session=session).scalar_one_or_none() + + def consume_stream_reset(self, reset_id: str, job_id: str, session: Optional[Session] = None): + """Mark the reset request as consumed by the given job""" + + stmt = ( + update(StreamReset) + .where(StreamReset.id == reset_id) + .values(consumed_at=datetime.now(tz=UTC), consumed_by_job_id=job_id) + .execution_options(synchronize_session="fetch") + ) + self._execute(stmt, session=session) + + def cancel_pending_stream_resets( + self, name: str, source_name: str, stream_name: str, session: Optional[Session] = None + ) -> int: + """Retire every pending reset request for the given stream, return how many were retired""" + + # consumed_at without a consumed_by_job_id is what distinguishes a cancelled request from one + # a run actually picked up. + stmt = ( + update(StreamReset) + .where( + StreamReset.name == name, + StreamReset.source_name == source_name, + StreamReset.stream_name == stream_name, + StreamReset.consumed_at.is_(None), + ) + .values(consumed_at=datetime.now(tz=UTC)) + .execution_options(synchronize_session="fetch") + ) + return self._execute(stmt, session=session).rowcount + #### SOURCE CURSOR #### def create_source_cursor( diff --git a/bizon/engine/backend/backend.py b/bizon/engine/backend/backend.py index 8166109..d992629 100644 --- a/bizon/engine/backend/backend.py +++ b/bizon/engine/backend/backend.py @@ -4,7 +4,14 @@ from sqlalchemy.orm import Session from .config import AbstractBackendConfig, AbstractBackendConfigDetails, BackendTypes -from .models import CursorStatus, DestinationCursor, JobStatus, SourceCursor, StreamJob +from .models import ( + CursorStatus, + DestinationCursor, + JobStatus, + SourceCursor, + StreamJob, + StreamReset, +) class AbstractBackend(ABC): @@ -69,6 +76,39 @@ def get_last_successful_stream_job( """Get the last successful job for the given source and stream name""" pass + #### STREAM RESET #### + + @abstractmethod + def create_stream_reset( + self, name: str, source_name: str, stream_name: str, session: Optional[Session] = None + ) -> StreamReset: + """Record a pending reset request for the given stream and return it""" + pass + + @abstractmethod + def get_pending_stream_reset( + self, name: str, source_name: str, stream_name: str, session: Optional[Session] = None + ) -> Optional[StreamReset]: + """Get the most recent reset request for the given stream that no run has consumed yet""" + pass + + @abstractmethod + def get_stream_reset_by_job_id(self, job_id: str, session: Optional[Session] = None) -> Optional[StreamReset]: + """Get the reset request consumed by the given job, if that job is a reset job""" + pass + + @abstractmethod + def consume_stream_reset(self, reset_id: str, job_id: str, session: Optional[Session] = None): + """Mark the reset request as consumed by the given job""" + pass + + @abstractmethod + def cancel_pending_stream_resets( + self, name: str, source_name: str, stream_name: str, session: Optional[Session] = None + ) -> int: + """Retire every pending reset request for the given stream, return how many were retired""" + pass + @abstractmethod def create_source_cursor( self, diff --git a/bizon/engine/backend/models.py b/bizon/engine/backend/models.py index b9f1a1b..baac3de 100644 --- a/bizon/engine/backend/models.py +++ b/bizon/engine/backend/models.py @@ -9,6 +9,7 @@ TABLE_STREAM_INFO = "stream_jobs" TABLE_SOURCE_CURSOR = "source_cursors" TABLE_DESTINATION_CURSOR = "destination_cursors" +TABLE_STREAM_RESET = "stream_resets" def generate_uuid(): @@ -109,3 +110,33 @@ class DestinationCursor(Base): pagination = Column( String, nullable=True, default=None, doc="Pagination source information from latest written buffer" ) + + +class StreamReset(Base): + """A request to reset an incremental stream: re-fetch it in full and replace the destination table. + + Lives in its own table rather than as a column on `stream_jobs` because `create_all_tables()` only + creates missing tables, so a new table is migration-free for existing backends while a new column + would not be. + + A row is pending until a run consumes it. `consumed_by_job_id` is what makes a reset survive a + crash: the retry recognises the in-flight job as a reset instead of degrading to an append. + """ + + __tablename__ = TABLE_STREAM_RESET + + id = Column(String(100), primary_key=True, default=generate_uuid, doc="Unique identifier for the reset request") + name = Column(String(100), nullable=False, doc="Name of the configuration, must be unique for a given pipeline") + source_name = Column(String(100), nullable=False, doc="Name of the source") + stream_name = Column(String(100), nullable=False, doc="Name of the stream") + requested_at = Column( + DateTime, nullable=False, default=lambda: datetime.now(tz=UTC), doc="Timestamp when the reset was requested" + ) + consumed_at = Column(DateTime, nullable=True, default=None, doc="Timestamp when a run picked up this reset request") + consumed_by_job_id = Column( + String(100), nullable=True, default=None, doc="Id of the job running this reset request" + ) + + def __repr__(self): + state = f"consumed by {self.consumed_by_job_id}" if self.consumed_at else "pending" + return f"" diff --git a/bizon/engine/engine.py b/bizon/engine/engine.py index 5b5537c..3123cec 100644 --- a/bizon/engine/engine.py +++ b/bizon/engine/engine.py @@ -19,14 +19,20 @@ def replace_env_variables_in_config(config: dict) -> dict: return config +def resolve_config(config: dict) -> dict: + """Resolve every reference in a raw config dict, leaving it ready for BizonConfig validation.""" + + # Replace legacy BIZON_ENV_ whole-value env references (kept for backwards compat) + config = replace_env_variables_in_config(config=config) + + # Resolve gsm:// / env:// references (whole-value and inline ${...}) + return resolve_references_in_config(config=config) + + class RunnerFactory: @staticmethod def create_from_config_dict(config: dict) -> AbstractRunner: - # Replace legacy BIZON_ENV_ whole-value env references (kept for backwards compat) - config = replace_env_variables_in_config(config=config) - - # Resolve gsm:// / env:// references (whole-value and inline ${...}) - config = resolve_references_in_config(config=config) + config = resolve_config(config=config) bizon_config = BizonConfig.model_validate(obj=config) diff --git a/bizon/engine/pipeline/producer.py b/bizon/engine/pipeline/producer.py index 5abd519..d30e2f1 100644 --- a/bizon/engine/pipeline/producer.py +++ b/bizon/engine/pipeline/producer.py @@ -136,7 +136,13 @@ def run( source_incremental_state = None is_incremental = self.bizon_config.source.sync_mode == SourceSyncModes.INCREMENTAL - if is_incremental: + if is_incremental and self.bizon_config.source.reset: + # Stream reset: deliberately ignore the watermark and re-pull everything. The destination + # replaces its table for this run, and the next run picks this job up as its new watermark. + logger.info("Stream reset: re-fetching the full stream, the destination table will be replaced.") + is_incremental = False + + elif is_incremental: # Get the last successful job to determine last_run timestamp last_successful_job = self.backend.get_last_successful_stream_job( name=self.bizon_config.name, diff --git a/bizon/engine/runner/runner.py b/bizon/engine/runner/runner.py index fa7d125..360bd95 100644 --- a/bizon/engine/runner/runner.py +++ b/bizon/engine/runner/runner.py @@ -132,6 +132,60 @@ def get_monitoring_client(sync_metadata: SyncMetadata, bizon_config: BizonConfig """Return the monitoring client instance""" return MonitorFactory.get_monitor(sync_metadata, bizon_config.monitoring) + @staticmethod + def resolve_reset(bizon_config: BizonConfig, backend: AbstractBackend, resuming_reset: bool) -> bool: + """Decide whether this run is a stream reset: re-fetch in full, then replace the table. + + A reset can be asked for in three ways, all converging here: the `--reset` CLI flag, + `source.reset` in the config, or a pending `stream_resets` marker written by + `bizon stream reset` (the only one that reaches a run whose command line is fixed by a + scheduler). + """ + if bizon_config.source.sync_mode != SourceSyncModes.INCREMENTAL: + if bizon_config.source.reset: + logger.warning( + f"source.reset is set but sync_mode is {bizon_config.source.sync_mode.value}, " + "there is no incremental state to reset - ignoring." + ) + return False + + if resuming_reset: + logger.info("Resuming an in-flight stream reset.") + return True + + if backend.get_pending_stream_reset( + name=bizon_config.name, + source_name=bizon_config.source.name, + stream_name=bizon_config.source.stream, + ): + logger.info("Found a pending stream reset request.") + return True + + return bizon_config.source.reset + + @staticmethod + def bind_stream_reset_to_job(bizon_config: BizonConfig, backend: AbstractBackend, job_id: str): + """Make sure the reset job has a consumed marker row pointing at it. + + This is the invariant that lets a crashed reset be retried: the next run recognises the + in-flight job as a reset instead of falling back to an incremental append. The `--reset` + flag path has no marker of its own, so one is created here. + """ + if backend.get_stream_reset_by_job_id(job_id=job_id): + return + + stream_reset = backend.get_pending_stream_reset( + name=bizon_config.name, + source_name=bizon_config.source.name, + stream_name=bizon_config.source.stream, + ) or backend.create_stream_reset( + name=bizon_config.name, + source_name=bizon_config.source.name, + stream_name=bizon_config.source.stream, + ) + + backend.consume_stream_reset(reset_id=stream_reset.id, job_id=job_id) + @staticmethod def get_or_create_job( bizon_config: BizonConfig, @@ -200,14 +254,38 @@ def init_job(bizon_config: BizonConfig, config: dict, **kwargs) -> StreamJob: logger.error(f"Error while connecting to source: {connection_error}") raise ConnectionError(f"Error while connecting to source: {connection_error}") + # Resolve the reset before touching the job: a reset that is not already in flight must start + # from iteration 0, so it needs a fresh job rather than the running one. + running_job = backend.get_running_stream_job( + name=bizon_config.name, + source_name=bizon_config.source.name, + stream_name=bizon_config.source.stream, + ) + resuming_reset = bool(running_job and backend.get_stream_reset_by_job_id(job_id=running_job.id)) + is_reset = AbstractRunner.resolve_reset( + bizon_config=bizon_config, backend=backend, resuming_reset=resuming_reset + ) + # Get or create the job, if force_ignore_checkpoint, we cancel the existing job and create a new one job = AbstractRunner.get_or_create_job( bizon_config=bizon_config, backend=backend, source=source, - force_create=bizon_config.source.force_ignore_checkpoint, + force_create=bizon_config.source.force_ignore_checkpoint or (is_reset and not resuming_reset), ) + if is_reset: + AbstractRunner.bind_stream_reset_to_job(bizon_config=bizon_config, backend=backend, job_id=job.id) + logger.info( + f"Stream reset for job {job.id}: the full stream will be re-fetched and the destination " + "table replaced. Incremental resumes from this run." + ) + + # Producer and consumer are handed these very objects (see the runner adapters), so setting the + # flag once here is what carries the reset to both sides of the pipeline. + bizon_config.source.reset = is_reset + config.setdefault("source", {})["reset"] = is_reset + # Set job status to running backend.update_stream_job_status(job_id=job.id, job_status=JobStatus.RUNNING) diff --git a/bizon/source/config.py b/bizon/source/config.py index 784d7a6..9c07d29 100644 --- a/bizon/source/config.py +++ b/bizon/source/config.py @@ -53,6 +53,12 @@ class SourceConfig(BaseModel, ABC): default=False, ) + reset: bool = Field( + description="Re-fetch the whole stream and replace the destination table for this run, then resume " + "incremental from it. Only meaningful with sync_mode: incremental.", + default=False, + ) + authentication: Optional[AuthConfig] = Field( description="Configuration for the authentication", default=None, diff --git a/tests/cli/test_cli_stream_reset.py b/tests/cli/test_cli_stream_reset.py new file mode 100644 index 0000000..18aa9f6 --- /dev/null +++ b/tests/cli/test_cli_stream_reset.py @@ -0,0 +1,158 @@ +import os +import tempfile + +import pytest +from click.testing import CliRunner + +from bizon.cli.main import cli +from bizon.engine.backend.adapters.sqlalchemy.backend import SQLAlchemyBackend +from bizon.engine.backend.adapters.sqlalchemy.config import ( + SQLiteConfigDetails, + SQLiteSQLAlchemyConfig, +) +from bizon.engine.backend.config import BackendTypes + +CONFIG_TEMPLATE = """ +name: test_reset_pipeline + +source: + name: dummy + stream: creatures + sync_mode: {sync_mode} + cursor_field: updated_at + authentication: + type: api_key + params: + token: dummy_key + +destination: + name: logger + config: + dummy: dummy + +engine: + backend: + type: sqlite + config: + database: {database} + schema: not_used + syncCursorInDBEvery: 2 +""" + + +@pytest.fixture +def sqlite_database(tmp_path): + """A file-backed sqlite database: the marker must outlive the CLI process that wrote it.""" + return str(tmp_path / "bizon_reset_test") + + +@pytest.fixture +def backend(sqlite_database) -> SQLAlchemyBackend: + return SQLAlchemyBackend( + config=SQLiteSQLAlchemyConfig( + type=BackendTypes.SQLITE, + config=SQLiteConfigDetails(database=sqlite_database, schema="not_used", syncCursorInDBEvery=2), + ).config, + type=BackendTypes.SQLITE, + ) + + +def write_config(sqlite_database: str, sync_mode: str = "incremental") -> str: + with tempfile.NamedTemporaryFile(mode="w", suffix=".yml", delete=False) as temp: + temp.write(CONFIG_TEMPLATE.format(sync_mode=sync_mode, database=sqlite_database)) + return temp.name + + +def pending(backend: SQLAlchemyBackend, stream_name: str = "creatures"): + return backend.get_pending_stream_reset(name="test_reset_pipeline", source_name="dummy", stream_name=stream_name) + + +def test_reset_records_a_pending_request(sqlite_database, backend): + config_path = write_config(sqlite_database) + + result = CliRunner().invoke(cli, ["stream", "reset", config_path]) + + assert result.exit_code == 0, result.output + assert pending(backend) is not None + + os.unlink(config_path) + + +def test_reset_is_idempotent(sqlite_database, backend): + """Asking twice must not queue two resets - the next run would reset, then reset again.""" + config_path = write_config(sqlite_database) + runner = CliRunner() + + runner.invoke(cli, ["stream", "reset", config_path]) + result = runner.invoke(cli, ["stream", "reset", config_path]) + + assert result.exit_code == 0, result.output + assert "already pending" in result.output + + os.unlink(config_path) + + +def test_reset_can_be_cancelled(sqlite_database, backend): + config_path = write_config(sqlite_database) + runner = CliRunner() + + runner.invoke(cli, ["stream", "reset", config_path]) + result = runner.invoke(cli, ["stream", "reset", config_path, "--cancel"]) + + assert result.exit_code == 0, result.output + assert pending(backend) is None + + os.unlink(config_path) + + +def test_reset_is_rejected_for_non_incremental_streams(sqlite_database, backend): + config_path = write_config(sqlite_database, sync_mode="full_refresh") + + result = CliRunner().invoke(cli, ["stream", "reset", config_path]) + + assert result.exit_code != 0 + assert "Only incremental streams can be reset" in result.output + + os.unlink(config_path) + + +def test_stream_option_overrides_the_config(sqlite_database, backend): + """One templated config, many streams: --stream picks which one to reset.""" + config_path = write_config(sqlite_database) + + result = CliRunner().invoke(cli, ["stream", "reset", config_path, "--stream", "pokemons"]) + + assert result.exit_code == 0, result.output + assert pending(backend, "pokemons") is not None + # The config's own stream must not be touched. + assert pending(backend, "creatures") is None + + os.unlink(config_path) + + +def test_stream_option_resets_are_independent(sqlite_database, backend): + """Resetting one stream must never fan out to another under the same pipeline name.""" + config_path = write_config(sqlite_database) + runner = CliRunner() + + runner.invoke(cli, ["stream", "reset", config_path, "--stream", "pokemons"]) + runner.invoke(cli, ["stream", "reset", config_path, "--stream", "pokemons", "--cancel"]) + + runner.invoke(cli, ["stream", "reset", config_path, "--stream", "berries"]) + + assert pending(backend, "pokemons") is None + assert pending(backend, "berries") is not None + + os.unlink(config_path) + + +def test_never_run_stream_is_flagged(sqlite_database, backend): + """A typo'd --stream would otherwise queue a reset that silently never fires.""" + config_path = write_config(sqlite_database) + + result = CliRunner().invoke(cli, ["stream", "reset", config_path, "--stream", "typoo"]) + + assert result.exit_code == 0, result.output + assert "no previous successful run found" in result.output + + os.unlink(config_path) diff --git a/tests/common/test_reset_config.py b/tests/common/test_reset_config.py new file mode 100644 index 0000000..b97f64c --- /dev/null +++ b/tests/common/test_reset_config.py @@ -0,0 +1,91 @@ +"""Tests for the `source.reset` guard on BizonConfig.""" + +import pytest +import yaml +from pydantic import ValidationError + +from bizon.common.models import BizonConfig, SyncMetadata +from bizon.source.config import SourceSyncModes + +CONFIG_TEMPLATE = """ +name: test_reset_pipeline + +source: + name: dummy + stream: creatures + sync_mode: {sync_mode} + cursor_field: updated_at + reset: {reset} + authentication: + type: api_key + params: + token: dummy_key + +destination: + name: {destination} + config: {destination_config} +""" + +DESTINATION_CONFIGS = { + "logger": "{dummy: dummy}", + "bigquery": "{project_id: p, dataset_id: d, gcs_buffer_bucket: b}", + "file": "{format: json, destination_id: /tmp/out.json}", + "bigquery_streaming_v2": "{project_id: p, dataset_id: d}", + "bigquery_streaming": "{project_id: p, dataset_id: d}", +} + + +def build_config(destination: str, reset: bool = True, sync_mode: str = "incremental") -> BizonConfig: + raw = CONFIG_TEMPLATE.format( + reset=str(reset).lower(), + sync_mode=sync_mode, + destination=destination, + destination_config=DESTINATION_CONFIGS[destination], + ) + return BizonConfig.model_validate(obj=yaml.safe_load(raw)) + + +@pytest.mark.parametrize("destination", ["bigquery", "logger", "file", "bigquery_streaming_v2"]) +def test_reset_is_accepted_by_destinations_that_can_replace(destination): + """A reset arrives as full_refresh, so any working full-refresh path supports it for free.""" + assert build_config(destination).source.reset is True + + +def test_reset_is_rejected_by_destinations_that_would_append(): + """bigquery_streaming has no finalize() and no staging table, so it appends even on full refresh.""" + with pytest.raises(ValidationError, match="source.reset is not supported"): + build_config("bigquery_streaming") + + +def test_unsupported_destination_is_unaffected_without_reset(): + assert build_config("bigquery_streaming", reset=False).source.reset is False + + +@pytest.mark.parametrize("sync_mode", ["full_refresh", "stream"]) +def test_non_incremental_runs_are_not_rejected(sync_mode): + """The runner ignores reset outside incremental, so validation must not reject it either.""" + assert build_config("bigquery_streaming", sync_mode=sync_mode).source.reset is True + + +class TestSyncMetadataSyncMode: + """A reset reaches destinations as a full refresh, which is how they replace the table.""" + + def _sync_metadata(self, sync_mode: str, reset: bool) -> SyncMetadata: + config = build_config("bigquery", reset=reset, sync_mode=sync_mode) + return SyncMetadata.from_bizon_config(job_id="job_1", config=config) + + def test_reset_incremental_is_materialized_as_full_refresh(self): + assert self._sync_metadata("incremental", reset=True).sync_mode == SourceSyncModes.FULL_REFRESH + + @pytest.mark.parametrize("sync_mode", ["full_refresh", "incremental", "stream"]) + def test_without_reset_the_sync_mode_is_passed_through(self, sync_mode): + assert self._sync_metadata(sync_mode, reset=False).sync_mode == sync_mode + + @pytest.mark.parametrize("sync_mode", ["full_refresh", "stream"]) + def test_reset_does_not_divert_other_sync_modes(self, sync_mode): + assert self._sync_metadata(sync_mode, reset=True).sync_mode == sync_mode + + @pytest.mark.parametrize("reset", [True, False]) + def test_reset_flag_is_always_carried_through(self, reset): + """Destinations still need the raw flag, e.g. to drop a stale temp table exactly once.""" + assert self._sync_metadata("incremental", reset=reset).reset is reset diff --git a/tests/connectors/destinations/bigquery/test_bigquery_incremental.py b/tests/connectors/destinations/bigquery/test_bigquery_incremental.py index 74cdf04..577edff 100644 --- a/tests/connectors/destinations/bigquery/test_bigquery_incremental.py +++ b/tests/connectors/destinations/bigquery/test_bigquery_incremental.py @@ -3,8 +3,9 @@ from unittest.mock import MagicMock, patch import pytest +import yaml -from bizon.common.models import SyncMetadata +from bizon.common.models import BizonConfig, SyncMetadata from bizon.connectors.destinations.bigquery.src.config import ( BigQueryConfigDetails, GCSBufferFormat, @@ -38,7 +39,7 @@ def bigquery_config(): ) -def create_sync_metadata(sync_mode: SourceSyncModes) -> SyncMetadata: +def create_sync_metadata(sync_mode: SourceSyncModes, reset: bool = False) -> SyncMetadata: """Create SyncMetadata with specified sync mode.""" return SyncMetadata( name="test_pipeline", @@ -48,9 +49,30 @@ def create_sync_metadata(sync_mode: SourceSyncModes) -> SyncMetadata: destination_name="bigquery", destination_alias="bigquery", sync_mode=sync_mode.value, + reset=reset, ) +RESET_CONFIG = """ +name: test_pipeline +source: + name: dummy + stream: creatures + sync_mode: incremental + cursor_field: updated_at + reset: {reset} + authentication: {{type: api_key, params: {{token: t}}}} +destination: + name: bigquery + config: {{project_id: test-project, dataset_id: test_dataset, gcs_buffer_bucket: test-bucket}} +""" + + +def build_reset_config(reset: bool = True) -> BizonConfig: + """An incremental config with `reset` set, as the runner would hand it to the destination.""" + return BizonConfig.model_validate(obj=yaml.safe_load(RESET_CONFIG.format(reset=str(reset).lower()))) + + class TestBigQueryTempTableId: """Test cases for temp_table_id property.""" @@ -224,3 +246,86 @@ def test_finalize_stream(self, bigquery_config, mock_bq_client, mock_gcs_client) # STREAM mode writes directly to the final table: no query, no copy mock_query.assert_not_called() mock_copy.assert_not_called() + + +class TestBigQueryStreamReset: + """Test cases for stream reset: an incremental job that replaces the table for one run.""" + + def _destination(self, bigquery_config, backend=None, reset=True): + # Built from a real config through from_bizon_config, so these exercise the actual wiring: + # an `incremental` + `reset` config is what has to reach the destination as a full refresh. + return BigQueryDestination( + sync_metadata=SyncMetadata.from_bizon_config(job_id="test_job_123", config=build_reset_config(reset=reset)), + config=bigquery_config, + backend=backend or MagicMock(), + source_callback=MagicMock(), + monitor=MagicMock(), + ) + + def test_temp_table_id_uses_full_refresh_staging(self, bigquery_config, mock_bq_client, mock_gcs_client): + """A reset stages into the full-refresh temp table, since it publishes with WRITE_TRUNCATE.""" + destination = self._destination(bigquery_config) + + assert destination.temp_table_id == f"{destination.table_id}_temp" + + def test_finalize_replaces_table(self, bigquery_config, mock_bq_client, mock_gcs_client): + """A reset must replace the table (WRITE_TRUNCATE), not append to it like a normal incremental.""" + from google.cloud import bigquery + + destination = self._destination(bigquery_config) + + mock_copy = MagicMock() + destination.bq_client.copy_table = mock_copy + destination.bq_client.query = MagicMock() + destination.bq_client.get_table = MagicMock() + destination.bq_client.delete_table = MagicMock() + + assert destination.finalize() is True + + destination.bq_client.query.assert_not_called() + _, kwargs = mock_copy.call_args + assert kwargs["job_config"].write_disposition == bigquery.WriteDisposition.WRITE_TRUNCATE + + def test_stale_temp_table_is_dropped_on_a_fresh_reset(self, bigquery_config, mock_bq_client, mock_gcs_client): + """Rows left by an earlier crashed run must not survive into the replaced table.""" + backend = MagicMock() + backend.get_last_cursor_by_job_id.return_value = None + destination = self._destination(bigquery_config, backend=backend) + destination.bq_client.delete_table = MagicMock() + + destination._ensure_clean_temp_table() + + destination.bq_client.delete_table.assert_called_once_with(destination.temp_table_id, not_found_ok=True) + + def test_temp_table_is_kept_when_resuming_a_crashed_reset(self, bigquery_config, mock_bq_client, mock_gcs_client): + """The producer resumes from the last cursor, so already-written iterations must be kept.""" + backend = MagicMock() + backend.get_last_cursor_by_job_id.return_value = MagicMock(to_source_iteration=4) + destination = self._destination(bigquery_config, backend=backend) + destination.bq_client.delete_table = MagicMock() + + destination._ensure_clean_temp_table() + + destination.bq_client.delete_table.assert_not_called() + + def test_temp_table_is_only_dropped_once(self, bigquery_config, mock_bq_client, mock_gcs_client): + """Both write paths call the guard on every flush; only the first may drop the table.""" + backend = MagicMock() + backend.get_last_cursor_by_job_id.return_value = None + destination = self._destination(bigquery_config, backend=backend) + destination.bq_client.delete_table = MagicMock() + + destination._ensure_clean_temp_table() + destination._ensure_clean_temp_table() + + destination.bq_client.delete_table.assert_called_once() + + def test_non_reset_run_never_drops_its_temp_table(self, bigquery_config, mock_bq_client, mock_gcs_client): + """A plain incremental appends into `_incremental` across runs and must leave it alone.""" + destination = self._destination(bigquery_config, reset=False) + destination.bq_client.delete_table = MagicMock() + + destination._ensure_clean_temp_table() + + assert destination.temp_table_id == f"{destination.table_id}_incremental" + destination.bq_client.delete_table.assert_not_called() diff --git a/tests/engine/backend/test_backend.py b/tests/engine/backend/test_backend.py index 66f6978..2363adc 100644 --- a/tests/engine/backend/test_backend.py +++ b/tests/engine/backend/test_backend.py @@ -336,3 +336,77 @@ def test_number_of_rows_written(backend: SQLAlchemyBackend, session: Session): success=True, ) assert backend.get_number_of_written_rows_for_job(job_id=new_job.id, session=session) == 15 + + +@pytest.mark.parametrize( + "backend,session", + [ + (pytest.lazy_fixture("my_pg_backend"), pytest.lazy_fixture("pg_db_session")), + (pytest.lazy_fixture("my_sqlite_backend"), pytest.lazy_fixture("sqlite_db_session")), + ], +) +def test_stream_reset_request_and_consume(backend: SQLAlchemyBackend, session: Session): + backend.create_all_tables() + + stream_name = f"stream_{uuid.uuid4().hex}" + stream = {"name": "testjob", "source_name": "sourcetest", "stream_name": stream_name} + + assert backend.get_pending_stream_reset(session=session, **stream) is None + + stream_reset = backend.create_stream_reset(session=session, **stream) + + pending = backend.get_pending_stream_reset(session=session, **stream) + assert pending is not None + assert pending.id == stream_reset.id + + job_id = uuid.uuid4().hex + backend.consume_stream_reset(reset_id=stream_reset.id, job_id=job_id, session=session) + + # A consumed request must never be picked up twice, but stays attached to the job running it so a + # crashed reset can be recognised as a reset on retry. + assert backend.get_pending_stream_reset(session=session, **stream) is None + consumed = backend.get_stream_reset_by_job_id(job_id=job_id, session=session) + assert consumed is not None + assert consumed.id == stream_reset.id + + +@pytest.mark.parametrize( + "backend,session", + [ + (pytest.lazy_fixture("my_pg_backend"), pytest.lazy_fixture("pg_db_session")), + (pytest.lazy_fixture("my_sqlite_backend"), pytest.lazy_fixture("sqlite_db_session")), + ], +) +def test_stream_reset_is_scoped_to_its_stream(backend: SQLAlchemyBackend, session: Session): + backend.create_all_tables() + + requested = f"stream_{uuid.uuid4().hex}" + other = f"stream_{uuid.uuid4().hex}" + + backend.create_stream_reset(name="testjob", source_name="sourcetest", stream_name=requested, session=session) + + assert ( + backend.get_pending_stream_reset(name="testjob", source_name="sourcetest", stream_name=other, session=session) + is None + ) + + +@pytest.mark.parametrize( + "backend,session", + [ + (pytest.lazy_fixture("my_pg_backend"), pytest.lazy_fixture("pg_db_session")), + (pytest.lazy_fixture("my_sqlite_backend"), pytest.lazy_fixture("sqlite_db_session")), + ], +) +def test_cancel_pending_stream_resets(backend: SQLAlchemyBackend, session: Session): + backend.create_all_tables() + + stream_name = f"stream_{uuid.uuid4().hex}" + stream = {"name": "testjob", "source_name": "sourcetest", "stream_name": stream_name} + + backend.create_stream_reset(session=session, **stream) + backend.create_stream_reset(session=session, **stream) + + assert backend.cancel_pending_stream_resets(session=session, **stream) == 2 + assert backend.get_pending_stream_reset(session=session, **stream) is None + assert backend.cancel_pending_stream_resets(session=session, **stream) == 0 diff --git a/tests/engine/test_producer_incremental.py b/tests/engine/test_producer_incremental.py index 8a2e4e1..8368c17 100644 --- a/tests/engine/test_producer_incremental.py +++ b/tests/engine/test_producer_incremental.py @@ -1,6 +1,8 @@ import os from datetime import datetime from queue import Queue +from threading import Event +from unittest.mock import MagicMock import pytest from pytz import UTC @@ -10,7 +12,7 @@ from bizon.engine.engine import RunnerFactory from bizon.engine.pipeline.producer import Producer from bizon.source.config import SourceSyncModes -from bizon.source.models import SourceIncrementalState +from bizon.source.models import SourceIncrementalState, SourceIteration @pytest.fixture(scope="function") @@ -130,3 +132,48 @@ def test_source_incremental_state_default_values(): assert state.last_run == now assert state.state == {} assert state.cursor_field is None + + +@pytest.fixture(scope="function") +def started_job(incremental_producer: Producer, sqlite_db_session) -> StreamJob: + """Create the job the producer under test is running.""" + return incremental_producer.backend.create_stream_job( + name=incremental_producer.bizon_config.name, + source_name=incremental_producer.source.config.name, + stream_name=incremental_producer.source.config.stream, + sync_mode=SourceSyncModes.INCREMENTAL.value, + job_status=JobStatus.STARTED, + session=sqlite_db_session, + ) + + +def _stub_source_fetches(producer: Producer): + """Stub both fetch methods with a single terminal iteration, so run() exits after one loop.""" + terminal = SourceIteration(records=[], next_pagination={}) + producer.source.get = MagicMock(return_value=terminal) + producer.source.get_records_after = MagicMock(return_value=terminal) + + +def test_incremental_uses_watermark_without_reset( + incremental_producer: Producer, previous_successful_job: StreamJob, started_job: StreamJob +): + """Baseline: with a previous successful job and no reset, the producer fetches incrementally.""" + _stub_source_fetches(incremental_producer) + + incremental_producer.run(job_id=started_job.id, stop_event=Event()) + + incremental_producer.source.get_records_after.assert_called_once() + incremental_producer.source.get.assert_not_called() + + +def test_reset_ignores_the_watermark( + incremental_producer: Producer, previous_successful_job: StreamJob, started_job: StreamJob +): + """A reset re-fetches the whole stream, even though a watermark is available.""" + incremental_producer.bizon_config.source.reset = True + _stub_source_fetches(incremental_producer) + + incremental_producer.run(job_id=started_job.id, stop_event=Event()) + + incremental_producer.source.get.assert_called_once() + incremental_producer.source.get_records_after.assert_not_called() diff --git a/tests/engine/test_runner_reset.py b/tests/engine/test_runner_reset.py new file mode 100644 index 0000000..c3a29ac --- /dev/null +++ b/tests/engine/test_runner_reset.py @@ -0,0 +1,217 @@ +"""Tests for stream reset resolution in the runner. + +`init_job` is the single place where "is this run a reset?" is decided: it runs in the parent before +the producer and consumer are submitted, and mutates the config both of them are handed. +""" + +import pytest +import yaml + +from bizon.engine.backend.adapters.sqlalchemy.backend import SQLAlchemyBackend +from bizon.engine.backend.models import JobStatus +from bizon.engine.engine import RunnerFactory +from bizon.engine.runner.runner import AbstractRunner +from bizon.source.config import SourceSyncModes + +BIZON_CONFIG_DUMMY_INCREMENTAL = """ +name: test_reset_job + +source: + name: dummy + stream: creatures + sync_mode: incremental + cursor_field: updated_at + authentication: + type: api_key + params: + token: dummy_key + +destination: + name: logger + config: + dummy: dummy + +engine: + backend: + type: sqlite_in_memory + config: + database: not_used + schema: not_used + syncCursorInDBEvery: 400 + runner: + log_level: INFO +""" + + +def build_runner(**source_overrides): + config = yaml.safe_load(BIZON_CONFIG_DUMMY_INCREMENTAL) + config["source"].update(source_overrides) + return RunnerFactory.create_from_config_dict(config) + + +@pytest.fixture(scope="function") +def backend(my_sqlite_backend: SQLAlchemyBackend) -> SQLAlchemyBackend: + my_sqlite_backend.create_all_tables() + return my_sqlite_backend + + +def resolve(runner, backend: SQLAlchemyBackend, resuming_reset: bool = False) -> bool: + return AbstractRunner.resolve_reset( + bizon_config=runner.bizon_config, backend=backend, resuming_reset=resuming_reset + ) + + +def request_reset(runner, backend: SQLAlchemyBackend): + return backend.create_stream_reset( + name=runner.bizon_config.name, + source_name=runner.bizon_config.source.name, + stream_name=runner.bizon_config.source.stream, + ) + + +class TestResolveReset: + def test_plain_incremental_run_is_not_a_reset(self, backend): + assert resolve(build_runner(), backend) is False + + def test_reset_flag_is_honoured(self, backend): + assert resolve(build_runner(reset=True), backend) is True + + def test_pending_request_is_picked_up_without_the_flag(self, backend): + """This is what makes `bizon stream reset` work for a scheduled `bizon run config.yml`.""" + runner = build_runner() + request_reset(runner, backend) + + assert resolve(runner, backend) is True + + def test_in_flight_reset_is_resumed(self, backend): + """A crashed reset must retry as a reset, not degrade into an incremental append.""" + assert resolve(build_runner(), backend, resuming_reset=True) is True + + def test_reset_is_ignored_for_non_incremental_sync_modes(self, backend): + """There is no incremental state to reset, and full refresh already replaces the table.""" + runner = build_runner(reset=True, sync_mode=SourceSyncModes.FULL_REFRESH.value) + + assert resolve(runner, backend) is False + + def test_pending_request_is_ignored_for_non_incremental_sync_modes(self, backend): + runner = build_runner(sync_mode=SourceSyncModes.FULL_REFRESH.value) + request_reset(runner, backend) + + assert resolve(runner, backend) is False + + +class TestBindStreamResetToJob: + """Every reset job must end up with a consumed marker row pointing at it.""" + + def _job(self, runner, backend, status=JobStatus.RUNNING): + return backend.create_stream_job( + name=runner.bizon_config.name, + source_name=runner.bizon_config.source.name, + stream_name=runner.bizon_config.source.stream, + sync_mode=SourceSyncModes.INCREMENTAL.value, + job_status=status, + ) + + def test_flag_path_creates_and_consumes_a_marker(self, backend): + """`--reset` has no marker of its own, so one is created to make the run recoverable.""" + runner = build_runner(reset=True) + job = self._job(runner, backend) + + AbstractRunner.bind_stream_reset_to_job(bizon_config=runner.bizon_config, backend=backend, job_id=job.id) + + assert backend.get_stream_reset_by_job_id(job_id=job.id) is not None + + def test_pending_marker_is_consumed_rather_than_duplicated(self, backend): + runner = build_runner() + stream_reset = request_reset(runner, backend) + job = self._job(runner, backend) + + AbstractRunner.bind_stream_reset_to_job(bizon_config=runner.bizon_config, backend=backend, job_id=job.id) + + bound = backend.get_stream_reset_by_job_id(job_id=job.id) + assert bound.id == stream_reset.id + # Consumed, so the next run does not reset all over again. + assert ( + backend.get_pending_stream_reset( + name=runner.bizon_config.name, + source_name=runner.bizon_config.source.name, + stream_name=runner.bizon_config.source.stream, + ) + is None + ) + + def test_binding_is_idempotent(self, backend): + """Re-running against an already-bound job must not consume a second request.""" + runner = build_runner(reset=True) + job = self._job(runner, backend) + + AbstractRunner.bind_stream_reset_to_job(bizon_config=runner.bizon_config, backend=backend, job_id=job.id) + first = backend.get_stream_reset_by_job_id(job_id=job.id) + + request_reset(runner, backend) + AbstractRunner.bind_stream_reset_to_job(bizon_config=runner.bizon_config, backend=backend, job_id=job.id) + + assert backend.get_stream_reset_by_job_id(job_id=job.id).id == first.id + + +class TestInitJob: + def test_pending_request_flips_the_config_for_producer_and_consumer(self, backend, monkeypatch): + """Producer and consumer are handed these objects, so the flag must land on both.""" + runner = build_runner() + request_reset(runner, backend) + monkeypatch.setattr(AbstractRunner, "get_backend", staticmethod(lambda **kwargs: backend)) + + job = AbstractRunner.init_job(bizon_config=runner.bizon_config, config=runner.config) + + assert runner.bizon_config.source.reset is True + assert runner.config["source"]["reset"] is True + assert backend.get_stream_reset_by_job_id(job_id=job.id) is not None + # The job itself stays incremental so it becomes the next run's watermark. + assert job.sync_mode == SourceSyncModes.INCREMENTAL.value + + def test_plain_run_leaves_the_flag_off(self, backend, monkeypatch): + runner = build_runner() + monkeypatch.setattr(AbstractRunner, "get_backend", staticmethod(lambda **kwargs: backend)) + + job = AbstractRunner.init_job(bizon_config=runner.bizon_config, config=runner.config) + + assert runner.bizon_config.source.reset is False + assert backend.get_stream_reset_by_job_id(job_id=job.id) is None + + def test_reset_starts_a_fresh_job_instead_of_resuming(self, backend, monkeypatch): + """A reset re-fetches from iteration 0, so it must not adopt a half-finished job.""" + runner = build_runner() + running_job = backend.create_stream_job( + name=runner.bizon_config.name, + source_name=runner.bizon_config.source.name, + stream_name=runner.bizon_config.source.stream, + sync_mode=SourceSyncModes.INCREMENTAL.value, + job_status=JobStatus.RUNNING, + ) + request_reset(runner, backend) + monkeypatch.setattr(AbstractRunner, "get_backend", staticmethod(lambda **kwargs: backend)) + + job = AbstractRunner.init_job(bizon_config=runner.bizon_config, config=runner.config) + + assert job.id != running_job.id + assert backend.get_stream_job_by_id(job_id=running_job.id).status == JobStatus.CANCELED + + def test_crashed_reset_resumes_the_same_job(self, backend, monkeypatch): + """The retry has no flag and no pending marker: the bound job is what keeps it a reset.""" + runner = build_runner() + running_job = backend.create_stream_job( + name=runner.bizon_config.name, + source_name=runner.bizon_config.source.name, + stream_name=runner.bizon_config.source.stream, + sync_mode=SourceSyncModes.INCREMENTAL.value, + job_status=JobStatus.RUNNING, + ) + AbstractRunner.bind_stream_reset_to_job( + bizon_config=runner.bizon_config, backend=backend, job_id=running_job.id + ) + monkeypatch.setattr(AbstractRunner, "get_backend", staticmethod(lambda **kwargs: backend)) + + job = AbstractRunner.init_job(bizon_config=runner.bizon_config, config=runner.config) + + assert job.id == running_job.id + assert runner.bizon_config.source.reset is True