From 82ebb1931e7f5848442cca8048e1db13643a5cf8 Mon Sep 17 00:00:00 2001 From: pkuwkl Date: Fri, 14 Aug 2026 16:38:00 +0800 Subject: [PATCH 1/2] feat(etl): add agent-native observable pipelines Add whole-run and micro-batch async ETL scaffolds with atomic local run state, machine-readable receipts, explicit dry-run semantics, and agent-facing progress. Document the architecture boundaries, examples, import contracts, and failure/cancellation observation behavior. Closes #147. --- .agents/skills/quantmind-dev/SKILL.md | 2 +- .../references/develop-components.md | 31 +- .claude/skills/quantmind-dev/SKILL.md | 2 +- .../references/develop-components.md | 31 +- .gitignore | 2 + AGENTS.md | 8 + contexts/CONTEXT_MAP.md | 1 + contexts/design/README.md | 1 + contexts/design/operations/etl.md | 136 +++ contexts/usage/README.md | 5 +- docs/README.md | 9 +- docs/etl.md | 144 +++ examples/etl/batch_local_artifacts.py | 132 +++ examples/etl/local_artifact.py | 106 ++ pyproject.toml | 33 +- quantmind/etl/__init__.py | 19 + quantmind/etl/_batch.py | 453 +++++++ quantmind/etl/_pipeline.py | 272 +++++ quantmind/etl/_record.py | 590 +++++++++ tests/etl/__init__.py | 0 tests/etl/test_batch.py | 1053 +++++++++++++++++ tests/etl/test_examples.py | 144 +++ tests/etl/test_pipeline.py | 784 ++++++++++++ 23 files changed, 3939 insertions(+), 19 deletions(-) create mode 100644 contexts/design/operations/etl.md create mode 100644 docs/etl.md create mode 100644 examples/etl/batch_local_artifacts.py create mode 100644 examples/etl/local_artifact.py create mode 100644 quantmind/etl/__init__.py create mode 100644 quantmind/etl/_batch.py create mode 100644 quantmind/etl/_pipeline.py create mode 100644 quantmind/etl/_record.py create mode 100644 tests/etl/__init__.py create mode 100644 tests/etl/test_batch.py create mode 100644 tests/etl/test_examples.py create mode 100644 tests/etl/test_pipeline.py diff --git a/.agents/skills/quantmind-dev/SKILL.md b/.agents/skills/quantmind-dev/SKILL.md index 50391fa..97aa6ff 100644 --- a/.agents/skills/quantmind-dev/SKILL.md +++ b/.agents/skills/quantmind-dev/SKILL.md @@ -1,6 +1,6 @@ --- name: quantmind-dev -description: Contributor workflow for the QuantMind codebase. Covers contributor setup (environment + hooks), filing issues, commit format, pull request format, and component development across quantmind/ modules (knowledge, configs, preprocess, rag, flows, mind, utils) with tests, examples, and verification. Use when setting up as a contributor, filing an issue, committing, opening a PR, or implementing/refactoring QuantMind code. +description: Contributor workflow for the QuantMind codebase. Covers contributor setup (environment + hooks), filing issues, commit format, pull request format, and component development across quantmind/ modules (etl, knowledge, configs, preprocess, rag, flows, mind, utils) with tests, examples, and verification. Use when setting up as a contributor, filing an issue, committing, opening a PR, or implementing/refactoring QuantMind code. --- # QuantMind Dev diff --git a/.agents/skills/quantmind-dev/references/develop-components.md b/.agents/skills/quantmind-dev/references/develop-components.md index 71bd98b..f158543 100644 --- a/.agents/skills/quantmind-dev/references/develop-components.md +++ b/.agents/skills/quantmind-dev/references/develop-components.md @@ -39,6 +39,7 @@ apply throughout. | Module | May import from `quantmind.*` | |--------|-------------------------------| +| `quantmind/etl/` | nothing (independent leaf) | | `quantmind/utils/` | nothing (leaf) | | `quantmind/knowledge/` | nothing (leaf) | | `quantmind/configs/` | `knowledge` only | @@ -46,7 +47,26 @@ apply throughout. | `quantmind/rag/` | `preprocess` only | | `quantmind/library/` | `knowledge` only | | `quantmind/mind/` | `knowledge`, `configs`, `utils` (retrieval is library-free; the `library` edge is reserved for the future collection path, not single-tree `retrieve`) | -| `quantmind/flows/`, `quantmind/magic.py` | apex — may import all of the above | +| `quantmind/flows/`, `quantmind/magic.py` | apex — may import domain layers above, but not the independent `etl` scaffold | + +### `quantmind/etl/` — observable whole-run and micro-batch ETL + +- Bind exactly three async stage callables to `ETLPipeline` for one whole-run + delivery. Use the parallel `BatchETLPipeline` when an async producer yields + business batches that each pass through transform and load. Never switch + execution shape by inspecting a callable's return value, and never hide batch + loads inside a whole-run transform. +- Use composition rather than an ABC, subclass tree, or inheritance between the + two pipeline classes. Keep batch execution strictly serial unless a later + observation contract explicitly represents simultaneously active stages. +- Keep it independent of every other `quantmind.*` package. Existing flows do + not inherit it; their pure `input → artifact` contract remains unchanged. +- Report only real completed work through `PipelineContext.progress()`. In batch + mode, only a load that returns successfully increments the completed-batch + count; partial-write safety remains the business load's responsibility. The + scaffold owns its local lifecycle snapshots; do not add custom run-state + files, a CLI, heartbeat, scheduler, retry policy, checkpoint/resume, or + workflow engine. See `contexts/design/operations/etl.md`. ### `quantmind/knowledge/` — data standard @@ -144,10 +164,13 @@ apply throughout. A public operation is complete only when all of these agree: 1. A stage and name consistent with `contexts/design/operations/naming.md`. -2. Typed input and config models, exported from `quantmind.configs`. +2. Typed input and config models, when the operation has them, exported from + the canonical owning package (`quantmind.configs` for flow configs, + `quantmind.etl` for ETL run contracts, or another explicit owner). 3. One intent-oriented async function, small service class, or document-scoped - handle exported from `quantmind.flows`, with its result contract exported - from the canonical owning layer. + handle exported from its canonical owning package (`quantmind.flows`, + `quantmind.etl`, `quantmind.library`, etc.), with its result contract + exported from the same owning layer. 4. Offline success and failure tests for the public callable, plus a magic-introspection test when a function follows the `(input, *, cfg)` convention. diff --git a/.claude/skills/quantmind-dev/SKILL.md b/.claude/skills/quantmind-dev/SKILL.md index 50391fa..97aa6ff 100644 --- a/.claude/skills/quantmind-dev/SKILL.md +++ b/.claude/skills/quantmind-dev/SKILL.md @@ -1,6 +1,6 @@ --- name: quantmind-dev -description: Contributor workflow for the QuantMind codebase. Covers contributor setup (environment + hooks), filing issues, commit format, pull request format, and component development across quantmind/ modules (knowledge, configs, preprocess, rag, flows, mind, utils) with tests, examples, and verification. Use when setting up as a contributor, filing an issue, committing, opening a PR, or implementing/refactoring QuantMind code. +description: Contributor workflow for the QuantMind codebase. Covers contributor setup (environment + hooks), filing issues, commit format, pull request format, and component development across quantmind/ modules (etl, knowledge, configs, preprocess, rag, flows, mind, utils) with tests, examples, and verification. Use when setting up as a contributor, filing an issue, committing, opening a PR, or implementing/refactoring QuantMind code. --- # QuantMind Dev diff --git a/.claude/skills/quantmind-dev/references/develop-components.md b/.claude/skills/quantmind-dev/references/develop-components.md index 71bd98b..f158543 100644 --- a/.claude/skills/quantmind-dev/references/develop-components.md +++ b/.claude/skills/quantmind-dev/references/develop-components.md @@ -39,6 +39,7 @@ apply throughout. | Module | May import from `quantmind.*` | |--------|-------------------------------| +| `quantmind/etl/` | nothing (independent leaf) | | `quantmind/utils/` | nothing (leaf) | | `quantmind/knowledge/` | nothing (leaf) | | `quantmind/configs/` | `knowledge` only | @@ -46,7 +47,26 @@ apply throughout. | `quantmind/rag/` | `preprocess` only | | `quantmind/library/` | `knowledge` only | | `quantmind/mind/` | `knowledge`, `configs`, `utils` (retrieval is library-free; the `library` edge is reserved for the future collection path, not single-tree `retrieve`) | -| `quantmind/flows/`, `quantmind/magic.py` | apex — may import all of the above | +| `quantmind/flows/`, `quantmind/magic.py` | apex — may import domain layers above, but not the independent `etl` scaffold | + +### `quantmind/etl/` — observable whole-run and micro-batch ETL + +- Bind exactly three async stage callables to `ETLPipeline` for one whole-run + delivery. Use the parallel `BatchETLPipeline` when an async producer yields + business batches that each pass through transform and load. Never switch + execution shape by inspecting a callable's return value, and never hide batch + loads inside a whole-run transform. +- Use composition rather than an ABC, subclass tree, or inheritance between the + two pipeline classes. Keep batch execution strictly serial unless a later + observation contract explicitly represents simultaneously active stages. +- Keep it independent of every other `quantmind.*` package. Existing flows do + not inherit it; their pure `input → artifact` contract remains unchanged. +- Report only real completed work through `PipelineContext.progress()`. In batch + mode, only a load that returns successfully increments the completed-batch + count; partial-write safety remains the business load's responsibility. The + scaffold owns its local lifecycle snapshots; do not add custom run-state + files, a CLI, heartbeat, scheduler, retry policy, checkpoint/resume, or + workflow engine. See `contexts/design/operations/etl.md`. ### `quantmind/knowledge/` — data standard @@ -144,10 +164,13 @@ apply throughout. A public operation is complete only when all of these agree: 1. A stage and name consistent with `contexts/design/operations/naming.md`. -2. Typed input and config models, exported from `quantmind.configs`. +2. Typed input and config models, when the operation has them, exported from + the canonical owning package (`quantmind.configs` for flow configs, + `quantmind.etl` for ETL run contracts, or another explicit owner). 3. One intent-oriented async function, small service class, or document-scoped - handle exported from `quantmind.flows`, with its result contract exported - from the canonical owning layer. + handle exported from its canonical owning package (`quantmind.flows`, + `quantmind.etl`, `quantmind.library`, etc.), with its result contract + exported from the same owning layer. 4. Offline success and failure tests for the public callable, plus a magic-introspection test when a function follows the `(input, *, cfg)` convention. diff --git a/.gitignore b/.gitignore index 08813f2..44f2d48 100644 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,7 @@ docs/superpowers/ # Coverage artifacts (generated by pytest --cov) .coverage +.coverage.* htmlcov/ coverage.xml @@ -36,5 +37,6 @@ coverage.xml .DS_Store # Ephemeral local scratch (temp dirs, e2e harness output) +.quant-mind/ temp/ tmp/ diff --git a/AGENTS.md b/AGENTS.md index accb431..2b65df6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -46,6 +46,7 @@ The canonical, always-current statement lives in | Module | Role | |--------|------| +| `quantmind/etl/` | Stdlib-only, observable whole-run and micro-batch `extract → transform → load` authoring scaffolds — independent leaf | | `quantmind/knowledge/` | Pydantic data standard (`FlattenKnowledge` / `TreeKnowledge` / `GraphKnowledge`) — dependency leaf | | `quantmind/library/` | Local persistence and semantic retrieval for canonical knowledge — depends only on `knowledge` | | `quantmind/configs/` | Operation cfg + typed input models or unions (`BaseFlowCfg`, `NewsWindow`, `PaperInput`) — depends only on `knowledge` | @@ -131,6 +132,13 @@ the user explicitly authorizes it — fix the underlying issue instead. `BaseKnowledge`. Accept modest redundancy to keep artifacts self-contained. Half-finished intermediates stay component seams, not public flows. See `contexts/design/operations/orchestration.md`. +10. **ETL execution stays separate from flows** — use + `quantmind.etl.ETLPipeline` for one whole-run delivery and + `BatchETLPipeline` for repeated business-batch deliveries. Call + `ctx.progress()` in long loops so the next Agent can observe real completed + work; do not write custom run-state files. Both are composition-based + authoring scaffolds, not base classes for `quantmind.flows`. See + `contexts/design/operations/etl.md`. ## Tests and Examples diff --git a/contexts/CONTEXT_MAP.md b/contexts/CONTEXT_MAP.md index d68abcc..74ab179 100644 --- a/contexts/CONTEXT_MAP.md +++ b/contexts/CONTEXT_MAP.md @@ -28,6 +28,7 @@ contexts/ │ ├── library/local.md ← LocalKnowledgeLibrary storage and retrieval │ ├── mind/retrieval.md ← page-preserving structure tree + agentic retrieval │ ├── operations/ +│ │ ├── etl.md ← whole-run vs micro-batch ETL execution and observation │ │ ├── naming.md ← public operation naming rules │ │ └── orchestration.md ← pipelines vs components (altitude) │ ├── preprocess/pdf.md ← page-aware ParsedDocument diff --git a/contexts/design/README.md b/contexts/design/README.md index 4d9ecee..34af075 100644 --- a/contexts/design/README.md +++ b/contexts/design/README.md @@ -28,6 +28,7 @@ This directory records QuantMind engineering decisions. Use it to understand whi | Mind | [Build and retrieve from a page-preserving structure tree](mind/retrieval.md) | | Operations | [Public operation naming](operations/naming.md) | | Operations | [Orchestration and construction altitude](operations/orchestration.md) | +| Operations | [Observable whole-run and micro-batch ETL](operations/etl.md) | | Utils | [Cross-provider structured output](utils/structured_output.md) | | Utils | [Collect per-run token and timing usage from SDK traces](utils/usage.md) | diff --git a/contexts/design/operations/etl.md b/contexts/design/operations/etl.md new file mode 100644 index 0000000..c55be37 --- /dev/null +++ b/contexts/design/operations/etl.md @@ -0,0 +1,136 @@ +# Author observable whole-run and micro-batch ETL + +## Quick Summary + +- **Purpose**: Define the two explicit execution shapes in `quantmind.etl`, their delivery semantics, and the local observation contract shared by both. +- **Read when**: Choosing between `ETLPipeline` and `BatchETLPipeline`, changing their execution or run-record behavior, or adopting either scaffold in a consuming repository. +- **Owner**: `quantmind.etl`; concrete domain pipelines remain in consuming repositories or focused examples. +- **Status**: Current. `ETLPipeline` and `BatchETLPipeline` are implemented together in the active ETL change. + +## Contents + +- [Keep two execution shapes explicit](#keep-two-execution-shapes-explicit) +- [Run one delivery through ETLPipeline](#run-one-delivery-through-etlpipeline) +- [Deliver bounded batches through BatchETLPipeline](#deliver-bounded-batches-through-batchetlpipeline) +- [Make dry-run an explicit run property](#make-dry-run-an-explicit-run-property) +- [Observe lifecycle and real progress locally](#observe-lifecycle-and-real-progress-locally) +- [Distinguish staging from delivery](#distinguish-staging-from-delivery) +- [Choose a scaffold by delivery cardinality](#choose-a-scaffold-by-delivery-cardinality) +- [Keep execution policy in the authored pipeline](#keep-execution-policy-in-the-authored-pipeline) +- [Defer concurrent and durable orchestration](#defer-concurrent-and-durable-orchestration) + +## Keep two execution shapes explicit + +`quantmind.etl` is a stdlib-only authoring and observation scaffold, not a workflow engine. It exposes two parallel composition-based classes because delivery cardinality changes execution and observation semantics: + +- `ETLPipeline` performs one whole-run `extract → transform → load` and returns that single load result. +- `BatchETLPipeline` repeatedly pulls one business-defined batch, transforms it, and loads it before pulling the next batch. It returns a bounded summary rather than retaining every load result. + +Neither class inherits the other. A callable returning an async iterable never makes `ETLPipeline` switch modes implicitly: return-value sniffing would misclassify legitimate values and make the run-record contract conditional. Both classes reuse a package-private local record layer because atomic snapshots, sparse JSONL events, and progress coalescing now have two real callers. + +Existing `quantmind.flows` remain pure `input → self-contained artifact` operations and neither inherit nor depend on these ETL scaffolds. Concrete ingestion pipelines belong in their consuming repository; `quantmind.etl` contains only reusable execution and observation behavior. + +## Run one delivery through ETLPipeline + +`ETLPipeline` binds three async callables and invokes each exactly once: + +```text +source → extract → extracted value → transform → final value → load → result +``` + +It is the smallest shape when the operation has one delivery boundary. Each run is one-shot. The pipeline instance stores only its stable name and callables; input, intermediate values, result, stage, and progress remain run-local and are never persisted automatically. + +## Deliver bounded batches through BatchETLPipeline + +`BatchETLPipeline` binds an async batch producer plus async transform and load callables. The framework owns a strictly serial loop: + +```text +pull batch 1 as extract → transform batch 1 → load batch 1 +pull batch 2 as extract → transform batch 2 → load batch 2 +... +``` + +The authored extractor decides batch boundaries and yields them lazily. The framework does not automatically slice inputs. It sets `stage="extract"` while awaiting the next yielded batch, then switches to `transform` and `load` for that same one-based batch index. It never waits for every transformed batch in memory, and v-next does not overlap extraction or transformation of batch N+1 with loading batch N. + +A batch becomes completed only after its load callable returns successfully. A failed or cancelled load may already have produced business side effects; the scaffold cannot prove transactionality. The load implementation must make one batch an idempotent or atomic delivery unit so rerunning the whole run can safely skip or repeat it. The framework does not checkpoint a cursor or resume inside a prior run. + +Each successful load returns optional non-negative count deltas. The runner adds those deltas into a final `BatchRunSummary` and discards the per-batch return value, preserving bounded memory. The summary reports only facts the framework observed: successfully completed batch count and accumulated declared counts. + +An optional known batch total is an assertion, not an estimate. Yielding more batches than declared, or exhausting the producer before the declared total completes, fails the run. Unknown totals remain `null`. + +## Make dry-run an explicit run property + +Both pipeline shapes require callers to choose `dry_run` when they create a run: + +```python +run = pipeline.create_run(source, dry_run=dry_run) +``` + +There is deliberately no `False` default. A production script or CLI must obtain the value from a runtime option and pass that variable so switching modes never requires editing pipeline code; tests, notebooks, and fixed-purpose one-off scripts may pass a Boolean literal. The run-level value is immutable, shared by every stage and batch, and available as the read-only `PipelineContext.dry_run` property. Stage authors pass it through to the repository, gateway, publisher, or other capability that owns the actual mutation decision instead of scattering temporary switches through orchestration code. + +Dry-run executes the complete stage shape. The scaffold never skips `load`: a dry-run load still validates the would-be delivery, computes planned counts, reports progress, and may expose errors visible only at the delivery boundary. With `ctx.dry_run=True`, however, the authored pipeline must prevent every persistent business mutation, including database changes, storage or artifact writes, staging and checkpoint writes, publishing, queueing, webhooks, and final delivery. Reads, fetching, parsing, normalization, pruning, AI processing, previews, planned counts, validation, and QuantMind's own local run-observation files remain allowed. AI processing may still incur cost and rate usage; dry-run is not a zero-cost mode. + +One load signature serves both modes, so its output type must honestly represent both. In dry-run it returns a planned path, planned summary, or `None`, never a reference that implies a nonexistent artifact was delivered. Dry-run progress metrics and batch count deltas use a `planned_*` prefix or an equally explicit planned name; names that assert real delivery, such as `rows_written`, are reserved for normal runs. + +`dry_run` is a formal top-level run field, not configuration metadata. It appears in `run.json` from the initial `created` snapshot onward and in the creation receipt, while the schema IDs remain at v1 because readers already tolerate additive fields. `config_summary["dry_run"]` is rejected with guidance to use `create_run(..., dry_run=...)`, preventing two competing sources of truth. An observer reads the top-level flag before interpreting the terminal state: a dry-run `succeeded` means the plan completed validation, not that data was delivered. + +This flag establishes a narrow admission rule for future execution-semantic parameters. A run-level flag enters `create_run()` only when it both changes external side-effect semantics and is necessary for an observer to interpret the run correctly. `dry_run` satisfies both conditions; tuning such as `verbose`, `fast_mode`, or `force` remains authored-pipeline configuration rather than expanding the framework signature. + +## Observe lifecycle and real progress locally + +`create_run()` allocates a Run ID and atomically writes `state="created"` before returning. `receipt()` provides that Run ID and the absolute `run.json` path as one JSON line, which a wrapper prints and flushes before `execute()`. + +Both shapes keep an atomic latest snapshot in `run.json` and a sparse lifecycle journal in `events.jsonl`. They use distinct schema IDs so an observer never has to infer execution shape from values: + +- whole-run snapshots: `quantmind.etl.run/v1`; +- whole-run events: `quantmind.etl.event/v1`; +- micro-batch snapshots: `quantmind.etl.batch-run/v1`, with `batch.index`, successfully completed load count `batch.completed`, and optional `batch.total`; +- micro-batch events: `quantmind.etl.batch-event/v1`. + +`run.json` is the authoritative state: after a terminal snapshot commits, the terminal JSONL append is best-effort and cannot change the `execute()` outcome, while failure and cancellation observation errors preserve the original exception. + +In a normal micro-batch run, completed loads represent delivery; in dry-run they represent validated planned deliveries. + +Micro-batch events record run lifecycle, explicitly reported stage progress, and one `batch_completed` event per meaningful delivery unit or validated planned delivery in dry-run. They do not emit stage-start/stage-complete pairs for every batch. Batch size is expected to represent tens or hundreds of meaningful delivery units; producing thousands of tiny batches is an authored-pipeline configuration problem rather than a reason to turn the sparse journal into arbitrary logging. + +`PipelineContext.progress()` always means real finished work. Its strict monotonicity and known-total rules apply within one active stage for `ETLPipeline`, and within one active `(batch index, stage)` for `BatchETLPipeline`. Every stage or batch switch clears the previous progress snapshot before user code runs. In batch mode, `ctx.batch_index` identifies the current one-based batch; it is `None` for whole-run ETL. + +Child tasks created inside the active stage inherit its progress scope and may report real completion, subject to the same monotonic counter. A task that outlives that stage or batch keeps the old scope and its later progress call is rejected, so a stale producer cannot write into the next batch's snapshot. + +Inputs, complete configuration, intermediate values, individual load results, tracebacks, headers, response bodies, and arbitrary logs are not written. Only the caller's explicit JSON-scalar `config_summary`, framework lifecycle fields, limited error type/message, and explicit progress/count measurements appear. + +## Distinguish staging from delivery + +`load` marks the delivery boundary. In a normal run, after it succeeds the pipeline's intended downstream consumer may treat that output as delivered; in dry-run it validates or plans the same boundary without delivering. A database call is not automatically a load. Extract or transform may perform staging writes when they persist an intermediate for recovery or reuse while keeping it unavailable to formal downstream consumers. + +Staging writes are permitted in any stage of a normal run when they are idempotent and their real completion is visible through progress. Dry-run must plan or validate them without writing. They remain business behavior: the scaffold provides no transaction, rollback, exactly-once, or artifact-management guarantee. + +If a batch write makes the final product consumable, it is a real batch load even when the target table is named `raw`. Repeated `transform(batch) → load(batch)` delivery belongs in `BatchETLPipeline`; hiding those loads inside a whole-run transform would make `run.json.stage` dishonest. + +## Choose a scaffold by delivery cardinality + +Use this selection test: + +| Delivery shape | Scaffold | +|---|---| +| One final delivery after all processing | `ETLPipeline` | +| Repeated bounded deliveries, one per business batch | `BatchETLPipeline` | +| Intermediate durable writes that are not consumable products | Keep them as idempotent staging inside the owning stage | +| An intermediate is itself an independently consumed product | Split into two pipelines and pass run A's artifact to run B | +| Named steps need independent retry/resume or form a graph | Neither current scaffold; evaluate a later ETL DAG only after that requirement is concrete | + +A write followed by another decision does not alone imply a workflow. Ask whether the write is staging or delivery, and how many delivery units the run intentionally publishes. + +## Keep execution policy in the authored pipeline + +The scaffold records what runs; it does not select concurrency, timeout, retry, transaction, or idempotency policy. Stage-local timeouts fail the stage when surfaced as `TimeoutError`; cancelling `execute()` from an outer deadline records `cancelled` on a best-effort basis and re-raises `CancelledError`. Retry remains explicit and must respect side-effect safety. + +Batch extractor cleanup is intentionally not a hidden bounded policy. If failure or cancellation happens after a batch has already been yielded, the runner does not proactively start the extractor's async `aclose()` cleanup: once such an awaitable is running on the current event loop, asyncio cannot force it to stop if it swallows cancellation, and leaving it detached can block `asyncio.run()` shutdown. If cancellation lands while awaiting the next batch from a normal async generator, Python's own generator cancellation semantics still run cleanup for that active await. Cleanup after yielded batches remains an authored-pipeline responsibility. + +Cancellation stops at the active await. A batch runner neither drains queued work nor rolls back prior completed batches. `batch.completed` counts only load calls that returned successfully, so an observer can distinguish confirmed completed batches from the current uncertain batch. Recovery is a new run plus business idempotency or ready-skip logic. + +## Defer concurrent and durable orchestration + +The current package deliberately excludes cross-batch pipelining, multiple simultaneously active stages, automatic checkpoint/resume, retry policy, heartbeat, ETA, retention, scheduling, queues, database ledgers, dashboards, arbitrary logging, artifact management, DAGs, streaming records, and exactly-once claims. + +Cross-batch pipelining may improve throughput, but it would make a single `stage` field false because transform and load could be active simultaneously. It requires a separate observation design rather than a small optimization to `BatchETLPipeline`. A future ETL DAG remains a parallel class inside `quantmind.etl`, sharing only proven package-private recording mechanics and never inheriting either current pipeline class. diff --git a/contexts/usage/README.md b/contexts/usage/README.md index 5fd7e8b..d3f47ba 100644 --- a/contexts/usage/README.md +++ b/contexts/usage/README.md @@ -5,7 +5,7 @@ - **Purpose**: Route library users and agents to current public operations, inputs, results, examples, and guides. - **Read when**: Calling QuantMind as a library or selecting a supported public operation. - **Load next**: Start with the component row that matches the requested operation; do not load unrelated component designs. -- **Import rule**: Import inputs and configs from `quantmind.configs`, operations from `quantmind.flows`, and result types from the package shown in the component catalog. +- **Import rule**: Import public callables from the owning package shown in the component catalog; ETL scaffolds come from `quantmind.etl`, flow operations from `quantmind.flows`, inputs and configs from `quantmind.configs`, and result types from their canonical component package. ## Contents @@ -23,9 +23,10 @@ Use this index when calling QuantMind as a library. These links point to the cur | Source-first paper flow | [Paper flow design](../design/flow/paper.md) | | News collection | [News design and behavior](../design/flow/news.md) | | Search local knowledge by meaning | [Library guide](../../docs/library.md) and [focused example](../../examples/library/README.md) | +| Observable ETL scaffolds | [ETL guide](../../docs/etl.md) and [focused examples](../../examples/etl/) | | Runnable operation examples | [`examples/flows/`](../../examples/flows/) | | Focused preprocessing examples | [`examples/preprocess/`](../../examples/preprocess/) | ## Where to Import From -Import public inputs and configs from `quantmind.configs`, public operations from `quantmind.flows`, and result types from the package shown in the component catalog. +Use the public component catalog as the import authority for each capability. Observable ETL scaffolds are exported from `quantmind.etl`, public flow operations and builders from `quantmind.flows`, public inputs and configs from `quantmind.configs`, cognitive services from `quantmind.mind`, and result contracts from the canonical layer named in the catalog. diff --git a/docs/README.md b/docs/README.md index 862e00f..eef2c01 100644 --- a/docs/README.md +++ b/docs/README.md @@ -14,6 +14,8 @@ harness. | Operation | Import | Input and config | Result | Example | Design or guide | |---|---|---|---|---|---| +| Observable whole-run ETL | `quantmind.etl.ETLPipeline` | Three async stage callables; `create_run(source, dry_run=..., config_summary=...)` | Single load result plus local `run.json` / `events.jsonl` with top-level `dry_run` | [Write a local artifact](../examples/etl/local_artifact.py) | [Observable ETL guide](etl.md) | +| Observable micro-batch ETL | `quantmind.etl.BatchETLPipeline` | Async batch producer plus transform/load; `create_run(source, dry_run=..., total_batches=...)` | `BatchRunSummary` plus local batch run records with top-level `dry_run` | [Write batch local artifacts](../examples/etl/batch_local_artifacts.py) | [Observable ETL guide](etl.md) | | Source-first paper flow | `quantmind.flows.PaperFlow` | `PaperFlow(PaperSemanticCfg)`; `build()`: `PaperInput` | `PaperSemanticResult` | [Persist and search a paper](../examples/flows/paper.py) | [Paper flow design](../contexts/design/flow/paper.md) | | Paper structure build | `quantmind.flows.PaperFlow` | `PaperFlow(PaperStructureCfg)`; `build()`: `PaperInput` | `PaperStructureTree` (self-contained) | [Build and retrieve](../examples/mind/paper_structure_retrieval.py) | [Structure retrieval design](../contexts/design/mind/retrieval.md) | | Reasoning-based retrieval (agentic) | `quantmind.mind.AgenticRetriever` | `AgenticRetriever(RetrievalCfg)`; `retrieve()`: one `StructureTree` + question (no library) | `list[RetrievalEvidence]` | [Build and retrieve](../examples/mind/paper_structure_retrieval.py) | [Structure retrieval design](../contexts/design/mind/retrieval.md) | @@ -22,9 +24,10 @@ harness. | Local semantic search | `quantmind.library.LocalKnowledgeLibrary` | `BaseKnowledge` or `PaperSemanticResult`, `SemanticQuery` | `list[SemanticHit]` | [Library example](../examples/library/README.md) | [Library guide](library.md) | | Page-aware document RAG | `quantmind.rag.chunk_parsed_document`, `quantmind.rag.retrieve_parsed_document` | `ParsedDocument`, splitter config, and query | `tuple[ParsedDocumentHit, ...]` | [Paper RAG](../examples/rag/paper.py) | [Document RAG design](../contexts/design/rag/document.md) | -Import public inputs and configs from `quantmind.configs`, flow operations and -builders from `quantmind.flows`, and cognitive services from `quantmind.mind`. -Import result contracts from the canonical layer shown in the catalog. +Import observable ETL authoring primitives from `quantmind.etl`, public inputs +and configs from `quantmind.configs`, flow operations and builders from +`quantmind.flows`, and cognitive services from `quantmind.mind`. Import result +contracts from the canonical layer shown in the catalog. ## Public-Network Sources diff --git a/docs/etl.md b/docs/etl.md new file mode 100644 index 0000000..7ae9901 --- /dev/null +++ b/docs/etl.md @@ -0,0 +1,144 @@ +# Observable ETL authoring + +`quantmind.etl` provides two small async authoring scaffolds with local machine-readable run state: + +- `ETLPipeline` performs one whole-run `extract → transform → load` and returns the single load result. +- `BatchETLPipeline` lazily pulls one business batch, transforms it, and loads it before pulling the next batch. It returns an aggregate summary without retaining every batch. + +Both require callers to choose `dry_run` when creating a run. Dry-run executes the same stages, including `load`, but authored mutation boundaries must plan or validate without persistent business writes. Concrete domain pipelines belong in the consuming repository; `quantmind.etl` contains only the reusable scaffold. See the [ETL design](../contexts/design/operations/etl.md) for delivery, staging, dry-run, recovery, and selection rules. + +## Run one whole-run delivery + +Bind three async callables when all processing leads to one final delivery: + +```python +from quantmind.etl import ETLPipeline, PipelineContext + + +async def extract(source: str, *, ctx: PipelineContext) -> list[str]: + rows = source.splitlines() + await ctx.progress(len(rows), total=len(rows)) + return rows + + +async def transform( + rows: list[str], *, ctx: PipelineContext +) -> list[str]: + return [row.strip() for row in rows] + + +async def load(rows: list[str], *, ctx: PipelineContext) -> dict[str, int]: + if ctx.dry_run: + return {"planned_rows": len(rows)} + return {"rows_written": len(rows)} + + +pipeline = ETLPipeline( + "line-loader", + extract=extract, + transform=transform, + load=load, +) +dry_run = False +run = pipeline.create_run( + "one\ntwo", + dry_run=dry_run, + config_summary={"drop_empty": False}, +) +print(run.receipt(), flush=True) +result = await run.execute() +``` + +See the runnable [local artifact example](../examples/etl/local_artifact.py) for the smallest complete whole-run path. + +## Deliver bounded batches + +Use `BatchETLPipeline` when each business-defined batch should be transformed and delivered before the next batch is pulled: + +```python +from collections.abc import AsyncIterator, Mapping + +from quantmind.etl import BatchETLPipeline, PipelineContext + + +async def extract( + rows: list[str], *, ctx: PipelineContext +) -> AsyncIterator[list[str]]: + batch_size = 2 + for start in range(0, len(rows), batch_size): + batch = rows[start : start + batch_size] + await ctx.progress(len(batch), total=len(batch)) + yield batch + + +async def transform( + batch: list[str], *, ctx: PipelineContext +) -> list[str]: + return [row.strip().upper() for row in batch] + + +async def load( + batch: list[str], *, ctx: PipelineContext +) -> Mapping[str, int]: + if ctx.dry_run: + return {"planned_records": len(batch)} + # Persist this batch idempotently, then report count deltas. + return {"records_written": len(batch)} + + +pipeline = BatchETLPipeline( + "batched-line-loader", + extract=extract, + transform=transform, + load=load, +) +dry_run = False +run = pipeline.create_run( + ["one", "two", "three"], + dry_run=dry_run, + total_batches=2, + config_summary={"batch_size": 2}, +) +print(run.receipt(), flush=True) +summary = await run.execute() +``` + +The framework owns a strictly serial loop. While it awaits the next yielded batch, `run.json.stage` is `extract`; that batch then moves through `transform` and `load` under one one-based `batch.index`. Only a load call that returns successfully increments `batch.completed`. The next batch is not pulled early, so there is no cross-batch pipelining or hidden buffer. + +`load` may return `None` or JSON-style non-negative integer count deltas. The runner sums those deltas into `BatchRunSummary.counts` and discards each per-batch result. `total_batches` is optional; when supplied, it is an assertion that fails if extraction yields too many or too few batches. + +See [batch local artifacts](../examples/etl/batch_local_artifacts.py) for a network-free, bounded-memory example with idempotent local loads. + +## Implement dry-run honestly + +`create_run(source, *, dry_run=...)` has no default. Production scripts should read the value from a runtime option such as `--dry-run` and pass that variable, so switching modes never requires editing stage code. + +Every stage receives the same read-only `ctx.dry_run` value. Pass it to the repository, gateway, publisher, or storage adapter that owns the mutation decision. The scaffold does not skip `load`; a dry-run load validates the would-be delivery, reports `planned_*` progress or batch counts, and returns a planned path, planned summary, or `None` rather than a reference to a nonexistent artifact. Dry-run forbids persistent business mutation, including staging/checkpoint writes, but still allows reads, parsing, validation, previews, planned counts, and QuantMind's local run-observation files. AI processing may still run and incur cost or rate usage. + +## Read local run state + +`create_run()` atomically writes `state="created"` before it returns. Snapshots, sparse events, and the receipt include top-level `dry_run`; do not duplicate it inside `config_summary`. The receipt is one JSON line containing the Run ID, `dry_run`, and the absolute `status_file`. By default, each run lives under: + +```text +/.quant-mind/etl-pipeline-runs//run.json +/.quant-mind/etl-pipeline-runs//events.jsonl +``` + +The four canonical schema IDs are: + +- whole-run snapshots: `quantmind.etl.run/v1`; +- whole-run events: `quantmind.etl.event/v1`; +- micro-batch snapshots: `quantmind.etl.batch-run/v1`; +- micro-batch events: `quantmind.etl.batch-event/v1`. + +Read `run.json` for the latest snapshot. Check its top-level `dry_run` value before interpreting the result: dry-run `succeeded` means the planned delivery passed validation, not that data was delivered. Its states are `created`, `running`, `succeeded`, `failed`, and `cancelled`. It includes the current macro stage, the executing process's PID, timestamps, safe config summary, latest progress, and a limited error type/message. A batch run additionally includes its current batch index, successfully completed batch count, and optional total. The PID is only a local process-liveness hint; it is not a heartbeat and does not prove forward progress or success. + +Call `await ctx.progress(...)` inside meaningful long loops. `completed` means work that really finished and must strictly increase within the current stage, or within the current `(batch, stage)` for batch ETL. Use `total=None` until the total is known. Starting another stage or batch clears the previous progress before user code runs. `ctx.batch_index` is the active one-based index for batch ETL and `None` for whole-run ETL. + +Async child tasks created inside the active stage may report progress with the same monotonic counter. Do not leave progress-reporting tasks running after the stage returns: an inherited scope from a completed stage or prior batch is rejected. + +Every progress call updates the atomic snapshot. `events.jsonl` keeps only sparse lifecycle events and coalesces progress events emitted within one second; it is not a general logging API. Batch runs emit one `batch_completed` event per batch whose `load` call returned successfully rather than stage-start/stage-complete pairs for every cycle. In a normal run that means the batch was delivered; in a dry-run it means the planned delivery validated. + +Only the explicit JSON-scalar `config_summary` allowlist is persisted. Inputs, full configs, intermediate values, individual load results, final results, tracebacks, locals, headers, and response bodies are not serialized automatically. + +Every run handle is one-shot. A stage exception records `failed` and is re-raised. Cancellation records `cancelled` on a best-effort basis and then re-raises `asyncio.CancelledError`. Prior successful batch loads are not rolled back. When a batch run fails or is cancelled after a batch has been yielded, the runner does not start the extractor's async `aclose()` cleanup, because a coroutine that ignores cancellation cannot be forcibly bounded on the same event loop and would block `asyncio.run()` shutdown. If cancellation lands while the runner is awaiting the next batch from a normal async generator, Python's own generator cancellation semantics still run that active `finally` cleanup. Concurrency, timeout, retry, extractor cleanup after yielded batches, batch transactionality, and sink idempotency remain explicit responsibilities of the authored pipeline. diff --git a/examples/etl/batch_local_artifacts.py b/examples/etl/batch_local_artifacts.py new file mode 100644 index 0000000..9cd1b4e --- /dev/null +++ b/examples/etl/batch_local_artifacts.py @@ -0,0 +1,132 @@ +"""Run a tiny observable micro-batch ETL that writes local artifacts.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +from pathlib import Path + +from quantmind.etl import BatchETLPipeline, PipelineContext + +BATCH_SIZE = 2 + + +async def extract(source: list[str], *, ctx: PipelineContext): + """Yield small in-memory batches and report each batch before yielding.""" + for offset in range(0, len(source), BATCH_SIZE): + batch = source[offset : offset + BATCH_SIZE] + await ctx.progress( + len(batch), + total=len(batch), + message=f"Prepared rows {offset + 1}-{offset + len(batch)}", + ) + yield batch + + +async def transform( + batch: list[str], *, ctx: PipelineContext +) -> list[dict[str, object]]: + """Build deterministic records for one batch.""" + await ctx.progress(1, total=1, message="Normalized batch") + return [ + { + "text": line.strip(), + "characters": len(line.strip()), + } + for line in batch + if line.strip() + ] + + +async def run_example(*, dry_run: bool) -> dict[str, object]: + """Create the run, print its receipt, then execute all batches.""" + artifact_dir = Path.cwd() / ".quant-mind" / "etl-batch-example" + + async def load( + records: list[dict[str, object]], *, ctx: PipelineContext + ) -> dict[str, int]: + batch_index = ctx.batch_index + if batch_index is None: + raise RuntimeError("batch load requires a batch index") + target = artifact_dir / f"batch-{batch_index:03d}.json" + content = json.dumps(records, indent=2, sort_keys=True) + "\n" + byte_count = len(content.encode("utf-8")) + if ctx.dry_run: + await ctx.progress( + 1, + total=1, + message="Planned batch artifact", + metrics={ + "planned_artifacts": 1, + "planned_bytes": byte_count, + "planned_records": len(records), + }, + ) + return { + "planned_artifacts": 1, + "planned_records": len(records), + } + + artifact_dir.mkdir(parents=True, exist_ok=True) + temporary = target.with_suffix(".json.tmp") + await asyncio.to_thread(temporary.write_text, content, encoding="utf-8") + temporary.replace(target) + await ctx.progress( + 1, + total=1, + message="Wrote batch artifact", + metrics={ + "artifacts_written": 1, + "bytes_written": byte_count, + "records_written": len(records), + }, + ) + return {"artifacts_written": 1, "records_written": len(records)} + + source = ["alpha", " beta ", "", "gamma"] + pipeline = BatchETLPipeline( + "batch-local-artifacts", + extract=extract, + transform=transform, + load=load, + ) + run = pipeline.create_run( + source, + dry_run=dry_run, + config_summary={ + "artifact_format": "json", + "batch_size": BATCH_SIZE, + }, + total_batches=math.ceil(len(source) / BATCH_SIZE), + ) + print(run.receipt(), flush=True) + summary = await run.execute() + return { + "dry_run": dry_run, + "completed_batches": summary.completed_batches, + "counts": dict(summary.counts), + } + + +def _parse_args() -> argparse.Namespace: + """Parse the runtime run-mode option.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--dry-run", + action="store_true", + help="plan and validate without creating example batch artifacts", + ) + return parser.parse_args() + + +async def main() -> None: + """Run the example in normal or dry-run mode.""" + args = _parse_args() + summary = await run_example(dry_run=bool(args.dry_run)) + print("summary=" + json.dumps(summary, sort_keys=True)) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/etl/local_artifact.py b/examples/etl/local_artifact.py new file mode 100644 index 0000000..3058861 --- /dev/null +++ b/examples/etl/local_artifact.py @@ -0,0 +1,106 @@ +"""Run a tiny observable ETL that writes one idempotent local artifact.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +from pathlib import Path + +from quantmind.etl import ETLPipeline, PipelineContext + +LoadResult = dict[str, object] + + +async def extract(source: Path, *, ctx: PipelineContext) -> str: + """Read one local text source.""" + text = await asyncio.to_thread(source.read_text, encoding="utf-8") + await ctx.progress(1, total=1, message="Read local source") + return text + + +async def transform(text: str, *, ctx: PipelineContext) -> dict[str, object]: + """Build a deterministic summary.""" + lines = [line for line in text.splitlines() if line.strip()] + return {"non_empty_lines": len(lines), "characters": len(text)} + + +async def run_example(*, dry_run: bool) -> LoadResult: + """Create the run, print its receipt, then execute all stages.""" + artifact = Path.cwd() / ".quant-mind" / "etl-example" / "artifact.json" + + async def load( + summary: dict[str, object], *, ctx: PipelineContext + ) -> LoadResult: + content = json.dumps(summary, indent=2, sort_keys=True) + "\n" + byte_count = len(content.encode("utf-8")) + if ctx.dry_run: + await ctx.progress( + 1, + total=1, + message="Planned local artifact", + metrics={ + "planned_artifacts": 1, + "planned_bytes": byte_count, + }, + ) + return { + "dry_run": True, + "planned_artifact": str(artifact), + "planned_artifacts": 1, + "planned_bytes": byte_count, + } + + artifact.parent.mkdir(parents=True, exist_ok=True) + await asyncio.to_thread(artifact.write_text, content, encoding="utf-8") + await ctx.progress( + 1, + total=1, + message="Wrote local artifact", + metrics={ + "artifacts_written": 1, + "bytes_written": byte_count, + }, + ) + return { + "dry_run": False, + "artifact": str(artifact), + "artifacts_written": 1, + "bytes_written": byte_count, + } + + pipeline = ETLPipeline( + "local-source-summary", + extract=extract, + transform=transform, + load=load, + ) + run = pipeline.create_run( + Path(__file__), + dry_run=dry_run, + config_summary={"artifact_format": "json"}, + ) + print(run.receipt(), flush=True) + return await run.execute() + + +def _parse_args() -> argparse.Namespace: + """Parse the runtime run-mode option.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--dry-run", + action="store_true", + help="plan and validate without creating the example artifact", + ) + return parser.parse_args() + + +async def main() -> None: + """Run the example in normal or dry-run mode.""" + args = _parse_args() + result = await run_example(dry_run=bool(args.dry_run)) + print("result=" + json.dumps(result, sort_keys=True)) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/pyproject.toml b/pyproject.toml index 26e9c30..688e41f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,20 +120,38 @@ reportIncompatibleVariableOverride = "none" # ---------------------------------------------------------------------------- # import-linter: architectural boundary contracts # ---------------------------------------------------------------------------- -# Encodes the target architecture: utils, knowledge, and preprocess are -# leaves; configs and library depend only on knowledge; rag depends only on -# preprocess; flows + magic is the apex layer. The transitional packages -# (config/, flow/, llm/, models/) remain forbidden as a tripwire. +# Encodes the target architecture: etl is an independent stdlib-only leaf; +# utils, knowledge, and preprocess are leaves; configs and library depend only +# on knowledge; rag depends only on preprocess; flows + magic is the apex +# layer. The transitional packages (config/, flow/, llm/, models/) remain +# forbidden as a tripwire. [tool.importlinter] root_packages = ["quantmind"] +[[tool.importlinter.contracts]] +name = "etl is independent (stdlib only)" +type = "forbidden" +source_modules = ["quantmind.etl"] +forbidden_modules = [ + "quantmind.configs", + "quantmind.flows", + "quantmind.knowledge", + "quantmind.library", + "quantmind.magic", + "quantmind.mind", + "quantmind.preprocess", + "quantmind.rag", + "quantmind.utils", +] + [[tool.importlinter.contracts]] name = "utils is a leaf (no inbound deps from quantmind packages)" type = "forbidden" source_modules = ["quantmind.utils"] forbidden_modules = [ "quantmind.configs", + "quantmind.etl", "quantmind.flows", "quantmind.knowledge", "quantmind.library", @@ -149,6 +167,7 @@ type = "forbidden" source_modules = ["quantmind.knowledge"] forbidden_modules = [ "quantmind.configs", + "quantmind.etl", "quantmind.flows", "quantmind.library", "quantmind.magic", @@ -163,6 +182,7 @@ name = "configs only depends on knowledge" type = "forbidden" source_modules = ["quantmind.configs"] forbidden_modules = [ + "quantmind.etl", "quantmind.flows", "quantmind.library", "quantmind.magic", @@ -177,6 +197,7 @@ type = "forbidden" source_modules = ["quantmind.preprocess"] forbidden_modules = [ "quantmind.configs", + "quantmind.etl", "quantmind.flows", "quantmind.knowledge", "quantmind.library", @@ -192,6 +213,7 @@ source_modules = ["quantmind.library"] forbidden_modules = [ "quantmind.config", "quantmind.configs", + "quantmind.etl", "quantmind.flow", "quantmind.flows", "quantmind.llm", @@ -210,6 +232,7 @@ source_modules = ["quantmind.rag"] forbidden_modules = [ "quantmind.config", "quantmind.configs", + "quantmind.etl", "quantmind.flow", "quantmind.flows", "quantmind.knowledge", @@ -227,6 +250,7 @@ type = "forbidden" source_modules = ["quantmind.mind"] forbidden_modules = [ "quantmind.config", + "quantmind.etl", "quantmind.flow", "quantmind.flows", "quantmind.llm", @@ -247,6 +271,7 @@ source_modules = [ # future code re-introducing them under the same names. forbidden_modules = [ "quantmind.config", + "quantmind.etl", "quantmind.flow", "quantmind.llm", "quantmind.models", diff --git a/quantmind/etl/__init__.py b/quantmind/etl/__init__.py new file mode 100644 index 0000000..ce86070 --- /dev/null +++ b/quantmind/etl/__init__.py @@ -0,0 +1,19 @@ +"""Observable whole-run and micro-batch ETL authoring primitives.""" + +from quantmind.etl._batch import ( + BatchETLPipeline, + BatchPipelineRun, + BatchRunSummary, +) +from quantmind.etl._pipeline import ETLPipeline, PipelineRun +from quantmind.etl._record import JsonScalar, PipelineContext + +__all__ = [ + "BatchETLPipeline", + "BatchPipelineRun", + "BatchRunSummary", + "ETLPipeline", + "JsonScalar", + "PipelineContext", + "PipelineRun", +] diff --git a/quantmind/etl/_batch.py b/quantmind/etl/_batch.py new file mode 100644 index 0000000..55770c4 --- /dev/null +++ b/quantmind/etl/_batch.py @@ -0,0 +1,453 @@ +"""Micro-batch observable ``extract -> transform -> load`` execution.""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncIterable, Awaitable, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Any, Generic, Protocol, TypeAlias, TypeVar, cast + +from quantmind.etl._record import ( + JsonScalar, + PipelineContext, + RunRecord, + copy_scalar_mapping, + validate_dry_run, +) + +InputT = TypeVar("InputT") +ExtractedBatchT = TypeVar("ExtractedBatchT") +TransformedBatchT = TypeVar("TransformedBatchT") +_StageInputT = TypeVar("_StageInputT", contravariant=True) +_StageOutputT = TypeVar("_StageOutputT", covariant=True) +_BatchExtractInputT = TypeVar("_BatchExtractInputT", contravariant=True) +_BatchExtractOutputT = TypeVar("_BatchExtractOutputT", covariant=True) +_BatchLoadInputT = TypeVar("_BatchLoadInputT", contravariant=True) + +CountDeltas: TypeAlias = Mapping[str, int] | None + + +class _BatchExtractCallable( + Protocol[_BatchExtractInputT, _BatchExtractOutputT] +): + def __call__( + self, + value: _BatchExtractInputT, + /, + *, + ctx: PipelineContext, + ) -> AsyncIterable[_BatchExtractOutputT]: ... + + +class _StageCallable(Protocol[_StageInputT, _StageOutputT]): + def __call__( + self, + value: _StageInputT, + /, + *, + ctx: PipelineContext, + ) -> Awaitable[_StageOutputT]: ... + + +class _BatchLoadCallable(Protocol[_BatchLoadInputT]): + def __call__( + self, + value: _BatchLoadInputT, + /, + *, + ctx: PipelineContext, + ) -> Awaitable[CountDeltas]: ... + + +@dataclass(slots=True) +class _BatchSnapshotState: + index: int | None + completed: int + total: int | None + + def as_json(self) -> dict[str, object]: + return { + "index": self.index, + "completed": self.completed, + "total": self.total, + } + + +@dataclass(frozen=True, slots=True) +class BatchRunSummary: + """Bounded-memory summary returned by a successful batch ETL run.""" + + completed_batches: int + counts: Mapping[str, int] + + def __post_init__(self) -> None: + if type(self.completed_batches) is not int: + raise TypeError("completed_batches must be an integer") + if self.completed_batches < 0: + raise ValueError("completed_batches must be >= 0") + object.__setattr__( + self, + "counts", + MappingProxyType(_copy_count_mapping(self.counts, "counts")), + ) + + +class BatchETLPipeline(Generic[InputT, ExtractedBatchT, TransformedBatchT]): + """Bind async callables into a strict serial micro-batch ETL operation. + + Every authored batch ETL should honor the required ``dry_run`` run option: + callers pass ``create_run(..., dry_run=dry_run)`` once, and each extract, + transform, and load scope reads the same immutable ``ctx.dry_run`` value. + Dry-run still pulls, transforms, and loads every batch serially; the + authored ``load`` must validate or plan without persistent business writes, + including staging/checkpoint writes, and return ``planned_*`` count deltas + instead of delivered counts. AI processing may still run and incur cost or + rate usage. + """ + + def __init__( + self, + name: str, + *, + extract: _BatchExtractCallable[InputT, ExtractedBatchT], + transform: _StageCallable[ExtractedBatchT, TransformedBatchT], + load: _BatchLoadCallable[TransformedBatchT], + ) -> None: + if not isinstance(name, str) or not name.strip(): + raise ValueError("pipeline name must be a non-empty string") + if ( + not callable(extract) + or not callable(transform) + or not callable(load) + ): + raise TypeError("extract, transform, and load must be callable") + self._name = name + self._extract = extract + self._transform = transform + self._load = load + + @property + def name(self) -> str: + """The stable name written to each local run record.""" + return self._name + + def create_run( + self, + source: InputT, + *, + dry_run: bool, + config_summary: Mapping[str, JsonScalar] | None = None, + run_root: Path | None = None, + total_batches: int | None = None, + ) -> BatchPipelineRun[ExtractedBatchT, TransformedBatchT]: + """Create a one-shot micro-batch run. + + ``dry_run`` is required because it changes side-effect semantics and is + written as a top-level field from the initial ``created`` snapshot. + Do not put ``dry_run`` in ``config_summary``. Dry-run runs the complete + batch stage shape, including ``load`` for every yielded batch; stage + authors read ``ctx.dry_run`` and forward it to the dependency that owns + mutation decisions. In dry-run, those dependencies must not create + persistent business mutations, including staging/checkpoint writes or + final delivery, and count deltas should use ``planned_*`` names. AI + processing is allowed and may incur cost. + + Args: + source: Runtime input retained in memory until execution. + dry_run: Explicit run mode. ``True`` plans and validates without + persistent business delivery; ``False`` allows authored + mutation boundaries to commit. + config_summary: Explicit allowlist of safe JSON-scalar settings. + run_root: Directory that will contain the run directory. The + default is ``/.quant-mind/etl-pipeline-runs``. + total_batches: Optional known batch count. Unknown is recorded as + ``null`` until the extractor is exhausted. + + Returns: + A one-shot batch run whose ``status_file`` already exists in + ``created`` state. + """ + root = ( + Path.cwd() / ".quant-mind" / "etl-pipeline-runs" + if run_root is None + else Path(run_root) + ).resolve() + return BatchPipelineRun._create( + pipeline_name=self._name, + source=source, + extract=self._extract, + transform=self._transform, + load=self._load, + dry_run=validate_dry_run(dry_run), + config_summary=copy_scalar_mapping( + config_summary, field_name="config_summary" + ), + run_root=root, + total_batches=_validate_total_batches(total_batches), + ) + + +class BatchPipelineRun(Generic[ExtractedBatchT, TransformedBatchT]): + """One-shot execution handle for a serial micro-batch ETL run.""" + + def __init__( + self, + *, + record: RunRecord, + batch: _BatchSnapshotState, + source: object, + extract: _BatchExtractCallable[Any, ExtractedBatchT], + transform: _StageCallable[ExtractedBatchT, TransformedBatchT], + load: _BatchLoadCallable[TransformedBatchT], + ) -> None: + self._record = record + self._batch = batch + self._source = source + self._extract = extract + self._transform = transform + self._load = load + self._counts: dict[str, int] = {} + self._executed = False + + @classmethod + def _create( + cls, + *, + pipeline_name: str, + source: object, + extract: _BatchExtractCallable[Any, ExtractedBatchT], + transform: _StageCallable[ExtractedBatchT, TransformedBatchT], + load: _BatchLoadCallable[TransformedBatchT], + dry_run: bool, + config_summary: dict[str, JsonScalar], + run_root: Path, + total_batches: int | None, + ) -> BatchPipelineRun[ExtractedBatchT, TransformedBatchT]: + batch = _BatchSnapshotState( + index=None, completed=0, total=total_batches + ) + record = RunRecord.create( + pipeline_name=pipeline_name, + dry_run=dry_run, + config_summary=config_summary, + run_root=run_root, + snapshot_schema="quantmind.etl.batch-run/v1", + event_schema="quantmind.etl.batch-event/v1", + snapshot_extra=lambda: {"batch": batch.as_json()}, + ) + return cls( + record=record, + batch=batch, + source=source, + extract=extract, + transform=transform, + load=load, + ) + + @property + def id(self) -> str: + """The unique local run ID.""" + return self._record.id + + @property + def dry_run(self) -> bool: + """Whether this run plans and validates without business delivery.""" + return self._record.dry_run + + @property + def status_file(self) -> Path: + """Absolute path to the latest atomic ``run.json`` snapshot.""" + return self._record.status_file + + @property + def events_file(self) -> Path: + """Absolute path to the sparse lifecycle ``events.jsonl`` record.""" + return self._record.events_file + + def receipt(self) -> str: + """Return the run ID and absolute status path as one JSON line.""" + return json.dumps( + { + "event": "etl_batch_run_created", + "run_id": self._record.id, + "dry_run": self._record.dry_run, + "status_file": str(self._record.status_file), + }, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + ) + + async def execute(self) -> BatchRunSummary: + """Run every extracted batch through ``transform`` then ``load``. + + The executor is strictly serial: it never starts batch ``N+1`` before + batch ``N`` has loaded successfully. Cancellation and stage exceptions + are recorded and the original exception is re-raised. After a yielded + batch fails or is cancelled, user extractor ``aclose()`` is not started + because an async cleanup coroutine cannot be forcibly bounded on the + active event loop. + ``dry_run=True`` does not skip ``load``; the authored load uses + ``ctx.dry_run`` to validate or plan without committing, and reports + planned count deltas honestly. + """ + if self._executed: + raise RuntimeError( + "BatchPipelineRun.execute() can only be called once" + ) + self._executed = True + self._record.start_running() + try: + self._record.append_event("run_started") + except Exception: + pass + + try: + extract_context = self._record.new_inactive_context(stage="extract") + batches = self._extract(self._source, ctx=extract_context) + try: + iterator = batches.__aiter__() + except AttributeError as exc: + raise TypeError("extract must return an AsyncIterable") from exc + + await self._execute_batches(iterator, extract_context) + except asyncio.CancelledError: + try: + self._record.finish("cancelled", "run_cancelled") + except Exception: + pass + raise + except Exception as exc: + self._record.set_error(exc) + try: + self._record.finish("failed", "run_failed") + except Exception: + pass + raise + + self._batch.index = None + self._record.stage = "load" if self._batch.completed else None + self._record.clear_progress() + self._record.finish("succeeded", "run_succeeded") + return BatchRunSummary( + completed_batches=self._batch.completed, + counts=self._counts, + ) + + async def _execute_batches( + self, + iterator: object, + extract_context: PipelineContext, + ) -> None: + while True: + batch_index = self._batch.completed + 1 + self._batch.index = batch_index + context = self._record.start_context( + "extract", batch_index=batch_index, context=extract_context + ) + try: + try: + extracted = await anext(cast(Any, iterator)) + except StopAsyncIteration: + context._close() + self._record.complete_context(context) + if ( + self._batch.total is not None + and self._batch.completed != self._batch.total + ): + raise ValueError( + "completed batches must equal total_batches" + ) from None + self._batch.index = None + return + finally: + context._close() + + self._record.complete_context(context) + if ( + self._batch.total is not None + and batch_index > self._batch.total + ): + raise ValueError( + "extract yielded more batches than total_batches" + ) + + transformed = await self._transform_batch(extracted, batch_index) + counts = await self._load_batch(transformed, batch_index) + _merge_counts(self._counts, counts) + self._batch.completed += 1 + self._record.write_snapshot() + self._record.append_event( + "batch_completed", + stage="load", + batch=self._batch.as_json(), + counts=counts, + ) + + async def _transform_batch( + self, extracted: ExtractedBatchT, batch_index: int + ) -> TransformedBatchT: + context = self._record.start_context( + "transform", batch_index=batch_index + ) + try: + return await self._transform(extracted, ctx=context) + finally: + context._close() + if self._record.active_context is context: + self._record.complete_context(context) + + async def _load_batch( + self, transformed: TransformedBatchT, batch_index: int + ) -> dict[str, int]: + context = self._record.start_context("load", batch_index=batch_index) + try: + result = await self._load(transformed, ctx=context) + return _copy_count_deltas(result) + finally: + context._close() + if self._record.active_context is context: + self._record.complete_context(context) + + +def _validate_total_batches(value: int | None) -> int | None: + if value is None: + return None + if type(value) is not int: + raise TypeError("total_batches must be an integer or None") + if value < 0: + raise ValueError("total_batches must be >= 0") + return value + + +def _copy_count_deltas(value: CountDeltas) -> dict[str, int]: + if value is None: + return {} + return _copy_count_mapping(value, "load counts") + + +def _copy_count_mapping( + value: Mapping[str, int], + field_name: str, +) -> dict[str, int]: + copied: dict[str, int] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"{field_name} keys must be strings") + if type(item) is not int: + raise TypeError(f"{field_name}[{key!r}] must be an integer") + if item < 0: + raise ValueError(f"{field_name}[{key!r}] must be >= 0") + copied[key] = item + return copied + + +def _merge_counts( + totals: dict[str, int], + deltas: Mapping[str, int], +) -> None: + for key, value in deltas.items(): + totals[key] = totals.get(key, 0) + value diff --git a/quantmind/etl/_pipeline.py b/quantmind/etl/_pipeline.py new file mode 100644 index 0000000..ef89f64 --- /dev/null +++ b/quantmind/etl/_pipeline.py @@ -0,0 +1,272 @@ +"""Async, observable-by-default ``extract -> transform -> load`` execution.""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import Awaitable, Mapping +from pathlib import Path +from typing import Any, Generic, Protocol, TypeAlias, TypeVar, cast + +from quantmind.etl._record import ( + JsonScalar, + PipelineContext, + RunRecord, + StageName, + copy_scalar_mapping, + validate_dry_run, +) + +InputT = TypeVar("InputT") +ExtractedT = TypeVar("ExtractedT") +TransformedT = TypeVar("TransformedT") +OutputT = TypeVar("OutputT") +_StageInputT = TypeVar("_StageInputT", contravariant=True) +_StageOutputT = TypeVar("_StageOutputT", covariant=True) + + +class _StageCallable(Protocol[_StageInputT, _StageOutputT]): + def __call__( + self, + value: _StageInputT, + /, + *, + ctx: PipelineContext, + ) -> Awaitable[_StageOutputT]: ... + + +_Stage: TypeAlias = _StageCallable[Any, Any] + + +class ETLPipeline(Generic[InputT, ExtractedT, TransformedT, OutputT]): + """Bind three async callables into one fixed, observable ETL operation. + + This class uses composition instead of an inheritance hierarchy. Coding + agents implement the three stage functions and bind them once here; every + call to :meth:`create_run` then owns independent run-specific state. + Every authored ETL should honor the required ``dry_run`` run option: + callers pass ``create_run(..., dry_run=dry_run)`` and stage functions read + ``ctx.dry_run`` to pass the mode to the real mutation boundary. Dry-run + still executes ``extract``, ``transform``, and ``load``; the authored + ``load`` must validate or plan without persistent business writes, + including staging/checkpoint writes, and return an honest planned result + rather than a reference to a nonexistent delivery. AI processing may still + run and incur cost or rate usage. + """ + + def __init__( + self, + name: str, + *, + extract: _StageCallable[InputT, ExtractedT], + transform: _StageCallable[ExtractedT, TransformedT], + load: _StageCallable[TransformedT, OutputT], + ) -> None: + if not isinstance(name, str) or not name.strip(): + raise ValueError("pipeline name must be a non-empty string") + if ( + not callable(extract) + or not callable(transform) + or not callable(load) + ): + raise TypeError("extract, transform, and load must be callable") + self._name = name + self._extract = extract + self._transform = transform + self._load = load + + @property + def name(self) -> str: + """The stable name written to each local run record.""" + return self._name + + def create_run( + self, + source: InputT, + *, + dry_run: bool, + config_summary: Mapping[str, JsonScalar] | None = None, + run_root: Path | None = None, + ) -> PipelineRun[OutputT]: + """Create a one-shot run and atomically write its initial snapshot. + + Only ``config_summary`` is persisted. The source, stage values, full + configuration, and eventual result are kept out of the run record. + ``dry_run`` is required because it changes side-effect semantics and is + written as a top-level field from the initial ``created`` snapshot. + Do not put ``dry_run`` in ``config_summary``. + + Dry-run runs the complete stage shape, including ``load``. Stage + authors read ``ctx.dry_run`` and forward it to repositories, gateways, + publishers, or other mutation-capable dependencies. In dry-run, those + dependencies must not create persistent business mutations, including + staging/checkpoint writes or final delivery; they may return planned + paths, planned summaries, or ``None``. Use ``planned_*`` progress + metrics in dry-run. AI processing is allowed and may incur cost. + + Args: + source: Runtime input retained in memory until execution. + dry_run: Explicit run mode. ``True`` plans and validates without + persistent business delivery; ``False`` allows authored + mutation boundaries to commit. + config_summary: Explicit allowlist of safe JSON-scalar settings. + run_root: Directory that will contain the run directory. The + default is ``/.quant-mind/etl-pipeline-runs``. + + Returns: + A one-shot run whose ``status_file`` already exists in ``created`` + state. + """ + root = ( + Path.cwd() / ".quant-mind" / "etl-pipeline-runs" + if run_root is None + else Path(run_root) + ).resolve() + return PipelineRun._create( + pipeline_name=self._name, + source=source, + stages=(self._extract, self._transform, self._load), + dry_run=validate_dry_run(dry_run), + config_summary=copy_scalar_mapping( + config_summary, field_name="config_summary" + ), + run_root=root, + ) + + +class PipelineRun(Generic[OutputT]): + """One-shot execution handle with an atomically updated local snapshot.""" + + def __init__( + self, + *, + record: RunRecord, + source: object, + stages: tuple[_Stage, _Stage, _Stage], + ) -> None: + self._record = record + self._source = source + self._stages = stages + self._executed = False + + @classmethod + def _create( + cls, + *, + pipeline_name: str, + source: object, + stages: tuple[_Stage, _Stage, _Stage], + dry_run: bool, + config_summary: dict[str, JsonScalar], + run_root: Path, + ) -> PipelineRun[Any]: + record = RunRecord.create( + pipeline_name=pipeline_name, + dry_run=dry_run, + config_summary=config_summary, + run_root=run_root, + snapshot_schema="quantmind.etl.run/v1", + event_schema="quantmind.etl.event/v1", + ) + return cls( + record=record, + source=source, + stages=stages, + ) + + @property + def id(self) -> str: + """The unique local run ID.""" + return self._record.id + + @property + def dry_run(self) -> bool: + """Whether this run plans and validates without business delivery.""" + return self._record.dry_run + + @property + def status_file(self) -> Path: + """Absolute path to the latest atomic ``run.json`` snapshot.""" + return self._record.status_file + + @property + def events_file(self) -> Path: + """Absolute path to the sparse lifecycle ``events.jsonl`` record.""" + return self._record.events_file + + def receipt(self) -> str: + """Return the run ID and absolute status path as one JSON line.""" + return json.dumps( + { + "event": "etl_run_created", + "run_id": self._record.id, + "dry_run": self._record.dry_run, + "status_file": str(self._record.status_file), + }, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + ) + + async def execute(self) -> OutputT: + """Execute ``extract -> transform -> load`` exactly once. + + ``dry_run=True`` never changes the stage sequence; the stage functions + are responsible for honoring ``ctx.dry_run`` at mutation boundaries and + returning an honest planned value from ``load``. + + Cancellation is recorded as a terminal state on a best-effort basis, + then the original :class:`asyncio.CancelledError` is re-raised. + + Raises: + RuntimeError: If this handle has already been executed. + Exception: Any exception raised by a stage after recording failure. + asyncio.CancelledError: Re-raised after recording cancellation. + """ + if self._executed: + raise RuntimeError("PipelineRun.execute() can only be called once") + self._executed = True + self._record.start_running() + try: + self._record.append_event("run_started") + except Exception: + pass + + value: object = self._source + try: + for stage, operation in zip( + cast( + tuple[StageName, StageName, StageName], + ( + "extract", + "transform", + "load", + ), + ), + self._stages, + strict=True, + ): + context = self._record.start_context( + stage, event="stage_started" + ) + try: + value = await operation(value, ctx=context) + finally: + context._close() + self._record.complete_context(context, event="stage_completed") + except asyncio.CancelledError: + try: + self._record.finish("cancelled", "run_cancelled") + except Exception: + pass + raise + except Exception as exc: + self._record.set_error(exc) + try: + self._record.finish("failed", "run_failed") + except Exception: + pass + raise + + self._record.finish("succeeded", "run_succeeded") + return cast(OutputT, value) diff --git a/quantmind/etl/_record.py b/quantmind/etl/_record.py new file mode 100644 index 0000000..4eeb195 --- /dev/null +++ b/quantmind/etl/_record.py @@ -0,0 +1,590 @@ +"""Shared local run recording mechanics for observable ETL runs.""" + +from __future__ import annotations + +import contextvars +import json +import math +import os +import time +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Literal, TypeAlias + +JsonScalar: TypeAlias = None | bool | int | float | str +StageName: TypeAlias = Literal["extract", "transform", "load"] +RunState: TypeAlias = Literal[ + "created", "running", "succeeded", "failed", "cancelled" +] + +_PROGRESS_EVENT_INTERVAL_SECONDS = 1.0 +_MAX_ERROR_MESSAGE_LENGTH = 1000 +_CURRENT_PROGRESS_SCOPE: contextvars.ContextVar[tuple[int, int] | None] = ( + contextvars.ContextVar("quantmind_etl_progress_scope", default=None) +) +_RESERVED_CONFIG_SUMMARY_DRY_RUN_MESSAGE = ( + "config_summary['dry_run'] is reserved; use create_run(..., dry_run=...)" +) + + +def _utc_now() -> str: + """Return a compact UTC timestamp suitable for local run records.""" + value = datetime.now(timezone.utc).replace(microsecond=0) + return value.isoformat().replace("+00:00", "Z") + + +def _monotonic_seconds() -> float: + return time.monotonic() + + +def _new_run_id(timestamp: str) -> str: + compact_timestamp = timestamp.replace("-", "").replace(":", "") + return f"qmr_{compact_timestamp}_{uuid.uuid4().hex[:8]}" + + +def copy_scalar_mapping( + value: Mapping[str, JsonScalar] | None, + *, + field_name: str, +) -> dict[str, JsonScalar]: + """Copy and validate an explicit JSON-scalar allowlist.""" + if value is None: + return {} + + copied: dict[str, JsonScalar] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"{field_name} keys must be strings") + if field_name == "config_summary" and key == "dry_run": + raise ValueError(_RESERVED_CONFIG_SUMMARY_DRY_RUN_MESSAGE) + if item is not None and not isinstance(item, (bool, int, float, str)): + raise TypeError(f"{field_name}[{key!r}] must be a JSON scalar") + if isinstance(item, float) and not math.isfinite(item): + raise ValueError(f"{field_name}[{key!r}] must be finite") + copied[key] = item + return copied + + +def validate_dry_run(value: bool) -> bool: + """Validate the explicit run-level dry-run flag.""" + if type(value) is not bool: + raise TypeError("dry_run must be a boolean") + return value + + +def _write_json_atomic(path: Path, value: Mapping[str, object]) -> None: + temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + try: + with temporary.open("x", encoding="utf-8") as stream: + json.dump( + _json_safe_value(value), + stream, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + stream.write("\n") + stream.flush() + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _append_json_line(path: Path, value: Mapping[str, object]) -> None: + with path.open("a", encoding="utf-8") as stream: + json.dump( + _json_safe_value(value), + stream, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + stream.write("\n") + stream.flush() + + +def _json_safe_string(value: str) -> str: + return value.encode("utf-8", errors="backslashreplace").decode("utf-8") + + +def _json_safe_value(value: object) -> object: + if isinstance(value, str): + return _json_safe_string(value) + if isinstance(value, Mapping): + return { + _json_safe_string(key) if isinstance(key, str) else key: ( + _json_safe_value(item) + ) + for key, item in value.items() + } + if isinstance(value, list): + return [_json_safe_value(item) for item in value] + if isinstance(value, tuple): + return [_json_safe_value(item) for item in value] + return value + + +def _exception_message(exc: BaseException) -> str: + try: + message = str(exc) + except Exception as message_error: + message = ( + f"" + ) + return _json_safe_string(message)[:_MAX_ERROR_MESSAGE_LENGTH] + + +@dataclass(frozen=True, slots=True) +class Progress: + """Validated stage-local progress snapshot.""" + + completed: int + total: int | None + message: str | None + metrics: dict[str, JsonScalar] + + def as_json(self) -> dict[str, object]: + """Return the JSON representation written to snapshots and events.""" + return { + "completed": self.completed, + "total": self.total, + "message": self.message, + "metrics": self.metrics, + } + + +@dataclass(frozen=True, slots=True) +class _PendingProgress: + progress: Progress + occurred_at: str + + +class PipelineContext: + """Stage-local context used to read run mode and report completed work. + + ``completed`` is the amount of work that has finished, not work that was + queued or discovered. It must strictly increase within one active + ``(batch_index, stage)`` scope. Pass ``total=None`` while the total is + unknown; once a scope reports a known total, later updates for that scope + must keep reporting a known total. + + Every ETL pipeline should implement dry-run behavior. Callers choose it + once with ``create_run(..., dry_run=dry_run)``; all whole-run stages and all + batch-stage scopes then read the same immutable ``ctx.dry_run`` value. + Stage authors pass that flag to the repository, gateway, publisher, or + other capability that owns mutation decisions. Dry-run still executes every + stage, including ``load``, but authored code must prevent persistent + business mutations such as database writes, storage artifacts, staging or + checkpoint files, publish/queue/webhook calls, and final delivery. Use + ``planned_*`` names for dry-run progress metrics and count deltas. AI + processing may still run and incur cost or rate usage; dry-run is not a + zero-cost mode. + """ + + def __init__( + self, + *, + run_id: str, + dry_run: bool, + stage: StageName, + report: Callable[[PipelineContext, Progress], None], + batch_index: int | None = None, + active: bool = True, + ) -> None: + self._run_id = run_id + self._dry_run = dry_run + self._stage = stage + self._batch_index = batch_index + self._report = report + self._previous: Progress | None = None + self._active = False + self._activation_epoch = 0 + self._scope_token: contextvars.Token[tuple[int, int] | None] | None = ( + None + ) + if active: + self._activate(stage=stage, batch_index=batch_index) + + @property + def run_id(self) -> str: + """The ID of the run that owns this stage.""" + return self._run_id + + @property + def dry_run(self) -> bool: + """Whether this run is planning and validating without delivery.""" + return self._dry_run + + @property + def stage(self) -> str: + """The fixed macro stage currently being executed.""" + return self._stage + + @property + def batch_index(self) -> int | None: + """The current 1-based batch index, or ``None`` for whole-run ETL.""" + return self._batch_index + + async def progress( + self, + completed: int, + *, + total: int | None = None, + message: str | None = None, + metrics: Mapping[str, JsonScalar] | None = None, + ) -> None: + """Record real completed work for the current stage. + + Args: + completed: Finished work. It must be a non-negative integer and + strictly increase within this active ``(batch, stage)`` scope. + total: Known total work, or ``None`` while it is unknown. + message: Optional short, safe progress description. + metrics: Optional allowlisted JSON-scalar measurements. + + Raises: + RuntimeError: If the stage has already ended. + TypeError: If progress values have invalid types. + ValueError: If progress moves backwards or contradicts ``total``. + """ + if not self._active: + raise RuntimeError( + "progress cannot be reported after the stage ends" + ) + if _CURRENT_PROGRESS_SCOPE.get() != ( + id(self), + self._activation_epoch, + ): + raise RuntimeError( + "progress can only be reported by the active stage scope" + ) + if type(completed) is not int: + raise TypeError("completed must be an integer") + if completed < 0: + raise ValueError("completed must be >= 0") + if total is not None and type(total) is not int: + raise TypeError("total must be an integer or None") + if total is not None and total < completed: + raise ValueError("total must be >= completed") + if message is not None and not isinstance(message, str): + raise TypeError("message must be a string or None") + + previous = self._previous + if previous is not None and completed <= previous.completed: + raise ValueError("completed must strictly increase within a stage") + if ( + previous is not None + and previous.total is not None + and total is None + ): + raise ValueError("total cannot become unknown after it was known") + + progress = Progress( + completed=completed, + total=total, + message=message, + metrics=copy_scalar_mapping(metrics, field_name="metrics"), + ) + self._report(self, progress) + self._previous = progress + + def _activate(self, *, stage: StageName, batch_index: int | None) -> None: + self._stage = stage + self._batch_index = batch_index + self._previous = None + self._active = True + self._activation_epoch += 1 + self._scope_token = _CURRENT_PROGRESS_SCOPE.set( + (id(self), self._activation_epoch) + ) + + def _close(self) -> None: + self._active = False + token = self._scope_token + self._scope_token = None + if token is not None: + _CURRENT_PROGRESS_SCOPE.reset(token) + + +class RunRecord: + """Atomic snapshot and sparse event writer shared by ETL run shapes.""" + + def __init__( + self, + *, + run_id: str, + pipeline_name: str, + dry_run: bool, + config_summary: dict[str, JsonScalar], + run_directory: Path, + created_at: str, + snapshot_schema: str, + event_schema: str, + snapshot_extra: Callable[[], Mapping[str, object]] | None = None, + ) -> None: + self.id = run_id + self.pipeline_name = pipeline_name + self.dry_run = dry_run + self.config_summary = config_summary + self.run_directory = run_directory + self.status_file = run_directory / "run.json" + self.events_file = run_directory / "events.jsonl" + self.created_at = created_at + self.updated_at = created_at + self.snapshot_schema = snapshot_schema + self.event_schema = event_schema + self.snapshot_extra = snapshot_extra + self.state: RunState = "created" + self.stage: StageName | None = None + self.pid: int | None = None + self.progress: Progress | None = None + self.error: dict[str, str] | None = None + self.active_context: PipelineContext | None = None + self._last_progress_event_at: float | None = None + self._pending_progress: _PendingProgress | None = None + + @classmethod + def create( + cls, + *, + pipeline_name: str, + dry_run: bool, + config_summary: dict[str, JsonScalar], + run_root: Path, + snapshot_schema: str, + event_schema: str, + snapshot_extra: Callable[[], Mapping[str, object]] | None = None, + ) -> RunRecord: + """Create the run directory and write the initial snapshot.""" + created_at = _utc_now() + run_id = _new_run_id(created_at) + run_root.mkdir(parents=True, exist_ok=True) + run_directory = run_root / run_id + run_directory.mkdir(mode=0o700) + events_file = run_directory / "events.jsonl" + events_file.touch(exist_ok=False) + + record = cls( + run_id=run_id, + pipeline_name=pipeline_name, + dry_run=dry_run, + config_summary=config_summary, + run_directory=run_directory, + created_at=created_at, + snapshot_schema=snapshot_schema, + event_schema=event_schema, + snapshot_extra=snapshot_extra, + ) + record.write_snapshot() + return record + + def start_running(self) -> None: + """Mark the run as running and write the first live snapshot.""" + self.state = "running" + self.pid = os.getpid() + self.updated_at = _utc_now() + self.write_snapshot() + + def new_inactive_context( + self, + *, + stage: StageName, + batch_index: int | None = None, + ) -> PipelineContext: + """Create a reusable inactive context for lazy async iterators.""" + return PipelineContext( + run_id=self.id, + dry_run=self.dry_run, + stage=stage, + batch_index=batch_index, + report=self.record_progress, + active=False, + ) + + def start_context( + self, + stage: StageName, + *, + batch_index: int | None = None, + context: PipelineContext | None = None, + event: str | None = None, + ) -> PipelineContext: + """Start a stage or batch-stage context and reset progress state.""" + if self.active_context is not None: + raise RuntimeError("another stage context is already active") + self.stage = stage + self.progress = None + self._pending_progress = None + self._last_progress_event_at = None + self.updated_at = _utc_now() + + if context is None: + context = PipelineContext( + run_id=self.id, + dry_run=self.dry_run, + stage=stage, + batch_index=batch_index, + report=self.record_progress, + ) + else: + context._activate(stage=stage, batch_index=batch_index) + + self.active_context = context + self.write_snapshot() + if event is not None: + self.append_event(event, stage=stage) + return context + + def complete_context( + self, + context: PipelineContext, + *, + event: str | None = None, + ) -> None: + """Flush progress and mark a stage or batch-stage context complete.""" + if self.active_context is not context: + raise RuntimeError("stage context is no longer active") + self.flush_pending_progress() + self.updated_at = _utc_now() + self.write_snapshot() + if event is not None: + self.append_event(event, stage=context.stage) + self.active_context = None + + def record_progress( + self, context: PipelineContext, progress: Progress + ) -> None: + """Merge progress into the snapshot and coalesce progress events.""" + if self.state != "running" or self.active_context is not context: + raise RuntimeError( + "progress can only be reported by the active stage" + ) + occurred_at = _utc_now() + self.progress = progress + self.updated_at = occurred_at + self.write_snapshot() + + now = _monotonic_seconds() + last = self._last_progress_event_at + if last is None or now - last >= _PROGRESS_EVENT_INTERVAL_SECONDS: + self._append_progress_event( + context, progress, occurred_at=occurred_at + ) + self._last_progress_event_at = now + self._pending_progress = None + else: + self._pending_progress = _PendingProgress(progress, occurred_at) + + def flush_pending_progress(self) -> None: + """Append the last coalesced progress update, if one exists.""" + pending = self._pending_progress + context = self.active_context + if pending is None or context is None: + return + self._append_progress_event( + context, + pending.progress, + occurred_at=pending.occurred_at, + ) + self._pending_progress = None + + def set_error(self, exc: BaseException) -> None: + """Store a bounded, JSON-safe error summary.""" + self.error = { + "type": type(exc).__name__, + "message": _exception_message(exc), + } + + def clear_progress(self) -> None: + """Clear progress before a terminal snapshot that changes scope.""" + self.progress = None + self._pending_progress = None + self._last_progress_event_at = None + + def finish(self, state: RunState, event: str) -> None: + """Write the terminal snapshot and best-effort lifecycle event.""" + if self.active_context is not None: + self.active_context._close() + try: + self.flush_pending_progress() + except Exception: + self._pending_progress = None + self.active_context = None + self.state = state + self.updated_at = _utc_now() + self.write_snapshot() + try: + if self.error is None: + self.append_event(event, stage=self.stage) + else: + self.append_event(event, stage=self.stage, error=self.error) + except Exception: + pass + + def write_snapshot(self) -> None: + """Atomically write the current ``run.json`` snapshot.""" + value: dict[str, object] = { + "schema": self.snapshot_schema, + "run_id": self.id, + "pipeline": self.pipeline_name, + "dry_run": self.dry_run, + "state": self.state, + "stage": self.stage, + "pid": self.pid, + "created_at": self.created_at, + "updated_at": self.updated_at, + "config_summary": self.config_summary, + "progress": ( + None if self.progress is None else self.progress.as_json() + ), + "error": self.error, + } + if self.snapshot_extra is not None: + value.update(self.snapshot_extra()) + _write_json_atomic(self.status_file, value) + + def append_event( + self, + event: str, + *, + stage: str | None = None, + occurred_at: str | None = None, + **fields: object, + ) -> None: + """Append one sparse JSON event.""" + value: dict[str, object] = { + "schema": self.event_schema, + "event": event, + "run_id": self.id, + "dry_run": self.dry_run, + "occurred_at": occurred_at or _utc_now(), + } + if stage is not None: + value["stage"] = stage + value.update(fields) + _append_json_line(self.events_file, value) + + def _append_progress_event( + self, + context: PipelineContext, + progress: Progress, + *, + occurred_at: str, + ) -> None: + fields: dict[str, object] = { + "progress": progress.as_json(), + } + if context.batch_index is not None: + fields["batch_index"] = context.batch_index + self.append_event( + "stage_progress", + stage=context.stage, + occurred_at=occurred_at, + **fields, + ) diff --git a/tests/etl/__init__.py b/tests/etl/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/etl/test_batch.py b/tests/etl/test_batch.py new file mode 100644 index 0000000..59a009c --- /dev/null +++ b/tests/etl/test_batch.py @@ -0,0 +1,1053 @@ +"""Tests for the observable micro-batch ETL scaffold.""" + +import asyncio +import json +import os +import subprocess +import sys +import tempfile +import textwrap +import unittest +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import quantmind.etl._record as record_module +from quantmind.etl import BatchETLPipeline, PipelineContext + + +def _read_json(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def _read_events(path: Path) -> list[dict[str, Any]]: + return [ + json.loads(line) + for line in path.read_text(encoding="utf-8").splitlines() + ] + + +def _json_safe_text(value: str) -> str: + return value.encode("utf-8", errors="backslashreplace").decode("utf-8") + + +class BatchPipelineRunTests(unittest.IsolatedAsyncioTestCase): + def setUp(self) -> None: + self._temporary_directory = tempfile.TemporaryDirectory() + self.run_root = Path(self._temporary_directory.name) + + def tearDown(self) -> None: + self._temporary_directory.cleanup() + + @staticmethod + def _pipeline( + *, + extract=None, + transform=None, + load=None, + ) -> BatchETLPipeline[list[int], int, int]: + async def default_extract(source: list[int], *, ctx: PipelineContext): + for item in source: + yield item + + async def default_transform(batch: int, *, ctx: PipelineContext) -> int: + return batch * 10 + + async def default_load( + batch: int, *, ctx: PipelineContext + ) -> dict[str, int]: + return {"batches": 1, "items": batch // 10} + + return BatchETLPipeline( + "test-batch-pipeline", + extract=extract or default_extract, + transform=transform or default_transform, + load=load or default_load, + ) + + async def test_create_run_writes_batch_snapshot_and_receipt(self) -> None: + run = self._pipeline().create_run( + [1, 2], + dry_run=True, + run_root=self.run_root, + config_summary={"window_days": 7}, + total_batches=2, + ) + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["schema"], "quantmind.etl.batch-run/v1") + self.assertTrue(snapshot["dry_run"]) + self.assertEqual(snapshot["state"], "created") + self.assertIsNone(snapshot["stage"]) + self.assertEqual( + snapshot["batch"], {"index": None, "completed": 0, "total": 2} + ) + self.assertEqual(snapshot["config_summary"], {"window_days": 7}) + self.assertEqual(run.events_file.read_text(encoding="utf-8"), "") + + receipt = json.loads(run.receipt()) + self.assertEqual(receipt["event"], "etl_batch_run_created") + self.assertEqual(receipt["run_id"], run.id) + self.assertTrue(receipt["dry_run"]) + self.assertEqual(receipt["status_file"], str(run.status_file)) + self.assertTrue(Path(receipt["status_file"]).is_absolute()) + + async def test_success_runs_each_batch_serially_with_honest_snapshot( + self, + ) -> None: + run_holder = [] + calls: list[tuple[str, int | None, dict[str, object] | None]] = [] + generator_started = False + + async def extract(source: list[int], *, ctx: PipelineContext): + nonlocal generator_started + generator_started = True + for item in source: + snapshot = _read_json(run_holder[0].status_file) + calls.append( + ( + "extract", + ctx.batch_index, + snapshot["batch"], + ) + ) + self.assertEqual(snapshot["stage"], "extract") + self.assertIsNone(snapshot["progress"]) + self.assertFalse(ctx.dry_run) + yield item + + async def transform(batch: int, *, ctx: PipelineContext) -> int: + snapshot = _read_json(run_holder[0].status_file) + calls.append(("transform", ctx.batch_index, snapshot["batch"])) + self.assertEqual(snapshot["stage"], "transform") + self.assertIsNone(snapshot["progress"]) + self.assertFalse(ctx.dry_run) + return batch * 10 + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + snapshot = _read_json(run_holder[0].status_file) + calls.append(("load", ctx.batch_index, snapshot["batch"])) + self.assertEqual(snapshot["stage"], "load") + self.assertIsNone(snapshot["progress"]) + self.assertFalse(ctx.dry_run) + return {"batches": 1, "items": batch // 10} + + run = self._pipeline( + extract=extract, transform=transform, load=load + ).create_run( + [1, 2], + dry_run=False, + run_root=self.run_root, + total_batches=2, + ) + run_holder.append(run) + self.assertFalse(generator_started) + + summary = await run.execute() + + self.assertEqual(summary.completed_batches, 2) + self.assertEqual(dict(summary.counts), {"batches": 2, "items": 3}) + with self.assertRaises(TypeError): + summary.counts["items"] = 99 # type: ignore[index] + self.assertEqual( + calls, + [ + ("extract", 1, {"index": 1, "completed": 0, "total": 2}), + ("transform", 1, {"index": 1, "completed": 0, "total": 2}), + ("load", 1, {"index": 1, "completed": 0, "total": 2}), + ("extract", 2, {"index": 2, "completed": 1, "total": 2}), + ("transform", 2, {"index": 2, "completed": 1, "total": 2}), + ("load", 2, {"index": 2, "completed": 1, "total": 2}), + ], + ) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "succeeded") + self.assertFalse(snapshot["dry_run"]) + self.assertEqual(snapshot["stage"], "load") + self.assertEqual( + snapshot["batch"], {"index": None, "completed": 2, "total": 2} + ) + self.assertEqual( + [event["event"] for event in _read_events(run.events_file)], + [ + "run_started", + "batch_completed", + "batch_completed", + "run_succeeded", + ], + ) + + async def test_every_event_records_schema_and_dry_run(self) -> None: + async def transform(batch: int, *, ctx: PipelineContext) -> int: + await ctx.progress(1, total=1, message="transformed") + return batch * 10 + + run = self._pipeline(transform=transform).create_run( + [1], dry_run=True, run_root=self.run_root, total_batches=1 + ) + + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["schema"], "quantmind.etl.batch-run/v1") + self.assertTrue(snapshot["dry_run"]) + events = _read_events(run.events_file) + self.assertGreater(len(events), 0) + for event in events: + self.assertEqual(event["schema"], "quantmind.etl.batch-event/v1") + self.assertEqual(event["run_id"], run.id) + self.assertTrue(event["dry_run"]) + + async def test_dry_run_reaches_every_batch_stage_scope( + self, + ) -> None: + calls: list[tuple[str, int | None, bool]] = [] + + async def extract(source: list[int], *, ctx: PipelineContext): + for item in source: + calls.append(("extract", ctx.batch_index, ctx.dry_run)) + yield item + + async def transform(batch: int, *, ctx: PipelineContext) -> int: + calls.append(("transform", ctx.batch_index, ctx.dry_run)) + return batch * 10 + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + calls.append(("load", ctx.batch_index, ctx.dry_run)) + await ctx.progress( + 1, + total=1, + metrics={"planned_records": batch // 10}, + ) + return {"planned_batches": 1, "planned_records": batch // 10} + + run = self._pipeline( + extract=extract, transform=transform, load=load + ).create_run( + [1, 2], + dry_run=True, + run_root=self.run_root, + total_batches=2, + ) + + summary = await run.execute() + + self.assertTrue(run.dry_run) + self.assertEqual(summary.completed_batches, 2) + self.assertEqual( + dict(summary.counts), + {"planned_batches": 2, "planned_records": 3}, + ) + self.assertEqual( + calls, + [ + ("extract", 1, True), + ("transform", 1, True), + ("load", 1, True), + ("extract", 2, True), + ("transform", 2, True), + ("load", 2, True), + ], + ) + snapshot = _read_json(run.status_file) + self.assertTrue(snapshot["dry_run"]) + + async def test_json_records_escape_surrogate_strings_recursively( + self, + ) -> None: + unsafe = os.fsdecode(b"\xff") + safe = _json_safe_text(unsafe) + + async def transform(batch: int, *, ctx: PipelineContext) -> int: + await ctx.progress( + 1, + total=1, + message=unsafe, + metrics={unsafe: unsafe}, + ) + return batch * 10 + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + return {unsafe: 1} + + run = self._pipeline(transform=transform, load=load).create_run( + [1], + dry_run=False, + run_root=self.run_root, + config_summary={unsafe: unsafe}, + total_batches=1, + ) + + summary = await run.execute() + + self.assertEqual(dict(summary.counts), {unsafe: 1}) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "succeeded") + self.assertEqual(snapshot["config_summary"], {safe: safe}) + progress_events = [ + event + for event in _read_events(run.events_file) + if event["event"] == "stage_progress" + ] + self.assertEqual(progress_events[-1]["progress"]["message"], safe) + self.assertEqual( + progress_events[-1]["progress"]["metrics"], {safe: safe} + ) + batch_events = [ + event + for event in _read_events(run.events_file) + if event["event"] == "batch_completed" + ] + self.assertEqual(batch_events[-1]["counts"], {safe: 1}) + + async def test_unknown_total_succeeds_with_null_total(self) -> None: + run = self._pipeline().create_run( + [1, 2], dry_run=False, run_root=self.run_root + ) + + summary = await run.execute() + + self.assertEqual(summary.completed_batches, 2) + self.assertEqual( + _read_json(run.status_file)["batch"], + {"index": None, "completed": 2, "total": None}, + ) + + async def test_known_total_rejects_too_few_batches(self) -> None: + run = self._pipeline().create_run( + [1, 2], + dry_run=False, + run_root=self.run_root, + total_batches=3, + ) + + with self.assertRaisesRegex( + ValueError, "completed batches must equal total_batches" + ): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["stage"], "extract") + self.assertEqual( + snapshot["batch"], {"index": 3, "completed": 2, "total": 3} + ) + + async def test_known_total_rejects_too_many_batches(self) -> None: + run = self._pipeline().create_run( + [1, 2], + dry_run=False, + run_root=self.run_root, + total_batches=1, + ) + + with self.assertRaisesRegex( + ValueError, "extract yielded more batches than total_batches" + ): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["stage"], "extract") + self.assertEqual( + snapshot["batch"], {"index": 2, "completed": 1, "total": 1} + ) + + async def test_progress_resets_per_batch_stage_and_events_include_batch( + self, + ) -> None: + async def extract(source: list[int], *, ctx: PipelineContext): + for item in source: + await ctx.progress(1, total=1, message="fetched") + yield item + + async def transform(batch: int, *, ctx: PipelineContext) -> int: + await ctx.progress(1, total=1, message="transformed") + return batch + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + await ctx.progress(1, total=1, message="loaded") + return {"batches": 1} + + run = self._pipeline( + extract=extract, transform=transform, load=load + ).create_run([1, 2], dry_run=False, run_root=self.run_root) + + await run.execute() + + progress_events = [ + event + for event in _read_events(run.events_file) + if event["event"] == "stage_progress" + ] + self.assertEqual( + [ + (event["stage"], event["batch_index"]) + for event in progress_events + ], + [ + ("extract", 1), + ("transform", 1), + ("load", 1), + ("extract", 2), + ("transform", 2), + ("load", 2), + ], + ) + self.assertTrue( + all( + event["progress"]["completed"] == 1 for event in progress_events + ) + ) + + async def test_current_batch_stage_child_tasks_can_report_progress( + self, + ) -> None: + async def transform(batch: int, *, ctx: PipelineContext) -> int: + async def first_worker() -> None: + await ctx.progress(1, total=2) + + async def second_worker() -> None: + await asyncio.sleep(0) + await ctx.progress(2, total=2) + + await asyncio.gather(first_worker(), second_worker()) + return batch + + run = self._pipeline(transform=transform).create_run( + [1], dry_run=False, run_root=self.run_root + ) + + await run.execute() + + progress_events = [ + event + for event in _read_events(run.events_file) + if event["event"] == "stage_progress" + ] + self.assertEqual( + [ + ( + event["stage"], + event["batch_index"], + event["progress"]["completed"], + ) + for event in progress_events + ], + [("transform", 1, 1), ("transform", 1, 2)], + ) + + async def test_transform_failure_on_batch_two_preserves_completed_one( + self, + ) -> None: + async def transform(batch: int, *, ctx: PipelineContext) -> int: + if batch == 2: + raise LookupError("bad transform") + return batch + + run = self._pipeline(transform=transform).create_run( + [1, 2], dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(LookupError, "bad transform"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["stage"], "transform") + self.assertEqual( + snapshot["batch"], {"index": 2, "completed": 1, "total": None} + ) + self.assertEqual( + [ + event["event"] + for event in _read_events(run.events_file) + if event["event"] == "batch_completed" + ], + ["batch_completed"], + ) + + async def test_load_failure_on_batch_two_preserves_completed_one( + self, + ) -> None: + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + if batch == 20: + raise OSError("sink rejected") + return {"batches": 1} + + run = self._pipeline(load=load).create_run( + [1, 2], dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(OSError, "sink rejected"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["stage"], "load") + self.assertEqual( + snapshot["batch"], {"index": 2, "completed": 1, "total": None} + ) + + async def test_cancellation_on_batch_two_preserves_completed_one_without_aclose( + self, + ) -> None: + entered = asyncio.Event() + closed = False + + async def extract(source: list[int], *, ctx: PipelineContext): + nonlocal closed + try: + for item in source: + yield item + finally: + closed = True + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + if batch == 20: + entered.set() + await asyncio.Event().wait() + return {"batches": 1} + + run = self._pipeline(extract=extract, load=load).create_run( + [1, 2, 3], dry_run=False, run_root=self.run_root + ) + task = asyncio.create_task(run.execute()) + await entered.wait() + + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + snapshot = _read_json(run.status_file) + self.assertFalse(closed) + self.assertEqual(snapshot["state"], "cancelled") + self.assertEqual(snapshot["stage"], "load") + self.assertEqual( + snapshot["batch"], {"index": 2, "completed": 1, "total": None} + ) + + async def test_cancellation_during_extract_anext_runs_generator_cleanup( + self, + ) -> None: + entered = asyncio.Event() + closed = False + + async def extract(source: list[int], *, ctx: PipelineContext): + nonlocal closed + try: + entered.set() + await asyncio.Event().wait() + yield source[0] + finally: + closed = True + + run = self._pipeline(extract=extract).create_run( + [1], dry_run=False, run_root=self.run_root + ) + task = asyncio.create_task(run.execute()) + await entered.wait() + + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + snapshot = _read_json(run.status_file) + self.assertTrue(closed) + self.assertEqual(snapshot["state"], "cancelled") + self.assertEqual(snapshot["stage"], "extract") + self.assertEqual( + snapshot["batch"], {"index": 1, "completed": 0, "total": None} + ) + + async def test_zero_batches_succeeds_without_batch_completed_event( + self, + ) -> None: + run = self._pipeline().create_run( + [], dry_run=False, run_root=self.run_root, total_batches=0 + ) + + summary = await run.execute() + + self.assertEqual(summary.completed_batches, 0) + self.assertEqual(dict(summary.counts), {}) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "succeeded") + self.assertIsNone(snapshot["stage"]) + self.assertEqual( + snapshot["batch"], {"index": None, "completed": 0, "total": 0} + ) + self.assertEqual( + [event["event"] for event in _read_events(run.events_file)], + ["run_started", "run_succeeded"], + ) + + async def test_success_after_final_extract_progress_clears_progress( + self, + ) -> None: + async def extract(source: list[int], *, ctx: PipelineContext): + yield source[0] + await ctx.progress( + 1, + total=1, + message="checked for another batch", + ) + + run = self._pipeline(extract=extract).create_run( + [1], dry_run=False, run_root=self.run_root, total_batches=1 + ) + + summary = await run.execute() + + self.assertEqual(summary.completed_batches, 1) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "succeeded") + self.assertEqual(snapshot["stage"], "load") + self.assertIsNone(snapshot["progress"]) + self.assertEqual( + snapshot["batch"], {"index": None, "completed": 1, "total": 1} + ) + + async def test_zero_batch_success_after_extract_progress_clears_stage( + self, + ) -> None: + async def extract(source: list[int], *, ctx: PipelineContext): + await ctx.progress(1, total=1, message="confirmed empty") + if False: + yield source[0] + + run = self._pipeline(extract=extract).create_run( + [], dry_run=False, run_root=self.run_root, total_batches=0 + ) + + summary = await run.execute() + + self.assertEqual(summary.completed_batches, 0) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "succeeded") + self.assertIsNone(snapshot["stage"]) + self.assertIsNone(snapshot["progress"]) + self.assertEqual( + snapshot["batch"], {"index": None, "completed": 0, "total": 0} + ) + + async def test_execute_is_one_shot(self) -> None: + run = self._pipeline().create_run( + [1], dry_run=False, run_root=self.run_root + ) + await run.execute() + + with self.assertRaisesRegex(RuntimeError, "only be called once"): + await run.execute() + + async def test_extract_context_is_inactive_after_anext_returns( + self, + ) -> None: + captured_contexts: list[PipelineContext] = [] + + async def extract(source: list[int], *, ctx: PipelineContext): + captured_contexts.append(ctx) + yield source[0] + + async def transform(batch: int, *, ctx: PipelineContext) -> int: + with self.assertRaisesRegex(RuntimeError, "after the stage ends"): + await captured_contexts[0].progress(1, total=1) + return batch + + run = self._pipeline(extract=extract, transform=transform).create_run( + [1], dry_run=False, run_root=self.run_root + ) + + await run.execute() + + with self.assertRaisesRegex(RuntimeError, "after the stage ends"): + await captured_contexts[0].progress(1, total=1) + + async def test_stale_background_extract_context_cannot_report_later( + self, + ) -> None: + release = asyncio.Event() + background_errors: list[str] = [] + + async def extract(source: list[int], *, ctx: PipelineContext): + async def report_later() -> None: + await release.wait() + try: + await ctx.progress(1, total=1) + except RuntimeError as exc: + background_errors.append(str(exc)) + + task = asyncio.create_task(report_later()) + yield source[0] + release.set() + await asyncio.sleep(0) + await ctx.progress(1, total=1) + yield source[1] + await task + + run = self._pipeline(extract=extract).create_run( + [1, 2], dry_run=False, run_root=self.run_root + ) + + await run.execute() + + self.assertEqual( + background_errors, + ["progress can only be reported by the active stage scope"], + ) + + async def test_load_count_validation_fails_before_batch_completion( + self, + ) -> None: + async def load(batch: int, *, ctx: PipelineContext): + return {"items": True} + + run = self._pipeline(load=load).create_run( + [1], dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(TypeError, "must be an integer"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual( + snapshot["batch"], {"index": 1, "completed": 0, "total": None} + ) + + async def test_load_failure_does_not_start_extractor_aclose( + self, + ) -> None: + class ClosingExtractor: + def __init__(self) -> None: + self.next_value = 1 + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self) -> int: + if self.next_value > 3: + raise StopAsyncIteration + value = self.next_value + self.next_value += 1 + return value + + async def aclose(self) -> None: + self.closed = True + raise RuntimeError("cleanup failed") + + extractor = ClosingExtractor() + + def extract(source: list[int], *, ctx: PipelineContext): + return extractor + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + if batch == 20: + raise LookupError("original load failure") + return {"batches": 1} + + run = self._pipeline(extract=extract, load=load).create_run( + [1, 2, 3], dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(LookupError, "original load failure"): + await run.execute() + + self.assertFalse(extractor.closed) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["error"]["type"], "LookupError") + + async def test_hanging_extractor_aclose_is_not_started_on_load_failure( + self, + ) -> None: + class HangingExtractor: + def __init__(self) -> None: + self.next_value = 1 + self.close_started = False + + def __aiter__(self): + return self + + async def __anext__(self) -> int: + if self.next_value > 3: + raise StopAsyncIteration + value = self.next_value + self.next_value += 1 + return value + + async def aclose(self) -> None: + self.close_started = True + await asyncio.Event().wait() + + extractor = HangingExtractor() + + def extract(source: list[int], *, ctx: PipelineContext): + return extractor + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + if batch == 20: + raise LookupError("original load failure") + return {"batches": 1} + + run = self._pipeline(extract=extract, load=load).create_run( + [1, 2, 3], dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(LookupError, "original load failure"): + await run.execute() + + self.assertFalse(extractor.close_started) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["error"]["type"], "LookupError") + + async def test_asyncio_run_exits_when_extractor_aclose_swallows_cancel( + self, + ) -> None: + script = r""" +import asyncio +import tempfile +from pathlib import Path + +import quantmind.etl._batch as batch_module +from quantmind.etl import BatchETLPipeline, PipelineContext + +if hasattr(batch_module, "_EXTRACTOR_CLOSE_TIMEOUT_SECONDS"): + batch_module._EXTRACTOR_CLOSE_TIMEOUT_SECONDS = 0.01 + + +class StubbornExtractor: + def __init__(self) -> None: + self.next_value = 1 + + def __aiter__(self): + return self + + async def __anext__(self) -> int: + if self.next_value > 2: + raise StopAsyncIteration + value = self.next_value + self.next_value += 1 + return value + + async def aclose(self) -> None: + while True: + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + pass + + +async def transform(batch: int, *, ctx: PipelineContext) -> int: + return batch * 10 + + +async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + if batch == 20: + raise LookupError("original load failure") + return {"batches": 1} + + +async def main() -> None: + extractor = StubbornExtractor() + + def extract(source: list[int], *, ctx: PipelineContext): + return extractor + + with tempfile.TemporaryDirectory() as directory: + pipeline = BatchETLPipeline( + "shutdown-regression", + extract=extract, + transform=transform, + load=load, + ) + run = pipeline.create_run( + [1, 2], + dry_run=False, + run_root=Path(directory), + ) + try: + await run.execute() + except LookupError: + pass + + +asyncio.run(main()) +""" + completed = subprocess.run( + [sys.executable, "-c", textwrap.dedent(script)], + check=False, + cwd=Path(__file__).resolve().parents[2], + text=True, + capture_output=True, + timeout=3, + ) + + self.assertEqual( + completed.returncode, + 0, + f"stdout={completed.stdout}\nstderr={completed.stderr}", + ) + + async def test_hanging_extractor_aclose_is_not_started_on_cancellation( + self, + ) -> None: + entered = asyncio.Event() + + class HangingExtractor: + def __init__(self) -> None: + self.next_value = 1 + self.close_started = False + + def __aiter__(self): + return self + + async def __anext__(self) -> int: + if self.next_value > 3: + raise StopAsyncIteration + value = self.next_value + self.next_value += 1 + return value + + async def aclose(self) -> None: + self.close_started = True + await asyncio.Event().wait() + + extractor = HangingExtractor() + + def extract(source: list[int], *, ctx: PipelineContext): + return extractor + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + if batch == 20: + entered.set() + await asyncio.Event().wait() + return {"batches": 1} + + run = self._pipeline(extract=extract, load=load).create_run( + [1, 2, 3], dry_run=False, run_root=self.run_root + ) + task = asyncio.create_task(run.execute()) + await entered.wait() + + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + self.assertFalse(extractor.close_started) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "cancelled") + + async def test_failed_terminal_event_append_does_not_mask_batch_error( + self, + ) -> None: + original_append = record_module._append_json_line + + def append_event(path: Path, value: dict[str, object]) -> None: + if value["event"] == "run_failed": + raise OSError("event sink unavailable") + original_append(path, value) + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + raise LookupError("batch load failed") + + run = self._pipeline(load=load).create_run( + [1], dry_run=False, run_root=self.run_root + ) + + with patch( + "quantmind.etl._record._append_json_line", + side_effect=append_event, + ): + with self.assertRaisesRegex(LookupError, "batch load failed"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["error"]["type"], "LookupError") + + async def test_run_started_event_append_failure_does_not_leave_running( + self, + ) -> None: + original_append = record_module._append_json_line + + def append_event(path: Path, value: dict[str, object]) -> None: + if value["event"] == "run_started": + raise OSError("event sink unavailable") + original_append(path, value) + + async def load(batch: int, *, ctx: PipelineContext) -> dict[str, int]: + raise LookupError("batch load failed") + + run = self._pipeline(load=load).create_run( + [1], dry_run=False, run_root=self.run_root + ) + + with patch( + "quantmind.etl._record._append_json_line", + side_effect=append_event, + ): + with self.assertRaisesRegex(LookupError, "batch load failed"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["error"]["type"], "LookupError") + self.assertNotEqual(snapshot["state"], "running") + + +class BatchPipelineValidationTests(unittest.TestCase): + def test_pipeline_name_must_not_be_empty(self) -> None: + async def extract(value, *, ctx): + yield value + + async def stage(value, *, ctx): + return value + + with self.assertRaises(ValueError): + BatchETLPipeline("", extract=extract, transform=stage, load=stage) + + def test_total_batches_rejects_bool_and_negative_values(self) -> None: + async def extract(value, *, ctx): + yield value + + async def stage(value, *, ctx): + return value + + pipeline = BatchETLPipeline( + "safe-summary", extract=extract, transform=stage, load=stage + ) + with tempfile.TemporaryDirectory() as directory: + with self.assertRaisesRegex(TypeError, "dry_run"): + pipeline.create_run("alpha", run_root=Path(directory)) + with self.assertRaisesRegex(TypeError, "dry_run must be"): + pipeline.create_run( + "alpha", + dry_run=None, + run_root=Path(directory), + ) + with self.assertRaisesRegex( + ValueError, r"create_run\(..., dry_run=...\)" + ): + pipeline.create_run( + "alpha", + dry_run=False, + run_root=Path(directory), + config_summary={"dry_run": False}, + ) + with self.assertRaises(TypeError): + pipeline.create_run( + "alpha", + dry_run=False, + run_root=Path(directory), + total_batches=True, + ) + with self.assertRaises(ValueError): + pipeline.create_run( + "alpha", + dry_run=False, + run_root=Path(directory), + total_batches=-1, + ) diff --git a/tests/etl/test_examples.py b/tests/etl/test_examples.py new file mode 100644 index 0000000..c3665b1 --- /dev/null +++ b/tests/etl/test_examples.py @@ -0,0 +1,144 @@ +"""Tests for the runnable ETL dry-run examples.""" + +import io +import json +import os +import tempfile +import unittest +from collections.abc import Iterator +from contextlib import contextmanager, redirect_stdout +from pathlib import Path + +from examples.etl import batch_local_artifacts, local_artifact + + +@contextmanager +def _temporary_cwd(path: Path) -> Iterator[None]: + previous = Path.cwd() + os.chdir(path) + try: + yield + finally: + os.chdir(previous) + + +async def _run_without_stdout(coro): + with redirect_stdout(io.StringIO()): + return await coro + + +def _read_run_snapshot(root: Path) -> dict[str, object]: + snapshots = sorted( + (root / ".quant-mind" / "etl-pipeline-runs").glob("*/run.json") + ) + if len(snapshots) != 1: + raise AssertionError(f"expected one run snapshot, found {snapshots}") + return json.loads(snapshots[0].read_text(encoding="utf-8")) + + +class ETLExampleTests(unittest.IsolatedAsyncioTestCase): + async def test_whole_run_example_dry_run_plans_without_artifact( + self, + ) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + with _temporary_cwd(root): + result = await _run_without_stdout( + local_artifact.run_example(dry_run=True) + ) + + artifact = root / ".quant-mind" / "etl-example" / "artifact.json" + self.assertFalse(artifact.exists()) + self.assertTrue(result["dry_run"]) + self.assertEqual(result["planned_artifacts"], 1) + self.assertIn("planned_artifact", result) + + snapshot = _read_run_snapshot(root) + self.assertTrue(snapshot["dry_run"]) + self.assertEqual(snapshot["state"], "succeeded") + self.assertEqual(snapshot["stage"], "load") + self.assertEqual( + snapshot["progress"], + { + "completed": 1, + "total": 1, + "message": "Planned local artifact", + "metrics": { + "planned_artifacts": 1, + "planned_bytes": result["planned_bytes"], + }, + }, + ) + + async def test_whole_run_example_normal_creates_artifact(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + with _temporary_cwd(root): + result = await _run_without_stdout( + local_artifact.run_example(dry_run=False) + ) + + artifact = root / ".quant-mind" / "etl-example" / "artifact.json" + self.assertTrue(artifact.exists()) + self.assertFalse(result["dry_run"]) + self.assertEqual( + Path(str(result["artifact"])).resolve(), artifact.resolve() + ) + self.assertEqual(result["artifacts_written"], 1) + + snapshot = _read_run_snapshot(root) + self.assertFalse(snapshot["dry_run"]) + self.assertEqual(snapshot["state"], "succeeded") + + async def test_batch_example_dry_run_plans_without_artifacts(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + with _temporary_cwd(root): + result = await _run_without_stdout( + batch_local_artifacts.run_example(dry_run=True) + ) + + artifact_dir = root / ".quant-mind" / "etl-batch-example" + self.assertFalse(artifact_dir.exists()) + self.assertTrue(result["dry_run"]) + self.assertEqual(result["completed_batches"], 2) + self.assertEqual( + result["counts"], + { + "planned_artifacts": 2, + "planned_records": 3, + }, + ) + + snapshot = _read_run_snapshot(root) + self.assertTrue(snapshot["dry_run"]) + self.assertEqual(snapshot["state"], "succeeded") + self.assertEqual( + snapshot["batch"], + {"index": None, "completed": 2, "total": 2}, + ) + + async def test_batch_example_normal_creates_artifacts(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + with _temporary_cwd(root): + result = await _run_without_stdout( + batch_local_artifacts.run_example(dry_run=False) + ) + + artifact_dir = root / ".quant-mind" / "etl-batch-example" + self.assertTrue((artifact_dir / "batch-001.json").exists()) + self.assertTrue((artifact_dir / "batch-002.json").exists()) + self.assertFalse(result["dry_run"]) + self.assertEqual(result["completed_batches"], 2) + self.assertEqual( + result["counts"], + { + "artifacts_written": 2, + "records_written": 3, + }, + ) + + snapshot = _read_run_snapshot(root) + self.assertFalse(snapshot["dry_run"]) + self.assertEqual(snapshot["state"], "succeeded") diff --git a/tests/etl/test_pipeline.py b/tests/etl/test_pipeline.py new file mode 100644 index 0000000..0b2dc63 --- /dev/null +++ b/tests/etl/test_pipeline.py @@ -0,0 +1,784 @@ +"""Tests for the observable ETL pipeline scaffold.""" + +import asyncio +import json +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import quantmind.etl._record as record_module +from quantmind.etl import ETLPipeline, PipelineContext + + +def _read_json(path: Path) -> dict[str, object]: + return json.loads(path.read_text(encoding="utf-8")) + + +def _read_events(path: Path) -> list[dict[str, object]]: + return [ + json.loads(line) + for line in path.read_text(encoding="utf-8").splitlines() + ] + + +def _json_safe_text(value: str) -> str: + return value.encode("utf-8", errors="backslashreplace").decode("utf-8") + + +class PipelineRunTests(unittest.IsolatedAsyncioTestCase): + def setUp(self) -> None: + self._temporary_directory = tempfile.TemporaryDirectory() + self.run_root = Path(self._temporary_directory.name) + + def tearDown(self) -> None: + self._temporary_directory.cleanup() + + @staticmethod + def _pipeline( + *, + extract=None, + transform=None, + load=None, + ) -> ETLPipeline[str, str, str, str]: + async def default_extract(source: str, *, ctx: PipelineContext) -> str: + return source + + async def default_transform( + extracted: str, *, ctx: PipelineContext + ) -> str: + return extracted.upper() + + async def default_load( + transformed: str, *, ctx: PipelineContext + ) -> str: + return f"loaded:{transformed}" + + return ETLPipeline( + "test-pipeline", + extract=extract or default_extract, + transform=transform or default_transform, + load=load or default_load, + ) + + async def test_create_run_writes_created_snapshot_before_return( + self, + ) -> None: + source = "private-source" + run = self._pipeline().create_run( + source, + dry_run=True, + run_root=self.run_root, + config_summary={"window_days": 30}, + ) + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["schema"], "quantmind.etl.run/v1") + self.assertTrue(snapshot["dry_run"]) + self.assertTrue(run.status_file.is_absolute()) + self.assertEqual(snapshot["state"], "created") + self.assertIsNone(snapshot["stage"]) + self.assertIsNone(snapshot["pid"]) + self.assertEqual(snapshot["config_summary"], {"window_days": 30}) + self.assertNotIn(source, run.status_file.read_text(encoding="utf-8")) + self.assertEqual(run.events_file.read_text(encoding="utf-8"), "") + + async def test_success_runs_fixed_stages_and_records_lifecycle( + self, + ) -> None: + calls: list[str] = [] + + async def extract(source: str, *, ctx: PipelineContext) -> str: + calls.append(ctx.stage) + self.assertIsNone(ctx.batch_index) + self.assertFalse(ctx.dry_run) + await ctx.progress(1, total=1, message="read") + return source + + async def transform(value: str, *, ctx: PipelineContext) -> str: + calls.append(ctx.stage) + self.assertIsNone(ctx.batch_index) + self.assertFalse(ctx.dry_run) + return value.upper() + + async def load(value: str, *, ctx: PipelineContext) -> str: + calls.append(ctx.stage) + self.assertIsNone(ctx.batch_index) + self.assertFalse(ctx.dry_run) + return f"loaded:{value}" + + run = self._pipeline( + extract=extract, + transform=transform, + load=load, + ).create_run("alpha", dry_run=False, run_root=self.run_root) + + result = await run.execute() + + self.assertEqual(result, "loaded:ALPHA") + self.assertEqual(calls, ["extract", "transform", "load"]) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "succeeded") + self.assertFalse(snapshot["dry_run"]) + self.assertEqual(snapshot["stage"], "load") + self.assertEqual(snapshot["pid"], os.getpid()) + self.assertIsNone(snapshot["error"]) + self.assertNotIn(result, run.status_file.read_text(encoding="utf-8")) + self.assertEqual( + [event["event"] for event in _read_events(run.events_file)], + [ + "run_started", + "stage_started", + "stage_progress", + "stage_completed", + "stage_started", + "stage_completed", + "stage_started", + "stage_completed", + "run_succeeded", + ], + ) + + async def test_every_event_records_schema_and_dry_run(self) -> None: + async def extract(source: str, *, ctx: PipelineContext) -> str: + await ctx.progress(1, total=1, message="read") + return source + + run = self._pipeline(extract=extract).create_run( + "alpha", dry_run=True, run_root=self.run_root + ) + + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["schema"], "quantmind.etl.run/v1") + self.assertTrue(snapshot["dry_run"]) + events = _read_events(run.events_file) + self.assertGreater(len(events), 0) + for event in events: + self.assertEqual(event["schema"], "quantmind.etl.event/v1") + self.assertEqual(event["run_id"], run.id) + self.assertTrue(event["dry_run"]) + + async def test_dry_run_reaches_every_stage_and_still_calls_load( + self, + ) -> None: + calls: list[tuple[str, bool]] = [] + + async def extract(source: str, *, ctx: PipelineContext) -> str: + calls.append((ctx.stage, ctx.dry_run)) + return source + + async def transform(value: str, *, ctx: PipelineContext) -> str: + calls.append((ctx.stage, ctx.dry_run)) + return value.upper() + + async def load(value: str, *, ctx: PipelineContext) -> str: + calls.append((ctx.stage, ctx.dry_run)) + await ctx.progress( + 1, + total=1, + metrics={"planned_records": 1}, + ) + return f"planned:{value}" + + run = self._pipeline( + extract=extract, + transform=transform, + load=load, + ).create_run("alpha", dry_run=True, run_root=self.run_root) + + result = await run.execute() + + self.assertEqual(result, "planned:ALPHA") + self.assertEqual( + calls, + [("extract", True), ("transform", True), ("load", True)], + ) + snapshot = _read_json(run.status_file) + self.assertTrue(snapshot["dry_run"]) + self.assertEqual( + snapshot["progress"], + { + "completed": 1, + "total": 1, + "message": None, + "metrics": {"planned_records": 1}, + }, + ) + + async def test_json_records_escape_surrogate_strings_recursively( + self, + ) -> None: + unsafe = os.fsdecode(b"\xff") + safe = _json_safe_text(unsafe) + + async def load(value: str, *, ctx: PipelineContext) -> str: + await ctx.progress( + 1, + total=1, + message=unsafe, + metrics={unsafe: unsafe}, + ) + return value + + run = self._pipeline(load=load).create_run( + "alpha", + dry_run=False, + run_root=self.run_root, + config_summary={unsafe: unsafe}, + ) + + await run.execute() + + snapshot_text = run.status_file.read_text(encoding="utf-8") + snapshot = json.loads(snapshot_text) + self.assertEqual(snapshot["state"], "succeeded") + self.assertEqual(snapshot["config_summary"], {safe: safe}) + self.assertEqual(snapshot["progress"]["message"], safe) + self.assertEqual(snapshot["progress"]["metrics"], {safe: safe}) + progress_events = [ + event + for event in _read_events(run.events_file) + if event["event"] == "stage_progress" + ] + self.assertEqual(progress_events[-1]["progress"]["message"], safe) + self.assertEqual( + progress_events[-1]["progress"]["metrics"], {safe: safe} + ) + + async def test_failure_after_surrogate_progress_records_failed_json( + self, + ) -> None: + unsafe = os.fsdecode(b"\xff") + safe = _json_safe_text(unsafe) + + async def transform(value: str, *, ctx: PipelineContext) -> str: + await ctx.progress( + 1, + total=1, + message=unsafe, + metrics={unsafe: unsafe}, + ) + raise LookupError("stage failed") + + run = self._pipeline(transform=transform).create_run( + "alpha", + dry_run=False, + run_root=self.run_root, + config_summary={unsafe: unsafe}, + ) + + with self.assertRaisesRegex(LookupError, "stage failed"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["config_summary"], {safe: safe}) + self.assertEqual(snapshot["progress"]["message"], safe) + self.assertEqual(snapshot["progress"]["metrics"], {safe: safe}) + self.assertEqual(snapshot["error"]["type"], "LookupError") + self.assertEqual( + _read_events(run.events_file)[-1]["event"], "run_failed" + ) + + async def test_normal_and_dry_runs_are_independent(self) -> None: + async def load(value: str, *, ctx: PipelineContext) -> str: + if ctx.dry_run: + return f"planned:{value}" + return f"loaded:{value}" + + pipeline = self._pipeline(load=load) + normal = pipeline.create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + dry = pipeline.create_run("beta", dry_run=True, run_root=self.run_root) + + normal_result = await normal.execute() + dry_result = await dry.execute() + + self.assertNotEqual(normal.id, dry.id) + self.assertFalse(normal.dry_run) + self.assertTrue(dry.dry_run) + self.assertEqual(normal_result, "loaded:ALPHA") + self.assertEqual(dry_result, "planned:BETA") + normal_snapshot = _read_json(normal.status_file) + dry_snapshot = _read_json(dry.status_file) + self.assertEqual(normal_snapshot["state"], "succeeded") + self.assertEqual(dry_snapshot["state"], "succeeded") + self.assertFalse(normal_snapshot["dry_run"]) + self.assertTrue(dry_snapshot["dry_run"]) + + async def test_failure_records_limited_error_and_reraises(self) -> None: + async def fail(value: str, *, ctx: PipelineContext) -> str: + raise LookupError("missing item") + + run = self._pipeline(transform=fail).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(LookupError, "missing item"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["stage"], "transform") + self.assertEqual( + snapshot["error"], + {"message": "missing item", "type": "LookupError"}, + ) + self.assertEqual( + _read_events(run.events_file)[-1]["event"], "run_failed" + ) + + async def test_surrogate_error_message_records_json_safely( + self, + ) -> None: + class SurrogateMessageError(Exception): + def __str__(self) -> str: + return "bad value \ud800" + + original_error = SurrogateMessageError() + + async def fail(value: str, *, ctx: PipelineContext) -> str: + raise original_error + + run = self._pipeline(transform=fail).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with self.assertRaises(SurrogateMessageError) as raised: + await run.execute() + + self.assertIs(raised.exception, original_error) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual( + snapshot["error"], + {"message": "bad value \\ud800", "type": "SurrogateMessageError"}, + ) + self.assertEqual( + _read_events(run.events_file)[-1]["error"], + {"message": "bad value \\ud800", "type": "SurrogateMessageError"}, + ) + + async def test_failed_terminal_event_append_does_not_mask_stage_error( + self, + ) -> None: + original_append = record_module._append_json_line + + def append_event(path: Path, value: dict[str, object]) -> None: + if value["event"] == "run_failed": + raise OSError("event sink unavailable") + original_append(path, value) + + async def fail(value: str, *, ctx: PipelineContext) -> str: + raise LookupError("stage failed") + + run = self._pipeline(transform=fail).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with patch( + "quantmind.etl._record._append_json_line", + side_effect=append_event, + ): + with self.assertRaisesRegex(LookupError, "stage failed"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["error"]["type"], "LookupError") + + async def test_run_started_event_append_failure_does_not_leave_running( + self, + ) -> None: + original_append = record_module._append_json_line + + def append_event(path: Path, value: dict[str, object]) -> None: + if value["event"] == "run_started": + raise OSError("event sink unavailable") + original_append(path, value) + + async def fail(value: str, *, ctx: PipelineContext) -> str: + raise LookupError("stage failed") + + run = self._pipeline(transform=fail).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with patch( + "quantmind.etl._record._append_json_line", + side_effect=append_event, + ): + with self.assertRaisesRegex(LookupError, "stage failed"): + await run.execute() + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["error"]["type"], "LookupError") + self.assertNotEqual(snapshot["state"], "running") + + async def test_failed_terminal_snapshot_write_does_not_mask_stage_error( + self, + ) -> None: + original_write = record_module._write_json_atomic + original_error = LookupError("stage failed") + + def write_snapshot(path: Path, value: dict[str, object]) -> None: + if value["state"] == "failed": + raise OSError("snapshot unavailable") + original_write(path, value) + + async def fail(value: str, *, ctx: PipelineContext) -> str: + raise original_error + + run = self._pipeline(transform=fail).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with patch( + "quantmind.etl._record._write_json_atomic", + side_effect=write_snapshot, + ): + with self.assertRaises(LookupError) as raised: + await run.execute() + + self.assertIs(raised.exception, original_error) + snapshot = _read_json(run.status_file) + self.assertNotEqual(snapshot["state"], "succeeded") + + async def test_failed_terminal_snapshot_survives_pending_progress_event_error( + self, + ) -> None: + original_append = record_module._append_json_line + original_error = LookupError("stage failed") + + def append_event(path: Path, value: dict[str, object]) -> None: + if ( + value["event"] == "stage_progress" + and value["progress"]["completed"] == 2 + ): + raise OSError("progress journal unavailable") + original_append(path, value) + + async def fail_after_pending_progress( + value: str, *, ctx: PipelineContext + ) -> str: + await ctx.progress(1, total=3) + await ctx.progress(2, total=3) + raise original_error + + run = self._pipeline(extract=fail_after_pending_progress).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with ( + patch( + "quantmind.etl._record._append_json_line", + side_effect=append_event, + ), + patch( + "quantmind.etl._record._monotonic_seconds", + side_effect=[10.0, 10.1], + ), + ): + with self.assertRaises(LookupError) as raised: + await run.execute() + + self.assertIs(raised.exception, original_error) + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "failed") + self.assertEqual(snapshot["error"]["type"], "LookupError") + + async def test_cancelled_terminal_snapshot_write_preserves_cancelled_error( + self, + ) -> None: + original_write = record_module._write_json_atomic + entered = asyncio.Event() + + def write_snapshot(path: Path, value: dict[str, object]) -> None: + if value["state"] == "cancelled": + raise OSError("snapshot unavailable") + original_write(path, value) + + async def wait_forever(value: str, *, ctx: PipelineContext) -> str: + entered.set() + await asyncio.Event().wait() + return value + + run = self._pipeline(extract=wait_forever).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + task = asyncio.create_task(run.execute()) + await entered.wait() + + with patch( + "quantmind.etl._record._write_json_atomic", + side_effect=write_snapshot, + ): + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + snapshot = _read_json(run.status_file) + self.assertNotEqual(snapshot["state"], "succeeded") + + async def test_succeeded_snapshot_makes_terminal_event_best_effort( + self, + ) -> None: + original_append = record_module._append_json_line + + def append_event(path: Path, value: dict[str, object]) -> None: + if value["event"] == "run_succeeded": + raise OSError("event sink unavailable") + original_append(path, value) + + run = self._pipeline().create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with patch( + "quantmind.etl._record._append_json_line", + side_effect=append_event, + ): + result = await run.execute() + + self.assertEqual(result, "loaded:ALPHA") + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "succeeded") + self.assertEqual( + [event["event"] for event in _read_events(run.events_file)][-1], + "stage_completed", + ) + + async def test_cancellation_records_terminal_state_then_reraises( + self, + ) -> None: + entered = asyncio.Event() + + async def wait_forever(value: str, *, ctx: PipelineContext) -> str: + entered.set() + await asyncio.Event().wait() + return value + + run = self._pipeline(extract=wait_forever).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + task = asyncio.create_task(run.execute()) + await entered.wait() + + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + snapshot = _read_json(run.status_file) + self.assertEqual(snapshot["state"], "cancelled") + self.assertEqual(snapshot["stage"], "extract") + self.assertIsNone(snapshot["error"]) + self.assertEqual( + _read_events(run.events_file)[-1]["event"], "run_cancelled" + ) + + async def test_receipt_is_one_line_json_with_absolute_status_file( + self, + ) -> None: + run = self._pipeline().create_run( + "alpha", dry_run=True, run_root=self.run_root + ) + + receipt = run.receipt() + parsed = json.loads(receipt) + self.assertNotIn("\n", receipt) + self.assertEqual(parsed["event"], "etl_run_created") + self.assertEqual(parsed["run_id"], run.id) + self.assertTrue(parsed["dry_run"]) + self.assertEqual(parsed["status_file"], str(run.status_file)) + self.assertTrue(Path(parsed["status_file"]).is_absolute()) + + async def test_new_stage_clears_previous_stage_progress(self) -> None: + run_holder = [] + observed: list[tuple[str, object]] = [] + + async def extract(source: str, *, ctx: PipelineContext) -> str: + await ctx.progress(2, total=2) + return source + + async def transform(value: str, *, ctx: PipelineContext) -> str: + snapshot = _read_json(run_holder[0].status_file) + observed.append((str(snapshot["stage"]), snapshot["progress"])) + await ctx.progress(1, total=None) + return value + + async def load(value: str, *, ctx: PipelineContext) -> str: + snapshot = _read_json(run_holder[0].status_file) + observed.append((str(snapshot["stage"]), snapshot["progress"])) + return value + + run = self._pipeline( + extract=extract, + transform=transform, + load=load, + ).create_run("alpha", dry_run=False, run_root=self.run_root) + run_holder.append(run) + + await run.execute() + + self.assertEqual(observed, [("transform", None), ("load", None)]) + + async def test_progress_requires_strict_real_completion(self) -> None: + async def invalid(source: str, *, ctx: PipelineContext) -> str: + await ctx.progress(1, total=3) + await ctx.progress(1, total=3) + return source + + run = self._pipeline(extract=invalid).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(ValueError, "strictly increase"): + await run.execute() + self.assertEqual(_read_json(run.status_file)["state"], "failed") + + async def test_progress_rejects_completed_above_total(self) -> None: + async def invalid(source: str, *, ctx: PipelineContext) -> str: + await ctx.progress(2, total=1) + return source + + run = self._pipeline(extract=invalid).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(ValueError, "total must be >= completed"): + await run.execute() + + async def test_progress_total_none_means_unknown_until_known(self) -> None: + async def invalid(source: str, *, ctx: PipelineContext) -> str: + await ctx.progress(1, total=None) + await ctx.progress(2, total=3) + await ctx.progress(3, total=None) + return source + + run = self._pipeline(extract=invalid).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with self.assertRaisesRegex(ValueError, "cannot become unknown"): + await run.execute() + + async def test_stage_child_tasks_can_report_progress(self) -> None: + async def extract(source: str, *, ctx: PipelineContext) -> str: + async def first_worker() -> None: + await ctx.progress(1, total=2) + + async def second_worker() -> None: + await asyncio.sleep(0) + await ctx.progress(2, total=2) + + await asyncio.gather(first_worker(), second_worker()) + return source + + run = self._pipeline(extract=extract).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + await run.execute() + + progress_events = [ + event + for event in _read_events(run.events_file) + if event["event"] == "stage_progress" + ] + self.assertEqual( + [event["progress"]["completed"] for event in progress_events], + [1, 2], + ) + + async def test_progress_events_are_coalesced_but_snapshot_is_current( + self, + ) -> None: + async def extract(source: str, *, ctx: PipelineContext) -> str: + await ctx.progress(1, total=3) + await ctx.progress(2, total=3) + await ctx.progress(3, total=3) + return source + + run = self._pipeline(extract=extract).create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + + with patch( + "quantmind.etl._record._monotonic_seconds", + side_effect=[10.0, 10.1, 10.2], + ): + await run.execute() + + progress_events = [ + event + for event in _read_events(run.events_file) + if event["event"] == "stage_progress" + ] + self.assertEqual(len(progress_events), 2) + self.assertEqual( + [event["progress"]["completed"] for event in progress_events], + [1, 3], + ) + + async def test_execute_is_one_shot(self) -> None: + run = self._pipeline().create_run( + "alpha", dry_run=False, run_root=self.run_root + ) + await run.execute() + + with self.assertRaisesRegex(RuntimeError, "only be called once"): + await run.execute() + + +class PipelineValidationTests(unittest.TestCase): + def test_pipeline_name_must_not_be_empty(self) -> None: + async def stage(value, *, ctx): + return value + + with self.assertRaises(ValueError): + ETLPipeline("", extract=stage, transform=stage, load=stage) + + def test_config_summary_is_a_json_scalar_allowlist(self) -> None: + async def stage(value, *, ctx): + return value + + pipeline = ETLPipeline( + "safe-summary", extract=stage, transform=stage, load=stage + ) + with tempfile.TemporaryDirectory() as directory: + with self.assertRaisesRegex(TypeError, "dry_run"): + pipeline.create_run("alpha", run_root=Path(directory)) + with self.assertRaisesRegex(TypeError, "dry_run must be"): + pipeline.create_run( + "alpha", + dry_run=None, + run_root=Path(directory), + ) + with self.assertRaisesRegex( + ValueError, r"create_run\(..., dry_run=...\)" + ): + pipeline.create_run( + "alpha", + dry_run=False, + run_root=Path(directory), + config_summary={"dry_run": True}, + ) + with self.assertRaises(TypeError): + pipeline.create_run( + "alpha", + dry_run=False, + run_root=Path(directory), + config_summary={"secret": {"nested": "value"}}, + ) + with self.assertRaises(ValueError): + pipeline.create_run( + "alpha", + dry_run=False, + run_root=Path(directory), + config_summary={"ratio": float("nan")}, + ) From efde686214539207be456d6e0e41c05e9f171b0b Mon Sep 17 00:00:00 2001 From: pkuwkl Date: Fri, 14 Aug 2026 16:54:22 +0800 Subject: [PATCH 2/2] docs(etl): clarify staging delivery boundary Document that staging stays inside its owning extract or transform callable, remains idempotent and observable, and is suppressed during dry-run rather than becoming a fourth framework stage. --- contexts/design/operations/etl.md | 4 +++- docs/etl.md | 28 ++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/contexts/design/operations/etl.md b/contexts/design/operations/etl.md index c55be37..d3fb8a5 100644 --- a/contexts/design/operations/etl.md +++ b/contexts/design/operations/etl.md @@ -101,9 +101,11 @@ Inputs, complete configuration, intermediate values, individual load results, tr ## Distinguish staging from delivery +Staging is intentionally not a fourth framework stage, and the scaffold exposes no `stage()` or `staging()` API. The authored `extract` or `transform` operation that creates and owns an intermediate also owns any staging write needed to recover or reuse it. + `load` marks the delivery boundary. In a normal run, after it succeeds the pipeline's intended downstream consumer may treat that output as delivered; in dry-run it validates or plans the same boundary without delivering. A database call is not automatically a load. Extract or transform may perform staging writes when they persist an intermediate for recovery or reuse while keeping it unavailable to formal downstream consumers. -Staging writes are permitted in any stage of a normal run when they are idempotent and their real completion is visible through progress. Dry-run must plan or validate them without writing. They remain business behavior: the scaffold provides no transaction, rollback, exactly-once, or artifact-management guarantee. +Staging writes are permitted inside the owning extract or transform stage of a normal run when they are idempotent and their real completion is visible through progress. Dry-run must plan or validate them without writing. They remain business behavior: the scaffold provides no transaction, rollback, exactly-once, or artifact-management guarantee. If a batch write makes the final product consumable, it is a real batch load even when the target table is named `raw`. Repeated `transform(batch) → load(batch)` delivery belongs in `BatchETLPipeline`; hiding those loads inside a whole-run transform would make `run.json.stage` dishonest. diff --git a/docs/etl.md b/docs/etl.md index 7ae9901..12ae961 100644 --- a/docs/etl.md +++ b/docs/etl.md @@ -109,6 +109,34 @@ The framework owns a strictly serial loop. While it awaits the next yielded batc See [batch local artifacts](../examples/etl/batch_local_artifacts.py) for a network-free, bounded-memory example with idempotent local loads. +## Handle staging inside the owning stage + +Decide whether a write is staging or load by downstream consumability, not by whether it touches a database. A staging write persists an intermediate for recovery or reuse while keeping it unavailable to formal downstream consumers; `load` is the boundary after which the intended consumer may treat the result as delivered. + +Staging is not a fourth framework stage. Put it inside the `extract` or `transform` callable that creates and owns the intermediate, make the write idempotent, report its real completion through `ctx.progress()`, and suppress the mutation during dry-run. The framework observes the authored stage but does not manage staging artifacts, transactions, rollback, or recovery. + +```python +async def extract(source: str, *, ctx: PipelineContext) -> list[str]: + raw_records = await fetch_raw_records(source) + if ctx.dry_run: + await validate_raw_records(raw_records) + metrics = {"planned_raw_records": len(raw_records)} + message = "raw staging planned" + else: + await raw_store.upsert_many(raw_records) # Idempotent staging write. + metrics = {"raw_records_staged": len(raw_records)} + message = "raw records staged" + await ctx.progress( + len(raw_records), + total=len(raw_records), + message=message, + metrics=metrics, + ) + return raw_records +``` + +If that write already makes a batch consumable, it is a real `load(batch)` even when the table or object is named `raw`. If an intermediate has an independent downstream consumer, model it as the delivered output of a separate pipeline. See [Distinguish staging from delivery](../contexts/design/operations/etl.md#distinguish-staging-from-delivery) for the canonical decision rule. + ## Implement dry-run honestly `create_run(source, *, dry_run=...)` has no default. Production scripts should read the value from a runtime option such as `--dry-run` and pass that variable, so switching modes never requires editing stage code.