Skip to main content

faucet_core/
pipeline.rs

1//! Source-to-sink pipeline orchestration.
2//!
3//! The [`Pipeline`] struct connects any [`Source`] to any
4//! [`Sink`] and handles moving data between them.
5//!
6//! # Batch mode
7//!
8//! Fetches all records from the source, then writes them to the sink in one
9//! shot.  Supports incremental replication (returns a bookmark for the next
10//! run).
11//!
12//! ```rust,no_run
13//! use faucet_core::{Pipeline, Source, Sink};
14//! # async fn example(source: impl Source, sink: impl Sink) -> Result<(), faucet_core::FaucetError> {
15//! let result = Pipeline::new(&source, &sink).run().await?;
16//! println!("wrote {} records", result.records_written);
17//! // Persist result.bookmark for the next incremental run
18//! # Ok(())
19//! # }
20//! ```
21//!
22//! # Streaming mode
23//!
24//! Writes records page-by-page as they arrive from a source's
25//! [`stream_pages`](crate::Source::stream_pages) implementation, keeping
26//! memory usage bounded.  [`Pipeline::run`] uses this internally; callers
27//! that have already assembled a [`StreamPage`] stream can drive it directly
28//! via [`run_stream`].
29//!
30//! ```rust,no_run
31//! use faucet_core::{run_stream, RunStreamOptions, Sink, StreamPage, FaucetError};
32//! use futures_core::Stream;
33//! # async fn example(
34//! #     pages: impl Stream<Item = Result<StreamPage, FaucetError>> + Unpin,
35//! #     sink: impl Sink,
36//! # ) -> Result<(), FaucetError> {
37//! let result = run_stream(pages, &sink, RunStreamOptions::new()).await?;
38//! # Ok(())
39//! # }
40//! ```
41
42use crate::dlq::{DlqConfig, DlqStats};
43use crate::error::FaucetError;
44use crate::observability::RunStreamOptions;
45use crate::state::{StateStore, validate_state_key};
46use crate::traits::{Sink, Source};
47use futures_core::Stream;
48use serde_json::Value;
49use std::pin::Pin;
50use std::sync::Arc;
51
52/// Default page size used when a caller does not specify one.
53///
54/// Sources are free to override this from their own config when implementing
55/// [`Source::stream_pages`]; the value passed
56/// from the pipeline acts as a hint when no source-side preference exists.
57pub const DEFAULT_BATCH_SIZE: usize = 1000;
58
59/// Hard upper bound on `batch_size`. Values above this (other than the
60/// special `0` "no batching" sentinel) are rejected at config validation
61/// time to prevent accidental O(total) buffering in the default
62/// implementation of [`Source::stream_pages`].
63pub const MAX_BATCH_SIZE: usize = 1_000_000;
64
65/// Validate a `batch_size` value against the global constraints.
66///
67/// `batch_size = 0` is the **opt-out-of-batching sentinel**: sources and
68/// sinks should treat it as "emit / accept the entire result set in one
69/// page." This is useful for small lookup tables or for sinks (e.g. SQL
70/// `COPY`, BigQuery load jobs) that prefer one large request to many small
71/// ones. Any non-zero value above [`MAX_BATCH_SIZE`] is rejected to prevent
72/// accidental unbounded buffering through a typo.
73///
74/// Returns the unchanged value on success. Returns `FaucetError::Config`
75/// only for values strictly greater than [`MAX_BATCH_SIZE`].
76pub fn validate_batch_size(batch_size: usize) -> Result<usize, FaucetError> {
77    if batch_size > MAX_BATCH_SIZE {
78        return Err(FaucetError::Config(format!(
79            "batch_size {batch_size} exceeds maximum {MAX_BATCH_SIZE} \
80             (use 0 to opt out of batching entirely)"
81        )));
82    }
83    Ok(batch_size)
84}
85
86/// One page emitted by [`Source::stream_pages`].
87///
88/// `records` is the chunk of records for this page. `bookmark` is `Some` only
89/// when the source has a durable checkpoint to advance — most sources emit
90/// `Some` only on the final page (max-replication-value semantics); CDC-style
91/// sources emit `Some` per committed transaction. The pipeline flushes the
92/// sink and persists the bookmark every time a page carries one, so a
93/// mid-stream crash never advances past records the sink has not durably
94/// written.
95#[derive(Debug, Clone, Default)]
96pub struct StreamPage {
97    /// Records to write to the sink for this page.
98    pub records: Vec<Value>,
99    /// Optional bookmark to checkpoint after this page is durably written.
100    pub bookmark: Option<Value>,
101}
102
103/// Result of a pipeline run.
104#[derive(Debug, Clone)]
105pub struct PipelineResult {
106    /// Total number of records written to the sink.
107    pub records_written: usize,
108    /// Bookmark value for incremental replication.
109    ///
110    /// `Some(value)` when the source returned a bookmark on its final
111    /// (or, for streaming CDC sources, most recent) page. Persist this and
112    /// pass it back as `start_replication_value` on the next run; this is
113    /// handled automatically when a [`StateStore`] is attached via
114    /// [`Pipeline::with_state_store`].
115    pub bookmark: Option<Value>,
116    /// DLQ counters. `None` when no DLQ is configured.
117    pub dlq: Option<DlqStats>,
118}
119
120/// A pipeline that moves data from a [`Source`] to a [`Sink`].
121///
122/// The pipeline is generic over the source and sink types — any combination
123/// of connectors works as long as they implement the respective traits.
124pub struct Pipeline<'a, So: Source + ?Sized, Si: Sink + ?Sized> {
125    source: &'a So,
126    sink: &'a Si,
127    state_store: Option<Arc<dyn StateStore>>,
128    name: Option<String>,
129    row: Option<String>,
130    run_id: Option<String>,
131    dlq: Option<DlqConfig>,
132    #[cfg(feature = "quality")]
133    quality: Option<Arc<crate::quality::CompiledQuality>>,
134    adaptive: Option<crate::adaptive::AdaptiveBatchConfig>,
135    cancel: Option<tokio_util::sync::CancellationToken>,
136    delivery: crate::idempotency::DeliveryMode,
137}
138
139impl<'a, So: Source + ?Sized, Si: Sink + ?Sized> Pipeline<'a, So, Si> {
140    /// Create a new pipeline from a source and a sink.
141    pub fn new(source: &'a So, sink: &'a Si) -> Self {
142        Self {
143            source,
144            sink,
145            state_store: None,
146            name: None,
147            row: None,
148            run_id: None,
149            dlq: None,
150            #[cfg(feature = "quality")]
151            quality: None,
152            adaptive: None,
153            cancel: None,
154            delivery: crate::idempotency::DeliveryMode::AtLeastOnce,
155        }
156    }
157
158    /// Attach a [`StateStore`] for persistent incremental-replication bookmarks.
159    ///
160    /// When configured, `run()` will:
161    /// 1. Read any previously stored bookmark at the source's
162    ///    [`state_key`](Source::state_key) and call
163    ///    [`apply_start_bookmark`](Source::apply_start_bookmark) on the source
164    ///    so it can resume from that point.
165    /// 2. Run the fetch + write as usual.
166    /// 3. Persist the new bookmark **only after** the sink confirms the
167    ///    batch was written and flushed.
168    ///
169    /// Sources that do not return a [`state_key`](Source::state_key) are
170    /// unaffected — the store is consulted only when the source opts in.
171    pub fn with_state_store(mut self, store: Arc<dyn StateStore>) -> Self {
172        self.state_store = Some(store);
173        self
174    }
175
176    /// Set the pipeline name used in spans and metric labels.
177    /// Defaults to `"unnamed"` when unset.
178    pub fn with_name(mut self, name: impl Into<String>) -> Self {
179        self.name = Some(name.into());
180        self
181    }
182
183    /// Set the matrix row id used in spans and metric labels.
184    /// Defaults to `""` (Prometheus treats empty labels as absent).
185    pub fn with_row(mut self, row: impl Into<String>) -> Self {
186        self.row = Some(row.into());
187        self
188    }
189
190    /// Set an explicit run id (UUIDv7-shaped). When unset, `Pipeline::run`
191    /// generates one. Used only as a tracing span attribute — never a metric
192    /// label.
193    pub fn with_run_id(mut self, run_id: impl Into<String>) -> Self {
194        self.run_id = Some(run_id.into());
195        self
196    }
197
198    /// Attach a DLQ for per-row failure routing.
199    pub fn with_dlq(mut self, dlq: DlqConfig) -> Self {
200        self.dlq = Some(dlq);
201        self
202    }
203
204    /// Attach a compiled quality spec. Checks run after transforms, before the
205    /// sink, per page.
206    #[cfg(feature = "quality")]
207    pub fn with_quality(mut self, quality: Arc<crate::quality::CompiledQuality>) -> Self {
208        self.quality = Some(quality);
209        self
210    }
211
212    /// Attach an adaptive batch-size controller (opt-in). When `enabled`, the
213    /// pipeline reslices each source page into sub-batches whose size the
214    /// controller tunes from observed sink latency + error rate.
215    pub fn with_adaptive(mut self, cfg: crate::adaptive::AdaptiveBatchConfig) -> Self {
216        self.adaptive = Some(cfg);
217        self
218    }
219
220    /// Attach a cancellation token. When cancelled mid-run, the streaming loop
221    /// stops at the next page boundary, flushes the sink(s) so buffered output
222    /// (e.g. a Parquet footer) is durable, and returns the partial result
223    /// instead of leaving the file unreadable (#146 H16).
224    pub fn with_cancel(mut self, cancel: tokio_util::sync::CancellationToken) -> Self {
225        self.cancel = Some(cancel);
226        self
227    }
228
229    /// Set the delivery guarantee. `ExactlyOnce` requires a state store, an
230    /// idempotent sink (`Sink::supports_idempotent_writes`), and a
231    /// deterministic-replay source — otherwise `run` returns
232    /// `FaucetError::Config`.
233    pub fn with_delivery(mut self, mode: crate::idempotency::DeliveryMode) -> Self {
234        self.delivery = mode;
235        self
236    }
237
238    /// Run the pipeline in streaming mode.
239    ///
240    /// 1. Loads the stored bookmark and pushes it to the source (if a state
241    ///    store is configured and the source returns a `state_key`).
242    /// 2. Drives [`Source::stream_pages`] with [`DEFAULT_BATCH_SIZE`],
243    ///    writing each page to the sink as it arrives via
244    ///    [`Sink::write_batch`].
245    /// 3. Whenever a page carries `Some(bookmark)`, flushes the sink and
246    ///    persists the bookmark to the state store before polling the next
247    ///    page. This makes per-page CDC checkpointing automatic.
248    /// 4. Flushes the sink one final time after the stream completes.
249    /// 5. Returns a [`PipelineResult`] with the total count and the last
250    ///    bookmark observed.
251    pub async fn run(&self) -> Result<PipelineResult, FaucetError> {
252        use crate::observability::{
253            DurationGuard, InstrumentedSink, InstrumentedSource, InstrumentedStateStore, Labels,
254        };
255        use metrics::{Label, SharedString, counter, gauge};
256        use tracing::Instrument;
257
258        // Resolve identity for this run.
259        let name = self.name.clone().unwrap_or_else(|| "unnamed".to_string());
260        let row = self.row.clone().unwrap_or_default();
261        let run_id = self
262            .run_id
263            .clone()
264            .unwrap_or_else(|| uuid::Uuid::now_v7().to_string());
265        let obs_labels = Labels::new(name.clone(), row.clone(), run_id.clone());
266
267        // Wrap source, sink, state-store.
268        let wrapped_source = InstrumentedSource::new(self.source, obs_labels.clone());
269        let wrapped_sink = InstrumentedSink::new(self.sink, obs_labels.clone());
270        let wrapped_state_store: Option<Arc<dyn StateStore>> = self.state_store.as_ref().map(|s| {
271            Arc::new(InstrumentedStateStore::new(
272                Arc::clone(s),
273                obs_labels.clone(),
274            )) as Arc<dyn StateStore>
275        });
276
277        // Pipeline-level span. Use .instrument(span) on the inner future so
278        // the span correctly enters/exits across awaits.
279        let span = tracing::info_span!(
280            "faucet.pipeline.run",
281            pipeline = %name,
282            row = %row,
283            run_id = %run_id,
284            source = %wrapped_source.connector_name(),
285            sink = %wrapped_sink.connector_name(),
286        );
287
288        // Per-pipeline metric labels (pipeline + row).
289        let base_labels: Vec<Label> = vec![
290            Label::new("pipeline", SharedString::from(name.clone())),
291            Label::new("row", SharedString::from(row.clone())),
292        ];
293        let run_labels: Vec<Label> = {
294            let mut v = base_labels.clone();
295            v.push(Label::new(
296                "source",
297                SharedString::from(wrapped_source.connector_name().to_string()),
298            ));
299            v.push(Label::new(
300                "sink",
301                SharedString::from(wrapped_sink.connector_name().to_string()),
302            ));
303            v
304        };
305
306        // RAII guard so the in-flight gauge stays consistent even on cancellation.
307        struct InFlightGuard(Vec<Label>);
308        impl Drop for InFlightGuard {
309            fn drop(&mut self) {
310                gauge!("faucet_pipeline_in_flight", self.0.clone()).decrement(1.0);
311            }
312        }
313        gauge!("faucet_pipeline_in_flight", base_labels.clone()).increment(1.0);
314        let _in_flight = InFlightGuard(base_labels.clone());
315
316        // Stamp the start time so dashboards can compute uptime for long-running
317        // (streaming / CDC) pipelines where `*_run_duration_seconds` never fires.
318        let start_unix = std::time::SystemTime::now()
319            .duration_since(std::time::UNIX_EPOCH)
320            .map(|d| d.as_secs_f64())
321            .unwrap_or(0.0);
322        gauge!(
323            "faucet_pipeline_start_time_unix_seconds",
324            base_labels.clone()
325        )
326        .set(start_unix);
327
328        // Histogram timer for the whole run.
329        let _run_timer =
330            DurationGuard::new("faucet_pipeline_run_duration_seconds", run_labels.clone());
331
332        // Run inside the span.
333        let result = async {
334            // Bookmark resume — goes through the wrapped state store so the
335            // get is instrumented too.
336            let state_key = self.source.state_key();
337            let mut start_seq = 0u64;
338            if let (Some(store), Some(key)) = (wrapped_state_store.as_ref(), state_key.as_ref()) {
339                validate_state_key(key)?;
340                if let Some(prior) = store.get(key).await? {
341                    if self.delivery == crate::idempotency::DeliveryMode::ExactlyOnce {
342                        let (bookmark, seq) = crate::idempotency::unwrap_state(&prior);
343                        start_seq = seq;
344                        if let Some(bm) = bookmark {
345                            wrapped_source.apply_start_bookmark(bm).await?;
346                        }
347                    } else {
348                        wrapped_source.apply_start_bookmark(prior).await?;
349                    }
350                }
351            }
352
353            let ctx = std::collections::HashMap::new();
354            let pages = wrapped_source.stream_pages(&ctx, DEFAULT_BATCH_SIZE);
355
356            let mut opts = RunStreamOptions::new()
357                .with_name(name.clone())
358                .with_row(row.clone())
359                .with_run_id(run_id.clone());
360            if let (Some(store), Some(key)) = (wrapped_state_store.clone(), state_key) {
361                opts = opts.with_state(store, key);
362            }
363            if let Some(dlq) = self.dlq.clone() {
364                opts = opts.with_dlq(dlq);
365            }
366            #[cfg(feature = "quality")]
367            if let Some(q) = self.quality.clone() {
368                opts = opts.with_quality(q);
369            }
370            if let Some(ad) = self.adaptive.clone() {
371                opts = opts.with_adaptive(ad);
372            }
373            if let Some(cancel) = self.cancel.clone() {
374                opts = opts.with_cancel(cancel);
375            }
376            opts = opts.with_delivery(self.delivery).with_start_seq(start_seq);
377
378            run_stream(pages, &wrapped_sink, opts).await
379        }
380        .instrument(span)
381        .await;
382
383        // Final run-counter increment. On error, also attach a `kind` label
384        // (matching the FaucetError variant) so dashboards can break out failed
385        // runs by error type without spelunking the *_errors_total surfaces.
386        let status = if result.is_ok() { "ok" } else { "err" };
387        let mut final_labels = run_labels;
388        final_labels.push(Label::new("status", SharedString::const_str(status)));
389        if let Err(ref e) = result {
390            final_labels.push(Label::new(
391                "kind",
392                SharedString::const_str(crate::observability::decorator::error_kind(e)),
393            ));
394        }
395        counter!("faucet_pipeline_runs_total", final_labels).increment(1);
396
397        result
398    }
399}
400
401/// Run a streaming pipeline, writing each [`StreamPage`] to the sink as it
402/// arrives and persisting bookmarks per page.
403///
404/// This keeps memory usage bounded — only one page of records is held at a
405/// time. The stream comes from [`Source::stream_pages`] (or any
406/// `Stream<Item = Result<StreamPage, FaucetError>>` a caller assembles
407/// directly).
408///
409/// Bookmark semantics: whenever a page carries `Some(bookmark)`, the sink is
410/// flushed and the bookmark is persisted (when `state_store` and `state_key`
411/// are both `Some`) before the next page is polled. Sources that only know
412/// their bookmark after seeing every record emit `Some` on the final page;
413/// CDC-style sources emit `Some` per committed transaction and get
414/// per-transaction durability automatically.
415///
416/// Returns the cumulative [`PipelineResult`] — `records_written` is the sum
417/// across all pages and `bookmark` is the last per-page bookmark observed.
418pub async fn run_stream<S, Si>(
419    mut pages: S,
420    sink: &Si,
421    options: RunStreamOptions,
422) -> Result<PipelineResult, FaucetError>
423where
424    S: Stream<Item = Result<StreamPage, FaucetError>> + Unpin,
425    Si: Sink + ?Sized,
426{
427    use crate::dlq::{DlqStats, OnBatchError, build_envelope};
428
429    let state_store = options.state_store.clone();
430    let state_key = options.state_key.clone();
431    let pipeline_name = options.pipeline_name.unwrap_or_else(|| "unnamed".into());
432    let row = options.row.unwrap_or_default();
433    let run_id = options.run_id.unwrap_or_default();
434    let dlq = options.dlq.clone();
435    let cancel = options.cancel.clone();
436
437    #[cfg(feature = "quality")]
438    let quality = options.quality.clone();
439
440    // Fail fast: quarantine requires a DLQ sink.
441    #[cfg(feature = "quality")]
442    if let Some(q) = quality.as_ref()
443        && q.requires_dlq()
444        && dlq.is_none()
445    {
446        return Err(FaucetError::Config(
447            "quality: on_failure 'quarantine'/'quarantine_batch' requires a DLQ sink".into(),
448        ));
449    }
450
451    if let Some(key) = state_key.as_ref() {
452        validate_state_key(key)?;
453    }
454
455    // ── Exactly-once gates + resume ──────────────────────────────────────────
456    let exactly_once = options.delivery == crate::idempotency::DeliveryMode::ExactlyOnce;
457    if exactly_once {
458        if !sink.supports_idempotent_writes() {
459            return Err(FaucetError::Config(format!(
460                "delivery: exactly_once requires an idempotent sink, but '{}' does not support it",
461                sink.connector_name()
462            )));
463        }
464        if state_store.is_none() || state_key.is_none() {
465            return Err(FaucetError::Config(
466                "delivery: exactly_once requires a state store".into(),
467            ));
468        }
469        if dlq.is_some() {
470            return Err(FaucetError::Config(
471                "delivery: exactly_once is not compatible with a DLQ in this version".into(),
472            ));
473        }
474    }
475    let scope = state_key.clone().unwrap_or_default();
476    let mut next_seq = options.start_seq;
477    let committed_seq = if exactly_once {
478        sink.last_committed_token(&scope)
479            .await?
480            .and_then(|t| crate::idempotency::parse_token(&t))
481            .unwrap_or(0)
482    } else {
483        0
484    };
485
486    let mut records_written = 0usize;
487    let mut last_bookmark: Option<Value> = None;
488    let mut dlq_stats = DlqStats::default();
489
490    let adaptive_cfg = options.adaptive.clone().filter(|c| c.enabled);
491    // Validate at the core boundary so library callers of `run_stream` (not
492    // just the CLI, which validates earlier) reject an invalid adaptive config
493    // — e.g. the rejected `respect_source_max=false` knob — up front.
494    if let Some(cfg) = adaptive_cfg.as_ref() {
495        cfg.validate()?;
496    }
497    let mut controller: Option<crate::adaptive::AimdController> = None;
498    let mut warned_noop_sink = false;
499
500    let sink_name = sink.connector_name();
501    let dlq_sink_name = dlq.as_ref().map(|d| d.sink.connector_name()).unwrap_or("");
502
503    // Drive the streaming loop inside an inner future so that EVERY early exit
504    // (a source error, a `?`-propagated write/flush/state failure, or a DLQ
505    // budget overflow) funnels through one place. On any error we best-effort
506    // flush the sinks before propagating, so a buffered sink that only commits
507    // on flush — Parquet writes its footer there; without it the whole file is
508    // unreadable — does not lose everything written so far (#78/#3).
509    // Set when the loop exits because the cancellation token fired (vs. the
510    // stream ending naturally). Either way we fall through to the success-path
511    // flush below, so a buffered sink (Parquet footer, S3 multipart) is made
512    // durable — the difference from a dropped future, which flushes nothing.
513    let mut cancelled = false;
514    let loop_result: Result<(), FaucetError> = async {
515        loop {
516            // Poll the next page, but if a cancellation token is wired, race it
517            // so a cancel between pages stops the run promptly and cleanly
518            // (#146 H16). `biased` checks cancellation first each iteration.
519            let page = match &cancel {
520                Some(token) => tokio::select! {
521                    biased;
522                    _ = token.cancelled() => {
523                        cancelled = true;
524                        break;
525                    }
526                    p = std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)) => p,
527                },
528                None => std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)).await,
529            };
530            match page {
531                Some(Ok(page)) => {
532                    if page.records.is_empty() && page.bookmark.is_none() {
533                        continue;
534                    }
535
536                    // ── Quality pass (after transforms, before sink) ─────────
537                    #[cfg(feature = "quality")]
538                    let (records, quality_envelopes): (Vec<Value>, Vec<Value>) =
539                        if let Some(q) = quality.as_ref() {
540                            let labels =
541                                crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
542                            let outcome = crate::observability::instrumented_apply_quality(
543                                page.records,
544                                q,
545                                &labels,
546                            )?;
547                            let envelopes: Vec<Value> = outcome
548                                .quarantined
549                                .iter()
550                                .map(|qr| {
551                                    let err = FaucetError::QualityFailure {
552                                        check: qr.check.to_string(),
553                                        message: qr.message.clone(),
554                                    };
555                                    // `record_index` is the position within the PAGE
556                                    // (the frozen envelope contract), not the index in
557                                    // the quarantine list (#146 R).
558                                    build_envelope(
559                                        &qr.record,
560                                        &err,
561                                        sink_name,
562                                        &pipeline_name,
563                                        &row,
564                                        qr.page_index,
565                                    )
566                                })
567                                .collect();
568                            (outcome.survivors, envelopes)
569                        } else {
570                            (page.records, Vec::new())
571                        };
572                    #[cfg(not(feature = "quality"))]
573                    let (records, quality_envelopes): (Vec<Value>, Vec<Value>) =
574                        (page.records, Vec::new());
575
576                    let page = StreamPage {
577                        records,
578                        bookmark: page.bookmark,
579                    };
580
581                    if let Some(ref dlq_cfg) = dlq {
582                        // ── DLQ-enabled path ───────────────────────────────────
583                        use crate::dlq::DlqReason;
584                        use metrics::{Label, SharedString, counter};
585                        let metric_labels: Vec<Label> = vec![
586                            Label::new("pipeline", SharedString::from(pipeline_name.clone())),
587                            Label::new("row", SharedString::from(row.clone())),
588                            Label::new("connector", SharedString::from(sink_name.to_string())),
589                            Label::new(
590                                "dlq_connector",
591                                SharedString::from(dlq_sink_name.to_string()),
592                            ),
593                        ];
594                        let span = tracing::info_span!(
595                            "faucet.dlq.route",
596                            pipeline = %pipeline_name,
597                            row = %row,
598                            run_id = %run_id,
599                            connector = %sink_name,
600                            dlq_connector = %dlq_sink_name,
601                        );
602                        let _enter = span.enter();
603
604                        // Reslice the page into sub-batches driven by the
605                        // adaptive controller (or write the whole page in one
606                        // shot when adaptive is disabled — same as before).
607                        let mut envelopes: Vec<Value> = Vec::new();
608                        let mut page_success = 0usize;
609                        let mut outer_err_recovered = false;
610                        // True if any chunk reported genuine per-row sink `Err`s
611                        // (as opposed to a chunk wholly synthesized from an outer
612                        // error under `DlqAll`). Drives the `partial` label when a
613                        // resliced page mixes the two failure modes.
614                        let mut had_per_row_sink_failure = false;
615                        let records_len = page.records.len();
616                        let mut offset = 0usize;
617                        while offset < records_len {
618                            let size = match adaptive_cfg.as_ref() {
619                                Some(cfg) => {
620                                    let ctrl = controller.get_or_insert_with(|| {
621                                        crate::adaptive::AimdController::new(cfg, records_len)
622                                    });
623                                    ctrl.current().max(1).min(records_len - offset)
624                                }
625                                None => records_len - offset, // whole page = today's behavior
626                            };
627                            if adaptive_cfg.is_some() {
628                                maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
629                            }
630                            let chunk = &page.records[offset..offset + size];
631                            let t0 = std::time::Instant::now();
632                            let chunk_outcomes_result = sink.write_batch_partial(chunk).await;
633                            let latency = t0.elapsed();
634                            // `chunk_synthesized` is true only when this chunk's
635                            // outcomes were fabricated from a single outer
636                            // `write_batch_partial` error under `DlqAll` — as
637                            // opposed to genuine per-row `Err`s the sink
638                            // reported. Tracking it per chunk keeps the page
639                            // `reason` label accurate when adaptive reslicing
640                            // mixes a synthesized chunk with partial-failure
641                            // chunks on the same page.
642                            let (chunk_outcomes, chunk_synthesized): (
643                                Vec<crate::RowOutcome>,
644                                bool,
645                            ) = match chunk_outcomes_result {
646                                Ok(o) => (o, false),
647                                Err(e) => match dlq_cfg.on_batch_error {
648                                    OnBatchError::Propagate => return Err(e),
649                                    OnBatchError::DlqAll => {
650                                        outer_err_recovered = true;
651                                        let msg = e.to_string();
652                                        let synth = (0..chunk.len())
653                                            .map(|_| Err(FaucetError::Sink(msg.clone())))
654                                            .collect();
655                                        (synth, true)
656                                    }
657                                },
658                            };
659                            let mut chunk_errors = 0usize;
660                            for (j, outcome) in chunk_outcomes.iter().enumerate() {
661                                match outcome {
662                                    Ok(()) => page_success += 1,
663                                    Err(err) => {
664                                        chunk_errors += 1;
665                                        if !chunk_synthesized {
666                                            had_per_row_sink_failure = true;
667                                        }
668                                        envelopes.push(build_envelope(
669                                            &chunk[j],
670                                            err,
671                                            sink_name,
672                                            &pipeline_name,
673                                            &row,
674                                            offset + j,
675                                        ));
676                                    }
677                                }
678                            }
679                            if let Some(ctrl) = controller.as_mut() {
680                                let adj = ctrl.observe(crate::adaptive::Observation {
681                                    batch_len: chunk.len(),
682                                    errors: chunk_errors,
683                                    latency,
684                                });
685                                emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
686                            }
687                            offset += size;
688                        }
689                        // Quality-quarantined records share the DLQ budget/write.
690                        // Capture the quality count BEFORE the splice — the splice
691                        // moves `quality_envelopes`, so its length is unavailable
692                        // afterward. Used below to pick the page `reason` label.
693                        #[cfg(feature = "quality")]
694                        let quality_count = quality_envelopes.len();
695                        #[cfg(not(feature = "quality"))]
696                        let quality_count = 0usize;
697                        envelopes.splice(0..0, quality_envelopes);
698                        let page_failures = envelopes.len();
699
700                        // Budget checks. `write_batch_partial` above already
701                        // committed this page's survivors to the main sink, so
702                        // we must NOT abort here: returning now would strand
703                        // those committed survivors without advancing the
704                        // bookmark (they would re-deliver on the next run) and
705                        // drop this page's failures before they reach the DLQ
706                        // (#146 M4). Instead, record the budget error, finish
707                        // committing the page below (route failures to the DLQ,
708                        // flush, persist the bookmark), and abort only once the
709                        // page is fully durable. The failed rows that crossed
710                        // the threshold are still written to the DLQ — losing
711                        // them would be strictly worse than the small overshoot.
712                        let mut budget_error: Option<FaucetError> = None;
713                        if let Some(limit) = dlq_cfg.max_failures_per_page
714                            && page_failures > limit
715                        {
716                            let mut lbl = metric_labels.clone();
717                            lbl.retain(|l| l.key() != "dlq_connector");
718                            lbl.push(Label::new("scope", SharedString::const_str("per_page")));
719                            counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
720                            budget_error = Some(FaucetError::Sink(format!(
721                                "DLQ per-page budget exceeded: {page_failures} > {limit}"
722                            )));
723                        }
724                        let new_total = dlq_stats.records_dlq + page_failures;
725                        if budget_error.is_none()
726                            && let Some(limit) = dlq_cfg.max_failures_total
727                            && new_total > limit
728                        {
729                            let mut lbl = metric_labels.clone();
730                            lbl.retain(|l| l.key() != "dlq_connector");
731                            lbl.push(Label::new("scope", SharedString::const_str("total")));
732                            counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
733                            budget_error = Some(FaucetError::Sink(format!(
734                                "DLQ total budget exceeded: {new_total} > {limit}"
735                            )));
736                        }
737
738                        // Write to DLQ sink. Errors here are fatal, no recursion.
739                        if !envelopes.is_empty() {
740                            let _dlq_write_timer = crate::observability::DurationGuard::new(
741                                "faucet_sink_dlq_write_duration_seconds",
742                                metric_labels.clone(),
743                            );
744                            dlq_cfg.sink.write_batch(&envelopes).await.map_err(|e| {
745                                let mut lbl = metric_labels.clone();
746                                lbl.push(Label::new(
747                                    "kind",
748                                    SharedString::const_str(
749                                        crate::observability::decorator::error_kind(&e),
750                                    ),
751                                ));
752                                counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
753                                FaucetError::Sink(format!("DLQ sink write failed: {e}"))
754                            })?;
755                            dlq_stats.records_dlq += page_failures;
756                            dlq_stats.pages_with_failures += 1;
757
758                            // Page `reason` label, 3-way (precedence: partial > dlq_all > quality):
759                            //  - `partial`  — at least one chunk reported genuine
760                            //    per-row sink `Err`s. Checked FIRST so a resliced
761                            //    page that mixes a synthesized chunk (DlqAll) with
762                            //    partial-failure chunks is labeled `partial` — the
763                            //    real per-row failure dominates. (For a
764                            //    non-resliced page this is equivalent to the old
765                            //    `page_failures > quality_count` test, since a
766                            //    single chunk is either all-synthesized or all
767                            //    per-row.)
768                            //  - `dlq_all`  — every sink-side failure on the page
769                            //    was synthesized from an outer `write_batch_partial`
770                            //    error (OnBatchError::DlqAll); no genuine per-row
771                            //    failures occurred.
772                            //  - `quality`  — every envelope is quality-sourced
773                            //    (no sink-side failures on this page).
774                            // The per-row quality volume is separately exposed via
775                            // `faucet_quality_records_quarantined_total`.
776                            let reason_label = if had_per_row_sink_failure {
777                                DlqReason::Partial.as_str()
778                            } else if outer_err_recovered {
779                                DlqReason::DlqAll.as_str()
780                            } else if page_failures > quality_count {
781                                DlqReason::Partial.as_str()
782                            } else {
783                                DlqReason::Quality.as_str()
784                            };
785                            counter!("faucet_sink_dlq_records_total", metric_labels.clone())
786                                .increment(page_failures as u64);
787                            let mut page_labels = metric_labels.clone();
788                            page_labels
789                                .push(Label::new("reason", SharedString::const_str(reason_label)));
790                            counter!("faucet_sink_dlq_pages_total", page_labels).increment(1);
791                        }
792
793                        records_written += page_success;
794
795                        if let Some(bookmark) = page.bookmark {
796                            sink.flush().await?;
797                            let _dlq_flush_timer = crate::observability::DurationGuard::new(
798                                "faucet_sink_dlq_flush_duration_seconds",
799                                metric_labels.clone(),
800                            );
801                            dlq_cfg.sink.flush().await.map_err(|e| {
802                                let mut lbl = metric_labels.clone();
803                                lbl.push(Label::new(
804                                    "kind",
805                                    SharedString::const_str(
806                                        crate::observability::decorator::error_kind(&e),
807                                    ),
808                                ));
809                                counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
810                                FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
811                            })?;
812                            let bm_labels =
813                                crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
814                            crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
815                            if let (Some(store), Some(key)) =
816                                (state_store.as_ref(), state_key.as_ref())
817                            {
818                                store.put(key, &bookmark).await?;
819                            }
820                            last_bookmark = Some(bookmark);
821                        }
822
823                        // The page is now durable — survivors committed to the
824                        // main sink, failures routed to the DLQ, and (if the
825                        // page carried one) the bookmark persisted. Honor a
826                        // deferred DLQ-budget abort now, so the run still stops
827                        // as a circuit breaker but never re-delivers this
828                        // already-committed page (#146 M4).
829                        if let Some(e) = budget_error {
830                            return Err(e);
831                        }
832                    } else if exactly_once {
833                        // ── Exactly-once path ──────────────────────────────────
834                        // A token is issued only for bookmark-carrying pages, so
835                        // (seq, bookmark) advance together and realign on resume.
836                        if let Some(bookmark) = page.bookmark {
837                            next_seq += 1;
838                            let token = crate::idempotency::format_token(next_seq);
839                            if next_seq <= committed_seq {
840                                // Sink already durably committed this page. Skip
841                                // the write; advance state so a later crash does
842                                // not re-skip it.
843                                use metrics::{Label, SharedString, counter};
844                                let skip_labels: Vec<Label> = vec![
845                                    Label::new(
846                                        "pipeline",
847                                        SharedString::from(pipeline_name.clone()),
848                                    ),
849                                    Label::new("row", SharedString::from(row.clone())),
850                                ];
851                                counter!("faucet_pipeline_pages_skipped_total", skip_labels)
852                                    .increment(1);
853                            } else {
854                                records_written += sink
855                                    .write_batch_idempotent(&page.records, &scope, &token)
856                                    .await?;
857                            }
858                            sink.flush().await?;
859                            let bm_labels =
860                                crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
861                            crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
862                            if let (Some(store), Some(key)) =
863                                (state_store.as_ref(), state_key.as_ref())
864                            {
865                                store
866                                    .put(
867                                        key,
868                                        &crate::idempotency::wrap_state(Some(&bookmark), next_seq),
869                                    )
870                                    .await?;
871                            }
872                            last_bookmark = Some(bookmark);
873                        } else if !page.records.is_empty() {
874                            // No bookmark → not individually checkpointed; write
875                            // as-is (rare for EO sources, which bookmark every
876                            // page). Stays at-least-once for this page.
877                            records_written += sink.write_batch(&page.records).await?;
878                        }
879                    } else {
880                        // ── DLQ-disabled path (today's behaviour) ──────────────
881                        debug_assert!(
882                            quality_envelopes.is_empty(),
883                            "quality quarantine without DLQ should have been rejected at run start"
884                        );
885                        if !page.records.is_empty() {
886                            if let Some(cfg) = adaptive_cfg.as_ref() {
887                                let ctrl = controller.get_or_insert_with(|| {
888                                    crate::adaptive::AimdController::new(cfg, page.records.len())
889                                });
890                                maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
891                                let mut offset = 0;
892                                while offset < page.records.len() {
893                                    let size =
894                                        ctrl.current().max(1).min(page.records.len() - offset);
895                                    let chunk = &page.records[offset..offset + size];
896                                    let t0 = std::time::Instant::now();
897                                    let n = sink.write_batch(chunk).await?;
898                                    let latency = t0.elapsed();
899                                    records_written += n;
900                                    offset += size;
901                                    let adj = ctrl.observe(crate::adaptive::Observation {
902                                        batch_len: chunk.len(),
903                                        errors: 0,
904                                        latency,
905                                    });
906                                    emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
907                                }
908                            } else {
909                                records_written += sink.write_batch(&page.records).await?;
910                            }
911                        }
912                        if let Some(bookmark) = page.bookmark {
913                            sink.flush().await?;
914                            let bm_labels =
915                                crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
916                            crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
917                            if let (Some(store), Some(key)) =
918                                (state_store.as_ref(), state_key.as_ref())
919                            {
920                                store.put(key, &bookmark).await?;
921                            }
922                            last_bookmark = Some(bookmark);
923                        }
924                    }
925                }
926                Some(Err(e)) => return Err(e),
927                None => break,
928            }
929        }
930        Ok(())
931    }
932    .await;
933
934    // Error/early-return unwind: best-effort flush so any buffered output is
935    // made durable, then propagate the ORIGINAL error. Flush errors here are
936    // logged and swallowed — the source/sink error that triggered the unwind
937    // is the meaningful one to surface. DLQ is flushed first (mirroring the
938    // success path below): its records are only ever written here, whereas the
939    // next run re-reads post-bookmark records from the source.
940    if let Err(e) = loop_result {
941        if let Some(ref dlq_cfg) = dlq
942            && let Err(flush_err) = dlq_cfg.sink.flush().await
943        {
944            tracing::warn!(
945                error = %flush_err,
946                "DLQ sink flush failed during error unwind; original error preserved"
947            );
948        }
949        if let Err(flush_err) = sink.flush().await {
950            tracing::warn!(
951                error = %flush_err,
952                "sink flush failed during error unwind; original error preserved"
953            );
954        }
955        return Err(e);
956    }
957
958    // Flush the DLQ sink BEFORE the main sink so quarantined records are made
959    // durable even if the main sink's final flush fails. The next run will
960    // re-read post-bookmark records from the source and re-route any that
961    // would have fallen out of the main sink's unflushed buffer; DLQ records,
962    // by contrast, are only ever written here and would otherwise be lost.
963    if let Some(ref dlq_cfg) = dlq {
964        let final_metric_labels: Vec<metrics::Label> = vec![
965            metrics::Label::new(
966                "pipeline",
967                metrics::SharedString::from(pipeline_name.clone()),
968            ),
969            metrics::Label::new("row", metrics::SharedString::from(row.clone())),
970            metrics::Label::new(
971                "connector",
972                metrics::SharedString::from(sink_name.to_string()),
973            ),
974            metrics::Label::new(
975                "dlq_connector",
976                metrics::SharedString::from(dlq_sink_name.to_string()),
977            ),
978        ];
979        let _final_dlq_flush_timer = crate::observability::DurationGuard::new(
980            "faucet_sink_dlq_flush_duration_seconds",
981            final_metric_labels.clone(),
982        );
983        dlq_cfg.sink.flush().await.map_err(|e| {
984            let mut lbl = final_metric_labels.clone();
985            lbl.push(metrics::Label::new(
986                "kind",
987                metrics::SharedString::const_str(crate::observability::decorator::error_kind(&e)),
988            ));
989            metrics::counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
990            FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
991        })?;
992    }
993    sink.flush().await?;
994
995    if cancelled {
996        tracing::info!(
997            records_written,
998            "pipeline run cancelled cooperatively; sink flushed (partial output is durable)"
999        );
1000    }
1001
1002    tracing::info!(
1003        records_written,
1004        cancelled,
1005        has_bookmark = last_bookmark.is_some(),
1006        persisted = state_store.is_some() && state_key.is_some() && last_bookmark.is_some(),
1007        dlq_records = dlq_stats.records_dlq,
1008        "pipeline streaming run complete"
1009    );
1010
1011    Ok(PipelineResult {
1012        records_written,
1013        bookmark: last_bookmark,
1014        dlq: dlq.is_some().then_some(dlq_stats),
1015    })
1016}
1017
1018/// Emit the adaptive controller's current state + any adjustment as metrics.
1019/// Labels are `pipeline,row` only (the controller is pipeline-scoped).
1020fn emit_adaptive_metrics(
1021    ctrl: &crate::adaptive::AimdController,
1022    adj: Option<crate::adaptive::Adjustment>,
1023    pipeline: &str,
1024    row: &str,
1025) {
1026    use metrics::{Label, SharedString, counter, gauge};
1027    let base = vec![
1028        Label::new("pipeline", SharedString::from(pipeline.to_string())),
1029        Label::new("row", SharedString::from(row.to_string())),
1030    ];
1031    gauge!("faucet_pipeline_adaptive_batch_size", base.clone()).set(ctrl.current() as f64);
1032    gauge!(
1033        "faucet_pipeline_adaptive_batch_cooldown_active",
1034        base.clone()
1035    )
1036    .set(if ctrl.cooldown_active() { 1.0 } else { 0.0 });
1037    if let Some(p50) = ctrl.p50_latency_ms() {
1038        gauge!(
1039            "faucet_pipeline_adaptive_batch_p50_latency_ms",
1040            base.clone()
1041        )
1042        .set(p50 as f64);
1043    }
1044    if let Some(a) = adj {
1045        let mut lbl = base;
1046        lbl.push(Label::new(
1047            "direction",
1048            SharedString::const_str(a.direction.as_str()),
1049        ));
1050        lbl.push(Label::new(
1051            "reason",
1052            SharedString::const_str(a.reason.as_str()),
1053        ));
1054        counter!("faucet_pipeline_adaptive_batch_adjustments_total", lbl).increment(1);
1055    }
1056}
1057
1058/// One-shot info when adaptive sizing targets a per-record sink that ignores
1059/// `batch_size` (its adjustments are harmless no-ops).
1060fn maybe_warn_noop_sink(sink_name: &str, warned: &mut bool) {
1061    if !*warned && matches!(sink_name, "jsonl" | "csv" | "stdout") {
1062        tracing::info!(
1063            sink = sink_name,
1064            "adaptive batch sizing is a no-op for this per-record sink"
1065        );
1066        *warned = true;
1067    }
1068}
1069
1070#[cfg(test)]
1071mod tests {
1072    use super::*;
1073    use async_trait::async_trait;
1074    use serde_json::json;
1075
1076    // ── Mock Source ──────────────────────────────────────────────────────────
1077
1078    struct MockSource(Vec<Value>);
1079
1080    #[async_trait]
1081    impl Source for MockSource {
1082        async fn fetch_with_context(
1083            &self,
1084            _context: &std::collections::HashMap<String, Value>,
1085        ) -> Result<Vec<Value>, FaucetError> {
1086            Ok(self.0.clone())
1087        }
1088    }
1089
1090    struct IncrementalSource {
1091        records: Vec<Value>,
1092        bookmark: Value,
1093    }
1094
1095    #[async_trait]
1096    impl Source for IncrementalSource {
1097        async fn fetch_with_context(
1098            &self,
1099            _context: &std::collections::HashMap<String, Value>,
1100        ) -> Result<Vec<Value>, FaucetError> {
1101            Ok(self.records.clone())
1102        }
1103        async fn fetch_with_context_incremental(
1104            &self,
1105            _context: &std::collections::HashMap<String, Value>,
1106        ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
1107            Ok((self.records.clone(), Some(self.bookmark.clone())))
1108        }
1109    }
1110
1111    struct FailingSource;
1112
1113    #[async_trait]
1114    impl Source for FailingSource {
1115        async fn fetch_with_context(
1116            &self,
1117            _context: &std::collections::HashMap<String, Value>,
1118        ) -> Result<Vec<Value>, FaucetError> {
1119            Err(FaucetError::Auth("no credentials".into()))
1120        }
1121    }
1122
1123    // ── Mock Sink ───────────────────────────────────────────────────────────
1124
1125    struct MockSink(std::sync::Mutex<Vec<Value>>);
1126
1127    impl MockSink {
1128        fn new() -> Self {
1129            Self(std::sync::Mutex::new(Vec::new()))
1130        }
1131        fn written(&self) -> Vec<Value> {
1132            self.0.lock().unwrap().clone()
1133        }
1134    }
1135
1136    #[async_trait]
1137    impl Sink for MockSink {
1138        async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1139            self.0.lock().unwrap().extend(records.iter().cloned());
1140            Ok(records.len())
1141        }
1142    }
1143
1144    struct FailingSink;
1145
1146    #[async_trait]
1147    impl Sink for FailingSink {
1148        async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
1149            Err(FaucetError::Sink("write failed".into()))
1150        }
1151    }
1152
1153    /// Records writes and how many times `flush` was called. Used to assert the
1154    /// pipeline flushes the sink on the error/early-return path so partial
1155    /// output (e.g. a Parquet footer) is made durable before the error
1156    /// propagates.
1157    struct FlushTrackingSink {
1158        written: std::sync::Mutex<Vec<Value>>,
1159        flush_count: std::sync::atomic::AtomicUsize,
1160    }
1161
1162    impl FlushTrackingSink {
1163        fn new() -> Self {
1164            Self {
1165                written: std::sync::Mutex::new(Vec::new()),
1166                flush_count: std::sync::atomic::AtomicUsize::new(0),
1167            }
1168        }
1169        fn written(&self) -> Vec<Value> {
1170            self.written.lock().unwrap().clone()
1171        }
1172        fn flush_count(&self) -> usize {
1173            self.flush_count.load(std::sync::atomic::Ordering::SeqCst)
1174        }
1175    }
1176
1177    #[async_trait]
1178    impl Sink for FlushTrackingSink {
1179        async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1180            self.written.lock().unwrap().extend(records.iter().cloned());
1181            Ok(records.len())
1182        }
1183        async fn flush(&self) -> Result<(), FaucetError> {
1184            self.flush_count
1185                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1186            Ok(())
1187        }
1188    }
1189
1190    // ── Exactly-once test doubles ────────────────────────────────────────────
1191
1192    /// In-memory sink that commits rows + a per-scope token atomically.
1193    struct IdempotentMockSink {
1194        rows: std::sync::Mutex<Vec<Value>>,
1195        tokens: std::sync::Mutex<std::collections::HashMap<String, String>>,
1196    }
1197    impl IdempotentMockSink {
1198        fn new() -> Self {
1199            Self {
1200                rows: std::sync::Mutex::new(Vec::new()),
1201                tokens: std::sync::Mutex::new(std::collections::HashMap::new()),
1202            }
1203        }
1204        fn rows(&self) -> Vec<Value> {
1205            self.rows.lock().unwrap().clone()
1206        }
1207    }
1208    #[async_trait]
1209    impl Sink for IdempotentMockSink {
1210        async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1211            self.rows.lock().unwrap().extend(records.iter().cloned());
1212            Ok(records.len())
1213        }
1214        fn supports_idempotent_writes(&self) -> bool {
1215            true
1216        }
1217        async fn write_batch_idempotent(
1218            &self,
1219            records: &[Value],
1220            scope: &str,
1221            token: &str,
1222        ) -> Result<usize, FaucetError> {
1223            self.rows.lock().unwrap().extend(records.iter().cloned());
1224            self.tokens
1225                .lock()
1226                .unwrap()
1227                .insert(scope.to_string(), token.to_string());
1228            Ok(records.len())
1229        }
1230        async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
1231            Ok(self.tokens.lock().unwrap().get(scope).cloned())
1232        }
1233    }
1234
1235    fn eo_opts(store: Arc<dyn StateStore>, key: &str, start_seq: u64) -> RunStreamOptions {
1236        RunStreamOptions::new()
1237            .with_state(store, key)
1238            .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
1239            .with_start_seq(start_seq)
1240    }
1241
1242    #[tokio::test]
1243    async fn exactly_once_writes_pages_and_persists_wrapped_state() {
1244        let pages = vec![
1245            Ok(StreamPage {
1246                records: vec![json!({"id": 1})],
1247                bookmark: Some(json!("b1")),
1248            }),
1249            Ok(StreamPage {
1250                records: vec![json!({"id": 2})],
1251                bookmark: Some(json!("b2")),
1252            }),
1253        ];
1254        let sink = IdempotentMockSink::new();
1255        let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
1256        let r = run_stream(
1257            futures::stream::iter(pages),
1258            &sink,
1259            eo_opts(store.clone(), "k", 0),
1260        )
1261        .await
1262        .unwrap();
1263        assert_eq!(r.records_written, 2);
1264        let (bm, seq) = crate::idempotency::unwrap_state(&store.get("k").await.unwrap().unwrap());
1265        assert_eq!(bm, Some(json!("b2")));
1266        assert_eq!(seq, 2);
1267        assert_eq!(
1268            sink.last_committed_token("k").await.unwrap(),
1269            Some(crate::idempotency::format_token(2))
1270        );
1271    }
1272
1273    #[tokio::test]
1274    async fn exactly_once_skips_already_committed_pages_on_resume() {
1275        let sink = IdempotentMockSink::new();
1276        // Run 1: commit page seq 1 directly (simulate crash: state lost).
1277        sink.write_batch_idempotent(
1278            &[json!({"id": 1})],
1279            "k",
1280            &crate::idempotency::format_token(1),
1281        )
1282        .await
1283        .unwrap();
1284        assert_eq!(sink.rows().len(), 1);
1285        // Run 2 (resume): fresh state, full replay. Page 1 must be skipped.
1286        let pages = vec![
1287            Ok(StreamPage {
1288                records: vec![json!({"id": 1})],
1289                bookmark: Some(json!("b1")),
1290            }),
1291            Ok(StreamPage {
1292                records: vec![json!({"id": 2})],
1293                bookmark: Some(json!("b2")),
1294            }),
1295        ];
1296        let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
1297        let r = run_stream(futures::stream::iter(pages), &sink, eo_opts(store, "k", 0))
1298            .await
1299            .unwrap();
1300        assert_eq!(r.records_written, 1);
1301        let rows = sink.rows();
1302        assert_eq!(
1303            rows.len(),
1304            2,
1305            "exactly one row per id — no duplicate of id=1"
1306        );
1307        assert_eq!(rows.iter().filter(|v| v["id"] == 1).count(), 1);
1308    }
1309
1310    #[tokio::test]
1311    async fn exactly_once_rejects_non_idempotent_sink() {
1312        let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
1313        let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
1314        let r = run_stream(
1315            futures::stream::iter(pages),
1316            &MockSink::new(),
1317            eo_opts(store, "k", 0),
1318        )
1319        .await;
1320        assert!(matches!(r, Err(FaucetError::Config(_))));
1321    }
1322
1323    #[tokio::test]
1324    async fn exactly_once_rejects_missing_state_store() {
1325        let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
1326        let opts =
1327            RunStreamOptions::new().with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce);
1328        let r = run_stream(
1329            futures::stream::iter(pages),
1330            &IdempotentMockSink::new(),
1331            opts,
1332        )
1333        .await;
1334        assert!(matches!(r, Err(FaucetError::Config(_))));
1335    }
1336
1337    // ── StreamPage / batch_size tests ───────────────────────────────────────
1338
1339    #[test]
1340    fn stream_page_constructs() {
1341        let page = StreamPage {
1342            records: vec![json!({"id": 1})],
1343            bookmark: Some(json!("2026-05-18")),
1344        };
1345        assert_eq!(page.records.len(), 1);
1346        assert_eq!(page.bookmark, Some(json!("2026-05-18")));
1347    }
1348
1349    #[test]
1350    fn validate_batch_size_accepts_zero_as_no_batching_sentinel() {
1351        // 0 means "do not batch — emit/accept the whole result set in one page".
1352        assert_eq!(validate_batch_size(0).unwrap(), 0);
1353    }
1354
1355    #[test]
1356    fn validate_batch_size_rejects_too_large() {
1357        let err = validate_batch_size(MAX_BATCH_SIZE + 1).unwrap_err();
1358        assert!(matches!(err, FaucetError::Config(_)));
1359    }
1360
1361    #[test]
1362    fn validate_batch_size_accepts_one() {
1363        assert_eq!(validate_batch_size(1).unwrap(), 1);
1364    }
1365
1366    #[test]
1367    fn validate_batch_size_accepts_max() {
1368        assert_eq!(validate_batch_size(MAX_BATCH_SIZE).unwrap(), MAX_BATCH_SIZE);
1369    }
1370
1371    // Compile-time invariant: DEFAULT_BATCH_SIZE must be within [1, MAX_BATCH_SIZE].
1372    const _: () = {
1373        assert!(DEFAULT_BATCH_SIZE >= 1);
1374        assert!(DEFAULT_BATCH_SIZE <= MAX_BATCH_SIZE);
1375    };
1376
1377    // ── Batch mode tests ────────────────────────────────────────────────────
1378
1379    #[tokio::test]
1380    async fn batch_pipeline_writes_all_records() {
1381        let source = MockSource(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
1382        let sink = MockSink::new();
1383
1384        let result = Pipeline::new(&source, &sink).run().await.unwrap();
1385
1386        assert_eq!(result.records_written, 3);
1387        assert!(result.bookmark.is_none());
1388        assert_eq!(sink.written().len(), 3);
1389    }
1390
1391    #[tokio::test]
1392    async fn batch_pipeline_returns_bookmark() {
1393        let source = IncrementalSource {
1394            records: vec![json!({"id": 1, "ts": "2024-12-01"})],
1395            bookmark: json!("2024-12-01"),
1396        };
1397        let sink = MockSink::new();
1398
1399        let result = Pipeline::new(&source, &sink).run().await.unwrap();
1400
1401        assert_eq!(result.records_written, 1);
1402        assert_eq!(result.bookmark, Some(json!("2024-12-01")));
1403    }
1404
1405    #[tokio::test]
1406    async fn batch_pipeline_empty_source() {
1407        let source = MockSource(vec![]);
1408        let sink = MockSink::new();
1409
1410        let result = Pipeline::new(&source, &sink).run().await.unwrap();
1411
1412        assert_eq!(result.records_written, 0);
1413        assert!(sink.written().is_empty());
1414    }
1415
1416    #[tokio::test]
1417    async fn batch_pipeline_source_error_propagates() {
1418        let source = FailingSource;
1419        let sink = MockSink::new();
1420
1421        let result = Pipeline::new(&source, &sink).run().await;
1422        assert!(result.is_err());
1423        assert!(sink.written().is_empty());
1424    }
1425
1426    #[tokio::test]
1427    async fn batch_pipeline_sink_error_propagates() {
1428        let source = MockSource(vec![json!({"id": 1})]);
1429        let sink = FailingSink;
1430
1431        let result = Pipeline::new(&source, &sink).run().await;
1432        assert!(result.is_err());
1433    }
1434
1435    #[tokio::test]
1436    async fn batch_pipeline_with_trait_objects() {
1437        let source: Box<dyn Source> = Box::new(MockSource(vec![json!({"id": 1})]));
1438        let sink: Box<dyn Sink> = Box::new(MockSink::new());
1439
1440        let result = Pipeline::new(source.as_ref(), sink.as_ref())
1441            .run()
1442            .await
1443            .unwrap();
1444
1445        assert_eq!(result.records_written, 1);
1446    }
1447
1448    // ── Streaming mode tests ────────────────────────────────────────────────
1449
1450    #[tokio::test]
1451    async fn stream_pipeline_writes_pages() {
1452        let pages: Vec<Result<StreamPage, FaucetError>> = vec![
1453            Ok(StreamPage {
1454                records: vec![json!({"id": 1}), json!({"id": 2})],
1455                bookmark: None,
1456            }),
1457            Ok(StreamPage {
1458                records: vec![json!({"id": 3})],
1459                bookmark: None,
1460            }),
1461        ];
1462        let stream = futures::stream::iter(pages);
1463        let sink = MockSink::new();
1464
1465        let result = run_stream(stream, &sink, RunStreamOptions::new())
1466            .await
1467            .unwrap();
1468
1469        assert_eq!(result.records_written, 3);
1470        assert!(result.bookmark.is_none());
1471        assert_eq!(sink.written().len(), 3);
1472    }
1473
1474    #[tokio::test]
1475    async fn stream_pipeline_flushes_sink_on_source_error() {
1476        // Regression for #78/#3: a mid-stream source error must not skip the
1477        // sink flush. Without flushing, a buffered sink (e.g. Parquet, whose
1478        // footer is only written on flush) loses everything written so far.
1479        let pages: Vec<Result<StreamPage, FaucetError>> = vec![
1480            Ok(StreamPage {
1481                records: vec![json!({"id": 1}), json!({"id": 2})],
1482                bookmark: None,
1483            }),
1484            Err(FaucetError::Source("transient blip mid-stream".into())),
1485        ];
1486        let stream = futures::stream::iter(pages);
1487        let sink = FlushTrackingSink::new();
1488
1489        let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
1490
1491        // The original source error must still propagate.
1492        assert!(matches!(result, Err(FaucetError::Source(_))));
1493        // The good page must have been written before the error.
1494        assert_eq!(sink.written().len(), 2);
1495        // Crucially, the sink must have been flushed on the error path.
1496        assert!(
1497            sink.flush_count() >= 1,
1498            "sink must be flushed on the error path so partial output is durable"
1499        );
1500    }
1501
1502    #[tokio::test]
1503    async fn stream_pipeline_flushes_sink_on_cancel() {
1504        // #146 H16: a cooperative cancellation mid-run must stop polling, flush
1505        // the sink (so a Parquet footer / S3 multipart is completed rather than
1506        // orphaned), and return the partial result — NOT drop the run future,
1507        // which would flush nothing.
1508        use tokio_util::sync::CancellationToken;
1509
1510        // One page, then block forever — the only way out is the cancel token.
1511        let stream = Box::pin(async_stream::stream! {
1512            yield Ok(StreamPage {
1513                records: vec![json!({"id": 1}), json!({"id": 2})],
1514                bookmark: None,
1515            });
1516            futures::future::pending::<()>().await;
1517        });
1518        let sink = FlushTrackingSink::new();
1519
1520        let token = CancellationToken::new();
1521        let canceller = token.clone();
1522        tokio::spawn(async move {
1523            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1524            canceller.cancel();
1525        });
1526
1527        let result = run_stream(stream, &sink, RunStreamOptions::new().with_cancel(token))
1528            .await
1529            .expect("a cooperative cancel returns Ok with the partial result");
1530
1531        // The page written before cancellation survives, and the sink was
1532        // flushed so that output is durable.
1533        assert_eq!(result.records_written, 2);
1534        assert_eq!(sink.written().len(), 2);
1535        assert!(
1536            sink.flush_count() >= 1,
1537            "sink must be flushed on the cancel path so partial output is durable"
1538        );
1539    }
1540
1541    #[tokio::test]
1542    async fn stream_pipeline_empty() {
1543        let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
1544        let stream = futures::stream::iter(pages);
1545        let sink = MockSink::new();
1546
1547        let result = run_stream(stream, &sink, RunStreamOptions::new())
1548            .await
1549            .unwrap();
1550
1551        assert_eq!(result.records_written, 0);
1552    }
1553
1554    #[tokio::test]
1555    async fn stream_pipeline_skips_empty_pages() {
1556        let pages: Vec<Result<StreamPage, FaucetError>> = vec![
1557            Ok(StreamPage {
1558                records: vec![json!({"id": 1})],
1559                bookmark: None,
1560            }),
1561            Ok(StreamPage {
1562                records: vec![],
1563                bookmark: None,
1564            }),
1565            Ok(StreamPage {
1566                records: vec![json!({"id": 2})],
1567                bookmark: None,
1568            }),
1569        ];
1570        let stream = futures::stream::iter(pages);
1571        let sink = MockSink::new();
1572
1573        let result = run_stream(stream, &sink, RunStreamOptions::new())
1574            .await
1575            .unwrap();
1576
1577        assert_eq!(result.records_written, 2);
1578    }
1579
1580    #[tokio::test]
1581    async fn stream_pipeline_error_in_page_propagates() {
1582        let pages: Vec<Result<StreamPage, FaucetError>> = vec![
1583            Ok(StreamPage {
1584                records: vec![json!({"id": 1})],
1585                bookmark: None,
1586            }),
1587            Err(FaucetError::HttpStatus {
1588                status: 500,
1589                url: "https://example.com".into(),
1590                body: "Internal Server Error".into(),
1591            }),
1592        ];
1593        let stream = futures::stream::iter(pages);
1594        let sink = MockSink::new();
1595
1596        let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
1597        assert!(result.is_err());
1598        // First page was written before the error
1599        assert_eq!(sink.written().len(), 1);
1600    }
1601
1602    #[tokio::test]
1603    async fn stream_pipeline_sink_error_propagates() {
1604        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
1605            records: vec![json!({"id": 1})],
1606            bookmark: None,
1607        })];
1608        let stream = futures::stream::iter(pages);
1609        let sink = FailingSink;
1610
1611        let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
1612        assert!(result.is_err());
1613    }
1614
1615    #[tokio::test]
1616    async fn stream_pipeline_with_trait_object_sink() {
1617        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
1618            records: vec![json!({"id": 1})],
1619            bookmark: None,
1620        })];
1621        let stream = futures::stream::iter(pages);
1622        let sink: Box<dyn Sink> = Box::new(MockSink::new());
1623
1624        let result = run_stream(stream, sink.as_ref(), RunStreamOptions::new())
1625            .await
1626            .unwrap();
1627        assert_eq!(result.records_written, 1);
1628    }
1629
1630    #[tokio::test]
1631    async fn stream_pipeline_persists_bookmark_when_page_carries_one() {
1632        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
1633        let pages: Vec<Result<StreamPage, FaucetError>> = vec![
1634            Ok(StreamPage {
1635                records: vec![json!({"id": 1})],
1636                bookmark: None,
1637            }),
1638            Ok(StreamPage {
1639                records: vec![json!({"id": 2})],
1640                bookmark: Some(json!("checkpoint-final")),
1641            }),
1642        ];
1643        let stream = futures::stream::iter(pages);
1644        let sink = MockSink::new();
1645
1646        let result = run_stream(
1647            stream,
1648            &sink,
1649            RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
1650        )
1651        .await
1652        .unwrap();
1653
1654        assert_eq!(result.records_written, 2);
1655        assert_eq!(result.bookmark, Some(json!("checkpoint-final")));
1656        assert_eq!(
1657            store.get("k").await.unwrap(),
1658            Some(json!("checkpoint-final"))
1659        );
1660    }
1661
1662    #[tokio::test]
1663    async fn stream_pipeline_persists_per_page_bookmarks() {
1664        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
1665        let pages: Vec<Result<StreamPage, FaucetError>> = vec![
1666            Ok(StreamPage {
1667                records: vec![json!({"id": 1})],
1668                bookmark: Some(json!("tx-1")),
1669            }),
1670            Ok(StreamPage {
1671                records: vec![json!({"id": 2})],
1672                bookmark: Some(json!("tx-2")),
1673            }),
1674        ];
1675        let stream = futures::stream::iter(pages);
1676        let sink = MockSink::new();
1677
1678        run_stream(
1679            stream,
1680            &sink,
1681            RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
1682        )
1683        .await
1684        .unwrap();
1685
1686        // Latest per-page bookmark wins.
1687        assert_eq!(store.get("k").await.unwrap(), Some(json!("tx-2")));
1688    }
1689
1690    // ── State-store integration tests ───────────────────────────────────────
1691
1692    use crate::state::{FileStateStore, MemoryStateStore, StateStore};
1693    use std::sync::Arc;
1694    use tempfile::TempDir;
1695
1696    /// Source that opts into state persistence. It records the bookmark it
1697    /// received via `apply_start_bookmark` so tests can verify the pipeline
1698    /// pushed the stored value back into it on resume.
1699    struct StatefulSource {
1700        key: String,
1701        records: Vec<Value>,
1702        new_bookmark: Value,
1703        seen_bookmark: std::sync::Mutex<Option<Value>>,
1704    }
1705
1706    impl StatefulSource {
1707        fn new(key: &str, records: Vec<Value>, new_bookmark: Value) -> Self {
1708            Self {
1709                key: key.into(),
1710                records,
1711                new_bookmark,
1712                seen_bookmark: std::sync::Mutex::new(None),
1713            }
1714        }
1715        fn observed_start(&self) -> Option<Value> {
1716            self.seen_bookmark.lock().unwrap().clone()
1717        }
1718    }
1719
1720    #[async_trait]
1721    impl Source for StatefulSource {
1722        async fn fetch_with_context(
1723            &self,
1724            _ctx: &std::collections::HashMap<String, Value>,
1725        ) -> Result<Vec<Value>, FaucetError> {
1726            Ok(self.records.clone())
1727        }
1728        async fn fetch_with_context_incremental(
1729            &self,
1730            _ctx: &std::collections::HashMap<String, Value>,
1731        ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
1732            Ok((self.records.clone(), Some(self.new_bookmark.clone())))
1733        }
1734        fn state_key(&self) -> Option<String> {
1735            Some(self.key.clone())
1736        }
1737        async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
1738            *self.seen_bookmark.lock().unwrap() = Some(bookmark);
1739            Ok(())
1740        }
1741    }
1742
1743    #[tokio::test]
1744    async fn pipeline_with_state_store_persists_bookmark_after_sink() {
1745        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
1746        let source = StatefulSource::new(
1747            "github_issues",
1748            vec![json!({"id": 1, "ts": "2026-05-01"})],
1749            json!("2026-05-01"),
1750        );
1751        let sink = MockSink::new();
1752        let result = Pipeline::new(&source, &sink)
1753            .with_state_store(Arc::clone(&store))
1754            .run()
1755            .await
1756            .unwrap();
1757
1758        assert_eq!(result.records_written, 1);
1759        assert_eq!(result.bookmark, Some(json!("2026-05-01")));
1760        // Stored value matches what the source returned.
1761        let stored = store.get("github_issues").await.unwrap();
1762        assert_eq!(stored, Some(json!("2026-05-01")));
1763    }
1764
1765    #[tokio::test]
1766    async fn pipeline_with_state_store_resumes_from_stored_bookmark() {
1767        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
1768        store
1769            .put("github_issues", &json!("2026-04-30"))
1770            .await
1771            .unwrap();
1772
1773        let source =
1774            StatefulSource::new("github_issues", vec![json!({"id": 2})], json!("2026-05-01"));
1775        let sink = MockSink::new();
1776        Pipeline::new(&source, &sink)
1777            .with_state_store(Arc::clone(&store))
1778            .run()
1779            .await
1780            .unwrap();
1781
1782        // The pipeline pushed the previously-stored bookmark back into the source.
1783        assert_eq!(source.observed_start(), Some(json!("2026-04-30")));
1784        // And then overwrote it with the new value from this run.
1785        assert_eq!(
1786            store.get("github_issues").await.unwrap(),
1787            Some(json!("2026-05-01"))
1788        );
1789    }
1790
1791    #[tokio::test]
1792    async fn pipeline_with_state_store_does_not_persist_when_sink_fails() {
1793        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
1794        let source = StatefulSource::new("k", vec![json!({"id": 1})], json!("2026-05-01"));
1795        let sink = FailingSink;
1796
1797        let result = Pipeline::new(&source, &sink)
1798            .with_state_store(Arc::clone(&store))
1799            .run()
1800            .await;
1801        assert!(result.is_err());
1802        assert!(store.get("k").await.unwrap().is_none());
1803    }
1804
1805    #[tokio::test]
1806    async fn pipeline_with_state_store_no_state_key_means_no_persist() {
1807        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
1808        let source = IncrementalSource {
1809            records: vec![json!({"id": 1})],
1810            bookmark: json!("ignored"),
1811        };
1812        let sink = MockSink::new();
1813        Pipeline::new(&source, &sink)
1814            .with_state_store(Arc::clone(&store))
1815            .run()
1816            .await
1817            .unwrap();
1818        // IncrementalSource doesn't override state_key, so nothing was persisted.
1819        // Cross-check that no keys exist by trying a likely one.
1820        assert!(store.get("anything").await.unwrap().is_none());
1821    }
1822
1823    #[tokio::test]
1824    async fn pipeline_with_state_store_skips_persist_when_bookmark_is_none() {
1825        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
1826        struct NoBookmarkSource;
1827        #[async_trait]
1828        impl Source for NoBookmarkSource {
1829            async fn fetch_with_context(
1830                &self,
1831                _ctx: &std::collections::HashMap<String, Value>,
1832            ) -> Result<Vec<Value>, FaucetError> {
1833                Ok(vec![json!({"id": 1})])
1834            }
1835            fn state_key(&self) -> Option<String> {
1836                Some("k".into())
1837            }
1838        }
1839        let source = NoBookmarkSource;
1840        let sink = MockSink::new();
1841        Pipeline::new(&source, &sink)
1842            .with_state_store(Arc::clone(&store))
1843            .run()
1844            .await
1845            .unwrap();
1846        assert!(store.get("k").await.unwrap().is_none());
1847    }
1848
1849    // ── Pipeline::run drives stream_pages ──────────────────────────────────
1850
1851    /// A source with a custom `stream_pages` impl that yields three pages.
1852    /// Used to prove `Pipeline::run` drives the streaming path.
1853    struct PagedSource;
1854
1855    #[async_trait]
1856    impl Source for PagedSource {
1857        async fn fetch_with_context(
1858            &self,
1859            _ctx: &std::collections::HashMap<String, Value>,
1860        ) -> Result<Vec<Value>, FaucetError> {
1861            // Should never be called when stream_pages is overridden.
1862            unreachable!("Pipeline::run must drive stream_pages, not fetch_with_context");
1863        }
1864        fn stream_pages<'a>(
1865            &'a self,
1866            _ctx: &'a std::collections::HashMap<String, Value>,
1867            _batch_size: usize,
1868        ) -> std::pin::Pin<
1869            Box<dyn futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>,
1870        > {
1871            Box::pin(async_stream::try_stream! {
1872                yield StreamPage { records: vec![json!({"i": 1})], bookmark: None };
1873                yield StreamPage { records: vec![json!({"i": 2})], bookmark: None };
1874                yield StreamPage { records: vec![json!({"i": 3})], bookmark: Some(json!("final")) };
1875            })
1876        }
1877    }
1878
1879    /// Sink that counts how many distinct write_batch calls happen.
1880    struct CountingSink {
1881        calls: std::sync::Mutex<Vec<usize>>,
1882    }
1883
1884    impl CountingSink {
1885        fn new() -> Self {
1886            Self {
1887                calls: std::sync::Mutex::new(Vec::new()),
1888            }
1889        }
1890        fn call_count(&self) -> usize {
1891            self.calls.lock().unwrap().len()
1892        }
1893    }
1894
1895    #[async_trait]
1896    impl Sink for CountingSink {
1897        async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1898            self.calls.lock().unwrap().push(records.len());
1899            Ok(records.len())
1900        }
1901    }
1902
1903    #[tokio::test]
1904    async fn pipeline_run_drives_stream_pages() {
1905        let source = PagedSource;
1906        let sink = CountingSink::new();
1907
1908        let result = Pipeline::new(&source, &sink).run().await.unwrap();
1909
1910        // Three pages of one record each → three sink calls, three records.
1911        assert_eq!(sink.call_count(), 3);
1912        assert_eq!(result.records_written, 3);
1913        assert_eq!(result.bookmark, Some(json!("final")));
1914    }
1915
1916    #[tokio::test]
1917    async fn pipeline_with_file_state_store_round_trips_across_runs() {
1918        let dir = TempDir::new().unwrap();
1919        let store: Arc<dyn StateStore> = Arc::new(FileStateStore::new(dir.path()));
1920
1921        // Run 1: nothing stored yet, persist new bookmark.
1922        let s1 = StatefulSource::new("k", vec![json!({"i": 1})], json!("v1"));
1923        let sink1 = MockSink::new();
1924        Pipeline::new(&s1, &sink1)
1925            .with_state_store(Arc::clone(&store))
1926            .run()
1927            .await
1928            .unwrap();
1929        assert_eq!(s1.observed_start(), None);
1930        assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
1931
1932        // Run 2: resume from v1, persist v2.
1933        let s2 = StatefulSource::new("k", vec![json!({"i": 2})], json!("v2"));
1934        let sink2 = MockSink::new();
1935        Pipeline::new(&s2, &sink2)
1936            .with_state_store(Arc::clone(&store))
1937            .run()
1938            .await
1939            .unwrap();
1940        assert_eq!(s2.observed_start(), Some(json!("v1")));
1941        assert_eq!(store.get("k").await.unwrap(), Some(json!("v2")));
1942    }
1943
1944    #[tokio::test]
1945    #[allow(clippy::await_holding_lock)]
1946    async fn pipeline_run_increments_runs_total() {
1947        use crate::observability::decorator::source_tests::{LOCK, snapshotter};
1948        use metrics_util::debugging::DebugValue;
1949        let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
1950        let snap = snapshotter();
1951
1952        let source = MockSource(vec![json!({"i": 1})]);
1953        let sink = MockSink::new();
1954        let _ = Pipeline::new(&source, &sink)
1955            .with_name("test-pipeline")
1956            .with_row("rowA")
1957            .run()
1958            .await
1959            .unwrap();
1960
1961        let snapshot = snap.snapshot();
1962        let found = snapshot.into_vec().into_iter().any(
1963            |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
1964                key.key().name() == "faucet_pipeline_runs_total"
1965                    && key.key().labels().any(|l: &metrics::Label| {
1966                        l.key() == "pipeline" && l.value() == "test-pipeline"
1967                    })
1968                    && key
1969                        .key()
1970                        .labels()
1971                        .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowA")
1972                    && key
1973                        .key()
1974                        .labels()
1975                        .any(|l: &metrics::Label| l.key() == "status" && l.value() == "ok")
1976                    && matches!(v, DebugValue::Counter(c) if c >= 1)
1977            },
1978        );
1979        assert!(
1980            found,
1981            "expected faucet_pipeline_runs_total{{pipeline=test-pipeline, row=rowA, status=ok}}"
1982        );
1983    }
1984
1985    #[tokio::test]
1986    #[allow(clippy::await_holding_lock)]
1987    async fn pipeline_failure_attaches_kind_label_to_runs_total() {
1988        use crate::observability::decorator::source_tests::{LOCK, snapshotter};
1989        use metrics_util::debugging::DebugValue;
1990        let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
1991        let snap = snapshotter();
1992
1993        let source = FailingSource;
1994        let sink = MockSink::new();
1995        let _ = Pipeline::new(&source, &sink)
1996            .with_name("err-pipeline")
1997            .with_row("rowE")
1998            .run()
1999            .await;
2000
2001        let snapshot = snap.snapshot();
2002        let found = snapshot.into_vec().into_iter().any(
2003            |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
2004                key.key().name() == "faucet_pipeline_runs_total"
2005                    && key.key().labels().any(|l: &metrics::Label| {
2006                        l.key() == "pipeline" && l.value() == "err-pipeline"
2007                    })
2008                    && key
2009                        .key()
2010                        .labels()
2011                        .any(|l: &metrics::Label| l.key() == "status" && l.value() == "err")
2012                    && key
2013                        .key()
2014                        .labels()
2015                        .any(|l: &metrics::Label| l.key() == "kind" && l.value() == "Auth")
2016                    && matches!(v, DebugValue::Counter(c) if c >= 1)
2017            },
2018        );
2019        assert!(
2020            found,
2021            "expected faucet_pipeline_runs_total{{status=err, kind=Auth}} for failing source"
2022        );
2023    }
2024
2025    #[tokio::test]
2026    #[allow(clippy::await_holding_lock)]
2027    async fn pipeline_run_emits_start_time_gauge() {
2028        use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2029        use metrics_util::debugging::DebugValue;
2030        let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2031        let snap = snapshotter();
2032
2033        let source = MockSource(vec![json!({"i": 1})]);
2034        let sink = MockSink::new();
2035        let before = std::time::SystemTime::now()
2036            .duration_since(std::time::UNIX_EPOCH)
2037            .map(|d| d.as_secs_f64())
2038            .unwrap_or(0.0);
2039        let _ = Pipeline::new(&source, &sink)
2040            .with_name("start-time-pipeline")
2041            .with_row("rowS")
2042            .run()
2043            .await
2044            .unwrap();
2045
2046        let snapshot = snap.snapshot();
2047        let found = snapshot.into_vec().into_iter().any(
2048            |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
2049                if key.key().name() != "faucet_pipeline_start_time_unix_seconds" {
2050                    return false;
2051                }
2052                let labels_match = key.key().labels().any(|l: &metrics::Label| {
2053                    l.key() == "pipeline" && l.value() == "start-time-pipeline"
2054                }) && key
2055                    .key()
2056                    .labels()
2057                    .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowS");
2058                if !labels_match {
2059                    return false;
2060                }
2061                matches!(v, DebugValue::Gauge(g) if g.into_inner() >= before)
2062            },
2063        );
2064        assert!(
2065            found,
2066            "expected faucet_pipeline_start_time_unix_seconds gauge >= test-start timestamp"
2067        );
2068    }
2069
2070    #[tokio::test]
2071    #[allow(clippy::await_holding_lock)]
2072    async fn register_build_info_sets_version_gauge() {
2073        use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2074        use metrics_util::debugging::DebugValue;
2075        let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2076        let snap = snapshotter();
2077
2078        crate::observability::register_build_info();
2079
2080        let snapshot = snap.snapshot();
2081        let found = snapshot.into_vec().into_iter().any(
2082            |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
2083                key.key().name() == "faucet_build_info"
2084                    && key.key().labels().any(|l: &metrics::Label| {
2085                        l.key() == "version" && l.value() == env!("CARGO_PKG_VERSION")
2086                    })
2087                    && matches!(v, DebugValue::Gauge(g) if (g.into_inner() - 1.0).abs() < f64::EPSILON)
2088            },
2089        );
2090        assert!(
2091            found,
2092            "expected faucet_build_info{{version=CARGO_PKG_VERSION}} = 1.0 after register_build_info()"
2093        );
2094    }
2095
2096    // ── DLQ routing tests ──────────────────────────────────────────────────
2097
2098    use crate::dlq::{DlqConfig, OnBatchError};
2099
2100    /// Sink that returns mixed per-row outcomes: failure indices come from
2101    /// the constructor; everything else succeeds. Captures the rows that
2102    /// *would* have committed to the main sink.
2103    struct PartialSink {
2104        fail_indices: std::sync::Mutex<Vec<usize>>,
2105        committed: std::sync::Mutex<Vec<Value>>,
2106    }
2107
2108    impl PartialSink {
2109        fn new(fail_indices: Vec<usize>) -> Self {
2110            Self {
2111                fail_indices: std::sync::Mutex::new(fail_indices),
2112                committed: std::sync::Mutex::new(Vec::new()),
2113            }
2114        }
2115    }
2116
2117    #[async_trait]
2118    impl Sink for PartialSink {
2119        async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
2120            unreachable!("PartialSink only overrides write_batch_partial");
2121        }
2122        async fn write_batch_partial(
2123            &self,
2124            records: &[Value],
2125        ) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
2126            let fails: std::collections::HashSet<usize> =
2127                self.fail_indices.lock().unwrap().iter().copied().collect();
2128            let mut outcomes = Vec::with_capacity(records.len());
2129            for (i, rec) in records.iter().enumerate() {
2130                if fails.contains(&i) {
2131                    outcomes.push(Err(FaucetError::Sink(format!("row {i} rejected"))));
2132                } else {
2133                    self.committed.lock().unwrap().push(rec.clone());
2134                    outcomes.push(Ok(()));
2135                }
2136            }
2137            Ok(outcomes)
2138        }
2139    }
2140
2141    #[tokio::test]
2142    async fn dlq_routes_only_failed_rows_for_partial_success_sink() {
2143        let main = PartialSink::new(vec![1, 3]); // 4 rows, indices 1 and 3 fail
2144        let dlq = std::sync::Arc::new(MockSink::new());
2145        let dlq_cfg = DlqConfig::new(dlq.clone());
2146
2147        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2148            records: (0..4).map(|i| json!({"i": i})).collect(),
2149            bookmark: None,
2150        })];
2151        let stream = futures::stream::iter(pages);
2152        let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg))
2153            .await
2154            .unwrap();
2155
2156        assert_eq!(result.records_written, 2); // 0 and 2 committed
2157        assert_eq!(main.committed.lock().unwrap().len(), 2);
2158        let envelopes = dlq.0.lock().unwrap();
2159        assert_eq!(envelopes.len(), 2);
2160        assert_eq!(envelopes[0]["payload"]["i"], 1);
2161        assert_eq!(envelopes[0]["record_index"], 1);
2162        assert_eq!(envelopes[1]["payload"]["i"], 3);
2163        assert_eq!(envelopes[1]["record_index"], 3);
2164        let stats = result.dlq.unwrap();
2165        assert_eq!(stats.records_dlq, 2);
2166        assert_eq!(stats.pages_with_failures, 1);
2167    }
2168
2169    #[tokio::test]
2170    async fn dlq_propagate_policy_bubbles_outer_err() {
2171        let main = FailingSink;
2172        let dlq = std::sync::Arc::new(MockSink::new());
2173        let mut dlq_cfg = DlqConfig::new(dlq.clone());
2174        dlq_cfg.on_batch_error = OnBatchError::Propagate;
2175
2176        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2177            records: vec![json!({"i": 0}), json!({"i": 1})],
2178            bookmark: Some(json!("v1")),
2179        })];
2180        let stream = futures::stream::iter(pages);
2181        let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2182        let result = run_stream(
2183            stream,
2184            &main,
2185            RunStreamOptions::new()
2186                .with_dlq(dlq_cfg)
2187                .with_state(std::sync::Arc::clone(&store), "k"),
2188        )
2189        .await;
2190        assert!(matches!(result, Err(FaucetError::Sink(_))));
2191        assert!(dlq.0.lock().unwrap().is_empty());
2192        // Bookmark must NOT be persisted on a propagated failure.
2193        assert!(store.get("k").await.unwrap().is_none());
2194    }
2195
2196    #[tokio::test]
2197    async fn dlq_dlq_all_policy_routes_every_row_on_outer_err() {
2198        let main = FailingSink;
2199        let dlq = std::sync::Arc::new(MockSink::new());
2200        let mut dlq_cfg = DlqConfig::new(dlq.clone());
2201        dlq_cfg.on_batch_error = OnBatchError::DlqAll;
2202
2203        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2204            records: vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})],
2205            bookmark: Some(json!("v1")),
2206        })];
2207        let stream = futures::stream::iter(pages);
2208        let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2209        let result = run_stream(
2210            stream,
2211            &main,
2212            RunStreamOptions::new()
2213                .with_dlq(dlq_cfg)
2214                .with_state(std::sync::Arc::clone(&store), "k"),
2215        )
2216        .await
2217        .unwrap();
2218        assert_eq!(result.records_written, 0);
2219        {
2220            let envelopes = dlq.0.lock().unwrap();
2221            assert_eq!(envelopes.len(), 3);
2222            // Every envelope's error.message includes the underlying message.
2223            for env in envelopes.iter() {
2224                let msg = env["error"]["message"].as_str().unwrap();
2225                assert!(msg.contains("write failed"), "got: {msg}");
2226            }
2227        }
2228        assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
2229        assert_eq!(result.dlq.unwrap().records_dlq, 3);
2230    }
2231
2232    #[tokio::test]
2233    async fn dlq_per_page_budget_exceeded_aborts() {
2234        let main = PartialSink::new(vec![0, 1, 2]);
2235        let dlq = std::sync::Arc::new(MockSink::new());
2236        let mut dlq_cfg = DlqConfig::new(dlq.clone());
2237        dlq_cfg.max_failures_per_page = Some(2);
2238
2239        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2240            records: (0..3).map(|i| json!({"i": i})).collect(),
2241            bookmark: None,
2242        })];
2243        let stream = futures::stream::iter(pages);
2244        let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
2245        assert!(
2246            matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
2247            "got: {result:?}"
2248        );
2249    }
2250
2251    #[tokio::test]
2252    async fn dlq_total_budget_exceeded_aborts_on_later_page() {
2253        let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2254            Ok(StreamPage {
2255                records: (0..3).map(|i| json!({"i": i})).collect(),
2256                bookmark: None,
2257            }),
2258            Ok(StreamPage {
2259                records: (3..6).map(|i| json!({"i": i})).collect(),
2260                bookmark: None,
2261            }),
2262        ];
2263        // Fail every row across both pages.
2264        let main = PartialSink::new(vec![0, 1, 2]); // applied per page
2265        let dlq = std::sync::Arc::new(MockSink::new());
2266        let mut dlq_cfg = DlqConfig::new(dlq.clone());
2267        dlq_cfg.max_failures_total = Some(4);
2268
2269        let stream = futures::stream::iter(pages);
2270        let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
2271        assert!(
2272            matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
2273            "got: {result:?}"
2274        );
2275    }
2276
2277    #[tokio::test]
2278    async fn dlq_per_page_budget_exceeded_commits_page_before_aborting() {
2279        // M4 (#146): write_batch_partial already commits the survivors to the
2280        // main sink. When the per-page budget then trips, the run must finish
2281        // committing the page — route its failures to the DLQ and persist the
2282        // bookmark — BEFORE aborting, so the committed survivors do NOT
2283        // re-deliver on the next run and the failed rows are not lost.
2284        let main = PartialSink::new(vec![1, 2]); // rows 1,2 fail; row 0 commits
2285        let dlq = std::sync::Arc::new(MockSink::new());
2286        let mut dlq_cfg = DlqConfig::new(dlq.clone());
2287        dlq_cfg.max_failures_per_page = Some(1); // 2 failures > 1 → trips
2288
2289        let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2290        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2291            records: (0..3).map(|i| json!({ "i": i })).collect(),
2292            bookmark: Some(json!("v1")),
2293        })];
2294        let stream = futures::stream::iter(pages);
2295        let result = run_stream(
2296            stream,
2297            &main,
2298            RunStreamOptions::new()
2299                .with_dlq(dlq_cfg)
2300                .with_state(std::sync::Arc::clone(&store), "k"),
2301        )
2302        .await;
2303
2304        // Run still aborts with the budget error.
2305        assert!(
2306            matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
2307            "got: {result:?}"
2308        );
2309        // The survivor (row 0) was committed to the main sink.
2310        assert_eq!(main.committed.lock().unwrap().len(), 1);
2311        // The two failures were routed to the DLQ (not lost on abort).
2312        assert_eq!(dlq.0.lock().unwrap().len(), 2);
2313        // The bookmark was persisted, so the survivor will NOT re-deliver.
2314        assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
2315    }
2316
2317    #[tokio::test]
2318    async fn dlq_total_budget_exceeded_commits_tripping_page_before_aborting() {
2319        // M4 (#146): same guarantee for the cumulative total budget — the page
2320        // that crosses the threshold is committed fully (failures→DLQ, bookmark
2321        // persisted) before the run aborts.
2322        let main = PartialSink::new(vec![1, 2]); // rows 1,2 fail; row 0 commits
2323        let dlq = std::sync::Arc::new(MockSink::new());
2324        let mut dlq_cfg = DlqConfig::new(dlq.clone());
2325        dlq_cfg.max_failures_total = Some(1); // 2 failures > 1 → trips
2326
2327        let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2328        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2329            records: (0..3).map(|i| json!({ "i": i })).collect(),
2330            bookmark: Some(json!("v1")),
2331        })];
2332        let stream = futures::stream::iter(pages);
2333        let result = run_stream(
2334            stream,
2335            &main,
2336            RunStreamOptions::new()
2337                .with_dlq(dlq_cfg)
2338                .with_state(std::sync::Arc::clone(&store), "k"),
2339        )
2340        .await;
2341
2342        assert!(
2343            matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
2344            "got: {result:?}"
2345        );
2346        assert_eq!(main.committed.lock().unwrap().len(), 1);
2347        assert_eq!(dlq.0.lock().unwrap().len(), 2);
2348        assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
2349    }
2350
2351    /// DLQ sink that always fails. Used to assert the router does not
2352    /// recurse into itself.
2353    struct FailingDlqSink;
2354    #[async_trait]
2355    impl Sink for FailingDlqSink {
2356        async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
2357            Err(FaucetError::Sink("dlq disk full".into()))
2358        }
2359    }
2360
2361    /// DLQ sink that succeeds on write but fails on flush. Used to assert
2362    /// the router wraps DLQ flush errors and bails without persisting the
2363    /// bookmark.
2364    struct FailingFlushDlqSink {
2365        written: std::sync::Mutex<Vec<Value>>,
2366    }
2367    impl FailingFlushDlqSink {
2368        fn new() -> Self {
2369            Self {
2370                written: std::sync::Mutex::new(Vec::new()),
2371            }
2372        }
2373    }
2374    #[async_trait]
2375    impl Sink for FailingFlushDlqSink {
2376        async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2377            self.written.lock().unwrap().extend(records.iter().cloned());
2378            Ok(records.len())
2379        }
2380        async fn flush(&self) -> Result<(), FaucetError> {
2381            Err(FaucetError::Sink("dlq flush failed".into()))
2382        }
2383    }
2384
2385    #[tokio::test]
2386    async fn dlq_sink_failure_is_fatal_no_recursion() {
2387        let main = PartialSink::new(vec![0]);
2388        let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingDlqSink);
2389        let dlq_cfg = DlqConfig::new(dlq);
2390
2391        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2392            records: vec![json!({"i": 0}), json!({"i": 1})],
2393            bookmark: Some(json!("v1")),
2394        })];
2395        let stream = futures::stream::iter(pages);
2396        let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2397        let result = run_stream(
2398            stream,
2399            &main,
2400            RunStreamOptions::new()
2401                .with_dlq(dlq_cfg)
2402                .with_state(std::sync::Arc::clone(&store), "k"),
2403        )
2404        .await;
2405        assert!(
2406            matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink write failed")),
2407            "got: {result:?}"
2408        );
2409        assert!(store.get("k").await.unwrap().is_none());
2410    }
2411
2412    #[tokio::test]
2413    async fn dlq_bookmark_advances_only_after_both_flushes() {
2414        let main = PartialSink::new(vec![1]); // row 1 fails, row 0 commits
2415        let dlq = std::sync::Arc::new(MockSink::new());
2416        let dlq_cfg = DlqConfig::new(dlq.clone());
2417
2418        let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2419        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2420            records: vec![json!({"i": 0}), json!({"i": 1})],
2421            bookmark: Some(json!("v1")),
2422        })];
2423        let stream = futures::stream::iter(pages);
2424        run_stream(
2425            stream,
2426            &main,
2427            RunStreamOptions::new()
2428                .with_dlq(dlq_cfg)
2429                .with_state(std::sync::Arc::clone(&store), "k"),
2430        )
2431        .await
2432        .unwrap();
2433        assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
2434        assert_eq!(dlq.0.lock().unwrap().len(), 1);
2435        assert_eq!(main.committed.lock().unwrap().len(), 1);
2436    }
2437
2438    #[tokio::test]
2439    async fn dlq_disabled_pipeline_behaves_identically_to_today() {
2440        // Regression guard: omitting DLQ keeps existing behavior bit-identical.
2441        let main = MockSink::new();
2442        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2443            records: vec![json!({"i": 0}), json!({"i": 1})],
2444            bookmark: None,
2445        })];
2446        let stream = futures::stream::iter(pages);
2447        let result = run_stream(stream, &main, RunStreamOptions::new())
2448            .await
2449            .unwrap();
2450        assert_eq!(result.records_written, 2);
2451        assert!(result.dlq.is_none());
2452    }
2453
2454    #[tokio::test]
2455    async fn dlq_per_page_flush_failure_is_fatal_and_blocks_bookmark() {
2456        // Per-page flush path: page carries a bookmark, row 1 fails, the
2457        // DLQ write succeeds but the DLQ flush at the bookmark gate errors.
2458        // The pipeline must bail with "DLQ sink flush failed" and the
2459        // bookmark must NOT be persisted.
2460        let main = PartialSink::new(vec![1]);
2461        let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
2462        let dlq_cfg = DlqConfig::new(dlq);
2463
2464        let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2465        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2466            records: vec![json!({"i": 0}), json!({"i": 1})],
2467            bookmark: Some(json!("v1")),
2468        })];
2469        let stream = futures::stream::iter(pages);
2470        let result = run_stream(
2471            stream,
2472            &main,
2473            RunStreamOptions::new()
2474                .with_dlq(dlq_cfg)
2475                .with_state(std::sync::Arc::clone(&store), "k"),
2476        )
2477        .await;
2478        assert!(
2479            matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
2480            "got: {result:?}"
2481        );
2482        assert!(store.get("k").await.unwrap().is_none());
2483    }
2484
2485    #[tokio::test]
2486    async fn dlq_end_of_stream_flush_failure_is_fatal() {
2487        // End-of-stream flush path: no page carries a bookmark, but DLQ
2488        // received envelopes during the run. The final post-loop flush of
2489        // the DLQ sink errors. The pipeline must bail with "DLQ sink flush
2490        // failed".
2491        let main = PartialSink::new(vec![1]);
2492        let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
2493        let dlq_cfg = DlqConfig::new(dlq);
2494
2495        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2496            records: vec![json!({"i": 0}), json!({"i": 1})],
2497            bookmark: None,
2498        })];
2499        let stream = futures::stream::iter(pages);
2500        let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
2501        assert!(
2502            matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
2503            "got: {result:?}"
2504        );
2505    }
2506
2507    #[tokio::test]
2508    #[allow(clippy::await_holding_lock)]
2509    async fn dlq_emits_records_total_and_pages_total() {
2510        use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2511        use metrics_util::debugging::DebugValue;
2512
2513        let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2514        let snap = snapshotter();
2515
2516        let source = MockSource(vec![json!({"i": 0}), json!({"i": 1})]);
2517        let main = PartialSink::new(vec![1]);
2518        let dlq = std::sync::Arc::new(MockSink::new());
2519        let _ = Pipeline::new(&source, &main)
2520            .with_name("p_dlq_metrics")
2521            .with_row("r1")
2522            .with_dlq(DlqConfig::new(dlq.clone()))
2523            .run()
2524            .await
2525            .unwrap();
2526
2527        let snapshot = snap.snapshot();
2528        let mut saw_records = false;
2529        let mut saw_pages = false;
2530        for (k, _u, _d, v) in snapshot.into_vec() {
2531            let key = k.key();
2532            let labels = key.labels().collect::<Vec<_>>();
2533            let has = |k: &str, v: &str| labels.iter().any(|l| l.key() == k && l.value() == v);
2534            if key.name() == "faucet_sink_dlq_records_total"
2535                && has("pipeline", "p_dlq_metrics")
2536                && has("row", "r1")
2537                && matches!(v, DebugValue::Counter(c) if c >= 1)
2538            {
2539                saw_records = true;
2540            }
2541            if key.name() == "faucet_sink_dlq_pages_total"
2542                && has("pipeline", "p_dlq_metrics")
2543                && matches!(v, DebugValue::Counter(c) if c >= 1)
2544            {
2545                saw_pages = true;
2546            }
2547        }
2548        assert!(saw_records, "faucet_sink_dlq_records_total not emitted");
2549        assert!(saw_pages, "faucet_sink_dlq_pages_total not emitted");
2550    }
2551
2552    #[tokio::test]
2553    #[allow(clippy::await_holding_lock)]
2554    async fn dlq_budget_exceeded_emits_counter() {
2555        use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2556        use metrics_util::debugging::DebugValue;
2557
2558        let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2559        let snap = snapshotter();
2560
2561        let source = MockSource((0..3).map(|i| json!({"i": i})).collect());
2562        let main = PartialSink::new(vec![0, 1, 2]);
2563        let dlq = std::sync::Arc::new(MockSink::new());
2564        let mut cfg = DlqConfig::new(dlq);
2565        cfg.max_failures_per_page = Some(1);
2566        let _ = Pipeline::new(&source, &main)
2567            .with_name("p_budget")
2568            .with_dlq(cfg)
2569            .run()
2570            .await;
2571
2572        let snapshot = snap.snapshot();
2573        let saw = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
2574            k.key().name() == "faucet_sink_dlq_budget_exceeded_total"
2575                && k.key()
2576                    .labels()
2577                    .any(|l| l.key() == "scope" && l.value() == "per_page")
2578                && matches!(v, DebugValue::Counter(c) if c >= 1)
2579        });
2580        assert!(saw, "faucet_sink_dlq_budget_exceeded_total not emitted");
2581    }
2582
2583    #[tokio::test]
2584    async fn pipeline_run_with_dlq_routes_partial_failures_end_to_end() {
2585        // Source: 3 records. Main sink: fails index 1. DLQ: in-memory.
2586        let source = MockSource(vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})]);
2587        let main = PartialSink::new(vec![1]);
2588        let dlq = std::sync::Arc::new(MockSink::new());
2589
2590        let result = Pipeline::new(&source, &main)
2591            .with_dlq(DlqConfig::new(dlq.clone()))
2592            .run()
2593            .await
2594            .unwrap();
2595
2596        assert_eq!(result.records_written, 2);
2597        let stats = result.dlq.unwrap();
2598        assert_eq!(stats.records_dlq, 1);
2599        {
2600            let dlq_records = dlq.0.lock().unwrap();
2601            assert_eq!(dlq_records.len(), 1);
2602        }
2603    }
2604
2605    // ── Quality routing tests ──────────────────────────────────────────────
2606
2607    #[cfg(feature = "quality")]
2608    #[tokio::test]
2609    async fn quality_quarantines_to_dlq_and_writes_survivors() {
2610        use crate::dlq::DlqConfig;
2611        use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
2612
2613        let main = Arc::new(MockSink::new());
2614        let dlq_sink = Arc::new(MockSink::new());
2615        let spec = QualitySpec {
2616            record: vec![RecordCheck::NotNull {
2617                field: "id".into(),
2618                treat_missing_as_null: true,
2619                on_failure: OnFailure::Quarantine,
2620            }],
2621            batch: vec![],
2622        };
2623        let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
2624        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2625            records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
2626            bookmark: None,
2627        })];
2628        let opts = RunStreamOptions::new()
2629            .with_dlq(DlqConfig::new(dlq_sink.clone()))
2630            .with_quality(quality);
2631        let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
2632            .await
2633            .unwrap();
2634
2635        assert_eq!(result.records_written, 2); // survivors
2636        assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
2637        // one quarantined record reached the DLQ with a QualityFailure envelope
2638        let dlq = dlq_sink.written();
2639        assert_eq!(dlq.len(), 1);
2640        assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
2641        assert_eq!(result.dlq.unwrap().records_dlq, 1);
2642    }
2643
2644    #[cfg(feature = "quality")]
2645    #[tokio::test]
2646    #[allow(clippy::await_holding_lock)]
2647    async fn quality_only_page_emits_quality_reason() {
2648        // A page whose DLQ traffic is entirely quality-sourced (the main sink
2649        // accepts every survivor) must label `faucet_sink_dlq_pages_total`
2650        // with `reason="quality"`, not `partial`.
2651        use crate::dlq::DlqConfig;
2652        use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2653        use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
2654        use metrics_util::debugging::DebugValue;
2655
2656        let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2657        let snap = snapshotter();
2658
2659        // MockSink accepts everything → no sink-side failures, so the only DLQ
2660        // traffic comes from the quality quarantine.
2661        let main = Arc::new(MockSink::new());
2662        let dlq_sink = Arc::new(MockSink::new());
2663        let spec = QualitySpec {
2664            record: vec![RecordCheck::NotNull {
2665                field: "id".into(),
2666                treat_missing_as_null: true,
2667                on_failure: OnFailure::Quarantine,
2668            }],
2669            batch: vec![],
2670        };
2671        let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
2672        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2673            records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
2674            bookmark: None,
2675        })];
2676        let opts = RunStreamOptions::new()
2677            .with_name("p_quality_reason")
2678            .with_dlq(DlqConfig::new(dlq_sink.clone()))
2679            .with_quality(quality);
2680        let _ = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
2681            .await
2682            .unwrap();
2683
2684        let snapshot = snap.snapshot();
2685        let saw_quality_reason = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
2686            k.key().name() == "faucet_sink_dlq_pages_total"
2687                && k.key()
2688                    .labels()
2689                    .any(|l| l.key() == "pipeline" && l.value() == "p_quality_reason")
2690                && k.key()
2691                    .labels()
2692                    .any(|l| l.key() == "reason" && l.value() == "quality")
2693                && matches!(v, DebugValue::Counter(c) if c >= 1)
2694        });
2695        assert!(
2696            saw_quality_reason,
2697            "expected faucet_sink_dlq_pages_total with reason=\"quality\""
2698        );
2699    }
2700
2701    #[cfg(feature = "quality")]
2702    #[tokio::test]
2703    async fn quality_abort_fails_run() {
2704        use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
2705        let main = MockSink::new();
2706        let spec = QualitySpec {
2707            record: vec![],
2708            batch: vec![BatchCheck::RowCount {
2709                min: Some(5),
2710                max: None,
2711                on_failure: OnFailure::Abort,
2712            }],
2713        };
2714        let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
2715        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2716            records: vec![json!({"id": 1})],
2717            bookmark: None,
2718        })];
2719        let opts = RunStreamOptions::new().with_quality(quality);
2720        let result = run_stream(futures::stream::iter(pages), &main, opts).await;
2721        assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
2722    }
2723
2724    #[cfg(feature = "quality")]
2725    #[tokio::test]
2726    async fn quality_quarantine_without_dlq_is_rejected() {
2727        use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
2728        let main = MockSink::new();
2729        let spec = QualitySpec {
2730            record: vec![RecordCheck::NotNull {
2731                field: "id".into(),
2732                treat_missing_as_null: true,
2733                on_failure: OnFailure::Quarantine,
2734            }],
2735            batch: vec![],
2736        };
2737        let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
2738        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2739            records: vec![json!({"id": null})],
2740            bookmark: None,
2741        })];
2742        // No .with_dlq(...) — must be rejected up front.
2743        let opts = RunStreamOptions::new().with_quality(quality);
2744        let result = run_stream(futures::stream::iter(pages), &main, opts).await;
2745        assert!(matches!(result, Err(FaucetError::Config(_))));
2746    }
2747
2748    /// Sink whose write_batch_partial fails every Nth record; drives the
2749    /// error-rate signal. Requires a DLQ in run_stream.
2750    struct FlakySink {
2751        every: usize,
2752        calls: std::sync::Mutex<Vec<usize>>,
2753    }
2754    impl FlakySink {
2755        fn new(every: usize) -> Self {
2756            Self {
2757                every,
2758                calls: std::sync::Mutex::new(Vec::new()),
2759            }
2760        }
2761        fn call_sizes(&self) -> Vec<usize> {
2762            self.calls.lock().unwrap().clone()
2763        }
2764    }
2765    #[async_trait]
2766    impl Sink for FlakySink {
2767        async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2768            Ok(records.len())
2769        }
2770        async fn write_batch_partial(
2771            &self,
2772            records: &[Value],
2773        ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
2774            self.calls.lock().unwrap().push(records.len());
2775            Ok(records
2776                .iter()
2777                .enumerate()
2778                .map(|(i, _)| {
2779                    if (i + 1) % self.every == 0 {
2780                        Err(FaucetError::Sink("synthetic".into()))
2781                    } else {
2782                        Ok(())
2783                    }
2784                })
2785                .collect())
2786        }
2787    }
2788
2789    #[tokio::test]
2790    async fn adaptive_shrinks_under_errors_on_dlq_path() {
2791        use crate::adaptive::AdaptiveBatchConfig;
2792        use crate::dlq::{DlqConfig, OnBatchError};
2793        // Three pages of 400 records each, matching the pattern used by
2794        // adaptive_shrinks_under_latency_target_then_smaller_chunks. After
2795        // page 1 (single 400-record chunk, 25% error rate > threshold 0.1),
2796        // the controller shrinks and subsequent pages get smaller sub-batches.
2797        let mk = || StreamPage {
2798            records: (0..400).map(|i| json!({"i": i})).collect(),
2799            bookmark: None,
2800        };
2801        let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
2802        let sink = FlakySink::new(4); // 25% error rate > threshold 0.1
2803        let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
2804        let dlq = DlqConfig {
2805            sink: dlq_sink,
2806            on_batch_error: OnBatchError::Propagate,
2807            max_failures_per_page: None,
2808            max_failures_total: None,
2809            include_original_payload: true,
2810        };
2811        let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
2812            "enabled": true, "min": 50, "max": 400,
2813            "decrease_factor": 0.5, "cooldown_batches": 0, "error_threshold": 0.1
2814        }))
2815        .unwrap();
2816        let opts = RunStreamOptions::new().with_dlq(dlq).with_adaptive(cfg);
2817        let result = run_stream(stream, &sink, opts).await.unwrap();
2818        // 3 × 400 = 1200 records total; FlakySink(4) fails every 4th record
2819        // per-chunk (floor(n/4)), so exact counts depend on chunk sizes due to
2820        // integer arithmetic. With the controller shrinking under 25% error
2821        // rate: page 1 = one 400-record chunk (300 written, 100 DLQ); pages
2822        // 2–3 = smaller sub-batches; overall >≈75% of 1200 commit and ~25% go
2823        // to the DLQ.
2824        assert!(
2825            result.records_written >= 900,
2826            "expected ≥900 written, got {}",
2827            result.records_written
2828        );
2829        let sizes = sink.call_sizes();
2830        assert_eq!(sizes[0], 400, "first chunk is the full page");
2831        assert!(
2832            sizes.last().unwrap() < &400,
2833            "controller should shrink under errors: {sizes:?}"
2834        );
2835        assert!(
2836            result.dlq.unwrap().records_dlq >= 250,
2837            "expected ≥250 DLQ records"
2838        );
2839    }
2840
2841    // ── Adaptive batch-size tests ──────────────────────────────────────────
2842
2843    /// A sink that records each write_batch call's size and reports a fixed
2844    /// per-call latency, so we can assert the adaptive controller resliced.
2845    struct RecordingSink {
2846        calls: std::sync::Mutex<Vec<usize>>,
2847        latency: std::time::Duration,
2848    }
2849    impl RecordingSink {
2850        fn new(latency_ms: u64) -> Self {
2851            Self {
2852                calls: std::sync::Mutex::new(Vec::new()),
2853                latency: std::time::Duration::from_millis(latency_ms),
2854            }
2855        }
2856        fn call_sizes(&self) -> Vec<usize> {
2857            self.calls.lock().unwrap().clone()
2858        }
2859    }
2860    #[async_trait]
2861    impl Sink for RecordingSink {
2862        async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2863            tokio::time::sleep(self.latency).await;
2864            self.calls.lock().unwrap().push(records.len());
2865            Ok(records.len())
2866        }
2867    }
2868
2869    #[tokio::test]
2870    async fn adaptive_reslices_non_dlq_page_into_subbatches() {
2871        use crate::adaptive::AdaptiveBatchConfig;
2872        let page = StreamPage {
2873            records: (0..1000).map(|i| json!({ "i": i })).collect(),
2874            bookmark: None,
2875        };
2876        let stream = futures::stream::iter(vec![Ok(page)]);
2877        let sink = RecordingSink::new(0);
2878        let cfg: AdaptiveBatchConfig =
2879            serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
2880        let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
2881            .await
2882            .unwrap();
2883        assert_eq!(result.records_written, 1000);
2884        // current starts at min(max, page_len)=1000 → one chunk (no regression).
2885        assert_eq!(sink.call_sizes(), vec![1000]);
2886    }
2887
2888    #[tokio::test]
2889    async fn adaptive_shrinks_under_latency_target_then_smaller_chunks() {
2890        use crate::adaptive::AdaptiveBatchConfig;
2891        let mk = || StreamPage {
2892            records: (0..400).map(|i| json!({"i": i})).collect(),
2893            bookmark: None,
2894        };
2895        let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
2896        let sink = RecordingSink::new(50);
2897        let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
2898            "enabled": true, "min": 50, "max": 400,
2899            "decrease_factor": 0.5, "cooldown_batches": 0,
2900            "target_latency_ms": 10, "latency_window": 1
2901        }))
2902        .unwrap();
2903        let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
2904            .await
2905            .unwrap();
2906        assert_eq!(result.records_written, 1200);
2907        let sizes = sink.call_sizes();
2908        assert_eq!(sizes[0], 400);
2909        assert!(
2910            sizes.last().unwrap() < &400,
2911            "controller should have shrunk: {sizes:?}"
2912        );
2913    }
2914
2915    #[tokio::test]
2916    #[allow(clippy::await_holding_lock)]
2917    async fn adaptive_emits_batch_size_and_adjustments_metrics() {
2918        // Mirror the same LOCK+snapshotter pattern used by
2919        // `pipeline_run_increments_runs_total` and `dlq_emits_records_total_and_pages_total`.
2920        use crate::adaptive::AdaptiveBatchConfig;
2921        use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2922        use metrics_util::debugging::DebugValue;
2923
2924        let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2925        let snap = snapshotter();
2926
2927        // Three pages of 400 records. RecordingSink(50) reports 50ms latency.
2928        // Config: target_latency_ms=10, latency_window=1, cooldown_batches=0 so
2929        // p50 (50ms) > 10*1.2=12ms on every batch → controller shrinks each time,
2930        // guaranteeing at least one `faucet_pipeline_adaptive_batch_adjustments_total`
2931        // is emitted.
2932        let mk = || StreamPage {
2933            records: (0..400).map(|i| json!({"i": i})).collect(),
2934            bookmark: None,
2935        };
2936        let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
2937        let sink = RecordingSink::new(50);
2938        let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
2939            "enabled": true, "min": 50, "max": 400,
2940            "decrease_factor": 0.5, "cooldown_batches": 0,
2941            "target_latency_ms": 10, "latency_window": 1
2942        }))
2943        .unwrap();
2944
2945        let _ = run_stream(
2946            stream,
2947            &sink,
2948            RunStreamOptions::new()
2949                .with_adaptive(cfg)
2950                .with_name("p")
2951                .with_row("r"),
2952        )
2953        .await
2954        .unwrap();
2955
2956        let snapshot = snap.snapshot();
2957        let mut saw_batch_size = false;
2958        let mut saw_adjustments = false;
2959        for (k, _u, _d, v) in snapshot.into_vec() {
2960            let key = k.key();
2961            let labels = key.labels().collect::<Vec<_>>();
2962            let has = |k: &str, val: &str| labels.iter().any(|l| l.key() == k && l.value() == val);
2963
2964            if key.name() == "faucet_pipeline_adaptive_batch_size"
2965                && has("pipeline", "p")
2966                && has("row", "r")
2967                && matches!(v, DebugValue::Gauge(_))
2968            {
2969                saw_batch_size = true;
2970            }
2971            if key.name() == "faucet_pipeline_adaptive_batch_adjustments_total"
2972                && has("pipeline", "p")
2973                && has("row", "r")
2974                && matches!(v, DebugValue::Counter(c) if c >= 1)
2975            {
2976                saw_adjustments = true;
2977            }
2978        }
2979        assert!(
2980            saw_batch_size,
2981            "expected faucet_pipeline_adaptive_batch_size gauge with pipeline=p, row=r"
2982        );
2983        assert!(
2984            saw_adjustments,
2985            "expected faucet_pipeline_adaptive_batch_adjustments_total counter with pipeline=p, row=r"
2986        );
2987    }
2988
2989    // ── run_stream: exactly-once + DLQ incompatibility gate ─────────────────
2990
2991    #[tokio::test]
2992    async fn exactly_once_rejects_dlq() {
2993        // Exactly-once must reject a configured DLQ (incompatible in this
2994        // version): the gate fires before any page is polled.
2995        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2996        let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
2997        let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2998        let opts = eo_opts(store, "k", 0).with_dlq(DlqConfig::new(dlq_sink));
2999        let r = run_stream(
3000            futures::stream::iter(pages),
3001            &IdempotentMockSink::new(),
3002            opts,
3003        )
3004        .await;
3005        assert!(
3006            matches!(&r, Err(FaucetError::Config(m)) if m.contains("not compatible with a DLQ")),
3007            "got: {r:?}"
3008        );
3009    }
3010
3011    // ── run_stream: exactly-once page with records but no bookmark ──────────
3012
3013    #[tokio::test]
3014    async fn exactly_once_writes_unbookmarked_page_at_least_once() {
3015        // Under exactly-once, a page that carries records but NO bookmark is
3016        // not individually checkpointed: it falls through to a plain
3017        // `write_batch` (at-least-once for that page) and is NOT idempotently
3018        // tokened.
3019        let sink = IdempotentMockSink::new();
3020        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3021        let pages = vec![Ok(StreamPage {
3022            records: vec![json!({"id": 1}), json!({"id": 2})],
3023            bookmark: None,
3024        })];
3025        let r = run_stream(
3026            futures::stream::iter(pages),
3027            &sink,
3028            eo_opts(store.clone(), "k", 0),
3029        )
3030        .await
3031        .unwrap();
3032        assert_eq!(r.records_written, 2);
3033        assert_eq!(r.bookmark, None);
3034        // Rows were written via the plain (non-idempotent) write_batch path, so
3035        // no commit token was recorded for the scope.
3036        assert_eq!(sink.last_committed_token("k").await.unwrap(), None);
3037        // Nothing was persisted to the state store (no bookmark to checkpoint).
3038        assert!(store.get("k").await.unwrap().is_none());
3039        assert_eq!(sink.rows(), vec![json!({"id": 1}), json!({"id": 2})]);
3040    }
3041
3042    // ── DLQ-with-bookmark success path: persist bookmark after routing ──────
3043
3044    #[tokio::test]
3045    async fn dlq_with_bookmark_persists_after_routing_failures() {
3046        // A bookmark-carrying page whose main sink reports one per-row failure:
3047        // survivors commit, the failed row reaches the DLQ, the DLQ is flushed,
3048        // and the bookmark is persisted to the state store.
3049        let main = PartialSink::new(vec![1]); // 2 rows, index 1 fails
3050        let dlq = std::sync::Arc::new(MockSink::new());
3051        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3052        let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3053            records: vec![json!({"i": 0}), json!({"i": 1})],
3054            bookmark: Some(json!("ckpt")),
3055        })];
3056        let result = run_stream(
3057            futures::stream::iter(pages),
3058            &main,
3059            RunStreamOptions::new()
3060                .with_dlq(DlqConfig::new(dlq.clone()))
3061                .with_state(Arc::clone(&store), "k"),
3062        )
3063        .await
3064        .unwrap();
3065
3066        assert_eq!(result.records_written, 1); // row 0 committed
3067        assert_eq!(result.bookmark, Some(json!("ckpt")));
3068        // Bookmark was persisted after the page was made durable.
3069        assert_eq!(store.get("k").await.unwrap(), Some(json!("ckpt")));
3070        // The single failed row reached the DLQ.
3071        let envelopes = dlq.0.lock().unwrap();
3072        assert_eq!(envelopes.len(), 1);
3073        assert_eq!(envelopes[0]["payload"]["i"], 1);
3074    }
3075
3076    // ── run_stream drives exactly-once with resume from start_seq ───────────
3077
3078    #[tokio::test]
3079    async fn exactly_once_resumes_sequence_from_start_seq() {
3080        // A resume run starts at start_seq (the persisted seq) and continues
3081        // numbering tokens from there; the next bookmark-carrying page is
3082        // committed at seq = start_seq + 1.
3083        let sink = IdempotentMockSink::new();
3084        let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3085        let pages = vec![Ok(StreamPage {
3086            records: vec![json!({"id": 9})],
3087            bookmark: Some(json!("bm-after-resume")),
3088        })];
3089        let r = run_stream(
3090            futures::stream::iter(pages),
3091            &sink,
3092            eo_opts(store.clone(), "eo_key", 7),
3093        )
3094        .await
3095        .unwrap();
3096
3097        assert_eq!(r.records_written, 1);
3098        let (bm, seq) =
3099            crate::idempotency::unwrap_state(&store.get("eo_key").await.unwrap().unwrap());
3100        assert_eq!(bm, Some(json!("bm-after-resume")));
3101        assert_eq!(seq, 8, "sequence resumes at start_seq + 1");
3102        assert_eq!(
3103            sink.last_committed_token("eo_key").await.unwrap(),
3104            Some(crate::idempotency::format_token(8))
3105        );
3106    }
3107
3108    // ── Pipeline::run with quality wired through the builder ────────────────
3109
3110    #[cfg(feature = "quality")]
3111    #[tokio::test]
3112    async fn pipeline_run_with_quality_aborts_on_failed_batch_check() {
3113        use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
3114        let source = MockSource(vec![json!({"id": 1})]);
3115        let main = MockSink::new();
3116        let spec = QualitySpec {
3117            record: vec![],
3118            batch: vec![BatchCheck::RowCount {
3119                min: Some(5),
3120                max: None,
3121                on_failure: OnFailure::Abort,
3122            }],
3123        };
3124        let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3125        let result = Pipeline::new(&source, &main)
3126            .with_quality(quality)
3127            .run()
3128            .await;
3129        assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
3130        // The abort fired before the sink committed.
3131        assert!(main.written().is_empty());
3132    }
3133
3134    // ── Pipeline::run with adaptive wired through the builder ───────────────
3135
3136    #[tokio::test]
3137    async fn pipeline_run_with_adaptive_reslices_page() {
3138        use crate::adaptive::AdaptiveBatchConfig;
3139        let source = MockSource((0..1000).map(|i| json!({ "i": i })).collect());
3140        let sink = MockSink::new();
3141        let cfg: AdaptiveBatchConfig =
3142            serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
3143        let result = Pipeline::new(&source, &sink)
3144            .with_adaptive(cfg)
3145            .run()
3146            .await
3147            .unwrap();
3148        assert_eq!(result.records_written, 1000);
3149        assert_eq!(sink.written().len(), 1000);
3150    }
3151
3152    // ── Adaptive no-op warning for per-record sinks (jsonl/csv/stdout) ──────
3153
3154    #[tokio::test]
3155    async fn adaptive_noop_sink_name_is_handled() {
3156        use crate::adaptive::AdaptiveBatchConfig;
3157
3158        // A sink whose connector_name() is one of the per-record names triggers
3159        // the one-shot "no-op for this per-record sink" info path.
3160        struct JsonlNamedSink(std::sync::Mutex<usize>);
3161        #[async_trait]
3162        impl Sink for JsonlNamedSink {
3163            async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3164                *self.0.lock().unwrap() += records.len();
3165                Ok(records.len())
3166            }
3167            fn connector_name(&self) -> &'static str {
3168                "jsonl"
3169            }
3170        }
3171
3172        let page = StreamPage {
3173            records: (0..10).map(|i| json!({ "i": i })).collect(),
3174            bookmark: None,
3175        };
3176        let stream = futures::stream::iter(vec![Ok(page)]);
3177        let sink = JsonlNamedSink(std::sync::Mutex::new(0));
3178        let cfg: AdaptiveBatchConfig =
3179            serde_json::from_value(json!({"enabled": true, "min": 5, "max": 10})).unwrap();
3180        let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
3181            .await
3182            .unwrap();
3183        assert_eq!(result.records_written, 10);
3184        assert_eq!(*sink.0.lock().unwrap(), 10);
3185    }
3186}