1use 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
52pub const DEFAULT_BATCH_SIZE: usize = 1000;
58
59pub const MAX_BATCH_SIZE: usize = 1_000_000;
64
65pub 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#[derive(Debug, Clone, Default)]
96pub struct StreamPage {
97 pub records: Vec<Value>,
99 pub bookmark: Option<Value>,
101}
102
103#[derive(Debug, Clone)]
105pub struct PipelineResult {
106 pub records_written: usize,
108 pub bookmark: Option<Value>,
116 pub dlq: Option<DlqStats>,
118}
119
120pub 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 #[cfg(feature = "contract")]
135 contract: Option<Arc<crate::contract::CompiledContract>>,
136 #[cfg(feature = "masking")]
137 masking: Option<Arc<crate::masking::CompiledMasking>>,
138 adaptive: Option<crate::adaptive::AdaptiveBatchConfig>,
139 cancel: Option<tokio_util::sync::CancellationToken>,
140 delivery: crate::idempotency::DeliveryMode,
141 resilience: Option<crate::resilience::ResiliencePolicy>,
142 schema_drift: Option<crate::drift::SchemaDriftPolicy>,
143}
144
145impl<'a, So: Source + ?Sized, Si: Sink + ?Sized> Pipeline<'a, So, Si> {
146 pub fn new(source: &'a So, sink: &'a Si) -> Self {
148 Self {
149 source,
150 sink,
151 state_store: None,
152 name: None,
153 row: None,
154 run_id: None,
155 dlq: None,
156 #[cfg(feature = "quality")]
157 quality: None,
158 #[cfg(feature = "contract")]
159 contract: None,
160 #[cfg(feature = "masking")]
161 masking: None,
162 adaptive: None,
163 cancel: None,
164 delivery: crate::idempotency::DeliveryMode::AtLeastOnce,
165 resilience: None,
166 schema_drift: None,
167 }
168 }
169
170 pub fn with_state_store(mut self, store: Arc<dyn StateStore>) -> Self {
184 self.state_store = Some(store);
185 self
186 }
187
188 pub fn with_name(mut self, name: impl Into<String>) -> Self {
191 self.name = Some(name.into());
192 self
193 }
194
195 pub fn with_row(mut self, row: impl Into<String>) -> Self {
198 self.row = Some(row.into());
199 self
200 }
201
202 pub fn with_run_id(mut self, run_id: impl Into<String>) -> Self {
206 self.run_id = Some(run_id.into());
207 self
208 }
209
210 pub fn with_dlq(mut self, dlq: DlqConfig) -> Self {
212 self.dlq = Some(dlq);
213 self
214 }
215
216 #[cfg(feature = "quality")]
219 pub fn with_quality(mut self, quality: Arc<crate::quality::CompiledQuality>) -> Self {
220 self.quality = Some(quality);
221 self
222 }
223
224 #[cfg(feature = "contract")]
227 pub fn with_contract(mut self, contract: Arc<crate::contract::CompiledContract>) -> Self {
228 self.contract = Some(contract);
229 self
230 }
231
232 #[cfg(feature = "masking")]
236 pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
237 self.masking = Some(masking);
238 self
239 }
240
241 pub fn with_adaptive(mut self, cfg: crate::adaptive::AdaptiveBatchConfig) -> Self {
245 self.adaptive = Some(cfg);
246 self
247 }
248
249 pub fn with_cancel(mut self, cancel: tokio_util::sync::CancellationToken) -> Self {
254 self.cancel = Some(cancel);
255 self
256 }
257
258 pub fn with_delivery(mut self, mode: crate::idempotency::DeliveryMode) -> Self {
263 self.delivery = mode;
264 self
265 }
266
267 pub fn with_resilience(mut self, policy: crate::resilience::ResiliencePolicy) -> Self {
269 self.resilience = Some(policy);
270 self
271 }
272
273 pub fn with_schema_drift(mut self, policy: crate::drift::SchemaDriftPolicy) -> Self {
276 self.schema_drift = Some(policy);
277 self
278 }
279
280 pub async fn run(&self) -> Result<PipelineResult, FaucetError> {
294 use crate::observability::{
295 DurationGuard, InstrumentedSink, InstrumentedSource, InstrumentedStateStore, Labels,
296 };
297 use metrics::{Label, SharedString, counter, gauge};
298 use tracing::Instrument;
299
300 let name = self.name.clone().unwrap_or_else(|| "unnamed".to_string());
302 let row = self.row.clone().unwrap_or_default();
303 let run_id = self
304 .run_id
305 .clone()
306 .unwrap_or_else(|| uuid::Uuid::now_v7().to_string());
307 let obs_labels = Labels::new(name.clone(), row.clone(), run_id.clone());
308
309 let wrapped_source = InstrumentedSource::new(self.source, obs_labels.clone());
311 let wrapped_sink = InstrumentedSink::new(self.sink, obs_labels.clone());
312 let wrapped_state_store: Option<Arc<dyn StateStore>> = self.state_store.as_ref().map(|s| {
313 Arc::new(InstrumentedStateStore::new(
314 Arc::clone(s),
315 obs_labels.clone(),
316 )) as Arc<dyn StateStore>
317 });
318
319 let span = tracing::info_span!(
322 "faucet.pipeline.run",
323 pipeline = %name,
324 row = %row,
325 run_id = %run_id,
326 source = %wrapped_source.connector_name(),
327 sink = %wrapped_sink.connector_name(),
328 );
329
330 let base_labels: Vec<Label> = vec![
332 Label::new("pipeline", SharedString::from(name.clone())),
333 Label::new("row", SharedString::from(row.clone())),
334 ];
335 let run_labels: Vec<Label> = {
336 let mut v = base_labels.clone();
337 v.push(Label::new(
338 "source",
339 SharedString::from(wrapped_source.connector_name().to_string()),
340 ));
341 v.push(Label::new(
342 "sink",
343 SharedString::from(wrapped_sink.connector_name().to_string()),
344 ));
345 v
346 };
347
348 struct InFlightGuard(Vec<Label>);
350 impl Drop for InFlightGuard {
351 fn drop(&mut self) {
352 gauge!("faucet_pipeline_in_flight", self.0.clone()).decrement(1.0);
353 }
354 }
355 gauge!("faucet_pipeline_in_flight", base_labels.clone()).increment(1.0);
356 let _in_flight = InFlightGuard(base_labels.clone());
357
358 let start_unix = std::time::SystemTime::now()
361 .duration_since(std::time::UNIX_EPOCH)
362 .map(|d| d.as_secs_f64())
363 .unwrap_or(0.0);
364 gauge!(
365 "faucet_pipeline_start_time_unix_seconds",
366 base_labels.clone()
367 )
368 .set(start_unix);
369
370 let _run_timer =
372 DurationGuard::new("faucet_pipeline_run_duration_seconds", run_labels.clone());
373
374 let result = async {
376 let state_key = self.source.state_key();
379 let mut start_seq = 0u64;
380 if let (Some(store), Some(key)) = (wrapped_state_store.as_ref(), state_key.as_ref()) {
381 validate_state_key(key)?;
382 if let Some(prior) = store.get(key).await? {
383 if self.delivery == crate::idempotency::DeliveryMode::ExactlyOnce {
384 let (bookmark, seq) = crate::idempotency::unwrap_state(&prior);
385 start_seq = seq;
386 if let Some(bm) = bookmark {
387 wrapped_source.apply_start_bookmark(bm).await?;
388 }
389 } else {
390 wrapped_source.apply_start_bookmark(prior).await?;
391 }
392 }
393 }
394
395 if self.delivery == crate::idempotency::DeliveryMode::ExactlyOnce
406 && wrapped_sink.supports_idempotent_writes()
407 && wrapped_source.replay_guarantee()
408 == crate::idempotency::ReplayGuarantee::Deterministic
409 && let Some(key) = state_key.as_ref()
410 && let Some(token) = wrapped_sink.last_committed_token(key).await?
411 && let Some((sink_seq, Some(bm))) = crate::idempotency::parse_token_parts(&token)
412 && sink_seq > start_seq
413 {
414 wrapped_source.apply_start_bookmark(bm).await?;
415 start_seq = sink_seq;
416 }
417
418 #[cfg(feature = "arrow")]
426 {
427 let columnar_ok = wrapped_source.supports_columnar()
428 && wrapped_sink.supports_columnar()
429 && self.dlq.is_none()
430 && self.delivery == crate::idempotency::DeliveryMode::AtLeastOnce
431 && self.schema_drift.is_none()
432 && self.adaptive.is_none()
433 && self.resilience.is_none();
434 #[cfg(feature = "quality")]
435 let columnar_ok = columnar_ok && self.quality.is_none();
436 #[cfg(feature = "contract")]
437 let columnar_ok = columnar_ok && self.contract.is_none();
438 #[cfg(feature = "masking")]
439 let columnar_ok = columnar_ok && self.masking.is_none();
440 if columnar_ok {
441 let state = match (wrapped_state_store.clone(), state_key.clone()) {
442 (Some(store), Some(key)) => Some((store, key)),
443 _ => None,
444 };
445 return run_stream_columnar(
446 &wrapped_source,
447 &wrapped_sink,
448 state,
449 self.cancel.clone(),
450 &name,
451 &row,
452 &run_id,
453 )
454 .await;
455 }
456 }
457
458 let ctx = std::collections::HashMap::new();
459 let pages = wrapped_source.stream_pages(&ctx, DEFAULT_BATCH_SIZE);
460
461 let mut opts = RunStreamOptions::new()
462 .with_name(name.clone())
463 .with_row(row.clone())
464 .with_run_id(run_id.clone());
465 if let (Some(store), Some(key)) = (wrapped_state_store.clone(), state_key) {
466 opts = opts.with_state(store, key);
467 }
468 if let Some(dlq) = self.dlq.clone() {
469 opts = opts.with_dlq(dlq);
470 }
471 #[cfg(feature = "quality")]
472 if let Some(q) = self.quality.clone() {
473 opts = opts.with_quality(q);
474 }
475 #[cfg(feature = "contract")]
476 if let Some(c) = self.contract.clone() {
477 opts = opts.with_contract(c);
478 }
479 #[cfg(feature = "masking")]
480 if let Some(m) = self.masking.clone() {
481 opts = opts.with_masking(m);
482 }
483 if let Some(ad) = self.adaptive.clone() {
484 opts = opts.with_adaptive(ad);
485 }
486 if let Some(cancel) = self.cancel.clone() {
487 opts = opts.with_cancel(cancel);
488 }
489 if let Some(policy) = self.resilience.clone() {
490 opts = opts.with_resilience(policy);
491 }
492 if let Some(p) = self.schema_drift {
493 opts = opts.with_schema_drift(p);
494 }
495 opts = opts
496 .with_delivery(self.delivery)
497 .with_start_seq(start_seq)
498 .with_replay_guarantee(wrapped_source.replay_guarantee());
499
500 run_stream(pages, &wrapped_sink, opts).await
501 }
502 .instrument(span)
503 .await;
504
505 let status = if result.is_ok() { "ok" } else { "err" };
509 let mut final_labels = run_labels;
510 final_labels.push(Label::new("status", SharedString::const_str(status)));
511 if let Err(ref e) = result {
512 final_labels.push(Label::new(
513 "kind",
514 SharedString::const_str(crate::observability::decorator::error_kind(e)),
515 ));
516 }
517 counter!("faucet_pipeline_runs_total", final_labels).increment(1);
518
519 result
520 }
521}
522
523#[cfg(feature = "arrow")]
534async fn run_stream_columnar<S, Si>(
535 source: &S,
536 sink: &Si,
537 state: Option<(Arc<dyn StateStore>, String)>,
538 cancel: Option<tokio_util::sync::CancellationToken>,
539 pipeline: &str,
540 row: &str,
541 run_id: &str,
542) -> Result<PipelineResult, FaucetError>
543where
544 S: crate::Source + ?Sized,
545 Si: Sink + ?Sized,
546{
547 use futures::StreamExt;
548 use metrics::{Label, SharedString, counter};
549
550 let labels = |connector: &str| -> Vec<Label> {
551 vec![
552 Label::new("pipeline", SharedString::from(pipeline.to_string())),
553 Label::new("row", SharedString::from(row.to_string())),
554 Label::new("connector", SharedString::from(connector.to_string())),
555 ]
556 };
557 let src_labels = labels(source.connector_name());
558 let sink_labels = labels(sink.connector_name());
559 let _ = run_id; let ctx = std::collections::HashMap::new();
562 let mut batches = source.stream_batches(&ctx, DEFAULT_BATCH_SIZE);
563 let mut records_written = 0usize;
564 let mut last_bookmark: Option<Value> = None;
565
566 loop {
567 let page = match &cancel {
570 Some(token) => {
571 tokio::select! {
572 biased;
573 _ = token.cancelled() => break,
574 p = batches.next() => p,
575 }
576 }
577 None => batches.next().await,
578 };
579 let Some(page) = page else { break };
580 let page = page?;
581 let rows = page.num_rows();
582 counter!("faucet_source_records_total", src_labels.clone()).increment(rows as u64);
583
584 if rows > 0 {
585 let n = sink.write_batch_columnar(&page.batch).await?;
586 records_written += n;
587 counter!("faucet_sink_records_total", sink_labels.clone()).increment(n as u64);
588 counter!("faucet_sink_writes_total", sink_labels.clone()).increment(1);
589 }
590
591 if let Some(bm) = page.bookmark {
594 sink.flush().await?;
595 if let Some((store, key)) = state.as_ref() {
596 store.put(key, &bm).await?;
597 }
598 last_bookmark = Some(bm);
599 }
600 }
601
602 sink.flush().await?;
604 Ok(PipelineResult {
605 records_written,
606 bookmark: last_bookmark,
607 dlq: None,
608 })
609}
610
611pub async fn run_stream<S, Si>(
629 mut pages: S,
630 sink: &Si,
631 options: RunStreamOptions,
632) -> Result<PipelineResult, FaucetError>
633where
634 S: Stream<Item = Result<StreamPage, FaucetError>> + Unpin,
635 Si: Sink + ?Sized,
636{
637 use crate::dlq::{DlqReason, DlqStats, OnBatchError, build_envelope};
638
639 let state_store = options.state_store.clone();
640 let state_key = options.state_key.clone();
641 let pipeline_name = options.pipeline_name.unwrap_or_else(|| "unnamed".into());
642 let row = options.row.unwrap_or_default();
643 let run_id = options.run_id.unwrap_or_default();
644 let dlq = options.dlq.clone();
645 let cancel = options.cancel.clone();
646
647 #[cfg(feature = "quality")]
648 let quality = options.quality.clone();
649
650 #[cfg(feature = "quality")]
652 if let Some(q) = quality.as_ref()
653 && q.requires_dlq()
654 && dlq.is_none()
655 {
656 return Err(FaucetError::Config(
657 "quality: on_failure 'quarantine'/'quarantine_batch' requires a DLQ sink".into(),
658 ));
659 }
660
661 #[cfg(feature = "contract")]
662 let contract = options.contract.clone();
663 #[cfg(feature = "contract")]
665 if let Some(c) = contract.as_ref()
666 && c.requires_dlq()
667 && dlq.is_none()
668 {
669 return Err(FaucetError::Config(
670 "contract: on_breach 'quarantine' requires a DLQ sink".into(),
671 ));
672 }
673 #[cfg(feature = "contract")]
675 let mut warned_contract_breach = false;
676
677 #[cfg(feature = "masking")]
681 let masking = options.masking.clone();
682
683 let schema_drift = options.schema_drift;
685 if let Some(p) = schema_drift.as_ref()
687 && p.requires_dlq()
688 && dlq.is_none()
689 {
690 return Err(FaucetError::Config(
691 "schema: on_drift 'quarantine' (or on_incompatible 'quarantine') requires a DLQ sink"
692 .into(),
693 ));
694 }
695 let mut dest_schema_cache: Option<Option<Value>> = None;
699 let mut warned_drift_inert = false;
700
701 if let Some(key) = state_key.as_ref() {
702 validate_state_key(key)?;
703 }
704
705 let mechanism: Option<crate::idempotency::EffectivelyOnceMechanism> =
713 if options.delivery == crate::idempotency::DeliveryMode::ExactlyOnce {
714 let replay_ok = options
715 .replay
716 .is_none_or(|r| r == crate::idempotency::ReplayGuarantee::Deterministic);
717 if sink.supports_idempotent_writes() && replay_ok {
718 if state_store.is_none() || state_key.is_none() {
719 return Err(FaucetError::Config(
720 "delivery: exactly_once (atomic watermark) requires a state store".into(),
721 ));
722 }
723 if dlq.is_some() {
724 return Err(FaucetError::Config(
725 "delivery: exactly_once (atomic watermark) is not compatible with a DLQ \
726 in this version"
727 .into(),
728 ));
729 }
730 Some(crate::idempotency::EffectivelyOnceMechanism::AtomicWatermark)
731 } else if sink.dedups_by_key() {
732 Some(crate::idempotency::EffectivelyOnceMechanism::KeyedUpsert)
733 } else if sink.supports_idempotent_writes() {
734 return Err(FaucetError::Config(format!(
737 "delivery: exactly_once — the source does not replay deterministically from \
738 a bookmark, so the atomic-watermark mechanism cannot be used; configure \
739 `write_mode: upsert` with a `key` on sink '{}' for keyed-upsert \
740 effectively-once instead",
741 sink.connector_name()
742 )));
743 } else {
744 return Err(FaucetError::Config(format!(
745 "delivery: exactly_once requires an idempotent (atomic-watermark) sink or a \
746 sink configured to dedup by key (`write_mode: upsert` + `key`), but '{}' \
747 provides neither",
748 sink.connector_name()
749 )));
750 }
751 } else {
752 None
753 };
754 let exactly_once =
758 mechanism == Some(crate::idempotency::EffectivelyOnceMechanism::AtomicWatermark);
759 let scope = state_key.clone().unwrap_or_default();
760 let mut next_seq = options.start_seq;
761 let committed_seq = if exactly_once {
762 sink.last_committed_token(&scope)
763 .await?
764 .and_then(|t| crate::idempotency::parse_token(&t))
765 .unwrap_or(0)
766 } else {
767 0
768 };
769
770 let mut records_written = 0usize;
771 let mut last_bookmark: Option<Value> = None;
772 let mut dlq_stats = DlqStats::default();
773
774 let adaptive_cfg = options.adaptive.clone().filter(|c| c.enabled);
775 if let Some(cfg) = adaptive_cfg.as_ref() {
779 cfg.validate()?;
780 }
781 let mut controller: Option<crate::adaptive::AimdController> = None;
782 let mut warned_noop_sink = false;
783 let mut warned_poison_drop = false;
785
786 let sink_name = sink.connector_name();
787 let dlq_sink_name = dlq.as_ref().map(|d| d.sink.connector_name()).unwrap_or("");
788
789 let mut cancelled = false;
800
801 let resilience = options.resilience.clone();
808 let retry_policy = resilience.as_ref().map(|r| r.retry.clone());
809 let mut breaker = resilience
810 .as_ref()
811 .and_then(|r| r.circuit_breaker)
812 .map(|cb| {
813 (
814 crate::resilience::CircuitBreaker::new(cb.consecutive_failures),
815 cb.cooldown,
816 )
817 });
818 let poison = resilience.as_ref().and_then(|r| r.poison);
820
821 macro_rules! with_retry {
834 ($op_label:literal, $op:expr) => {
835 match &retry_policy {
836 Some(p) => {
837 let m = crate::resilience::RetryMetrics {
838 pipeline: pipeline_name.to_string(),
839 row: row.to_string(),
840 op: $op_label,
841 };
842 crate::resilience::execute_with_policy_metered(p, cancel.as_ref(), &m, || $op)
843 .await
844 }
845 None => $op.await,
846 }
847 };
848 }
849
850 macro_rules! with_retry_write {
861 ($op_label:literal, $op:expr) => {
862 if retry_policy.is_some() && sink.supports_idempotent_writes() {
863 with_retry!($op_label, $op)
864 } else {
865 $op.await
866 }
867 };
868 }
869
870 let loop_result: Result<(), FaucetError> = async {
871 loop {
872 let page = match &cancel {
876 Some(token) => tokio::select! {
877 biased;
878 _ = token.cancelled() => {
879 cancelled = true;
880 break;
881 }
882 p = std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)) => p,
883 },
884 None => std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)).await,
885 };
886 match page {
887 Some(Ok(page)) => {
888 if page.records.is_empty() && page.bookmark.is_none() {
889 continue;
890 }
891
892 #[cfg(feature = "masking")]
898 let page = if let Some(m) = masking.as_ref() {
899 let labels =
900 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
901 let outcome = crate::observability::instrumented_apply_masking(
902 page.records,
903 m,
904 &labels,
905 );
906 StreamPage {
907 records: outcome.records,
908 bookmark: page.bookmark,
909 }
910 } else {
911 page
912 };
913
914 let page_len = page.records.len();
921
922 #[cfg(feature = "quality")]
924 let (records, quality_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
925 if let Some(q) = quality.as_ref() {
926 let labels =
927 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
928 let outcome = crate::observability::instrumented_apply_quality(
929 page.records,
930 q,
931 &labels,
932 )?;
933 let quarantined_idx: std::collections::HashSet<usize> =
934 outcome.quarantined.iter().map(|qr| qr.page_index).collect();
935 let envelopes: Vec<Value> = outcome
936 .quarantined
937 .iter()
938 .map(|qr| {
939 let err = FaucetError::QualityFailure {
940 check: qr.check.to_string(),
941 message: qr.message.clone(),
942 };
943 build_envelope(
947 &qr.record,
948 &err,
949 DlqReason::Quality,
950 sink_name,
951 &pipeline_name,
952 &row,
953 qr.page_index,
954 )
955 })
956 .collect();
957 let survivor_idx: Vec<usize> =
958 (0..page_len).filter(|i| !quarantined_idx.contains(i)).collect();
959 (outcome.survivors, envelopes, survivor_idx)
960 } else {
961 (page.records, Vec::new(), (0..page_len).collect())
962 };
963 #[cfg(not(feature = "quality"))]
964 let (records, quality_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
965 (page.records, Vec::new(), (0..page_len).collect());
966
967 #[cfg(feature = "contract")]
974 let (records, contract_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
975 if let Some(c) = contract.as_ref() {
976 let labels =
977 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
978 let outcome = crate::observability::instrumented_apply_contract(
979 records, c, &labels,
980 )?;
981 if !outcome.warned.is_empty() && !warned_contract_breach {
982 tracing::warn!(
983 version = %c.version,
984 breaches = outcome.warned.len(),
985 first = %outcome.warned[0].describe(),
986 "contract: breaching records written unchanged \
987 (on_breach=warn); this warning fires once per run"
988 );
989 warned_contract_breach = true;
990 }
991 let envelopes: Vec<Value> = outcome
992 .quarantined
993 .iter()
994 .map(|vr| {
995 let err = FaucetError::ContractViolation {
996 version: c.version.clone(),
997 message: vr.violation.describe(),
998 };
999 let page_index = page_indices
1009 .get(vr.violation.page_index)
1010 .copied()
1011 .unwrap_or(vr.violation.page_index);
1012 build_envelope(
1013 &vr.record,
1014 &err,
1015 DlqReason::Contract,
1016 sink_name,
1017 &pipeline_name,
1018 &row,
1019 page_index,
1020 )
1021 })
1022 .collect();
1023 let contract_quarantined: std::collections::HashSet<usize> = outcome
1029 .quarantined
1030 .iter()
1031 .map(|vr| vr.violation.page_index)
1032 .collect();
1033 let survivor_idx: Vec<usize> = page_indices
1034 .iter()
1035 .enumerate()
1036 .filter(|(pos, _)| !contract_quarantined.contains(pos))
1037 .map(|(_, orig)| *orig)
1038 .collect();
1039 (outcome.survivors, envelopes, survivor_idx)
1040 } else {
1041 (records, Vec::new(), page_indices)
1042 };
1043
1044 let mut drift_envelopes: Vec<Value> = Vec::new();
1046 let (records, drift_abort): (Vec<Value>, Option<FaucetError>) =
1047 if let Some(policy) = schema_drift.as_ref().filter(|_| !records.is_empty()) {
1048 if dest_schema_cache.is_none() {
1050 dest_schema_cache = Some(sink.current_schema().await?);
1051 }
1052 let dest = dest_schema_cache.as_ref().and_then(|o| o.as_ref());
1053 match dest {
1054 None => {
1055 if !warned_drift_inert {
1056 tracing::info!(
1057 connector = sink_name,
1058 "schema-drift: sink reports no destination schema; \
1059 drift handling is inert this run"
1060 );
1061 warned_drift_inert = true;
1062 }
1063 (records, None)
1064 }
1065 Some(dest) => {
1066 let inferred = crate::schema::infer_schema(&records);
1067 let diff = crate::drift::diff_schema(
1068 dest,
1069 &inferred,
1070 policy.allow_widening,
1071 );
1072 if diff.is_empty() {
1073 (records, None)
1074 } else {
1075 let dest_owned = dest.clone();
1079 apply_drift_policy(
1080 policy,
1081 &diff,
1082 &dest_owned,
1083 records,
1084 &page_indices,
1085 sink,
1086 sink_name,
1087 &pipeline_name,
1088 &row,
1089 &mut dest_schema_cache,
1090 &mut drift_envelopes,
1091 )
1092 .await?
1093 }
1094 }
1095 }
1096 } else {
1097 (records, None)
1098 };
1099 let quality_envelopes = {
1103 let mut q = quality_envelopes;
1104 #[cfg(feature = "contract")]
1105 q.extend(contract_envelopes);
1106 q.append(&mut drift_envelopes);
1107 q
1108 };
1109 let mut drift_abort = drift_abort;
1119 if dlq.is_none()
1120 && let Some(e) = drift_abort.take()
1121 {
1122 return Err(e);
1123 }
1124
1125 let page = StreamPage {
1126 records,
1127 bookmark: page.bookmark,
1128 };
1129
1130 if let Some(ref dlq_cfg) = dlq {
1131 use metrics::{Label, SharedString, counter};
1133 let metric_labels: Vec<Label> = vec![
1134 Label::new("pipeline", SharedString::from(pipeline_name.clone())),
1135 Label::new("row", SharedString::from(row.clone())),
1136 Label::new("connector", SharedString::from(sink_name.to_string())),
1137 Label::new(
1138 "dlq_connector",
1139 SharedString::from(dlq_sink_name.to_string()),
1140 ),
1141 ];
1142 let span = tracing::info_span!(
1143 "faucet.dlq.route",
1144 pipeline = %pipeline_name,
1145 row = %row,
1146 run_id = %run_id,
1147 connector = %sink_name,
1148 dlq_connector = %dlq_sink_name,
1149 );
1150 let _enter = span.enter();
1151
1152 let mut envelopes: Vec<Value> = Vec::new();
1156 let mut page_success = 0usize;
1157 let mut outer_err_recovered = false;
1158 let mut had_per_row_sink_failure = false;
1163 let records_len = page.records.len();
1164 let mut offset = 0usize;
1165 while offset < records_len {
1166 let size = match adaptive_cfg.as_ref() {
1167 Some(cfg) => {
1168 let ctrl = controller.get_or_insert_with(|| {
1169 crate::adaptive::AimdController::new(cfg, records_len)
1170 });
1171 ctrl.current().max(1).min(records_len - offset)
1172 }
1173 None => records_len - offset, };
1175 if adaptive_cfg.is_some() {
1176 maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
1177 }
1178 let chunk = &page.records[offset..offset + size];
1179 let t0 = std::time::Instant::now();
1180 let chunk_outcomes_result =
1186 with_retry_write!("sink_write", sink.write_batch_partial(chunk));
1187 let latency = t0.elapsed();
1188 let (mut chunk_outcomes, chunk_synthesized): (
1197 Vec<crate::RowOutcome>,
1198 bool,
1199 ) = match chunk_outcomes_result {
1200 Ok(o) => (o, false),
1201 Err(e) => match dlq_cfg.on_batch_error {
1202 OnBatchError::Propagate => return Err(e),
1203 OnBatchError::DlqAll => {
1204 outer_err_recovered = true;
1205 let msg = e.to_string();
1206 let synth = (0..chunk.len())
1207 .map(|_| Err(FaucetError::Sink(msg.clone())))
1208 .collect();
1209 (synth, true)
1210 }
1211 },
1212 };
1213
1214 if let Some(pp) = poison
1223 && !chunk_synthesized
1224 {
1225 let mut attempt = 1u32; while attempt < pp.max_row_attempts {
1227 let failing: Vec<usize> = chunk_outcomes
1228 .iter()
1229 .enumerate()
1230 .filter_map(|(j, o)| match o {
1231 Err(e)
1232 if retry_policy
1233 .as_ref()
1234 .map(|p| p.is_retriable(e))
1235 .unwrap_or(false) =>
1236 {
1237 Some(j)
1238 }
1239 _ => None,
1240 })
1241 .collect();
1242 if failing.is_empty() {
1243 break;
1244 }
1245 let subset: Vec<Value> =
1246 failing.iter().map(|&j| chunk[j].clone()).collect();
1247 let retried = sink.write_batch_partial(&subset).await?;
1255 let mut retried = retried.into_iter();
1259 for &j in failing.iter() {
1260 chunk_outcomes[j] = retried.next().unwrap_or(Ok(()));
1261 }
1262 attempt += 1;
1263 }
1264 }
1265
1266 let mut chunk_errors = 0usize;
1267 let mut poison_dlq = 0u64;
1273 let mut poison_drop = 0u64;
1274 for (j, outcome) in chunk_outcomes.iter().enumerate() {
1275 match outcome {
1276 Ok(()) => page_success += 1,
1277 Err(err) => {
1278 let action = poison
1283 .map(|pp| pp.action)
1284 .unwrap_or(crate::resilience::PoisonAction::Dlq);
1285 match action {
1286 crate::resilience::PoisonAction::Fail => {
1287 crate::observability::resilience::poison_rows(
1288 &pipeline_name,
1289 &row,
1290 "fail",
1291 1,
1292 );
1293 return Err(FaucetError::Sink(format!(
1294 "poison-pill row failed permanently: {err}"
1295 )));
1296 }
1297 crate::resilience::PoisonAction::Drop => {
1298 poison_drop += 1;
1301 if !warned_poison_drop {
1302 tracing::warn!(
1303 "poison-pill: dropping permanently-failing row(s) (action=drop); this warning fires once per run"
1304 );
1305 warned_poison_drop = true;
1306 }
1307 }
1308 crate::resilience::PoisonAction::Dlq => {
1309 poison_dlq += 1;
1310 chunk_errors += 1;
1311 if !chunk_synthesized {
1312 had_per_row_sink_failure = true;
1313 }
1314 envelopes.push(build_envelope(
1315 &chunk[j],
1316 err,
1317 DlqReason::Partial,
1318 sink_name,
1319 &pipeline_name,
1320 &row,
1321 offset + j,
1322 ));
1323 }
1324 }
1325 }
1326 }
1327 }
1328 if poison.is_some() {
1332 crate::observability::resilience::poison_rows(
1333 &pipeline_name,
1334 &row,
1335 "dlq",
1336 poison_dlq,
1337 );
1338 crate::observability::resilience::poison_rows(
1339 &pipeline_name,
1340 &row,
1341 "drop",
1342 poison_drop,
1343 );
1344 }
1345 if let Some(ctrl) = controller.as_mut() {
1346 let adj = ctrl.observe(crate::adaptive::Observation {
1347 batch_len: chunk.len(),
1348 errors: chunk_errors,
1349 latency,
1350 });
1351 emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
1352 }
1353 offset += size;
1354 }
1355 #[cfg(feature = "quality")]
1360 let quality_count = quality_envelopes.len();
1361 #[cfg(not(feature = "quality"))]
1362 let quality_count = 0usize;
1363 envelopes.splice(0..0, quality_envelopes);
1364 let page_failures = envelopes.len();
1365
1366 let mut budget_error: Option<FaucetError> = None;
1379 let mut circuit_error: Option<FaucetError> = None;
1388 if let Some((b, cooldown)) = breaker.as_mut() {
1389 if records_len > 0 && page_success == 0 {
1390 if b.record_failure() {
1391 crate::observability::resilience::circuit_opened(
1392 &pipeline_name,
1393 &row,
1394 );
1395 circuit_error = Some(FaucetError::CircuitOpen {
1396 failures: b.consecutive(),
1397 cooldown: *cooldown,
1398 });
1399 }
1400 } else if page_success > 0 {
1401 b.record_success();
1402 }
1403 }
1404 if let Some(limit) = dlq_cfg.max_failures_per_page
1405 && page_failures > limit
1406 {
1407 let mut lbl = metric_labels.clone();
1408 lbl.retain(|l| l.key() != "dlq_connector");
1409 lbl.push(Label::new("scope", SharedString::const_str("per_page")));
1410 counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
1411 budget_error = Some(FaucetError::Sink(format!(
1412 "DLQ per-page budget exceeded: {page_failures} > {limit}"
1413 )));
1414 }
1415 let new_total = dlq_stats.records_dlq + page_failures;
1416 if budget_error.is_none()
1417 && let Some(limit) = dlq_cfg.max_failures_total
1418 && new_total > limit
1419 {
1420 let mut lbl = metric_labels.clone();
1421 lbl.retain(|l| l.key() != "dlq_connector");
1422 lbl.push(Label::new("scope", SharedString::const_str("total")));
1423 counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
1424 budget_error = Some(FaucetError::Sink(format!(
1425 "DLQ total budget exceeded: {new_total} > {limit}"
1426 )));
1427 }
1428
1429 if !envelopes.is_empty() {
1431 let _dlq_write_timer = crate::observability::DurationGuard::new(
1432 "faucet_sink_dlq_write_duration_seconds",
1433 metric_labels.clone(),
1434 );
1435 dlq_cfg.sink.write_batch(&envelopes).await.map_err(|e| {
1436 let mut lbl = metric_labels.clone();
1437 lbl.push(Label::new(
1438 "kind",
1439 SharedString::const_str(
1440 crate::observability::decorator::error_kind(&e),
1441 ),
1442 ));
1443 counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1444 FaucetError::Sink(format!("DLQ sink write failed: {e}"))
1445 })?;
1446 dlq_stats.records_dlq += page_failures;
1447 dlq_stats.pages_with_failures += 1;
1448
1449 let reason_label = if had_per_row_sink_failure {
1468 DlqReason::Partial.as_str()
1469 } else if outer_err_recovered {
1470 DlqReason::DlqAll.as_str()
1471 } else if page_failures > quality_count {
1472 DlqReason::Partial.as_str()
1473 } else {
1474 DlqReason::Quality.as_str()
1475 };
1476 counter!("faucet_sink_dlq_records_total", metric_labels.clone())
1477 .increment(page_failures as u64);
1478 let mut page_labels = metric_labels.clone();
1479 page_labels
1480 .push(Label::new("reason", SharedString::const_str(reason_label)));
1481 counter!("faucet_sink_dlq_pages_total", page_labels).increment(1);
1482 }
1483
1484 records_written += page_success;
1485
1486 if let Some(bookmark) = page.bookmark {
1487 with_retry!("flush", sink.flush())?;
1493 let _dlq_flush_timer = crate::observability::DurationGuard::new(
1494 "faucet_sink_dlq_flush_duration_seconds",
1495 metric_labels.clone(),
1496 );
1497 with_retry!("flush", dlq_cfg.sink.flush()).map_err(|e| {
1498 let mut lbl = metric_labels.clone();
1499 lbl.push(Label::new(
1500 "kind",
1501 SharedString::const_str(
1502 crate::observability::decorator::error_kind(&e),
1503 ),
1504 ));
1505 counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1506 FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
1507 })?;
1508 let bm_labels =
1509 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1510 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1511 if let (Some(store), Some(key)) =
1512 (state_store.as_ref(), state_key.as_ref())
1513 {
1514 with_retry!("state_put", store.put(key, &bookmark))?;
1515 }
1516 last_bookmark = Some(bookmark);
1517 }
1518
1519 if let Some(e) = budget_error {
1526 return Err(e);
1527 }
1528 if let Some(e) = circuit_error {
1530 return Err(e);
1531 }
1532 if let Some(e) = drift_abort {
1537 return Err(e);
1538 }
1539 } else if exactly_once {
1540 if let Some(bookmark) = page.bookmark {
1544 next_seq += 1;
1545 let token = crate::idempotency::format_token_with_bookmark(
1552 next_seq,
1553 Some(&bookmark),
1554 );
1555 if next_seq <= committed_seq {
1556 use metrics::{Label, SharedString, counter};
1560 let skip_labels: Vec<Label> = vec![
1561 Label::new(
1562 "pipeline",
1563 SharedString::from(pipeline_name.clone()),
1564 ),
1565 Label::new("row", SharedString::from(row.clone())),
1566 ];
1567 counter!("faucet_pipeline_pages_skipped_total", skip_labels)
1568 .increment(1);
1569 } else {
1570 records_written += with_retry!(
1571 "sink_write",
1572 sink.write_batch_idempotent(&page.records, &scope, &token)
1573 )?;
1574 }
1575 with_retry!("flush", sink.flush())?;
1576 let bm_labels =
1577 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1578 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1579 if let (Some(store), Some(key)) =
1580 (state_store.as_ref(), state_key.as_ref())
1581 {
1582 let wrapped =
1583 crate::idempotency::wrap_state(Some(&bookmark), next_seq);
1584 with_retry!("state_put", store.put(key, &wrapped))?;
1585 }
1586 last_bookmark = Some(bookmark);
1587 } else if !page.records.is_empty() {
1588 records_written +=
1592 with_retry_write!("sink_write", sink.write_batch(&page.records))?;
1593 }
1594 } else {
1595 debug_assert!(
1597 quality_envelopes.is_empty(),
1598 "quality quarantine without DLQ should have been rejected at run start"
1599 );
1600 if !page.records.is_empty() {
1601 if let Some(cfg) = adaptive_cfg.as_ref() {
1602 let ctrl = controller.get_or_insert_with(|| {
1603 crate::adaptive::AimdController::new(cfg, page.records.len())
1604 });
1605 maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
1606 let mut offset = 0;
1607 while offset < page.records.len() {
1608 let size =
1609 ctrl.current().max(1).min(page.records.len() - offset);
1610 let chunk = &page.records[offset..offset + size];
1611 let t0 = std::time::Instant::now();
1612 let n = with_retry_write!("sink_write", sink.write_batch(chunk))?;
1613 let latency = t0.elapsed();
1614 records_written += n;
1615 offset += size;
1616 let adj = ctrl.observe(crate::adaptive::Observation {
1617 batch_len: chunk.len(),
1618 errors: 0,
1619 latency,
1620 });
1621 emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
1622 }
1623 } else {
1624 records_written +=
1625 with_retry_write!("sink_write", sink.write_batch(&page.records))?;
1626 }
1627 }
1628 if let Some(bookmark) = page.bookmark {
1629 with_retry!("flush", sink.flush())?;
1630 let bm_labels =
1631 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1632 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1633 if let (Some(store), Some(key)) =
1634 (state_store.as_ref(), state_key.as_ref())
1635 {
1636 with_retry!("state_put", store.put(key, &bookmark))?;
1637 }
1638 last_bookmark = Some(bookmark);
1639 }
1640 }
1641 }
1642 Some(Err(e)) => return Err(e),
1643 None => break,
1644 }
1645 }
1646 Ok(())
1647 }
1648 .await;
1649
1650 if let Err(e) = loop_result {
1657 if let Some(ref dlq_cfg) = dlq
1658 && let Err(flush_err) = dlq_cfg.sink.flush().await
1659 {
1660 tracing::warn!(
1661 error = %flush_err,
1662 "DLQ sink flush failed during error unwind; original error preserved"
1663 );
1664 }
1665 if let Err(flush_err) = sink.flush().await {
1666 tracing::warn!(
1667 error = %flush_err,
1668 "sink flush failed during error unwind; original error preserved"
1669 );
1670 }
1671 return Err(e);
1672 }
1673
1674 if let Some(ref dlq_cfg) = dlq {
1680 let final_metric_labels: Vec<metrics::Label> = vec![
1681 metrics::Label::new(
1682 "pipeline",
1683 metrics::SharedString::from(pipeline_name.clone()),
1684 ),
1685 metrics::Label::new("row", metrics::SharedString::from(row.clone())),
1686 metrics::Label::new(
1687 "connector",
1688 metrics::SharedString::from(sink_name.to_string()),
1689 ),
1690 metrics::Label::new(
1691 "dlq_connector",
1692 metrics::SharedString::from(dlq_sink_name.to_string()),
1693 ),
1694 ];
1695 let _final_dlq_flush_timer = crate::observability::DurationGuard::new(
1696 "faucet_sink_dlq_flush_duration_seconds",
1697 final_metric_labels.clone(),
1698 );
1699 dlq_cfg.sink.flush().await.map_err(|e| {
1700 let mut lbl = final_metric_labels.clone();
1701 lbl.push(metrics::Label::new(
1702 "kind",
1703 metrics::SharedString::const_str(crate::observability::decorator::error_kind(&e)),
1704 ));
1705 metrics::counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1706 FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
1707 })?;
1708 }
1709 sink.flush().await?;
1710
1711 if cancelled {
1712 tracing::info!(
1713 records_written,
1714 "pipeline run cancelled cooperatively; sink flushed (partial output is durable)"
1715 );
1716 }
1717
1718 tracing::info!(
1719 records_written,
1720 cancelled,
1721 has_bookmark = last_bookmark.is_some(),
1722 persisted = state_store.is_some() && state_key.is_some() && last_bookmark.is_some(),
1723 dlq_records = dlq_stats.records_dlq,
1724 "pipeline streaming run complete"
1725 );
1726
1727 Ok(PipelineResult {
1728 records_written,
1729 bookmark: last_bookmark,
1730 dlq: dlq.is_some().then_some(dlq_stats),
1731 })
1732}
1733
1734fn emit_adaptive_metrics(
1737 ctrl: &crate::adaptive::AimdController,
1738 adj: Option<crate::adaptive::Adjustment>,
1739 pipeline: &str,
1740 row: &str,
1741) {
1742 use metrics::{Label, SharedString, counter, gauge};
1743 let base = vec![
1744 Label::new("pipeline", SharedString::from(pipeline.to_string())),
1745 Label::new("row", SharedString::from(row.to_string())),
1746 ];
1747 gauge!("faucet_pipeline_adaptive_batch_size", base.clone()).set(ctrl.current() as f64);
1748 gauge!(
1749 "faucet_pipeline_adaptive_batch_cooldown_active",
1750 base.clone()
1751 )
1752 .set(if ctrl.cooldown_active() { 1.0 } else { 0.0 });
1753 if let Some(p50) = ctrl.p50_latency_ms() {
1754 gauge!(
1755 "faucet_pipeline_adaptive_batch_p50_latency_ms",
1756 base.clone()
1757 )
1758 .set(p50 as f64);
1759 }
1760 if let Some(a) = adj {
1761 let mut lbl = base;
1762 lbl.push(Label::new(
1763 "direction",
1764 SharedString::const_str(a.direction.as_str()),
1765 ));
1766 lbl.push(Label::new(
1767 "reason",
1768 SharedString::const_str(a.reason.as_str()),
1769 ));
1770 counter!("faucet_pipeline_adaptive_batch_adjustments_total", lbl).increment(1);
1771 }
1772}
1773
1774fn maybe_warn_noop_sink(sink_name: &str, warned: &mut bool) {
1777 if !*warned && matches!(sink_name, "jsonl" | "csv" | "stdout") {
1778 tracing::info!(
1779 sink = sink_name,
1780 "adaptive batch sizing is a no-op for this per-record sink"
1781 );
1782 *warned = true;
1783 }
1784}
1785
1786#[allow(clippy::too_many_arguments)]
1794async fn apply_drift_policy<Si: Sink + ?Sized>(
1795 policy: &crate::drift::SchemaDriftPolicy,
1796 diff: &crate::drift::SchemaDiff,
1797 dest: &Value,
1798 records: Vec<Value>,
1799 page_indices: &[usize],
1800 sink: &Si,
1801 sink_name: &str,
1802 pipeline_name: &str,
1803 row: &str,
1804 dest_schema_cache: &mut Option<Option<Value>>,
1805 drift_envelopes: &mut Vec<Value>,
1806) -> Result<(Vec<Value>, Option<FaucetError>), FaucetError> {
1807 use crate::drift::{OnDrift, OnIncompatible};
1808 use crate::observability::schema_drift as emit_drift;
1809
1810 let mode = match policy.on_drift {
1811 OnDrift::Warn => "warn",
1812 OnDrift::Ignore => "ignore",
1813 OnDrift::Quarantine => "quarantine",
1814 OnDrift::Fail => "fail",
1815 OnDrift::Evolve => "evolve",
1816 };
1817 emit_drift(
1818 pipeline_name,
1819 row,
1820 sink_name,
1821 mode,
1822 "added",
1823 diff.additions.len() as u64,
1824 );
1825 emit_drift(
1826 pipeline_name,
1827 row,
1828 sink_name,
1829 mode,
1830 "widened",
1831 diff.widenings.len() as u64,
1832 );
1833 emit_drift(
1834 pipeline_name,
1835 row,
1836 sink_name,
1837 mode,
1838 "narrowed",
1839 diff.incompatible.len() as u64,
1840 );
1841 emit_drift(
1842 pipeline_name,
1843 row,
1844 sink_name,
1845 mode,
1846 "dropped",
1847 diff.droppable_required.len() as u64,
1848 );
1849
1850 match policy.on_drift {
1851 OnDrift::Warn => {
1852 tracing::warn!(
1853 connector = sink_name,
1854 columns = ?diff.changed_columns(),
1855 "schema-drift detected (on_drift=warn); writing page unchanged"
1856 );
1857 Ok((records, None))
1858 }
1859 OnDrift::Fail => Ok((
1860 records,
1861 Some(FaucetError::SchemaDrift {
1862 columns: diff.changed_columns(),
1863 message: "schema drift detected (on_drift=fail)".to_string(),
1864 }),
1865 )),
1866 OnDrift::Ignore => {
1867 let allowed: std::collections::HashSet<String> = dest
1869 .get("properties")
1870 .and_then(|p| p.as_object())
1871 .map(|m| m.keys().cloned().collect())
1872 .unwrap_or_default();
1873 let trimmed = records
1874 .into_iter()
1875 .map(|r| match r {
1876 Value::Object(map) => Value::Object(
1877 map.into_iter()
1878 .filter(|(k, _)| allowed.contains(k))
1879 .collect(),
1880 ),
1881 other => other,
1882 })
1883 .collect();
1884 Ok((trimmed, None))
1885 }
1886 OnDrift::Quarantine => {
1887 let (kept, env) =
1888 quarantine_drift_rows(diff, records, page_indices, sink_name, pipeline_name, row);
1889 drift_envelopes.extend(env);
1890 Ok((kept, None))
1891 }
1892 OnDrift::Evolve => {
1893 let evolution = crate::drift::SchemaEvolution {
1894 additions: diff.additions.clone(),
1895 widenings: diff
1896 .widenings
1897 .iter()
1898 .filter(|c| {
1899 c.from
1900 .as_ref()
1901 .map(|f| crate::drift::base_widened(f, &c.to))
1902 .unwrap_or(false)
1903 })
1904 .cloned()
1905 .collect(),
1906 relax_nullability: diff
1907 .droppable_required
1908 .iter()
1909 .filter(|_| policy.relax_nullability_on_missing)
1914 .cloned()
1915 .chain(
1916 diff.widenings
1917 .iter()
1918 .filter(|c| {
1919 c.from
1920 .as_ref()
1921 .map(|f| crate::drift::adds_null(f, &c.to))
1922 .unwrap_or(false)
1923 })
1924 .map(|c| c.name.clone()),
1925 )
1926 .collect(),
1927 };
1928 if !evolution.is_empty() {
1929 sink.evolve_schema(&evolution).await?;
1930 *dest_schema_cache = Some(sink.current_schema().await?);
1933 }
1934 if diff.incompatible.is_empty() {
1936 Ok((records, None))
1937 } else {
1938 match policy.on_incompatible {
1939 OnIncompatible::Fail => Ok((
1940 records,
1941 Some(FaucetError::SchemaDrift {
1942 columns: diff.incompatible.iter().map(|c| c.name.clone()).collect(),
1943 message: "incompatible type change cannot be auto-evolved \
1944 (on_incompatible=fail)"
1945 .into(),
1946 }),
1947 )),
1948 OnIncompatible::Quarantine => {
1949 let incompat_only = crate::drift::SchemaDiff {
1951 incompatible: diff.incompatible.clone(),
1952 ..Default::default()
1953 };
1954 let (kept, env) = quarantine_drift_rows(
1955 &incompat_only,
1956 records,
1957 page_indices,
1958 sink_name,
1959 pipeline_name,
1960 row,
1961 );
1962 drift_envelopes.extend(env);
1963 Ok((kept, None))
1964 }
1965 }
1966 }
1967 }
1968 }
1969}
1970
1971fn quarantine_drift_rows(
1981 diff: &crate::drift::SchemaDiff,
1982 records: Vec<Value>,
1983 page_indices: &[usize],
1984 sink_name: &str,
1985 pipeline_name: &str,
1986 row: &str,
1987) -> (Vec<Value>, Vec<Value>) {
1988 use crate::dlq::{DlqReason, build_envelope};
1989 let present_cols: std::collections::HashSet<&str> = diff
1991 .additions
1992 .iter()
1993 .chain(&diff.widenings)
1994 .chain(&diff.incompatible)
1995 .map(|c| c.name.as_str())
1996 .collect();
1997 let required_cols: std::collections::HashSet<&str> =
1999 diff.droppable_required.iter().map(|s| s.as_str()).collect();
2000 let mut kept = Vec::new();
2001 let mut envelopes = Vec::new();
2002 for (idx, rec) in records.into_iter().enumerate() {
2003 let exhibits = rec
2004 .as_object()
2005 .map(|m| {
2006 m.keys().any(|k| present_cols.contains(k.as_str()))
2007 || required_cols.iter().any(|c| !m.contains_key(*c))
2008 })
2009 .unwrap_or(false);
2010 if exhibits {
2011 let err = FaucetError::SchemaDrift {
2012 columns: diff.changed_columns(),
2013 message: "row exhibits schema drift (on_drift=quarantine)".into(),
2014 };
2015 let page_index = page_indices.get(idx).copied().unwrap_or(idx);
2018 envelopes.push(build_envelope(
2019 &rec,
2020 &err,
2021 DlqReason::SchemaDrift,
2022 sink_name,
2023 pipeline_name,
2024 row,
2025 page_index,
2026 ));
2027 } else {
2028 kept.push(rec);
2029 }
2030 }
2031 (kept, envelopes)
2032}
2033
2034#[cfg(test)]
2035mod tests {
2036 use super::*;
2037 use async_trait::async_trait;
2038 use serde_json::json;
2039
2040 struct MockSource(Vec<Value>);
2043
2044 #[async_trait]
2045 impl Source for MockSource {
2046 async fn fetch_with_context(
2047 &self,
2048 _context: &std::collections::HashMap<String, Value>,
2049 ) -> Result<Vec<Value>, FaucetError> {
2050 Ok(self.0.clone())
2051 }
2052 }
2053
2054 struct IncrementalSource {
2055 records: Vec<Value>,
2056 bookmark: Value,
2057 }
2058
2059 #[async_trait]
2060 impl Source for IncrementalSource {
2061 async fn fetch_with_context(
2062 &self,
2063 _context: &std::collections::HashMap<String, Value>,
2064 ) -> Result<Vec<Value>, FaucetError> {
2065 Ok(self.records.clone())
2066 }
2067 async fn fetch_with_context_incremental(
2068 &self,
2069 _context: &std::collections::HashMap<String, Value>,
2070 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
2071 Ok((self.records.clone(), Some(self.bookmark.clone())))
2072 }
2073 }
2074
2075 struct FailingSource;
2076
2077 #[async_trait]
2078 impl Source for FailingSource {
2079 async fn fetch_with_context(
2080 &self,
2081 _context: &std::collections::HashMap<String, Value>,
2082 ) -> Result<Vec<Value>, FaucetError> {
2083 Err(FaucetError::Auth("no credentials".into()))
2084 }
2085 }
2086
2087 struct MockSink(std::sync::Mutex<Vec<Value>>);
2090
2091 impl MockSink {
2092 fn new() -> Self {
2093 Self(std::sync::Mutex::new(Vec::new()))
2094 }
2095 fn written(&self) -> Vec<Value> {
2096 self.0.lock().unwrap().clone()
2097 }
2098 }
2099
2100 #[async_trait]
2101 impl Sink for MockSink {
2102 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2103 self.0.lock().unwrap().extend(records.iter().cloned());
2104 Ok(records.len())
2105 }
2106 }
2107
2108 struct FailingSink;
2109
2110 #[async_trait]
2111 impl Sink for FailingSink {
2112 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
2113 Err(FaucetError::Sink("write failed".into()))
2114 }
2115 }
2116
2117 struct FlushTrackingSink {
2122 written: std::sync::Mutex<Vec<Value>>,
2123 flush_count: std::sync::atomic::AtomicUsize,
2124 }
2125
2126 impl FlushTrackingSink {
2127 fn new() -> Self {
2128 Self {
2129 written: std::sync::Mutex::new(Vec::new()),
2130 flush_count: std::sync::atomic::AtomicUsize::new(0),
2131 }
2132 }
2133 fn written(&self) -> Vec<Value> {
2134 self.written.lock().unwrap().clone()
2135 }
2136 fn flush_count(&self) -> usize {
2137 self.flush_count.load(std::sync::atomic::Ordering::SeqCst)
2138 }
2139 }
2140
2141 #[async_trait]
2142 impl Sink for FlushTrackingSink {
2143 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2144 self.written.lock().unwrap().extend(records.iter().cloned());
2145 Ok(records.len())
2146 }
2147 async fn flush(&self) -> Result<(), FaucetError> {
2148 self.flush_count
2149 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2150 Ok(())
2151 }
2152 }
2153
2154 struct IdempotentMockSink {
2158 rows: std::sync::Mutex<Vec<Value>>,
2159 tokens: std::sync::Mutex<std::collections::HashMap<String, String>>,
2160 }
2161 impl IdempotentMockSink {
2162 fn new() -> Self {
2163 Self {
2164 rows: std::sync::Mutex::new(Vec::new()),
2165 tokens: std::sync::Mutex::new(std::collections::HashMap::new()),
2166 }
2167 }
2168 fn rows(&self) -> Vec<Value> {
2169 self.rows.lock().unwrap().clone()
2170 }
2171 }
2172 #[async_trait]
2173 impl Sink for IdempotentMockSink {
2174 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2175 self.rows.lock().unwrap().extend(records.iter().cloned());
2176 Ok(records.len())
2177 }
2178 fn supports_idempotent_writes(&self) -> bool {
2179 true
2180 }
2181 async fn write_batch_idempotent(
2182 &self,
2183 records: &[Value],
2184 scope: &str,
2185 token: &str,
2186 ) -> Result<usize, FaucetError> {
2187 self.rows.lock().unwrap().extend(records.iter().cloned());
2188 self.tokens
2189 .lock()
2190 .unwrap()
2191 .insert(scope.to_string(), token.to_string());
2192 Ok(records.len())
2193 }
2194 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
2195 Ok(self.tokens.lock().unwrap().get(scope).cloned())
2196 }
2197 }
2198
2199 fn eo_opts(store: Arc<dyn StateStore>, key: &str, start_seq: u64) -> RunStreamOptions {
2200 RunStreamOptions::new()
2201 .with_state(store, key)
2202 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2203 .with_start_seq(start_seq)
2204 }
2205
2206 #[tokio::test]
2207 async fn exactly_once_writes_pages_and_persists_wrapped_state() {
2208 let pages = vec![
2209 Ok(StreamPage {
2210 records: vec![json!({"id": 1})],
2211 bookmark: Some(json!("b1")),
2212 }),
2213 Ok(StreamPage {
2214 records: vec![json!({"id": 2})],
2215 bookmark: Some(json!("b2")),
2216 }),
2217 ];
2218 let sink = IdempotentMockSink::new();
2219 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2220 let r = run_stream(
2221 futures::stream::iter(pages),
2222 &sink,
2223 eo_opts(store.clone(), "k", 0),
2224 )
2225 .await
2226 .unwrap();
2227 assert_eq!(r.records_written, 2);
2228 let (bm, seq) = crate::idempotency::unwrap_state(&store.get("k").await.unwrap().unwrap());
2229 assert_eq!(bm, Some(json!("b2")));
2230 assert_eq!(seq, 2);
2231 let token = sink.last_committed_token("k").await.unwrap().unwrap();
2234 assert_eq!(
2235 crate::idempotency::parse_token_parts(&token),
2236 Some((2, Some(json!("b2"))))
2237 );
2238 }
2239
2240 #[tokio::test]
2241 async fn exactly_once_skips_already_committed_pages_on_resume() {
2242 let sink = IdempotentMockSink::new();
2243 sink.write_batch_idempotent(
2245 &[json!({"id": 1})],
2246 "k",
2247 &crate::idempotency::format_token(1),
2248 )
2249 .await
2250 .unwrap();
2251 assert_eq!(sink.rows().len(), 1);
2252 let pages = vec![
2254 Ok(StreamPage {
2255 records: vec![json!({"id": 1})],
2256 bookmark: Some(json!("b1")),
2257 }),
2258 Ok(StreamPage {
2259 records: vec![json!({"id": 2})],
2260 bookmark: Some(json!("b2")),
2261 }),
2262 ];
2263 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2264 let r = run_stream(futures::stream::iter(pages), &sink, eo_opts(store, "k", 0))
2265 .await
2266 .unwrap();
2267 assert_eq!(r.records_written, 1);
2268 let rows = sink.rows();
2269 assert_eq!(
2270 rows.len(),
2271 2,
2272 "exactly one row per id — no duplicate of id=1"
2273 );
2274 assert_eq!(rows.iter().filter(|v| v["id"] == 1).count(), 1);
2275 }
2276
2277 #[tokio::test]
2278 async fn exactly_once_rejects_non_idempotent_sink() {
2279 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2280 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2281 let r = run_stream(
2282 futures::stream::iter(pages),
2283 &MockSink::new(),
2284 eo_opts(store, "k", 0),
2285 )
2286 .await;
2287 assert!(matches!(r, Err(FaucetError::Config(_))));
2288 }
2289
2290 #[tokio::test]
2291 async fn exactly_once_rejects_missing_state_store() {
2292 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2293 let opts =
2294 RunStreamOptions::new().with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce);
2295 let r = run_stream(
2296 futures::stream::iter(pages),
2297 &IdempotentMockSink::new(),
2298 opts,
2299 )
2300 .await;
2301 assert!(matches!(r, Err(FaucetError::Config(_))));
2302 }
2303
2304 struct KeyedMockSink(MockSink);
2307 #[async_trait]
2308 impl Sink for KeyedMockSink {
2309 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2310 self.0.write_batch(records).await
2311 }
2312 fn dedups_by_key(&self) -> bool {
2313 true
2314 }
2315 fn supported_write_modes(&self) -> &'static [crate::write_mode::WriteMode] {
2316 &[
2317 crate::write_mode::WriteMode::Append,
2318 crate::write_mode::WriteMode::Upsert,
2319 crate::write_mode::WriteMode::Delete,
2320 ]
2321 }
2322 }
2323
2324 #[tokio::test]
2325 async fn exactly_once_keyed_upsert_mechanism_uses_plain_write_path() {
2326 let sink = KeyedMockSink(MockSink::new());
2331 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2332 let pages = vec![Ok(StreamPage {
2333 records: vec![json!({"id": 1})],
2334 bookmark: Some(json!("b1")),
2335 })];
2336 let opts = eo_opts(store.clone(), "k", 0)
2337 .with_replay_guarantee(crate::idempotency::ReplayGuarantee::NonDeterministic);
2338 let r = run_stream(futures::stream::iter(pages), &sink, opts)
2339 .await
2340 .unwrap();
2341 assert_eq!(r.records_written, 1);
2342 assert_eq!(sink.0.written(), vec![json!({"id": 1})]);
2343 assert_eq!(store.get("k").await.unwrap(), Some(json!("b1")));
2345 }
2346
2347 #[tokio::test]
2348 async fn exactly_once_rejects_non_deterministic_source_without_keyed_dedup() {
2349 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2354 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2355 let opts = eo_opts(store, "k", 0)
2356 .with_replay_guarantee(crate::idempotency::ReplayGuarantee::NonDeterministic);
2357 let r = run_stream(
2358 futures::stream::iter(pages),
2359 &IdempotentMockSink::new(),
2360 opts,
2361 )
2362 .await;
2363 match r {
2364 Err(FaucetError::Config(msg)) => {
2365 assert!(msg.contains("write_mode: upsert"), "hint present: {msg}")
2366 }
2367 other => panic!("expected Config error, got {other:?}"),
2368 }
2369 }
2370
2371 #[tokio::test]
2372 async fn exactly_once_rejects_plain_sink_with_mechanism_hint() {
2373 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2374 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2375 let r = run_stream(
2376 futures::stream::iter(pages),
2377 &MockSink::new(),
2378 eo_opts(store, "k", 0),
2379 )
2380 .await;
2381 match r {
2382 Err(FaucetError::Config(msg)) => assert!(
2383 msg.contains("provides neither"),
2384 "names both mechanisms: {msg}"
2385 ),
2386 other => panic!("expected Config error, got {other:?}"),
2387 }
2388 }
2389
2390 struct AnchorRecordingSource {
2393 applied: std::sync::Mutex<Vec<Value>>,
2394 }
2395 #[async_trait]
2396 impl Source for AnchorRecordingSource {
2397 async fn fetch_with_context(
2398 &self,
2399 _context: &std::collections::HashMap<String, Value>,
2400 ) -> Result<Vec<Value>, FaucetError> {
2401 Ok(vec![json!({"id": 10})])
2402 }
2403 async fn fetch_with_context_incremental(
2404 &self,
2405 _context: &std::collections::HashMap<String, Value>,
2406 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
2407 Ok((vec![json!({"id": 10})], Some(json!("after"))))
2408 }
2409 fn state_key(&self) -> Option<String> {
2410 Some("anchor_key".to_string())
2411 }
2412 async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
2413 self.applied.lock().unwrap().push(bookmark);
2414 Ok(())
2415 }
2416 fn supports_exactly_once(&self) -> bool {
2417 true
2418 }
2419 }
2420
2421 #[tokio::test]
2422 async fn pipeline_run_anchors_resume_at_sink_embedded_bookmark() {
2423 let source = AnchorRecordingSource {
2430 applied: std::sync::Mutex::new(Vec::new()),
2431 };
2432 let sink = IdempotentMockSink::new();
2433 sink.tokens.lock().unwrap().insert(
2434 "anchor_key".to_string(),
2435 crate::idempotency::format_token_with_bookmark(5, Some(&json!("sink-pos"))),
2436 );
2437 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2438 store
2439 .put(
2440 "anchor_key",
2441 &crate::idempotency::wrap_state(Some(&json!("state-pos")), 4),
2442 )
2443 .await
2444 .unwrap();
2445
2446 let r = Pipeline::new(&source, &sink)
2447 .with_state_store(Arc::clone(&store))
2448 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2449 .run()
2450 .await
2451 .unwrap();
2452
2453 assert_eq!(
2454 *source.applied.lock().unwrap(),
2455 vec![json!("state-pos"), json!("sink-pos")],
2456 "state bookmark applied, then overridden by the sink anchor"
2457 );
2458 assert_eq!(r.records_written, 1);
2461 let token = sink
2462 .last_committed_token("anchor_key")
2463 .await
2464 .unwrap()
2465 .unwrap();
2466 assert_eq!(
2467 crate::idempotency::parse_token_parts(&token),
2468 Some((6, Some(json!("after"))))
2469 );
2470 let (bm, seq) =
2471 crate::idempotency::unwrap_state(&store.get("anchor_key").await.unwrap().unwrap());
2472 assert_eq!((bm, seq), (Some(json!("after")), 6));
2473 }
2474
2475 #[tokio::test]
2476 async fn pipeline_run_ignores_sink_token_behind_state_seq() {
2477 let source = AnchorRecordingSource {
2480 applied: std::sync::Mutex::new(Vec::new()),
2481 };
2482 let sink = IdempotentMockSink::new();
2483 sink.tokens.lock().unwrap().insert(
2484 "anchor_key".to_string(),
2485 crate::idempotency::format_token_with_bookmark(4, Some(&json!("sink-pos"))),
2486 );
2487 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2488 store
2489 .put(
2490 "anchor_key",
2491 &crate::idempotency::wrap_state(Some(&json!("state-pos")), 4),
2492 )
2493 .await
2494 .unwrap();
2495 Pipeline::new(&source, &sink)
2496 .with_state_store(Arc::clone(&store))
2497 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2498 .run()
2499 .await
2500 .unwrap();
2501 assert_eq!(*source.applied.lock().unwrap(), vec![json!("state-pos")]);
2502 }
2503
2504 #[tokio::test]
2505 async fn pipeline_run_legacy_bare_token_falls_back_to_skip_path() {
2506 let source = AnchorRecordingSource {
2510 applied: std::sync::Mutex::new(Vec::new()),
2511 };
2512 let sink = IdempotentMockSink::new();
2513 sink.tokens.lock().unwrap().insert(
2514 "anchor_key".to_string(),
2515 crate::idempotency::format_token(1),
2516 );
2517 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2518 let r = Pipeline::new(&source, &sink)
2519 .with_state_store(Arc::clone(&store))
2520 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2521 .run()
2522 .await
2523 .unwrap();
2524 assert!(source.applied.lock().unwrap().is_empty());
2525 assert_eq!(r.records_written, 0, "page 1 already committed → skipped");
2526 assert!(sink.rows().is_empty());
2527 }
2528
2529 #[test]
2532 fn stream_page_constructs() {
2533 let page = StreamPage {
2534 records: vec![json!({"id": 1})],
2535 bookmark: Some(json!("2026-05-18")),
2536 };
2537 assert_eq!(page.records.len(), 1);
2538 assert_eq!(page.bookmark, Some(json!("2026-05-18")));
2539 }
2540
2541 #[test]
2542 fn validate_batch_size_accepts_zero_as_no_batching_sentinel() {
2543 assert_eq!(validate_batch_size(0).unwrap(), 0);
2545 }
2546
2547 #[test]
2548 fn validate_batch_size_rejects_too_large() {
2549 let err = validate_batch_size(MAX_BATCH_SIZE + 1).unwrap_err();
2550 assert!(matches!(err, FaucetError::Config(_)));
2551 }
2552
2553 #[test]
2554 fn validate_batch_size_accepts_one() {
2555 assert_eq!(validate_batch_size(1).unwrap(), 1);
2556 }
2557
2558 #[test]
2559 fn validate_batch_size_accepts_max() {
2560 assert_eq!(validate_batch_size(MAX_BATCH_SIZE).unwrap(), MAX_BATCH_SIZE);
2561 }
2562
2563 const _: () = {
2565 assert!(DEFAULT_BATCH_SIZE >= 1);
2566 assert!(DEFAULT_BATCH_SIZE <= MAX_BATCH_SIZE);
2567 };
2568
2569 #[tokio::test]
2572 async fn batch_pipeline_writes_all_records() {
2573 let source = MockSource(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
2574 let sink = MockSink::new();
2575
2576 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2577
2578 assert_eq!(result.records_written, 3);
2579 assert!(result.bookmark.is_none());
2580 assert_eq!(sink.written().len(), 3);
2581 }
2582
2583 #[tokio::test]
2584 async fn batch_pipeline_returns_bookmark() {
2585 let source = IncrementalSource {
2586 records: vec![json!({"id": 1, "ts": "2024-12-01"})],
2587 bookmark: json!("2024-12-01"),
2588 };
2589 let sink = MockSink::new();
2590
2591 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2592
2593 assert_eq!(result.records_written, 1);
2594 assert_eq!(result.bookmark, Some(json!("2024-12-01")));
2595 }
2596
2597 #[tokio::test]
2598 async fn batch_pipeline_empty_source() {
2599 let source = MockSource(vec![]);
2600 let sink = MockSink::new();
2601
2602 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2603
2604 assert_eq!(result.records_written, 0);
2605 assert!(sink.written().is_empty());
2606 }
2607
2608 #[tokio::test]
2609 async fn batch_pipeline_source_error_propagates() {
2610 let source = FailingSource;
2611 let sink = MockSink::new();
2612
2613 let result = Pipeline::new(&source, &sink).run().await;
2614 assert!(result.is_err());
2615 assert!(sink.written().is_empty());
2616 }
2617
2618 #[tokio::test]
2619 async fn batch_pipeline_sink_error_propagates() {
2620 let source = MockSource(vec![json!({"id": 1})]);
2621 let sink = FailingSink;
2622
2623 let result = Pipeline::new(&source, &sink).run().await;
2624 assert!(result.is_err());
2625 }
2626
2627 #[tokio::test]
2628 async fn batch_pipeline_with_trait_objects() {
2629 let source: Box<dyn Source> = Box::new(MockSource(vec![json!({"id": 1})]));
2630 let sink: Box<dyn Sink> = Box::new(MockSink::new());
2631
2632 let result = Pipeline::new(source.as_ref(), sink.as_ref())
2633 .run()
2634 .await
2635 .unwrap();
2636
2637 assert_eq!(result.records_written, 1);
2638 }
2639
2640 #[tokio::test]
2643 async fn stream_pipeline_writes_pages() {
2644 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2645 Ok(StreamPage {
2646 records: vec![json!({"id": 1}), json!({"id": 2})],
2647 bookmark: None,
2648 }),
2649 Ok(StreamPage {
2650 records: vec![json!({"id": 3})],
2651 bookmark: None,
2652 }),
2653 ];
2654 let stream = futures::stream::iter(pages);
2655 let sink = MockSink::new();
2656
2657 let result = run_stream(stream, &sink, RunStreamOptions::new())
2658 .await
2659 .unwrap();
2660
2661 assert_eq!(result.records_written, 3);
2662 assert!(result.bookmark.is_none());
2663 assert_eq!(sink.written().len(), 3);
2664 }
2665
2666 #[tokio::test]
2667 async fn stream_pipeline_flushes_sink_on_source_error() {
2668 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2672 Ok(StreamPage {
2673 records: vec![json!({"id": 1}), json!({"id": 2})],
2674 bookmark: None,
2675 }),
2676 Err(FaucetError::Source("transient blip mid-stream".into())),
2677 ];
2678 let stream = futures::stream::iter(pages);
2679 let sink = FlushTrackingSink::new();
2680
2681 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2682
2683 assert!(matches!(result, Err(FaucetError::Source(_))));
2685 assert_eq!(sink.written().len(), 2);
2687 assert!(
2689 sink.flush_count() >= 1,
2690 "sink must be flushed on the error path so partial output is durable"
2691 );
2692 }
2693
2694 #[tokio::test]
2695 async fn stream_pipeline_flushes_sink_on_cancel() {
2696 use tokio_util::sync::CancellationToken;
2701
2702 let stream = Box::pin(async_stream::stream! {
2704 yield Ok(StreamPage {
2705 records: vec![json!({"id": 1}), json!({"id": 2})],
2706 bookmark: None,
2707 });
2708 futures::future::pending::<()>().await;
2709 });
2710 let sink = FlushTrackingSink::new();
2711
2712 let token = CancellationToken::new();
2713 let canceller = token.clone();
2714 tokio::spawn(async move {
2715 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
2716 canceller.cancel();
2717 });
2718
2719 let result = run_stream(stream, &sink, RunStreamOptions::new().with_cancel(token))
2720 .await
2721 .expect("a cooperative cancel returns Ok with the partial result");
2722
2723 assert_eq!(result.records_written, 2);
2726 assert_eq!(sink.written().len(), 2);
2727 assert!(
2728 sink.flush_count() >= 1,
2729 "sink must be flushed on the cancel path so partial output is durable"
2730 );
2731 }
2732
2733 #[tokio::test]
2734 async fn stream_pipeline_empty() {
2735 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2736 let stream = futures::stream::iter(pages);
2737 let sink = MockSink::new();
2738
2739 let result = run_stream(stream, &sink, RunStreamOptions::new())
2740 .await
2741 .unwrap();
2742
2743 assert_eq!(result.records_written, 0);
2744 }
2745
2746 #[tokio::test]
2747 async fn stream_pipeline_skips_empty_pages() {
2748 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2749 Ok(StreamPage {
2750 records: vec![json!({"id": 1})],
2751 bookmark: None,
2752 }),
2753 Ok(StreamPage {
2754 records: vec![],
2755 bookmark: None,
2756 }),
2757 Ok(StreamPage {
2758 records: vec![json!({"id": 2})],
2759 bookmark: None,
2760 }),
2761 ];
2762 let stream = futures::stream::iter(pages);
2763 let sink = MockSink::new();
2764
2765 let result = run_stream(stream, &sink, RunStreamOptions::new())
2766 .await
2767 .unwrap();
2768
2769 assert_eq!(result.records_written, 2);
2770 }
2771
2772 #[tokio::test]
2773 async fn stream_pipeline_error_in_page_propagates() {
2774 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2775 Ok(StreamPage {
2776 records: vec![json!({"id": 1})],
2777 bookmark: None,
2778 }),
2779 Err(FaucetError::HttpStatus {
2780 status: 500,
2781 url: "https://example.com".into(),
2782 body: "Internal Server Error".into(),
2783 }),
2784 ];
2785 let stream = futures::stream::iter(pages);
2786 let sink = MockSink::new();
2787
2788 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2789 assert!(result.is_err());
2790 assert_eq!(sink.written().len(), 1);
2792 }
2793
2794 #[tokio::test]
2795 async fn stream_pipeline_sink_error_propagates() {
2796 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2797 records: vec![json!({"id": 1})],
2798 bookmark: None,
2799 })];
2800 let stream = futures::stream::iter(pages);
2801 let sink = FailingSink;
2802
2803 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2804 assert!(result.is_err());
2805 }
2806
2807 #[tokio::test]
2808 async fn stream_pipeline_with_trait_object_sink() {
2809 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2810 records: vec![json!({"id": 1})],
2811 bookmark: None,
2812 })];
2813 let stream = futures::stream::iter(pages);
2814 let sink: Box<dyn Sink> = Box::new(MockSink::new());
2815
2816 let result = run_stream(stream, sink.as_ref(), RunStreamOptions::new())
2817 .await
2818 .unwrap();
2819 assert_eq!(result.records_written, 1);
2820 }
2821
2822 #[tokio::test]
2823 async fn stream_pipeline_persists_bookmark_when_page_carries_one() {
2824 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2825 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2826 Ok(StreamPage {
2827 records: vec![json!({"id": 1})],
2828 bookmark: None,
2829 }),
2830 Ok(StreamPage {
2831 records: vec![json!({"id": 2})],
2832 bookmark: Some(json!("checkpoint-final")),
2833 }),
2834 ];
2835 let stream = futures::stream::iter(pages);
2836 let sink = MockSink::new();
2837
2838 let result = run_stream(
2839 stream,
2840 &sink,
2841 RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
2842 )
2843 .await
2844 .unwrap();
2845
2846 assert_eq!(result.records_written, 2);
2847 assert_eq!(result.bookmark, Some(json!("checkpoint-final")));
2848 assert_eq!(
2849 store.get("k").await.unwrap(),
2850 Some(json!("checkpoint-final"))
2851 );
2852 }
2853
2854 #[tokio::test]
2855 async fn stream_pipeline_persists_per_page_bookmarks() {
2856 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2857 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2858 Ok(StreamPage {
2859 records: vec![json!({"id": 1})],
2860 bookmark: Some(json!("tx-1")),
2861 }),
2862 Ok(StreamPage {
2863 records: vec![json!({"id": 2})],
2864 bookmark: Some(json!("tx-2")),
2865 }),
2866 ];
2867 let stream = futures::stream::iter(pages);
2868 let sink = MockSink::new();
2869
2870 run_stream(
2871 stream,
2872 &sink,
2873 RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
2874 )
2875 .await
2876 .unwrap();
2877
2878 assert_eq!(store.get("k").await.unwrap(), Some(json!("tx-2")));
2880 }
2881
2882 use crate::state::{FileStateStore, MemoryStateStore, StateStore};
2885 use std::sync::Arc;
2886 use tempfile::TempDir;
2887
2888 struct StatefulSource {
2892 key: String,
2893 records: Vec<Value>,
2894 new_bookmark: Value,
2895 seen_bookmark: std::sync::Mutex<Option<Value>>,
2896 }
2897
2898 impl StatefulSource {
2899 fn new(key: &str, records: Vec<Value>, new_bookmark: Value) -> Self {
2900 Self {
2901 key: key.into(),
2902 records,
2903 new_bookmark,
2904 seen_bookmark: std::sync::Mutex::new(None),
2905 }
2906 }
2907 fn observed_start(&self) -> Option<Value> {
2908 self.seen_bookmark.lock().unwrap().clone()
2909 }
2910 }
2911
2912 #[async_trait]
2913 impl Source for StatefulSource {
2914 async fn fetch_with_context(
2915 &self,
2916 _ctx: &std::collections::HashMap<String, Value>,
2917 ) -> Result<Vec<Value>, FaucetError> {
2918 Ok(self.records.clone())
2919 }
2920 async fn fetch_with_context_incremental(
2921 &self,
2922 _ctx: &std::collections::HashMap<String, Value>,
2923 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
2924 Ok((self.records.clone(), Some(self.new_bookmark.clone())))
2925 }
2926 fn state_key(&self) -> Option<String> {
2927 Some(self.key.clone())
2928 }
2929 async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
2930 *self.seen_bookmark.lock().unwrap() = Some(bookmark);
2931 Ok(())
2932 }
2933 }
2934
2935 #[tokio::test]
2936 async fn pipeline_with_state_store_persists_bookmark_after_sink() {
2937 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2938 let source = StatefulSource::new(
2939 "github_issues",
2940 vec![json!({"id": 1, "ts": "2026-05-01"})],
2941 json!("2026-05-01"),
2942 );
2943 let sink = MockSink::new();
2944 let result = Pipeline::new(&source, &sink)
2945 .with_state_store(Arc::clone(&store))
2946 .run()
2947 .await
2948 .unwrap();
2949
2950 assert_eq!(result.records_written, 1);
2951 assert_eq!(result.bookmark, Some(json!("2026-05-01")));
2952 let stored = store.get("github_issues").await.unwrap();
2954 assert_eq!(stored, Some(json!("2026-05-01")));
2955 }
2956
2957 #[tokio::test]
2958 async fn pipeline_with_state_store_resumes_from_stored_bookmark() {
2959 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2960 store
2961 .put("github_issues", &json!("2026-04-30"))
2962 .await
2963 .unwrap();
2964
2965 let source =
2966 StatefulSource::new("github_issues", vec![json!({"id": 2})], json!("2026-05-01"));
2967 let sink = MockSink::new();
2968 Pipeline::new(&source, &sink)
2969 .with_state_store(Arc::clone(&store))
2970 .run()
2971 .await
2972 .unwrap();
2973
2974 assert_eq!(source.observed_start(), Some(json!("2026-04-30")));
2976 assert_eq!(
2978 store.get("github_issues").await.unwrap(),
2979 Some(json!("2026-05-01"))
2980 );
2981 }
2982
2983 #[tokio::test]
2984 async fn pipeline_with_state_store_does_not_persist_when_sink_fails() {
2985 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2986 let source = StatefulSource::new("k", vec![json!({"id": 1})], json!("2026-05-01"));
2987 let sink = FailingSink;
2988
2989 let result = Pipeline::new(&source, &sink)
2990 .with_state_store(Arc::clone(&store))
2991 .run()
2992 .await;
2993 assert!(result.is_err());
2994 assert!(store.get("k").await.unwrap().is_none());
2995 }
2996
2997 #[tokio::test]
2998 async fn pipeline_with_state_store_no_state_key_means_no_persist() {
2999 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3000 let source = IncrementalSource {
3001 records: vec![json!({"id": 1})],
3002 bookmark: json!("ignored"),
3003 };
3004 let sink = MockSink::new();
3005 Pipeline::new(&source, &sink)
3006 .with_state_store(Arc::clone(&store))
3007 .run()
3008 .await
3009 .unwrap();
3010 assert!(store.get("anything").await.unwrap().is_none());
3013 }
3014
3015 #[tokio::test]
3016 async fn pipeline_with_state_store_skips_persist_when_bookmark_is_none() {
3017 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3018 struct NoBookmarkSource;
3019 #[async_trait]
3020 impl Source for NoBookmarkSource {
3021 async fn fetch_with_context(
3022 &self,
3023 _ctx: &std::collections::HashMap<String, Value>,
3024 ) -> Result<Vec<Value>, FaucetError> {
3025 Ok(vec![json!({"id": 1})])
3026 }
3027 fn state_key(&self) -> Option<String> {
3028 Some("k".into())
3029 }
3030 }
3031 let source = NoBookmarkSource;
3032 let sink = MockSink::new();
3033 Pipeline::new(&source, &sink)
3034 .with_state_store(Arc::clone(&store))
3035 .run()
3036 .await
3037 .unwrap();
3038 assert!(store.get("k").await.unwrap().is_none());
3039 }
3040
3041 struct PagedSource;
3046
3047 #[async_trait]
3048 impl Source for PagedSource {
3049 async fn fetch_with_context(
3050 &self,
3051 _ctx: &std::collections::HashMap<String, Value>,
3052 ) -> Result<Vec<Value>, FaucetError> {
3053 unreachable!("Pipeline::run must drive stream_pages, not fetch_with_context");
3055 }
3056 fn stream_pages<'a>(
3057 &'a self,
3058 _ctx: &'a std::collections::HashMap<String, Value>,
3059 _batch_size: usize,
3060 ) -> std::pin::Pin<
3061 Box<dyn futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>,
3062 > {
3063 Box::pin(async_stream::try_stream! {
3064 yield StreamPage { records: vec![json!({"i": 1})], bookmark: None };
3065 yield StreamPage { records: vec![json!({"i": 2})], bookmark: None };
3066 yield StreamPage { records: vec![json!({"i": 3})], bookmark: Some(json!("final")) };
3067 })
3068 }
3069 }
3070
3071 struct CountingSink {
3073 calls: std::sync::Mutex<Vec<usize>>,
3074 }
3075
3076 impl CountingSink {
3077 fn new() -> Self {
3078 Self {
3079 calls: std::sync::Mutex::new(Vec::new()),
3080 }
3081 }
3082 fn call_count(&self) -> usize {
3083 self.calls.lock().unwrap().len()
3084 }
3085 }
3086
3087 #[async_trait]
3088 impl Sink for CountingSink {
3089 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3090 self.calls.lock().unwrap().push(records.len());
3091 Ok(records.len())
3092 }
3093 }
3094
3095 #[tokio::test]
3096 async fn pipeline_run_drives_stream_pages() {
3097 let source = PagedSource;
3098 let sink = CountingSink::new();
3099
3100 let result = Pipeline::new(&source, &sink).run().await.unwrap();
3101
3102 assert_eq!(sink.call_count(), 3);
3104 assert_eq!(result.records_written, 3);
3105 assert_eq!(result.bookmark, Some(json!("final")));
3106 }
3107
3108 #[tokio::test]
3109 async fn pipeline_with_file_state_store_round_trips_across_runs() {
3110 let dir = TempDir::new().unwrap();
3111 let store: Arc<dyn StateStore> = Arc::new(FileStateStore::new(dir.path()));
3112
3113 let s1 = StatefulSource::new("k", vec![json!({"i": 1})], json!("v1"));
3115 let sink1 = MockSink::new();
3116 Pipeline::new(&s1, &sink1)
3117 .with_state_store(Arc::clone(&store))
3118 .run()
3119 .await
3120 .unwrap();
3121 assert_eq!(s1.observed_start(), None);
3122 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3123
3124 let s2 = StatefulSource::new("k", vec![json!({"i": 2})], json!("v2"));
3126 let sink2 = MockSink::new();
3127 Pipeline::new(&s2, &sink2)
3128 .with_state_store(Arc::clone(&store))
3129 .run()
3130 .await
3131 .unwrap();
3132 assert_eq!(s2.observed_start(), Some(json!("v1")));
3133 assert_eq!(store.get("k").await.unwrap(), Some(json!("v2")));
3134 }
3135
3136 #[tokio::test]
3137 #[allow(clippy::await_holding_lock)]
3138 async fn pipeline_run_increments_runs_total() {
3139 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3140 use metrics_util::debugging::DebugValue;
3141 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3142 let snap = snapshotter();
3143
3144 let source = MockSource(vec![json!({"i": 1})]);
3145 let sink = MockSink::new();
3146 let _ = Pipeline::new(&source, &sink)
3147 .with_name("test-pipeline")
3148 .with_row("rowA")
3149 .run()
3150 .await
3151 .unwrap();
3152
3153 let snapshot = snap.snapshot();
3154 let found = snapshot.into_vec().into_iter().any(
3155 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3156 key.key().name() == "faucet_pipeline_runs_total"
3157 && key.key().labels().any(|l: &metrics::Label| {
3158 l.key() == "pipeline" && l.value() == "test-pipeline"
3159 })
3160 && key
3161 .key()
3162 .labels()
3163 .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowA")
3164 && key
3165 .key()
3166 .labels()
3167 .any(|l: &metrics::Label| l.key() == "status" && l.value() == "ok")
3168 && matches!(v, DebugValue::Counter(c) if c >= 1)
3169 },
3170 );
3171 assert!(
3172 found,
3173 "expected faucet_pipeline_runs_total{{pipeline=test-pipeline, row=rowA, status=ok}}"
3174 );
3175 }
3176
3177 #[tokio::test]
3178 #[allow(clippy::await_holding_lock)]
3179 async fn pipeline_failure_attaches_kind_label_to_runs_total() {
3180 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3181 use metrics_util::debugging::DebugValue;
3182 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3183 let snap = snapshotter();
3184
3185 let source = FailingSource;
3186 let sink = MockSink::new();
3187 let _ = Pipeline::new(&source, &sink)
3188 .with_name("err-pipeline")
3189 .with_row("rowE")
3190 .run()
3191 .await;
3192
3193 let snapshot = snap.snapshot();
3194 let found = snapshot.into_vec().into_iter().any(
3195 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3196 key.key().name() == "faucet_pipeline_runs_total"
3197 && key.key().labels().any(|l: &metrics::Label| {
3198 l.key() == "pipeline" && l.value() == "err-pipeline"
3199 })
3200 && key
3201 .key()
3202 .labels()
3203 .any(|l: &metrics::Label| l.key() == "status" && l.value() == "err")
3204 && key
3205 .key()
3206 .labels()
3207 .any(|l: &metrics::Label| l.key() == "kind" && l.value() == "Auth")
3208 && matches!(v, DebugValue::Counter(c) if c >= 1)
3209 },
3210 );
3211 assert!(
3212 found,
3213 "expected faucet_pipeline_runs_total{{status=err, kind=Auth}} for failing source"
3214 );
3215 }
3216
3217 #[tokio::test]
3218 #[allow(clippy::await_holding_lock)]
3219 async fn pipeline_run_emits_start_time_gauge() {
3220 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3221 use metrics_util::debugging::DebugValue;
3222 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3223 let snap = snapshotter();
3224
3225 let source = MockSource(vec![json!({"i": 1})]);
3226 let sink = MockSink::new();
3227 let before = std::time::SystemTime::now()
3228 .duration_since(std::time::UNIX_EPOCH)
3229 .map(|d| d.as_secs_f64())
3230 .unwrap_or(0.0);
3231 let _ = Pipeline::new(&source, &sink)
3232 .with_name("start-time-pipeline")
3233 .with_row("rowS")
3234 .run()
3235 .await
3236 .unwrap();
3237
3238 let snapshot = snap.snapshot();
3239 let found = snapshot.into_vec().into_iter().any(
3240 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3241 if key.key().name() != "faucet_pipeline_start_time_unix_seconds" {
3242 return false;
3243 }
3244 let labels_match = key.key().labels().any(|l: &metrics::Label| {
3245 l.key() == "pipeline" && l.value() == "start-time-pipeline"
3246 }) && key
3247 .key()
3248 .labels()
3249 .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowS");
3250 if !labels_match {
3251 return false;
3252 }
3253 matches!(v, DebugValue::Gauge(g) if g.into_inner() >= before)
3254 },
3255 );
3256 assert!(
3257 found,
3258 "expected faucet_pipeline_start_time_unix_seconds gauge >= test-start timestamp"
3259 );
3260 }
3261
3262 #[tokio::test]
3263 #[allow(clippy::await_holding_lock)]
3264 async fn register_build_info_sets_version_gauge() {
3265 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3266 use metrics_util::debugging::DebugValue;
3267 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3268 let snap = snapshotter();
3269
3270 crate::observability::register_build_info();
3271
3272 let snapshot = snap.snapshot();
3273 let found = snapshot.into_vec().into_iter().any(
3274 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3275 key.key().name() == "faucet_build_info"
3276 && key.key().labels().any(|l: &metrics::Label| {
3277 l.key() == "version" && l.value() == env!("CARGO_PKG_VERSION")
3278 })
3279 && matches!(v, DebugValue::Gauge(g) if (g.into_inner() - 1.0).abs() < f64::EPSILON)
3280 },
3281 );
3282 assert!(
3283 found,
3284 "expected faucet_build_info{{version=CARGO_PKG_VERSION}} = 1.0 after register_build_info()"
3285 );
3286 }
3287
3288 use crate::dlq::{DlqConfig, OnBatchError};
3291
3292 struct PartialSink {
3296 fail_indices: std::sync::Mutex<Vec<usize>>,
3297 committed: std::sync::Mutex<Vec<Value>>,
3298 }
3299
3300 impl PartialSink {
3301 fn new(fail_indices: Vec<usize>) -> Self {
3302 Self {
3303 fail_indices: std::sync::Mutex::new(fail_indices),
3304 committed: std::sync::Mutex::new(Vec::new()),
3305 }
3306 }
3307 }
3308
3309 #[async_trait]
3310 impl Sink for PartialSink {
3311 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
3312 unreachable!("PartialSink only overrides write_batch_partial");
3313 }
3314 async fn write_batch_partial(
3315 &self,
3316 records: &[Value],
3317 ) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
3318 let fails: std::collections::HashSet<usize> =
3319 self.fail_indices.lock().unwrap().iter().copied().collect();
3320 let mut outcomes = Vec::with_capacity(records.len());
3321 for (i, rec) in records.iter().enumerate() {
3322 if fails.contains(&i) {
3323 outcomes.push(Err(FaucetError::Sink(format!("row {i} rejected"))));
3324 } else {
3325 self.committed.lock().unwrap().push(rec.clone());
3326 outcomes.push(Ok(()));
3327 }
3328 }
3329 Ok(outcomes)
3330 }
3331 }
3332
3333 #[tokio::test]
3334 async fn dlq_routes_only_failed_rows_for_partial_success_sink() {
3335 let main = PartialSink::new(vec![1, 3]); let dlq = std::sync::Arc::new(MockSink::new());
3337 let dlq_cfg = DlqConfig::new(dlq.clone());
3338
3339 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3340 records: (0..4).map(|i| json!({"i": i})).collect(),
3341 bookmark: None,
3342 })];
3343 let stream = futures::stream::iter(pages);
3344 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg))
3345 .await
3346 .unwrap();
3347
3348 assert_eq!(result.records_written, 2); assert_eq!(main.committed.lock().unwrap().len(), 2);
3350 let envelopes = dlq.0.lock().unwrap();
3351 assert_eq!(envelopes.len(), 2);
3352 assert_eq!(envelopes[0]["payload"]["i"], 1);
3353 assert_eq!(envelopes[0]["record_index"], 1);
3354 assert_eq!(envelopes[1]["payload"]["i"], 3);
3355 assert_eq!(envelopes[1]["record_index"], 3);
3356 let stats = result.dlq.unwrap();
3357 assert_eq!(stats.records_dlq, 2);
3358 assert_eq!(stats.pages_with_failures, 1);
3359 }
3360
3361 #[cfg(feature = "masking")]
3362 #[tokio::test]
3363 async fn masking_runs_before_the_sink() {
3364 use crate::masking::{CompiledMasking, MaskingSpec};
3365 let sink = MockSink::new();
3366 let spec: MaskingSpec = serde_json::from_value(json!({
3367 "rules": [{ "match": { "value_detector": "email" },
3368 "action": { "type": "redact" } }]
3369 }))
3370 .unwrap();
3371 let m = std::sync::Arc::new(CompiledMasking::compile(&spec).unwrap());
3372 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3373 records: vec![json!({"email": "a@b.com", "name": "Al"})],
3374 bookmark: None,
3375 })];
3376 run_stream(
3377 futures::stream::iter(pages),
3378 &sink,
3379 RunStreamOptions::new().with_masking(m),
3380 )
3381 .await
3382 .unwrap();
3383 assert_eq!(sink.written()[0], json!({"email": "***", "name": "Al"}));
3384 }
3385
3386 #[cfg(feature = "masking")]
3387 #[tokio::test]
3388 async fn masking_applies_before_the_dlq_envelope() {
3389 use crate::masking::{CompiledMasking, MaskingSpec};
3393 let main = PartialSink::new(vec![0]); let dlq = std::sync::Arc::new(MockSink::new());
3395 let spec: MaskingSpec = serde_json::from_value(json!({
3396 "rules": [{ "match": { "value_detector": "email" },
3397 "action": { "type": "redact" } }]
3398 }))
3399 .unwrap();
3400 let m = std::sync::Arc::new(CompiledMasking::compile(&spec).unwrap());
3401 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3402 records: vec![
3403 json!({"email": "secret@x.com"}),
3404 json!({"email": "ok@y.com"}),
3405 ],
3406 bookmark: None,
3407 })];
3408 let opts = RunStreamOptions::new()
3409 .with_masking(m)
3410 .with_dlq(DlqConfig::new(dlq.clone()));
3411 run_stream(futures::stream::iter(pages), &main, opts)
3412 .await
3413 .unwrap();
3414 let env = dlq.0.lock().unwrap();
3416 assert_eq!(env.len(), 1);
3417 assert_eq!(
3418 env[0]["payload"]["email"], "***",
3419 "the DLQ payload must be masked, not raw PII"
3420 );
3421 assert_eq!(main.committed.lock().unwrap()[0]["email"], "***");
3422 }
3423
3424 #[tokio::test]
3425 async fn dlq_propagate_policy_bubbles_outer_err() {
3426 let main = FailingSink;
3427 let dlq = std::sync::Arc::new(MockSink::new());
3428 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3429 dlq_cfg.on_batch_error = OnBatchError::Propagate;
3430
3431 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3432 records: vec![json!({"i": 0}), json!({"i": 1})],
3433 bookmark: Some(json!("v1")),
3434 })];
3435 let stream = futures::stream::iter(pages);
3436 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3437 let result = run_stream(
3438 stream,
3439 &main,
3440 RunStreamOptions::new()
3441 .with_dlq(dlq_cfg)
3442 .with_state(std::sync::Arc::clone(&store), "k"),
3443 )
3444 .await;
3445 assert!(matches!(result, Err(FaucetError::Sink(_))));
3446 assert!(dlq.0.lock().unwrap().is_empty());
3447 assert!(store.get("k").await.unwrap().is_none());
3449 }
3450
3451 #[tokio::test]
3452 async fn dlq_dlq_all_policy_routes_every_row_on_outer_err() {
3453 let main = FailingSink;
3454 let dlq = std::sync::Arc::new(MockSink::new());
3455 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3456 dlq_cfg.on_batch_error = OnBatchError::DlqAll;
3457
3458 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3459 records: vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})],
3460 bookmark: Some(json!("v1")),
3461 })];
3462 let stream = futures::stream::iter(pages);
3463 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3464 let result = run_stream(
3465 stream,
3466 &main,
3467 RunStreamOptions::new()
3468 .with_dlq(dlq_cfg)
3469 .with_state(std::sync::Arc::clone(&store), "k"),
3470 )
3471 .await
3472 .unwrap();
3473 assert_eq!(result.records_written, 0);
3474 {
3475 let envelopes = dlq.0.lock().unwrap();
3476 assert_eq!(envelopes.len(), 3);
3477 for env in envelopes.iter() {
3479 let msg = env["error"]["message"].as_str().unwrap();
3480 assert!(msg.contains("write failed"), "got: {msg}");
3481 }
3482 }
3483 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3484 assert_eq!(result.dlq.unwrap().records_dlq, 3);
3485 }
3486
3487 #[tokio::test]
3488 async fn dlq_per_page_budget_exceeded_aborts() {
3489 let main = PartialSink::new(vec![0, 1, 2]);
3490 let dlq = std::sync::Arc::new(MockSink::new());
3491 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3492 dlq_cfg.max_failures_per_page = Some(2);
3493
3494 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3495 records: (0..3).map(|i| json!({"i": i})).collect(),
3496 bookmark: None,
3497 })];
3498 let stream = futures::stream::iter(pages);
3499 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3500 assert!(
3501 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
3502 "got: {result:?}"
3503 );
3504 }
3505
3506 #[tokio::test]
3507 async fn dlq_total_budget_exceeded_aborts_on_later_page() {
3508 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
3509 Ok(StreamPage {
3510 records: (0..3).map(|i| json!({"i": i})).collect(),
3511 bookmark: None,
3512 }),
3513 Ok(StreamPage {
3514 records: (3..6).map(|i| json!({"i": i})).collect(),
3515 bookmark: None,
3516 }),
3517 ];
3518 let main = PartialSink::new(vec![0, 1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3521 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3522 dlq_cfg.max_failures_total = Some(4);
3523
3524 let stream = futures::stream::iter(pages);
3525 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3526 assert!(
3527 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
3528 "got: {result:?}"
3529 );
3530 }
3531
3532 #[tokio::test]
3533 async fn dlq_per_page_budget_exceeded_commits_page_before_aborting() {
3534 let main = PartialSink::new(vec![1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3541 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3542 dlq_cfg.max_failures_per_page = Some(1); let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3545 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3546 records: (0..3).map(|i| json!({ "i": i })).collect(),
3547 bookmark: Some(json!("v1")),
3548 })];
3549 let stream = futures::stream::iter(pages);
3550 let result = run_stream(
3551 stream,
3552 &main,
3553 RunStreamOptions::new()
3554 .with_dlq(dlq_cfg)
3555 .with_state(std::sync::Arc::clone(&store), "k"),
3556 )
3557 .await;
3558
3559 assert!(
3561 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
3562 "got: {result:?}"
3563 );
3564 assert_eq!(main.committed.lock().unwrap().len(), 1);
3566 assert_eq!(dlq.0.lock().unwrap().len(), 2);
3568 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3570 }
3571
3572 #[tokio::test]
3573 async fn dlq_total_budget_exceeded_commits_tripping_page_before_aborting() {
3574 let main = PartialSink::new(vec![1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3579 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3580 dlq_cfg.max_failures_total = Some(1); let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3583 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3584 records: (0..3).map(|i| json!({ "i": i })).collect(),
3585 bookmark: Some(json!("v1")),
3586 })];
3587 let stream = futures::stream::iter(pages);
3588 let result = run_stream(
3589 stream,
3590 &main,
3591 RunStreamOptions::new()
3592 .with_dlq(dlq_cfg)
3593 .with_state(std::sync::Arc::clone(&store), "k"),
3594 )
3595 .await;
3596
3597 assert!(
3598 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
3599 "got: {result:?}"
3600 );
3601 assert_eq!(main.committed.lock().unwrap().len(), 1);
3602 assert_eq!(dlq.0.lock().unwrap().len(), 2);
3603 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3604 }
3605
3606 struct FailingDlqSink;
3609 #[async_trait]
3610 impl Sink for FailingDlqSink {
3611 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
3612 Err(FaucetError::Sink("dlq disk full".into()))
3613 }
3614 }
3615
3616 struct FailingFlushDlqSink {
3620 written: std::sync::Mutex<Vec<Value>>,
3621 }
3622 impl FailingFlushDlqSink {
3623 fn new() -> Self {
3624 Self {
3625 written: std::sync::Mutex::new(Vec::new()),
3626 }
3627 }
3628 }
3629 #[async_trait]
3630 impl Sink for FailingFlushDlqSink {
3631 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3632 self.written.lock().unwrap().extend(records.iter().cloned());
3633 Ok(records.len())
3634 }
3635 async fn flush(&self) -> Result<(), FaucetError> {
3636 Err(FaucetError::Sink("dlq flush failed".into()))
3637 }
3638 }
3639
3640 #[tokio::test]
3641 async fn dlq_sink_failure_is_fatal_no_recursion() {
3642 let main = PartialSink::new(vec![0]);
3643 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingDlqSink);
3644 let dlq_cfg = DlqConfig::new(dlq);
3645
3646 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3647 records: vec![json!({"i": 0}), json!({"i": 1})],
3648 bookmark: Some(json!("v1")),
3649 })];
3650 let stream = futures::stream::iter(pages);
3651 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3652 let result = run_stream(
3653 stream,
3654 &main,
3655 RunStreamOptions::new()
3656 .with_dlq(dlq_cfg)
3657 .with_state(std::sync::Arc::clone(&store), "k"),
3658 )
3659 .await;
3660 assert!(
3661 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink write failed")),
3662 "got: {result:?}"
3663 );
3664 assert!(store.get("k").await.unwrap().is_none());
3665 }
3666
3667 #[tokio::test]
3668 async fn dlq_bookmark_advances_only_after_both_flushes() {
3669 let main = PartialSink::new(vec![1]); let dlq = std::sync::Arc::new(MockSink::new());
3671 let dlq_cfg = DlqConfig::new(dlq.clone());
3672
3673 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3674 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3675 records: vec![json!({"i": 0}), json!({"i": 1})],
3676 bookmark: Some(json!("v1")),
3677 })];
3678 let stream = futures::stream::iter(pages);
3679 run_stream(
3680 stream,
3681 &main,
3682 RunStreamOptions::new()
3683 .with_dlq(dlq_cfg)
3684 .with_state(std::sync::Arc::clone(&store), "k"),
3685 )
3686 .await
3687 .unwrap();
3688 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3689 assert_eq!(dlq.0.lock().unwrap().len(), 1);
3690 assert_eq!(main.committed.lock().unwrap().len(), 1);
3691 }
3692
3693 #[tokio::test]
3694 async fn dlq_disabled_pipeline_behaves_identically_to_today() {
3695 let main = MockSink::new();
3697 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3698 records: vec![json!({"i": 0}), json!({"i": 1})],
3699 bookmark: None,
3700 })];
3701 let stream = futures::stream::iter(pages);
3702 let result = run_stream(stream, &main, RunStreamOptions::new())
3703 .await
3704 .unwrap();
3705 assert_eq!(result.records_written, 2);
3706 assert!(result.dlq.is_none());
3707 }
3708
3709 #[tokio::test]
3710 async fn dlq_per_page_flush_failure_is_fatal_and_blocks_bookmark() {
3711 let main = PartialSink::new(vec![1]);
3716 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
3717 let dlq_cfg = DlqConfig::new(dlq);
3718
3719 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3720 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3721 records: vec![json!({"i": 0}), json!({"i": 1})],
3722 bookmark: Some(json!("v1")),
3723 })];
3724 let stream = futures::stream::iter(pages);
3725 let result = run_stream(
3726 stream,
3727 &main,
3728 RunStreamOptions::new()
3729 .with_dlq(dlq_cfg)
3730 .with_state(std::sync::Arc::clone(&store), "k"),
3731 )
3732 .await;
3733 assert!(
3734 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
3735 "got: {result:?}"
3736 );
3737 assert!(store.get("k").await.unwrap().is_none());
3738 }
3739
3740 #[tokio::test]
3741 async fn dlq_end_of_stream_flush_failure_is_fatal() {
3742 let main = PartialSink::new(vec![1]);
3747 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
3748 let dlq_cfg = DlqConfig::new(dlq);
3749
3750 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3751 records: vec![json!({"i": 0}), json!({"i": 1})],
3752 bookmark: None,
3753 })];
3754 let stream = futures::stream::iter(pages);
3755 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3756 assert!(
3757 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
3758 "got: {result:?}"
3759 );
3760 }
3761
3762 #[tokio::test]
3763 #[allow(clippy::await_holding_lock)]
3764 async fn dlq_emits_records_total_and_pages_total() {
3765 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3766 use metrics_util::debugging::DebugValue;
3767
3768 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3769 let snap = snapshotter();
3770
3771 let source = MockSource(vec![json!({"i": 0}), json!({"i": 1})]);
3772 let main = PartialSink::new(vec![1]);
3773 let dlq = std::sync::Arc::new(MockSink::new());
3774 let _ = Pipeline::new(&source, &main)
3775 .with_name("p_dlq_metrics")
3776 .with_row("r1")
3777 .with_dlq(DlqConfig::new(dlq.clone()))
3778 .run()
3779 .await
3780 .unwrap();
3781
3782 let snapshot = snap.snapshot();
3783 let mut saw_records = false;
3784 let mut saw_pages = false;
3785 for (k, _u, _d, v) in snapshot.into_vec() {
3786 let key = k.key();
3787 let labels = key.labels().collect::<Vec<_>>();
3788 let has = |k: &str, v: &str| labels.iter().any(|l| l.key() == k && l.value() == v);
3789 if key.name() == "faucet_sink_dlq_records_total"
3790 && has("pipeline", "p_dlq_metrics")
3791 && has("row", "r1")
3792 && matches!(v, DebugValue::Counter(c) if c >= 1)
3793 {
3794 saw_records = true;
3795 }
3796 if key.name() == "faucet_sink_dlq_pages_total"
3797 && has("pipeline", "p_dlq_metrics")
3798 && matches!(v, DebugValue::Counter(c) if c >= 1)
3799 {
3800 saw_pages = true;
3801 }
3802 }
3803 assert!(saw_records, "faucet_sink_dlq_records_total not emitted");
3804 assert!(saw_pages, "faucet_sink_dlq_pages_total not emitted");
3805 }
3806
3807 #[tokio::test]
3808 #[allow(clippy::await_holding_lock)]
3809 async fn dlq_budget_exceeded_emits_counter() {
3810 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3811 use metrics_util::debugging::DebugValue;
3812
3813 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3814 let snap = snapshotter();
3815
3816 let source = MockSource((0..3).map(|i| json!({"i": i})).collect());
3817 let main = PartialSink::new(vec![0, 1, 2]);
3818 let dlq = std::sync::Arc::new(MockSink::new());
3819 let mut cfg = DlqConfig::new(dlq);
3820 cfg.max_failures_per_page = Some(1);
3821 let _ = Pipeline::new(&source, &main)
3822 .with_name("p_budget")
3823 .with_dlq(cfg)
3824 .run()
3825 .await;
3826
3827 let snapshot = snap.snapshot();
3828 let saw = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
3829 k.key().name() == "faucet_sink_dlq_budget_exceeded_total"
3830 && k.key()
3831 .labels()
3832 .any(|l| l.key() == "scope" && l.value() == "per_page")
3833 && matches!(v, DebugValue::Counter(c) if c >= 1)
3834 });
3835 assert!(saw, "faucet_sink_dlq_budget_exceeded_total not emitted");
3836 }
3837
3838 #[tokio::test]
3839 async fn pipeline_run_with_dlq_routes_partial_failures_end_to_end() {
3840 let source = MockSource(vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})]);
3842 let main = PartialSink::new(vec![1]);
3843 let dlq = std::sync::Arc::new(MockSink::new());
3844
3845 let result = Pipeline::new(&source, &main)
3846 .with_dlq(DlqConfig::new(dlq.clone()))
3847 .run()
3848 .await
3849 .unwrap();
3850
3851 assert_eq!(result.records_written, 2);
3852 let stats = result.dlq.unwrap();
3853 assert_eq!(stats.records_dlq, 1);
3854 {
3855 let dlq_records = dlq.0.lock().unwrap();
3856 assert_eq!(dlq_records.len(), 1);
3857 }
3858 }
3859
3860 #[cfg(feature = "quality")]
3863 #[tokio::test]
3864 async fn quality_quarantines_to_dlq_and_writes_survivors() {
3865 use crate::dlq::DlqConfig;
3866 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3867
3868 let main = Arc::new(MockSink::new());
3869 let dlq_sink = Arc::new(MockSink::new());
3870 let spec = QualitySpec {
3871 record: vec![RecordCheck::NotNull {
3872 field: "id".into(),
3873 treat_missing_as_null: true,
3874 on_failure: OnFailure::Quarantine,
3875 }],
3876 batch: vec![],
3877 };
3878 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3879 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3880 records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
3881 bookmark: None,
3882 })];
3883 let opts = RunStreamOptions::new()
3884 .with_dlq(DlqConfig::new(dlq_sink.clone()))
3885 .with_quality(quality);
3886 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
3887 .await
3888 .unwrap();
3889
3890 assert_eq!(result.records_written, 2); assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
3892 let dlq = dlq_sink.written();
3894 assert_eq!(dlq.len(), 1);
3895 assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
3896 assert_eq!(result.dlq.unwrap().records_dlq, 1);
3897 }
3898
3899 #[cfg(feature = "quality")]
3900 #[tokio::test]
3901 #[allow(clippy::await_holding_lock)]
3902 async fn quality_only_page_emits_quality_reason() {
3903 use crate::dlq::DlqConfig;
3907 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3908 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3909 use metrics_util::debugging::DebugValue;
3910
3911 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3912 let snap = snapshotter();
3913
3914 let main = Arc::new(MockSink::new());
3917 let dlq_sink = Arc::new(MockSink::new());
3918 let spec = QualitySpec {
3919 record: vec![RecordCheck::NotNull {
3920 field: "id".into(),
3921 treat_missing_as_null: true,
3922 on_failure: OnFailure::Quarantine,
3923 }],
3924 batch: vec![],
3925 };
3926 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3927 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3928 records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
3929 bookmark: None,
3930 })];
3931 let opts = RunStreamOptions::new()
3932 .with_name("p_quality_reason")
3933 .with_dlq(DlqConfig::new(dlq_sink.clone()))
3934 .with_quality(quality);
3935 let _ = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
3936 .await
3937 .unwrap();
3938
3939 let snapshot = snap.snapshot();
3940 let saw_quality_reason = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
3941 k.key().name() == "faucet_sink_dlq_pages_total"
3942 && k.key()
3943 .labels()
3944 .any(|l| l.key() == "pipeline" && l.value() == "p_quality_reason")
3945 && k.key()
3946 .labels()
3947 .any(|l| l.key() == "reason" && l.value() == "quality")
3948 && matches!(v, DebugValue::Counter(c) if c >= 1)
3949 });
3950 assert!(
3951 saw_quality_reason,
3952 "expected faucet_sink_dlq_pages_total with reason=\"quality\""
3953 );
3954 }
3955
3956 #[cfg(feature = "quality")]
3957 #[tokio::test]
3958 async fn quality_abort_fails_run() {
3959 use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
3960 let main = MockSink::new();
3961 let spec = QualitySpec {
3962 record: vec![],
3963 batch: vec![BatchCheck::RowCount {
3964 min: Some(5),
3965 max: None,
3966 on_failure: OnFailure::Abort,
3967 }],
3968 };
3969 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3970 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3971 records: vec![json!({"id": 1})],
3972 bookmark: None,
3973 })];
3974 let opts = RunStreamOptions::new().with_quality(quality);
3975 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
3976 assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
3977 }
3978
3979 #[cfg(feature = "quality")]
3980 #[tokio::test]
3981 async fn quality_quarantine_without_dlq_is_rejected() {
3982 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3983 let main = MockSink::new();
3984 let spec = QualitySpec {
3985 record: vec![RecordCheck::NotNull {
3986 field: "id".into(),
3987 treat_missing_as_null: true,
3988 on_failure: OnFailure::Quarantine,
3989 }],
3990 batch: vec![],
3991 };
3992 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3993 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3994 records: vec![json!({"id": null})],
3995 bookmark: None,
3996 })];
3997 let opts = RunStreamOptions::new().with_quality(quality);
3999 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4000 assert!(matches!(result, Err(FaucetError::Config(_))));
4001 }
4002
4003 #[cfg(feature = "contract")]
4004 fn compiled_contract(on_breach: &str) -> Arc<crate::contract::CompiledContract> {
4005 let spec: crate::contract::ContractSpec = serde_json::from_value(json!({
4006 "version": "1.0.0",
4007 "on_breach": on_breach,
4008 "fields": [{ "name": "id", "type": "integer" }]
4009 }))
4010 .unwrap();
4011 Arc::new(crate::contract::CompiledContract::compile(&spec).unwrap())
4012 }
4013
4014 #[cfg(feature = "contract")]
4015 #[tokio::test]
4016 async fn contract_quarantines_to_dlq_and_writes_survivors() {
4017 use crate::dlq::DlqConfig;
4018 let main = Arc::new(MockSink::new());
4019 let dlq_sink = Arc::new(MockSink::new());
4020 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4021 records: vec![json!({"id": 1}), json!({"id": "bad"}), json!({"id": 3})],
4022 bookmark: None,
4023 })];
4024 let opts = RunStreamOptions::new()
4025 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4026 .with_contract(compiled_contract("quarantine"));
4027 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4028 .await
4029 .unwrap();
4030
4031 assert_eq!(result.records_written, 2);
4032 assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
4033 let dlq = dlq_sink.written();
4034 assert_eq!(dlq.len(), 1);
4035 assert_eq!(dlq[0]["error"]["kind"], "ContractViolation");
4036 assert_eq!(dlq[0]["payload"], json!({"id": "bad"}));
4037 assert_eq!(dlq[0]["record_index"], 1);
4039 assert_eq!(result.dlq.unwrap().records_dlq, 1);
4040 }
4041
4042 #[cfg(feature = "contract")]
4043 #[tokio::test]
4044 async fn contract_fail_aborts_run_and_writes_nothing() {
4045 let main = MockSink::new();
4046 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4047 records: vec![json!({"id": 1}), json!({"id": "bad"})],
4048 bookmark: None,
4049 })];
4050 let opts = RunStreamOptions::new().with_contract(compiled_contract("fail"));
4051 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4052 match result {
4053 Err(FaucetError::ContractViolation { version, message }) => {
4054 assert_eq!(version, "1.0.0");
4055 assert!(message.contains("id"), "message: {message}");
4056 }
4057 other => panic!("expected ContractViolation, got {other:?}"),
4058 }
4059 assert!(
4060 main.written().is_empty(),
4061 "a contract fail must not commit any of the page's records"
4062 );
4063 }
4064
4065 #[cfg(feature = "contract")]
4066 #[tokio::test]
4067 async fn contract_warn_writes_everything() {
4068 let main = MockSink::new();
4069 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4070 records: vec![json!({"id": 1}), json!({"id": "bad"})],
4071 bookmark: None,
4072 })];
4073 let opts = RunStreamOptions::new().with_contract(compiled_contract("warn"));
4074 let result = run_stream(futures::stream::iter(pages), &main, opts)
4075 .await
4076 .unwrap();
4077 assert_eq!(result.records_written, 2);
4078 assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": "bad"})]);
4079 }
4080
4081 #[cfg(feature = "contract")]
4082 #[tokio::test]
4083 async fn contract_quarantine_without_dlq_is_rejected() {
4084 let main = MockSink::new();
4085 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4086 records: vec![json!({"id": 1})],
4087 bookmark: None,
4088 })];
4089 let opts = RunStreamOptions::new().with_contract(compiled_contract("quarantine"));
4091 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4092 assert!(matches!(result, Err(FaucetError::Config(_))));
4093 }
4094
4095 #[cfg(all(feature = "contract", feature = "quality"))]
4096 #[tokio::test]
4097 async fn contract_runs_after_quality_and_shares_dlq() {
4098 use crate::dlq::DlqConfig;
4102 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
4103 let main = Arc::new(MockSink::new());
4104 let dlq_sink = Arc::new(MockSink::new());
4105 let quality = Arc::new(
4106 CompiledQuality::compile(&QualitySpec {
4107 record: vec![RecordCheck::NotNull {
4108 field: "id".into(),
4109 treat_missing_as_null: true,
4110 on_failure: OnFailure::Quarantine,
4111 }],
4112 batch: vec![],
4113 })
4114 .unwrap(),
4115 );
4116 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4117 records: vec![json!({"id": null}), json!({"id": "bad"}), json!({"id": 3})],
4118 bookmark: None,
4119 })];
4120 let opts = RunStreamOptions::new()
4121 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4122 .with_quality(quality)
4123 .with_contract(compiled_contract("quarantine"));
4124 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4125 .await
4126 .unwrap();
4127
4128 assert_eq!(result.records_written, 1);
4129 assert_eq!(main.written(), vec![json!({"id": 3})]);
4130 let dlq = dlq_sink.written();
4131 assert_eq!(dlq.len(), 2);
4132 let kinds: Vec<&str> = dlq
4133 .iter()
4134 .map(|e| e["error"]["kind"].as_str().unwrap())
4135 .collect();
4136 assert!(kinds.contains(&"QualityFailure"), "kinds: {kinds:?}");
4137 assert!(kinds.contains(&"ContractViolation"), "kinds: {kinds:?}");
4138 assert_eq!(result.dlq.unwrap().records_dlq, 2);
4139
4140 let contract_env = dlq
4146 .iter()
4147 .find(|e| e["error"]["kind"] == "ContractViolation")
4148 .expect("a contract envelope");
4149 assert_eq!(
4150 contract_env["record_index"], 1,
4151 "contract record_index must be the original page position: {contract_env}"
4152 );
4153 let quality_env = dlq
4154 .iter()
4155 .find(|e| e["error"]["kind"] == "QualityFailure")
4156 .expect("a quality envelope");
4157 assert_eq!(
4158 quality_env["record_index"], 0,
4159 "null id was at page index 0"
4160 );
4161 }
4162
4163 struct FlakySink {
4166 every: usize,
4167 calls: std::sync::Mutex<Vec<usize>>,
4168 }
4169 impl FlakySink {
4170 fn new(every: usize) -> Self {
4171 Self {
4172 every,
4173 calls: std::sync::Mutex::new(Vec::new()),
4174 }
4175 }
4176 fn call_sizes(&self) -> Vec<usize> {
4177 self.calls.lock().unwrap().clone()
4178 }
4179 }
4180 #[async_trait]
4181 impl Sink for FlakySink {
4182 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4183 Ok(records.len())
4184 }
4185 async fn write_batch_partial(
4186 &self,
4187 records: &[Value],
4188 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4189 self.calls.lock().unwrap().push(records.len());
4190 Ok(records
4191 .iter()
4192 .enumerate()
4193 .map(|(i, _)| {
4194 if (i + 1) % self.every == 0 {
4195 Err(FaucetError::Sink("synthetic".into()))
4196 } else {
4197 Ok(())
4198 }
4199 })
4200 .collect())
4201 }
4202 }
4203
4204 #[tokio::test]
4205 async fn adaptive_shrinks_under_errors_on_dlq_path() {
4206 use crate::adaptive::AdaptiveBatchConfig;
4207 use crate::dlq::{DlqConfig, OnBatchError};
4208 let mk = || StreamPage {
4213 records: (0..400).map(|i| json!({"i": i})).collect(),
4214 bookmark: None,
4215 };
4216 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4217 let sink = FlakySink::new(4); let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
4219 let dlq = DlqConfig {
4220 sink: dlq_sink,
4221 on_batch_error: OnBatchError::Propagate,
4222 max_failures_per_page: None,
4223 max_failures_total: None,
4224 include_original_payload: true,
4225 };
4226 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4227 "enabled": true, "min": 50, "max": 400,
4228 "decrease_factor": 0.5, "cooldown_batches": 0, "error_threshold": 0.1
4229 }))
4230 .unwrap();
4231 let opts = RunStreamOptions::new().with_dlq(dlq).with_adaptive(cfg);
4232 let result = run_stream(stream, &sink, opts).await.unwrap();
4233 assert!(
4240 result.records_written >= 900,
4241 "expected ≥900 written, got {}",
4242 result.records_written
4243 );
4244 let sizes = sink.call_sizes();
4245 assert_eq!(sizes[0], 400, "first chunk is the full page");
4246 assert!(
4247 sizes.last().unwrap() < &400,
4248 "controller should shrink under errors: {sizes:?}"
4249 );
4250 assert!(
4251 result.dlq.unwrap().records_dlq >= 250,
4252 "expected ≥250 DLQ records"
4253 );
4254 }
4255
4256 struct RecordingSink {
4261 calls: std::sync::Mutex<Vec<usize>>,
4262 latency: std::time::Duration,
4263 }
4264 impl RecordingSink {
4265 fn new(latency_ms: u64) -> Self {
4266 Self {
4267 calls: std::sync::Mutex::new(Vec::new()),
4268 latency: std::time::Duration::from_millis(latency_ms),
4269 }
4270 }
4271 fn call_sizes(&self) -> Vec<usize> {
4272 self.calls.lock().unwrap().clone()
4273 }
4274 }
4275 #[async_trait]
4276 impl Sink for RecordingSink {
4277 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4278 tokio::time::sleep(self.latency).await;
4279 self.calls.lock().unwrap().push(records.len());
4280 Ok(records.len())
4281 }
4282 }
4283
4284 #[tokio::test]
4285 async fn adaptive_reslices_non_dlq_page_into_subbatches() {
4286 use crate::adaptive::AdaptiveBatchConfig;
4287 let page = StreamPage {
4288 records: (0..1000).map(|i| json!({ "i": i })).collect(),
4289 bookmark: None,
4290 };
4291 let stream = futures::stream::iter(vec![Ok(page)]);
4292 let sink = RecordingSink::new(0);
4293 let cfg: AdaptiveBatchConfig =
4294 serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
4295 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4296 .await
4297 .unwrap();
4298 assert_eq!(result.records_written, 1000);
4299 assert_eq!(sink.call_sizes(), vec![1000]);
4301 }
4302
4303 #[tokio::test]
4304 async fn adaptive_shrinks_under_latency_target_then_smaller_chunks() {
4305 use crate::adaptive::AdaptiveBatchConfig;
4306 let mk = || StreamPage {
4307 records: (0..400).map(|i| json!({"i": i})).collect(),
4308 bookmark: None,
4309 };
4310 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4311 let sink = RecordingSink::new(50);
4312 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4313 "enabled": true, "min": 50, "max": 400,
4314 "decrease_factor": 0.5, "cooldown_batches": 0,
4315 "target_latency_ms": 10, "latency_window": 1
4316 }))
4317 .unwrap();
4318 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4319 .await
4320 .unwrap();
4321 assert_eq!(result.records_written, 1200);
4322 let sizes = sink.call_sizes();
4323 assert_eq!(sizes[0], 400);
4324 assert!(
4325 sizes.last().unwrap() < &400,
4326 "controller should have shrunk: {sizes:?}"
4327 );
4328 }
4329
4330 #[tokio::test]
4331 #[allow(clippy::await_holding_lock)]
4332 async fn adaptive_emits_batch_size_and_adjustments_metrics() {
4333 use crate::adaptive::AdaptiveBatchConfig;
4336 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
4337 use metrics_util::debugging::DebugValue;
4338
4339 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
4340 let snap = snapshotter();
4341
4342 let mk = || StreamPage {
4348 records: (0..400).map(|i| json!({"i": i})).collect(),
4349 bookmark: None,
4350 };
4351 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4352 let sink = RecordingSink::new(50);
4353 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4354 "enabled": true, "min": 50, "max": 400,
4355 "decrease_factor": 0.5, "cooldown_batches": 0,
4356 "target_latency_ms": 10, "latency_window": 1
4357 }))
4358 .unwrap();
4359
4360 let _ = run_stream(
4361 stream,
4362 &sink,
4363 RunStreamOptions::new()
4364 .with_adaptive(cfg)
4365 .with_name("p")
4366 .with_row("r"),
4367 )
4368 .await
4369 .unwrap();
4370
4371 let snapshot = snap.snapshot();
4372 let mut saw_batch_size = false;
4373 let mut saw_adjustments = false;
4374 for (k, _u, _d, v) in snapshot.into_vec() {
4375 let key = k.key();
4376 let labels = key.labels().collect::<Vec<_>>();
4377 let has = |k: &str, val: &str| labels.iter().any(|l| l.key() == k && l.value() == val);
4378
4379 if key.name() == "faucet_pipeline_adaptive_batch_size"
4380 && has("pipeline", "p")
4381 && has("row", "r")
4382 && matches!(v, DebugValue::Gauge(_))
4383 {
4384 saw_batch_size = true;
4385 }
4386 if key.name() == "faucet_pipeline_adaptive_batch_adjustments_total"
4387 && has("pipeline", "p")
4388 && has("row", "r")
4389 && matches!(v, DebugValue::Counter(c) if c >= 1)
4390 {
4391 saw_adjustments = true;
4392 }
4393 }
4394 assert!(
4395 saw_batch_size,
4396 "expected faucet_pipeline_adaptive_batch_size gauge with pipeline=p, row=r"
4397 );
4398 assert!(
4399 saw_adjustments,
4400 "expected faucet_pipeline_adaptive_batch_adjustments_total counter with pipeline=p, row=r"
4401 );
4402 }
4403
4404 #[tokio::test]
4407 async fn exactly_once_rejects_dlq() {
4408 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4411 let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
4412 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
4413 let opts = eo_opts(store, "k", 0).with_dlq(DlqConfig::new(dlq_sink));
4414 let r = run_stream(
4415 futures::stream::iter(pages),
4416 &IdempotentMockSink::new(),
4417 opts,
4418 )
4419 .await;
4420 assert!(
4421 matches!(&r, Err(FaucetError::Config(m)) if m.contains("not compatible with a DLQ")),
4422 "got: {r:?}"
4423 );
4424 }
4425
4426 #[tokio::test]
4429 async fn exactly_once_writes_unbookmarked_page_at_least_once() {
4430 let sink = IdempotentMockSink::new();
4435 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4436 let pages = vec![Ok(StreamPage {
4437 records: vec![json!({"id": 1}), json!({"id": 2})],
4438 bookmark: None,
4439 })];
4440 let r = run_stream(
4441 futures::stream::iter(pages),
4442 &sink,
4443 eo_opts(store.clone(), "k", 0),
4444 )
4445 .await
4446 .unwrap();
4447 assert_eq!(r.records_written, 2);
4448 assert_eq!(r.bookmark, None);
4449 assert_eq!(sink.last_committed_token("k").await.unwrap(), None);
4452 assert!(store.get("k").await.unwrap().is_none());
4454 assert_eq!(sink.rows(), vec![json!({"id": 1}), json!({"id": 2})]);
4455 }
4456
4457 #[tokio::test]
4460 async fn dlq_with_bookmark_persists_after_routing_failures() {
4461 let main = PartialSink::new(vec![1]); let dlq = std::sync::Arc::new(MockSink::new());
4466 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4467 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4468 records: vec![json!({"i": 0}), json!({"i": 1})],
4469 bookmark: Some(json!("ckpt")),
4470 })];
4471 let result = run_stream(
4472 futures::stream::iter(pages),
4473 &main,
4474 RunStreamOptions::new()
4475 .with_dlq(DlqConfig::new(dlq.clone()))
4476 .with_state(Arc::clone(&store), "k"),
4477 )
4478 .await
4479 .unwrap();
4480
4481 assert_eq!(result.records_written, 1); assert_eq!(result.bookmark, Some(json!("ckpt")));
4483 assert_eq!(store.get("k").await.unwrap(), Some(json!("ckpt")));
4485 let envelopes = dlq.0.lock().unwrap();
4487 assert_eq!(envelopes.len(), 1);
4488 assert_eq!(envelopes[0]["payload"]["i"], 1);
4489 }
4490
4491 #[tokio::test]
4494 async fn exactly_once_resumes_sequence_from_start_seq() {
4495 let sink = IdempotentMockSink::new();
4499 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4500 let pages = vec![Ok(StreamPage {
4501 records: vec![json!({"id": 9})],
4502 bookmark: Some(json!("bm-after-resume")),
4503 })];
4504 let r = run_stream(
4505 futures::stream::iter(pages),
4506 &sink,
4507 eo_opts(store.clone(), "eo_key", 7),
4508 )
4509 .await
4510 .unwrap();
4511
4512 assert_eq!(r.records_written, 1);
4513 let (bm, seq) =
4514 crate::idempotency::unwrap_state(&store.get("eo_key").await.unwrap().unwrap());
4515 assert_eq!(bm, Some(json!("bm-after-resume")));
4516 assert_eq!(seq, 8, "sequence resumes at start_seq + 1");
4517 let token = sink.last_committed_token("eo_key").await.unwrap().unwrap();
4518 assert_eq!(
4519 crate::idempotency::parse_token_parts(&token),
4520 Some((8, Some(json!("bm-after-resume"))))
4521 );
4522 }
4523
4524 #[cfg(feature = "quality")]
4527 #[tokio::test]
4528 async fn pipeline_run_with_quality_aborts_on_failed_batch_check() {
4529 use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
4530 let source = MockSource(vec![json!({"id": 1})]);
4531 let main = MockSink::new();
4532 let spec = QualitySpec {
4533 record: vec![],
4534 batch: vec![BatchCheck::RowCount {
4535 min: Some(5),
4536 max: None,
4537 on_failure: OnFailure::Abort,
4538 }],
4539 };
4540 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4541 let result = Pipeline::new(&source, &main)
4542 .with_quality(quality)
4543 .run()
4544 .await;
4545 assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
4546 assert!(main.written().is_empty());
4548 }
4549
4550 #[tokio::test]
4553 async fn pipeline_run_with_adaptive_reslices_page() {
4554 use crate::adaptive::AdaptiveBatchConfig;
4555 let source = MockSource((0..1000).map(|i| json!({ "i": i })).collect());
4556 let sink = MockSink::new();
4557 let cfg: AdaptiveBatchConfig =
4558 serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
4559 let result = Pipeline::new(&source, &sink)
4560 .with_adaptive(cfg)
4561 .run()
4562 .await
4563 .unwrap();
4564 assert_eq!(result.records_written, 1000);
4565 assert_eq!(sink.written().len(), 1000);
4566 }
4567
4568 #[tokio::test]
4571 async fn adaptive_noop_sink_name_is_handled() {
4572 use crate::adaptive::AdaptiveBatchConfig;
4573
4574 struct JsonlNamedSink(std::sync::Mutex<usize>);
4577 #[async_trait]
4578 impl Sink for JsonlNamedSink {
4579 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4580 *self.0.lock().unwrap() += records.len();
4581 Ok(records.len())
4582 }
4583 fn connector_name(&self) -> &'static str {
4584 "jsonl"
4585 }
4586 }
4587
4588 let page = StreamPage {
4589 records: (0..10).map(|i| json!({ "i": i })).collect(),
4590 bookmark: None,
4591 };
4592 let stream = futures::stream::iter(vec![Ok(page)]);
4593 let sink = JsonlNamedSink(std::sync::Mutex::new(0));
4594 let cfg: AdaptiveBatchConfig =
4595 serde_json::from_value(json!({"enabled": true, "min": 5, "max": 10})).unwrap();
4596 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4597 .await
4598 .unwrap();
4599 assert_eq!(result.records_written, 10);
4600 assert_eq!(*sink.0.lock().unwrap(), 10);
4601 }
4602
4603 #[tokio::test]
4604 async fn resilience_retries_transient_sink_write() {
4605 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4606 use std::sync::Arc;
4607 use std::sync::atomic::{AtomicU32, Ordering};
4608 use std::time::Duration;
4609
4610 struct TransientFlakySink {
4611 attempts: Arc<AtomicU32>,
4612 written: Arc<AtomicU32>,
4613 }
4614 #[async_trait::async_trait]
4615 impl Sink for TransientFlakySink {
4616 async fn write_batch(
4617 &self,
4618 records: &[serde_json::Value],
4619 ) -> Result<usize, FaucetError> {
4620 let n = self.attempts.fetch_add(1, Ordering::SeqCst);
4621 if n < 2 {
4622 return Err(FaucetError::HttpStatus {
4623 status: 503,
4624 url: "u".into(),
4625 body: "".into(),
4626 });
4627 }
4628 self.written
4629 .fetch_add(records.len() as u32, Ordering::SeqCst);
4630 Ok(records.len())
4631 }
4632 async fn flush(&self) -> Result<(), FaucetError> {
4633 Ok(())
4634 }
4635 fn supports_idempotent_writes(&self) -> bool {
4637 true
4638 }
4639 }
4640
4641 let attempts = Arc::new(AtomicU32::new(0));
4642 let written = Arc::new(AtomicU32::new(0));
4643 let sink = TransientFlakySink {
4644 attempts: attempts.clone(),
4645 written: written.clone(),
4646 };
4647 let pages = futures::stream::iter(vec![Ok(StreamPage {
4648 records: vec![serde_json::json!({"a": 1})],
4649 bookmark: None,
4650 })]);
4651 let policy = ResiliencePolicy {
4652 retry: RetryPolicy {
4653 max_attempts: 5,
4654 backoff: BackoffKind::None,
4655 base: Duration::ZERO,
4656 max: Duration::ZERO,
4657 jitter: false,
4658 ..RetryPolicy::default()
4659 },
4660 ..ResiliencePolicy::default()
4661 };
4662 let res = run_stream(
4663 pages,
4664 &sink,
4665 RunStreamOptions::new().with_resilience(policy),
4666 )
4667 .await
4668 .unwrap();
4669 assert_eq!(written.load(Ordering::SeqCst), 1);
4670 assert_eq!(attempts.load(Ordering::SeqCst), 3);
4671 assert_eq!(res.records_written, 1);
4672 }
4673
4674 #[tokio::test]
4675 async fn resilience_does_not_retry_non_idempotent_write_batch() {
4676 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4680 use std::sync::Arc;
4681 use std::sync::atomic::{AtomicU32, Ordering};
4682 use std::time::Duration;
4683
4684 struct NonIdempotentFlakySink {
4685 attempts: Arc<AtomicU32>,
4686 }
4687 #[async_trait::async_trait]
4688 impl Sink for NonIdempotentFlakySink {
4689 async fn write_batch(&self, _r: &[Value]) -> Result<usize, FaucetError> {
4690 self.attempts.fetch_add(1, Ordering::SeqCst);
4691 Err(FaucetError::HttpStatus {
4692 status: 503,
4693 url: "u".into(),
4694 body: "".into(),
4695 })
4696 }
4697 async fn flush(&self) -> Result<(), FaucetError> {
4698 Ok(())
4699 }
4700 }
4702
4703 let attempts = Arc::new(AtomicU32::new(0));
4704 let sink = NonIdempotentFlakySink {
4705 attempts: attempts.clone(),
4706 };
4707 let pages = futures::stream::iter(vec![Ok(StreamPage {
4708 records: vec![json!({"a": 1})],
4709 bookmark: None,
4710 })]);
4711 let policy = ResiliencePolicy {
4712 retry: RetryPolicy {
4713 max_attempts: 5,
4714 backoff: BackoffKind::None,
4715 base: Duration::ZERO,
4716 max: Duration::ZERO,
4717 jitter: false,
4718 ..RetryPolicy::default()
4719 },
4720 ..ResiliencePolicy::default()
4721 };
4722 let err = run_stream(
4723 pages,
4724 &sink,
4725 RunStreamOptions::new().with_resilience(policy),
4726 )
4727 .await
4728 .unwrap_err();
4729 assert!(matches!(err, FaucetError::HttpStatus { status: 503, .. }));
4730 assert_eq!(
4731 attempts.load(Ordering::SeqCst),
4732 1,
4733 "non-idempotent write_batch must be attempted exactly once (no retry)"
4734 );
4735 }
4736
4737 #[tokio::test]
4738 async fn resilience_circuit_opens_after_consecutive_failed_pages() {
4739 use crate::dlq::{DlqConfig, OnBatchError};
4740 use crate::resilience::{BackoffKind, CircuitBreakerConfig, ResiliencePolicy, RetryPolicy};
4741 use std::sync::Arc;
4742 use std::time::Duration;
4743
4744 struct DeadSink;
4746 #[async_trait]
4747 impl Sink for DeadSink {
4748 async fn write_batch(&self, _r: &[Value]) -> Result<usize, FaucetError> {
4749 Err(FaucetError::Sink("down".into()))
4750 }
4751 async fn flush(&self) -> Result<(), FaucetError> {
4752 Ok(())
4753 }
4754 }
4755 struct NullSink;
4757 #[async_trait]
4758 impl Sink for NullSink {
4759 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4760 Ok(r.len())
4761 }
4762 async fn flush(&self) -> Result<(), FaucetError> {
4763 Ok(())
4764 }
4765 }
4766
4767 let pages = futures::stream::iter((0..10).map(|i| {
4768 Ok(StreamPage {
4769 records: vec![json!({"i": i})],
4770 bookmark: None,
4771 })
4772 }));
4773 let dlq = DlqConfig {
4774 on_batch_error: OnBatchError::DlqAll,
4775 ..DlqConfig::new(Arc::new(NullSink))
4776 };
4777 let policy = ResiliencePolicy {
4778 retry: RetryPolicy {
4779 max_attempts: 1,
4780 backoff: BackoffKind::None,
4781 base: Duration::ZERO,
4782 max: Duration::ZERO,
4783 jitter: false,
4784 ..RetryPolicy::default()
4785 },
4786 circuit_breaker: Some(CircuitBreakerConfig {
4787 consecutive_failures: 3,
4788 cooldown: Duration::from_secs(60),
4789 }),
4790 poison: None,
4791 };
4792 let err = run_stream(
4793 pages,
4794 &DeadSink,
4795 RunStreamOptions::new()
4796 .with_dlq(dlq)
4797 .with_resilience(policy),
4798 )
4799 .await
4800 .unwrap_err();
4801 assert!(
4802 matches!(err, FaucetError::CircuitOpen { failures: 3, .. }),
4803 "got {err:?}"
4804 );
4805 }
4806
4807 #[tokio::test]
4808 async fn resilience_poison_retries_then_dlqs_failing_row() {
4809 use crate::dlq::DlqConfig;
4810 use crate::resilience::{
4811 BackoffKind, PoisonAction, PoisonPolicy, ResiliencePolicy, RetryPolicy,
4812 };
4813 use std::sync::{Arc, Mutex};
4814 use std::time::Duration;
4815
4816 struct PickySink {
4819 bad_attempts: Arc<Mutex<u32>>,
4820 }
4821 #[async_trait]
4822 impl Sink for PickySink {
4823 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4824 Ok(r.len())
4825 }
4826 async fn write_batch_partial(
4827 &self,
4828 records: &[Value],
4829 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4830 Ok(records
4831 .iter()
4832 .map(|rec| {
4833 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
4834 *self.bad_attempts.lock().unwrap() += 1;
4835 Err(FaucetError::HttpStatus {
4836 status: 503,
4837 url: "u".into(),
4838 body: "".into(),
4839 })
4840 } else {
4841 Ok(())
4842 }
4843 })
4844 .collect())
4845 }
4846 async fn flush(&self) -> Result<(), FaucetError> {
4847 Ok(())
4848 }
4849 }
4850 struct CaptureSink(Arc<Mutex<Vec<Value>>>);
4851 #[async_trait]
4852 impl Sink for CaptureSink {
4853 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4854 self.0.lock().unwrap().extend_from_slice(r);
4855 Ok(r.len())
4856 }
4857 async fn flush(&self) -> Result<(), FaucetError> {
4858 Ok(())
4859 }
4860 }
4861
4862 let captured = Arc::new(Mutex::new(Vec::new()));
4863 let bad_attempts = Arc::new(Mutex::new(0u32));
4864 let sink = PickySink {
4865 bad_attempts: bad_attempts.clone(),
4866 };
4867 let pages = futures::stream::iter(vec![Ok(StreamPage {
4868 records: vec![json!({"ok": 1}), json!({"bad": true})],
4869 bookmark: None,
4870 })]);
4871 let policy = ResiliencePolicy {
4872 retry: RetryPolicy {
4873 max_attempts: 1,
4874 backoff: BackoffKind::None,
4875 base: Duration::ZERO,
4876 max: Duration::ZERO,
4877 jitter: false,
4878 ..RetryPolicy::default()
4879 },
4880 circuit_breaker: None,
4881 poison: Some(PoisonPolicy {
4882 max_row_attempts: 3,
4883 action: PoisonAction::Dlq,
4884 }),
4885 };
4886 let res = run_stream(
4887 pages,
4888 &sink,
4889 RunStreamOptions::new()
4890 .with_dlq(DlqConfig::new(Arc::new(CaptureSink(captured.clone()))))
4891 .with_resilience(policy),
4892 )
4893 .await
4894 .unwrap();
4895
4896 assert_eq!(
4897 *bad_attempts.lock().unwrap(),
4898 3,
4899 "bad row tried max_row_attempts times"
4900 );
4901 assert_eq!(res.records_written, 1, "the ok row");
4902 let dlq = captured.lock().unwrap();
4903 assert_eq!(dlq.len(), 1, "one row to DLQ");
4904 assert_eq!(dlq[0]["payload"]["bad"], json!(true));
4905 }
4906
4907 #[tokio::test]
4908 async fn poison_loop_does_not_nest_resilience_retry_on_subset_resubmit() {
4909 use crate::dlq::DlqConfig;
4917 use crate::resilience::{
4918 BackoffKind, PoisonAction, PoisonPolicy, ResiliencePolicy, RetryPolicy,
4919 };
4920 use std::sync::{Arc, Mutex};
4921 use std::time::Duration;
4922
4923 struct OuterErrOnSubsetSink {
4924 subset_calls: Arc<Mutex<u32>>,
4925 }
4926 #[async_trait]
4927 impl Sink for OuterErrOnSubsetSink {
4928 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4929 Ok(r.len())
4930 }
4931 async fn write_batch_partial(
4932 &self,
4933 records: &[Value],
4934 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4935 if records.len() == 1 {
4937 *self.subset_calls.lock().unwrap() += 1;
4938 return Err(FaucetError::HttpStatus {
4939 status: 503,
4940 url: "u".into(),
4941 body: "".into(),
4942 });
4943 }
4944 Ok(records
4945 .iter()
4946 .map(|rec| {
4947 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
4948 Err(FaucetError::HttpStatus {
4949 status: 503,
4950 url: "u".into(),
4951 body: "".into(),
4952 })
4953 } else {
4954 Ok(())
4955 }
4956 })
4957 .collect())
4958 }
4959 async fn flush(&self) -> Result<(), FaucetError> {
4960 Ok(())
4961 }
4962 fn supports_idempotent_writes(&self) -> bool {
4965 true
4966 }
4967 }
4968 struct CaptureSink(Arc<Mutex<Vec<Value>>>);
4969 #[async_trait]
4970 impl Sink for CaptureSink {
4971 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4972 self.0.lock().unwrap().extend_from_slice(r);
4973 Ok(r.len())
4974 }
4975 async fn flush(&self) -> Result<(), FaucetError> {
4976 Ok(())
4977 }
4978 }
4979
4980 let subset_calls = Arc::new(Mutex::new(0u32));
4981 let sink = OuterErrOnSubsetSink {
4982 subset_calls: subset_calls.clone(),
4983 };
4984 let pages = futures::stream::iter(vec![Ok(StreamPage {
4985 records: vec![json!({"ok": 1}), json!({"bad": true})],
4986 bookmark: None,
4987 })]);
4988 let policy = ResiliencePolicy {
4989 retry: RetryPolicy {
4990 max_attempts: 4, backoff: BackoffKind::None,
4992 base: Duration::ZERO,
4993 max: Duration::ZERO,
4994 jitter: false,
4995 ..RetryPolicy::default()
4996 },
4997 circuit_breaker: None,
4998 poison: Some(PoisonPolicy {
4999 max_row_attempts: 3,
5000 action: PoisonAction::Dlq,
5001 }),
5002 };
5003 let res = run_stream(
5004 pages,
5005 &sink,
5006 RunStreamOptions::new()
5007 .with_dlq(DlqConfig::new(Arc::new(CaptureSink(Arc::new(Mutex::new(
5008 Vec::new(),
5009 ))))))
5010 .with_resilience(policy),
5011 )
5012 .await;
5013
5014 assert!(matches!(
5016 res,
5017 Err(FaucetError::HttpStatus { status: 503, .. })
5018 ));
5019 assert_eq!(
5022 *subset_calls.lock().unwrap(),
5023 1,
5024 "poison subset resubmit must be bare (no nested resilience retry)"
5025 );
5026 }
5027
5028 #[tokio::test]
5029 async fn resilience_retries_transient_flush_and_state_on_dlq_path() {
5030 use crate::dlq::DlqConfig;
5037 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
5038 use crate::state::StateStore;
5039 use std::sync::Arc;
5040 use std::sync::atomic::{AtomicU32, Ordering};
5041 use std::time::Duration;
5042
5043 fn transient_503() -> FaucetError {
5044 FaucetError::HttpStatus {
5045 status: 503,
5046 url: "u".into(),
5047 body: "".into(),
5048 }
5049 }
5050
5051 struct FlakyFlushSink {
5054 flush_attempts: Arc<AtomicU32>,
5055 }
5056 #[async_trait]
5057 impl Sink for FlakyFlushSink {
5058 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5059 Ok(r.len())
5060 }
5061 async fn write_batch_partial(
5062 &self,
5063 records: &[Value],
5064 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
5065 Ok(records
5066 .iter()
5067 .map(|rec| {
5068 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
5069 Err(transient_503())
5070 } else {
5071 Ok(())
5072 }
5073 })
5074 .collect())
5075 }
5076 async fn flush(&self) -> Result<(), FaucetError> {
5077 let n = self.flush_attempts.fetch_add(1, Ordering::SeqCst);
5078 if n < 2 { Err(transient_503()) } else { Ok(()) }
5079 }
5080 }
5081 struct NullSink;
5082 #[async_trait]
5083 impl Sink for NullSink {
5084 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5085 Ok(r.len())
5086 }
5087 async fn flush(&self) -> Result<(), FaucetError> {
5088 Ok(())
5089 }
5090 }
5091
5092 struct FlakyStore {
5094 put_attempts: Arc<AtomicU32>,
5095 value: Arc<std::sync::Mutex<Option<Value>>>,
5096 }
5097 #[async_trait]
5098 impl StateStore for FlakyStore {
5099 async fn get(&self, _key: &str) -> Result<Option<Value>, FaucetError> {
5100 Ok(self.value.lock().unwrap().clone())
5101 }
5102 async fn put(&self, _key: &str, value: &Value) -> Result<(), FaucetError> {
5103 let n = self.put_attempts.fetch_add(1, Ordering::SeqCst);
5104 if n < 2 {
5105 return Err(transient_503());
5106 }
5107 *self.value.lock().unwrap() = Some(value.clone());
5108 Ok(())
5109 }
5110 async fn delete(&self, _key: &str) -> Result<(), FaucetError> {
5111 Ok(())
5112 }
5113 }
5114
5115 let flush_attempts = Arc::new(AtomicU32::new(0));
5116 let put_attempts = Arc::new(AtomicU32::new(0));
5117 let stored = Arc::new(std::sync::Mutex::new(None));
5118 let sink = FlakyFlushSink {
5119 flush_attempts: flush_attempts.clone(),
5120 };
5121 let store: Arc<dyn StateStore> = Arc::new(FlakyStore {
5122 put_attempts: put_attempts.clone(),
5123 value: stored.clone(),
5124 });
5125 let pages = futures::stream::iter(vec![Ok(StreamPage {
5126 records: vec![json!({"ok": 1}), json!({"bad": true})],
5127 bookmark: Some(json!({"cursor": 42})),
5128 })]);
5129 let policy = ResiliencePolicy {
5130 retry: RetryPolicy {
5131 max_attempts: 5,
5132 backoff: BackoffKind::None,
5133 base: Duration::ZERO,
5134 max: Duration::ZERO,
5135 jitter: false,
5136 ..RetryPolicy::default()
5137 },
5138 ..ResiliencePolicy::default()
5139 };
5140 let res = run_stream(
5141 pages,
5142 &sink,
5143 RunStreamOptions::new()
5144 .with_dlq(DlqConfig::new(Arc::new(NullSink)))
5145 .with_state(store, "k")
5146 .with_resilience(policy),
5147 )
5148 .await
5149 .unwrap();
5150
5151 assert_eq!(
5155 flush_attempts.load(Ordering::SeqCst),
5156 4,
5157 "page-gate flush retried past two transient failures (3) + 1 final flush"
5158 );
5159 assert_eq!(
5161 put_attempts.load(Ordering::SeqCst),
5162 3,
5163 "state put retried past two transient failures"
5164 );
5165 assert_eq!(*stored.lock().unwrap(), Some(json!({"cursor": 42})));
5166 assert_eq!(res.records_written, 1, "the ok row");
5167 }
5168
5169 #[tokio::test]
5170 #[allow(clippy::await_holding_lock)]
5171 async fn resilience_emits_retries_total_with_op_and_class_labels() {
5172 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
5176 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
5177 use metrics_util::debugging::DebugValue;
5178 use std::sync::Arc;
5179 use std::sync::atomic::{AtomicU32, Ordering};
5180 use std::time::Duration;
5181
5182 struct RetryProbeSink {
5183 attempts: Arc<AtomicU32>,
5184 }
5185 #[async_trait]
5186 impl Sink for RetryProbeSink {
5187 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5188 let n = self.attempts.fetch_add(1, Ordering::SeqCst);
5189 if n < 2 {
5190 return Err(FaucetError::HttpStatus {
5191 status: 503,
5192 url: "u".into(),
5193 body: "".into(),
5194 });
5195 }
5196 Ok(r.len())
5197 }
5198 async fn flush(&self) -> Result<(), FaucetError> {
5199 Ok(())
5200 }
5201 fn connector_name(&self) -> &'static str {
5204 "retry-probe"
5205 }
5206 fn supports_idempotent_writes(&self) -> bool {
5208 true
5209 }
5210 }
5211
5212 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
5213 let snap = snapshotter();
5214
5215 let attempts = Arc::new(AtomicU32::new(0));
5216 let sink = RetryProbeSink {
5217 attempts: attempts.clone(),
5218 };
5219 let pages = futures::stream::iter(vec![Ok(StreamPage {
5220 records: vec![json!({"a": 1})],
5221 bookmark: None,
5222 })]);
5223 let policy = ResiliencePolicy {
5224 retry: RetryPolicy {
5225 max_attempts: 5,
5226 backoff: BackoffKind::None,
5227 base: Duration::ZERO,
5228 max: Duration::ZERO,
5229 jitter: false,
5230 ..RetryPolicy::default()
5231 },
5232 ..ResiliencePolicy::default()
5233 };
5234 run_stream(
5235 pages,
5236 &sink,
5237 RunStreamOptions::new()
5238 .with_name("retry-metrics-pipeline")
5239 .with_resilience(policy),
5240 )
5241 .await
5242 .unwrap();
5243 assert_eq!(attempts.load(Ordering::SeqCst), 3);
5244
5245 let snapshot = snap.snapshot();
5246 let retries: u64 = snapshot
5247 .into_vec()
5248 .into_iter()
5249 .filter_map(|(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
5250 if key.key().name() == "faucet_resilience_retries_total"
5251 && key.key().labels().any(|l: &metrics::Label| {
5252 l.key() == "pipeline" && l.value() == "retry-metrics-pipeline"
5253 })
5254 && key
5255 .key()
5256 .labels()
5257 .any(|l: &metrics::Label| l.key() == "op" && l.value() == "sink_write")
5258 && key
5259 .key()
5260 .labels()
5261 .any(|l: &metrics::Label| l.key() == "class" && l.value() == "http_5xx")
5262 && let DebugValue::Counter(c) = v
5263 {
5264 Some(c)
5265 } else {
5266 None
5267 }
5268 })
5269 .sum();
5270 assert_eq!(
5271 retries, 2,
5272 "expected 2 retries (2 transient 503s) counted with op=sink_write, class=http_5xx"
5273 );
5274 }
5275
5276 struct SchemaSink {
5280 schema: Value,
5281 written: std::sync::Mutex<Vec<Value>>,
5282 evolutions: std::sync::Mutex<Vec<crate::drift::SchemaEvolution>>,
5283 evolvable: bool,
5284 }
5285 impl SchemaSink {
5286 fn new(schema: Value, evolvable: bool) -> Self {
5287 Self {
5288 schema,
5289 written: std::sync::Mutex::new(Vec::new()),
5290 evolutions: std::sync::Mutex::new(Vec::new()),
5291 evolvable,
5292 }
5293 }
5294 fn written(&self) -> Vec<Value> {
5295 self.written.lock().unwrap().clone()
5296 }
5297 }
5298 #[async_trait]
5299 impl Sink for SchemaSink {
5300 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
5301 self.written.lock().unwrap().extend(records.iter().cloned());
5302 Ok(records.len())
5303 }
5304 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
5305 Ok(Some(self.schema.clone()))
5306 }
5307 fn supports_schema_evolution(&self) -> bool {
5308 self.evolvable
5309 }
5310 async fn evolve_schema(
5311 &self,
5312 evo: &crate::drift::SchemaEvolution,
5313 ) -> Result<(), FaucetError> {
5314 if !self.evolvable {
5315 return Err(FaucetError::Sink("not evolvable".into()));
5316 }
5317 self.evolutions.lock().unwrap().push(evo.clone());
5318 Ok(())
5319 }
5320 }
5321
5322 fn drift_opts(policy: crate::drift::SchemaDriftPolicy) -> RunStreamOptions {
5323 RunStreamOptions::new().with_schema_drift(policy)
5324 }
5325
5326 fn one_page(
5327 records: Vec<Value>,
5328 ) -> impl futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Unpin {
5329 Box::pin(futures::stream::iter(vec![Ok(StreamPage {
5330 records,
5331 bookmark: None,
5332 })]))
5333 }
5334
5335 #[tokio::test]
5336 async fn drift_warn_writes_unchanged() {
5337 let sink = SchemaSink::new(
5338 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5339 false,
5340 );
5341 let policy = crate::drift::SchemaDriftPolicy {
5342 on_drift: crate::drift::OnDrift::Warn,
5343 allow_widening: true,
5344 on_incompatible: crate::drift::OnIncompatible::Fail,
5345 relax_nullability_on_missing: false,
5346 };
5347 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5348 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5349 assert_eq!(res.records_written, 1);
5350 assert_eq!(sink.written()[0], json!({"id": 1, "email": "a@x.com"}));
5352 }
5353
5354 #[tokio::test]
5355 async fn drift_ignore_strips_unknown_fields() {
5356 let sink = SchemaSink::new(
5357 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5358 false,
5359 );
5360 let policy = crate::drift::SchemaDriftPolicy {
5361 on_drift: crate::drift::OnDrift::Ignore,
5362 allow_widening: true,
5363 on_incompatible: crate::drift::OnIncompatible::Fail,
5364 relax_nullability_on_missing: false,
5365 };
5366 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5367 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5368 assert_eq!(res.records_written, 1);
5369 assert_eq!(
5370 sink.written()[0],
5371 json!({"id": 1}),
5372 "email must be stripped"
5373 );
5374 }
5375
5376 #[tokio::test]
5377 async fn drift_fail_raises_schema_drift() {
5378 let sink = SchemaSink::new(
5379 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5380 false,
5381 );
5382 let policy = crate::drift::SchemaDriftPolicy {
5383 on_drift: crate::drift::OnDrift::Fail,
5384 allow_widening: true,
5385 on_incompatible: crate::drift::OnIncompatible::Fail,
5386 relax_nullability_on_missing: false,
5387 };
5388 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5389 let err = run_stream(pages, &sink, drift_opts(policy))
5390 .await
5391 .unwrap_err();
5392 assert!(matches!(err, FaucetError::SchemaDrift { .. }));
5393 }
5394
5395 #[tokio::test]
5396 async fn drift_evolve_calls_sink_then_writes() {
5397 let sink = SchemaSink::new(
5398 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5399 true,
5400 );
5401 let policy = crate::drift::SchemaDriftPolicy {
5402 on_drift: crate::drift::OnDrift::Evolve,
5403 allow_widening: true,
5404 on_incompatible: crate::drift::OnIncompatible::Fail,
5405 relax_nullability_on_missing: false,
5406 };
5407 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5408 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5409 assert_eq!(res.records_written, 1);
5410 let evos = sink.evolutions.lock().unwrap();
5411 assert_eq!(evos.len(), 1);
5412 assert_eq!(evos[0].additions.len(), 1);
5413 assert_eq!(evos[0].additions[0].name, "email");
5414 assert_eq!(sink.written()[0], json!({"id": 1, "email": "a@x.com"}));
5416 }
5417
5418 #[tokio::test]
5419 async fn drift_evolve_does_not_relax_not_null_for_merely_absent_column() {
5420 let sink = SchemaSink::new(
5425 json!({"type":"object","properties":{
5426 "id":{"type":"integer"},
5427 "legacy":{"type":"string"}
5428 }}),
5429 true,
5430 );
5431 let policy = crate::drift::SchemaDriftPolicy {
5432 on_drift: crate::drift::OnDrift::Evolve,
5433 allow_widening: true,
5434 on_incompatible: crate::drift::OnIncompatible::Fail,
5435 relax_nullability_on_missing: false,
5436 };
5437 let pages = one_page(vec![json!({"id": 1})]);
5438 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5439 assert_eq!(res.records_written, 1);
5440 assert!(
5441 sink.evolutions.lock().unwrap().is_empty(),
5442 "NOT NULL must not be relaxed for a merely-absent column"
5443 );
5444 }
5445
5446 #[tokio::test]
5447 async fn drift_evolve_relaxes_absent_column_only_with_opt_in() {
5448 let sink = SchemaSink::new(
5450 json!({"type":"object","properties":{
5451 "id":{"type":"integer"},
5452 "legacy":{"type":"string"}
5453 }}),
5454 true,
5455 );
5456 let policy = crate::drift::SchemaDriftPolicy {
5457 on_drift: crate::drift::OnDrift::Evolve,
5458 allow_widening: true,
5459 on_incompatible: crate::drift::OnIncompatible::Fail,
5460 relax_nullability_on_missing: true,
5461 };
5462 let pages = one_page(vec![json!({"id": 1})]);
5463 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5464 assert_eq!(res.records_written, 1);
5465 let evos = sink.evolutions.lock().unwrap();
5466 assert_eq!(evos.len(), 1);
5467 assert_eq!(evos[0].relax_nullability, vec!["legacy".to_string()]);
5468 }
5469
5470 #[tokio::test]
5471 async fn drift_inert_when_sink_reports_no_schema() {
5472 let sink = MockSink::new();
5474 let policy = crate::drift::SchemaDriftPolicy {
5475 on_drift: crate::drift::OnDrift::Fail, allow_widening: true,
5477 on_incompatible: crate::drift::OnIncompatible::Fail,
5478 relax_nullability_on_missing: false,
5479 };
5480 let pages = one_page(vec![json!({"id": 1, "anything": true})]);
5481 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5482 assert_eq!(res.records_written, 1);
5483 }
5484
5485 #[tokio::test]
5486 async fn drift_quarantine_routes_drift_rows_to_dlq() {
5487 let sink = SchemaSink::new(
5488 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5489 false,
5490 );
5491 let dlq_sink = std::sync::Arc::new(MockSink::new());
5492 let policy = crate::drift::SchemaDriftPolicy {
5493 on_drift: crate::drift::OnDrift::Quarantine,
5494 allow_widening: true,
5495 on_incompatible: crate::drift::OnIncompatible::Fail,
5496 relax_nullability_on_missing: false,
5497 };
5498 let pages = one_page(vec![
5499 json!({"id": 1}), json!({"id": 2, "email": "x"}), ]);
5502 let opts = RunStreamOptions::new()
5503 .with_schema_drift(policy)
5504 .with_dlq(crate::dlq::DlqConfig::new(dlq_sink.clone()));
5505 let res = run_stream(pages, &sink, opts).await.unwrap();
5506 assert_eq!(res.records_written, 1, "only the conforming row is written");
5507 assert_eq!(sink.written(), vec![json!({"id": 1})]);
5508 let dlq = dlq_sink.written();
5510 assert_eq!(dlq.len(), 1);
5511 assert_eq!(dlq[0]["payload"], json!({"id": 2, "email": "x"}));
5512 assert_eq!(dlq[0]["error"]["kind"], "SchemaDrift");
5513 }
5514
5515 #[test]
5516 fn quarantine_drift_rows_covers_widening_and_droppable_required() {
5517 use crate::drift::{ColumnChange, SchemaDiff};
5518 let diff = SchemaDiff {
5520 additions: vec![],
5521 widenings: vec![ColumnChange {
5522 name: "amount".into(),
5523 from: Some(json!({"type":"integer"})),
5524 to: json!({"type":"number"}),
5525 }],
5526 incompatible: vec![],
5527 droppable_required: vec!["legacy".into()],
5528 };
5529 let records = vec![
5530 json!({"id": 1, "amount": 1.5, "legacy": "x"}), json!({"id": 2, "amount": 7, "legacy": "y"}), json!({"id": 3, "legacy": "z"}), json!({"id": 4}), ];
5535 let page_indices = vec![10, 11, 12, 13];
5538 let (kept, env) = quarantine_drift_rows(&diff, records, &page_indices, "sink", "pl", "");
5539 assert_eq!(kept, vec![json!({"id": 3, "legacy": "z"})]);
5540 assert_eq!(
5541 env.len(),
5542 3,
5543 "widening rows + the missing-required row quarantined"
5544 );
5545 let indices: Vec<i64> = env
5547 .iter()
5548 .map(|e| e["record_index"].as_i64().unwrap())
5549 .collect();
5550 assert_eq!(indices, vec![10, 11, 13]);
5551 }
5552
5553 #[tokio::test]
5554 async fn drift_quarantine_without_dlq_is_rejected() {
5555 let sink = SchemaSink::new(json!({"type":"object","properties":{}}), false);
5556 let policy = crate::drift::SchemaDriftPolicy {
5557 on_drift: crate::drift::OnDrift::Quarantine,
5558 allow_widening: true,
5559 on_incompatible: crate::drift::OnIncompatible::Fail,
5560 relax_nullability_on_missing: false,
5561 };
5562 let pages = one_page(vec![json!({"id": 1})]);
5563 let err = run_stream(pages, &sink, drift_opts(policy))
5564 .await
5565 .unwrap_err();
5566 assert!(matches!(err, FaucetError::Config(_)));
5567 }
5568
5569 #[cfg(feature = "quality")]
5575 #[tokio::test]
5576 async fn drift_fail_with_dlq_still_routes_quality_quarantine() {
5577 use crate::dlq::DlqConfig;
5578 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
5579
5580 let sink = SchemaSink::new(
5583 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5584 false,
5585 );
5586 let dlq_sink = std::sync::Arc::new(MockSink::new());
5587 let policy = crate::drift::SchemaDriftPolicy {
5588 on_drift: crate::drift::OnDrift::Fail,
5589 allow_widening: true,
5590 on_incompatible: crate::drift::OnIncompatible::Fail,
5591 relax_nullability_on_missing: false,
5592 };
5593 let spec = QualitySpec {
5594 record: vec![RecordCheck::NotNull {
5595 field: "name".into(),
5596 treat_missing_as_null: true,
5597 on_failure: OnFailure::Quarantine,
5598 }],
5599 batch: vec![],
5600 };
5601 let quality = std::sync::Arc::new(CompiledQuality::compile(&spec).unwrap());
5602 let pages = one_page(vec![
5603 json!({"id": 1, "name": "ok", "email": "a@x"}), json!({"id": 2, "email": "b@x"}), ]);
5606 let opts = RunStreamOptions::new()
5607 .with_schema_drift(policy)
5608 .with_quality(quality)
5609 .with_dlq(DlqConfig::new(dlq_sink.clone()));
5610 let err = run_stream(pages, &sink, opts).await.unwrap_err();
5611 assert!(
5613 matches!(err, FaucetError::SchemaDrift { .. }),
5614 "got {err:?}"
5615 );
5616 let dlq = dlq_sink.written();
5618 assert_eq!(
5619 dlq.len(),
5620 1,
5621 "quarantined row must reach the DLQ before abort"
5622 );
5623 assert_eq!(dlq[0]["payload"], json!({"id": 2, "email": "b@x"}));
5624 assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
5625 assert_eq!(
5627 sink.written(),
5628 vec![json!({"id": 1, "name": "ok", "email": "a@x"})]
5629 );
5630 }
5631}