From 17ed4b9a0154c736bbc690eb67326c1c4214273c Mon Sep 17 00:00:00 2001 From: Cu pid Date: Tue, 4 Aug 2026 16:47:48 +0100 Subject: [PATCH] feat: add durable metadata and single-node recovery --- CHANGELOG.md | 19 + PRODUCTION_READINESS.md | 19 + README.md | 23 +- loafer/adapters/metadata.py | 1104 +++++++++++++++++ loafer/adapters/metadata_schema.py | 321 +++++ loafer/adapters/object_storage.py | 142 +++ loafer/adapters/runtime.py | 122 ++ loafer/application/__init__.py | 11 + loafer/application/durable.py | 123 ++ loafer/application/service.py | 4 + loafer/cli.py | 56 +- loafer/core/run_state.py | 110 ++ loafer/data_plane.py | 94 +- loafer/engine.py | 5 + loafer/exceptions.py | 16 + loafer/metadata.py | 115 ++ loafer/ports/metadata.py | 140 +++ loafer/ports/object_storage.py | 32 + loafer/scheduler.py | 32 +- loafer/worker.py | 175 +++ .../references/architecture.md | 15 +- tests/e2e/test_durable_recovery.py | 161 +++ tests/integration/test_metadata_store.py | 80 ++ tests/unit/test_metadata_store.py | 312 +++++ tests/unit/test_object_storage.py | 38 + tests/unit/test_run_state.py | 42 + tests/unit/test_scheduler.py | 18 + 27 files changed, 3288 insertions(+), 41 deletions(-) create mode 100644 loafer/adapters/metadata.py create mode 100644 loafer/adapters/metadata_schema.py create mode 100644 loafer/adapters/object_storage.py create mode 100644 loafer/application/durable.py create mode 100644 loafer/core/run_state.py create mode 100644 loafer/metadata.py create mode 100644 loafer/ports/metadata.py create mode 100644 loafer/ports/object_storage.py create mode 100644 loafer/worker.py create mode 100644 tests/e2e/test_durable_recovery.py create mode 100644 tests/integration/test_metadata_store.py create mode 100644 tests/unit/test_metadata_store.py create mode 100644 tests/unit/test_object_storage.py create mode 100644 tests/unit/test_run_state.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 86e1fdc..1779032 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,14 @@ Notable changes to Loafer are documented here. This project follows ### Added +- Versioned SQLite/PostgreSQL metadata migrations for immutable pipeline versions, runs, stages, + partitions, batches, checkpoints, events, artifacts, schedules, and transactional outbox rows. +- Explicit run, stage, and batch state machines; idempotent run/schedule/cancel commands; leases, + monotonic fencing tokens, heartbeats, cooperative cancellation, and classified retries. +- Filesystem and in-memory object-storage adapters behind a shared port for logs, documents, + generated artifacts, and replayable temporary batch output. +- A separately runnable durable worker plus `loafer enqueue` and `loafer worker` commands, with + batch-artifact replay from the last committed source position after a worker crash. - A framework-independent application service with strict, JSON-roundtrippable contracts for run requests, execution plans, batch envelopes, events, snapshots, and results. - Runtime ports and local adapters for cancellation, checkpoints, secret resolution, event @@ -23,6 +31,11 @@ Notable changes to Loafer are documented here. This project follows ### Changed +- Scheduled callbacks now create durable idempotent run commands; pipeline execution happens only + in a separately started worker process. +- Bounded durable runs stage each transformed batch as an immutable object and commit its metadata, + checkpoint, event, and outbox record under the active fencing token before attempt-local target + writes and final publication. - The CLI, scheduler, and legacy Python runner now share the same application boundary while core execution orchestration remains independent of client frameworks. - Durable application contracts now exclude credentials, connector instances, iterators, provider @@ -44,6 +57,12 @@ Notable changes to Loafer are documented here. This project follows ### Known limitations +- SQLite metadata is restricted to the embedded profile with one scheduler and one worker; it does + not advertise high availability or distributed claims. PostgreSQL is the authoritative platform + profile. The bundled object-storage adapter is local filesystem storage, not a distributed blob + store. +- Durable checkpoint replay currently applies to declared row-local, single-partition runs with a + stable offset-ordered source. Materialized/global transforms still restart as whole runs. - MongoDB row-local runs remain rejected until a tested staging/merge protocol replaces direct partial batch effects. PostgreSQL append is intentionally at-least-once across an ambiguous target-commit/checkpoint gap; keyed upsert is the replay-safe merge mode. diff --git a/PRODUCTION_READINESS.md b/PRODUCTION_READINESS.md index 85d5b6f..8f49e07 100644 --- a/PRODUCTION_READINESS.md +++ b/PRODUCTION_READINESS.md @@ -432,6 +432,25 @@ Exit gate: - migrations work from an empty database and the previous supported schema; - SQLite and PostgreSQL contract tests pass for the capabilities each profile advertises. +**Current status:** implementation and single-node recovery verification complete for declared +row-local, single-partition runs. Versioned migrations persist immutable pipeline versions, runs, +stages, partitions, batches, events, artifacts, checkpoints, schedules, and outbox records behind +one SQLAlchemy metadata interface. PostgreSQL uses row locks for concurrent claims and is the +authoritative platform profile; SQLite is explicitly restricted to one scheduler and one worker +without HA or NATS claims. Run, stage, and batch state machines reject impossible transitions, +events allocate contiguous per-run sequences, run creation and schedule firing are idempotent, and +every worker mutation requires its active lease and fencing token. + +The bounded data plane now stages transformed batches in object storage before atomically +committing batch metadata, checkpoint, event, and outbox state. Worker-kill tests terminate after +each of three batch commit boundaries and prove that a newly fenced worker replays committed +artifacts, skips the durable source offset, and publishes the exact output. Empty-database, +previous-schema upgrade, rollback/re-apply, constraint, cancellation, schedule, and stale-token +contracts pass on SQLite; the migration, event, claim, and fencing contracts also pass against a +live PostgreSQL 16 database. Recovery remains limited to stable offset-ordered, row-local, +single-partition execution. Materialized/global transforms still restart as whole runs, and the +bundled filesystem object adapter is not a distributed object store. + ### Phase 4 — Build authentication, tenancy, and the control-plane API **Goal:** expose safe multi-tenant application use cases without running data work in HTTP diff --git a/README.md b/README.md index bd84204..21b302f 100644 --- a/README.md +++ b/README.md @@ -7,9 +7,10 @@ Define a source, transformation, and target; validate the pipeline; then run it or from a scheduler. Transformations can use SQL, custom Python, multi-step pipelines, or optional LLM-generated artifacts. -> **Project status:** Loafer currently ships as a CLI engine with a local scheduler/daemon. The -> multi-tenant API, distributed workers, web operations dashboard, and terminal dashboard are under -> active development. The `/studio` web route is a product preview, not a connected control plane. +> **Project status:** Loafer ships as a CLI engine with durable single-node scheduling and worker +> recovery. The multi-tenant API, distributed queue/workers, web operations dashboard, and terminal +> dashboard are under active development. The `/studio` web route is a product preview, not a +> connected control plane. ## What works today @@ -20,6 +21,8 @@ LLM-generated artifacts. - PostgreSQL and MongoDB upserts - Cursor-based incremental extraction with local state - Local scheduling, daemon management, run summaries, and logs +- SQLite/PostgreSQL run metadata, fenced worker leases, durable batch checkpoints, replayable + temporary output, and monotonic run events - Optional Gemini, OpenAI, Claude, and Qwen providers - Resource-limited Python transform subprocesses on Linux and macOS - Declared row-local ETL with bounded batches, per-batch validation, schema policies, @@ -247,6 +250,8 @@ implemented; `ocr_applied` remains `false` in provenance. ```text loafer run +loafer enqueue --command-key +loafer worker [--once] loafer validate loafer connectors loafer schedule @@ -260,6 +265,11 @@ loafer init Use `loafer --help` for command-specific options. +`loafer schedule` and the scheduler daemon only enqueue durable run commands; start `loafer worker` +as a separate process to execute them. The embedded profile defaults to SQLite under `~/.loafer` +and supports one scheduler and one worker. Set `LOAFER_METADATA_URL` to a PostgreSQL URL for the +authoritative platform store and `LOAFER_OBJECTS_PATH` to choose the local artifact root. + ## Self-hosted platform direction The production architecture separates clients, control plane, and data plane: @@ -275,9 +285,10 @@ The web dashboard and planned terminal dashboard will use the same API, permissi metrics, and logs. Workers will run independently so startups can deploy the stack on one host while larger installations can scale and isolate worker pools. -The full stack is not shipped yet. Until the API, durable queue, tenant authorization, worker -leases, and recovery tests exist, use the CLI/Docker path for bounded workloads and do not expose -Studio as a production operations surface. +The full stack is not shipped yet. Durable metadata, leases, fencing, outbox records, and +single-node bounded-batch recovery are implemented; the authenticated API, distributed queue, +tenant authorization, and distributed object store are not. Use the CLI/Docker path for bounded +workloads and do not expose Studio as a production operations surface. The planned web source uses Crawlee for Python with HTTP/Parsel and Playwright execution profiles. It will support bounded crawling, authorized authenticated sessions, JavaScript rendering, and diff --git a/loafer/adapters/metadata.py b/loafer/adapters/metadata.py new file mode 100644 index 0000000..e20bc26 --- /dev/null +++ b/loafer/adapters/metadata.py @@ -0,0 +1,1104 @@ +"""SQLAlchemy metadata adapter shared by PostgreSQL and local SQLite.""" + +from __future__ import annotations + +import hashlib +import uuid +from datetime import UTC, datetime, timedelta +from typing import Any + +from sqlalchemy import and_, event, insert, or_, select, update +from sqlalchemy.engine import Connection, Engine, RowMapping, create_engine +from sqlalchemy.exc import IntegrityError + +from loafer.adapters import metadata_schema as schema +from loafer.contracts import BatchEnvelope, Checkpoint +from loafer.core.run_state import ( + BatchState, + RetryCategory, + RunState, + StageState, + require_batch_transition, + require_run_transition, + require_stage_transition, +) +from loafer.exceptions import IdempotencyConflictError, MetadataError, StaleFenceError +from loafer.metadata import ( + BatchCommit, + OutboxRecord, + PipelineVersion, + RunLease, + RunRecord, + ScheduleRecord, + StoredArtifact, + StoredEvent, + utc_now, +) + +_ACTIVE_STATES = (RunState.CLAIMED.value, RunState.RUNNING.value, RunState.CANCELLING.value) + + +class SqlMetadataStore: + """Authoritative durable state with PostgreSQL and restricted SQLite profiles. + + PostgreSQL uses row locks for concurrent claims and event allocation. SQLite is + intentionally advertised only for one scheduler and one worker process. + """ + + def __init__( + self, + url: str, + *, + clock: Any = utc_now, + engine: Engine | None = None, + ) -> None: + self._clock = clock + self._engine = engine or create_engine(url, future=True) + if self._engine.dialect.name == "sqlite": + event.listen(self._engine, "connect", _enable_sqlite_foreign_keys) + + @property + def engine(self) -> Engine: + return self._engine + + @property + def profile(self) -> str: + return self._engine.dialect.name + + def close(self) -> None: + self._engine.dispose() + + def migrate(self, target_version: int | None = None) -> int: + return schema.migrate(self._engine, target_version) + + def register_pipeline_version( + self, + *, + workspace_id: str, + pipeline_key: str, + config_digest: str, + config: dict[str, Any], + ) -> PipelineVersion: + version_id = hashlib.sha256( + f"{workspace_id}\0{pipeline_key}\0{config_digest}".encode() + ).hexdigest()[:32] + now = self._clock() + with self._engine.begin() as connection: + row = ( + connection.execute( + select(schema.pipeline_versions).where( + schema.pipeline_versions.c.id == version_id + ) + ) + .mappings() + .one_or_none() + ) + if row is None: + connection.execute( + insert(schema.pipeline_versions).values( + id=version_id, + workspace_id=workspace_id, + pipeline_key=pipeline_key, + config_digest=config_digest, + config_json=config, + created_at=now, + ) + ) + row = ( + connection.execute( + select(schema.pipeline_versions).where( + schema.pipeline_versions.c.id == version_id + ) + ) + .mappings() + .one() + ) + elif row["config_json"] != config: + raise IdempotencyConflictError( + "pipeline version digest was reused for different configuration" + ) + return _pipeline_version(row) + + def get_pipeline_version(self, version_id: str) -> PipelineVersion: + with self._engine.connect() as connection: + row = ( + connection.execute( + select(schema.pipeline_versions).where( + schema.pipeline_versions.c.id == version_id + ) + ) + .mappings() + .one_or_none() + ) + if row is None: + raise MetadataError(f"pipeline version not found: {version_id}") + return _pipeline_version(row) + + def create_run( + self, + *, + workspace_id: str, + pipeline_version_id: str, + command_key: str, + run_id: str | None = None, + parent_run_id: str | None = None, + retry_category: RetryCategory | None = None, + ) -> RunRecord: + now = self._clock() + with self._engine.begin() as connection: + return self._create_run( + connection, + workspace_id=workspace_id, + pipeline_version_id=pipeline_version_id, + command_key=command_key, + run_id=run_id or uuid.uuid4().hex, + parent_run_id=parent_run_id, + retry_category=retry_category, + now=now, + ) + + def get_run(self, run_id: str) -> RunRecord: + with self._engine.connect() as connection: + return _run_record(self._run_row(connection, run_id)) + + def claim_run(self, worker_id: str, lease_for: timedelta) -> RunLease | None: + if lease_for <= timedelta(0): + raise ValueError("lease_for must be positive") + now = self._clock() + expires_at = now + lease_for + runnable = or_( + schema.runs.c.state == RunState.QUEUED.value, + and_( + schema.runs.c.state == RunState.RETRY_WAIT.value, + schema.runs.c.retry_at <= now, + ), + and_( + schema.runs.c.state.in_(_ACTIVE_STATES), + schema.runs.c.lease_expires_at <= now, + ), + ) + with self._engine.begin() as connection: + query = ( + select(schema.runs) + .where(runnable) + .order_by(schema.runs.c.created_at, schema.runs.c.id) + .limit(1) + ) + if self.profile == "postgresql": + query = query.with_for_update(skip_locked=True) + row = connection.execute(query).mappings().one_or_none() + if row is None: + return None + current_state = RunState(row["state"]) + if current_state == RunState.RETRY_WAIT: + require_run_transition(current_state, RunState.QUEUED) + current_state = RunState.QUEUED + elif current_state in { + RunState.CLAIMED, + RunState.RUNNING, + RunState.CANCELLING, + }: + require_run_transition(current_state, RunState.RETRY_WAIT) + require_run_transition(RunState.RETRY_WAIT, RunState.QUEUED) + current_state = RunState.QUEUED + require_run_transition(current_state, RunState.CLAIMED) + token = int(row["fencing_token"]) + 1 + attempt = int(row["attempt"]) + if row["state"] in _ACTIVE_STATES: + attempt += 1 + result = connection.execute( + update(schema.runs) + .where( + schema.runs.c.id == row["id"], + schema.runs.c.fencing_token == row["fencing_token"], + ) + .values( + state=RunState.CLAIMED.value, + attempt=attempt, + retry_category=( + RetryCategory.INFRASTRUCTURE.value + if row["state"] in _ACTIVE_STATES + else row["retry_category"] + ), + retry_at=None, + lease_owner=worker_id, + fencing_token=token, + lease_expires_at=expires_at, + heartbeat_at=now, + finished_at=None, + ) + ) + if result.rowcount != 1: + return None + claimed = self._run_row(connection, row["id"]) + self._append_event( + connection, + claimed, + "run.claimed", + {"worker_id": worker_id, "fencing_token": token, "attempt": attempt}, + now, + ) + run = _run_record(claimed) + return RunLease( + run=run, + worker_id=worker_id, + fencing_token=token, + expires_at=expires_at, + ) + + def heartbeat(self, lease: RunLease, lease_for: timedelta) -> RunLease: + if lease_for <= timedelta(0): + raise ValueError("lease_for must be positive") + now = self._clock() + expires_at = now + lease_for + with self._engine.begin() as connection: + row = self._require_fence(connection, lease, now) + result = connection.execute( + update(schema.runs) + .where( + schema.runs.c.id == lease.run.id, + schema.runs.c.fencing_token == lease.fencing_token, + schema.runs.c.lease_owner == lease.worker_id, + ) + .values(heartbeat_at=now, lease_expires_at=expires_at) + ) + if result.rowcount != 1: + raise StaleFenceError(f"stale fencing token for run {lease.run.id}") + updated = _run_record(self._run_row(connection, row["id"])) + return RunLease(updated, lease.worker_id, lease.fencing_token, expires_at) + + def transition_run( + self, + lease: RunLease, + target: RunState, + *, + error: dict[str, Any] | None = None, + retry_category: RetryCategory | None = None, + retry_at: datetime | None = None, + ) -> RunRecord: + now = self._clock() + with self._engine.begin() as connection: + row = self._require_fence(connection, lease, now) + current = RunState(row["state"]) + require_run_transition(current, target) + values: dict[str, Any] = {"state": target.value} + if target == RunState.RUNNING and row["started_at"] is None: + values["started_at"] = now + if target in {RunState.SUCCEEDED, RunState.FAILED, RunState.CANCELLED}: + values.update( + finished_at=now, + lease_owner=None, + lease_expires_at=None, + heartbeat_at=None, + ) + if target == RunState.RETRY_WAIT: + if retry_at is None or retry_category is None: + raise MetadataError("retry transitions require retry_at and retry_category") + values.update( + retry_at=retry_at, + retry_category=retry_category.value, + lease_owner=None, + lease_expires_at=None, + heartbeat_at=None, + ) + if error is not None: + values["error_json"] = error + result = connection.execute( + update(schema.runs) + .where( + schema.runs.c.id == lease.run.id, + schema.runs.c.fencing_token == lease.fencing_token, + ) + .values(**values) + ) + if result.rowcount != 1: + raise StaleFenceError(f"stale fencing token for run {lease.run.id}") + updated = self._run_row(connection, lease.run.id) + self._append_event( + connection, + updated, + f"run.{target.value}", + {"state": target.value, "error": error, "retry_category": retry_category}, + now, + ) + return _run_record(self._run_row(connection, lease.run.id)) + + def transition_stage( + self, + lease: RunLease, + stage_name: str, + target: StageState, + ) -> None: + now = self._clock() + with self._engine.begin() as connection: + run = self._require_fence(connection, lease, now) + stage_id = _stage_id(lease.run.id, stage_name, int(run["attempt"])) + row = ( + connection.execute(select(schema.stages).where(schema.stages.c.id == stage_id)) + .mappings() + .one_or_none() + ) + if row is None: + connection.execute( + insert(schema.stages).values( + id=stage_id, + run_id=lease.run.id, + name=stage_name, + state=StageState.PENDING.value, + attempt=run["attempt"], + created_at=now, + ) + ) + row = ( + connection.execute(select(schema.stages).where(schema.stages.c.id == stage_id)) + .mappings() + .one() + ) + current = StageState(row["state"]) + require_stage_transition(current, target) + values: dict[str, Any] = {"state": target.value} + if target == StageState.RUNNING and row["started_at"] is None: + values["started_at"] = now + if target in { + StageState.SUCCEEDED, + StageState.FAILED, + StageState.CANCELLED, + StageState.SKIPPED, + }: + values["finished_at"] = now + connection.execute( + update(schema.stages).where(schema.stages.c.id == stage_id).values(**values) + ) + self._append_event( + connection, + run, + f"stage.{target.value}", + {"stage": stage_name, "state": target.value}, + now, + stage_id=stage_id, + ) + + def append_event( + self, + lease: RunLease, + event_type: str, + payload: dict[str, Any], + ) -> StoredEvent: + now = self._clock() + with self._engine.begin() as connection: + run = self._require_fence(connection, lease, now) + return self._append_event(connection, run, event_type, payload, now) + + def commit_batch( + self, + lease: RunLease, + envelope: BatchEnvelope, + checkpoint: Checkpoint, + artifact: StoredArtifact, + ) -> BatchCommit: + if envelope.run_id != lease.run.id or checkpoint.run_id != lease.run.id: + raise MetadataError("batch, checkpoint, and lease must belong to the same run") + if checkpoint.batch_id != envelope.batch_id: + raise MetadataError("checkpoint batch_id must match the batch envelope") + now = self._clock() + with self._engine.begin() as connection: + run = self._require_fence(connection, lease, now) + stage_id = _stage_id(lease.run.id, envelope.stage_id, int(run["attempt"])) + self._ensure_stage(connection, run, stage_id, envelope.stage_id, now) + partition_id = _partition_id(lease.run.id, envelope.partition_id) + self._ensure_partition( + connection, + lease.run.id, + stage_id, + partition_id, + envelope.partition_id, + now, + ) + batch_id = _batch_id(lease.run.id, envelope.partition_id, envelope.batch_id) + existing = ( + connection.execute(select(schema.batches).where(schema.batches.c.id == batch_id)) + .mappings() + .one_or_none() + ) + if existing is not None: + saved = BatchEnvelope.model_validate(existing["envelope_json"]) + if saved.output_checksum != envelope.output_checksum: + raise IdempotencyConflictError( + f"batch {envelope.batch_id} was recommitted with different output" + ) + return self._batch_commit(connection, existing) + + self._ensure_artifact(connection, artifact) + envelope_json = envelope.model_dump(mode="json") + connection.execute( + insert(schema.batches).values( + id=batch_id, + run_id=lease.run.id, + partition_id=partition_id, + batch_key=envelope.batch_id, + state=BatchState.PENDING.value, + attempt=envelope.attempt, + envelope_json=envelope_json, + artifact_id=None, + created_at=now, + ) + ) + require_batch_transition(BatchState.PENDING, BatchState.RUNNING) + connection.execute( + update(schema.batches) + .where(schema.batches.c.id == batch_id) + .values(state=BatchState.RUNNING.value) + ) + require_batch_transition(BatchState.RUNNING, BatchState.COMMITTED) + connection.execute( + update(schema.batches) + .where(schema.batches.c.id == batch_id) + .values( + state=BatchState.COMMITTED.value, + artifact_id=artifact.id, + committed_at=checkpoint.committed_at, + ) + ) + connection.execute( + insert(schema.checkpoints).values( + id=checkpoint.checkpoint_id, + run_id=lease.run.id, + partition_id=partition_id, + batch_id=envelope.batch_id, + source_position_json=checkpoint.source_position, + fencing_token=lease.fencing_token, + committed_at=checkpoint.committed_at, + ) + ) + event = self._append_event( + connection, + run, + "batch.committed", + { + "batch_id": envelope.batch_id, + "partition_id": envelope.partition_id, + "checkpoint_id": checkpoint.checkpoint_id, + "artifact_id": artifact.id, + "rows_in": envelope.rows_in, + "rows_out": envelope.rows_out, + }, + now, + stage_id=stage_id, + batch_id=envelope.batch_id, + ) + connection.execute( + insert(schema.outbox).values( + id=uuid.uuid4().hex, + aggregate_type="run", + aggregate_id=lease.run.id, + event_type="batch.committed", + payload_json={ + "run_id": lease.run.id, + "sequence": event.sequence, + "batch_id": envelope.batch_id, + }, + available_at=now, + attempts=0, + ) + ) + row = ( + connection.execute(select(schema.batches).where(schema.batches.c.id == batch_id)) + .mappings() + .one() + ) + return self._batch_commit(connection, row) + + def list_batch_commits(self, run_id: str, partition_id: str) -> list[BatchCommit]: + internal_partition_id = _partition_id(run_id, partition_id) + with self._engine.connect() as connection: + rows = connection.execute( + select(schema.batches) + .where( + schema.batches.c.run_id == run_id, + schema.batches.c.partition_id == internal_partition_id, + schema.batches.c.state == BatchState.COMMITTED.value, + ) + .order_by(schema.batches.c.committed_at, schema.batches.c.batch_key) + ).mappings() + return [self._batch_commit(connection, row) for row in rows] + + def latest_checkpoint(self, run_id: str, partition_id: str) -> Checkpoint | None: + internal_partition_id = _partition_id(run_id, partition_id) + with self._engine.connect() as connection: + row = ( + connection.execute( + select(schema.checkpoints) + .where( + schema.checkpoints.c.run_id == run_id, + schema.checkpoints.c.partition_id == internal_partition_id, + ) + .order_by(schema.checkpoints.c.committed_at.desc()) + .limit(1) + ) + .mappings() + .one_or_none() + ) + return _checkpoint(row, partition_id) if row is not None else None + + def request_cancel(self, run_id: str) -> RunRecord: + now = self._clock() + with self._engine.begin() as connection: + row = self._run_row(connection, run_id, for_update=True) + state = RunState(row["state"]) + if state in {RunState.SUCCEEDED, RunState.FAILED, RunState.CANCELLED}: + return _run_record(row) + values: dict[str, Any] = {"cancel_requested": True} + if state in {RunState.QUEUED, RunState.RETRY_WAIT}: + require_run_transition(state, RunState.CANCELLED) + values.update(state=RunState.CANCELLED.value, finished_at=now) + elif state in {RunState.CLAIMED, RunState.RUNNING}: + require_run_transition(state, RunState.CANCELLING) + values["state"] = RunState.CANCELLING.value + connection.execute( + update(schema.runs).where(schema.runs.c.id == run_id).values(**values) + ) + updated = self._run_row(connection, run_id) + self._append_event( + connection, + updated, + "run.cancel_requested", + {"state": updated["state"]}, + now, + ) + return _run_record(self._run_row(connection, run_id)) + + def cancellation_requested(self, run_id: str) -> bool: + with self._engine.connect() as connection: + value = connection.execute( + select(schema.runs.c.cancel_requested).where(schema.runs.c.id == run_id) + ).scalar_one_or_none() + if value is None: + raise MetadataError(f"run not found: {run_id}") + return bool(value) + + def upsert_schedule(self, schedule: ScheduleRecord) -> ScheduleRecord: + with self._engine.begin() as connection: + existing = ( + connection.execute( + select(schema.schedules).where(schema.schedules.c.id == schedule.id) + ) + .mappings() + .one_or_none() + ) + values = _schedule_values(schedule) + if existing is None: + connection.execute(insert(schema.schedules).values(**values)) + else: + if existing["workspace_id"] != schedule.workspace_id: + raise IdempotencyConflictError("schedule id belongs to another workspace") + connection.execute( + update(schema.schedules) + .where(schema.schedules.c.id == schedule.id) + .values(**values) + ) + row = ( + connection.execute( + select(schema.schedules).where(schema.schedules.c.id == schedule.id) + ) + .mappings() + .one() + ) + return _schedule(row) + + def enqueue_due_schedules(self, now: datetime) -> list[RunRecord]: + created: list[RunRecord] = [] + with self._engine.begin() as connection: + query = ( + select(schema.schedules) + .where( + schema.schedules.c.enabled.is_(True), + schema.schedules.c.next_run_at <= now, + ) + .order_by(schema.schedules.c.next_run_at, schema.schedules.c.id) + ) + if self.profile == "postgresql": + query = query.with_for_update(skip_locked=True) + for item in connection.execute(query).mappings(): + due_at = _as_utc(item["next_run_at"]) + command_key = f"schedule:{item['id']}:{due_at.isoformat()}" + run_id = hashlib.sha256(command_key.encode()).hexdigest()[:24] + created.append( + self._create_run( + connection, + workspace_id=item["workspace_id"], + pipeline_version_id=item["pipeline_version_id"], + command_key=command_key, + run_id=run_id, + parent_run_id=None, + retry_category=None, + now=now, + ) + ) + connection.execute( + update(schema.schedules) + .where(schema.schedules.c.id == item["id"]) + .values(next_run_at=_next_fire(item, due_at), updated_at=now) + ) + return created + + def list_events(self, run_id: str, after: int = 0) -> list[StoredEvent]: + with self._engine.connect() as connection: + rows = connection.execute( + select(schema.run_events) + .where( + schema.run_events.c.run_id == run_id, + schema.run_events.c.sequence > after, + ) + .order_by(schema.run_events.c.sequence) + ).mappings() + return [_stored_event(row) for row in rows] + + def pending_outbox(self, limit: int = 100) -> list[OutboxRecord]: + with self._engine.connect() as connection: + rows = connection.execute( + select(schema.outbox) + .where( + schema.outbox.c.published_at.is_(None), + schema.outbox.c.available_at <= self._clock(), + ) + .order_by(schema.outbox.c.available_at, schema.outbox.c.id) + .limit(limit) + ).mappings() + return [_outbox(row) for row in rows] + + def mark_outbox_published(self, outbox_id: str, published_at: datetime) -> None: + with self._engine.begin() as connection: + connection.execute( + update(schema.outbox) + .where( + schema.outbox.c.id == outbox_id, + schema.outbox.c.published_at.is_(None), + ) + .values(published_at=published_at) + ) + + def _create_run( + self, + connection: Connection, + *, + workspace_id: str, + pipeline_version_id: str, + command_key: str, + run_id: str, + parent_run_id: str | None, + retry_category: RetryCategory | None, + now: datetime, + ) -> RunRecord: + existing = ( + connection.execute( + select(schema.runs).where( + schema.runs.c.workspace_id == workspace_id, + schema.runs.c.command_key == command_key, + ) + ) + .mappings() + .one_or_none() + ) + if existing is not None: + if existing["pipeline_version_id"] != pipeline_version_id: + raise IdempotencyConflictError( + "run command key was reused for another pipeline version" + ) + return _run_record(existing) + try: + connection.execute( + insert(schema.runs).values( + id=run_id, + workspace_id=workspace_id, + pipeline_version_id=pipeline_version_id, + command_key=command_key, + state=RunState.QUEUED.value, + attempt=0, + retry_category=retry_category.value if retry_category else None, + cancel_requested=False, + next_event_sequence=1, + fencing_token=0, + created_at=now, + parent_run_id=parent_run_id, + ) + ) + except IntegrityError as exc: + raise MetadataError(f"could not create run {run_id}: {exc}") from exc + row = self._run_row(connection, run_id) + event = self._append_event( + connection, + row, + "run.created", + {"pipeline_version_id": pipeline_version_id}, + now, + ) + connection.execute( + insert(schema.outbox).values( + id=uuid.uuid4().hex, + aggregate_type="run", + aggregate_id=run_id, + event_type="run.created", + payload_json={"run_id": run_id, "sequence": event.sequence}, + available_at=now, + attempts=0, + ) + ) + return _run_record(self._run_row(connection, run_id)) + + def _run_row( + self, + connection: Connection, + run_id: str, + *, + for_update: bool = False, + ) -> RowMapping: + query = select(schema.runs).where(schema.runs.c.id == run_id) + if for_update and self.profile == "postgresql": + query = query.with_for_update() + row = connection.execute(query).mappings().one_or_none() + if row is None: + raise MetadataError(f"run not found: {run_id}") + return row + + def _require_fence( + self, + connection: Connection, + lease: RunLease, + now: datetime, + ) -> RowMapping: + row = self._run_row(connection, lease.run.id, for_update=True) + expires_at = row["lease_expires_at"] + valid = ( + row["lease_owner"] == lease.worker_id + and int(row["fencing_token"]) == lease.fencing_token + and row["state"] in _ACTIVE_STATES + and expires_at is not None + and _as_utc(expires_at) > now + ) + if not valid: + raise StaleFenceError( + f"stale or expired fencing token {lease.fencing_token} for run {lease.run.id}" + ) + return row + + def _append_event( + self, + connection: Connection, + run: RowMapping, + event_type: str, + payload: dict[str, Any], + occurred_at: datetime, + *, + stage_id: str | None = None, + batch_id: str | None = None, + ) -> StoredEvent: + sequence = int(run["next_event_sequence"]) + event_id = _stable_id("event", str(run["id"]), str(sequence)) + connection.execute( + insert(schema.run_events).values( + id=event_id, + run_id=run["id"], + sequence=sequence, + event_type=event_type, + stage_id=stage_id, + batch_id=batch_id, + payload_json=_json(payload), + occurred_at=occurred_at, + ) + ) + result = connection.execute( + update(schema.runs) + .where( + schema.runs.c.id == run["id"], + schema.runs.c.next_event_sequence == sequence, + ) + .values(next_event_sequence=sequence + 1) + ) + if result.rowcount != 1: + raise MetadataError(f"event sequence allocation raced for run {run['id']}") + return StoredEvent(run["id"], sequence, event_type, _json(payload), occurred_at) + + @staticmethod + def _ensure_stage( + connection: Connection, + run: RowMapping, + stage_id: str, + stage_name: str, + now: datetime, + ) -> None: + exists = connection.execute( + select(schema.stages.c.id).where(schema.stages.c.id == stage_id) + ).scalar_one_or_none() + if exists is None: + require_stage_transition(StageState.PENDING, StageState.RUNNING) + connection.execute( + insert(schema.stages).values( + id=stage_id, + run_id=run["id"], + name=stage_name, + state=StageState.RUNNING.value, + attempt=run["attempt"], + created_at=now, + started_at=now, + ) + ) + + @staticmethod + def _ensure_partition( + connection: Connection, + run_id: str, + stage_id: str, + partition_id: str, + partition_key: str, + now: datetime, + ) -> None: + exists = connection.execute( + select(schema.partitions.c.id).where(schema.partitions.c.id == partition_id) + ).scalar_one_or_none() + if exists is None: + connection.execute( + insert(schema.partitions).values( + id=partition_id, + run_id=run_id, + stage_id=stage_id, + partition_key=partition_key, + created_at=now, + ) + ) + + @staticmethod + def _ensure_artifact(connection: Connection, artifact: StoredArtifact) -> None: + row = ( + connection.execute(select(schema.artifacts).where(schema.artifacts.c.id == artifact.id)) + .mappings() + .one_or_none() + ) + if row is None: + connection.execute( + insert(schema.artifacts).values( + id=artifact.id, + run_id=artifact.run_id, + kind=artifact.kind, + uri=artifact.uri, + checksum=artifact.checksum, + size_bytes=artifact.size_bytes, + metadata_json=artifact.metadata, + created_at=artifact.created_at, + ) + ) + elif row["checksum"] != artifact.checksum: + raise IdempotencyConflictError(f"artifact id {artifact.id} has conflicting content") + + @staticmethod + def _batch_commit(connection: Connection, row: RowMapping) -> BatchCommit: + envelope = BatchEnvelope.model_validate(row["envelope_json"]) + checkpoint_row = ( + connection.execute( + select(schema.checkpoints).where( + schema.checkpoints.c.run_id == row["run_id"], + schema.checkpoints.c.partition_id == row["partition_id"], + schema.checkpoints.c.batch_id == row["batch_key"], + ) + ) + .mappings() + .one() + ) + artifact_row = ( + connection.execute( + select(schema.artifacts).where(schema.artifacts.c.id == row["artifact_id"]) + ) + .mappings() + .one() + ) + return BatchCommit( + envelope, + _checkpoint(checkpoint_row, envelope.partition_id), + _artifact(artifact_row), + ) + + +def _enable_sqlite_foreign_keys(dbapi_connection: Any, _connection_record: Any) -> None: + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + + +def _pipeline_version(row: RowMapping) -> PipelineVersion: + return PipelineVersion( + id=row["id"], + workspace_id=row["workspace_id"], + pipeline_key=row["pipeline_key"], + config_digest=row["config_digest"], + config=dict(row["config_json"]), + created_at=_as_utc(row["created_at"]), + ) + + +def _run_record(row: RowMapping) -> RunRecord: + retry = row["retry_category"] + return RunRecord( + id=row["id"], + workspace_id=row["workspace_id"], + pipeline_version_id=row["pipeline_version_id"], + command_key=row["command_key"], + state=RunState(row["state"]), + attempt=int(row["attempt"]), + retry_category=RetryCategory(retry) if retry else None, + cancel_requested=bool(row["cancel_requested"]), + fencing_token=int(row["fencing_token"]), + lease_owner=row["lease_owner"], + lease_expires_at=_optional_utc(row["lease_expires_at"]), + heartbeat_at=_optional_utc(row["heartbeat_at"]), + created_at=_as_utc(row["created_at"]), + started_at=_optional_utc(row["started_at"]), + finished_at=_optional_utc(row["finished_at"]), + parent_run_id=row["parent_run_id"], + error=dict(row["error_json"]) if row["error_json"] else None, + ) + + +def _stored_event(row: RowMapping) -> StoredEvent: + return StoredEvent( + run_id=row["run_id"], + sequence=int(row["sequence"]), + event_type=row["event_type"], + payload=dict(row["payload_json"]), + occurred_at=_as_utc(row["occurred_at"]), + ) + + +def _artifact(row: RowMapping) -> StoredArtifact: + return StoredArtifact( + id=row["id"], + run_id=row["run_id"], + kind=row["kind"], + uri=row["uri"], + checksum=row["checksum"], + size_bytes=int(row["size_bytes"]), + metadata=dict(row["metadata_json"]), + created_at=_as_utc(row["created_at"]), + ) + + +def _checkpoint(row: RowMapping, partition_key: str) -> Checkpoint: + return Checkpoint( + checkpoint_id=row["id"], + run_id=row["run_id"], + partition_id=partition_key, + batch_id=row["batch_id"], + source_position=row["source_position_json"], + committed_at=_as_utc(row["committed_at"]), + ) + + +def _schedule(row: RowMapping) -> ScheduleRecord: + return ScheduleRecord( + id=row["id"], + workspace_id=row["workspace_id"], + pipeline_version_id=row["pipeline_version_id"], + trigger_kind=row["trigger_kind"], + trigger_spec=row["trigger_spec"], + timezone=row["timezone"], + enabled=bool(row["enabled"]), + next_run_at=_as_utc(row["next_run_at"]), + created_at=_as_utc(row["created_at"]), + updated_at=_as_utc(row["updated_at"]), + ) + + +def _schedule_values(item: ScheduleRecord) -> dict[str, Any]: + return { + "id": item.id, + "workspace_id": item.workspace_id, + "pipeline_version_id": item.pipeline_version_id, + "trigger_kind": item.trigger_kind, + "trigger_spec": item.trigger_spec, + "timezone": item.timezone, + "enabled": item.enabled, + "next_run_at": item.next_run_at, + "created_at": item.created_at, + "updated_at": item.updated_at, + } + + +def _outbox(row: RowMapping) -> OutboxRecord: + return OutboxRecord( + id=row["id"], + aggregate_type=row["aggregate_type"], + aggregate_id=row["aggregate_id"], + event_type=row["event_type"], + payload=dict(row["payload_json"]), + available_at=_as_utc(row["available_at"]), + published_at=_optional_utc(row["published_at"]), + attempts=int(row["attempts"]), + ) + + +def _stage_id(run_id: str, name: str, attempt: int) -> str: + return _stable_id("stage", run_id, name, str(attempt)) + + +def _partition_id(run_id: str, partition_key: str) -> str: + return _stable_id("partition", run_id, partition_key) + + +def _batch_id(run_id: str, partition_key: str, batch_key: str) -> str: + return _stable_id("batch", run_id, partition_key, batch_key) + + +def _stable_id(kind: str, *parts: str) -> str: + digest = hashlib.sha256("\0".join(parts).encode()).hexdigest()[:40] + return f"{kind}-{digest}" + + +def _json(value: Any) -> Any: + if isinstance(value, RetryCategory): + return value.value + if isinstance(value, dict): + return {str(key): _json(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json(item) for item in value] + if isinstance(value, datetime): + return value.isoformat() + return value + + +def _as_utc(value: datetime) -> datetime: + if value.tzinfo is None: + return value.replace(tzinfo=UTC) + return value.astimezone(UTC) + + +def _optional_utc(value: datetime | None) -> datetime | None: + return _as_utc(value) if value is not None else None + + +def _next_fire(schedule: RowMapping, previous: datetime) -> datetime: + if schedule["trigger_kind"] == "interval": + return previous + _parse_interval(schedule["trigger_spec"]) + from apscheduler.triggers.cron import CronTrigger + + trigger = CronTrigger.from_crontab(schedule["trigger_spec"], timezone=schedule["timezone"]) + next_fire = trigger.get_next_fire_time(previous, previous + timedelta(microseconds=1)) + if next_fire is None: + raise MetadataError(f"schedule {schedule['id']} has no next fire time") + return _as_utc(next_fire) + + +def _parse_interval(spec: str) -> timedelta: + units = { + "s": "seconds", + "m": "minutes", + "h": "hours", + "d": "days", + "w": "weeks", + } + if len(spec) < 2 or spec[-1] not in units: + raise MetadataError(f"invalid interval: {spec}") + try: + value = int(spec[:-1]) + except ValueError as exc: + raise MetadataError(f"invalid interval: {spec}") from exc + if value <= 0: + raise MetadataError(f"invalid interval: {spec}") + return timedelta(**{units[spec[-1]]: value}) diff --git a/loafer/adapters/metadata_schema.py b/loafer/adapters/metadata_schema.py new file mode 100644 index 0000000..947dfbe --- /dev/null +++ b/loafer/adapters/metadata_schema.py @@ -0,0 +1,321 @@ +"""Versioned SQLAlchemy Core migrations for Loafer metadata.""" + +from __future__ import annotations + +from collections.abc import Callable + +from sqlalchemy import ( + JSON, + BigInteger, + Boolean, + CheckConstraint, + Column, + DateTime, + ForeignKey, + Index, + Integer, + MetaData, + String, + Table, + Text, + UniqueConstraint, + delete, + func, + insert, + select, +) +from sqlalchemy.engine import Connection, Engine + +from loafer.core.run_state import BatchState, RetryCategory, RunState, StageState +from loafer.exceptions import MetadataError + +LATEST_SCHEMA_VERSION = 2 +metadata = MetaData() + + +def _values(enum: type[object]) -> str: + return ", ".join(f"'{item.value}'" for item in enum) # type: ignore[attr-defined] + + +schema_migrations = Table( + "loafer_schema_migrations", + metadata, + Column("version", Integer, primary_key=True), +) + +pipeline_versions = Table( + "loafer_pipeline_versions", + metadata, + Column("id", String(64), primary_key=True), + Column("workspace_id", String(64), nullable=False), + Column("pipeline_key", String(255), nullable=False), + Column("config_digest", String(64), nullable=False), + Column("config_json", JSON, nullable=False), + Column("created_at", DateTime(timezone=True), nullable=False), + UniqueConstraint( + "workspace_id", + "pipeline_key", + "config_digest", + name="uq_loafer_pipeline_version_digest", + ), +) + +runs = Table( + "loafer_runs", + metadata, + Column("id", String(64), primary_key=True), + Column("workspace_id", String(64), nullable=False), + Column( + "pipeline_version_id", + String(64), + ForeignKey("loafer_pipeline_versions.id", ondelete="RESTRICT"), + nullable=False, + ), + Column("command_key", String(255), nullable=False), + Column("state", String(32), nullable=False), + Column("attempt", Integer, nullable=False, default=0), + Column("retry_category", String(32)), + Column("retry_at", DateTime(timezone=True)), + Column("cancel_requested", Boolean, nullable=False, default=False), + Column("next_event_sequence", BigInteger, nullable=False, default=1), + Column("lease_owner", String(255)), + Column("fencing_token", BigInteger, nullable=False, default=0), + Column("lease_expires_at", DateTime(timezone=True)), + Column("heartbeat_at", DateTime(timezone=True)), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("started_at", DateTime(timezone=True)), + Column("finished_at", DateTime(timezone=True)), + Column("parent_run_id", String(64), ForeignKey("loafer_runs.id", ondelete="SET NULL")), + Column("error_json", JSON), + UniqueConstraint("workspace_id", "command_key", name="uq_loafer_run_command"), + CheckConstraint(f"state IN ({_values(RunState)})", name="ck_loafer_run_state"), + CheckConstraint( + f"retry_category IS NULL OR retry_category IN ({_values(RetryCategory)})", + name="ck_loafer_retry_category", + ), + CheckConstraint("attempt >= 0", name="ck_loafer_run_attempt"), + CheckConstraint("fencing_token >= 0", name="ck_loafer_run_fence"), +) + +stages = Table( + "loafer_stages", + metadata, + Column("id", String(96), primary_key=True), + Column("run_id", String(64), ForeignKey("loafer_runs.id", ondelete="CASCADE"), nullable=False), + Column("name", String(128), nullable=False), + Column("state", String(32), nullable=False), + Column("attempt", Integer, nullable=False, default=0), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("started_at", DateTime(timezone=True)), + Column("finished_at", DateTime(timezone=True)), + UniqueConstraint("run_id", "name", "attempt", name="uq_loafer_stage_attempt"), + CheckConstraint(f"state IN ({_values(StageState)})", name="ck_loafer_stage_state"), +) + +partitions = Table( + "loafer_partitions", + metadata, + Column("id", String(96), primary_key=True), + Column("run_id", String(64), ForeignKey("loafer_runs.id", ondelete="CASCADE"), nullable=False), + Column( + "stage_id", String(96), ForeignKey("loafer_stages.id", ondelete="CASCADE"), nullable=False + ), + Column("partition_key", String(255), nullable=False), + Column("created_at", DateTime(timezone=True), nullable=False), + UniqueConstraint("run_id", "partition_key", name="uq_loafer_partition_key"), +) + +artifacts = Table( + "loafer_artifacts", + metadata, + Column("id", String(64), primary_key=True), + Column("run_id", String(64), ForeignKey("loafer_runs.id", ondelete="CASCADE")), + Column("kind", String(64), nullable=False), + Column("uri", Text, nullable=False), + Column("checksum", String(64), nullable=False), + Column("size_bytes", BigInteger, nullable=False), + Column("metadata_json", JSON, nullable=False), + Column("created_at", DateTime(timezone=True), nullable=False), + UniqueConstraint("uri", name="uq_loafer_artifact_uri"), + CheckConstraint("size_bytes >= 0", name="ck_loafer_artifact_size"), +) + +batches = Table( + "loafer_batches", + metadata, + Column("id", String(128), primary_key=True), + Column("run_id", String(64), ForeignKey("loafer_runs.id", ondelete="CASCADE"), nullable=False), + Column( + "partition_id", + String(96), + ForeignKey("loafer_partitions.id", ondelete="CASCADE"), + nullable=False, + ), + Column("batch_key", String(128), nullable=False), + Column("state", String(32), nullable=False), + Column("attempt", Integer, nullable=False), + Column("envelope_json", JSON, nullable=False), + Column("artifact_id", String(64), ForeignKey("loafer_artifacts.id", ondelete="RESTRICT")), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("committed_at", DateTime(timezone=True)), + UniqueConstraint("run_id", "partition_id", "batch_key", name="uq_loafer_batch_key"), + CheckConstraint(f"state IN ({_values(BatchState)})", name="ck_loafer_batch_state"), + CheckConstraint("attempt >= 0", name="ck_loafer_batch_attempt"), +) + +checkpoints = Table( + "loafer_checkpoints", + metadata, + Column("id", String(64), primary_key=True), + Column("run_id", String(64), ForeignKey("loafer_runs.id", ondelete="CASCADE"), nullable=False), + Column( + "partition_id", + String(96), + ForeignKey("loafer_partitions.id", ondelete="CASCADE"), + nullable=False, + ), + Column("batch_id", String(128), nullable=False), + Column("source_position_json", JSON, nullable=False), + Column("fencing_token", BigInteger, nullable=False), + Column("committed_at", DateTime(timezone=True), nullable=False), + UniqueConstraint("run_id", "partition_id", "batch_id", name="uq_loafer_checkpoint_batch"), +) + +run_events = Table( + "loafer_run_events", + metadata, + Column("id", String(64), primary_key=True), + Column("run_id", String(64), ForeignKey("loafer_runs.id", ondelete="CASCADE"), nullable=False), + Column("sequence", BigInteger, nullable=False), + Column("event_type", String(128), nullable=False), + Column("stage_id", String(96), ForeignKey("loafer_stages.id", ondelete="SET NULL")), + Column("batch_id", String(128)), + Column("payload_json", JSON, nullable=False), + Column("occurred_at", DateTime(timezone=True), nullable=False), + UniqueConstraint("run_id", "sequence", name="uq_loafer_event_sequence"), + CheckConstraint("sequence > 0", name="ck_loafer_event_sequence"), +) + +schedules = Table( + "loafer_schedules", + metadata, + Column("id", String(64), primary_key=True), + Column("workspace_id", String(64), nullable=False), + Column( + "pipeline_version_id", + String(64), + ForeignKey("loafer_pipeline_versions.id", ondelete="RESTRICT"), + nullable=False, + ), + Column("trigger_kind", String(16), nullable=False), + Column("trigger_spec", String(255), nullable=False), + Column("timezone", String(64), nullable=False), + Column("enabled", Boolean, nullable=False), + Column("next_run_at", DateTime(timezone=True), nullable=False), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("updated_at", DateTime(timezone=True), nullable=False), + CheckConstraint("trigger_kind IN ('cron', 'interval')", name="ck_loafer_schedule_trigger"), +) + +outbox = Table( + "loafer_outbox", + metadata, + Column("id", String(64), primary_key=True), + Column("aggregate_type", String(64), nullable=False), + Column("aggregate_id", String(64), nullable=False), + Column("event_type", String(128), nullable=False), + Column("payload_json", JSON, nullable=False), + Column("available_at", DateTime(timezone=True), nullable=False), + Column("published_at", DateTime(timezone=True)), + Column("attempts", Integer, nullable=False, default=0), + CheckConstraint("attempts >= 0", name="ck_loafer_outbox_attempts"), +) + +_RUNNABLE_INDEX = Index( + "ix_loafer_runs_runnable", + runs.c.state, + runs.c.retry_at, + runs.c.created_at, +) +_CHECKPOINT_INDEX = Index( + "ix_loafer_checkpoints_latest", + checkpoints.c.run_id, + checkpoints.c.partition_id, + checkpoints.c.committed_at, +) +_OUTBOX_INDEX = Index( + "ix_loafer_outbox_pending", + outbox.c.published_at, + outbox.c.available_at, +) + +_V1_TABLES = ( + pipeline_versions, + runs, + stages, + partitions, + artifacts, + batches, + checkpoints, + run_events, + schedules, +) + + +def migrate(engine: Engine, target_version: int | None = None) -> int: + """Apply or revert checked-in migrations to the requested version.""" + target = LATEST_SCHEMA_VERSION if target_version is None else target_version + if target < 0 or target > LATEST_SCHEMA_VERSION: + raise MetadataError(f"unsupported metadata schema version: {target}") + + with engine.begin() as connection: + schema_migrations.create(connection, checkfirst=True) + current = ( + connection.execute(select(func.max(schema_migrations.c.version))).scalar_one() or 0 + ) + while current < target: + current += 1 + _UP[current](connection) + connection.execute(insert(schema_migrations).values(version=current)) + while current > target: + _DOWN[current](connection) + connection.execute( + delete(schema_migrations).where(schema_migrations.c.version == current) + ) + current -= 1 + return current + + +def current_version(engine: Engine) -> int: + with engine.connect() as connection: + if not engine.dialect.has_table(connection, schema_migrations.name): + return 0 + return connection.execute(select(func.max(schema_migrations.c.version))).scalar_one() or 0 + + +def _up_v1(connection: Connection) -> None: + for table in _V1_TABLES: + table.create(connection, checkfirst=True) + _RUNNABLE_INDEX.create(connection, checkfirst=True) + _CHECKPOINT_INDEX.create(connection, checkfirst=True) + + +def _down_v1(connection: Connection) -> None: + _CHECKPOINT_INDEX.drop(connection, checkfirst=True) + _RUNNABLE_INDEX.drop(connection, checkfirst=True) + for table in reversed(_V1_TABLES): + table.drop(connection, checkfirst=True) + + +def _up_v2(connection: Connection) -> None: + outbox.create(connection, checkfirst=True) + _OUTBOX_INDEX.create(connection, checkfirst=True) + + +def _down_v2(connection: Connection) -> None: + _OUTBOX_INDEX.drop(connection, checkfirst=True) + outbox.drop(connection, checkfirst=True) + + +_UP: dict[int, Callable[[Connection], None]] = {1: _up_v1, 2: _up_v2} +_DOWN: dict[int, Callable[[Connection], None]] = {1: _down_v1, 2: _down_v2} diff --git a/loafer/adapters/object_storage.py b/loafer/adapters/object_storage.py new file mode 100644 index 0000000..d58adc4 --- /dev/null +++ b/loafer/adapters/object_storage.py @@ -0,0 +1,142 @@ +"""Local and in-memory object-storage adapters.""" + +from __future__ import annotations + +import hashlib +import os +import tempfile +import uuid +from collections.abc import Iterable +from pathlib import Path, PurePosixPath +from urllib.parse import unquote, urlparse + +from loafer.exceptions import MetadataError +from loafer.metadata import StoredArtifact, utc_now + + +class FilesystemObjectStorage: + """Atomic local object storage for the embedded single-node profile.""" + + def __init__(self, root: str | Path) -> None: + self._root = Path(root).resolve() + self._root.mkdir(parents=True, exist_ok=True) + + def put( + self, + key: str, + content: bytes | Iterable[bytes], + *, + kind: str, + run_id: str | None = None, + metadata: dict[str, object] | None = None, + ) -> StoredArtifact: + destination = self._path_for_key(key) + destination.parent.mkdir(parents=True, exist_ok=True) + digest = hashlib.sha256() + size = 0 + handle = tempfile.NamedTemporaryFile( + mode="wb", + prefix=f".{destination.name}.", + suffix=".tmp", + dir=destination.parent, + delete=False, + ) + temporary = Path(handle.name) + try: + chunks = (content,) if isinstance(content, bytes) else content + for chunk in chunks: + if not isinstance(chunk, bytes): + raise TypeError("object storage content chunks must be bytes") + handle.write(chunk) + digest.update(chunk) + size += len(chunk) + handle.flush() + os.fsync(handle.fileno()) + handle.close() + os.replace(temporary, destination) + except Exception: + handle.close() + temporary.unlink(missing_ok=True) + raise + checksum = digest.hexdigest() + return StoredArtifact( + id=hashlib.sha256(f"{destination.as_uri()}\0{checksum}".encode()).hexdigest()[:32], + run_id=run_id, + kind=kind, + uri=destination.as_uri(), + checksum=checksum, + size_bytes=size, + metadata=dict(metadata or {}), + created_at=utc_now(), + ) + + def read(self, uri: str) -> bytes: + return self._path_for_uri(uri).read_bytes() + + def delete(self, uri: str) -> None: + self._path_for_uri(uri).unlink(missing_ok=True) + + def exists(self, uri: str) -> bool: + return self._path_for_uri(uri).is_file() + + def _path_for_key(self, key: str) -> Path: + logical = PurePosixPath(key) + if logical.is_absolute() or ".." in logical.parts or not logical.parts: + raise MetadataError(f"unsafe object key: {key}") + path = (self._root / Path(*logical.parts)).resolve() + if not path.is_relative_to(self._root): + raise MetadataError(f"unsafe object key: {key}") + return path + + def _path_for_uri(self, uri: str) -> Path: + parsed = urlparse(uri) + if parsed.scheme != "file" or parsed.netloc not in {"", "localhost"}: + raise MetadataError(f"unsupported local object URI: {uri}") + path = Path(unquote(parsed.path)).resolve() + if not path.is_relative_to(self._root): + raise MetadataError("object URI escapes configured storage root") + return path + + +class MemoryObjectStorage: + """Deterministic object-storage adapter for interface-level tests.""" + + def __init__(self) -> None: + self._objects: dict[str, bytes] = {} + + def put( + self, + key: str, + content: bytes | Iterable[bytes], + *, + kind: str, + run_id: str | None = None, + metadata: dict[str, object] | None = None, + ) -> StoredArtifact: + chunks = (content,) if isinstance(content, bytes) else content + payload = b"".join(chunks) + checksum = hashlib.sha256(payload).hexdigest() + uri = f"memory://{key}" + self._objects[uri] = payload + return StoredArtifact( + id=uuid.uuid5(uuid.NAMESPACE_URL, f"{uri}:{checksum}").hex, + run_id=run_id, + kind=kind, + uri=uri, + checksum=checksum, + size_bytes=len(payload), + metadata=dict(metadata or {}), + created_at=utc_now(), + ) + + def read(self, uri: str) -> bytes: + try: + return self._objects[uri] + except KeyError as exc: + raise FileNotFoundError(uri) from exc + + def delete(self, uri: str) -> None: + self._objects.pop(uri, None) + + def exists(self, uri: str) -> bool: + return uri in self._objects diff --git a/loafer/adapters/runtime.py b/loafer/adapters/runtime.py index 22c59aa..fd5da8f 100644 --- a/loafer/adapters/runtime.py +++ b/loafer/adapters/runtime.py @@ -2,9 +2,18 @@ from __future__ import annotations +import hashlib +import json import os +from datetime import date, datetime +from decimal import Decimal +from typing import Any +from uuid import UUID from loafer.contracts import Checkpoint, RunEvent +from loafer.metadata import BatchCommit, RunLease +from loafer.ports.metadata import MetadataStore +from loafer.ports.object_storage import ObjectStoragePort class NeverCancelled: @@ -51,3 +60,116 @@ def approve_transform(self, generated_code: str) -> bool: except (EOFError, KeyboardInterrupt): return False return answer in {"y", "yes"} + + +class MetadataCancellation: + """Read cooperative cancellation from the authoritative run record.""" + + def __init__(self, metadata: MetadataStore) -> None: + self._metadata = metadata + + def is_cancelled(self, run_id: str) -> bool: + return self._metadata.cancellation_requested(run_id) + + +class DurableBatchRecovery: + """Stage transformed batches as immutable objects before checkpointing.""" + + def __init__( + self, + metadata: MetadataStore, + objects: ObjectStoragePort, + lease: RunLease, + ) -> None: + self._metadata = metadata + self._objects = objects + self._lease = lease + + def restore(self, run_id: str, partition_id: str) -> list[BatchCommit]: + return self._metadata.list_batch_commits(run_id, partition_id) + + def read_rows(self, commit: BatchCommit) -> list[dict[str, Any]]: + payload = self._objects.read(commit.artifact.uri) + checksum = hashlib.sha256(payload).hexdigest() + if checksum != commit.artifact.checksum: + raise ValueError(f"recovery artifact checksum mismatch for {commit.envelope.batch_id}") + rows = [] + for line in payload.splitlines(): + value = json.loads(line, object_hook=_row_object_hook) + if not isinstance(value, dict): + raise ValueError("recovery artifact contains a non-object row") + rows.append(value) + return rows + + def commit(self, envelope: Any, rows: list[dict[str, Any]]) -> Checkpoint: + payload = b"".join( + json.dumps( + row, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + default=_row_json_default, + ).encode("utf-8") + + b"\n" + for row in rows + ) + artifact = self._objects.put( + ( + f"runs/{envelope.run_id}/partitions/{envelope.partition_id}/" + f"batches/{envelope.batch_id}.jsonl" + ), + payload, + kind="temporary_output", + run_id=envelope.run_id, + metadata={ + "batch_id": envelope.batch_id, + "rows": len(rows), + "format": "loafer-jsonl-v1", + }, + ) + checkpoint = Checkpoint( + checkpoint_id=hashlib.sha256( + (f"{envelope.run_id}\0{envelope.partition_id}\0{envelope.batch_id}").encode() + ).hexdigest()[:32], + run_id=envelope.run_id, + partition_id=envelope.partition_id, + batch_id=envelope.batch_id, + source_position=envelope.source_position_end, + committed_at=artifact.created_at, + ) + return self._metadata.commit_batch( + self._lease, + envelope, + checkpoint, + artifact, + ).checkpoint + + +def _row_json_default(value: Any) -> dict[str, str]: + if isinstance(value, datetime): + return {"__loafer_type__": "datetime", "value": value.isoformat()} + if isinstance(value, date): + return {"__loafer_type__": "date", "value": value.isoformat()} + if isinstance(value, Decimal): + return {"__loafer_type__": "decimal", "value": str(value)} + if isinstance(value, UUID): + return {"__loafer_type__": "uuid", "value": str(value)} + if isinstance(value, bytes): + return {"__loafer_type__": "bytes", "value": value.hex()} + raise TypeError(f"unsupported recovery value: {type(value).__name__}") + + +def _row_object_hook(value: dict[str, Any]) -> Any: + marker = value.get("__loafer_type__") + raw = value.get("value") + if marker == "datetime": + return datetime.fromisoformat(raw) + if marker == "date": + return date.fromisoformat(raw) + if marker == "decimal": + return Decimal(raw) + if marker == "uuid": + return UUID(raw) + if marker == "bytes": + return bytes.fromhex(raw) + return value diff --git a/loafer/application/__init__.py b/loafer/application/__init__.py index 7cceeec..b5d2e8e 100644 --- a/loafer/application/__init__.py +++ b/loafer/application/__init__.py @@ -1,5 +1,7 @@ """Versioned application boundary for Loafer clients.""" +from typing import Any + from loafer.application.local import get_local_application from loafer.application.service import LocalApplicationService, RunPipeline from loafer.contracts import ( @@ -17,6 +19,14 @@ ) from loafer.ports.runtime import CancellationPort, CheckpointPort, SecretResolver + +def enqueue_pipeline(*args: Any, **kwargs: Any) -> Any: + """Lazily create a durable run without coupling package import to worker setup.""" + from loafer.application.durable import enqueue_pipeline as enqueue + + return enqueue(*args, **kwargs) + + __all__ = [ "BatchEnvelope", "CancellationPort", @@ -34,5 +44,6 @@ "SecretResolver", "StageStatus", "ValidationResult", + "enqueue_pipeline", "get_local_application", ] diff --git a/loafer/application/durable.py b/loafer/application/durable.py new file mode 100644 index 0000000..5d40362 --- /dev/null +++ b/loafer/application/durable.py @@ -0,0 +1,123 @@ +"""Composition helpers for the durable local and PostgreSQL profiles.""" + +from __future__ import annotations + +import hashlib +import json +import os +import uuid +from pathlib import Path + +from loafer.adapters.metadata import SqlMetadataStore +from loafer.adapters.object_storage import FilesystemObjectStorage +from loafer.config import load_config +from loafer.metadata import PipelineVersion, RunRecord +from loafer.worker import DurableWorker + +_LOAFER_DIR = Path.home() / ".loafer" +_METADATA_PATH = _LOAFER_DIR / "metadata.db" +_OBJECTS_PATH = _LOAFER_DIR / "objects" + + +def default_metadata_url() -> str: + """Return PostgreSQL when configured, otherwise restricted local SQLite.""" + configured = os.environ.get("LOAFER_METADATA_URL") + if configured: + return configured + _LOAFER_DIR.mkdir(parents=True, exist_ok=True) + return f"sqlite:///{_METADATA_PATH}" + + +def get_metadata_store(url: str | None = None) -> SqlMetadataStore: + """Build and migrate the configured authoritative metadata adapter.""" + store = SqlMetadataStore(url or default_metadata_url()) + store.migrate() + return store + + +def get_object_storage(root: str | Path | None = None) -> FilesystemObjectStorage: + """Build the embedded filesystem object-store adapter.""" + configured = root or os.environ.get("LOAFER_OBJECTS_PATH") or _OBJECTS_PATH + return FilesystemObjectStorage(configured) + + +def enqueue_pipeline( + config_path: str | Path, + *, + command_key: str, + workspace_id: str = "local", + run_id: str | None = None, + metadata_url: str | None = None, +) -> RunRecord: + """Persist an immutable config version and idempotent run command.""" + version = register_pipeline_config( + config_path, + workspace_id=workspace_id, + metadata_url=metadata_url, + ) + return enqueue_registered_version( + version.id, + command_key=command_key, + workspace_id=workspace_id, + run_id=run_id, + metadata_url=metadata_url, + ) + + +def register_pipeline_config( + config_path: str | Path, + *, + workspace_id: str = "local", + metadata_url: str | None = None, +) -> PipelineVersion: + """Snapshot a resolved pipeline config as an immutable version.""" + resolved = Path(config_path).resolve() + config = load_config(resolved) + document = config.model_dump(mode="json") + rendered = json.dumps(document, sort_keys=True, separators=(",", ":")) + digest = hashlib.sha256(rendered.encode()).hexdigest() + store = get_metadata_store(metadata_url) + try: + return store.register_pipeline_version( + workspace_id=workspace_id, + pipeline_key=config.name or resolved.stem, + config_digest=digest, + config={"document": document, "source_path": str(resolved)}, + ) + finally: + store.close() + + +def enqueue_registered_version( + pipeline_version_id: str, + *, + command_key: str, + workspace_id: str = "local", + run_id: str | None = None, + metadata_url: str | None = None, +) -> RunRecord: + """Create an idempotent run for an already immutable version.""" + store = get_metadata_store(metadata_url) + try: + return store.create_run( + workspace_id=workspace_id, + pipeline_version_id=pipeline_version_id, + command_key=command_key, + run_id=run_id or uuid.uuid4().hex[:12], + ) + finally: + store.close() + + +def get_durable_worker( + *, + worker_id: str, + metadata_url: str | None = None, + object_root: str | Path | None = None, +) -> DurableWorker: + """Compose a worker process; callers own its long-running lifecycle.""" + return DurableWorker( + get_metadata_store(metadata_url), + get_object_storage(object_root), + worker_id=worker_id, + ) diff --git a/loafer/application/service.py b/loafer/application/service.py index 8a22acf..70175d4 100644 --- a/loafer/application/service.py +++ b/loafer/application/service.py @@ -27,6 +27,7 @@ from loafer.engine import ProviderFactory, stream_pipeline from loafer.exceptions import PipelineError from loafer.graph.state import PipelineState +from loafer.ports.metadata import BatchRecoveryPort from loafer.ports.runtime import ( CancellationPort, CheckpointPort, @@ -208,6 +209,7 @@ def __init__( events: EventPublisher, reviewer: ReviewPort, provider_factory: ProviderFactory | None = None, + recovery: BatchRecoveryPort | None = None, ) -> None: self._cancellation = cancellation self._checkpoints = checkpoints @@ -215,6 +217,7 @@ def __init__( self._events = events self._reviewer = reviewer self._provider_factory = provider_factory + self._recovery = recovery def create_plan(self, request: RunRequest) -> ExecutionPlan: """Validate a config and return its credential-free execution plan.""" @@ -296,6 +299,7 @@ def _stream_prepared( provider_factory=self._provider_factory, cancellation=self._cancellation, checkpoints=self._checkpoints, + recovery=self._recovery, ) sequence = 0 diff --git a/loafer/cli.py b/loafer/cli.py index 308c1c7..5aeb655 100644 --- a/loafer/cli.py +++ b/loafer/cli.py @@ -9,8 +9,11 @@ from __future__ import annotations import logging +import os import signal +import socket import time +import uuid from pathlib import Path from typing import Any @@ -91,6 +94,51 @@ def _main( """AI-assisted ETL and ELT pipelines from the command line.""" +@app.command("enqueue") +def enqueue_command( + config: Path = _config_arg, + command_key: str | None = typer.Option( + None, + "--command-key", + help="Idempotency key; repeated use returns the same durable run.", + ), +) -> None: + """Create a durable run for execution by a separate worker process.""" + from loafer.application.durable import enqueue_pipeline + + key = command_key or f"manual:{uuid.uuid4().hex}" + try: + run = enqueue_pipeline(config, command_key=key) + except Exception as exc: + err_console.print(f"[red]Could not enqueue pipeline:[/red] {exc}") + raise typer.Exit(1) from exc + console.print(f"[green]✓ Enqueued run[/green] {run.id}") + + +@app.command("worker") +def worker_command( + once: bool = typer.Option(False, "--once", help="Process at most one runnable job."), + worker_id: str | None = typer.Option(None, "--id", help="Stable worker identity."), +) -> None: + """Run the durable data-plane worker separately from the scheduler.""" + from loafer.application.durable import get_durable_worker + + identity = worker_id or f"{socket.gethostname()}-{os.getpid()}" + worker = get_durable_worker(worker_id=identity) + try: + if once: + run_id = worker.run_once() + if run_id is None: + console.print("[dim]No runnable jobs.[/dim]") + else: + console.print(f"[green]✓ Processed run[/green] {run_id}") + return + console.print(f"[green]Worker started[/green] {identity}") + worker.run_forever() + finally: + worker.close() + + # Animated stage loaders _STAGE_SPINNERS = { @@ -807,7 +855,10 @@ def schedule( console.print(f" Name: {pipeline_name}") console.print(f" Config: {config_file}") console.print(f" Trigger: {trigger_desc}") - console.print("\nRun [bold]loafer start[/bold] to begin executing scheduled jobs") + console.print( + "\nRun [bold]loafer start[/bold] to enqueue due jobs and " + "[bold]loafer worker[/bold] in a separate process to execute them" + ) @app.command() @@ -868,7 +919,7 @@ def list_schedules() -> None: def start( detached: bool = typer.Option(False, "--detached", "-d", help="Run in background"), ) -> None: - """Start the scheduler and run scheduled jobs.""" + """Start the scheduler that enqueues due run commands.""" from loafer.scheduler import PipelineScheduler if detached: @@ -883,6 +934,7 @@ def start( console.print("[green]✓ Scheduler started[/green]") console.print(f" Log: {log_path}") + console.print(" Execution: start `loafer worker` in a separate process") console.print("Press Ctrl+C to stop\n") jobs = scheduler.list_schedules() diff --git a/loafer/core/run_state.py b/loafer/core/run_state.py new file mode 100644 index 0000000..7ec0f9f --- /dev/null +++ b/loafer/core/run_state.py @@ -0,0 +1,110 @@ +"""Pure durable-execution state machines and retry policy.""" + +from __future__ import annotations + +from enum import StrEnum + +from loafer.exceptions import InvalidStateTransitionError + + +class RunState(StrEnum): + QUEUED = "queued" + CLAIMED = "claimed" + RUNNING = "running" + CANCELLING = "cancelling" + RETRY_WAIT = "retry_wait" + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + + +class StageState(StrEnum): + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + SKIPPED = "skipped" + + +class BatchState(StrEnum): + PENDING = "pending" + RUNNING = "running" + COMMITTED = "committed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class RetryCategory(StrEnum): + INFRASTRUCTURE = "infrastructure" + FAILED_BATCH = "failed_batch" + MANUAL_RERUN = "manual_rerun" + BACKFILL = "backfill" + + +_RUN_TRANSITIONS: dict[RunState, frozenset[RunState]] = { + RunState.QUEUED: frozenset({RunState.CLAIMED, RunState.CANCELLED}), + RunState.CLAIMED: frozenset( + {RunState.RUNNING, RunState.CANCELLING, RunState.RETRY_WAIT, RunState.FAILED} + ), + RunState.RUNNING: frozenset( + { + RunState.CANCELLING, + RunState.RETRY_WAIT, + RunState.SUCCEEDED, + RunState.FAILED, + RunState.CANCELLED, + } + ), + RunState.CANCELLING: frozenset({RunState.CANCELLED, RunState.RETRY_WAIT, RunState.FAILED}), + RunState.RETRY_WAIT: frozenset({RunState.QUEUED, RunState.CANCELLED}), + RunState.SUCCEEDED: frozenset(), + RunState.FAILED: frozenset(), + RunState.CANCELLED: frozenset(), +} + +_STAGE_TRANSITIONS: dict[StageState, frozenset[StageState]] = { + StageState.PENDING: frozenset({StageState.RUNNING, StageState.CANCELLED, StageState.SKIPPED}), + StageState.RUNNING: frozenset({StageState.SUCCEEDED, StageState.FAILED, StageState.CANCELLED}), + StageState.SUCCEEDED: frozenset(), + StageState.FAILED: frozenset(), + StageState.CANCELLED: frozenset(), + StageState.SKIPPED: frozenset(), +} + +_BATCH_TRANSITIONS: dict[BatchState, frozenset[BatchState]] = { + BatchState.PENDING: frozenset({BatchState.RUNNING, BatchState.CANCELLED}), + BatchState.RUNNING: frozenset({BatchState.COMMITTED, BatchState.FAILED, BatchState.CANCELLED}), + BatchState.COMMITTED: frozenset(), + BatchState.FAILED: frozenset(), + BatchState.CANCELLED: frozenset(), +} + + +def require_run_transition(current: RunState, target: RunState) -> None: + """Reject a non-idempotent transition not allowed by the run machine.""" + _require_transition("run", current, target, _RUN_TRANSITIONS) + + +def require_stage_transition(current: StageState, target: StageState) -> None: + """Reject a non-idempotent transition not allowed by the stage machine.""" + _require_transition("stage", current, target, _STAGE_TRANSITIONS) + + +def require_batch_transition(current: BatchState, target: BatchState) -> None: + """Reject a non-idempotent transition not allowed by the batch machine.""" + _require_transition("batch", current, target, _BATCH_TRANSITIONS) + + +def _require_transition( + machine: str, + current: StrEnum, + target: StrEnum, + transitions: dict[StrEnum, frozenset[StrEnum]], +) -> None: + if current == target: + return + if target not in transitions[current]: + raise InvalidStateTransitionError( + f"invalid {machine} state transition: {current.value} -> {target.value}" + ) diff --git a/loafer/data_plane.py b/loafer/data_plane.py index 4385885..fceff8f 100644 --- a/loafer/data_plane.py +++ b/loafer/data_plane.py @@ -33,6 +33,7 @@ if TYPE_CHECKING: from loafer.graph.state import PipelineState from loafer.ports.connector import SourceConnector, TargetConnector + from loafer.ports.metadata import BatchRecoveryPort from loafer.ports.runtime import CancellationPort, CheckpointPort _INCREMENTAL_PUSHDOWN_SOURCES = {"postgres", "mysql", "sqlite", "rest_api"} @@ -51,6 +52,7 @@ def stream_bounded_pipeline( dry_run: bool, cancellation: CancellationPort | None, checkpoints: CheckpointPort | None, + recovery: BatchRecoveryPort | None = None, ) -> Iterator[tuple[str, str, PipelineState]]: """Flow source batches through validation, transform, and staged publication.""" source: SourceConnector | None = None @@ -69,6 +71,9 @@ def stream_bounded_pipeline( quality_columns: dict[str, dict[str, int]] = {} quality_warnings: set[str] = set() destructive_seen: set[tuple[str, str]] = set() + recovered_rows_in = 0 + recovered_bytes_in = 0 + resume_offset = 0 state["is_streaming"] = True state["raw_data"] = [] @@ -98,6 +103,40 @@ def stream_bounded_pipeline( batch_index = 0 source_offset = 0 + if recovery is not None and not dry_run: + commits = recovery.restore(state["run_id"], _PARTITION_ID) + for commit in commits: + recovered_rows = recovery.read_rows(commit) + if target is not None: + _write_target(target, recovered_rows) + output_digest.update(recovered_rows) + if recovered_rows: + output_schema.apply(recovered_rows, config.execution.schema_drift) + recovered_rows_in += commit.envelope.rows_in + recovered_bytes_in += commit.envelope.bytes_in + resume_offset = max( + resume_offset, + _position_offset(commit.checkpoint.source_position) + 1, + ) + batch_index += 1 + last_envelope = commit.envelope + if commits: + source_offset = resume_offset + state["rows_extracted"] = recovered_rows_in + state["rows_transformed"] = output_digest.rows + state["rows_loaded"] = output_digest.rows + state["batches_completed"] = batch_index + state["bytes_in"] = recovered_bytes_in + state["bytes_out"] = output_digest.bytes + state["output_checksum"] = output_digest.checksum + state["last_batch_envelope"] = last_envelope + state.setdefault("warnings", []).append( + f"recovered {len(commits)} committed batch(es) from durable artifacts" + ) + yield ("recovery", "done", state) + + rows_to_skip = resume_offset + while True: _raise_if_cancelled(cancellation, state["run_id"]) read_started = time.monotonic() @@ -109,6 +148,10 @@ def stream_bounded_pipeline( stage_ms["extract"] += (time.monotonic() - read_started) * 1000 raw_rows = _filter_incremental_rows(config, state, raw_rows) + if rows_to_skip: + skipped = min(rows_to_skip, len(raw_rows)) + raw_rows = raw_rows[skipped:] + rows_to_skip -= skipped if not raw_rows: continue @@ -190,22 +233,8 @@ def stream_bounded_pipeline( stage_ms["transform"] += (time.monotonic() - transform_started) * 1000 _raise_if_cancelled(cancellation, state["run_id"]) - active_stage = "load" - load_started = time.monotonic() - written = len(transformed) if dry_run else _write_target(target, transformed) - stage_ms["load"] += (time.monotonic() - load_started) * 1000 - batch_output_digest = RollingRowsDigest() batch_output_digest.update(transformed) - output_digest.update(transformed) - state["rows_transformed"] = output_digest.rows - state["rows_loaded"] = state.get("rows_loaded", 0) + written - state["batches_completed"] = batch_index - state["bytes_in"] = input_digest.bytes - state["bytes_out"] = output_digest.bytes - state["input_checksum"] = input_digest.checksum - state["output_checksum"] = output_digest.checksum - state["validation_passed"] = True source_start = source_offset source_offset += len(raw_rows) @@ -230,16 +259,34 @@ def stream_bounded_pipeline( output_checksum=batch_output_digest.checksum, duration_ms=(time.monotonic() - batch_started) * 1000, ) + if recovery is not None and not dry_run: + recovery.commit(last_envelope, transformed) + + active_stage = "load" + load_started = time.monotonic() + written = len(transformed) if dry_run else _write_target(target, transformed) + stage_ms["load"] += (time.monotonic() - load_started) * 1000 + + output_digest.update(transformed) + state["rows_extracted"] = recovered_rows_in + input_digest.rows + state["rows_transformed"] = output_digest.rows + state["rows_loaded"] = state.get("rows_loaded", 0) + written + state["batches_completed"] = batch_index + state["bytes_in"] = recovered_bytes_in + input_digest.bytes + state["bytes_out"] = output_digest.bytes + state["input_checksum"] = input_digest.checksum if not recovered_rows_in else None + state["output_checksum"] = output_digest.checksum + state["validation_passed"] = True state["last_batch_envelope"] = last_envelope state["raw_data"] = [] state["transformed_data"] = [] yield ("batch", "done", state) - if input_digest.rows == 0: + if recovered_rows_in + input_digest.rows == 0: raise ValidationError("Source returned 0 rows — nothing to validate") state["validation_report"] = { - "rows_checked": input_digest.rows, + "rows_checked": recovered_rows_in + input_digest.rows, "rows_rejected": state.get("rows_rejected", 0), "columns": quality_columns, "hard_failures": [], @@ -262,7 +309,12 @@ def stream_bounded_pipeline( target.finalize() state["target_published"] = True - if checkpoints is not None and last_envelope is not None and not dry_run: + if ( + recovery is None + and checkpoints is not None + and last_envelope is not None + and not dry_run + ): checkpoints.save( Checkpoint( checkpoint_id=uuid.uuid4().hex, @@ -296,6 +348,14 @@ def stream_bounded_pipeline( source.disconnect() +def _position_offset(position: object) -> int: + if isinstance(position, dict): + value = position.get("offset") + if isinstance(value, int) and value >= 0: + return value + raise PipelineError(f"unsupported durable source position: {position!r}") + + def _source_connector(config: PipelineConfig, state: PipelineState) -> SourceConnector: incremental = config.incremental if incremental is None or config.source.type not in _INCREMENTAL_PUSHDOWN_SOURCES: diff --git a/loafer/engine.py b/loafer/engine.py index 28f241f..ce2f6dc 100644 --- a/loafer/engine.py +++ b/loafer/engine.py @@ -20,6 +20,7 @@ from collections.abc import Iterator from loafer.llm.base import LLMProvider + from loafer.ports.metadata import BatchRecoveryPort from loafer.ports.runtime import ( CancellationPort, CheckpointPort, @@ -201,6 +202,7 @@ def execute_pipeline( provider_factory: ProviderFactory | None = None, cancellation: CancellationPort | None = None, checkpoints: CheckpointPort | None = None, + recovery: BatchRecoveryPort | None = None, ) -> PipelineState: """Execute a validated ETL or ELT pipeline. @@ -245,6 +247,7 @@ def execute_pipeline( dry_run=dry_run, cancellation=cancellation, checkpoints=checkpoints, + recovery=recovery, ): pass except PipelineError: @@ -301,6 +304,7 @@ def stream_pipeline( provider_factory: ProviderFactory | None = None, cancellation: CancellationPort | None = None, checkpoints: CheckpointPort | None = None, + recovery: BatchRecoveryPort | None = None, ) -> Iterator[tuple[str, str, PipelineState]]: """Execute a validated pipeline and yield per-stage runtime updates. @@ -338,6 +342,7 @@ def stream_pipeline( dry_run=dry_run, cancellation=cancellation, checkpoints=checkpoints, + recovery=recovery, ) finally: if not dry_run and state.get("target_published", False): diff --git a/loafer/exceptions.py b/loafer/exceptions.py index b98f744..30dc9d6 100644 --- a/loafer/exceptions.py +++ b/loafer/exceptions.py @@ -58,3 +58,19 @@ class SchedulerError(LoaferError): class PipelineError(LoaferError): """Pipeline orchestration failure.""" + + +class MetadataError(LoaferError): + """Durable metadata operation failed.""" + + +class InvalidStateTransitionError(MetadataError): + """A durable run, stage, or batch transition is impossible.""" + + +class StaleFenceError(MetadataError): + """A worker attempted to write with an expired or superseded lease.""" + + +class IdempotencyConflictError(MetadataError): + """An idempotency key was reused for different immutable input.""" diff --git a/loafer/metadata.py b/loafer/metadata.py new file mode 100644 index 0000000..66ffd34 --- /dev/null +++ b/loafer/metadata.py @@ -0,0 +1,115 @@ +"""Serializable contracts for durable single-node execution metadata.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any + +from loafer.contracts import BatchEnvelope, Checkpoint +from loafer.core.run_state import RetryCategory, RunState + + +def utc_now() -> datetime: + return datetime.now(UTC) + + +@dataclass(frozen=True, slots=True) +class PipelineVersion: + id: str + workspace_id: str + pipeline_key: str + config_digest: str + config: dict[str, Any] + created_at: datetime + + +@dataclass(frozen=True, slots=True) +class RunRecord: + id: str + workspace_id: str + pipeline_version_id: str + command_key: str + state: RunState + attempt: int + retry_category: RetryCategory | None + cancel_requested: bool + fencing_token: int + lease_owner: str | None + lease_expires_at: datetime | None + heartbeat_at: datetime | None + created_at: datetime + started_at: datetime | None + finished_at: datetime | None + parent_run_id: str | None = None + error: dict[str, Any] | None = None + + +@dataclass(frozen=True, slots=True) +class RunLease: + run: RunRecord + worker_id: str + fencing_token: int + expires_at: datetime + + +@dataclass(frozen=True, slots=True) +class StoredEvent: + run_id: str + sequence: int + event_type: str + payload: dict[str, Any] + occurred_at: datetime + + +@dataclass(frozen=True, slots=True) +class StoredArtifact: + id: str + run_id: str | None + kind: str + uri: str + checksum: str + size_bytes: int + metadata: dict[str, Any] + created_at: datetime + + +@dataclass(frozen=True, slots=True) +class BatchCommit: + envelope: BatchEnvelope + checkpoint: Checkpoint + artifact: StoredArtifact + + +@dataclass(frozen=True, slots=True) +class RecoveredBatch: + envelope: BatchEnvelope + checkpoint: Checkpoint + artifact: StoredArtifact + rows: tuple[dict[str, Any], ...] = field(default_factory=tuple) + + +@dataclass(frozen=True, slots=True) +class ScheduleRecord: + id: str + workspace_id: str + pipeline_version_id: str + trigger_kind: str + trigger_spec: str + timezone: str + enabled: bool + next_run_at: datetime + created_at: datetime + updated_at: datetime + + +@dataclass(frozen=True, slots=True) +class OutboxRecord: + id: str + aggregate_type: str + aggregate_id: str + event_type: str + payload: dict[str, Any] + available_at: datetime + published_at: datetime | None + attempts: int diff --git a/loafer/ports/metadata.py b/loafer/ports/metadata.py new file mode 100644 index 0000000..e145e03 --- /dev/null +++ b/loafer/ports/metadata.py @@ -0,0 +1,140 @@ +"""Durable metadata seam for scheduler and worker processes.""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import datetime, timedelta +from typing import Any, Protocol + +from loafer.contracts import BatchEnvelope, Checkpoint +from loafer.core.run_state import RetryCategory, RunState, StageState +from loafer.metadata import ( + BatchCommit, + OutboxRecord, + PipelineVersion, + RunLease, + RunRecord, + ScheduleRecord, + StoredArtifact, + StoredEvent, +) + + +class MetadataStore(Protocol): + """Persist and advance one authoritative durable-execution model.""" + + def migrate(self, target_version: int | None = None) -> int: + """Move the metadata schema to a supported version and return it.""" + + def register_pipeline_version( + self, + *, + workspace_id: str, + pipeline_key: str, + config_digest: str, + config: dict[str, Any], + ) -> PipelineVersion: + """Return the immutable version for this exact config digest.""" + + def get_pipeline_version(self, version_id: str) -> PipelineVersion: + """Resolve the immutable configuration claimed by a worker.""" + + def create_run( + self, + *, + workspace_id: str, + pipeline_version_id: str, + command_key: str, + run_id: str | None = None, + parent_run_id: str | None = None, + retry_category: RetryCategory | None = None, + ) -> RunRecord: + """Create or return the run identified by an idempotent command.""" + + def get_run(self, run_id: str) -> RunRecord: + """Return one durable run.""" + + def claim_run(self, worker_id: str, lease_for: timedelta) -> RunLease | None: + """Claim the next runnable run and issue a new fencing token.""" + + def heartbeat(self, lease: RunLease, lease_for: timedelta) -> RunLease: + """Renew a current lease or reject its stale fencing token.""" + + def transition_run( + self, + lease: RunLease, + target: RunState, + *, + error: dict[str, Any] | None = None, + retry_category: RetryCategory | None = None, + retry_at: datetime | None = None, + ) -> RunRecord: + """Advance the run state under its active fence.""" + + def transition_stage( + self, + lease: RunLease, + stage_name: str, + target: StageState, + ) -> None: + """Create or advance a named stage under its active fence.""" + + def append_event( + self, + lease: RunLease, + event_type: str, + payload: dict[str, Any], + ) -> StoredEvent: + """Append one event and allocate its monotonic per-run sequence.""" + + def commit_batch( + self, + lease: RunLease, + envelope: BatchEnvelope, + checkpoint: Checkpoint, + artifact: StoredArtifact, + ) -> BatchCommit: + """Atomically record a committed batch, checkpoint, event, and outbox row.""" + + def list_batch_commits(self, run_id: str, partition_id: str) -> list[BatchCommit]: + """Return committed batches in source order for recovery.""" + + def latest_checkpoint(self, run_id: str, partition_id: str) -> Checkpoint | None: + """Return the last committed checkpoint for one partition.""" + + def request_cancel(self, run_id: str) -> RunRecord: + """Idempotently request cooperative cancellation.""" + + def cancellation_requested(self, run_id: str) -> bool: + """Return whether the worker should cancel at its next safe boundary.""" + + def upsert_schedule( + self, + schedule: ScheduleRecord, + ) -> ScheduleRecord: + """Idempotently create or replace a durable schedule.""" + + def enqueue_due_schedules(self, now: datetime) -> list[RunRecord]: + """Create idempotent run commands for due schedules and advance them.""" + + def list_events(self, run_id: str, after: int = 0) -> list[StoredEvent]: + """Return the append-only event stream after a sequence.""" + + def pending_outbox(self, limit: int = 100) -> Sequence[OutboxRecord]: + """Return unpublished transport records without exposing job contents.""" + + def mark_outbox_published(self, outbox_id: str, published_at: datetime) -> None: + """Idempotently mark a transport record as published.""" + + +class BatchRecoveryPort(Protocol): + """Make bounded batches durable and replay them after a worker crash.""" + + def restore(self, run_id: str, partition_id: str) -> list[BatchCommit]: + """Return durable batches in source order.""" + + def read_rows(self, commit: BatchCommit) -> list[dict[str, Any]]: + """Read the staged output rows for a committed batch.""" + + def commit(self, envelope: BatchEnvelope, rows: list[dict[str, Any]]) -> Checkpoint: + """Stage output and atomically advance the durable checkpoint.""" diff --git a/loafer/ports/object_storage.py b/loafer/ports/object_storage.py new file mode 100644 index 0000000..b5b9740 --- /dev/null +++ b/loafer/ports/object_storage.py @@ -0,0 +1,32 @@ +"""Object-storage seam for durable artifacts and temporary output.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Protocol + +from loafer.metadata import StoredArtifact + + +class ObjectStoragePort(Protocol): + """Store immutable binary objects behind stable logical keys.""" + + def put( + self, + key: str, + content: bytes | Iterable[bytes], + *, + kind: str, + run_id: str | None = None, + metadata: dict[str, object] | None = None, + ) -> StoredArtifact: + """Atomically store content and return its immutable descriptor.""" + + def read(self, uri: str) -> bytes: + """Read an object by the descriptor URI returned from ``put``.""" + + def delete(self, uri: str) -> None: + """Delete an object when retention policy permits it.""" + + def exists(self, uri: str) -> bool: + """Return whether an object currently exists.""" diff --git a/loafer/scheduler.py b/loafer/scheduler.py index 93a0ac2..7eaddfe 100644 --- a/loafer/scheduler.py +++ b/loafer/scheduler.py @@ -1,7 +1,8 @@ -"""Pipeline scheduler — APScheduler-based cron scheduling. +"""Pipeline scheduler — APScheduler-based durable command creation. -Manages recurring pipeline runs via cron or interval triggers. -Jobs are persisted in a SQLite store so they survive restarts. +Manages recurring pipeline commands via cron or interval triggers. The +scheduler never executes data work; a separate durable worker claims runs. +Triggers are persisted in a SQLite store so they survive restarts. """ from __future__ import annotations @@ -17,7 +18,7 @@ from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore from apscheduler.schedulers.background import BackgroundScheduler -from loafer.application import RunRequest, get_local_application +from loafer.application import enqueue_pipeline from loafer.exceptions import SchedulerError logger = logging.getLogger("loafer.scheduler") @@ -96,16 +97,17 @@ def _read_run_state() -> dict[str, Any]: return {} -def _run_pipeline_job(config_path: str, name: str = "") -> None: - """Execute a pipeline run, called by the scheduler.""" - run_id = uuid.uuid4().hex[:12] +def _run_pipeline_job(config_path: str, name: str = "", schedule_id: str = "") -> None: + """Create a durable run command; worker processes execute it separately.""" + occurrence = datetime.now(UTC).replace(microsecond=0).isoformat() + command_key = f"schedule:{schedule_id or config_path}:{occurrence}" display = f"{name} ({config_path})" if name else config_path - logger.info("Starting scheduled run %s for %s", run_id, display) + logger.info("Enqueuing scheduled run for %s", display) try: - get_local_application().run_pipeline.run(RunRequest(config_path=config_path, run_id=run_id)) - logger.info("Completed scheduled run %s", run_id) + run = enqueue_pipeline(config_path, command_key=command_key) + logger.info("Enqueued scheduled run %s", run.id) except Exception as exc: - logger.error("Scheduled run %s failed: %s", run_id, exc) + logger.error("Could not enqueue scheduled run for %s: %s", display, exc) raise @@ -201,7 +203,7 @@ def add_schedule( self._scheduler.add_job( _run_pipeline_job, trigger=trigger, - args=[config_path, name], + args=[config_path, name, job_id], id=job_id, replace_existing=replace, misfire_grace_time=300, @@ -331,11 +333,11 @@ def _parse_interval(self, spec: str) -> Any: def _on_job_executed(self, event: Any) -> None: """Log job execution events and record last-run state.""" if event.exception: - logger.error("Job %s failed: %s", event.job_id, event.exception) + logger.error("Job %s enqueue failed: %s", event.job_id, event.exception) _record_run(event.job_id, "failed") else: - logger.info("Job %s completed successfully", event.job_id) - _record_run(event.job_id, "success") + logger.info("Job %s enqueued successfully", event.job_id) + _record_run(event.job_id, "enqueued") def export_jobs(self, path: str | Path) -> None: """Export scheduled jobs to a JSON file.""" diff --git a/loafer/worker.py b/loafer/worker.py new file mode 100644 index 0000000..a05c789 --- /dev/null +++ b/loafer/worker.py @@ -0,0 +1,175 @@ +"""Single-node durable worker process.""" + +from __future__ import annotations + +import json +import tempfile +import time +from datetime import timedelta +from pathlib import Path +from typing import Any + +from loafer.adapters.runtime import ( + DurableBatchRecovery, + EnvironmentSecretResolver, + MetadataCancellation, + NullCheckpointStore, + NullEventPublisher, +) +from loafer.application.service import RunPipeline +from loafer.contracts import RunEvent, RunRequest, StageStatus +from loafer.core.run_state import RetryCategory, RunState, StageState +from loafer.exceptions import PipelineError +from loafer.metadata import RunLease +from loafer.ports.metadata import MetadataStore +from loafer.ports.object_storage import ObjectStoragePort + + +class RejectUnapprovedTransform: + """Workers never grant interactive approval to newly generated code.""" + + def approve_transform(self, generated_code: str) -> bool: + del generated_code + return False + + +class DurableWorker: + """Claim and execute immutable runs under leases and fencing tokens.""" + + def __init__( + self, + metadata: MetadataStore, + objects: ObjectStoragePort, + *, + worker_id: str, + lease_for: timedelta = timedelta(seconds=30), + retry_delay: timedelta = timedelta(seconds=5), + max_attempts: int = 3, + ) -> None: + if max_attempts < 1: + raise ValueError("max_attempts must be positive") + self._metadata = metadata + self._objects = objects + self._worker_id = worker_id + self._lease_for = lease_for + self._retry_delay = retry_delay + self._max_attempts = max_attempts + + def run_once(self) -> str | None: + """Execute at most one runnable job, returning its run ID.""" + lease = self._metadata.claim_run(self._worker_id, self._lease_for) + if lease is None: + return None + self.execute(lease) + return lease.run.id + + def execute(self, lease: RunLease) -> None: + """Execute one claimed run and persist all observable outcomes.""" + self._metadata.transition_run(lease, RunState.RUNNING) + version = self._metadata.get_pipeline_version(lease.run.pipeline_version_id) + config_document = version.config.get("document", version.config) + recovery = DurableBatchRecovery(self._metadata, self._objects, lease) + use_case = RunPipeline( + cancellation=MetadataCancellation(self._metadata), + checkpoints=NullCheckpointStore(), + secrets=EnvironmentSecretResolver(), + events=NullEventPublisher(), + reviewer=RejectUnapprovedTransform(), + recovery=recovery, + ) + + try: + with tempfile.TemporaryDirectory(prefix="loafer-run-") as directory: + config_path = Path(directory) / "pipeline.json" + config_path.write_text(json.dumps(config_document), encoding="utf-8") + request = RunRequest(config_path=str(config_path), run_id=lease.run.id) + for engine_event in use_case.stream(request): + lease = self._metadata.heartbeat(lease, self._lease_for) + self._record_engine_event(lease, engine_event) + self._metadata.transition_run(lease, RunState.SUCCEEDED) + except Exception as exc: + cancelled = self._metadata.cancellation_requested(lease.run.id) + if cancelled: + self._metadata.transition_run( + lease, + RunState.CANCELLED, + error={"type": type(exc).__name__, "message": str(exc)}, + ) + return + if lease.run.attempt + 1 < self._max_attempts: + checkpoint = self._metadata.latest_checkpoint(lease.run.id, "default") + category = ( + RetryCategory.FAILED_BATCH + if checkpoint is not None + else RetryCategory.INFRASTRUCTURE + ) + self._metadata.transition_run( + lease, + RunState.RETRY_WAIT, + error={"type": type(exc).__name__, "message": str(exc)}, + retry_category=category, + retry_at=_utc_after(self._retry_delay), + ) + return + self._metadata.transition_run( + lease, + RunState.FAILED, + error={"type": type(exc).__name__, "message": str(exc)}, + ) + if isinstance(exc, PipelineError): + return + raise + + def run_forever(self, poll_interval: float = 1.0) -> None: + """Poll for durable commands until the process is interrupted.""" + try: + while True: + if self.run_once() is None: + time.sleep(poll_interval) + except (KeyboardInterrupt, SystemExit): + return + + def close(self) -> None: + """Release adapter resources owned by this worker composition.""" + close = getattr(self._metadata, "close", None) + if close is not None: + close() + + def _record_engine_event(self, lease: RunLease, event: RunEvent) -> None: + if event.stage not in {"batch", "recovery"}: + if event.status in { + StageStatus.DONE, + StageStatus.FAILED, + StageStatus.CANCELLED, + }: + self._metadata.transition_stage( + lease, + event.stage, + StageState.RUNNING, + ) + self._metadata.transition_stage( + lease, + event.stage, + _stage_state(event.status), + ) + self._metadata.append_event( + lease, + f"engine.{event.stage}.{event.status.value}", + event.model_dump(mode="json", exclude={"sequence"}), + ) + + +def _stage_state(status: StageStatus) -> StageState: + return { + StageStatus.RUNNING: StageState.RUNNING, + StageStatus.DONE: StageState.SUCCEEDED, + StageStatus.SKIPPED: StageState.SKIPPED, + StageStatus.FAILED: StageState.FAILED, + StageStatus.CANCELLED: StageState.CANCELLED, + }[status] + + +def _utc_after(delta: timedelta) -> Any: + from loafer.metadata import utc_now + + return utc_now() + delta diff --git a/skills/loafer-engineering/references/architecture.md b/skills/loafer-engineering/references/architecture.md index ef62a98..e070ec0 100644 --- a/skills/loafer-engineering/references/architecture.md +++ b/skills/loafer-engineering/references/architecture.md @@ -50,6 +50,15 @@ YAML → ELT: load_raw target adapter → in-target SQL transform → engine persists the local incremental cursor after graph completion → application emits sanitized RunEvent / RunResult contracts + +Durable single-node mode adds a separate path around that application use case: + +```text +enqueue/scheduler → immutable pipeline version + idempotent run command + outbox +worker claim → expiring lease + fencing token → bounded engine execution +batch output object → transactional batch/checkpoint/event commit → attempt-local target +worker restart → replay committed objects → skip durable source offset → final publication +``` ``` `PipelineState` mixes configuration, data, execution metadata, live iterators, provider objects, and @@ -92,6 +101,9 @@ Verify before relying on it, but the repository currently contains: quarantine, checksums, cancellation, and atomic CSV/JSON publication. - Python transform subprocess sandboxing on supported operating systems. - CLI validation, connector listing, runs, scheduling, daemon management, logs, and initialization. +- Versioned SQLite/PostgreSQL metadata, state machines, sequenced events, leases/fencing, + idempotent commands, outbox records, filesystem object storage, and single-node bounded-batch + recovery through a separately runnable worker. - Unit, integration, end-to-end, smoke, and opt-in benchmark tests. ## Roadmap gaps @@ -108,7 +120,8 @@ Re-check the repository because this reference is a snapshot, not a substitute f - Undeclared/materialized AI, custom, SQL, and multi-step ETL still drain the source stream through `materialize_input_rows()`; only explicitly declared row-local work uses the bounded data plane. - The legacy graph path remains sample-validated; the bounded row-local path validates every batch. -- Local JSON watermark state is not sufficient for concurrent or distributed workers. +- Local JSON watermark state remains on the synchronous legacy path; durable worker runs use the + metadata checkpoint store, while SQLite is intentionally limited to one scheduler and worker. - CSV and JSON targets publish atomically, but the local watermark/state file does not yet use the same publication protocol. - PostgreSQL target/ELT identifiers are safely composed, but every future SQL adapter and diff --git a/tests/e2e/test_durable_recovery.py b/tests/e2e/test_durable_recovery.py new file mode 100644 index 0000000..28b220d --- /dev/null +++ b/tests/e2e/test_durable_recovery.py @@ -0,0 +1,161 @@ +"""Single-node crash recovery at the durable batch commit boundary.""" + +from __future__ import annotations + +import hashlib +import json +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any + +import pytest + +from loafer.adapters.metadata import SqlMetadataStore +from loafer.adapters.object_storage import FilesystemObjectStorage +from loafer.adapters.runtime import ( + DurableBatchRecovery, + EnvironmentSecretResolver, + MetadataCancellation, + NullCheckpointStore, + NullEventPublisher, +) +from loafer.application.service import RunPipeline +from loafer.config import load_config +from loafer.contracts import RunRequest +from loafer.core.run_state import RunState +from loafer.worker import DurableWorker, RejectUnapprovedTransform + + +class Clock: + def __init__(self) -> None: + self.now = datetime(2026, 8, 4, 12, 0, tzinfo=UTC) + + def __call__(self) -> datetime: + return self.now + + +class CrashAfterCommit: + def __init__(self, recovery: DurableBatchRecovery, crash_after: int) -> None: + self._recovery = recovery + self._crash_after = crash_after + self._commits = 0 + + def restore(self, run_id: str, partition_id: str) -> Any: + return self._recovery.restore(run_id, partition_id) + + def read_rows(self, commit: Any) -> Any: + return self._recovery.read_rows(commit) + + def commit(self, envelope: Any, rows: list[dict[str, Any]]) -> Any: + checkpoint = self._recovery.commit(envelope, rows) + self._commits += 1 + if self._commits == self._crash_after: + raise SystemExit("simulated worker kill after durable commit") + return checkpoint + + +def _config(tmp_path: Path) -> tuple[Path, Path]: + source = tmp_path / "input.csv" + source.write_text( + "id,name\n" + "".join(f"{index},name-{index}\n" for index in range(1, 6)), + encoding="utf-8", + ) + transform = tmp_path / "transform.py" + transform.write_text( + "def transform(data):\n return [{**row, 'name': row['name'].upper()} for row in data]\n", + encoding="utf-8", + ) + output = tmp_path / "output.json" + config = tmp_path / "pipeline.yaml" + config.write_text( + "\n".join( + [ + "name: durable-recovery", + "mode: etl", + "source:", + " type: csv", + f" path: {source}", + "target:", + " type: json", + f" path: {output}", + "transform:", + " type: custom", + f" path: {transform}", + "execution:", + " transform_class: row_local", + " schema_drift: fail", + "chunk_size: 2", + "", + ] + ), + encoding="utf-8", + ) + return config, output + + +@pytest.mark.parametrize("crash_after", [1, 2, 3]) +def test_reclaimed_worker_resumes_from_each_durable_batch( + tmp_path: Path, + crash_after: int, +) -> None: + config_path, output = _config(tmp_path) + clock = Clock() + metadata = SqlMetadataStore(f"sqlite:///{tmp_path / 'metadata.db'}", clock=clock) + metadata.migrate() + objects = FilesystemObjectStorage(tmp_path / "objects") + config_document = load_config(config_path).model_dump(mode="json") + digest = hashlib.sha256(json.dumps(config_document, sort_keys=True).encode()).hexdigest() + version = metadata.register_pipeline_version( + workspace_id="local", + pipeline_key="durable-recovery", + config_digest=digest, + config={"document": config_document}, + ) + metadata.create_run( + workspace_id="local", + pipeline_version_id=version.id, + command_key="manual:test", + run_id="recovery-run", + ) + stale_lease = metadata.claim_run("worker-a", timedelta(seconds=5)) + assert stale_lease is not None + metadata.transition_run(stale_lease, RunState.RUNNING) + crashing_recovery = CrashAfterCommit( + DurableBatchRecovery(metadata, objects, stale_lease), + crash_after, + ) + first_attempt = RunPipeline( + cancellation=MetadataCancellation(metadata), + checkpoints=NullCheckpointStore(), + secrets=EnvironmentSecretResolver(), + events=NullEventPublisher(), + reviewer=RejectUnapprovedTransform(), + recovery=crashing_recovery, + ) + + with pytest.raises(SystemExit, match="simulated worker kill"): + list(first_attempt.stream(RunRequest(config_path=str(config_path), run_id="recovery-run"))) + assert not output.exists() + checkpoint = metadata.latest_checkpoint("recovery-run", "default") + assert checkpoint is not None + assert checkpoint.batch_id == f"batch-{crash_after:08d}" + + clock.now += timedelta(seconds=6) + current_lease = metadata.claim_run("worker-b", timedelta(seconds=30)) + assert current_lease is not None + DurableWorker(metadata, objects, worker_id="worker-b").execute(current_lease) + + assert metadata.get_run("recovery-run").state is RunState.SUCCEEDED + published = json.loads(output.read_text(encoding="utf-8")) + assert [row["id"] for row in published] == ["1", "2", "3", "4", "5"] + assert [row["name"] for row in published] == [ + "NAME-1", + "NAME-2", + "NAME-3", + "NAME-4", + "NAME-5", + ] + assert len(metadata.list_batch_commits("recovery-run", "default")) == 3 + sequences = [event.sequence for event in metadata.list_events("recovery-run")] + assert sequences == list(range(1, len(sequences) + 1)) + metadata.close() diff --git a/tests/integration/test_metadata_store.py b/tests/integration/test_metadata_store.py new file mode 100644 index 0000000..125593b --- /dev/null +++ b/tests/integration/test_metadata_store.py @@ -0,0 +1,80 @@ +"""PostgreSQL contract tests for authoritative durable metadata.""" + +from __future__ import annotations + +from datetime import timedelta + +import pytest +from sqlalchemy import inspect + +from loafer.adapters.metadata import SqlMetadataStore +from loafer.core.run_state import RunState +from loafer.exceptions import StaleFenceError + +pytestmark = pytest.mark.integration + + +@pytest.fixture() +def metadata(postgres_url: str) -> SqlMetadataStore: + store = SqlMetadataStore(postgres_url) + store.migrate(0) + store.migrate() + try: + yield store + finally: + store.migrate(0) + store.close() + + +def test_postgres_empty_schema_and_previous_schema_upgrade( + postgres_url: str, +) -> None: + store = SqlMetadataStore(postgres_url) + try: + store.migrate(0) + assert store.migrate(1) == 1 + version = store.register_pipeline_version( + workspace_id="pg-workspace", + pipeline_key="customers", + config_digest="a" * 64, + config={"document": {"name": "customers"}}, + ) + + assert store.migrate() == 2 + assert store.get_pipeline_version(version.id).config_digest == "a" * 64 + assert "loafer_outbox" in inspect(store.engine).get_table_names() + + assert store.migrate(1) == 1 + assert "loafer_outbox" not in inspect(store.engine).get_table_names() + assert store.migrate() == 2 + finally: + store.migrate(0) + store.close() + + +def test_postgres_claims_are_fenced_and_events_are_monotonic( + metadata: SqlMetadataStore, +) -> None: + version = metadata.register_pipeline_version( + workspace_id="pg-workspace", + pipeline_key="orders", + config_digest="b" * 64, + config={"document": {"name": "orders"}}, + ) + metadata.create_run( + workspace_id="pg-workspace", + pipeline_version_id=version.id, + command_key="request-1", + run_id="pg-run-1", + ) + lease = metadata.claim_run("worker-a", timedelta(seconds=30)) + assert lease is not None + metadata.transition_run(lease, RunState.RUNNING) + metadata.append_event(lease, "worker.progress", {"rows": 10}) + metadata.transition_run(lease, RunState.SUCCEEDED) + + events = metadata.list_events("pg-run-1") + assert [event.sequence for event in events] == list(range(1, len(events) + 1)) + + with pytest.raises(StaleFenceError): + metadata.append_event(lease, "worker.late", {}) diff --git a/tests/unit/test_metadata_store.py b/tests/unit/test_metadata_store.py new file mode 100644 index 0000000..7ae6472 --- /dev/null +++ b/tests/unit/test_metadata_store.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from pathlib import Path + +import pytest +from sqlalchemy import inspect, select +from sqlalchemy.exc import IntegrityError + +from loafer.adapters import metadata_schema as schema +from loafer.adapters.metadata import SqlMetadataStore +from loafer.adapters.object_storage import MemoryObjectStorage +from loafer.adapters.runtime import DurableBatchRecovery +from loafer.contracts import BatchEnvelope +from loafer.core.run_state import RunState +from loafer.exceptions import IdempotencyConflictError, StaleFenceError +from loafer.metadata import ScheduleRecord + + +class Clock: + def __init__(self) -> None: + self.now = datetime(2026, 8, 4, tzinfo=UTC) + + def __call__(self) -> datetime: + return self.now + + def advance(self, **kwargs: int) -> None: + self.now += timedelta(**kwargs) + + +@pytest.fixture() +def store(tmp_path: Path) -> tuple[SqlMetadataStore, Clock]: + clock = Clock() + metadata = SqlMetadataStore(f"sqlite:///{tmp_path / 'metadata.db'}", clock=clock) + assert metadata.migrate() == 2 + try: + yield metadata, clock + finally: + metadata.close() + + +def _version(metadata: SqlMetadataStore) -> str: + return metadata.register_pipeline_version( + workspace_id="workspace-1", + pipeline_key="customers", + config_digest="a" * 64, + config={"document": {"name": "customers"}}, + ).id + + +def _envelope(run_id: str, batch_id: str = "batch-00000001") -> BatchEnvelope: + return BatchEnvelope( + run_id=run_id, + stage_id="load", + partition_id="default", + batch_id=batch_id, + attempt=0, + source_position_start={"offset": 0}, + source_position_end={"offset": 1}, + schema_version="schema-1", + rows_in=2, + rows_out=2, + rows_rejected=0, + bytes_in=10, + bytes_out=10, + output_checksum="b" * 64, + ) + + +def test_run_creation_is_idempotent_and_conflicts_are_rejected( + store: tuple[SqlMetadataStore, Clock], +) -> None: + metadata, _clock = store + version_id = _version(metadata) + + first = metadata.create_run( + workspace_id="workspace-1", + pipeline_version_id=version_id, + command_key="request-1", + run_id="run-1", + ) + repeated = metadata.create_run( + workspace_id="workspace-1", + pipeline_version_id=version_id, + command_key="request-1", + run_id="ignored", + ) + + assert first == repeated + assert first.state is RunState.QUEUED + assert [event.sequence for event in metadata.list_events("run-1")] == [1] + + other_version = metadata.register_pipeline_version( + workspace_id="workspace-1", + pipeline_key="customers", + config_digest="c" * 64, + config={"document": {"name": "changed"}}, + ) + with pytest.raises(IdempotencyConflictError): + metadata.create_run( + workspace_id="workspace-1", + pipeline_version_id=other_version.id, + command_key="request-1", + ) + + +def test_batch_commit_is_replayable_and_event_sequences_are_monotonic( + store: tuple[SqlMetadataStore, Clock], +) -> None: + metadata, _clock = store + run = metadata.create_run( + workspace_id="workspace-1", + pipeline_version_id=_version(metadata), + command_key="request-1", + run_id="run-1", + ) + lease = metadata.claim_run("worker-a", timedelta(seconds=30)) + assert lease is not None and lease.run.id == run.id + metadata.transition_run(lease, RunState.RUNNING) + recovery = DurableBatchRecovery(metadata, MemoryObjectStorage(), lease) + envelope = _envelope(run.id) + + checkpoint = recovery.commit(envelope, [{"id": 1}, {"id": 2}]) + repeated = recovery.commit(envelope, [{"id": 1}, {"id": 2}]) + commits = recovery.restore(run.id, "default") + + assert repeated == checkpoint + assert metadata.latest_checkpoint(run.id, "default") == checkpoint + assert recovery.read_rows(commits[0]) == [{"id": 1}, {"id": 2}] + sequences = [event.sequence for event in metadata.list_events(run.id)] + assert sequences == list(range(1, len(sequences) + 1)) + assert [event.event_type for event in metadata.list_events(run.id)].count( + "batch.committed" + ) == 1 + + +def test_expired_worker_is_fenced_after_reclaim( + store: tuple[SqlMetadataStore, Clock], +) -> None: + metadata, clock = store + metadata.create_run( + workspace_id="workspace-1", + pipeline_version_id=_version(metadata), + command_key="request-1", + run_id="run-1", + ) + stale = metadata.claim_run("worker-a", timedelta(seconds=5)) + assert stale is not None + metadata.transition_run(stale, RunState.RUNNING) + + clock.advance(seconds=6) + current = metadata.claim_run("worker-b", timedelta(seconds=30)) + assert current is not None + assert current.fencing_token == stale.fencing_token + 1 + + with pytest.raises(StaleFenceError): + metadata.append_event(stale, "worker.late", {}) + with pytest.raises(StaleFenceError): + DurableBatchRecovery(metadata, MemoryObjectStorage(), stale).commit( + _envelope("run-1"), [{"id": 1}] + ) + + +def test_migrations_upgrade_previous_schema_rollback_and_reapply(tmp_path: Path) -> None: + metadata = SqlMetadataStore(f"sqlite:///{tmp_path / 'migration.db'}") + try: + assert metadata.migrate(1) == 1 + version_id = _version(metadata) + + assert metadata.migrate() == 2 + assert metadata.get_pipeline_version(version_id).pipeline_key == "customers" + assert "loafer_outbox" in inspect(metadata.engine).get_table_names() + + assert metadata.migrate(1) == 1 + assert "loafer_outbox" not in inspect(metadata.engine).get_table_names() + assert metadata.get_pipeline_version(version_id).pipeline_key == "customers" + + assert metadata.migrate() == 2 + assert "loafer_outbox" in inspect(metadata.engine).get_table_names() + finally: + metadata.close() + + +def test_database_constraints_enforce_null_unique_foreign_key_and_check( + store: tuple[SqlMetadataStore, Clock], +) -> None: + metadata, clock = store + version_id = _version(metadata) + + with ( + pytest.raises(IntegrityError, match=r"NOT NULL constraint failed|not-null constraint"), + metadata.engine.begin() as connection, + ): + connection.execute( + schema.pipeline_versions.insert().values( + id="missing-key", + workspace_id="workspace-1", + config_digest="d" * 64, + config_json={}, + created_at=clock(), + ) + ) + + with ( + pytest.raises(IntegrityError, match=r"UNIQUE constraint failed|unique constraint"), + metadata.engine.begin() as connection, + ): + connection.execute( + schema.pipeline_versions.insert().values( + id="duplicate-version", + workspace_id="workspace-1", + pipeline_key="customers", + config_digest="a" * 64, + config_json={}, + created_at=clock(), + ) + ) + + with ( + pytest.raises( + IntegrityError, match=r"FOREIGN KEY constraint failed|foreign key constraint" + ), + metadata.engine.begin() as connection, + ): + connection.execute( + schema.runs.insert().values( + id="dangling", + workspace_id="workspace-1", + pipeline_version_id="does-not-exist", + command_key="dangling", + state=RunState.QUEUED.value, + attempt=0, + cancel_requested=False, + next_event_sequence=1, + fencing_token=0, + created_at=clock(), + ) + ) + + with ( + pytest.raises(IntegrityError, match=r"ck_loafer_run_state|CHECK constraint failed"), + metadata.engine.begin() as connection, + ): + connection.execute( + schema.runs.insert().values( + id="invalid", + workspace_id="workspace-1", + pipeline_version_id=version_id, + command_key="invalid", + state="teleported", + attempt=0, + cancel_requested=False, + next_event_sequence=1, + fencing_token=0, + created_at=clock(), + ) + ) + + with metadata.engine.connect() as connection: + assert ( + connection.execute( + select(schema.runs.c.id).where(schema.runs.c.id == "invalid") + ).scalar_one_or_none() + is None + ) + + +def test_due_schedule_creates_one_idempotent_command_and_advances( + store: tuple[SqlMetadataStore, Clock], +) -> None: + metadata, clock = store + version_id = _version(metadata) + schedule = ScheduleRecord( + id="hourly-customers", + workspace_id="workspace-1", + pipeline_version_id=version_id, + trigger_kind="interval", + trigger_spec="1h", + timezone="UTC", + enabled=True, + next_run_at=clock(), + created_at=clock(), + updated_at=clock(), + ) + metadata.upsert_schedule(schedule) + + first = metadata.enqueue_due_schedules(clock()) + repeated = metadata.enqueue_due_schedules(clock()) + + assert len(first) == 1 + assert repeated == [] + assert first[0].command_key == f"schedule:hourly-customers:{clock().isoformat()}" + assert [item.event_type for item in metadata.pending_outbox()] == ["run.created"] + + +def test_cancel_command_is_idempotent_before_claim( + store: tuple[SqlMetadataStore, Clock], +) -> None: + metadata, _clock = store + run = metadata.create_run( + workspace_id="workspace-1", + pipeline_version_id=_version(metadata), + command_key="request-cancel", + run_id="cancel-me", + ) + + cancelled = metadata.request_cancel(run.id) + repeated = metadata.request_cancel(run.id) + + assert cancelled.state is RunState.CANCELLED + assert repeated.state is RunState.CANCELLED + assert metadata.claim_run("worker-a", timedelta(seconds=30)) is None diff --git a/tests/unit/test_object_storage.py b/tests/unit/test_object_storage.py new file mode 100644 index 0000000..1e03076 --- /dev/null +++ b/tests/unit/test_object_storage.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from loafer.adapters.object_storage import FilesystemObjectStorage, MemoryObjectStorage +from loafer.exceptions import MetadataError + + +@pytest.mark.parametrize("adapter", ["filesystem", "memory"]) +def test_object_storage_contract(adapter: str, tmp_path: Path) -> None: + storage = ( + FilesystemObjectStorage(tmp_path / "objects") + if adapter == "filesystem" + else MemoryObjectStorage() + ) + artifact = storage.put( + "runs/run-1/logs/worker.log", + [b"first\n", b"second\n"], + kind="log", + run_id="run-1", + ) + + assert storage.exists(artifact.uri) + assert storage.read(artifact.uri) == b"first\nsecond\n" + assert artifact.size_bytes == 13 + assert len(artifact.checksum) == 64 + + storage.delete(artifact.uri) + assert not storage.exists(artifact.uri) + + +def test_filesystem_storage_rejects_path_escape(tmp_path: Path) -> None: + storage = FilesystemObjectStorage(tmp_path / "objects") + + with pytest.raises(MetadataError, match="unsafe object key"): + storage.put("../secret", b"nope", kind="artifact") diff --git a/tests/unit/test_run_state.py b/tests/unit/test_run_state.py new file mode 100644 index 0000000..67a653e --- /dev/null +++ b/tests/unit/test_run_state.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import pytest + +from loafer.core.run_state import ( + BatchState, + RunState, + StageState, + require_batch_transition, + require_run_transition, + require_stage_transition, +) +from loafer.exceptions import InvalidStateTransitionError + + +@pytest.mark.parametrize( + ("current", "target"), + [ + (RunState.QUEUED, RunState.CLAIMED), + (RunState.CLAIMED, RunState.RUNNING), + (RunState.RUNNING, RunState.SUCCEEDED), + (RunState.RUNNING, RunState.RETRY_WAIT), + ], +) +def test_valid_run_transitions(current: RunState, target: RunState) -> None: + require_run_transition(current, target) + + +def test_terminal_run_cannot_restart() -> None: + with pytest.raises(InvalidStateTransitionError, match="succeeded -> running"): + require_run_transition(RunState.SUCCEEDED, RunState.RUNNING) + + +def test_stage_and_batch_machines_reject_impossible_commits() -> None: + with pytest.raises(InvalidStateTransitionError, match="pending -> succeeded"): + require_stage_transition(StageState.PENDING, StageState.SUCCEEDED) + with pytest.raises(InvalidStateTransitionError, match="pending -> committed"): + require_batch_transition(BatchState.PENDING, BatchState.COMMITTED) + + +def test_idempotent_state_transition_is_allowed() -> None: + require_run_transition(RunState.RUNNING, RunState.RUNNING) diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index 79e1a6b..955807c 100644 --- a/tests/unit/test_scheduler.py +++ b/tests/unit/test_scheduler.py @@ -3,6 +3,7 @@ from __future__ import annotations from pathlib import Path +from types import SimpleNamespace from typing import Any import pytest @@ -10,6 +11,23 @@ from loafer.exceptions import SchedulerError +def test_scheduled_callback_only_enqueues_durable_command(monkeypatch: Any) -> None: + from loafer import scheduler as scheduler_module + + captured: dict[str, str] = {} + + def enqueue(config_path: str, *, command_key: str) -> Any: + captured.update(config_path=config_path, command_key=command_key) + return SimpleNamespace(id="durable-run-1") + + monkeypatch.setattr(scheduler_module, "enqueue_pipeline", enqueue) + + scheduler_module._run_pipeline_job("pipeline.yaml", "customers", "hourly") + + assert captured["config_path"] == "pipeline.yaml" + assert captured["command_key"].startswith("schedule:hourly:") + + class TestPipelineScheduler: """Tests for the PipelineScheduler class."""