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 cleanup: Option<std::sync::Arc<crate::cleanup::CleanupPolicy>>,
144 suppress_overwrite_lifecycle: bool,
152}
153
154impl<'a, So: Source + ?Sized, Si: Sink + ?Sized> Pipeline<'a, So, Si> {
155 pub fn new(source: &'a So, sink: &'a Si) -> Self {
157 Self {
158 source,
159 sink,
160 state_store: None,
161 name: None,
162 row: None,
163 run_id: None,
164 dlq: None,
165 #[cfg(feature = "quality")]
166 quality: None,
167 #[cfg(feature = "contract")]
168 contract: None,
169 #[cfg(feature = "masking")]
170 masking: None,
171 adaptive: None,
172 cancel: None,
173 delivery: crate::idempotency::DeliveryMode::AtLeastOnce,
174 resilience: None,
175 schema_drift: None,
176 cleanup: None,
177 suppress_overwrite_lifecycle: false,
178 }
179 }
180
181 pub fn with_suppress_overwrite(mut self, suppress: bool) -> Self {
190 self.suppress_overwrite_lifecycle = suppress;
191 self
192 }
193
194 pub fn with_state_store(mut self, store: Arc<dyn StateStore>) -> Self {
208 self.state_store = Some(store);
209 self
210 }
211
212 pub fn with_name(mut self, name: impl Into<String>) -> Self {
215 self.name = Some(name.into());
216 self
217 }
218
219 pub fn with_row(mut self, row: impl Into<String>) -> Self {
222 self.row = Some(row.into());
223 self
224 }
225
226 pub fn with_run_id(mut self, run_id: impl Into<String>) -> Self {
230 self.run_id = Some(run_id.into());
231 self
232 }
233
234 pub fn with_dlq(mut self, dlq: DlqConfig) -> Self {
236 self.dlq = Some(dlq);
237 self
238 }
239
240 #[cfg(feature = "quality")]
243 pub fn with_quality(mut self, quality: Arc<crate::quality::CompiledQuality>) -> Self {
244 self.quality = Some(quality);
245 self
246 }
247
248 #[cfg(feature = "contract")]
251 pub fn with_contract(mut self, contract: Arc<crate::contract::CompiledContract>) -> Self {
252 self.contract = Some(contract);
253 self
254 }
255
256 #[cfg(feature = "masking")]
260 pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
261 self.masking = Some(masking);
262 self
263 }
264
265 pub fn with_adaptive(mut self, cfg: crate::adaptive::AdaptiveBatchConfig) -> Self {
269 self.adaptive = Some(cfg);
270 self
271 }
272
273 pub fn with_cancel(mut self, cancel: tokio_util::sync::CancellationToken) -> Self {
278 self.cancel = Some(cancel);
279 self
280 }
281
282 pub fn with_delivery(mut self, mode: crate::idempotency::DeliveryMode) -> Self {
287 self.delivery = mode;
288 self
289 }
290
291 pub fn with_resilience(mut self, policy: crate::resilience::ResiliencePolicy) -> Self {
293 self.resilience = Some(policy);
294 self
295 }
296
297 pub fn with_schema_drift(mut self, policy: crate::drift::SchemaDriftPolicy) -> Self {
300 self.schema_drift = Some(policy);
301 self
302 }
303
304 pub fn with_cleanup(mut self, policy: std::sync::Arc<crate::cleanup::CleanupPolicy>) -> Self {
309 self.cleanup = Some(policy);
310 self
311 }
312
313 pub async fn run(&self) -> Result<PipelineResult, FaucetError> {
327 use crate::observability::{
328 DurationGuard, InstrumentedSink, InstrumentedSource, InstrumentedStateStore, Labels,
329 };
330 use metrics::{Label, SharedString, counter, gauge};
331 use tracing::Instrument;
332
333 let name = self.name.clone().unwrap_or_else(|| "unnamed".to_string());
335 let row = self.row.clone().unwrap_or_default();
336 let run_id = self
337 .run_id
338 .clone()
339 .unwrap_or_else(|| uuid::Uuid::now_v7().to_string());
340 let obs_labels = Labels::new(name.clone(), row.clone(), run_id.clone());
341
342 let wrapped_source = InstrumentedSource::new(self.source, obs_labels.clone());
344 let wrapped_sink = InstrumentedSink::new(self.sink, obs_labels.clone());
345 let wrapped_state_store: Option<Arc<dyn StateStore>> = self.state_store.as_ref().map(|s| {
346 Arc::new(InstrumentedStateStore::new(
347 Arc::clone(s),
348 obs_labels.clone(),
349 )) as Arc<dyn StateStore>
350 });
351
352 let span = tracing::info_span!(
355 "faucet.pipeline.run",
356 pipeline = %name,
357 row = %row,
358 run_id = %run_id,
359 source = %wrapped_source.connector_name(),
360 sink = %wrapped_sink.connector_name(),
361 );
362
363 let base_labels: Vec<Label> = vec![
365 Label::new("pipeline", SharedString::from(name.clone())),
366 Label::new("row", SharedString::from(row.clone())),
367 ];
368 let run_labels: Vec<Label> = {
369 let mut v = base_labels.clone();
370 v.push(Label::new(
371 "source",
372 SharedString::from(wrapped_source.connector_name().to_string()),
373 ));
374 v.push(Label::new(
375 "sink",
376 SharedString::from(wrapped_sink.connector_name().to_string()),
377 ));
378 v
379 };
380
381 struct InFlightGuard(Vec<Label>);
383 impl Drop for InFlightGuard {
384 fn drop(&mut self) {
385 gauge!("faucet_pipeline_in_flight", self.0.clone()).decrement(1.0);
386 }
387 }
388 gauge!("faucet_pipeline_in_flight", base_labels.clone()).increment(1.0);
389 let _in_flight = InFlightGuard(base_labels.clone());
390
391 let start_unix = std::time::SystemTime::now()
394 .duration_since(std::time::UNIX_EPOCH)
395 .map(|d| d.as_secs_f64())
396 .unwrap_or(0.0);
397 gauge!(
398 "faucet_pipeline_start_time_unix_seconds",
399 base_labels.clone()
400 )
401 .set(start_unix);
402
403 let _run_timer =
405 DurationGuard::new("faucet_pipeline_run_duration_seconds", run_labels.clone());
406
407 let result = async {
409 let state_key = self.source.state_key();
412 let mut start_seq = 0u64;
413 if let (Some(store), Some(key)) = (wrapped_state_store.as_ref(), state_key.as_ref()) {
414 validate_state_key(key)?;
415 if let Some(prior) = store.get(key).await? {
416 if self.delivery == crate::idempotency::DeliveryMode::ExactlyOnce {
417 let (bookmark, seq) = crate::idempotency::unwrap_state(&prior);
418 start_seq = seq;
419 if let Some(bm) = bookmark {
420 wrapped_source.apply_start_bookmark(bm).await?;
421 }
422 } else {
423 wrapped_source.apply_start_bookmark(prior).await?;
424 }
425 }
426 }
427
428 if self.delivery == crate::idempotency::DeliveryMode::ExactlyOnce
439 && wrapped_sink.supports_idempotent_writes()
440 && wrapped_source.replay_guarantee()
441 == crate::idempotency::ReplayGuarantee::Deterministic
442 && let Some(key) = state_key.as_ref()
443 && let Some(token) = wrapped_sink.last_committed_token(key).await?
444 && let Some((sink_seq, Some(bm))) = crate::idempotency::parse_token_parts(&token)
445 && sink_seq > start_seq
446 {
447 wrapped_source.apply_start_bookmark(bm).await?;
448 start_seq = sink_seq;
449 }
450
451 #[cfg(feature = "arrow")]
459 {
460 let columnar_ok = wrapped_source.supports_columnar()
461 && wrapped_sink.supports_columnar()
462 && self.dlq.is_none()
463 && self.delivery == crate::idempotency::DeliveryMode::AtLeastOnce
464 && self.schema_drift.is_none()
465 && self.adaptive.is_none()
466 && self.resilience.is_none()
467 && self.cleanup.is_none()
473 && !wrapped_sink.is_overwrite();
477 #[cfg(feature = "quality")]
478 let columnar_ok = columnar_ok && self.quality.is_none();
479 #[cfg(feature = "contract")]
480 let columnar_ok = columnar_ok && self.contract.is_none();
481 #[cfg(feature = "masking")]
482 let columnar_ok = columnar_ok && self.masking.is_none();
483 if columnar_ok {
484 let state = match (wrapped_state_store.clone(), state_key.clone()) {
485 (Some(store), Some(key)) => Some((store, key)),
486 _ => None,
487 };
488 return run_stream_columnar(
489 &wrapped_source,
490 &wrapped_sink,
491 state,
492 self.cancel.clone(),
493 &name,
494 &row,
495 &run_id,
496 )
497 .await;
498 }
499 }
500
501 let ctx = std::collections::HashMap::new();
502 let pages = wrapped_source.stream_pages(&ctx, DEFAULT_BATCH_SIZE);
503
504 let mut opts = RunStreamOptions::new()
505 .with_name(name.clone())
506 .with_row(row.clone())
507 .with_run_id(run_id.clone());
508 if let (Some(store), Some(key)) = (wrapped_state_store.clone(), state_key) {
509 opts = opts.with_state(store, key);
510 }
511 if let Some(dlq) = self.dlq.clone() {
512 opts = opts.with_dlq(dlq);
513 }
514 #[cfg(feature = "quality")]
515 if let Some(q) = self.quality.clone() {
516 opts = opts.with_quality(q);
517 }
518 #[cfg(feature = "contract")]
519 if let Some(c) = self.contract.clone() {
520 opts = opts.with_contract(c);
521 }
522 #[cfg(feature = "masking")]
523 if let Some(m) = self.masking.clone() {
524 opts = opts.with_masking(m);
525 }
526 if let Some(ad) = self.adaptive.clone() {
527 opts = opts.with_adaptive(ad);
528 }
529 if let Some(cancel) = self.cancel.clone() {
530 opts = opts.with_cancel(cancel);
531 }
532 if let Some(policy) = self.resilience.clone() {
533 opts = opts.with_resilience(policy);
534 }
535 if let Some(p) = self.schema_drift {
536 opts = opts.with_schema_drift(p);
537 }
538 opts = opts
539 .with_delivery(self.delivery)
540 .with_start_seq(start_seq)
541 .with_replay_guarantee(wrapped_source.replay_guarantee());
542
543 let overwriting = wrapped_sink.is_overwrite() && !self.suppress_overwrite_lifecycle;
562 if overwriting {
563 wrapped_sink.begin_overwrite().await?;
564 }
565
566 let run_result = match self.cleanup.clone() {
567 None => run_stream(pages, &wrapped_sink, opts).await,
568 Some(policy) => {
569 if !wrapped_sink.supports_cleanup() {
570 return Err(FaucetError::Config(format!(
571 "cleanup is configured but sink '{}' does not support scoped cleanup",
572 wrapped_sink.connector_name()
573 )));
574 }
575 let tracker = crate::cleanup::CleanupTracker::new(&wrapped_sink, &policy);
576 let result = run_stream(pages, &tracker, opts).await;
577 match result {
578 Ok(r) => {
579 let cancelled = self.cancel.as_ref().is_some_and(|c| c.is_cancelled());
580 if cancelled {
581 tracing::warn!(
582 "scoped cleanup skipped: the run was cancelled, so the \
583 set of written keys is incomplete and a delete would \
584 remove rows that were never read"
585 );
586 crate::observability::cleanup_run(&name, &row, "skipped_cancelled");
587 } else {
588 let tracked = tracker.tracked();
589 match tracker.finish(&policy).await {
590 Ok(deleted) => {
591 crate::observability::cleanup_deleted(
592 &name,
593 &row,
594 wrapped_sink.connector_name(),
595 deleted,
596 );
597 crate::observability::cleanup_run(&name, &row, "applied");
598 tracing::info!(
599 deleted,
600 kept = tracked,
601 "scoped cleanup complete: deleted destination rows \
602 in the claimed scope that this run did not write"
603 );
604 }
605 Err(e) => {
606 crate::observability::cleanup_run(
607 &name,
608 &row,
609 "refused_overflow",
610 );
611 return Err(e);
612 }
613 }
614 }
615 Ok(r)
616 }
617 Err(e) => Err(e),
618 }
619 }
620 };
621
622 if overwriting {
626 let cancelled = self.cancel.as_ref().is_some_and(|c| c.is_cancelled());
627 match &run_result {
628 Ok(_) if !cancelled => wrapped_sink.commit_overwrite().await?,
629 _ => {
630 if let Err(e) = wrapped_sink.abort_overwrite().await {
631 tracing::warn!(
632 error = %e,
633 "failed to discard overwrite staging after an unsuccessful \
634 or cancelled run; the destination is unchanged"
635 );
636 }
637 }
638 }
639 }
640 run_result
641 }
642 .instrument(span)
643 .await;
644
645 let status = if result.is_ok() { "ok" } else { "err" };
649 let mut final_labels = run_labels;
650 final_labels.push(Label::new("status", SharedString::const_str(status)));
651 if let Err(ref e) = result {
652 final_labels.push(Label::new(
653 "kind",
654 SharedString::const_str(crate::observability::decorator::error_kind(e)),
655 ));
656 }
657 counter!("faucet_pipeline_runs_total", final_labels).increment(1);
658
659 result
660 }
661}
662
663#[cfg(feature = "arrow")]
674async fn run_stream_columnar<S, Si>(
675 source: &S,
676 sink: &Si,
677 state: Option<(Arc<dyn StateStore>, String)>,
678 cancel: Option<tokio_util::sync::CancellationToken>,
679 pipeline: &str,
680 row: &str,
681 run_id: &str,
682) -> Result<PipelineResult, FaucetError>
683where
684 S: crate::Source + ?Sized,
685 Si: Sink + ?Sized,
686{
687 use futures::StreamExt;
688 use metrics::{Label, SharedString, counter};
689
690 let labels = |connector: &str| -> Vec<Label> {
691 vec![
692 Label::new("pipeline", SharedString::from(pipeline.to_string())),
693 Label::new("row", SharedString::from(row.to_string())),
694 Label::new("connector", SharedString::from(connector.to_string())),
695 ]
696 };
697 let src_labels = labels(source.connector_name());
698 let sink_labels = labels(sink.connector_name());
699 let _ = run_id; let ctx = std::collections::HashMap::new();
702 let mut batches = source.stream_batches(&ctx, DEFAULT_BATCH_SIZE);
703 let mut records_written = 0usize;
704 let mut last_bookmark: Option<Value> = None;
705
706 loop {
707 let page = match &cancel {
710 Some(token) => {
711 tokio::select! {
712 biased;
713 _ = token.cancelled() => break,
714 p = batches.next() => p,
715 }
716 }
717 None => batches.next().await,
718 };
719 let Some(page) = page else { break };
720 let page = page?;
721 let rows = page.num_rows();
722 counter!("faucet_source_records_total", src_labels.clone()).increment(rows as u64);
723
724 if rows > 0 {
725 let n = sink.write_batch_columnar(&page.batch).await?;
726 records_written += n;
727 counter!("faucet_sink_records_total", sink_labels.clone()).increment(n as u64);
728 counter!("faucet_sink_writes_total", sink_labels.clone()).increment(1);
729 }
730
731 if let Some(bm) = page.bookmark {
734 sink.flush().await?;
735 if let Some((store, key)) = state.as_ref() {
736 store.put(key, &bm).await?;
737 }
738 last_bookmark = Some(bm);
739 }
740 }
741
742 sink.flush().await?;
744 Ok(PipelineResult {
745 records_written,
746 bookmark: last_bookmark,
747 dlq: None,
748 })
749}
750
751pub async fn run_stream<S, Si>(
769 mut pages: S,
770 sink: &Si,
771 options: RunStreamOptions,
772) -> Result<PipelineResult, FaucetError>
773where
774 S: Stream<Item = Result<StreamPage, FaucetError>> + Unpin,
775 Si: Sink + ?Sized,
776{
777 use crate::dlq::{DlqReason, DlqStats, OnBatchError, build_envelope};
778
779 let state_store = options.state_store.clone();
780 let state_key = options.state_key.clone();
781 let pipeline_name = options.pipeline_name.unwrap_or_else(|| "unnamed".into());
782 let row = options.row.unwrap_or_default();
783 let run_id = options.run_id.unwrap_or_default();
784 let dlq = options.dlq.clone();
785 let cancel = options.cancel.clone();
786
787 #[cfg(feature = "quality")]
788 let quality = options.quality.clone();
789
790 #[cfg(feature = "quality")]
792 if let Some(q) = quality.as_ref()
793 && q.requires_dlq()
794 && dlq.is_none()
795 {
796 return Err(FaucetError::Config(
797 "quality: on_failure 'quarantine'/'quarantine_batch' requires a DLQ sink".into(),
798 ));
799 }
800
801 #[cfg(feature = "contract")]
802 let contract = options.contract.clone();
803 #[cfg(feature = "contract")]
805 if let Some(c) = contract.as_ref()
806 && c.requires_dlq()
807 && dlq.is_none()
808 {
809 return Err(FaucetError::Config(
810 "contract: on_breach 'quarantine' requires a DLQ sink".into(),
811 ));
812 }
813 #[cfg(feature = "contract")]
815 let mut warned_contract_breach = false;
816
817 #[cfg(feature = "masking")]
821 let masking = options.masking.clone();
822
823 let schema_drift = options.schema_drift;
825 if let Some(p) = schema_drift.as_ref()
827 && p.requires_dlq()
828 && dlq.is_none()
829 {
830 return Err(FaucetError::Config(
831 "schema: on_drift 'quarantine' (or on_incompatible 'quarantine') requires a DLQ sink"
832 .into(),
833 ));
834 }
835 let mut dest_schema_cache: Option<Option<Value>> = None;
839 let mut warned_drift_inert = false;
840
841 if let Some(key) = state_key.as_ref() {
842 validate_state_key(key)?;
843 }
844
845 let mechanism: Option<crate::idempotency::EffectivelyOnceMechanism> =
853 if options.delivery == crate::idempotency::DeliveryMode::ExactlyOnce {
854 let replay_ok = options
855 .replay
856 .is_none_or(|r| r == crate::idempotency::ReplayGuarantee::Deterministic);
857 if sink.supports_idempotent_writes() && replay_ok {
858 if state_store.is_none() || state_key.is_none() {
859 return Err(FaucetError::Config(
860 "delivery: exactly_once (atomic watermark) requires a state store".into(),
861 ));
862 }
863 if dlq.is_some() {
864 return Err(FaucetError::Config(
865 "delivery: exactly_once (atomic watermark) is not compatible with a DLQ \
866 in this version"
867 .into(),
868 ));
869 }
870 Some(crate::idempotency::EffectivelyOnceMechanism::AtomicWatermark)
871 } else if sink.dedups_by_key() {
872 Some(crate::idempotency::EffectivelyOnceMechanism::KeyedUpsert)
873 } else if sink.supports_idempotent_writes() {
874 return Err(FaucetError::Config(format!(
877 "delivery: exactly_once — the source does not replay deterministically from \
878 a bookmark, so the atomic-watermark mechanism cannot be used; configure \
879 `write_mode: upsert` with a `key` on sink '{}' for keyed-upsert \
880 effectively-once instead",
881 sink.connector_name()
882 )));
883 } else {
884 return Err(FaucetError::Config(format!(
885 "delivery: exactly_once requires an idempotent (atomic-watermark) sink or a \
886 sink configured to dedup by key (`write_mode: upsert` + `key`), but '{}' \
887 provides neither",
888 sink.connector_name()
889 )));
890 }
891 } else {
892 None
893 };
894 let exactly_once =
898 mechanism == Some(crate::idempotency::EffectivelyOnceMechanism::AtomicWatermark);
899 let scope = state_key.clone().unwrap_or_default();
900 let mut next_seq = options.start_seq;
901 let committed_seq = if exactly_once {
902 sink.last_committed_token(&scope)
903 .await?
904 .and_then(|t| crate::idempotency::parse_token(&t))
905 .unwrap_or(0)
906 } else {
907 0
908 };
909
910 let mut records_written = 0usize;
911 let mut last_bookmark: Option<Value> = None;
912 let mut dlq_stats = DlqStats::default();
913
914 let adaptive_cfg = options.adaptive.clone().filter(|c| c.enabled);
915 if let Some(cfg) = adaptive_cfg.as_ref() {
919 cfg.validate()?;
920 }
921 let mut controller: Option<crate::adaptive::AimdController> = None;
922 let mut warned_noop_sink = false;
923 let mut warned_poison_drop = false;
925
926 let sink_name = sink.connector_name();
927 let dlq_sink_name = dlq.as_ref().map(|d| d.sink.connector_name()).unwrap_or("");
928
929 let mut cancelled = false;
940
941 let resilience = options.resilience.clone();
948 let retry_policy = resilience.as_ref().map(|r| r.retry.clone());
949 let mut breaker = resilience
950 .as_ref()
951 .and_then(|r| r.circuit_breaker)
952 .map(|cb| {
953 (
954 crate::resilience::CircuitBreaker::new(cb.consecutive_failures),
955 cb.cooldown,
956 )
957 });
958 let poison = resilience.as_ref().and_then(|r| r.poison);
960
961 macro_rules! with_retry {
974 ($op_label:literal, $op:expr) => {
975 match &retry_policy {
976 Some(p) => {
977 let m = crate::resilience::RetryMetrics {
978 pipeline: pipeline_name.to_string(),
979 row: row.to_string(),
980 op: $op_label,
981 };
982 crate::resilience::execute_with_policy_metered(p, cancel.as_ref(), &m, || $op)
983 .await
984 }
985 None => $op.await,
986 }
987 };
988 }
989
990 macro_rules! with_retry_write {
1001 ($op_label:literal, $op:expr) => {
1002 if retry_policy.is_some() && sink.supports_idempotent_writes() {
1003 with_retry!($op_label, $op)
1004 } else {
1005 $op.await
1006 }
1007 };
1008 }
1009
1010 let loop_result: Result<(), FaucetError> = async {
1011 loop {
1012 let page = match &cancel {
1016 Some(token) => tokio::select! {
1017 biased;
1018 _ = token.cancelled() => {
1019 cancelled = true;
1020 break;
1021 }
1022 p = std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)) => p,
1023 },
1024 None => std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)).await,
1025 };
1026 match page {
1027 Some(Ok(page)) => {
1028 if page.records.is_empty() && page.bookmark.is_none() {
1029 continue;
1030 }
1031
1032 #[cfg(feature = "masking")]
1038 let page = if let Some(m) = masking.as_ref() {
1039 let labels =
1040 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1041 let outcome = crate::observability::instrumented_apply_masking(
1042 page.records,
1043 m,
1044 &labels,
1045 );
1046 StreamPage {
1047 records: outcome.records,
1048 bookmark: page.bookmark,
1049 }
1050 } else {
1051 page
1052 };
1053
1054 let page_len = page.records.len();
1061
1062 #[cfg(feature = "quality")]
1064 let (records, quality_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
1065 if let Some(q) = quality.as_ref() {
1066 let labels =
1067 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1068 let outcome = crate::observability::instrumented_apply_quality(
1069 page.records,
1070 q,
1071 &labels,
1072 )?;
1073 let quarantined_idx: std::collections::HashSet<usize> =
1074 outcome.quarantined.iter().map(|qr| qr.page_index).collect();
1075 let envelopes: Vec<Value> = outcome
1076 .quarantined
1077 .iter()
1078 .map(|qr| {
1079 let err = FaucetError::QualityFailure {
1080 check: qr.check.to_string(),
1081 message: qr.message.clone(),
1082 };
1083 build_envelope(
1087 &qr.record,
1088 &err,
1089 DlqReason::Quality,
1090 sink_name,
1091 &pipeline_name,
1092 &row,
1093 qr.page_index,
1094 )
1095 })
1096 .collect();
1097 let survivor_idx: Vec<usize> =
1098 (0..page_len).filter(|i| !quarantined_idx.contains(i)).collect();
1099 (outcome.survivors, envelopes, survivor_idx)
1100 } else {
1101 (page.records, Vec::new(), (0..page_len).collect())
1102 };
1103 #[cfg(not(feature = "quality"))]
1104 let (records, quality_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
1105 (page.records, Vec::new(), (0..page_len).collect());
1106
1107 #[cfg(feature = "contract")]
1114 let (records, contract_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
1115 if let Some(c) = contract.as_ref() {
1116 let labels =
1117 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1118 let outcome = crate::observability::instrumented_apply_contract(
1119 records, c, &labels,
1120 )?;
1121 if !outcome.warned.is_empty() && !warned_contract_breach {
1122 tracing::warn!(
1123 version = %c.version,
1124 breaches = outcome.warned.len(),
1125 first = %outcome.warned[0].describe(),
1126 "contract: breaching records written unchanged \
1127 (on_breach=warn); this warning fires once per run"
1128 );
1129 warned_contract_breach = true;
1130 }
1131 let envelopes: Vec<Value> = outcome
1132 .quarantined
1133 .iter()
1134 .map(|vr| {
1135 let err = FaucetError::ContractViolation {
1136 version: c.version.clone(),
1137 message: vr.violation.describe(),
1138 };
1139 let page_index = page_indices
1149 .get(vr.violation.page_index)
1150 .copied()
1151 .unwrap_or(vr.violation.page_index);
1152 build_envelope(
1153 &vr.record,
1154 &err,
1155 DlqReason::Contract,
1156 sink_name,
1157 &pipeline_name,
1158 &row,
1159 page_index,
1160 )
1161 })
1162 .collect();
1163 let contract_quarantined: std::collections::HashSet<usize> = outcome
1169 .quarantined
1170 .iter()
1171 .map(|vr| vr.violation.page_index)
1172 .collect();
1173 let survivor_idx: Vec<usize> = page_indices
1174 .iter()
1175 .enumerate()
1176 .filter(|(pos, _)| !contract_quarantined.contains(pos))
1177 .map(|(_, orig)| *orig)
1178 .collect();
1179 (outcome.survivors, envelopes, survivor_idx)
1180 } else {
1181 (records, Vec::new(), page_indices)
1182 };
1183
1184 let mut drift_envelopes: Vec<Value> = Vec::new();
1186 let (records, drift_abort): (Vec<Value>, Option<FaucetError>) =
1187 if let Some(policy) = schema_drift.as_ref().filter(|_| !records.is_empty()) {
1188 if dest_schema_cache.is_none() {
1190 dest_schema_cache = Some(sink.current_schema().await?);
1191 }
1192 let dest = dest_schema_cache.as_ref().and_then(|o| o.as_ref());
1193 match dest {
1194 None => {
1195 if !warned_drift_inert {
1196 tracing::info!(
1197 connector = sink_name,
1198 "schema-drift: sink reports no destination schema; \
1199 drift handling is inert this run"
1200 );
1201 warned_drift_inert = true;
1202 }
1203 (records, None)
1204 }
1205 Some(dest) => {
1206 let inferred = crate::schema::infer_schema(&records);
1207 let diff = crate::drift::diff_schema(
1208 dest,
1209 &inferred,
1210 policy.allow_widening,
1211 );
1212 if diff.is_empty() {
1213 (records, None)
1214 } else {
1215 let dest_owned = dest.clone();
1219 apply_drift_policy(
1220 policy,
1221 &diff,
1222 &dest_owned,
1223 records,
1224 &page_indices,
1225 sink,
1226 sink_name,
1227 &pipeline_name,
1228 &row,
1229 &mut dest_schema_cache,
1230 &mut drift_envelopes,
1231 )
1232 .await?
1233 }
1234 }
1235 }
1236 } else {
1237 (records, None)
1238 };
1239 let quality_envelopes = {
1243 let mut q = quality_envelopes;
1244 #[cfg(feature = "contract")]
1245 q.extend(contract_envelopes);
1246 q.append(&mut drift_envelopes);
1247 q
1248 };
1249 let mut drift_abort = drift_abort;
1259 if dlq.is_none()
1260 && let Some(e) = drift_abort.take()
1261 {
1262 return Err(e);
1263 }
1264
1265 let page = StreamPage {
1266 records,
1267 bookmark: page.bookmark,
1268 };
1269
1270 if let Some(ref dlq_cfg) = dlq {
1271 use metrics::{Label, SharedString, counter};
1273 let metric_labels: Vec<Label> = vec![
1274 Label::new("pipeline", SharedString::from(pipeline_name.clone())),
1275 Label::new("row", SharedString::from(row.clone())),
1276 Label::new("connector", SharedString::from(sink_name.to_string())),
1277 Label::new(
1278 "dlq_connector",
1279 SharedString::from(dlq_sink_name.to_string()),
1280 ),
1281 ];
1282 let span = tracing::info_span!(
1283 "faucet.dlq.route",
1284 pipeline = %pipeline_name,
1285 row = %row,
1286 run_id = %run_id,
1287 connector = %sink_name,
1288 dlq_connector = %dlq_sink_name,
1289 );
1290 let _enter = span.enter();
1291
1292 let mut envelopes: Vec<Value> = Vec::new();
1296 let mut page_success = 0usize;
1297 let mut outer_err_recovered = false;
1298 let mut had_per_row_sink_failure = false;
1303 let records_len = page.records.len();
1304 let mut offset = 0usize;
1305 while offset < records_len {
1306 let size = match adaptive_cfg.as_ref() {
1307 Some(cfg) => {
1308 let ctrl = controller.get_or_insert_with(|| {
1309 crate::adaptive::AimdController::new(cfg, records_len)
1310 });
1311 ctrl.current().max(1).min(records_len - offset)
1312 }
1313 None => records_len - offset, };
1315 if adaptive_cfg.is_some() {
1316 maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
1317 }
1318 let chunk = &page.records[offset..offset + size];
1319 let t0 = std::time::Instant::now();
1320 let chunk_outcomes_result =
1326 with_retry_write!("sink_write", sink.write_batch_partial(chunk));
1327 let latency = t0.elapsed();
1328 let (mut chunk_outcomes, chunk_synthesized): (
1337 Vec<crate::RowOutcome>,
1338 bool,
1339 ) = match chunk_outcomes_result {
1340 Ok(o) => (o, false),
1341 Err(e) => match dlq_cfg.on_batch_error {
1342 OnBatchError::Propagate => return Err(e),
1343 OnBatchError::DlqAll => {
1344 outer_err_recovered = true;
1345 let msg = e.to_string();
1346 let synth = (0..chunk.len())
1347 .map(|_| Err(FaucetError::Sink(msg.clone())))
1348 .collect();
1349 (synth, true)
1350 }
1351 },
1352 };
1353
1354 if let Some(pp) = poison
1363 && !chunk_synthesized
1364 {
1365 let mut attempt = 1u32; while attempt < pp.max_row_attempts {
1367 let failing: Vec<usize> = chunk_outcomes
1368 .iter()
1369 .enumerate()
1370 .filter_map(|(j, o)| match o {
1371 Err(e)
1372 if retry_policy
1373 .as_ref()
1374 .map(|p| p.is_retriable(e))
1375 .unwrap_or(false) =>
1376 {
1377 Some(j)
1378 }
1379 _ => None,
1380 })
1381 .collect();
1382 if failing.is_empty() {
1383 break;
1384 }
1385 let subset: Vec<Value> =
1386 failing.iter().map(|&j| chunk[j].clone()).collect();
1387 let retried = sink.write_batch_partial(&subset).await?;
1395 let mut retried = retried.into_iter();
1399 for &j in failing.iter() {
1400 chunk_outcomes[j] = retried.next().unwrap_or(Ok(()));
1401 }
1402 attempt += 1;
1403 }
1404 }
1405
1406 let mut chunk_errors = 0usize;
1407 let mut poison_dlq = 0u64;
1413 let mut poison_drop = 0u64;
1414 for (j, outcome) in chunk_outcomes.iter().enumerate() {
1415 match outcome {
1416 Ok(()) => page_success += 1,
1417 Err(err) => {
1418 let action = poison
1423 .map(|pp| pp.action)
1424 .unwrap_or(crate::resilience::PoisonAction::Dlq);
1425 match action {
1426 crate::resilience::PoisonAction::Fail => {
1427 crate::observability::resilience::poison_rows(
1428 &pipeline_name,
1429 &row,
1430 "fail",
1431 1,
1432 );
1433 return Err(FaucetError::Sink(format!(
1434 "poison-pill row failed permanently: {err}"
1435 )));
1436 }
1437 crate::resilience::PoisonAction::Drop => {
1438 poison_drop += 1;
1441 if !warned_poison_drop {
1442 tracing::warn!(
1443 "poison-pill: dropping permanently-failing row(s) (action=drop); this warning fires once per run"
1444 );
1445 warned_poison_drop = true;
1446 }
1447 }
1448 crate::resilience::PoisonAction::Dlq => {
1449 poison_dlq += 1;
1450 chunk_errors += 1;
1451 if !chunk_synthesized {
1452 had_per_row_sink_failure = true;
1453 }
1454 envelopes.push(build_envelope(
1455 &chunk[j],
1456 err,
1457 DlqReason::Partial,
1458 sink_name,
1459 &pipeline_name,
1460 &row,
1461 offset + j,
1462 ));
1463 }
1464 }
1465 }
1466 }
1467 }
1468 if poison.is_some() {
1472 crate::observability::resilience::poison_rows(
1473 &pipeline_name,
1474 &row,
1475 "dlq",
1476 poison_dlq,
1477 );
1478 crate::observability::resilience::poison_rows(
1479 &pipeline_name,
1480 &row,
1481 "drop",
1482 poison_drop,
1483 );
1484 }
1485 if let Some(ctrl) = controller.as_mut() {
1486 let adj = ctrl.observe(crate::adaptive::Observation {
1487 batch_len: chunk.len(),
1488 errors: chunk_errors,
1489 latency,
1490 });
1491 emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
1492 }
1493 offset += size;
1494 }
1495 #[cfg(feature = "quality")]
1500 let quality_count = quality_envelopes.len();
1501 #[cfg(not(feature = "quality"))]
1502 let quality_count = 0usize;
1503 envelopes.splice(0..0, quality_envelopes);
1504 let page_failures = envelopes.len();
1505
1506 let mut budget_error: Option<FaucetError> = None;
1519 let mut circuit_error: Option<FaucetError> = None;
1528 if let Some((b, cooldown)) = breaker.as_mut() {
1529 if records_len > 0 && page_success == 0 {
1530 if b.record_failure() {
1531 crate::observability::resilience::circuit_opened(
1532 &pipeline_name,
1533 &row,
1534 );
1535 circuit_error = Some(FaucetError::CircuitOpen {
1536 failures: b.consecutive(),
1537 cooldown: *cooldown,
1538 });
1539 }
1540 } else if page_success > 0 {
1541 b.record_success();
1542 }
1543 }
1544 if let Some(limit) = dlq_cfg.max_failures_per_page
1545 && page_failures > limit
1546 {
1547 let mut lbl = metric_labels.clone();
1548 lbl.retain(|l| l.key() != "dlq_connector");
1549 lbl.push(Label::new("scope", SharedString::const_str("per_page")));
1550 counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
1551 budget_error = Some(FaucetError::Sink(format!(
1552 "DLQ per-page budget exceeded: {page_failures} > {limit}"
1553 )));
1554 }
1555 let new_total = dlq_stats.records_dlq + page_failures;
1556 if budget_error.is_none()
1557 && let Some(limit) = dlq_cfg.max_failures_total
1558 && new_total > limit
1559 {
1560 let mut lbl = metric_labels.clone();
1561 lbl.retain(|l| l.key() != "dlq_connector");
1562 lbl.push(Label::new("scope", SharedString::const_str("total")));
1563 counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
1564 budget_error = Some(FaucetError::Sink(format!(
1565 "DLQ total budget exceeded: {new_total} > {limit}"
1566 )));
1567 }
1568
1569 if !envelopes.is_empty() {
1571 let _dlq_write_timer = crate::observability::DurationGuard::new(
1572 "faucet_sink_dlq_write_duration_seconds",
1573 metric_labels.clone(),
1574 );
1575 dlq_cfg.sink.write_batch(&envelopes).await.map_err(|e| {
1576 let mut lbl = metric_labels.clone();
1577 lbl.push(Label::new(
1578 "kind",
1579 SharedString::const_str(
1580 crate::observability::decorator::error_kind(&e),
1581 ),
1582 ));
1583 counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1584 FaucetError::Sink(format!("DLQ sink write failed: {e}"))
1585 })?;
1586 dlq_stats.records_dlq += page_failures;
1587 dlq_stats.pages_with_failures += 1;
1588
1589 let reason_label = if had_per_row_sink_failure {
1608 DlqReason::Partial.as_str()
1609 } else if outer_err_recovered {
1610 DlqReason::DlqAll.as_str()
1611 } else if page_failures > quality_count {
1612 DlqReason::Partial.as_str()
1613 } else {
1614 DlqReason::Quality.as_str()
1615 };
1616 counter!("faucet_sink_dlq_records_total", metric_labels.clone())
1617 .increment(page_failures as u64);
1618 let mut page_labels = metric_labels.clone();
1619 page_labels
1620 .push(Label::new("reason", SharedString::const_str(reason_label)));
1621 counter!("faucet_sink_dlq_pages_total", page_labels).increment(1);
1622 }
1623
1624 records_written += page_success;
1625
1626 if let Some(bookmark) = page.bookmark {
1627 with_retry!("flush", sink.flush())?;
1633 let _dlq_flush_timer = crate::observability::DurationGuard::new(
1634 "faucet_sink_dlq_flush_duration_seconds",
1635 metric_labels.clone(),
1636 );
1637 with_retry!("flush", dlq_cfg.sink.flush()).map_err(|e| {
1638 let mut lbl = metric_labels.clone();
1639 lbl.push(Label::new(
1640 "kind",
1641 SharedString::const_str(
1642 crate::observability::decorator::error_kind(&e),
1643 ),
1644 ));
1645 counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1646 FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
1647 })?;
1648 let bm_labels =
1649 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1650 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1651 if let (Some(store), Some(key)) =
1652 (state_store.as_ref(), state_key.as_ref())
1653 {
1654 with_retry!("state_put", store.put(key, &bookmark))?;
1655 }
1656 last_bookmark = Some(bookmark);
1657 }
1658
1659 if let Some(e) = budget_error {
1666 return Err(e);
1667 }
1668 if let Some(e) = circuit_error {
1670 return Err(e);
1671 }
1672 if let Some(e) = drift_abort {
1677 return Err(e);
1678 }
1679 } else if exactly_once {
1680 if let Some(bookmark) = page.bookmark {
1684 next_seq += 1;
1685 let token = crate::idempotency::format_token_with_bookmark(
1692 next_seq,
1693 Some(&bookmark),
1694 );
1695 if next_seq <= committed_seq {
1696 use metrics::{Label, SharedString, counter};
1700 let skip_labels: Vec<Label> = vec![
1701 Label::new(
1702 "pipeline",
1703 SharedString::from(pipeline_name.clone()),
1704 ),
1705 Label::new("row", SharedString::from(row.clone())),
1706 ];
1707 counter!("faucet_pipeline_pages_skipped_total", skip_labels)
1708 .increment(1);
1709 } else {
1710 records_written += with_retry!(
1711 "sink_write",
1712 sink.write_batch_idempotent(&page.records, &scope, &token)
1713 )?;
1714 }
1715 with_retry!("flush", sink.flush())?;
1716 let bm_labels =
1717 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1718 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1719 if let (Some(store), Some(key)) =
1720 (state_store.as_ref(), state_key.as_ref())
1721 {
1722 let wrapped =
1723 crate::idempotency::wrap_state(Some(&bookmark), next_seq);
1724 with_retry!("state_put", store.put(key, &wrapped))?;
1725 }
1726 last_bookmark = Some(bookmark);
1727 } else if !page.records.is_empty() {
1728 records_written +=
1732 with_retry_write!("sink_write", sink.write_batch(&page.records))?;
1733 }
1734 } else {
1735 debug_assert!(
1737 quality_envelopes.is_empty(),
1738 "quality quarantine without DLQ should have been rejected at run start"
1739 );
1740 if !page.records.is_empty() {
1741 if let Some(cfg) = adaptive_cfg.as_ref() {
1742 let ctrl = controller.get_or_insert_with(|| {
1743 crate::adaptive::AimdController::new(cfg, page.records.len())
1744 });
1745 maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
1746 let mut offset = 0;
1747 while offset < page.records.len() {
1748 let size =
1749 ctrl.current().max(1).min(page.records.len() - offset);
1750 let chunk = &page.records[offset..offset + size];
1751 let t0 = std::time::Instant::now();
1752 let n = with_retry_write!("sink_write", sink.write_batch(chunk))?;
1753 let latency = t0.elapsed();
1754 records_written += n;
1755 offset += size;
1756 let adj = ctrl.observe(crate::adaptive::Observation {
1757 batch_len: chunk.len(),
1758 errors: 0,
1759 latency,
1760 });
1761 emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
1762 }
1763 } else {
1764 records_written +=
1765 with_retry_write!("sink_write", sink.write_batch(&page.records))?;
1766 }
1767 }
1768 if let Some(bookmark) = page.bookmark {
1769 with_retry!("flush", sink.flush())?;
1770 let bm_labels =
1771 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1772 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1773 if let (Some(store), Some(key)) =
1774 (state_store.as_ref(), state_key.as_ref())
1775 {
1776 with_retry!("state_put", store.put(key, &bookmark))?;
1777 }
1778 last_bookmark = Some(bookmark);
1779 }
1780 }
1781 }
1782 Some(Err(e)) => return Err(e),
1783 None => break,
1784 }
1785 }
1786 Ok(())
1787 }
1788 .await;
1789
1790 if let Err(e) = loop_result {
1797 if let Some(ref dlq_cfg) = dlq
1798 && let Err(flush_err) = dlq_cfg.sink.flush().await
1799 {
1800 tracing::warn!(
1801 error = %flush_err,
1802 "DLQ sink flush failed during error unwind; original error preserved"
1803 );
1804 }
1805 if let Err(flush_err) = sink.flush().await {
1806 tracing::warn!(
1807 error = %flush_err,
1808 "sink flush failed during error unwind; original error preserved"
1809 );
1810 }
1811 return Err(e);
1812 }
1813
1814 if let Some(ref dlq_cfg) = dlq {
1820 let final_metric_labels: Vec<metrics::Label> = vec![
1821 metrics::Label::new(
1822 "pipeline",
1823 metrics::SharedString::from(pipeline_name.clone()),
1824 ),
1825 metrics::Label::new("row", metrics::SharedString::from(row.clone())),
1826 metrics::Label::new(
1827 "connector",
1828 metrics::SharedString::from(sink_name.to_string()),
1829 ),
1830 metrics::Label::new(
1831 "dlq_connector",
1832 metrics::SharedString::from(dlq_sink_name.to_string()),
1833 ),
1834 ];
1835 let _final_dlq_flush_timer = crate::observability::DurationGuard::new(
1836 "faucet_sink_dlq_flush_duration_seconds",
1837 final_metric_labels.clone(),
1838 );
1839 dlq_cfg.sink.flush().await.map_err(|e| {
1840 let mut lbl = final_metric_labels.clone();
1841 lbl.push(metrics::Label::new(
1842 "kind",
1843 metrics::SharedString::const_str(crate::observability::decorator::error_kind(&e)),
1844 ));
1845 metrics::counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1846 FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
1847 })?;
1848 }
1849 sink.flush().await?;
1850
1851 if cancelled {
1852 tracing::info!(
1853 records_written,
1854 "pipeline run cancelled cooperatively; sink flushed (partial output is durable)"
1855 );
1856 }
1857
1858 tracing::info!(
1859 records_written,
1860 cancelled,
1861 has_bookmark = last_bookmark.is_some(),
1862 persisted = state_store.is_some() && state_key.is_some() && last_bookmark.is_some(),
1863 dlq_records = dlq_stats.records_dlq,
1864 "pipeline streaming run complete"
1865 );
1866
1867 Ok(PipelineResult {
1868 records_written,
1869 bookmark: last_bookmark,
1870 dlq: dlq.is_some().then_some(dlq_stats),
1871 })
1872}
1873
1874fn emit_adaptive_metrics(
1877 ctrl: &crate::adaptive::AimdController,
1878 adj: Option<crate::adaptive::Adjustment>,
1879 pipeline: &str,
1880 row: &str,
1881) {
1882 use metrics::{Label, SharedString, counter, gauge};
1883 let base = vec![
1884 Label::new("pipeline", SharedString::from(pipeline.to_string())),
1885 Label::new("row", SharedString::from(row.to_string())),
1886 ];
1887 gauge!("faucet_pipeline_adaptive_batch_size", base.clone()).set(ctrl.current() as f64);
1888 gauge!(
1889 "faucet_pipeline_adaptive_batch_cooldown_active",
1890 base.clone()
1891 )
1892 .set(if ctrl.cooldown_active() { 1.0 } else { 0.0 });
1893 if let Some(p50) = ctrl.p50_latency_ms() {
1894 gauge!(
1895 "faucet_pipeline_adaptive_batch_p50_latency_ms",
1896 base.clone()
1897 )
1898 .set(p50 as f64);
1899 }
1900 if let Some(a) = adj {
1901 let mut lbl = base;
1902 lbl.push(Label::new(
1903 "direction",
1904 SharedString::const_str(a.direction.as_str()),
1905 ));
1906 lbl.push(Label::new(
1907 "reason",
1908 SharedString::const_str(a.reason.as_str()),
1909 ));
1910 counter!("faucet_pipeline_adaptive_batch_adjustments_total", lbl).increment(1);
1911 }
1912}
1913
1914fn maybe_warn_noop_sink(sink_name: &str, warned: &mut bool) {
1917 if !*warned && matches!(sink_name, "jsonl" | "csv" | "stdout") {
1918 tracing::info!(
1919 sink = sink_name,
1920 "adaptive batch sizing is a no-op for this per-record sink"
1921 );
1922 *warned = true;
1923 }
1924}
1925
1926#[allow(clippy::too_many_arguments)]
1934async fn apply_drift_policy<Si: Sink + ?Sized>(
1935 policy: &crate::drift::SchemaDriftPolicy,
1936 diff: &crate::drift::SchemaDiff,
1937 dest: &Value,
1938 records: Vec<Value>,
1939 page_indices: &[usize],
1940 sink: &Si,
1941 sink_name: &str,
1942 pipeline_name: &str,
1943 row: &str,
1944 dest_schema_cache: &mut Option<Option<Value>>,
1945 drift_envelopes: &mut Vec<Value>,
1946) -> Result<(Vec<Value>, Option<FaucetError>), FaucetError> {
1947 use crate::drift::{OnDrift, OnIncompatible};
1948 use crate::observability::schema_drift as emit_drift;
1949
1950 let mode = match policy.on_drift {
1951 OnDrift::Warn => "warn",
1952 OnDrift::Ignore => "ignore",
1953 OnDrift::Quarantine => "quarantine",
1954 OnDrift::Fail => "fail",
1955 OnDrift::Evolve => "evolve",
1956 };
1957 emit_drift(
1958 pipeline_name,
1959 row,
1960 sink_name,
1961 mode,
1962 "added",
1963 diff.additions.len() as u64,
1964 );
1965 emit_drift(
1966 pipeline_name,
1967 row,
1968 sink_name,
1969 mode,
1970 "widened",
1971 diff.widenings.len() as u64,
1972 );
1973 emit_drift(
1974 pipeline_name,
1975 row,
1976 sink_name,
1977 mode,
1978 "narrowed",
1979 diff.incompatible.len() as u64,
1980 );
1981 emit_drift(
1982 pipeline_name,
1983 row,
1984 sink_name,
1985 mode,
1986 "dropped",
1987 diff.droppable_required.len() as u64,
1988 );
1989
1990 match policy.on_drift {
1991 OnDrift::Warn => {
1992 tracing::warn!(
1993 connector = sink_name,
1994 columns = ?diff.changed_columns(),
1995 "schema-drift detected (on_drift=warn); writing page unchanged"
1996 );
1997 Ok((records, None))
1998 }
1999 OnDrift::Fail => Ok((
2000 records,
2001 Some(FaucetError::SchemaDrift {
2002 columns: diff.changed_columns(),
2003 message: "schema drift detected (on_drift=fail)".to_string(),
2004 }),
2005 )),
2006 OnDrift::Ignore => {
2007 let allowed: std::collections::HashSet<String> = dest
2009 .get("properties")
2010 .and_then(|p| p.as_object())
2011 .map(|m| m.keys().cloned().collect())
2012 .unwrap_or_default();
2013 let trimmed = records
2014 .into_iter()
2015 .map(|r| match r {
2016 Value::Object(map) => Value::Object(
2017 map.into_iter()
2018 .filter(|(k, _)| allowed.contains(k))
2019 .collect(),
2020 ),
2021 other => other,
2022 })
2023 .collect();
2024 Ok((trimmed, None))
2025 }
2026 OnDrift::Quarantine => {
2027 let (kept, env) =
2028 quarantine_drift_rows(diff, records, page_indices, sink_name, pipeline_name, row);
2029 drift_envelopes.extend(env);
2030 Ok((kept, None))
2031 }
2032 OnDrift::Evolve => {
2033 let evolution = crate::drift::SchemaEvolution {
2034 additions: diff.additions.clone(),
2035 widenings: diff
2036 .widenings
2037 .iter()
2038 .filter(|c| {
2039 c.from
2040 .as_ref()
2041 .map(|f| crate::drift::base_widened(f, &c.to))
2042 .unwrap_or(false)
2043 })
2044 .cloned()
2045 .collect(),
2046 relax_nullability: diff
2047 .droppable_required
2048 .iter()
2049 .filter(|_| policy.relax_nullability_on_missing)
2054 .cloned()
2055 .chain(
2056 diff.widenings
2057 .iter()
2058 .filter(|c| {
2059 c.from
2060 .as_ref()
2061 .map(|f| crate::drift::adds_null(f, &c.to))
2062 .unwrap_or(false)
2063 })
2064 .map(|c| c.name.clone()),
2065 )
2066 .collect(),
2067 };
2068 if !evolution.is_empty() {
2069 sink.evolve_schema(&evolution).await?;
2070 *dest_schema_cache = Some(sink.current_schema().await?);
2073 }
2074 if diff.incompatible.is_empty() {
2076 Ok((records, None))
2077 } else {
2078 match policy.on_incompatible {
2079 OnIncompatible::Fail => Ok((
2080 records,
2081 Some(FaucetError::SchemaDrift {
2082 columns: diff.incompatible.iter().map(|c| c.name.clone()).collect(),
2083 message: "incompatible type change cannot be auto-evolved \
2084 (on_incompatible=fail)"
2085 .into(),
2086 }),
2087 )),
2088 OnIncompatible::Quarantine => {
2089 let incompat_only = crate::drift::SchemaDiff {
2091 incompatible: diff.incompatible.clone(),
2092 ..Default::default()
2093 };
2094 let (kept, env) = quarantine_drift_rows(
2095 &incompat_only,
2096 records,
2097 page_indices,
2098 sink_name,
2099 pipeline_name,
2100 row,
2101 );
2102 drift_envelopes.extend(env);
2103 Ok((kept, None))
2104 }
2105 }
2106 }
2107 }
2108 }
2109}
2110
2111fn quarantine_drift_rows(
2121 diff: &crate::drift::SchemaDiff,
2122 records: Vec<Value>,
2123 page_indices: &[usize],
2124 sink_name: &str,
2125 pipeline_name: &str,
2126 row: &str,
2127) -> (Vec<Value>, Vec<Value>) {
2128 use crate::dlq::{DlqReason, build_envelope};
2129 let present_cols: std::collections::HashSet<&str> = diff
2131 .additions
2132 .iter()
2133 .chain(&diff.widenings)
2134 .chain(&diff.incompatible)
2135 .map(|c| c.name.as_str())
2136 .collect();
2137 let required_cols: std::collections::HashSet<&str> =
2139 diff.droppable_required.iter().map(|s| s.as_str()).collect();
2140 let mut kept = Vec::new();
2141 let mut envelopes = Vec::new();
2142 for (idx, rec) in records.into_iter().enumerate() {
2143 let exhibits = rec
2144 .as_object()
2145 .map(|m| {
2146 m.keys().any(|k| present_cols.contains(k.as_str()))
2147 || required_cols.iter().any(|c| !m.contains_key(*c))
2148 })
2149 .unwrap_or(false);
2150 if exhibits {
2151 let err = FaucetError::SchemaDrift {
2152 columns: diff.changed_columns(),
2153 message: "row exhibits schema drift (on_drift=quarantine)".into(),
2154 };
2155 let page_index = page_indices.get(idx).copied().unwrap_or(idx);
2158 envelopes.push(build_envelope(
2159 &rec,
2160 &err,
2161 DlqReason::SchemaDrift,
2162 sink_name,
2163 pipeline_name,
2164 row,
2165 page_index,
2166 ));
2167 } else {
2168 kept.push(rec);
2169 }
2170 }
2171 (kept, envelopes)
2172}
2173
2174#[cfg(test)]
2175mod tests {
2176 use super::*;
2177 use async_trait::async_trait;
2178 use serde_json::json;
2179
2180 struct MockSource(Vec<Value>);
2183
2184 #[async_trait]
2185 impl Source for MockSource {
2186 async fn fetch_with_context(
2187 &self,
2188 _context: &std::collections::HashMap<String, Value>,
2189 ) -> Result<Vec<Value>, FaucetError> {
2190 Ok(self.0.clone())
2191 }
2192 }
2193
2194 struct IncrementalSource {
2195 records: Vec<Value>,
2196 bookmark: Value,
2197 }
2198
2199 #[async_trait]
2200 impl Source for IncrementalSource {
2201 async fn fetch_with_context(
2202 &self,
2203 _context: &std::collections::HashMap<String, Value>,
2204 ) -> Result<Vec<Value>, FaucetError> {
2205 Ok(self.records.clone())
2206 }
2207 async fn fetch_with_context_incremental(
2208 &self,
2209 _context: &std::collections::HashMap<String, Value>,
2210 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
2211 Ok((self.records.clone(), Some(self.bookmark.clone())))
2212 }
2213 }
2214
2215 struct FailingSource;
2216
2217 #[async_trait]
2218 impl Source for FailingSource {
2219 async fn fetch_with_context(
2220 &self,
2221 _context: &std::collections::HashMap<String, Value>,
2222 ) -> Result<Vec<Value>, FaucetError> {
2223 Err(FaucetError::Auth("no credentials".into()))
2224 }
2225 }
2226
2227 struct MockSink(std::sync::Mutex<Vec<Value>>);
2230
2231 impl MockSink {
2232 fn new() -> Self {
2233 Self(std::sync::Mutex::new(Vec::new()))
2234 }
2235 fn written(&self) -> Vec<Value> {
2236 self.0.lock().unwrap().clone()
2237 }
2238 }
2239
2240 #[async_trait]
2241 impl Sink for MockSink {
2242 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2243 self.0.lock().unwrap().extend(records.iter().cloned());
2244 Ok(records.len())
2245 }
2246 }
2247
2248 struct FailingSink;
2249
2250 #[async_trait]
2251 impl Sink for FailingSink {
2252 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
2253 Err(FaucetError::Sink("write failed".into()))
2254 }
2255 }
2256
2257 struct OverwriteRecordingSink {
2261 events: Arc<std::sync::Mutex<Vec<String>>>,
2262 is_overwrite: bool,
2263 fail_write: bool,
2264 }
2265
2266 impl OverwriteRecordingSink {
2267 fn new(is_overwrite: bool, fail_write: bool) -> (Self, Arc<std::sync::Mutex<Vec<String>>>) {
2268 let events = Arc::new(std::sync::Mutex::new(Vec::new()));
2269 (
2270 Self {
2271 events: events.clone(),
2272 is_overwrite,
2273 fail_write,
2274 },
2275 events,
2276 )
2277 }
2278 fn log(&self, s: &str) {
2279 self.events.lock().unwrap().push(s.to_string());
2280 }
2281 }
2282
2283 #[async_trait]
2284 impl Sink for OverwriteRecordingSink {
2285 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2286 self.log(&format!("write:{}", records.len()));
2287 if self.fail_write {
2288 return Err(FaucetError::Sink("write failed".into()));
2289 }
2290 Ok(records.len())
2291 }
2292 fn is_overwrite(&self) -> bool {
2293 self.is_overwrite
2294 }
2295 async fn begin_overwrite(&self) -> Result<(), FaucetError> {
2296 self.log("begin");
2297 Ok(())
2298 }
2299 async fn commit_overwrite(&self) -> Result<(), FaucetError> {
2300 self.log("commit");
2301 Ok(())
2302 }
2303 async fn abort_overwrite(&self) -> Result<(), FaucetError> {
2304 self.log("abort");
2305 Ok(())
2306 }
2307 }
2308
2309 #[tokio::test]
2310 async fn overwrite_lifecycle_begins_then_commits_on_success() {
2311 let source = MockSource(vec![json!({"id": 1}), json!({"id": 2})]);
2312 let (sink, events) = OverwriteRecordingSink::new(true, false);
2313 let result = Pipeline::new(&source, &sink).run().await;
2314 assert!(result.is_ok());
2315 let log = events.lock().unwrap().clone();
2316 assert_eq!(log.first().map(String::as_str), Some("begin"));
2317 assert_eq!(log.last().map(String::as_str), Some("commit"));
2318 assert!(log.contains(&"write:2".to_string()));
2319 assert!(!log.contains(&"abort".to_string()));
2320 }
2321
2322 #[tokio::test]
2323 async fn overwrite_lifecycle_aborts_on_write_failure() {
2324 let source = MockSource(vec![json!({"id": 1})]);
2325 let (sink, events) = OverwriteRecordingSink::new(true, true);
2326 let result = Pipeline::new(&source, &sink).run().await;
2327 assert!(result.is_err(), "a failed write must fail the run");
2328 let log = events.lock().unwrap().clone();
2329 assert_eq!(log.first().map(String::as_str), Some("begin"));
2330 assert!(
2331 log.contains(&"abort".to_string()),
2332 "failure must trigger abort_overwrite: {log:?}"
2333 );
2334 assert!(
2335 !log.contains(&"commit".to_string()),
2336 "a failed run must never commit the swap: {log:?}"
2337 );
2338 }
2339
2340 #[tokio::test]
2341 async fn overwrite_lifecycle_aborts_on_cancel() {
2342 let source = MockSource(vec![json!({"id": 1})]);
2343 let (sink, events) = OverwriteRecordingSink::new(true, false);
2344 let token = crate::CancellationToken::new();
2345 token.cancel(); let result = Pipeline::new(&source, &sink).with_cancel(token).run().await;
2347 assert!(result.is_ok(), "a cooperative cancel returns Ok(partial)");
2348 let log = events.lock().unwrap().clone();
2349 assert_eq!(log.first().map(String::as_str), Some("begin"));
2350 assert!(
2351 log.contains(&"abort".to_string()),
2352 "a cancelled overwrite run must abort, not commit: {log:?}"
2353 );
2354 assert!(!log.contains(&"commit".to_string()));
2355 }
2356
2357 #[tokio::test]
2358 async fn non_overwrite_sink_skips_the_lifecycle() {
2359 let source = MockSource(vec![json!({"id": 1})]);
2360 let (sink, events) = OverwriteRecordingSink::new(false, false);
2361 let result = Pipeline::new(&source, &sink).run().await;
2362 assert!(result.is_ok());
2363 let log = events.lock().unwrap().clone();
2364 assert!(
2365 !log.iter()
2366 .any(|e| e == "begin" || e == "commit" || e == "abort"),
2367 "a non-overwrite sink must never see the overwrite lifecycle: {log:?}"
2368 );
2369 assert!(log.contains(&"write:1".to_string()));
2370 }
2371
2372 #[tokio::test]
2373 async fn suppressed_overwrite_writes_but_skips_the_lifecycle() {
2374 let source = MockSource(vec![json!({"id": 1}), json!({"id": 2})]);
2379 let (sink, events) = OverwriteRecordingSink::new(true, false);
2380 let result = Pipeline::new(&source, &sink)
2381 .with_suppress_overwrite(true)
2382 .run()
2383 .await;
2384 assert!(result.is_ok());
2385 let log = events.lock().unwrap().clone();
2386 assert!(
2387 log.contains(&"write:2".to_string()),
2388 "suppressed overwrite must still write: {log:?}"
2389 );
2390 assert!(
2391 !log.iter()
2392 .any(|e| e == "begin" || e == "commit" || e == "abort"),
2393 "suppressed overwrite must not drive the lifecycle: {log:?}"
2394 );
2395 }
2396
2397 #[tokio::test]
2398 async fn overwrite_abort_failure_is_logged_not_propagated() {
2399 struct AbortFailSink;
2404 #[async_trait]
2405 impl Sink for AbortFailSink {
2406 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
2407 Err(FaucetError::Sink("write failed".into()))
2408 }
2409 fn is_overwrite(&self) -> bool {
2410 true
2411 }
2412 async fn begin_overwrite(&self) -> Result<(), FaucetError> {
2413 Ok(())
2414 }
2415 async fn commit_overwrite(&self) -> Result<(), FaucetError> {
2416 Ok(())
2417 }
2418 async fn abort_overwrite(&self) -> Result<(), FaucetError> {
2419 Err(FaucetError::Sink("abort failed".into()))
2420 }
2421 }
2422 let source = MockSource(vec![json!({"id": 1})]);
2423 let sink = AbortFailSink;
2424 let err = Pipeline::new(&source, &sink)
2425 .run()
2426 .await
2427 .expect_err("a failed write must fail the run");
2428 assert!(
2429 err.to_string().contains("write failed"),
2430 "the run error must be the write failure, not the abort failure: {err}"
2431 );
2432 }
2433
2434 struct FlushTrackingSink {
2439 written: std::sync::Mutex<Vec<Value>>,
2440 flush_count: std::sync::atomic::AtomicUsize,
2441 }
2442
2443 impl FlushTrackingSink {
2444 fn new() -> Self {
2445 Self {
2446 written: std::sync::Mutex::new(Vec::new()),
2447 flush_count: std::sync::atomic::AtomicUsize::new(0),
2448 }
2449 }
2450 fn written(&self) -> Vec<Value> {
2451 self.written.lock().unwrap().clone()
2452 }
2453 fn flush_count(&self) -> usize {
2454 self.flush_count.load(std::sync::atomic::Ordering::SeqCst)
2455 }
2456 }
2457
2458 #[async_trait]
2459 impl Sink for FlushTrackingSink {
2460 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2461 self.written.lock().unwrap().extend(records.iter().cloned());
2462 Ok(records.len())
2463 }
2464 async fn flush(&self) -> Result<(), FaucetError> {
2465 self.flush_count
2466 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2467 Ok(())
2468 }
2469 }
2470
2471 struct IdempotentMockSink {
2475 rows: std::sync::Mutex<Vec<Value>>,
2476 tokens: std::sync::Mutex<std::collections::HashMap<String, String>>,
2477 }
2478 impl IdempotentMockSink {
2479 fn new() -> Self {
2480 Self {
2481 rows: std::sync::Mutex::new(Vec::new()),
2482 tokens: std::sync::Mutex::new(std::collections::HashMap::new()),
2483 }
2484 }
2485 fn rows(&self) -> Vec<Value> {
2486 self.rows.lock().unwrap().clone()
2487 }
2488 }
2489 #[async_trait]
2490 impl Sink for IdempotentMockSink {
2491 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2492 self.rows.lock().unwrap().extend(records.iter().cloned());
2493 Ok(records.len())
2494 }
2495 fn supports_idempotent_writes(&self) -> bool {
2496 true
2497 }
2498 async fn write_batch_idempotent(
2499 &self,
2500 records: &[Value],
2501 scope: &str,
2502 token: &str,
2503 ) -> Result<usize, FaucetError> {
2504 self.rows.lock().unwrap().extend(records.iter().cloned());
2505 self.tokens
2506 .lock()
2507 .unwrap()
2508 .insert(scope.to_string(), token.to_string());
2509 Ok(records.len())
2510 }
2511 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
2512 Ok(self.tokens.lock().unwrap().get(scope).cloned())
2513 }
2514 }
2515
2516 fn eo_opts(store: Arc<dyn StateStore>, key: &str, start_seq: u64) -> RunStreamOptions {
2517 RunStreamOptions::new()
2518 .with_state(store, key)
2519 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2520 .with_start_seq(start_seq)
2521 }
2522
2523 #[tokio::test]
2524 async fn exactly_once_writes_pages_and_persists_wrapped_state() {
2525 let pages = vec![
2526 Ok(StreamPage {
2527 records: vec![json!({"id": 1})],
2528 bookmark: Some(json!("b1")),
2529 }),
2530 Ok(StreamPage {
2531 records: vec![json!({"id": 2})],
2532 bookmark: Some(json!("b2")),
2533 }),
2534 ];
2535 let sink = IdempotentMockSink::new();
2536 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2537 let r = run_stream(
2538 futures::stream::iter(pages),
2539 &sink,
2540 eo_opts(store.clone(), "k", 0),
2541 )
2542 .await
2543 .unwrap();
2544 assert_eq!(r.records_written, 2);
2545 let (bm, seq) = crate::idempotency::unwrap_state(&store.get("k").await.unwrap().unwrap());
2546 assert_eq!(bm, Some(json!("b2")));
2547 assert_eq!(seq, 2);
2548 let token = sink.last_committed_token("k").await.unwrap().unwrap();
2551 assert_eq!(
2552 crate::idempotency::parse_token_parts(&token),
2553 Some((2, Some(json!("b2"))))
2554 );
2555 }
2556
2557 #[tokio::test]
2558 async fn exactly_once_skips_already_committed_pages_on_resume() {
2559 let sink = IdempotentMockSink::new();
2560 sink.write_batch_idempotent(
2562 &[json!({"id": 1})],
2563 "k",
2564 &crate::idempotency::format_token(1),
2565 )
2566 .await
2567 .unwrap();
2568 assert_eq!(sink.rows().len(), 1);
2569 let pages = vec![
2571 Ok(StreamPage {
2572 records: vec![json!({"id": 1})],
2573 bookmark: Some(json!("b1")),
2574 }),
2575 Ok(StreamPage {
2576 records: vec![json!({"id": 2})],
2577 bookmark: Some(json!("b2")),
2578 }),
2579 ];
2580 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2581 let r = run_stream(futures::stream::iter(pages), &sink, eo_opts(store, "k", 0))
2582 .await
2583 .unwrap();
2584 assert_eq!(r.records_written, 1);
2585 let rows = sink.rows();
2586 assert_eq!(
2587 rows.len(),
2588 2,
2589 "exactly one row per id — no duplicate of id=1"
2590 );
2591 assert_eq!(rows.iter().filter(|v| v["id"] == 1).count(), 1);
2592 }
2593
2594 #[tokio::test]
2595 async fn exactly_once_rejects_non_idempotent_sink() {
2596 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2597 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2598 let r = run_stream(
2599 futures::stream::iter(pages),
2600 &MockSink::new(),
2601 eo_opts(store, "k", 0),
2602 )
2603 .await;
2604 assert!(matches!(r, Err(FaucetError::Config(_))));
2605 }
2606
2607 #[tokio::test]
2608 async fn exactly_once_rejects_missing_state_store() {
2609 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2610 let opts =
2611 RunStreamOptions::new().with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce);
2612 let r = run_stream(
2613 futures::stream::iter(pages),
2614 &IdempotentMockSink::new(),
2615 opts,
2616 )
2617 .await;
2618 assert!(matches!(r, Err(FaucetError::Config(_))));
2619 }
2620
2621 struct KeyedMockSink(MockSink);
2624 #[async_trait]
2625 impl Sink for KeyedMockSink {
2626 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2627 self.0.write_batch(records).await
2628 }
2629 fn dedups_by_key(&self) -> bool {
2630 true
2631 }
2632 fn supported_write_modes(&self) -> &'static [crate::write_mode::WriteMode] {
2633 &[
2634 crate::write_mode::WriteMode::Append,
2635 crate::write_mode::WriteMode::Upsert,
2636 crate::write_mode::WriteMode::Delete,
2637 ]
2638 }
2639 }
2640
2641 #[tokio::test]
2642 async fn exactly_once_keyed_upsert_mechanism_uses_plain_write_path() {
2643 let sink = KeyedMockSink(MockSink::new());
2648 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2649 let pages = vec![Ok(StreamPage {
2650 records: vec![json!({"id": 1})],
2651 bookmark: Some(json!("b1")),
2652 })];
2653 let opts = eo_opts(store.clone(), "k", 0)
2654 .with_replay_guarantee(crate::idempotency::ReplayGuarantee::NonDeterministic);
2655 let r = run_stream(futures::stream::iter(pages), &sink, opts)
2656 .await
2657 .unwrap();
2658 assert_eq!(r.records_written, 1);
2659 assert_eq!(sink.0.written(), vec![json!({"id": 1})]);
2660 assert_eq!(store.get("k").await.unwrap(), Some(json!("b1")));
2662 }
2663
2664 #[tokio::test]
2665 async fn exactly_once_rejects_non_deterministic_source_without_keyed_dedup() {
2666 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2671 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2672 let opts = eo_opts(store, "k", 0)
2673 .with_replay_guarantee(crate::idempotency::ReplayGuarantee::NonDeterministic);
2674 let r = run_stream(
2675 futures::stream::iter(pages),
2676 &IdempotentMockSink::new(),
2677 opts,
2678 )
2679 .await;
2680 match r {
2681 Err(FaucetError::Config(msg)) => {
2682 assert!(msg.contains("write_mode: upsert"), "hint present: {msg}")
2683 }
2684 other => panic!("expected Config error, got {other:?}"),
2685 }
2686 }
2687
2688 #[tokio::test]
2689 async fn exactly_once_rejects_plain_sink_with_mechanism_hint() {
2690 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2691 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2692 let r = run_stream(
2693 futures::stream::iter(pages),
2694 &MockSink::new(),
2695 eo_opts(store, "k", 0),
2696 )
2697 .await;
2698 match r {
2699 Err(FaucetError::Config(msg)) => assert!(
2700 msg.contains("provides neither"),
2701 "names both mechanisms: {msg}"
2702 ),
2703 other => panic!("expected Config error, got {other:?}"),
2704 }
2705 }
2706
2707 struct AnchorRecordingSource {
2710 applied: std::sync::Mutex<Vec<Value>>,
2711 }
2712 #[async_trait]
2713 impl Source for AnchorRecordingSource {
2714 async fn fetch_with_context(
2715 &self,
2716 _context: &std::collections::HashMap<String, Value>,
2717 ) -> Result<Vec<Value>, FaucetError> {
2718 Ok(vec![json!({"id": 10})])
2719 }
2720 async fn fetch_with_context_incremental(
2721 &self,
2722 _context: &std::collections::HashMap<String, Value>,
2723 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
2724 Ok((vec![json!({"id": 10})], Some(json!("after"))))
2725 }
2726 fn state_key(&self) -> Option<String> {
2727 Some("anchor_key".to_string())
2728 }
2729 async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
2730 self.applied.lock().unwrap().push(bookmark);
2731 Ok(())
2732 }
2733 fn supports_exactly_once(&self) -> bool {
2734 true
2735 }
2736 }
2737
2738 #[tokio::test]
2739 async fn pipeline_run_anchors_resume_at_sink_embedded_bookmark() {
2740 let source = AnchorRecordingSource {
2747 applied: std::sync::Mutex::new(Vec::new()),
2748 };
2749 let sink = IdempotentMockSink::new();
2750 sink.tokens.lock().unwrap().insert(
2751 "anchor_key".to_string(),
2752 crate::idempotency::format_token_with_bookmark(5, Some(&json!("sink-pos"))),
2753 );
2754 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2755 store
2756 .put(
2757 "anchor_key",
2758 &crate::idempotency::wrap_state(Some(&json!("state-pos")), 4),
2759 )
2760 .await
2761 .unwrap();
2762
2763 let r = Pipeline::new(&source, &sink)
2764 .with_state_store(Arc::clone(&store))
2765 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2766 .run()
2767 .await
2768 .unwrap();
2769
2770 assert_eq!(
2771 *source.applied.lock().unwrap(),
2772 vec![json!("state-pos"), json!("sink-pos")],
2773 "state bookmark applied, then overridden by the sink anchor"
2774 );
2775 assert_eq!(r.records_written, 1);
2778 let token = sink
2779 .last_committed_token("anchor_key")
2780 .await
2781 .unwrap()
2782 .unwrap();
2783 assert_eq!(
2784 crate::idempotency::parse_token_parts(&token),
2785 Some((6, Some(json!("after"))))
2786 );
2787 let (bm, seq) =
2788 crate::idempotency::unwrap_state(&store.get("anchor_key").await.unwrap().unwrap());
2789 assert_eq!((bm, seq), (Some(json!("after")), 6));
2790 }
2791
2792 #[tokio::test]
2793 async fn pipeline_run_ignores_sink_token_behind_state_seq() {
2794 let source = AnchorRecordingSource {
2797 applied: std::sync::Mutex::new(Vec::new()),
2798 };
2799 let sink = IdempotentMockSink::new();
2800 sink.tokens.lock().unwrap().insert(
2801 "anchor_key".to_string(),
2802 crate::idempotency::format_token_with_bookmark(4, Some(&json!("sink-pos"))),
2803 );
2804 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2805 store
2806 .put(
2807 "anchor_key",
2808 &crate::idempotency::wrap_state(Some(&json!("state-pos")), 4),
2809 )
2810 .await
2811 .unwrap();
2812 Pipeline::new(&source, &sink)
2813 .with_state_store(Arc::clone(&store))
2814 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2815 .run()
2816 .await
2817 .unwrap();
2818 assert_eq!(*source.applied.lock().unwrap(), vec![json!("state-pos")]);
2819 }
2820
2821 #[tokio::test]
2822 async fn pipeline_run_legacy_bare_token_falls_back_to_skip_path() {
2823 let source = AnchorRecordingSource {
2827 applied: std::sync::Mutex::new(Vec::new()),
2828 };
2829 let sink = IdempotentMockSink::new();
2830 sink.tokens.lock().unwrap().insert(
2831 "anchor_key".to_string(),
2832 crate::idempotency::format_token(1),
2833 );
2834 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2835 let r = Pipeline::new(&source, &sink)
2836 .with_state_store(Arc::clone(&store))
2837 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2838 .run()
2839 .await
2840 .unwrap();
2841 assert!(source.applied.lock().unwrap().is_empty());
2842 assert_eq!(r.records_written, 0, "page 1 already committed → skipped");
2843 assert!(sink.rows().is_empty());
2844 }
2845
2846 #[test]
2849 fn stream_page_constructs() {
2850 let page = StreamPage {
2851 records: vec![json!({"id": 1})],
2852 bookmark: Some(json!("2026-05-18")),
2853 };
2854 assert_eq!(page.records.len(), 1);
2855 assert_eq!(page.bookmark, Some(json!("2026-05-18")));
2856 }
2857
2858 #[test]
2859 fn validate_batch_size_accepts_zero_as_no_batching_sentinel() {
2860 assert_eq!(validate_batch_size(0).unwrap(), 0);
2862 }
2863
2864 #[test]
2865 fn validate_batch_size_rejects_too_large() {
2866 let err = validate_batch_size(MAX_BATCH_SIZE + 1).unwrap_err();
2867 assert!(matches!(err, FaucetError::Config(_)));
2868 }
2869
2870 #[test]
2871 fn validate_batch_size_accepts_one() {
2872 assert_eq!(validate_batch_size(1).unwrap(), 1);
2873 }
2874
2875 #[test]
2876 fn validate_batch_size_accepts_max() {
2877 assert_eq!(validate_batch_size(MAX_BATCH_SIZE).unwrap(), MAX_BATCH_SIZE);
2878 }
2879
2880 const _: () = {
2882 assert!(DEFAULT_BATCH_SIZE >= 1);
2883 assert!(DEFAULT_BATCH_SIZE <= MAX_BATCH_SIZE);
2884 };
2885
2886 #[tokio::test]
2889 async fn batch_pipeline_writes_all_records() {
2890 let source = MockSource(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
2891 let sink = MockSink::new();
2892
2893 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2894
2895 assert_eq!(result.records_written, 3);
2896 assert!(result.bookmark.is_none());
2897 assert_eq!(sink.written().len(), 3);
2898 }
2899
2900 #[tokio::test]
2901 async fn batch_pipeline_returns_bookmark() {
2902 let source = IncrementalSource {
2903 records: vec![json!({"id": 1, "ts": "2024-12-01"})],
2904 bookmark: json!("2024-12-01"),
2905 };
2906 let sink = MockSink::new();
2907
2908 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2909
2910 assert_eq!(result.records_written, 1);
2911 assert_eq!(result.bookmark, Some(json!("2024-12-01")));
2912 }
2913
2914 #[tokio::test]
2915 async fn batch_pipeline_empty_source() {
2916 let source = MockSource(vec![]);
2917 let sink = MockSink::new();
2918
2919 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2920
2921 assert_eq!(result.records_written, 0);
2922 assert!(sink.written().is_empty());
2923 }
2924
2925 #[tokio::test]
2926 async fn batch_pipeline_source_error_propagates() {
2927 let source = FailingSource;
2928 let sink = MockSink::new();
2929
2930 let result = Pipeline::new(&source, &sink).run().await;
2931 assert!(result.is_err());
2932 assert!(sink.written().is_empty());
2933 }
2934
2935 #[tokio::test]
2936 async fn batch_pipeline_sink_error_propagates() {
2937 let source = MockSource(vec![json!({"id": 1})]);
2938 let sink = FailingSink;
2939
2940 let result = Pipeline::new(&source, &sink).run().await;
2941 assert!(result.is_err());
2942 }
2943
2944 #[tokio::test]
2945 async fn batch_pipeline_with_trait_objects() {
2946 let source: Box<dyn Source> = Box::new(MockSource(vec![json!({"id": 1})]));
2947 let sink: Box<dyn Sink> = Box::new(MockSink::new());
2948
2949 let result = Pipeline::new(source.as_ref(), sink.as_ref())
2950 .run()
2951 .await
2952 .unwrap();
2953
2954 assert_eq!(result.records_written, 1);
2955 }
2956
2957 #[tokio::test]
2960 async fn stream_pipeline_writes_pages() {
2961 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2962 Ok(StreamPage {
2963 records: vec![json!({"id": 1}), json!({"id": 2})],
2964 bookmark: None,
2965 }),
2966 Ok(StreamPage {
2967 records: vec![json!({"id": 3})],
2968 bookmark: None,
2969 }),
2970 ];
2971 let stream = futures::stream::iter(pages);
2972 let sink = MockSink::new();
2973
2974 let result = run_stream(stream, &sink, RunStreamOptions::new())
2975 .await
2976 .unwrap();
2977
2978 assert_eq!(result.records_written, 3);
2979 assert!(result.bookmark.is_none());
2980 assert_eq!(sink.written().len(), 3);
2981 }
2982
2983 #[tokio::test]
2984 async fn stream_pipeline_flushes_sink_on_source_error() {
2985 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2989 Ok(StreamPage {
2990 records: vec![json!({"id": 1}), json!({"id": 2})],
2991 bookmark: None,
2992 }),
2993 Err(FaucetError::Source("transient blip mid-stream".into())),
2994 ];
2995 let stream = futures::stream::iter(pages);
2996 let sink = FlushTrackingSink::new();
2997
2998 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2999
3000 assert!(matches!(result, Err(FaucetError::Source(_))));
3002 assert_eq!(sink.written().len(), 2);
3004 assert!(
3006 sink.flush_count() >= 1,
3007 "sink must be flushed on the error path so partial output is durable"
3008 );
3009 }
3010
3011 #[tokio::test]
3012 async fn stream_pipeline_flushes_sink_on_cancel() {
3013 use tokio_util::sync::CancellationToken;
3018
3019 let stream = Box::pin(async_stream::stream! {
3021 yield Ok(StreamPage {
3022 records: vec![json!({"id": 1}), json!({"id": 2})],
3023 bookmark: None,
3024 });
3025 futures::future::pending::<()>().await;
3026 });
3027 let sink = FlushTrackingSink::new();
3028
3029 let token = CancellationToken::new();
3030 let canceller = token.clone();
3031 tokio::spawn(async move {
3032 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
3033 canceller.cancel();
3034 });
3035
3036 let result = run_stream(stream, &sink, RunStreamOptions::new().with_cancel(token))
3037 .await
3038 .expect("a cooperative cancel returns Ok with the partial result");
3039
3040 assert_eq!(result.records_written, 2);
3043 assert_eq!(sink.written().len(), 2);
3044 assert!(
3045 sink.flush_count() >= 1,
3046 "sink must be flushed on the cancel path so partial output is durable"
3047 );
3048 }
3049
3050 #[tokio::test]
3051 async fn stream_pipeline_empty() {
3052 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
3053 let stream = futures::stream::iter(pages);
3054 let sink = MockSink::new();
3055
3056 let result = run_stream(stream, &sink, RunStreamOptions::new())
3057 .await
3058 .unwrap();
3059
3060 assert_eq!(result.records_written, 0);
3061 }
3062
3063 #[tokio::test]
3064 async fn stream_pipeline_skips_empty_pages() {
3065 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
3066 Ok(StreamPage {
3067 records: vec![json!({"id": 1})],
3068 bookmark: None,
3069 }),
3070 Ok(StreamPage {
3071 records: vec![],
3072 bookmark: None,
3073 }),
3074 Ok(StreamPage {
3075 records: vec![json!({"id": 2})],
3076 bookmark: None,
3077 }),
3078 ];
3079 let stream = futures::stream::iter(pages);
3080 let sink = MockSink::new();
3081
3082 let result = run_stream(stream, &sink, RunStreamOptions::new())
3083 .await
3084 .unwrap();
3085
3086 assert_eq!(result.records_written, 2);
3087 }
3088
3089 #[tokio::test]
3090 async fn stream_pipeline_error_in_page_propagates() {
3091 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
3092 Ok(StreamPage {
3093 records: vec![json!({"id": 1})],
3094 bookmark: None,
3095 }),
3096 Err(FaucetError::HttpStatus {
3097 status: 500,
3098 url: "https://example.com".into(),
3099 body: "Internal Server Error".into(),
3100 }),
3101 ];
3102 let stream = futures::stream::iter(pages);
3103 let sink = MockSink::new();
3104
3105 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
3106 assert!(result.is_err());
3107 assert_eq!(sink.written().len(), 1);
3109 }
3110
3111 #[tokio::test]
3112 async fn stream_pipeline_sink_error_propagates() {
3113 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3114 records: vec![json!({"id": 1})],
3115 bookmark: None,
3116 })];
3117 let stream = futures::stream::iter(pages);
3118 let sink = FailingSink;
3119
3120 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
3121 assert!(result.is_err());
3122 }
3123
3124 #[tokio::test]
3125 async fn stream_pipeline_with_trait_object_sink() {
3126 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3127 records: vec![json!({"id": 1})],
3128 bookmark: None,
3129 })];
3130 let stream = futures::stream::iter(pages);
3131 let sink: Box<dyn Sink> = Box::new(MockSink::new());
3132
3133 let result = run_stream(stream, sink.as_ref(), RunStreamOptions::new())
3134 .await
3135 .unwrap();
3136 assert_eq!(result.records_written, 1);
3137 }
3138
3139 #[tokio::test]
3140 async fn stream_pipeline_persists_bookmark_when_page_carries_one() {
3141 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3142 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
3143 Ok(StreamPage {
3144 records: vec![json!({"id": 1})],
3145 bookmark: None,
3146 }),
3147 Ok(StreamPage {
3148 records: vec![json!({"id": 2})],
3149 bookmark: Some(json!("checkpoint-final")),
3150 }),
3151 ];
3152 let stream = futures::stream::iter(pages);
3153 let sink = MockSink::new();
3154
3155 let result = run_stream(
3156 stream,
3157 &sink,
3158 RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
3159 )
3160 .await
3161 .unwrap();
3162
3163 assert_eq!(result.records_written, 2);
3164 assert_eq!(result.bookmark, Some(json!("checkpoint-final")));
3165 assert_eq!(
3166 store.get("k").await.unwrap(),
3167 Some(json!("checkpoint-final"))
3168 );
3169 }
3170
3171 #[tokio::test]
3172 async fn stream_pipeline_persists_per_page_bookmarks() {
3173 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3174 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
3175 Ok(StreamPage {
3176 records: vec![json!({"id": 1})],
3177 bookmark: Some(json!("tx-1")),
3178 }),
3179 Ok(StreamPage {
3180 records: vec![json!({"id": 2})],
3181 bookmark: Some(json!("tx-2")),
3182 }),
3183 ];
3184 let stream = futures::stream::iter(pages);
3185 let sink = MockSink::new();
3186
3187 run_stream(
3188 stream,
3189 &sink,
3190 RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
3191 )
3192 .await
3193 .unwrap();
3194
3195 assert_eq!(store.get("k").await.unwrap(), Some(json!("tx-2")));
3197 }
3198
3199 use crate::state::{FileStateStore, MemoryStateStore, StateStore};
3202 use std::sync::Arc;
3203 use tempfile::TempDir;
3204
3205 struct StatefulSource {
3209 key: String,
3210 records: Vec<Value>,
3211 new_bookmark: Value,
3212 seen_bookmark: std::sync::Mutex<Option<Value>>,
3213 }
3214
3215 impl StatefulSource {
3216 fn new(key: &str, records: Vec<Value>, new_bookmark: Value) -> Self {
3217 Self {
3218 key: key.into(),
3219 records,
3220 new_bookmark,
3221 seen_bookmark: std::sync::Mutex::new(None),
3222 }
3223 }
3224 fn observed_start(&self) -> Option<Value> {
3225 self.seen_bookmark.lock().unwrap().clone()
3226 }
3227 }
3228
3229 #[async_trait]
3230 impl Source for StatefulSource {
3231 async fn fetch_with_context(
3232 &self,
3233 _ctx: &std::collections::HashMap<String, Value>,
3234 ) -> Result<Vec<Value>, FaucetError> {
3235 Ok(self.records.clone())
3236 }
3237 async fn fetch_with_context_incremental(
3238 &self,
3239 _ctx: &std::collections::HashMap<String, Value>,
3240 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
3241 Ok((self.records.clone(), Some(self.new_bookmark.clone())))
3242 }
3243 fn state_key(&self) -> Option<String> {
3244 Some(self.key.clone())
3245 }
3246 async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
3247 *self.seen_bookmark.lock().unwrap() = Some(bookmark);
3248 Ok(())
3249 }
3250 }
3251
3252 #[tokio::test]
3253 async fn pipeline_with_state_store_persists_bookmark_after_sink() {
3254 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3255 let source = StatefulSource::new(
3256 "github_issues",
3257 vec![json!({"id": 1, "ts": "2026-05-01"})],
3258 json!("2026-05-01"),
3259 );
3260 let sink = MockSink::new();
3261 let result = Pipeline::new(&source, &sink)
3262 .with_state_store(Arc::clone(&store))
3263 .run()
3264 .await
3265 .unwrap();
3266
3267 assert_eq!(result.records_written, 1);
3268 assert_eq!(result.bookmark, Some(json!("2026-05-01")));
3269 let stored = store.get("github_issues").await.unwrap();
3271 assert_eq!(stored, Some(json!("2026-05-01")));
3272 }
3273
3274 #[tokio::test]
3275 async fn pipeline_with_state_store_resumes_from_stored_bookmark() {
3276 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3277 store
3278 .put("github_issues", &json!("2026-04-30"))
3279 .await
3280 .unwrap();
3281
3282 let source =
3283 StatefulSource::new("github_issues", vec![json!({"id": 2})], json!("2026-05-01"));
3284 let sink = MockSink::new();
3285 Pipeline::new(&source, &sink)
3286 .with_state_store(Arc::clone(&store))
3287 .run()
3288 .await
3289 .unwrap();
3290
3291 assert_eq!(source.observed_start(), Some(json!("2026-04-30")));
3293 assert_eq!(
3295 store.get("github_issues").await.unwrap(),
3296 Some(json!("2026-05-01"))
3297 );
3298 }
3299
3300 #[tokio::test]
3301 async fn pipeline_with_state_store_does_not_persist_when_sink_fails() {
3302 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3303 let source = StatefulSource::new("k", vec![json!({"id": 1})], json!("2026-05-01"));
3304 let sink = FailingSink;
3305
3306 let result = Pipeline::new(&source, &sink)
3307 .with_state_store(Arc::clone(&store))
3308 .run()
3309 .await;
3310 assert!(result.is_err());
3311 assert!(store.get("k").await.unwrap().is_none());
3312 }
3313
3314 #[tokio::test]
3315 async fn pipeline_with_state_store_no_state_key_means_no_persist() {
3316 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3317 let source = IncrementalSource {
3318 records: vec![json!({"id": 1})],
3319 bookmark: json!("ignored"),
3320 };
3321 let sink = MockSink::new();
3322 Pipeline::new(&source, &sink)
3323 .with_state_store(Arc::clone(&store))
3324 .run()
3325 .await
3326 .unwrap();
3327 assert!(store.get("anything").await.unwrap().is_none());
3330 }
3331
3332 #[tokio::test]
3333 async fn pipeline_with_state_store_skips_persist_when_bookmark_is_none() {
3334 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3335 struct NoBookmarkSource;
3336 #[async_trait]
3337 impl Source for NoBookmarkSource {
3338 async fn fetch_with_context(
3339 &self,
3340 _ctx: &std::collections::HashMap<String, Value>,
3341 ) -> Result<Vec<Value>, FaucetError> {
3342 Ok(vec![json!({"id": 1})])
3343 }
3344 fn state_key(&self) -> Option<String> {
3345 Some("k".into())
3346 }
3347 }
3348 let source = NoBookmarkSource;
3349 let sink = MockSink::new();
3350 Pipeline::new(&source, &sink)
3351 .with_state_store(Arc::clone(&store))
3352 .run()
3353 .await
3354 .unwrap();
3355 assert!(store.get("k").await.unwrap().is_none());
3356 }
3357
3358 struct PagedSource;
3363
3364 #[async_trait]
3365 impl Source for PagedSource {
3366 async fn fetch_with_context(
3367 &self,
3368 _ctx: &std::collections::HashMap<String, Value>,
3369 ) -> Result<Vec<Value>, FaucetError> {
3370 unreachable!("Pipeline::run must drive stream_pages, not fetch_with_context");
3372 }
3373 fn stream_pages<'a>(
3374 &'a self,
3375 _ctx: &'a std::collections::HashMap<String, Value>,
3376 _batch_size: usize,
3377 ) -> std::pin::Pin<
3378 Box<dyn futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>,
3379 > {
3380 Box::pin(async_stream::try_stream! {
3381 yield StreamPage { records: vec![json!({"i": 1})], bookmark: None };
3382 yield StreamPage { records: vec![json!({"i": 2})], bookmark: None };
3383 yield StreamPage { records: vec![json!({"i": 3})], bookmark: Some(json!("final")) };
3384 })
3385 }
3386 }
3387
3388 struct CountingSink {
3390 calls: std::sync::Mutex<Vec<usize>>,
3391 }
3392
3393 impl CountingSink {
3394 fn new() -> Self {
3395 Self {
3396 calls: std::sync::Mutex::new(Vec::new()),
3397 }
3398 }
3399 fn call_count(&self) -> usize {
3400 self.calls.lock().unwrap().len()
3401 }
3402 }
3403
3404 #[async_trait]
3405 impl Sink for CountingSink {
3406 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3407 self.calls.lock().unwrap().push(records.len());
3408 Ok(records.len())
3409 }
3410 }
3411
3412 #[tokio::test]
3413 async fn pipeline_run_drives_stream_pages() {
3414 let source = PagedSource;
3415 let sink = CountingSink::new();
3416
3417 let result = Pipeline::new(&source, &sink).run().await.unwrap();
3418
3419 assert_eq!(sink.call_count(), 3);
3421 assert_eq!(result.records_written, 3);
3422 assert_eq!(result.bookmark, Some(json!("final")));
3423 }
3424
3425 #[tokio::test]
3426 async fn pipeline_with_file_state_store_round_trips_across_runs() {
3427 let dir = TempDir::new().unwrap();
3428 let store: Arc<dyn StateStore> = Arc::new(FileStateStore::new(dir.path()));
3429
3430 let s1 = StatefulSource::new("k", vec![json!({"i": 1})], json!("v1"));
3432 let sink1 = MockSink::new();
3433 Pipeline::new(&s1, &sink1)
3434 .with_state_store(Arc::clone(&store))
3435 .run()
3436 .await
3437 .unwrap();
3438 assert_eq!(s1.observed_start(), None);
3439 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3440
3441 let s2 = StatefulSource::new("k", vec![json!({"i": 2})], json!("v2"));
3443 let sink2 = MockSink::new();
3444 Pipeline::new(&s2, &sink2)
3445 .with_state_store(Arc::clone(&store))
3446 .run()
3447 .await
3448 .unwrap();
3449 assert_eq!(s2.observed_start(), Some(json!("v1")));
3450 assert_eq!(store.get("k").await.unwrap(), Some(json!("v2")));
3451 }
3452
3453 #[tokio::test]
3454 #[allow(clippy::await_holding_lock)]
3455 async fn pipeline_run_increments_runs_total() {
3456 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3457 use metrics_util::debugging::DebugValue;
3458 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3459 let snap = snapshotter();
3460
3461 let source = MockSource(vec![json!({"i": 1})]);
3462 let sink = MockSink::new();
3463 let _ = Pipeline::new(&source, &sink)
3464 .with_name("test-pipeline")
3465 .with_row("rowA")
3466 .run()
3467 .await
3468 .unwrap();
3469
3470 let snapshot = snap.snapshot();
3471 let found = snapshot.into_vec().into_iter().any(
3472 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3473 key.key().name() == "faucet_pipeline_runs_total"
3474 && key.key().labels().any(|l: &metrics::Label| {
3475 l.key() == "pipeline" && l.value() == "test-pipeline"
3476 })
3477 && key
3478 .key()
3479 .labels()
3480 .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowA")
3481 && key
3482 .key()
3483 .labels()
3484 .any(|l: &metrics::Label| l.key() == "status" && l.value() == "ok")
3485 && matches!(v, DebugValue::Counter(c) if c >= 1)
3486 },
3487 );
3488 assert!(
3489 found,
3490 "expected faucet_pipeline_runs_total{{pipeline=test-pipeline, row=rowA, status=ok}}"
3491 );
3492 }
3493
3494 #[tokio::test]
3495 #[allow(clippy::await_holding_lock)]
3496 async fn pipeline_failure_attaches_kind_label_to_runs_total() {
3497 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3498 use metrics_util::debugging::DebugValue;
3499 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3500 let snap = snapshotter();
3501
3502 let source = FailingSource;
3503 let sink = MockSink::new();
3504 let _ = Pipeline::new(&source, &sink)
3505 .with_name("err-pipeline")
3506 .with_row("rowE")
3507 .run()
3508 .await;
3509
3510 let snapshot = snap.snapshot();
3511 let found = snapshot.into_vec().into_iter().any(
3512 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3513 key.key().name() == "faucet_pipeline_runs_total"
3514 && key.key().labels().any(|l: &metrics::Label| {
3515 l.key() == "pipeline" && l.value() == "err-pipeline"
3516 })
3517 && key
3518 .key()
3519 .labels()
3520 .any(|l: &metrics::Label| l.key() == "status" && l.value() == "err")
3521 && key
3522 .key()
3523 .labels()
3524 .any(|l: &metrics::Label| l.key() == "kind" && l.value() == "Auth")
3525 && matches!(v, DebugValue::Counter(c) if c >= 1)
3526 },
3527 );
3528 assert!(
3529 found,
3530 "expected faucet_pipeline_runs_total{{status=err, kind=Auth}} for failing source"
3531 );
3532 }
3533
3534 #[tokio::test]
3535 #[allow(clippy::await_holding_lock)]
3536 async fn pipeline_run_emits_start_time_gauge() {
3537 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3538 use metrics_util::debugging::DebugValue;
3539 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3540 let snap = snapshotter();
3541
3542 let source = MockSource(vec![json!({"i": 1})]);
3543 let sink = MockSink::new();
3544 let before = std::time::SystemTime::now()
3545 .duration_since(std::time::UNIX_EPOCH)
3546 .map(|d| d.as_secs_f64())
3547 .unwrap_or(0.0);
3548 let _ = Pipeline::new(&source, &sink)
3549 .with_name("start-time-pipeline")
3550 .with_row("rowS")
3551 .run()
3552 .await
3553 .unwrap();
3554
3555 let snapshot = snap.snapshot();
3556 let found = snapshot.into_vec().into_iter().any(
3557 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3558 if key.key().name() != "faucet_pipeline_start_time_unix_seconds" {
3559 return false;
3560 }
3561 let labels_match = key.key().labels().any(|l: &metrics::Label| {
3562 l.key() == "pipeline" && l.value() == "start-time-pipeline"
3563 }) && key
3564 .key()
3565 .labels()
3566 .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowS");
3567 if !labels_match {
3568 return false;
3569 }
3570 matches!(v, DebugValue::Gauge(g) if g.into_inner() >= before)
3571 },
3572 );
3573 assert!(
3574 found,
3575 "expected faucet_pipeline_start_time_unix_seconds gauge >= test-start timestamp"
3576 );
3577 }
3578
3579 #[tokio::test]
3580 #[allow(clippy::await_holding_lock)]
3581 async fn register_build_info_sets_version_gauge() {
3582 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3583 use metrics_util::debugging::DebugValue;
3584 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3585 let snap = snapshotter();
3586
3587 crate::observability::register_build_info();
3588
3589 let snapshot = snap.snapshot();
3590 let found = snapshot.into_vec().into_iter().any(
3591 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3592 key.key().name() == "faucet_build_info"
3593 && key.key().labels().any(|l: &metrics::Label| {
3594 l.key() == "version" && l.value() == env!("CARGO_PKG_VERSION")
3595 })
3596 && matches!(v, DebugValue::Gauge(g) if (g.into_inner() - 1.0).abs() < f64::EPSILON)
3597 },
3598 );
3599 assert!(
3600 found,
3601 "expected faucet_build_info{{version=CARGO_PKG_VERSION}} = 1.0 after register_build_info()"
3602 );
3603 }
3604
3605 use crate::dlq::{DlqConfig, OnBatchError};
3608
3609 struct PartialSink {
3613 fail_indices: std::sync::Mutex<Vec<usize>>,
3614 committed: std::sync::Mutex<Vec<Value>>,
3615 }
3616
3617 impl PartialSink {
3618 fn new(fail_indices: Vec<usize>) -> Self {
3619 Self {
3620 fail_indices: std::sync::Mutex::new(fail_indices),
3621 committed: std::sync::Mutex::new(Vec::new()),
3622 }
3623 }
3624 }
3625
3626 #[async_trait]
3627 impl Sink for PartialSink {
3628 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
3629 unreachable!("PartialSink only overrides write_batch_partial");
3630 }
3631 async fn write_batch_partial(
3632 &self,
3633 records: &[Value],
3634 ) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
3635 let fails: std::collections::HashSet<usize> =
3636 self.fail_indices.lock().unwrap().iter().copied().collect();
3637 let mut outcomes = Vec::with_capacity(records.len());
3638 for (i, rec) in records.iter().enumerate() {
3639 if fails.contains(&i) {
3640 outcomes.push(Err(FaucetError::Sink(format!("row {i} rejected"))));
3641 } else {
3642 self.committed.lock().unwrap().push(rec.clone());
3643 outcomes.push(Ok(()));
3644 }
3645 }
3646 Ok(outcomes)
3647 }
3648 }
3649
3650 #[tokio::test]
3651 async fn dlq_routes_only_failed_rows_for_partial_success_sink() {
3652 let main = PartialSink::new(vec![1, 3]); let dlq = std::sync::Arc::new(MockSink::new());
3654 let dlq_cfg = DlqConfig::new(dlq.clone());
3655
3656 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3657 records: (0..4).map(|i| json!({"i": i})).collect(),
3658 bookmark: None,
3659 })];
3660 let stream = futures::stream::iter(pages);
3661 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg))
3662 .await
3663 .unwrap();
3664
3665 assert_eq!(result.records_written, 2); assert_eq!(main.committed.lock().unwrap().len(), 2);
3667 let envelopes = dlq.0.lock().unwrap();
3668 assert_eq!(envelopes.len(), 2);
3669 assert_eq!(envelopes[0]["payload"]["i"], 1);
3670 assert_eq!(envelopes[0]["record_index"], 1);
3671 assert_eq!(envelopes[1]["payload"]["i"], 3);
3672 assert_eq!(envelopes[1]["record_index"], 3);
3673 let stats = result.dlq.unwrap();
3674 assert_eq!(stats.records_dlq, 2);
3675 assert_eq!(stats.pages_with_failures, 1);
3676 }
3677
3678 #[cfg(feature = "masking")]
3679 #[tokio::test]
3680 async fn masking_runs_before_the_sink() {
3681 use crate::masking::{CompiledMasking, MaskingSpec};
3682 let sink = MockSink::new();
3683 let spec: MaskingSpec = serde_json::from_value(json!({
3684 "rules": [{ "match": { "value_detector": "email" },
3685 "action": { "type": "redact" } }]
3686 }))
3687 .unwrap();
3688 let m = std::sync::Arc::new(CompiledMasking::compile(&spec).unwrap());
3689 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3690 records: vec![json!({"email": "a@b.com", "name": "Al"})],
3691 bookmark: None,
3692 })];
3693 run_stream(
3694 futures::stream::iter(pages),
3695 &sink,
3696 RunStreamOptions::new().with_masking(m),
3697 )
3698 .await
3699 .unwrap();
3700 assert_eq!(sink.written()[0], json!({"email": "***", "name": "Al"}));
3701 }
3702
3703 #[cfg(feature = "masking")]
3704 #[tokio::test]
3705 async fn masking_applies_before_the_dlq_envelope() {
3706 use crate::masking::{CompiledMasking, MaskingSpec};
3710 let main = PartialSink::new(vec![0]); let dlq = std::sync::Arc::new(MockSink::new());
3712 let spec: MaskingSpec = serde_json::from_value(json!({
3713 "rules": [{ "match": { "value_detector": "email" },
3714 "action": { "type": "redact" } }]
3715 }))
3716 .unwrap();
3717 let m = std::sync::Arc::new(CompiledMasking::compile(&spec).unwrap());
3718 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3719 records: vec![
3720 json!({"email": "secret@x.com"}),
3721 json!({"email": "ok@y.com"}),
3722 ],
3723 bookmark: None,
3724 })];
3725 let opts = RunStreamOptions::new()
3726 .with_masking(m)
3727 .with_dlq(DlqConfig::new(dlq.clone()));
3728 run_stream(futures::stream::iter(pages), &main, opts)
3729 .await
3730 .unwrap();
3731 let env = dlq.0.lock().unwrap();
3733 assert_eq!(env.len(), 1);
3734 assert_eq!(
3735 env[0]["payload"]["email"], "***",
3736 "the DLQ payload must be masked, not raw PII"
3737 );
3738 assert_eq!(main.committed.lock().unwrap()[0]["email"], "***");
3739 }
3740
3741 #[tokio::test]
3742 async fn dlq_propagate_policy_bubbles_outer_err() {
3743 let main = FailingSink;
3744 let dlq = std::sync::Arc::new(MockSink::new());
3745 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3746 dlq_cfg.on_batch_error = OnBatchError::Propagate;
3747
3748 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3749 records: vec![json!({"i": 0}), json!({"i": 1})],
3750 bookmark: Some(json!("v1")),
3751 })];
3752 let stream = futures::stream::iter(pages);
3753 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3754 let result = run_stream(
3755 stream,
3756 &main,
3757 RunStreamOptions::new()
3758 .with_dlq(dlq_cfg)
3759 .with_state(std::sync::Arc::clone(&store), "k"),
3760 )
3761 .await;
3762 assert!(matches!(result, Err(FaucetError::Sink(_))));
3763 assert!(dlq.0.lock().unwrap().is_empty());
3764 assert!(store.get("k").await.unwrap().is_none());
3766 }
3767
3768 #[tokio::test]
3769 async fn dlq_dlq_all_policy_routes_every_row_on_outer_err() {
3770 let main = FailingSink;
3771 let dlq = std::sync::Arc::new(MockSink::new());
3772 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3773 dlq_cfg.on_batch_error = OnBatchError::DlqAll;
3774
3775 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3776 records: vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})],
3777 bookmark: Some(json!("v1")),
3778 })];
3779 let stream = futures::stream::iter(pages);
3780 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3781 let result = run_stream(
3782 stream,
3783 &main,
3784 RunStreamOptions::new()
3785 .with_dlq(dlq_cfg)
3786 .with_state(std::sync::Arc::clone(&store), "k"),
3787 )
3788 .await
3789 .unwrap();
3790 assert_eq!(result.records_written, 0);
3791 {
3792 let envelopes = dlq.0.lock().unwrap();
3793 assert_eq!(envelopes.len(), 3);
3794 for env in envelopes.iter() {
3796 let msg = env["error"]["message"].as_str().unwrap();
3797 assert!(msg.contains("write failed"), "got: {msg}");
3798 }
3799 }
3800 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3801 assert_eq!(result.dlq.unwrap().records_dlq, 3);
3802 }
3803
3804 #[tokio::test]
3805 async fn dlq_per_page_budget_exceeded_aborts() {
3806 let main = PartialSink::new(vec![0, 1, 2]);
3807 let dlq = std::sync::Arc::new(MockSink::new());
3808 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3809 dlq_cfg.max_failures_per_page = Some(2);
3810
3811 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3812 records: (0..3).map(|i| json!({"i": i})).collect(),
3813 bookmark: None,
3814 })];
3815 let stream = futures::stream::iter(pages);
3816 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3817 assert!(
3818 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
3819 "got: {result:?}"
3820 );
3821 }
3822
3823 #[tokio::test]
3824 async fn dlq_total_budget_exceeded_aborts_on_later_page() {
3825 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
3826 Ok(StreamPage {
3827 records: (0..3).map(|i| json!({"i": i})).collect(),
3828 bookmark: None,
3829 }),
3830 Ok(StreamPage {
3831 records: (3..6).map(|i| json!({"i": i})).collect(),
3832 bookmark: None,
3833 }),
3834 ];
3835 let main = PartialSink::new(vec![0, 1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3838 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3839 dlq_cfg.max_failures_total = Some(4);
3840
3841 let stream = futures::stream::iter(pages);
3842 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3843 assert!(
3844 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
3845 "got: {result:?}"
3846 );
3847 }
3848
3849 #[tokio::test]
3850 async fn dlq_per_page_budget_exceeded_commits_page_before_aborting() {
3851 let main = PartialSink::new(vec![1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3858 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3859 dlq_cfg.max_failures_per_page = Some(1); let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3862 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3863 records: (0..3).map(|i| json!({ "i": i })).collect(),
3864 bookmark: Some(json!("v1")),
3865 })];
3866 let stream = futures::stream::iter(pages);
3867 let result = run_stream(
3868 stream,
3869 &main,
3870 RunStreamOptions::new()
3871 .with_dlq(dlq_cfg)
3872 .with_state(std::sync::Arc::clone(&store), "k"),
3873 )
3874 .await;
3875
3876 assert!(
3878 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
3879 "got: {result:?}"
3880 );
3881 assert_eq!(main.committed.lock().unwrap().len(), 1);
3883 assert_eq!(dlq.0.lock().unwrap().len(), 2);
3885 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3887 }
3888
3889 #[tokio::test]
3890 async fn dlq_total_budget_exceeded_commits_tripping_page_before_aborting() {
3891 let main = PartialSink::new(vec![1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3896 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3897 dlq_cfg.max_failures_total = Some(1); let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3900 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3901 records: (0..3).map(|i| json!({ "i": i })).collect(),
3902 bookmark: Some(json!("v1")),
3903 })];
3904 let stream = futures::stream::iter(pages);
3905 let result = run_stream(
3906 stream,
3907 &main,
3908 RunStreamOptions::new()
3909 .with_dlq(dlq_cfg)
3910 .with_state(std::sync::Arc::clone(&store), "k"),
3911 )
3912 .await;
3913
3914 assert!(
3915 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
3916 "got: {result:?}"
3917 );
3918 assert_eq!(main.committed.lock().unwrap().len(), 1);
3919 assert_eq!(dlq.0.lock().unwrap().len(), 2);
3920 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3921 }
3922
3923 struct FailingDlqSink;
3926 #[async_trait]
3927 impl Sink for FailingDlqSink {
3928 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
3929 Err(FaucetError::Sink("dlq disk full".into()))
3930 }
3931 }
3932
3933 struct FailingFlushDlqSink {
3937 written: std::sync::Mutex<Vec<Value>>,
3938 }
3939 impl FailingFlushDlqSink {
3940 fn new() -> Self {
3941 Self {
3942 written: std::sync::Mutex::new(Vec::new()),
3943 }
3944 }
3945 }
3946 #[async_trait]
3947 impl Sink for FailingFlushDlqSink {
3948 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3949 self.written.lock().unwrap().extend(records.iter().cloned());
3950 Ok(records.len())
3951 }
3952 async fn flush(&self) -> Result<(), FaucetError> {
3953 Err(FaucetError::Sink("dlq flush failed".into()))
3954 }
3955 }
3956
3957 #[tokio::test]
3958 async fn dlq_sink_failure_is_fatal_no_recursion() {
3959 let main = PartialSink::new(vec![0]);
3960 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingDlqSink);
3961 let dlq_cfg = DlqConfig::new(dlq);
3962
3963 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3964 records: vec![json!({"i": 0}), json!({"i": 1})],
3965 bookmark: Some(json!("v1")),
3966 })];
3967 let stream = futures::stream::iter(pages);
3968 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3969 let result = run_stream(
3970 stream,
3971 &main,
3972 RunStreamOptions::new()
3973 .with_dlq(dlq_cfg)
3974 .with_state(std::sync::Arc::clone(&store), "k"),
3975 )
3976 .await;
3977 assert!(
3978 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink write failed")),
3979 "got: {result:?}"
3980 );
3981 assert!(store.get("k").await.unwrap().is_none());
3982 }
3983
3984 #[tokio::test]
3985 async fn dlq_bookmark_advances_only_after_both_flushes() {
3986 let main = PartialSink::new(vec![1]); let dlq = std::sync::Arc::new(MockSink::new());
3988 let dlq_cfg = DlqConfig::new(dlq.clone());
3989
3990 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3991 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3992 records: vec![json!({"i": 0}), json!({"i": 1})],
3993 bookmark: Some(json!("v1")),
3994 })];
3995 let stream = futures::stream::iter(pages);
3996 run_stream(
3997 stream,
3998 &main,
3999 RunStreamOptions::new()
4000 .with_dlq(dlq_cfg)
4001 .with_state(std::sync::Arc::clone(&store), "k"),
4002 )
4003 .await
4004 .unwrap();
4005 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
4006 assert_eq!(dlq.0.lock().unwrap().len(), 1);
4007 assert_eq!(main.committed.lock().unwrap().len(), 1);
4008 }
4009
4010 #[tokio::test]
4011 async fn dlq_disabled_pipeline_behaves_identically_to_today() {
4012 let main = MockSink::new();
4014 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4015 records: vec![json!({"i": 0}), json!({"i": 1})],
4016 bookmark: None,
4017 })];
4018 let stream = futures::stream::iter(pages);
4019 let result = run_stream(stream, &main, RunStreamOptions::new())
4020 .await
4021 .unwrap();
4022 assert_eq!(result.records_written, 2);
4023 assert!(result.dlq.is_none());
4024 }
4025
4026 #[tokio::test]
4027 async fn dlq_per_page_flush_failure_is_fatal_and_blocks_bookmark() {
4028 let main = PartialSink::new(vec![1]);
4033 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
4034 let dlq_cfg = DlqConfig::new(dlq);
4035
4036 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
4037 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4038 records: vec![json!({"i": 0}), json!({"i": 1})],
4039 bookmark: Some(json!("v1")),
4040 })];
4041 let stream = futures::stream::iter(pages);
4042 let result = run_stream(
4043 stream,
4044 &main,
4045 RunStreamOptions::new()
4046 .with_dlq(dlq_cfg)
4047 .with_state(std::sync::Arc::clone(&store), "k"),
4048 )
4049 .await;
4050 assert!(
4051 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
4052 "got: {result:?}"
4053 );
4054 assert!(store.get("k").await.unwrap().is_none());
4055 }
4056
4057 #[tokio::test]
4058 async fn dlq_end_of_stream_flush_failure_is_fatal() {
4059 let main = PartialSink::new(vec![1]);
4064 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
4065 let dlq_cfg = DlqConfig::new(dlq);
4066
4067 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4068 records: vec![json!({"i": 0}), json!({"i": 1})],
4069 bookmark: None,
4070 })];
4071 let stream = futures::stream::iter(pages);
4072 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
4073 assert!(
4074 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
4075 "got: {result:?}"
4076 );
4077 }
4078
4079 #[tokio::test]
4080 #[allow(clippy::await_holding_lock)]
4081 async fn dlq_emits_records_total_and_pages_total() {
4082 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
4083 use metrics_util::debugging::DebugValue;
4084
4085 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
4086 let snap = snapshotter();
4087
4088 let source = MockSource(vec![json!({"i": 0}), json!({"i": 1})]);
4089 let main = PartialSink::new(vec![1]);
4090 let dlq = std::sync::Arc::new(MockSink::new());
4091 let _ = Pipeline::new(&source, &main)
4092 .with_name("p_dlq_metrics")
4093 .with_row("r1")
4094 .with_dlq(DlqConfig::new(dlq.clone()))
4095 .run()
4096 .await
4097 .unwrap();
4098
4099 let snapshot = snap.snapshot();
4100 let mut saw_records = false;
4101 let mut saw_pages = false;
4102 for (k, _u, _d, v) in snapshot.into_vec() {
4103 let key = k.key();
4104 let labels = key.labels().collect::<Vec<_>>();
4105 let has = |k: &str, v: &str| labels.iter().any(|l| l.key() == k && l.value() == v);
4106 if key.name() == "faucet_sink_dlq_records_total"
4107 && has("pipeline", "p_dlq_metrics")
4108 && has("row", "r1")
4109 && matches!(v, DebugValue::Counter(c) if c >= 1)
4110 {
4111 saw_records = true;
4112 }
4113 if key.name() == "faucet_sink_dlq_pages_total"
4114 && has("pipeline", "p_dlq_metrics")
4115 && matches!(v, DebugValue::Counter(c) if c >= 1)
4116 {
4117 saw_pages = true;
4118 }
4119 }
4120 assert!(saw_records, "faucet_sink_dlq_records_total not emitted");
4121 assert!(saw_pages, "faucet_sink_dlq_pages_total not emitted");
4122 }
4123
4124 #[tokio::test]
4125 #[allow(clippy::await_holding_lock)]
4126 async fn dlq_budget_exceeded_emits_counter() {
4127 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
4128 use metrics_util::debugging::DebugValue;
4129
4130 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
4131 let snap = snapshotter();
4132
4133 let source = MockSource((0..3).map(|i| json!({"i": i})).collect());
4134 let main = PartialSink::new(vec![0, 1, 2]);
4135 let dlq = std::sync::Arc::new(MockSink::new());
4136 let mut cfg = DlqConfig::new(dlq);
4137 cfg.max_failures_per_page = Some(1);
4138 let _ = Pipeline::new(&source, &main)
4139 .with_name("p_budget")
4140 .with_dlq(cfg)
4141 .run()
4142 .await;
4143
4144 let snapshot = snap.snapshot();
4145 let saw = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
4146 k.key().name() == "faucet_sink_dlq_budget_exceeded_total"
4147 && k.key()
4148 .labels()
4149 .any(|l| l.key() == "scope" && l.value() == "per_page")
4150 && matches!(v, DebugValue::Counter(c) if c >= 1)
4151 });
4152 assert!(saw, "faucet_sink_dlq_budget_exceeded_total not emitted");
4153 }
4154
4155 #[tokio::test]
4156 async fn pipeline_run_with_dlq_routes_partial_failures_end_to_end() {
4157 let source = MockSource(vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})]);
4159 let main = PartialSink::new(vec![1]);
4160 let dlq = std::sync::Arc::new(MockSink::new());
4161
4162 let result = Pipeline::new(&source, &main)
4163 .with_dlq(DlqConfig::new(dlq.clone()))
4164 .run()
4165 .await
4166 .unwrap();
4167
4168 assert_eq!(result.records_written, 2);
4169 let stats = result.dlq.unwrap();
4170 assert_eq!(stats.records_dlq, 1);
4171 {
4172 let dlq_records = dlq.0.lock().unwrap();
4173 assert_eq!(dlq_records.len(), 1);
4174 }
4175 }
4176
4177 #[cfg(feature = "quality")]
4180 #[tokio::test]
4181 async fn quality_quarantines_to_dlq_and_writes_survivors() {
4182 use crate::dlq::DlqConfig;
4183 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
4184
4185 let main = Arc::new(MockSink::new());
4186 let dlq_sink = Arc::new(MockSink::new());
4187 let spec = QualitySpec {
4188 record: vec![RecordCheck::NotNull {
4189 field: "id".into(),
4190 treat_missing_as_null: true,
4191 on_failure: OnFailure::Quarantine,
4192 }],
4193 batch: vec![],
4194 };
4195 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4196 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4197 records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
4198 bookmark: None,
4199 })];
4200 let opts = RunStreamOptions::new()
4201 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4202 .with_quality(quality);
4203 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4204 .await
4205 .unwrap();
4206
4207 assert_eq!(result.records_written, 2); assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
4209 let dlq = dlq_sink.written();
4211 assert_eq!(dlq.len(), 1);
4212 assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
4213 assert_eq!(result.dlq.unwrap().records_dlq, 1);
4214 }
4215
4216 #[cfg(feature = "quality")]
4217 #[tokio::test]
4218 #[allow(clippy::await_holding_lock)]
4219 async fn quality_only_page_emits_quality_reason() {
4220 use crate::dlq::DlqConfig;
4224 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
4225 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
4226 use metrics_util::debugging::DebugValue;
4227
4228 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
4229 let snap = snapshotter();
4230
4231 let main = Arc::new(MockSink::new());
4234 let dlq_sink = Arc::new(MockSink::new());
4235 let spec = QualitySpec {
4236 record: vec![RecordCheck::NotNull {
4237 field: "id".into(),
4238 treat_missing_as_null: true,
4239 on_failure: OnFailure::Quarantine,
4240 }],
4241 batch: vec![],
4242 };
4243 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4244 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4245 records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
4246 bookmark: None,
4247 })];
4248 let opts = RunStreamOptions::new()
4249 .with_name("p_quality_reason")
4250 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4251 .with_quality(quality);
4252 let _ = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4253 .await
4254 .unwrap();
4255
4256 let snapshot = snap.snapshot();
4257 let saw_quality_reason = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
4258 k.key().name() == "faucet_sink_dlq_pages_total"
4259 && k.key()
4260 .labels()
4261 .any(|l| l.key() == "pipeline" && l.value() == "p_quality_reason")
4262 && k.key()
4263 .labels()
4264 .any(|l| l.key() == "reason" && l.value() == "quality")
4265 && matches!(v, DebugValue::Counter(c) if c >= 1)
4266 });
4267 assert!(
4268 saw_quality_reason,
4269 "expected faucet_sink_dlq_pages_total with reason=\"quality\""
4270 );
4271 }
4272
4273 #[cfg(feature = "quality")]
4274 #[tokio::test]
4275 async fn quality_abort_fails_run() {
4276 use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
4277 let main = MockSink::new();
4278 let spec = QualitySpec {
4279 record: vec![],
4280 batch: vec![BatchCheck::RowCount {
4281 min: Some(5),
4282 max: None,
4283 on_failure: OnFailure::Abort,
4284 }],
4285 };
4286 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4287 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4288 records: vec![json!({"id": 1})],
4289 bookmark: None,
4290 })];
4291 let opts = RunStreamOptions::new().with_quality(quality);
4292 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4293 assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
4294 }
4295
4296 #[cfg(feature = "quality")]
4297 #[tokio::test]
4298 async fn quality_quarantine_without_dlq_is_rejected() {
4299 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
4300 let main = MockSink::new();
4301 let spec = QualitySpec {
4302 record: vec![RecordCheck::NotNull {
4303 field: "id".into(),
4304 treat_missing_as_null: true,
4305 on_failure: OnFailure::Quarantine,
4306 }],
4307 batch: vec![],
4308 };
4309 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4310 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4311 records: vec![json!({"id": null})],
4312 bookmark: None,
4313 })];
4314 let opts = RunStreamOptions::new().with_quality(quality);
4316 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4317 assert!(matches!(result, Err(FaucetError::Config(_))));
4318 }
4319
4320 #[cfg(feature = "contract")]
4321 fn compiled_contract(on_breach: &str) -> Arc<crate::contract::CompiledContract> {
4322 let spec: crate::contract::ContractSpec = serde_json::from_value(json!({
4323 "version": "1.0.0",
4324 "on_breach": on_breach,
4325 "fields": [{ "name": "id", "type": "integer" }]
4326 }))
4327 .unwrap();
4328 Arc::new(crate::contract::CompiledContract::compile(&spec).unwrap())
4329 }
4330
4331 #[cfg(feature = "contract")]
4332 #[tokio::test]
4333 async fn contract_quarantines_to_dlq_and_writes_survivors() {
4334 use crate::dlq::DlqConfig;
4335 let main = Arc::new(MockSink::new());
4336 let dlq_sink = Arc::new(MockSink::new());
4337 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4338 records: vec![json!({"id": 1}), json!({"id": "bad"}), json!({"id": 3})],
4339 bookmark: None,
4340 })];
4341 let opts = RunStreamOptions::new()
4342 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4343 .with_contract(compiled_contract("quarantine"));
4344 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4345 .await
4346 .unwrap();
4347
4348 assert_eq!(result.records_written, 2);
4349 assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
4350 let dlq = dlq_sink.written();
4351 assert_eq!(dlq.len(), 1);
4352 assert_eq!(dlq[0]["error"]["kind"], "ContractViolation");
4353 assert_eq!(dlq[0]["payload"], json!({"id": "bad"}));
4354 assert_eq!(dlq[0]["record_index"], 1);
4356 assert_eq!(result.dlq.unwrap().records_dlq, 1);
4357 }
4358
4359 #[cfg(feature = "contract")]
4360 #[tokio::test]
4361 async fn contract_fail_aborts_run_and_writes_nothing() {
4362 let main = MockSink::new();
4363 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4364 records: vec![json!({"id": 1}), json!({"id": "bad"})],
4365 bookmark: None,
4366 })];
4367 let opts = RunStreamOptions::new().with_contract(compiled_contract("fail"));
4368 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4369 match result {
4370 Err(FaucetError::ContractViolation { version, message }) => {
4371 assert_eq!(version, "1.0.0");
4372 assert!(message.contains("id"), "message: {message}");
4373 }
4374 other => panic!("expected ContractViolation, got {other:?}"),
4375 }
4376 assert!(
4377 main.written().is_empty(),
4378 "a contract fail must not commit any of the page's records"
4379 );
4380 }
4381
4382 #[cfg(feature = "contract")]
4383 #[tokio::test]
4384 async fn contract_warn_writes_everything() {
4385 let main = MockSink::new();
4386 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4387 records: vec![json!({"id": 1}), json!({"id": "bad"})],
4388 bookmark: None,
4389 })];
4390 let opts = RunStreamOptions::new().with_contract(compiled_contract("warn"));
4391 let result = run_stream(futures::stream::iter(pages), &main, opts)
4392 .await
4393 .unwrap();
4394 assert_eq!(result.records_written, 2);
4395 assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": "bad"})]);
4396 }
4397
4398 #[cfg(feature = "contract")]
4399 #[tokio::test]
4400 async fn contract_quarantine_without_dlq_is_rejected() {
4401 let main = MockSink::new();
4402 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4403 records: vec![json!({"id": 1})],
4404 bookmark: None,
4405 })];
4406 let opts = RunStreamOptions::new().with_contract(compiled_contract("quarantine"));
4408 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4409 assert!(matches!(result, Err(FaucetError::Config(_))));
4410 }
4411
4412 #[cfg(all(feature = "contract", feature = "quality"))]
4413 #[tokio::test]
4414 async fn contract_runs_after_quality_and_shares_dlq() {
4415 use crate::dlq::DlqConfig;
4419 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
4420 let main = Arc::new(MockSink::new());
4421 let dlq_sink = Arc::new(MockSink::new());
4422 let quality = Arc::new(
4423 CompiledQuality::compile(&QualitySpec {
4424 record: vec![RecordCheck::NotNull {
4425 field: "id".into(),
4426 treat_missing_as_null: true,
4427 on_failure: OnFailure::Quarantine,
4428 }],
4429 batch: vec![],
4430 })
4431 .unwrap(),
4432 );
4433 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4434 records: vec![json!({"id": null}), json!({"id": "bad"}), json!({"id": 3})],
4435 bookmark: None,
4436 })];
4437 let opts = RunStreamOptions::new()
4438 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4439 .with_quality(quality)
4440 .with_contract(compiled_contract("quarantine"));
4441 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4442 .await
4443 .unwrap();
4444
4445 assert_eq!(result.records_written, 1);
4446 assert_eq!(main.written(), vec![json!({"id": 3})]);
4447 let dlq = dlq_sink.written();
4448 assert_eq!(dlq.len(), 2);
4449 let kinds: Vec<&str> = dlq
4450 .iter()
4451 .map(|e| e["error"]["kind"].as_str().unwrap())
4452 .collect();
4453 assert!(kinds.contains(&"QualityFailure"), "kinds: {kinds:?}");
4454 assert!(kinds.contains(&"ContractViolation"), "kinds: {kinds:?}");
4455 assert_eq!(result.dlq.unwrap().records_dlq, 2);
4456
4457 let contract_env = dlq
4463 .iter()
4464 .find(|e| e["error"]["kind"] == "ContractViolation")
4465 .expect("a contract envelope");
4466 assert_eq!(
4467 contract_env["record_index"], 1,
4468 "contract record_index must be the original page position: {contract_env}"
4469 );
4470 let quality_env = dlq
4471 .iter()
4472 .find(|e| e["error"]["kind"] == "QualityFailure")
4473 .expect("a quality envelope");
4474 assert_eq!(
4475 quality_env["record_index"], 0,
4476 "null id was at page index 0"
4477 );
4478 }
4479
4480 struct FlakySink {
4483 every: usize,
4484 calls: std::sync::Mutex<Vec<usize>>,
4485 }
4486 impl FlakySink {
4487 fn new(every: usize) -> Self {
4488 Self {
4489 every,
4490 calls: std::sync::Mutex::new(Vec::new()),
4491 }
4492 }
4493 fn call_sizes(&self) -> Vec<usize> {
4494 self.calls.lock().unwrap().clone()
4495 }
4496 }
4497 #[async_trait]
4498 impl Sink for FlakySink {
4499 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4500 Ok(records.len())
4501 }
4502 async fn write_batch_partial(
4503 &self,
4504 records: &[Value],
4505 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4506 self.calls.lock().unwrap().push(records.len());
4507 Ok(records
4508 .iter()
4509 .enumerate()
4510 .map(|(i, _)| {
4511 if (i + 1) % self.every == 0 {
4512 Err(FaucetError::Sink("synthetic".into()))
4513 } else {
4514 Ok(())
4515 }
4516 })
4517 .collect())
4518 }
4519 }
4520
4521 #[tokio::test]
4522 async fn adaptive_shrinks_under_errors_on_dlq_path() {
4523 use crate::adaptive::AdaptiveBatchConfig;
4524 use crate::dlq::{DlqConfig, OnBatchError};
4525 let mk = || StreamPage {
4530 records: (0..400).map(|i| json!({"i": i})).collect(),
4531 bookmark: None,
4532 };
4533 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4534 let sink = FlakySink::new(4); let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
4536 let dlq = DlqConfig {
4537 sink: dlq_sink,
4538 on_batch_error: OnBatchError::Propagate,
4539 max_failures_per_page: None,
4540 max_failures_total: None,
4541 include_original_payload: true,
4542 };
4543 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4544 "enabled": true, "min": 50, "max": 400,
4545 "decrease_factor": 0.5, "cooldown_batches": 0, "error_threshold": 0.1
4546 }))
4547 .unwrap();
4548 let opts = RunStreamOptions::new().with_dlq(dlq).with_adaptive(cfg);
4549 let result = run_stream(stream, &sink, opts).await.unwrap();
4550 assert!(
4557 result.records_written >= 900,
4558 "expected ≥900 written, got {}",
4559 result.records_written
4560 );
4561 let sizes = sink.call_sizes();
4562 assert_eq!(sizes[0], 400, "first chunk is the full page");
4563 assert!(
4564 sizes.last().unwrap() < &400,
4565 "controller should shrink under errors: {sizes:?}"
4566 );
4567 assert!(
4568 result.dlq.unwrap().records_dlq >= 250,
4569 "expected ≥250 DLQ records"
4570 );
4571 }
4572
4573 struct RecordingSink {
4578 calls: std::sync::Mutex<Vec<usize>>,
4579 latency: std::time::Duration,
4580 }
4581 impl RecordingSink {
4582 fn new(latency_ms: u64) -> Self {
4583 Self {
4584 calls: std::sync::Mutex::new(Vec::new()),
4585 latency: std::time::Duration::from_millis(latency_ms),
4586 }
4587 }
4588 fn call_sizes(&self) -> Vec<usize> {
4589 self.calls.lock().unwrap().clone()
4590 }
4591 }
4592 #[async_trait]
4593 impl Sink for RecordingSink {
4594 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4595 tokio::time::sleep(self.latency).await;
4596 self.calls.lock().unwrap().push(records.len());
4597 Ok(records.len())
4598 }
4599 }
4600
4601 #[tokio::test]
4602 async fn adaptive_reslices_non_dlq_page_into_subbatches() {
4603 use crate::adaptive::AdaptiveBatchConfig;
4604 let page = StreamPage {
4605 records: (0..1000).map(|i| json!({ "i": i })).collect(),
4606 bookmark: None,
4607 };
4608 let stream = futures::stream::iter(vec![Ok(page)]);
4609 let sink = RecordingSink::new(0);
4610 let cfg: AdaptiveBatchConfig =
4611 serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
4612 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4613 .await
4614 .unwrap();
4615 assert_eq!(result.records_written, 1000);
4616 assert_eq!(sink.call_sizes(), vec![1000]);
4618 }
4619
4620 #[tokio::test]
4621 async fn adaptive_shrinks_under_latency_target_then_smaller_chunks() {
4622 use crate::adaptive::AdaptiveBatchConfig;
4623 let mk = || StreamPage {
4624 records: (0..400).map(|i| json!({"i": i})).collect(),
4625 bookmark: None,
4626 };
4627 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4628 let sink = RecordingSink::new(50);
4629 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4630 "enabled": true, "min": 50, "max": 400,
4631 "decrease_factor": 0.5, "cooldown_batches": 0,
4632 "target_latency_ms": 10, "latency_window": 1
4633 }))
4634 .unwrap();
4635 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4636 .await
4637 .unwrap();
4638 assert_eq!(result.records_written, 1200);
4639 let sizes = sink.call_sizes();
4640 assert_eq!(sizes[0], 400);
4641 assert!(
4642 sizes.last().unwrap() < &400,
4643 "controller should have shrunk: {sizes:?}"
4644 );
4645 }
4646
4647 #[tokio::test]
4648 #[allow(clippy::await_holding_lock)]
4649 async fn adaptive_emits_batch_size_and_adjustments_metrics() {
4650 use crate::adaptive::AdaptiveBatchConfig;
4653 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
4654 use metrics_util::debugging::DebugValue;
4655
4656 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
4657 let snap = snapshotter();
4658
4659 let mk = || StreamPage {
4665 records: (0..400).map(|i| json!({"i": i})).collect(),
4666 bookmark: None,
4667 };
4668 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4669 let sink = RecordingSink::new(50);
4670 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4671 "enabled": true, "min": 50, "max": 400,
4672 "decrease_factor": 0.5, "cooldown_batches": 0,
4673 "target_latency_ms": 10, "latency_window": 1
4674 }))
4675 .unwrap();
4676
4677 let _ = run_stream(
4678 stream,
4679 &sink,
4680 RunStreamOptions::new()
4681 .with_adaptive(cfg)
4682 .with_name("p")
4683 .with_row("r"),
4684 )
4685 .await
4686 .unwrap();
4687
4688 let snapshot = snap.snapshot();
4689 let mut saw_batch_size = false;
4690 let mut saw_adjustments = false;
4691 for (k, _u, _d, v) in snapshot.into_vec() {
4692 let key = k.key();
4693 let labels = key.labels().collect::<Vec<_>>();
4694 let has = |k: &str, val: &str| labels.iter().any(|l| l.key() == k && l.value() == val);
4695
4696 if key.name() == "faucet_pipeline_adaptive_batch_size"
4697 && has("pipeline", "p")
4698 && has("row", "r")
4699 && matches!(v, DebugValue::Gauge(_))
4700 {
4701 saw_batch_size = true;
4702 }
4703 if key.name() == "faucet_pipeline_adaptive_batch_adjustments_total"
4704 && has("pipeline", "p")
4705 && has("row", "r")
4706 && matches!(v, DebugValue::Counter(c) if c >= 1)
4707 {
4708 saw_adjustments = true;
4709 }
4710 }
4711 assert!(
4712 saw_batch_size,
4713 "expected faucet_pipeline_adaptive_batch_size gauge with pipeline=p, row=r"
4714 );
4715 assert!(
4716 saw_adjustments,
4717 "expected faucet_pipeline_adaptive_batch_adjustments_total counter with pipeline=p, row=r"
4718 );
4719 }
4720
4721 #[tokio::test]
4724 async fn exactly_once_rejects_dlq() {
4725 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4728 let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
4729 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
4730 let opts = eo_opts(store, "k", 0).with_dlq(DlqConfig::new(dlq_sink));
4731 let r = run_stream(
4732 futures::stream::iter(pages),
4733 &IdempotentMockSink::new(),
4734 opts,
4735 )
4736 .await;
4737 assert!(
4738 matches!(&r, Err(FaucetError::Config(m)) if m.contains("not compatible with a DLQ")),
4739 "got: {r:?}"
4740 );
4741 }
4742
4743 #[tokio::test]
4746 async fn exactly_once_writes_unbookmarked_page_at_least_once() {
4747 let sink = IdempotentMockSink::new();
4752 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4753 let pages = vec![Ok(StreamPage {
4754 records: vec![json!({"id": 1}), json!({"id": 2})],
4755 bookmark: None,
4756 })];
4757 let r = run_stream(
4758 futures::stream::iter(pages),
4759 &sink,
4760 eo_opts(store.clone(), "k", 0),
4761 )
4762 .await
4763 .unwrap();
4764 assert_eq!(r.records_written, 2);
4765 assert_eq!(r.bookmark, None);
4766 assert_eq!(sink.last_committed_token("k").await.unwrap(), None);
4769 assert!(store.get("k").await.unwrap().is_none());
4771 assert_eq!(sink.rows(), vec![json!({"id": 1}), json!({"id": 2})]);
4772 }
4773
4774 #[tokio::test]
4777 async fn dlq_with_bookmark_persists_after_routing_failures() {
4778 let main = PartialSink::new(vec![1]); let dlq = std::sync::Arc::new(MockSink::new());
4783 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4784 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4785 records: vec![json!({"i": 0}), json!({"i": 1})],
4786 bookmark: Some(json!("ckpt")),
4787 })];
4788 let result = run_stream(
4789 futures::stream::iter(pages),
4790 &main,
4791 RunStreamOptions::new()
4792 .with_dlq(DlqConfig::new(dlq.clone()))
4793 .with_state(Arc::clone(&store), "k"),
4794 )
4795 .await
4796 .unwrap();
4797
4798 assert_eq!(result.records_written, 1); assert_eq!(result.bookmark, Some(json!("ckpt")));
4800 assert_eq!(store.get("k").await.unwrap(), Some(json!("ckpt")));
4802 let envelopes = dlq.0.lock().unwrap();
4804 assert_eq!(envelopes.len(), 1);
4805 assert_eq!(envelopes[0]["payload"]["i"], 1);
4806 }
4807
4808 #[tokio::test]
4811 async fn exactly_once_resumes_sequence_from_start_seq() {
4812 let sink = IdempotentMockSink::new();
4816 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4817 let pages = vec![Ok(StreamPage {
4818 records: vec![json!({"id": 9})],
4819 bookmark: Some(json!("bm-after-resume")),
4820 })];
4821 let r = run_stream(
4822 futures::stream::iter(pages),
4823 &sink,
4824 eo_opts(store.clone(), "eo_key", 7),
4825 )
4826 .await
4827 .unwrap();
4828
4829 assert_eq!(r.records_written, 1);
4830 let (bm, seq) =
4831 crate::idempotency::unwrap_state(&store.get("eo_key").await.unwrap().unwrap());
4832 assert_eq!(bm, Some(json!("bm-after-resume")));
4833 assert_eq!(seq, 8, "sequence resumes at start_seq + 1");
4834 let token = sink.last_committed_token("eo_key").await.unwrap().unwrap();
4835 assert_eq!(
4836 crate::idempotency::parse_token_parts(&token),
4837 Some((8, Some(json!("bm-after-resume"))))
4838 );
4839 }
4840
4841 #[cfg(feature = "quality")]
4844 #[tokio::test]
4845 async fn pipeline_run_with_quality_aborts_on_failed_batch_check() {
4846 use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
4847 let source = MockSource(vec![json!({"id": 1})]);
4848 let main = MockSink::new();
4849 let spec = QualitySpec {
4850 record: vec![],
4851 batch: vec![BatchCheck::RowCount {
4852 min: Some(5),
4853 max: None,
4854 on_failure: OnFailure::Abort,
4855 }],
4856 };
4857 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4858 let result = Pipeline::new(&source, &main)
4859 .with_quality(quality)
4860 .run()
4861 .await;
4862 assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
4863 assert!(main.written().is_empty());
4865 }
4866
4867 #[tokio::test]
4870 async fn pipeline_run_with_adaptive_reslices_page() {
4871 use crate::adaptive::AdaptiveBatchConfig;
4872 let source = MockSource((0..1000).map(|i| json!({ "i": i })).collect());
4873 let sink = MockSink::new();
4874 let cfg: AdaptiveBatchConfig =
4875 serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
4876 let result = Pipeline::new(&source, &sink)
4877 .with_adaptive(cfg)
4878 .run()
4879 .await
4880 .unwrap();
4881 assert_eq!(result.records_written, 1000);
4882 assert_eq!(sink.written().len(), 1000);
4883 }
4884
4885 #[tokio::test]
4888 async fn adaptive_noop_sink_name_is_handled() {
4889 use crate::adaptive::AdaptiveBatchConfig;
4890
4891 struct JsonlNamedSink(std::sync::Mutex<usize>);
4894 #[async_trait]
4895 impl Sink for JsonlNamedSink {
4896 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4897 *self.0.lock().unwrap() += records.len();
4898 Ok(records.len())
4899 }
4900 fn connector_name(&self) -> &'static str {
4901 "jsonl"
4902 }
4903 }
4904
4905 let page = StreamPage {
4906 records: (0..10).map(|i| json!({ "i": i })).collect(),
4907 bookmark: None,
4908 };
4909 let stream = futures::stream::iter(vec![Ok(page)]);
4910 let sink = JsonlNamedSink(std::sync::Mutex::new(0));
4911 let cfg: AdaptiveBatchConfig =
4912 serde_json::from_value(json!({"enabled": true, "min": 5, "max": 10})).unwrap();
4913 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4914 .await
4915 .unwrap();
4916 assert_eq!(result.records_written, 10);
4917 assert_eq!(*sink.0.lock().unwrap(), 10);
4918 }
4919
4920 #[tokio::test]
4921 async fn resilience_retries_transient_sink_write() {
4922 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4923 use std::sync::Arc;
4924 use std::sync::atomic::{AtomicU32, Ordering};
4925 use std::time::Duration;
4926
4927 struct TransientFlakySink {
4928 attempts: Arc<AtomicU32>,
4929 written: Arc<AtomicU32>,
4930 }
4931 #[async_trait::async_trait]
4932 impl Sink for TransientFlakySink {
4933 async fn write_batch(
4934 &self,
4935 records: &[serde_json::Value],
4936 ) -> Result<usize, FaucetError> {
4937 let n = self.attempts.fetch_add(1, Ordering::SeqCst);
4938 if n < 2 {
4939 return Err(FaucetError::HttpStatus {
4940 status: 503,
4941 url: "u".into(),
4942 body: "".into(),
4943 });
4944 }
4945 self.written
4946 .fetch_add(records.len() as u32, Ordering::SeqCst);
4947 Ok(records.len())
4948 }
4949 async fn flush(&self) -> Result<(), FaucetError> {
4950 Ok(())
4951 }
4952 fn supports_idempotent_writes(&self) -> bool {
4954 true
4955 }
4956 }
4957
4958 let attempts = Arc::new(AtomicU32::new(0));
4959 let written = Arc::new(AtomicU32::new(0));
4960 let sink = TransientFlakySink {
4961 attempts: attempts.clone(),
4962 written: written.clone(),
4963 };
4964 let pages = futures::stream::iter(vec![Ok(StreamPage {
4965 records: vec![serde_json::json!({"a": 1})],
4966 bookmark: None,
4967 })]);
4968 let policy = ResiliencePolicy {
4969 retry: RetryPolicy {
4970 max_attempts: 5,
4971 backoff: BackoffKind::None,
4972 base: Duration::ZERO,
4973 max: Duration::ZERO,
4974 jitter: false,
4975 ..RetryPolicy::default()
4976 },
4977 ..ResiliencePolicy::default()
4978 };
4979 let res = run_stream(
4980 pages,
4981 &sink,
4982 RunStreamOptions::new().with_resilience(policy),
4983 )
4984 .await
4985 .unwrap();
4986 assert_eq!(written.load(Ordering::SeqCst), 1);
4987 assert_eq!(attempts.load(Ordering::SeqCst), 3);
4988 assert_eq!(res.records_written, 1);
4989 }
4990
4991 #[tokio::test]
4992 async fn resilience_does_not_retry_non_idempotent_write_batch() {
4993 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4997 use std::sync::Arc;
4998 use std::sync::atomic::{AtomicU32, Ordering};
4999 use std::time::Duration;
5000
5001 struct NonIdempotentFlakySink {
5002 attempts: Arc<AtomicU32>,
5003 }
5004 #[async_trait::async_trait]
5005 impl Sink for NonIdempotentFlakySink {
5006 async fn write_batch(&self, _r: &[Value]) -> Result<usize, FaucetError> {
5007 self.attempts.fetch_add(1, Ordering::SeqCst);
5008 Err(FaucetError::HttpStatus {
5009 status: 503,
5010 url: "u".into(),
5011 body: "".into(),
5012 })
5013 }
5014 async fn flush(&self) -> Result<(), FaucetError> {
5015 Ok(())
5016 }
5017 }
5019
5020 let attempts = Arc::new(AtomicU32::new(0));
5021 let sink = NonIdempotentFlakySink {
5022 attempts: attempts.clone(),
5023 };
5024 let pages = futures::stream::iter(vec![Ok(StreamPage {
5025 records: vec![json!({"a": 1})],
5026 bookmark: None,
5027 })]);
5028 let policy = ResiliencePolicy {
5029 retry: RetryPolicy {
5030 max_attempts: 5,
5031 backoff: BackoffKind::None,
5032 base: Duration::ZERO,
5033 max: Duration::ZERO,
5034 jitter: false,
5035 ..RetryPolicy::default()
5036 },
5037 ..ResiliencePolicy::default()
5038 };
5039 let err = run_stream(
5040 pages,
5041 &sink,
5042 RunStreamOptions::new().with_resilience(policy),
5043 )
5044 .await
5045 .unwrap_err();
5046 assert!(matches!(err, FaucetError::HttpStatus { status: 503, .. }));
5047 assert_eq!(
5048 attempts.load(Ordering::SeqCst),
5049 1,
5050 "non-idempotent write_batch must be attempted exactly once (no retry)"
5051 );
5052 }
5053
5054 #[tokio::test]
5055 async fn resilience_circuit_opens_after_consecutive_failed_pages() {
5056 use crate::dlq::{DlqConfig, OnBatchError};
5057 use crate::resilience::{BackoffKind, CircuitBreakerConfig, ResiliencePolicy, RetryPolicy};
5058 use std::sync::Arc;
5059 use std::time::Duration;
5060
5061 struct DeadSink;
5063 #[async_trait]
5064 impl Sink for DeadSink {
5065 async fn write_batch(&self, _r: &[Value]) -> Result<usize, FaucetError> {
5066 Err(FaucetError::Sink("down".into()))
5067 }
5068 async fn flush(&self) -> Result<(), FaucetError> {
5069 Ok(())
5070 }
5071 }
5072 struct NullSink;
5074 #[async_trait]
5075 impl Sink for NullSink {
5076 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5077 Ok(r.len())
5078 }
5079 async fn flush(&self) -> Result<(), FaucetError> {
5080 Ok(())
5081 }
5082 }
5083
5084 let pages = futures::stream::iter((0..10).map(|i| {
5085 Ok(StreamPage {
5086 records: vec![json!({"i": i})],
5087 bookmark: None,
5088 })
5089 }));
5090 let dlq = DlqConfig {
5091 on_batch_error: OnBatchError::DlqAll,
5092 ..DlqConfig::new(Arc::new(NullSink))
5093 };
5094 let policy = ResiliencePolicy {
5095 retry: RetryPolicy {
5096 max_attempts: 1,
5097 backoff: BackoffKind::None,
5098 base: Duration::ZERO,
5099 max: Duration::ZERO,
5100 jitter: false,
5101 ..RetryPolicy::default()
5102 },
5103 circuit_breaker: Some(CircuitBreakerConfig {
5104 consecutive_failures: 3,
5105 cooldown: Duration::from_secs(60),
5106 }),
5107 poison: None,
5108 };
5109 let err = run_stream(
5110 pages,
5111 &DeadSink,
5112 RunStreamOptions::new()
5113 .with_dlq(dlq)
5114 .with_resilience(policy),
5115 )
5116 .await
5117 .unwrap_err();
5118 assert!(
5119 matches!(err, FaucetError::CircuitOpen { failures: 3, .. }),
5120 "got {err:?}"
5121 );
5122 }
5123
5124 #[tokio::test]
5125 async fn resilience_poison_retries_then_dlqs_failing_row() {
5126 use crate::dlq::DlqConfig;
5127 use crate::resilience::{
5128 BackoffKind, PoisonAction, PoisonPolicy, ResiliencePolicy, RetryPolicy,
5129 };
5130 use std::sync::{Arc, Mutex};
5131 use std::time::Duration;
5132
5133 struct PickySink {
5136 bad_attempts: Arc<Mutex<u32>>,
5137 }
5138 #[async_trait]
5139 impl Sink for PickySink {
5140 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5141 Ok(r.len())
5142 }
5143 async fn write_batch_partial(
5144 &self,
5145 records: &[Value],
5146 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
5147 Ok(records
5148 .iter()
5149 .map(|rec| {
5150 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
5151 *self.bad_attempts.lock().unwrap() += 1;
5152 Err(FaucetError::HttpStatus {
5153 status: 503,
5154 url: "u".into(),
5155 body: "".into(),
5156 })
5157 } else {
5158 Ok(())
5159 }
5160 })
5161 .collect())
5162 }
5163 async fn flush(&self) -> Result<(), FaucetError> {
5164 Ok(())
5165 }
5166 }
5167 struct CaptureSink(Arc<Mutex<Vec<Value>>>);
5168 #[async_trait]
5169 impl Sink for CaptureSink {
5170 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5171 self.0.lock().unwrap().extend_from_slice(r);
5172 Ok(r.len())
5173 }
5174 async fn flush(&self) -> Result<(), FaucetError> {
5175 Ok(())
5176 }
5177 }
5178
5179 let captured = Arc::new(Mutex::new(Vec::new()));
5180 let bad_attempts = Arc::new(Mutex::new(0u32));
5181 let sink = PickySink {
5182 bad_attempts: bad_attempts.clone(),
5183 };
5184 let pages = futures::stream::iter(vec![Ok(StreamPage {
5185 records: vec![json!({"ok": 1}), json!({"bad": true})],
5186 bookmark: None,
5187 })]);
5188 let policy = ResiliencePolicy {
5189 retry: RetryPolicy {
5190 max_attempts: 1,
5191 backoff: BackoffKind::None,
5192 base: Duration::ZERO,
5193 max: Duration::ZERO,
5194 jitter: false,
5195 ..RetryPolicy::default()
5196 },
5197 circuit_breaker: None,
5198 poison: Some(PoisonPolicy {
5199 max_row_attempts: 3,
5200 action: PoisonAction::Dlq,
5201 }),
5202 };
5203 let res = run_stream(
5204 pages,
5205 &sink,
5206 RunStreamOptions::new()
5207 .with_dlq(DlqConfig::new(Arc::new(CaptureSink(captured.clone()))))
5208 .with_resilience(policy),
5209 )
5210 .await
5211 .unwrap();
5212
5213 assert_eq!(
5214 *bad_attempts.lock().unwrap(),
5215 3,
5216 "bad row tried max_row_attempts times"
5217 );
5218 assert_eq!(res.records_written, 1, "the ok row");
5219 let dlq = captured.lock().unwrap();
5220 assert_eq!(dlq.len(), 1, "one row to DLQ");
5221 assert_eq!(dlq[0]["payload"]["bad"], json!(true));
5222 }
5223
5224 #[tokio::test]
5225 async fn poison_loop_does_not_nest_resilience_retry_on_subset_resubmit() {
5226 use crate::dlq::DlqConfig;
5234 use crate::resilience::{
5235 BackoffKind, PoisonAction, PoisonPolicy, ResiliencePolicy, RetryPolicy,
5236 };
5237 use std::sync::{Arc, Mutex};
5238 use std::time::Duration;
5239
5240 struct OuterErrOnSubsetSink {
5241 subset_calls: Arc<Mutex<u32>>,
5242 }
5243 #[async_trait]
5244 impl Sink for OuterErrOnSubsetSink {
5245 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5246 Ok(r.len())
5247 }
5248 async fn write_batch_partial(
5249 &self,
5250 records: &[Value],
5251 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
5252 if records.len() == 1 {
5254 *self.subset_calls.lock().unwrap() += 1;
5255 return Err(FaucetError::HttpStatus {
5256 status: 503,
5257 url: "u".into(),
5258 body: "".into(),
5259 });
5260 }
5261 Ok(records
5262 .iter()
5263 .map(|rec| {
5264 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
5265 Err(FaucetError::HttpStatus {
5266 status: 503,
5267 url: "u".into(),
5268 body: "".into(),
5269 })
5270 } else {
5271 Ok(())
5272 }
5273 })
5274 .collect())
5275 }
5276 async fn flush(&self) -> Result<(), FaucetError> {
5277 Ok(())
5278 }
5279 fn supports_idempotent_writes(&self) -> bool {
5282 true
5283 }
5284 }
5285 struct CaptureSink(Arc<Mutex<Vec<Value>>>);
5286 #[async_trait]
5287 impl Sink for CaptureSink {
5288 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5289 self.0.lock().unwrap().extend_from_slice(r);
5290 Ok(r.len())
5291 }
5292 async fn flush(&self) -> Result<(), FaucetError> {
5293 Ok(())
5294 }
5295 }
5296
5297 let subset_calls = Arc::new(Mutex::new(0u32));
5298 let sink = OuterErrOnSubsetSink {
5299 subset_calls: subset_calls.clone(),
5300 };
5301 let pages = futures::stream::iter(vec![Ok(StreamPage {
5302 records: vec![json!({"ok": 1}), json!({"bad": true})],
5303 bookmark: None,
5304 })]);
5305 let policy = ResiliencePolicy {
5306 retry: RetryPolicy {
5307 max_attempts: 4, backoff: BackoffKind::None,
5309 base: Duration::ZERO,
5310 max: Duration::ZERO,
5311 jitter: false,
5312 ..RetryPolicy::default()
5313 },
5314 circuit_breaker: None,
5315 poison: Some(PoisonPolicy {
5316 max_row_attempts: 3,
5317 action: PoisonAction::Dlq,
5318 }),
5319 };
5320 let res = run_stream(
5321 pages,
5322 &sink,
5323 RunStreamOptions::new()
5324 .with_dlq(DlqConfig::new(Arc::new(CaptureSink(Arc::new(Mutex::new(
5325 Vec::new(),
5326 ))))))
5327 .with_resilience(policy),
5328 )
5329 .await;
5330
5331 assert!(matches!(
5333 res,
5334 Err(FaucetError::HttpStatus { status: 503, .. })
5335 ));
5336 assert_eq!(
5339 *subset_calls.lock().unwrap(),
5340 1,
5341 "poison subset resubmit must be bare (no nested resilience retry)"
5342 );
5343 }
5344
5345 #[tokio::test]
5346 async fn resilience_retries_transient_flush_and_state_on_dlq_path() {
5347 use crate::dlq::DlqConfig;
5354 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
5355 use crate::state::StateStore;
5356 use std::sync::Arc;
5357 use std::sync::atomic::{AtomicU32, Ordering};
5358 use std::time::Duration;
5359
5360 fn transient_503() -> FaucetError {
5361 FaucetError::HttpStatus {
5362 status: 503,
5363 url: "u".into(),
5364 body: "".into(),
5365 }
5366 }
5367
5368 struct FlakyFlushSink {
5371 flush_attempts: Arc<AtomicU32>,
5372 }
5373 #[async_trait]
5374 impl Sink for FlakyFlushSink {
5375 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5376 Ok(r.len())
5377 }
5378 async fn write_batch_partial(
5379 &self,
5380 records: &[Value],
5381 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
5382 Ok(records
5383 .iter()
5384 .map(|rec| {
5385 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
5386 Err(transient_503())
5387 } else {
5388 Ok(())
5389 }
5390 })
5391 .collect())
5392 }
5393 async fn flush(&self) -> Result<(), FaucetError> {
5394 let n = self.flush_attempts.fetch_add(1, Ordering::SeqCst);
5395 if n < 2 { Err(transient_503()) } else { Ok(()) }
5396 }
5397 }
5398 struct NullSink;
5399 #[async_trait]
5400 impl Sink for NullSink {
5401 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5402 Ok(r.len())
5403 }
5404 async fn flush(&self) -> Result<(), FaucetError> {
5405 Ok(())
5406 }
5407 }
5408
5409 struct FlakyStore {
5411 put_attempts: Arc<AtomicU32>,
5412 value: Arc<std::sync::Mutex<Option<Value>>>,
5413 }
5414 #[async_trait]
5415 impl StateStore for FlakyStore {
5416 async fn get(&self, _key: &str) -> Result<Option<Value>, FaucetError> {
5417 Ok(self.value.lock().unwrap().clone())
5418 }
5419 async fn put(&self, _key: &str, value: &Value) -> Result<(), FaucetError> {
5420 let n = self.put_attempts.fetch_add(1, Ordering::SeqCst);
5421 if n < 2 {
5422 return Err(transient_503());
5423 }
5424 *self.value.lock().unwrap() = Some(value.clone());
5425 Ok(())
5426 }
5427 async fn delete(&self, _key: &str) -> Result<(), FaucetError> {
5428 Ok(())
5429 }
5430 }
5431
5432 let flush_attempts = Arc::new(AtomicU32::new(0));
5433 let put_attempts = Arc::new(AtomicU32::new(0));
5434 let stored = Arc::new(std::sync::Mutex::new(None));
5435 let sink = FlakyFlushSink {
5436 flush_attempts: flush_attempts.clone(),
5437 };
5438 let store: Arc<dyn StateStore> = Arc::new(FlakyStore {
5439 put_attempts: put_attempts.clone(),
5440 value: stored.clone(),
5441 });
5442 let pages = futures::stream::iter(vec![Ok(StreamPage {
5443 records: vec![json!({"ok": 1}), json!({"bad": true})],
5444 bookmark: Some(json!({"cursor": 42})),
5445 })]);
5446 let policy = ResiliencePolicy {
5447 retry: RetryPolicy {
5448 max_attempts: 5,
5449 backoff: BackoffKind::None,
5450 base: Duration::ZERO,
5451 max: Duration::ZERO,
5452 jitter: false,
5453 ..RetryPolicy::default()
5454 },
5455 ..ResiliencePolicy::default()
5456 };
5457 let res = run_stream(
5458 pages,
5459 &sink,
5460 RunStreamOptions::new()
5461 .with_dlq(DlqConfig::new(Arc::new(NullSink)))
5462 .with_state(store, "k")
5463 .with_resilience(policy),
5464 )
5465 .await
5466 .unwrap();
5467
5468 assert_eq!(
5472 flush_attempts.load(Ordering::SeqCst),
5473 4,
5474 "page-gate flush retried past two transient failures (3) + 1 final flush"
5475 );
5476 assert_eq!(
5478 put_attempts.load(Ordering::SeqCst),
5479 3,
5480 "state put retried past two transient failures"
5481 );
5482 assert_eq!(*stored.lock().unwrap(), Some(json!({"cursor": 42})));
5483 assert_eq!(res.records_written, 1, "the ok row");
5484 }
5485
5486 #[tokio::test]
5487 #[allow(clippy::await_holding_lock)]
5488 async fn resilience_emits_retries_total_with_op_and_class_labels() {
5489 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
5493 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
5494 use metrics_util::debugging::DebugValue;
5495 use std::sync::Arc;
5496 use std::sync::atomic::{AtomicU32, Ordering};
5497 use std::time::Duration;
5498
5499 struct RetryProbeSink {
5500 attempts: Arc<AtomicU32>,
5501 }
5502 #[async_trait]
5503 impl Sink for RetryProbeSink {
5504 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5505 let n = self.attempts.fetch_add(1, Ordering::SeqCst);
5506 if n < 2 {
5507 return Err(FaucetError::HttpStatus {
5508 status: 503,
5509 url: "u".into(),
5510 body: "".into(),
5511 });
5512 }
5513 Ok(r.len())
5514 }
5515 async fn flush(&self) -> Result<(), FaucetError> {
5516 Ok(())
5517 }
5518 fn connector_name(&self) -> &'static str {
5521 "retry-probe"
5522 }
5523 fn supports_idempotent_writes(&self) -> bool {
5525 true
5526 }
5527 }
5528
5529 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
5530 let snap = snapshotter();
5531
5532 let attempts = Arc::new(AtomicU32::new(0));
5533 let sink = RetryProbeSink {
5534 attempts: attempts.clone(),
5535 };
5536 let pages = futures::stream::iter(vec![Ok(StreamPage {
5537 records: vec![json!({"a": 1})],
5538 bookmark: None,
5539 })]);
5540 let policy = ResiliencePolicy {
5541 retry: RetryPolicy {
5542 max_attempts: 5,
5543 backoff: BackoffKind::None,
5544 base: Duration::ZERO,
5545 max: Duration::ZERO,
5546 jitter: false,
5547 ..RetryPolicy::default()
5548 },
5549 ..ResiliencePolicy::default()
5550 };
5551 run_stream(
5552 pages,
5553 &sink,
5554 RunStreamOptions::new()
5555 .with_name("retry-metrics-pipeline")
5556 .with_resilience(policy),
5557 )
5558 .await
5559 .unwrap();
5560 assert_eq!(attempts.load(Ordering::SeqCst), 3);
5561
5562 let snapshot = snap.snapshot();
5563 let retries: u64 = snapshot
5564 .into_vec()
5565 .into_iter()
5566 .filter_map(|(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
5567 if key.key().name() == "faucet_resilience_retries_total"
5568 && key.key().labels().any(|l: &metrics::Label| {
5569 l.key() == "pipeline" && l.value() == "retry-metrics-pipeline"
5570 })
5571 && key
5572 .key()
5573 .labels()
5574 .any(|l: &metrics::Label| l.key() == "op" && l.value() == "sink_write")
5575 && key
5576 .key()
5577 .labels()
5578 .any(|l: &metrics::Label| l.key() == "class" && l.value() == "http_5xx")
5579 && let DebugValue::Counter(c) = v
5580 {
5581 Some(c)
5582 } else {
5583 None
5584 }
5585 })
5586 .sum();
5587 assert_eq!(
5588 retries, 2,
5589 "expected 2 retries (2 transient 503s) counted with op=sink_write, class=http_5xx"
5590 );
5591 }
5592
5593 struct SchemaSink {
5597 schema: Value,
5598 written: std::sync::Mutex<Vec<Value>>,
5599 evolutions: std::sync::Mutex<Vec<crate::drift::SchemaEvolution>>,
5600 evolvable: bool,
5601 }
5602 impl SchemaSink {
5603 fn new(schema: Value, evolvable: bool) -> Self {
5604 Self {
5605 schema,
5606 written: std::sync::Mutex::new(Vec::new()),
5607 evolutions: std::sync::Mutex::new(Vec::new()),
5608 evolvable,
5609 }
5610 }
5611 fn written(&self) -> Vec<Value> {
5612 self.written.lock().unwrap().clone()
5613 }
5614 }
5615 #[async_trait]
5616 impl Sink for SchemaSink {
5617 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
5618 self.written.lock().unwrap().extend(records.iter().cloned());
5619 Ok(records.len())
5620 }
5621 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
5622 Ok(Some(self.schema.clone()))
5623 }
5624 fn supports_schema_evolution(&self) -> bool {
5625 self.evolvable
5626 }
5627 async fn evolve_schema(
5628 &self,
5629 evo: &crate::drift::SchemaEvolution,
5630 ) -> Result<(), FaucetError> {
5631 if !self.evolvable {
5632 return Err(FaucetError::Sink("not evolvable".into()));
5633 }
5634 self.evolutions.lock().unwrap().push(evo.clone());
5635 Ok(())
5636 }
5637 }
5638
5639 fn drift_opts(policy: crate::drift::SchemaDriftPolicy) -> RunStreamOptions {
5640 RunStreamOptions::new().with_schema_drift(policy)
5641 }
5642
5643 fn one_page(
5644 records: Vec<Value>,
5645 ) -> impl futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Unpin {
5646 Box::pin(futures::stream::iter(vec![Ok(StreamPage {
5647 records,
5648 bookmark: None,
5649 })]))
5650 }
5651
5652 #[tokio::test]
5653 async fn drift_warn_writes_unchanged() {
5654 let sink = SchemaSink::new(
5655 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5656 false,
5657 );
5658 let policy = crate::drift::SchemaDriftPolicy {
5659 on_drift: crate::drift::OnDrift::Warn,
5660 allow_widening: true,
5661 on_incompatible: crate::drift::OnIncompatible::Fail,
5662 relax_nullability_on_missing: false,
5663 };
5664 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5665 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5666 assert_eq!(res.records_written, 1);
5667 assert_eq!(sink.written()[0], json!({"id": 1, "email": "a@x.com"}));
5669 }
5670
5671 #[tokio::test]
5672 async fn drift_ignore_strips_unknown_fields() {
5673 let sink = SchemaSink::new(
5674 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5675 false,
5676 );
5677 let policy = crate::drift::SchemaDriftPolicy {
5678 on_drift: crate::drift::OnDrift::Ignore,
5679 allow_widening: true,
5680 on_incompatible: crate::drift::OnIncompatible::Fail,
5681 relax_nullability_on_missing: false,
5682 };
5683 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5684 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5685 assert_eq!(res.records_written, 1);
5686 assert_eq!(
5687 sink.written()[0],
5688 json!({"id": 1}),
5689 "email must be stripped"
5690 );
5691 }
5692
5693 #[tokio::test]
5694 async fn drift_fail_raises_schema_drift() {
5695 let sink = SchemaSink::new(
5696 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5697 false,
5698 );
5699 let policy = crate::drift::SchemaDriftPolicy {
5700 on_drift: crate::drift::OnDrift::Fail,
5701 allow_widening: true,
5702 on_incompatible: crate::drift::OnIncompatible::Fail,
5703 relax_nullability_on_missing: false,
5704 };
5705 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5706 let err = run_stream(pages, &sink, drift_opts(policy))
5707 .await
5708 .unwrap_err();
5709 assert!(matches!(err, FaucetError::SchemaDrift { .. }));
5710 }
5711
5712 #[tokio::test]
5713 async fn drift_evolve_calls_sink_then_writes() {
5714 let sink = SchemaSink::new(
5715 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5716 true,
5717 );
5718 let policy = crate::drift::SchemaDriftPolicy {
5719 on_drift: crate::drift::OnDrift::Evolve,
5720 allow_widening: true,
5721 on_incompatible: crate::drift::OnIncompatible::Fail,
5722 relax_nullability_on_missing: false,
5723 };
5724 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5725 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5726 assert_eq!(res.records_written, 1);
5727 let evos = sink.evolutions.lock().unwrap();
5728 assert_eq!(evos.len(), 1);
5729 assert_eq!(evos[0].additions.len(), 1);
5730 assert_eq!(evos[0].additions[0].name, "email");
5731 assert_eq!(sink.written()[0], json!({"id": 1, "email": "a@x.com"}));
5733 }
5734
5735 #[tokio::test]
5736 async fn drift_evolve_does_not_relax_not_null_for_merely_absent_column() {
5737 let sink = SchemaSink::new(
5742 json!({"type":"object","properties":{
5743 "id":{"type":"integer"},
5744 "legacy":{"type":"string"}
5745 }}),
5746 true,
5747 );
5748 let policy = crate::drift::SchemaDriftPolicy {
5749 on_drift: crate::drift::OnDrift::Evolve,
5750 allow_widening: true,
5751 on_incompatible: crate::drift::OnIncompatible::Fail,
5752 relax_nullability_on_missing: false,
5753 };
5754 let pages = one_page(vec![json!({"id": 1})]);
5755 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5756 assert_eq!(res.records_written, 1);
5757 assert!(
5758 sink.evolutions.lock().unwrap().is_empty(),
5759 "NOT NULL must not be relaxed for a merely-absent column"
5760 );
5761 }
5762
5763 #[tokio::test]
5764 async fn drift_evolve_relaxes_absent_column_only_with_opt_in() {
5765 let sink = SchemaSink::new(
5767 json!({"type":"object","properties":{
5768 "id":{"type":"integer"},
5769 "legacy":{"type":"string"}
5770 }}),
5771 true,
5772 );
5773 let policy = crate::drift::SchemaDriftPolicy {
5774 on_drift: crate::drift::OnDrift::Evolve,
5775 allow_widening: true,
5776 on_incompatible: crate::drift::OnIncompatible::Fail,
5777 relax_nullability_on_missing: true,
5778 };
5779 let pages = one_page(vec![json!({"id": 1})]);
5780 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5781 assert_eq!(res.records_written, 1);
5782 let evos = sink.evolutions.lock().unwrap();
5783 assert_eq!(evos.len(), 1);
5784 assert_eq!(evos[0].relax_nullability, vec!["legacy".to_string()]);
5785 }
5786
5787 #[tokio::test]
5788 async fn drift_inert_when_sink_reports_no_schema() {
5789 let sink = MockSink::new();
5791 let policy = crate::drift::SchemaDriftPolicy {
5792 on_drift: crate::drift::OnDrift::Fail, allow_widening: true,
5794 on_incompatible: crate::drift::OnIncompatible::Fail,
5795 relax_nullability_on_missing: false,
5796 };
5797 let pages = one_page(vec![json!({"id": 1, "anything": true})]);
5798 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5799 assert_eq!(res.records_written, 1);
5800 }
5801
5802 #[tokio::test]
5803 async fn drift_quarantine_routes_drift_rows_to_dlq() {
5804 let sink = SchemaSink::new(
5805 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5806 false,
5807 );
5808 let dlq_sink = std::sync::Arc::new(MockSink::new());
5809 let policy = crate::drift::SchemaDriftPolicy {
5810 on_drift: crate::drift::OnDrift::Quarantine,
5811 allow_widening: true,
5812 on_incompatible: crate::drift::OnIncompatible::Fail,
5813 relax_nullability_on_missing: false,
5814 };
5815 let pages = one_page(vec![
5816 json!({"id": 1}), json!({"id": 2, "email": "x"}), ]);
5819 let opts = RunStreamOptions::new()
5820 .with_schema_drift(policy)
5821 .with_dlq(crate::dlq::DlqConfig::new(dlq_sink.clone()));
5822 let res = run_stream(pages, &sink, opts).await.unwrap();
5823 assert_eq!(res.records_written, 1, "only the conforming row is written");
5824 assert_eq!(sink.written(), vec![json!({"id": 1})]);
5825 let dlq = dlq_sink.written();
5827 assert_eq!(dlq.len(), 1);
5828 assert_eq!(dlq[0]["payload"], json!({"id": 2, "email": "x"}));
5829 assert_eq!(dlq[0]["error"]["kind"], "SchemaDrift");
5830 }
5831
5832 #[test]
5833 fn quarantine_drift_rows_covers_widening_and_droppable_required() {
5834 use crate::drift::{ColumnChange, SchemaDiff};
5835 let diff = SchemaDiff {
5837 additions: vec![],
5838 widenings: vec![ColumnChange {
5839 name: "amount".into(),
5840 from: Some(json!({"type":"integer"})),
5841 to: json!({"type":"number"}),
5842 }],
5843 incompatible: vec![],
5844 droppable_required: vec!["legacy".into()],
5845 };
5846 let records = vec![
5847 json!({"id": 1, "amount": 1.5, "legacy": "x"}), json!({"id": 2, "amount": 7, "legacy": "y"}), json!({"id": 3, "legacy": "z"}), json!({"id": 4}), ];
5852 let page_indices = vec![10, 11, 12, 13];
5855 let (kept, env) = quarantine_drift_rows(&diff, records, &page_indices, "sink", "pl", "");
5856 assert_eq!(kept, vec![json!({"id": 3, "legacy": "z"})]);
5857 assert_eq!(
5858 env.len(),
5859 3,
5860 "widening rows + the missing-required row quarantined"
5861 );
5862 let indices: Vec<i64> = env
5864 .iter()
5865 .map(|e| e["record_index"].as_i64().unwrap())
5866 .collect();
5867 assert_eq!(indices, vec![10, 11, 13]);
5868 }
5869
5870 #[tokio::test]
5871 async fn drift_quarantine_without_dlq_is_rejected() {
5872 let sink = SchemaSink::new(json!({"type":"object","properties":{}}), false);
5873 let policy = crate::drift::SchemaDriftPolicy {
5874 on_drift: crate::drift::OnDrift::Quarantine,
5875 allow_widening: true,
5876 on_incompatible: crate::drift::OnIncompatible::Fail,
5877 relax_nullability_on_missing: false,
5878 };
5879 let pages = one_page(vec![json!({"id": 1})]);
5880 let err = run_stream(pages, &sink, drift_opts(policy))
5881 .await
5882 .unwrap_err();
5883 assert!(matches!(err, FaucetError::Config(_)));
5884 }
5885
5886 #[cfg(feature = "quality")]
5892 #[tokio::test]
5893 async fn drift_fail_with_dlq_still_routes_quality_quarantine() {
5894 use crate::dlq::DlqConfig;
5895 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
5896
5897 let sink = SchemaSink::new(
5900 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5901 false,
5902 );
5903 let dlq_sink = std::sync::Arc::new(MockSink::new());
5904 let policy = crate::drift::SchemaDriftPolicy {
5905 on_drift: crate::drift::OnDrift::Fail,
5906 allow_widening: true,
5907 on_incompatible: crate::drift::OnIncompatible::Fail,
5908 relax_nullability_on_missing: false,
5909 };
5910 let spec = QualitySpec {
5911 record: vec![RecordCheck::NotNull {
5912 field: "name".into(),
5913 treat_missing_as_null: true,
5914 on_failure: OnFailure::Quarantine,
5915 }],
5916 batch: vec![],
5917 };
5918 let quality = std::sync::Arc::new(CompiledQuality::compile(&spec).unwrap());
5919 let pages = one_page(vec![
5920 json!({"id": 1, "name": "ok", "email": "a@x"}), json!({"id": 2, "email": "b@x"}), ]);
5923 let opts = RunStreamOptions::new()
5924 .with_schema_drift(policy)
5925 .with_quality(quality)
5926 .with_dlq(DlqConfig::new(dlq_sink.clone()));
5927 let err = run_stream(pages, &sink, opts).await.unwrap_err();
5928 assert!(
5930 matches!(err, FaucetError::SchemaDrift { .. }),
5931 "got {err:?}"
5932 );
5933 let dlq = dlq_sink.written();
5935 assert_eq!(
5936 dlq.len(),
5937 1,
5938 "quarantined row must reach the DLQ before abort"
5939 );
5940 assert_eq!(dlq[0]["payload"], json!({"id": 2, "email": "b@x"}));
5941 assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
5942 assert_eq!(
5944 sink.written(),
5945 vec![json!({"id": 1, "name": "ok", "email": "a@x"})]
5946 );
5947 }
5948}
5949
5950#[cfg(test)]
5951mod cleanup_tests {
5952 use super::*;
5957 use crate::cleanup::{CleanupPolicy, CleanupTracker, SeenKeys};
5958 use async_trait::async_trait;
5959 use serde_json::json;
5960 use std::collections::BTreeMap;
5961 use std::sync::{Arc, Mutex};
5962 use tokio_util::sync::CancellationToken;
5963
5964 #[derive(Default)]
5967 struct CleanupSink {
5968 written: Mutex<Vec<Value>>,
5969 cleanup_calls: Mutex<Vec<Vec<String>>>,
5970 supports: bool,
5971 fail_writes: bool,
5972 }
5973
5974 impl CleanupSink {
5975 fn new() -> Self {
5976 Self {
5977 supports: true,
5978 ..Default::default()
5979 }
5980 }
5981 fn unsupported() -> Self {
5982 Self::default()
5983 }
5984 fn failing() -> Self {
5985 Self {
5986 supports: true,
5987 fail_writes: true,
5988 ..Default::default()
5989 }
5990 }
5991 fn calls(&self) -> Vec<Vec<String>> {
5992 self.cleanup_calls.lock().unwrap().clone()
5993 }
5994 }
5995
5996 #[async_trait]
5997 impl Sink for CleanupSink {
5998 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
5999 if self.fail_writes {
6000 return Err(FaucetError::Sink("boom".into()));
6001 }
6002 self.written.lock().unwrap().extend(records.iter().cloned());
6003 Ok(records.len())
6004 }
6005 fn supports_cleanup(&self) -> bool {
6006 self.supports
6007 }
6008 async fn cleanup_scope(
6009 &self,
6010 scope: &BTreeMap<String, Value>,
6011 seen: &SeenKeys,
6012 ) -> Result<u64, FaucetError> {
6013 assert!(
6014 !scope.is_empty(),
6015 "the pipeline must never pass an empty scope"
6016 );
6017 let mut ids: Vec<String> = seen
6018 .keys()
6019 .iter()
6020 .map(|k| {
6021 k.0.iter()
6022 .map(|(_, v)| v.to_string())
6023 .collect::<Vec<_>>()
6024 .join("|")
6025 })
6026 .collect();
6027 ids.sort();
6028 self.cleanup_calls.lock().unwrap().push(ids);
6029 Ok(3)
6030 }
6031 }
6032
6033 fn policy(max_keys: usize) -> Arc<CleanupPolicy> {
6034 Arc::new(
6035 CleanupPolicy::new(
6036 BTreeMap::from([("contact_id".to_string(), json!(7))]),
6037 vec!["id".to_string()],
6038 max_keys,
6039 )
6040 .unwrap(),
6041 )
6042 }
6043
6044 struct VecSource(Vec<Value>);
6046
6047 #[async_trait]
6048 impl Source for VecSource {
6049 async fn fetch_with_context(
6050 &self,
6051 _context: &std::collections::HashMap<String, Value>,
6052 ) -> Result<Vec<Value>, FaucetError> {
6053 Ok(self.0.clone())
6054 }
6055 }
6056
6057 fn src(records: Vec<Value>) -> VecSource {
6058 VecSource(records)
6059 }
6060
6061 #[tokio::test]
6064 async fn tracker_counts_rows_that_fail_into_the_dlq() {
6065 let sink = CleanupSink::new();
6069 let p = policy(100);
6070 let tracker = CleanupTracker::new(&sink, &p);
6071 let page = vec![json!({"id": 1}), json!({"id": 2})];
6072 tracker.write_batch_partial(&page).await.unwrap();
6074 assert_eq!(tracker.tracked(), 2);
6075 }
6076
6077 #[tokio::test]
6078 async fn tracker_refuses_to_delete_after_an_overflow() {
6079 let sink = CleanupSink::new();
6080 let p = policy(1);
6081 let tracker = CleanupTracker::new(&sink, &p);
6082 tracker
6083 .write_batch(&[json!({"id": 1}), json!({"id": 2})])
6084 .await
6085 .unwrap();
6086 let err = tracker
6087 .finish(&p)
6088 .await
6089 .expect_err("an overflowed tracker must refuse");
6090 assert!(err.to_string().contains("Nothing was deleted"), "{err}");
6091 assert!(
6092 sink.calls().is_empty(),
6093 "the sink must never be asked to delete"
6094 );
6095 }
6096
6097 #[tokio::test]
6100 async fn fires_once_at_a_natural_end_with_every_written_key() {
6101 let sink = CleanupSink::new();
6102 let source = src(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
6103 let result = Pipeline::new(&source, &sink)
6104 .with_cleanup(policy(100))
6105 .run()
6106 .await
6107 .expect("run succeeds");
6108 assert_eq!(result.records_written, 3);
6109 let calls = sink.calls();
6110 assert_eq!(calls.len(), 1, "exactly one cleanup per run");
6111 assert_eq!(calls[0], vec!["1", "2", "3"]);
6112 }
6113
6114 #[tokio::test]
6115 async fn empty_result_set_still_fires_with_no_keys() {
6116 let sink = CleanupSink::new();
6120 let source = src(vec![]);
6121 Pipeline::new(&source, &sink)
6122 .with_cleanup(policy(100))
6123 .run()
6124 .await
6125 .unwrap();
6126 let calls = sink.calls();
6127 assert_eq!(calls.len(), 1, "must still fire on an empty fetch");
6128 assert!(
6129 calls[0].is_empty(),
6130 "no keys written → delete the whole scope"
6131 );
6132 }
6133
6134 #[tokio::test]
6135 async fn skipped_when_the_run_was_cancelled() {
6136 let sink = CleanupSink::new();
6137 let source = src(vec![json!({"id": 1})]);
6138 let token = CancellationToken::new();
6139 token.cancel();
6140 Pipeline::new(&source, &sink)
6141 .with_cleanup(policy(100))
6142 .with_cancel(token)
6143 .run()
6144 .await
6145 .expect("a cancelled run still returns Ok with partial output flushed");
6146 assert!(
6147 sink.calls().is_empty(),
6148 "cleanup must not run after a cancellation"
6149 );
6150 }
6151
6152 #[tokio::test]
6153 async fn not_fired_when_the_run_failed() {
6154 let sink = CleanupSink::failing();
6155 let source = src(vec![json!({"id": 1})]);
6156 let err = Pipeline::new(&source, &sink)
6157 .with_cleanup(policy(100))
6158 .run()
6159 .await
6160 .expect_err("the write failure must propagate");
6161 assert!(err.to_string().contains("boom"), "{err}");
6162 assert!(
6163 sink.calls().is_empty(),
6164 "a failed run never wrote the authoritative set"
6165 );
6166 }
6167
6168 #[tokio::test]
6169 async fn refuses_and_deletes_nothing_when_the_key_ceiling_is_breached() {
6170 let sink = CleanupSink::new();
6171 let source = src(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
6172 let err = Pipeline::new(&source, &sink)
6173 .with_cleanup(policy(2))
6174 .run()
6175 .await
6176 .expect_err("must fail rather than partially delete");
6177 assert!(err.to_string().contains("Nothing was deleted"), "{err}");
6178 assert!(sink.calls().is_empty());
6179 }
6180
6181 #[tokio::test]
6182 async fn rejects_a_sink_that_does_not_support_cleanup() {
6183 let sink = CleanupSink::unsupported();
6184 let source = src(vec![json!({"id": 1})]);
6185 let err = Pipeline::new(&source, &sink)
6186 .with_cleanup(policy(100))
6187 .run()
6188 .await
6189 .expect_err("unsupported sink must be rejected");
6190 assert!(
6191 err.to_string().contains("does not support scoped cleanup"),
6192 "{err}"
6193 );
6194 assert!(
6195 sink.written.lock().unwrap().is_empty(),
6196 "must fail before writing anything"
6197 );
6198 }
6199
6200 #[tokio::test]
6201 async fn no_policy_leaves_the_write_path_untouched() {
6202 let sink = CleanupSink::new();
6203 let source = src(vec![json!({"id": 1})]);
6204 let result = Pipeline::new(&source, &sink).run().await.unwrap();
6205 assert_eq!(result.records_written, 1);
6206 assert!(sink.calls().is_empty(), "no policy → no cleanup");
6207 }
6208
6209 #[tokio::test]
6210 async fn rows_missing_the_key_are_not_tracked_but_do_not_break_the_pass() {
6211 let sink = CleanupSink::new();
6212 let source = src(vec![json!({"id": 1}), json!({"no_key": true})]);
6213 Pipeline::new(&source, &sink)
6214 .with_cleanup(policy(100))
6215 .run()
6216 .await
6217 .unwrap();
6218 assert_eq!(sink.calls()[0], vec!["1"]);
6219 }
6220}