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}
145
146impl<'a, So: Source + ?Sized, Si: Sink + ?Sized> Pipeline<'a, So, Si> {
147 pub fn new(source: &'a So, sink: &'a Si) -> Self {
149 Self {
150 source,
151 sink,
152 state_store: None,
153 name: None,
154 row: None,
155 run_id: None,
156 dlq: None,
157 #[cfg(feature = "quality")]
158 quality: None,
159 #[cfg(feature = "contract")]
160 contract: None,
161 #[cfg(feature = "masking")]
162 masking: None,
163 adaptive: None,
164 cancel: None,
165 delivery: crate::idempotency::DeliveryMode::AtLeastOnce,
166 resilience: None,
167 schema_drift: None,
168 cleanup: None,
169 }
170 }
171
172 pub fn with_state_store(mut self, store: Arc<dyn StateStore>) -> Self {
186 self.state_store = Some(store);
187 self
188 }
189
190 pub fn with_name(mut self, name: impl Into<String>) -> Self {
193 self.name = Some(name.into());
194 self
195 }
196
197 pub fn with_row(mut self, row: impl Into<String>) -> Self {
200 self.row = Some(row.into());
201 self
202 }
203
204 pub fn with_run_id(mut self, run_id: impl Into<String>) -> Self {
208 self.run_id = Some(run_id.into());
209 self
210 }
211
212 pub fn with_dlq(mut self, dlq: DlqConfig) -> Self {
214 self.dlq = Some(dlq);
215 self
216 }
217
218 #[cfg(feature = "quality")]
221 pub fn with_quality(mut self, quality: Arc<crate::quality::CompiledQuality>) -> Self {
222 self.quality = Some(quality);
223 self
224 }
225
226 #[cfg(feature = "contract")]
229 pub fn with_contract(mut self, contract: Arc<crate::contract::CompiledContract>) -> Self {
230 self.contract = Some(contract);
231 self
232 }
233
234 #[cfg(feature = "masking")]
238 pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
239 self.masking = Some(masking);
240 self
241 }
242
243 pub fn with_adaptive(mut self, cfg: crate::adaptive::AdaptiveBatchConfig) -> Self {
247 self.adaptive = Some(cfg);
248 self
249 }
250
251 pub fn with_cancel(mut self, cancel: tokio_util::sync::CancellationToken) -> Self {
256 self.cancel = Some(cancel);
257 self
258 }
259
260 pub fn with_delivery(mut self, mode: crate::idempotency::DeliveryMode) -> Self {
265 self.delivery = mode;
266 self
267 }
268
269 pub fn with_resilience(mut self, policy: crate::resilience::ResiliencePolicy) -> Self {
271 self.resilience = Some(policy);
272 self
273 }
274
275 pub fn with_schema_drift(mut self, policy: crate::drift::SchemaDriftPolicy) -> Self {
278 self.schema_drift = Some(policy);
279 self
280 }
281
282 pub fn with_cleanup(mut self, policy: std::sync::Arc<crate::cleanup::CleanupPolicy>) -> Self {
287 self.cleanup = Some(policy);
288 self
289 }
290
291 pub async fn run(&self) -> Result<PipelineResult, FaucetError> {
305 use crate::observability::{
306 DurationGuard, InstrumentedSink, InstrumentedSource, InstrumentedStateStore, Labels,
307 };
308 use metrics::{Label, SharedString, counter, gauge};
309 use tracing::Instrument;
310
311 let name = self.name.clone().unwrap_or_else(|| "unnamed".to_string());
313 let row = self.row.clone().unwrap_or_default();
314 let run_id = self
315 .run_id
316 .clone()
317 .unwrap_or_else(|| uuid::Uuid::now_v7().to_string());
318 let obs_labels = Labels::new(name.clone(), row.clone(), run_id.clone());
319
320 let wrapped_source = InstrumentedSource::new(self.source, obs_labels.clone());
322 let wrapped_sink = InstrumentedSink::new(self.sink, obs_labels.clone());
323 let wrapped_state_store: Option<Arc<dyn StateStore>> = self.state_store.as_ref().map(|s| {
324 Arc::new(InstrumentedStateStore::new(
325 Arc::clone(s),
326 obs_labels.clone(),
327 )) as Arc<dyn StateStore>
328 });
329
330 let span = tracing::info_span!(
333 "faucet.pipeline.run",
334 pipeline = %name,
335 row = %row,
336 run_id = %run_id,
337 source = %wrapped_source.connector_name(),
338 sink = %wrapped_sink.connector_name(),
339 );
340
341 let base_labels: Vec<Label> = vec![
343 Label::new("pipeline", SharedString::from(name.clone())),
344 Label::new("row", SharedString::from(row.clone())),
345 ];
346 let run_labels: Vec<Label> = {
347 let mut v = base_labels.clone();
348 v.push(Label::new(
349 "source",
350 SharedString::from(wrapped_source.connector_name().to_string()),
351 ));
352 v.push(Label::new(
353 "sink",
354 SharedString::from(wrapped_sink.connector_name().to_string()),
355 ));
356 v
357 };
358
359 struct InFlightGuard(Vec<Label>);
361 impl Drop for InFlightGuard {
362 fn drop(&mut self) {
363 gauge!("faucet_pipeline_in_flight", self.0.clone()).decrement(1.0);
364 }
365 }
366 gauge!("faucet_pipeline_in_flight", base_labels.clone()).increment(1.0);
367 let _in_flight = InFlightGuard(base_labels.clone());
368
369 let start_unix = std::time::SystemTime::now()
372 .duration_since(std::time::UNIX_EPOCH)
373 .map(|d| d.as_secs_f64())
374 .unwrap_or(0.0);
375 gauge!(
376 "faucet_pipeline_start_time_unix_seconds",
377 base_labels.clone()
378 )
379 .set(start_unix);
380
381 let _run_timer =
383 DurationGuard::new("faucet_pipeline_run_duration_seconds", run_labels.clone());
384
385 let result = async {
387 let state_key = self.source.state_key();
390 let mut start_seq = 0u64;
391 if let (Some(store), Some(key)) = (wrapped_state_store.as_ref(), state_key.as_ref()) {
392 validate_state_key(key)?;
393 if let Some(prior) = store.get(key).await? {
394 if self.delivery == crate::idempotency::DeliveryMode::ExactlyOnce {
395 let (bookmark, seq) = crate::idempotency::unwrap_state(&prior);
396 start_seq = seq;
397 if let Some(bm) = bookmark {
398 wrapped_source.apply_start_bookmark(bm).await?;
399 }
400 } else {
401 wrapped_source.apply_start_bookmark(prior).await?;
402 }
403 }
404 }
405
406 if self.delivery == crate::idempotency::DeliveryMode::ExactlyOnce
417 && wrapped_sink.supports_idempotent_writes()
418 && wrapped_source.replay_guarantee()
419 == crate::idempotency::ReplayGuarantee::Deterministic
420 && let Some(key) = state_key.as_ref()
421 && let Some(token) = wrapped_sink.last_committed_token(key).await?
422 && let Some((sink_seq, Some(bm))) = crate::idempotency::parse_token_parts(&token)
423 && sink_seq > start_seq
424 {
425 wrapped_source.apply_start_bookmark(bm).await?;
426 start_seq = sink_seq;
427 }
428
429 #[cfg(feature = "arrow")]
437 {
438 let columnar_ok = wrapped_source.supports_columnar()
439 && wrapped_sink.supports_columnar()
440 && self.dlq.is_none()
441 && self.delivery == crate::idempotency::DeliveryMode::AtLeastOnce
442 && self.schema_drift.is_none()
443 && self.adaptive.is_none()
444 && self.resilience.is_none()
445 && self.cleanup.is_none();
451 #[cfg(feature = "quality")]
452 let columnar_ok = columnar_ok && self.quality.is_none();
453 #[cfg(feature = "contract")]
454 let columnar_ok = columnar_ok && self.contract.is_none();
455 #[cfg(feature = "masking")]
456 let columnar_ok = columnar_ok && self.masking.is_none();
457 if columnar_ok {
458 let state = match (wrapped_state_store.clone(), state_key.clone()) {
459 (Some(store), Some(key)) => Some((store, key)),
460 _ => None,
461 };
462 return run_stream_columnar(
463 &wrapped_source,
464 &wrapped_sink,
465 state,
466 self.cancel.clone(),
467 &name,
468 &row,
469 &run_id,
470 )
471 .await;
472 }
473 }
474
475 let ctx = std::collections::HashMap::new();
476 let pages = wrapped_source.stream_pages(&ctx, DEFAULT_BATCH_SIZE);
477
478 let mut opts = RunStreamOptions::new()
479 .with_name(name.clone())
480 .with_row(row.clone())
481 .with_run_id(run_id.clone());
482 if let (Some(store), Some(key)) = (wrapped_state_store.clone(), state_key) {
483 opts = opts.with_state(store, key);
484 }
485 if let Some(dlq) = self.dlq.clone() {
486 opts = opts.with_dlq(dlq);
487 }
488 #[cfg(feature = "quality")]
489 if let Some(q) = self.quality.clone() {
490 opts = opts.with_quality(q);
491 }
492 #[cfg(feature = "contract")]
493 if let Some(c) = self.contract.clone() {
494 opts = opts.with_contract(c);
495 }
496 #[cfg(feature = "masking")]
497 if let Some(m) = self.masking.clone() {
498 opts = opts.with_masking(m);
499 }
500 if let Some(ad) = self.adaptive.clone() {
501 opts = opts.with_adaptive(ad);
502 }
503 if let Some(cancel) = self.cancel.clone() {
504 opts = opts.with_cancel(cancel);
505 }
506 if let Some(policy) = self.resilience.clone() {
507 opts = opts.with_resilience(policy);
508 }
509 if let Some(p) = self.schema_drift {
510 opts = opts.with_schema_drift(p);
511 }
512 opts = opts
513 .with_delivery(self.delivery)
514 .with_start_seq(start_seq)
515 .with_replay_guarantee(wrapped_source.replay_guarantee());
516
517 match self.cleanup.clone() {
528 None => run_stream(pages, &wrapped_sink, opts).await,
529 Some(policy) => {
530 if !wrapped_sink.supports_cleanup() {
531 return Err(FaucetError::Config(format!(
532 "cleanup is configured but sink '{}' does not support scoped cleanup",
533 wrapped_sink.connector_name()
534 )));
535 }
536 let tracker = crate::cleanup::CleanupTracker::new(&wrapped_sink, &policy);
537 let result = run_stream(pages, &tracker, opts).await;
538 match result {
539 Ok(r) => {
540 let cancelled = self.cancel.as_ref().is_some_and(|c| c.is_cancelled());
541 if cancelled {
542 tracing::warn!(
543 "scoped cleanup skipped: the run was cancelled, so the \
544 set of written keys is incomplete and a delete would \
545 remove rows that were never read"
546 );
547 crate::observability::cleanup_run(&name, &row, "skipped_cancelled");
548 } else {
549 let tracked = tracker.tracked();
550 match tracker.finish(&policy).await {
551 Ok(deleted) => {
552 crate::observability::cleanup_deleted(
553 &name,
554 &row,
555 wrapped_sink.connector_name(),
556 deleted,
557 );
558 crate::observability::cleanup_run(&name, &row, "applied");
559 tracing::info!(
560 deleted,
561 kept = tracked,
562 "scoped cleanup complete: deleted destination rows \
563 in the claimed scope that this run did not write"
564 );
565 }
566 Err(e) => {
567 crate::observability::cleanup_run(
568 &name,
569 &row,
570 "refused_overflow",
571 );
572 return Err(e);
573 }
574 }
575 }
576 Ok(r)
577 }
578 Err(e) => Err(e),
579 }
580 }
581 }
582 }
583 .instrument(span)
584 .await;
585
586 let status = if result.is_ok() { "ok" } else { "err" };
590 let mut final_labels = run_labels;
591 final_labels.push(Label::new("status", SharedString::const_str(status)));
592 if let Err(ref e) = result {
593 final_labels.push(Label::new(
594 "kind",
595 SharedString::const_str(crate::observability::decorator::error_kind(e)),
596 ));
597 }
598 counter!("faucet_pipeline_runs_total", final_labels).increment(1);
599
600 result
601 }
602}
603
604#[cfg(feature = "arrow")]
615async fn run_stream_columnar<S, Si>(
616 source: &S,
617 sink: &Si,
618 state: Option<(Arc<dyn StateStore>, String)>,
619 cancel: Option<tokio_util::sync::CancellationToken>,
620 pipeline: &str,
621 row: &str,
622 run_id: &str,
623) -> Result<PipelineResult, FaucetError>
624where
625 S: crate::Source + ?Sized,
626 Si: Sink + ?Sized,
627{
628 use futures::StreamExt;
629 use metrics::{Label, SharedString, counter};
630
631 let labels = |connector: &str| -> Vec<Label> {
632 vec![
633 Label::new("pipeline", SharedString::from(pipeline.to_string())),
634 Label::new("row", SharedString::from(row.to_string())),
635 Label::new("connector", SharedString::from(connector.to_string())),
636 ]
637 };
638 let src_labels = labels(source.connector_name());
639 let sink_labels = labels(sink.connector_name());
640 let _ = run_id; let ctx = std::collections::HashMap::new();
643 let mut batches = source.stream_batches(&ctx, DEFAULT_BATCH_SIZE);
644 let mut records_written = 0usize;
645 let mut last_bookmark: Option<Value> = None;
646
647 loop {
648 let page = match &cancel {
651 Some(token) => {
652 tokio::select! {
653 biased;
654 _ = token.cancelled() => break,
655 p = batches.next() => p,
656 }
657 }
658 None => batches.next().await,
659 };
660 let Some(page) = page else { break };
661 let page = page?;
662 let rows = page.num_rows();
663 counter!("faucet_source_records_total", src_labels.clone()).increment(rows as u64);
664
665 if rows > 0 {
666 let n = sink.write_batch_columnar(&page.batch).await?;
667 records_written += n;
668 counter!("faucet_sink_records_total", sink_labels.clone()).increment(n as u64);
669 counter!("faucet_sink_writes_total", sink_labels.clone()).increment(1);
670 }
671
672 if let Some(bm) = page.bookmark {
675 sink.flush().await?;
676 if let Some((store, key)) = state.as_ref() {
677 store.put(key, &bm).await?;
678 }
679 last_bookmark = Some(bm);
680 }
681 }
682
683 sink.flush().await?;
685 Ok(PipelineResult {
686 records_written,
687 bookmark: last_bookmark,
688 dlq: None,
689 })
690}
691
692pub async fn run_stream<S, Si>(
710 mut pages: S,
711 sink: &Si,
712 options: RunStreamOptions,
713) -> Result<PipelineResult, FaucetError>
714where
715 S: Stream<Item = Result<StreamPage, FaucetError>> + Unpin,
716 Si: Sink + ?Sized,
717{
718 use crate::dlq::{DlqReason, DlqStats, OnBatchError, build_envelope};
719
720 let state_store = options.state_store.clone();
721 let state_key = options.state_key.clone();
722 let pipeline_name = options.pipeline_name.unwrap_or_else(|| "unnamed".into());
723 let row = options.row.unwrap_or_default();
724 let run_id = options.run_id.unwrap_or_default();
725 let dlq = options.dlq.clone();
726 let cancel = options.cancel.clone();
727
728 #[cfg(feature = "quality")]
729 let quality = options.quality.clone();
730
731 #[cfg(feature = "quality")]
733 if let Some(q) = quality.as_ref()
734 && q.requires_dlq()
735 && dlq.is_none()
736 {
737 return Err(FaucetError::Config(
738 "quality: on_failure 'quarantine'/'quarantine_batch' requires a DLQ sink".into(),
739 ));
740 }
741
742 #[cfg(feature = "contract")]
743 let contract = options.contract.clone();
744 #[cfg(feature = "contract")]
746 if let Some(c) = contract.as_ref()
747 && c.requires_dlq()
748 && dlq.is_none()
749 {
750 return Err(FaucetError::Config(
751 "contract: on_breach 'quarantine' requires a DLQ sink".into(),
752 ));
753 }
754 #[cfg(feature = "contract")]
756 let mut warned_contract_breach = false;
757
758 #[cfg(feature = "masking")]
762 let masking = options.masking.clone();
763
764 let schema_drift = options.schema_drift;
766 if let Some(p) = schema_drift.as_ref()
768 && p.requires_dlq()
769 && dlq.is_none()
770 {
771 return Err(FaucetError::Config(
772 "schema: on_drift 'quarantine' (or on_incompatible 'quarantine') requires a DLQ sink"
773 .into(),
774 ));
775 }
776 let mut dest_schema_cache: Option<Option<Value>> = None;
780 let mut warned_drift_inert = false;
781
782 if let Some(key) = state_key.as_ref() {
783 validate_state_key(key)?;
784 }
785
786 let mechanism: Option<crate::idempotency::EffectivelyOnceMechanism> =
794 if options.delivery == crate::idempotency::DeliveryMode::ExactlyOnce {
795 let replay_ok = options
796 .replay
797 .is_none_or(|r| r == crate::idempotency::ReplayGuarantee::Deterministic);
798 if sink.supports_idempotent_writes() && replay_ok {
799 if state_store.is_none() || state_key.is_none() {
800 return Err(FaucetError::Config(
801 "delivery: exactly_once (atomic watermark) requires a state store".into(),
802 ));
803 }
804 if dlq.is_some() {
805 return Err(FaucetError::Config(
806 "delivery: exactly_once (atomic watermark) is not compatible with a DLQ \
807 in this version"
808 .into(),
809 ));
810 }
811 Some(crate::idempotency::EffectivelyOnceMechanism::AtomicWatermark)
812 } else if sink.dedups_by_key() {
813 Some(crate::idempotency::EffectivelyOnceMechanism::KeyedUpsert)
814 } else if sink.supports_idempotent_writes() {
815 return Err(FaucetError::Config(format!(
818 "delivery: exactly_once — the source does not replay deterministically from \
819 a bookmark, so the atomic-watermark mechanism cannot be used; configure \
820 `write_mode: upsert` with a `key` on sink '{}' for keyed-upsert \
821 effectively-once instead",
822 sink.connector_name()
823 )));
824 } else {
825 return Err(FaucetError::Config(format!(
826 "delivery: exactly_once requires an idempotent (atomic-watermark) sink or a \
827 sink configured to dedup by key (`write_mode: upsert` + `key`), but '{}' \
828 provides neither",
829 sink.connector_name()
830 )));
831 }
832 } else {
833 None
834 };
835 let exactly_once =
839 mechanism == Some(crate::idempotency::EffectivelyOnceMechanism::AtomicWatermark);
840 let scope = state_key.clone().unwrap_or_default();
841 let mut next_seq = options.start_seq;
842 let committed_seq = if exactly_once {
843 sink.last_committed_token(&scope)
844 .await?
845 .and_then(|t| crate::idempotency::parse_token(&t))
846 .unwrap_or(0)
847 } else {
848 0
849 };
850
851 let mut records_written = 0usize;
852 let mut last_bookmark: Option<Value> = None;
853 let mut dlq_stats = DlqStats::default();
854
855 let adaptive_cfg = options.adaptive.clone().filter(|c| c.enabled);
856 if let Some(cfg) = adaptive_cfg.as_ref() {
860 cfg.validate()?;
861 }
862 let mut controller: Option<crate::adaptive::AimdController> = None;
863 let mut warned_noop_sink = false;
864 let mut warned_poison_drop = false;
866
867 let sink_name = sink.connector_name();
868 let dlq_sink_name = dlq.as_ref().map(|d| d.sink.connector_name()).unwrap_or("");
869
870 let mut cancelled = false;
881
882 let resilience = options.resilience.clone();
889 let retry_policy = resilience.as_ref().map(|r| r.retry.clone());
890 let mut breaker = resilience
891 .as_ref()
892 .and_then(|r| r.circuit_breaker)
893 .map(|cb| {
894 (
895 crate::resilience::CircuitBreaker::new(cb.consecutive_failures),
896 cb.cooldown,
897 )
898 });
899 let poison = resilience.as_ref().and_then(|r| r.poison);
901
902 macro_rules! with_retry {
915 ($op_label:literal, $op:expr) => {
916 match &retry_policy {
917 Some(p) => {
918 let m = crate::resilience::RetryMetrics {
919 pipeline: pipeline_name.to_string(),
920 row: row.to_string(),
921 op: $op_label,
922 };
923 crate::resilience::execute_with_policy_metered(p, cancel.as_ref(), &m, || $op)
924 .await
925 }
926 None => $op.await,
927 }
928 };
929 }
930
931 macro_rules! with_retry_write {
942 ($op_label:literal, $op:expr) => {
943 if retry_policy.is_some() && sink.supports_idempotent_writes() {
944 with_retry!($op_label, $op)
945 } else {
946 $op.await
947 }
948 };
949 }
950
951 let loop_result: Result<(), FaucetError> = async {
952 loop {
953 let page = match &cancel {
957 Some(token) => tokio::select! {
958 biased;
959 _ = token.cancelled() => {
960 cancelled = true;
961 break;
962 }
963 p = std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)) => p,
964 },
965 None => std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)).await,
966 };
967 match page {
968 Some(Ok(page)) => {
969 if page.records.is_empty() && page.bookmark.is_none() {
970 continue;
971 }
972
973 #[cfg(feature = "masking")]
979 let page = if let Some(m) = masking.as_ref() {
980 let labels =
981 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
982 let outcome = crate::observability::instrumented_apply_masking(
983 page.records,
984 m,
985 &labels,
986 );
987 StreamPage {
988 records: outcome.records,
989 bookmark: page.bookmark,
990 }
991 } else {
992 page
993 };
994
995 let page_len = page.records.len();
1002
1003 #[cfg(feature = "quality")]
1005 let (records, quality_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
1006 if let Some(q) = quality.as_ref() {
1007 let labels =
1008 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1009 let outcome = crate::observability::instrumented_apply_quality(
1010 page.records,
1011 q,
1012 &labels,
1013 )?;
1014 let quarantined_idx: std::collections::HashSet<usize> =
1015 outcome.quarantined.iter().map(|qr| qr.page_index).collect();
1016 let envelopes: Vec<Value> = outcome
1017 .quarantined
1018 .iter()
1019 .map(|qr| {
1020 let err = FaucetError::QualityFailure {
1021 check: qr.check.to_string(),
1022 message: qr.message.clone(),
1023 };
1024 build_envelope(
1028 &qr.record,
1029 &err,
1030 DlqReason::Quality,
1031 sink_name,
1032 &pipeline_name,
1033 &row,
1034 qr.page_index,
1035 )
1036 })
1037 .collect();
1038 let survivor_idx: Vec<usize> =
1039 (0..page_len).filter(|i| !quarantined_idx.contains(i)).collect();
1040 (outcome.survivors, envelopes, survivor_idx)
1041 } else {
1042 (page.records, Vec::new(), (0..page_len).collect())
1043 };
1044 #[cfg(not(feature = "quality"))]
1045 let (records, quality_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
1046 (page.records, Vec::new(), (0..page_len).collect());
1047
1048 #[cfg(feature = "contract")]
1055 let (records, contract_envelopes, page_indices): (Vec<Value>, Vec<Value>, Vec<usize>) =
1056 if let Some(c) = contract.as_ref() {
1057 let labels =
1058 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1059 let outcome = crate::observability::instrumented_apply_contract(
1060 records, c, &labels,
1061 )?;
1062 if !outcome.warned.is_empty() && !warned_contract_breach {
1063 tracing::warn!(
1064 version = %c.version,
1065 breaches = outcome.warned.len(),
1066 first = %outcome.warned[0].describe(),
1067 "contract: breaching records written unchanged \
1068 (on_breach=warn); this warning fires once per run"
1069 );
1070 warned_contract_breach = true;
1071 }
1072 let envelopes: Vec<Value> = outcome
1073 .quarantined
1074 .iter()
1075 .map(|vr| {
1076 let err = FaucetError::ContractViolation {
1077 version: c.version.clone(),
1078 message: vr.violation.describe(),
1079 };
1080 let page_index = page_indices
1090 .get(vr.violation.page_index)
1091 .copied()
1092 .unwrap_or(vr.violation.page_index);
1093 build_envelope(
1094 &vr.record,
1095 &err,
1096 DlqReason::Contract,
1097 sink_name,
1098 &pipeline_name,
1099 &row,
1100 page_index,
1101 )
1102 })
1103 .collect();
1104 let contract_quarantined: std::collections::HashSet<usize> = outcome
1110 .quarantined
1111 .iter()
1112 .map(|vr| vr.violation.page_index)
1113 .collect();
1114 let survivor_idx: Vec<usize> = page_indices
1115 .iter()
1116 .enumerate()
1117 .filter(|(pos, _)| !contract_quarantined.contains(pos))
1118 .map(|(_, orig)| *orig)
1119 .collect();
1120 (outcome.survivors, envelopes, survivor_idx)
1121 } else {
1122 (records, Vec::new(), page_indices)
1123 };
1124
1125 let mut drift_envelopes: Vec<Value> = Vec::new();
1127 let (records, drift_abort): (Vec<Value>, Option<FaucetError>) =
1128 if let Some(policy) = schema_drift.as_ref().filter(|_| !records.is_empty()) {
1129 if dest_schema_cache.is_none() {
1131 dest_schema_cache = Some(sink.current_schema().await?);
1132 }
1133 let dest = dest_schema_cache.as_ref().and_then(|o| o.as_ref());
1134 match dest {
1135 None => {
1136 if !warned_drift_inert {
1137 tracing::info!(
1138 connector = sink_name,
1139 "schema-drift: sink reports no destination schema; \
1140 drift handling is inert this run"
1141 );
1142 warned_drift_inert = true;
1143 }
1144 (records, None)
1145 }
1146 Some(dest) => {
1147 let inferred = crate::schema::infer_schema(&records);
1148 let diff = crate::drift::diff_schema(
1149 dest,
1150 &inferred,
1151 policy.allow_widening,
1152 );
1153 if diff.is_empty() {
1154 (records, None)
1155 } else {
1156 let dest_owned = dest.clone();
1160 apply_drift_policy(
1161 policy,
1162 &diff,
1163 &dest_owned,
1164 records,
1165 &page_indices,
1166 sink,
1167 sink_name,
1168 &pipeline_name,
1169 &row,
1170 &mut dest_schema_cache,
1171 &mut drift_envelopes,
1172 )
1173 .await?
1174 }
1175 }
1176 }
1177 } else {
1178 (records, None)
1179 };
1180 let quality_envelopes = {
1184 let mut q = quality_envelopes;
1185 #[cfg(feature = "contract")]
1186 q.extend(contract_envelopes);
1187 q.append(&mut drift_envelopes);
1188 q
1189 };
1190 let mut drift_abort = drift_abort;
1200 if dlq.is_none()
1201 && let Some(e) = drift_abort.take()
1202 {
1203 return Err(e);
1204 }
1205
1206 let page = StreamPage {
1207 records,
1208 bookmark: page.bookmark,
1209 };
1210
1211 if let Some(ref dlq_cfg) = dlq {
1212 use metrics::{Label, SharedString, counter};
1214 let metric_labels: Vec<Label> = vec![
1215 Label::new("pipeline", SharedString::from(pipeline_name.clone())),
1216 Label::new("row", SharedString::from(row.clone())),
1217 Label::new("connector", SharedString::from(sink_name.to_string())),
1218 Label::new(
1219 "dlq_connector",
1220 SharedString::from(dlq_sink_name.to_string()),
1221 ),
1222 ];
1223 let span = tracing::info_span!(
1224 "faucet.dlq.route",
1225 pipeline = %pipeline_name,
1226 row = %row,
1227 run_id = %run_id,
1228 connector = %sink_name,
1229 dlq_connector = %dlq_sink_name,
1230 );
1231 let _enter = span.enter();
1232
1233 let mut envelopes: Vec<Value> = Vec::new();
1237 let mut page_success = 0usize;
1238 let mut outer_err_recovered = false;
1239 let mut had_per_row_sink_failure = false;
1244 let records_len = page.records.len();
1245 let mut offset = 0usize;
1246 while offset < records_len {
1247 let size = match adaptive_cfg.as_ref() {
1248 Some(cfg) => {
1249 let ctrl = controller.get_or_insert_with(|| {
1250 crate::adaptive::AimdController::new(cfg, records_len)
1251 });
1252 ctrl.current().max(1).min(records_len - offset)
1253 }
1254 None => records_len - offset, };
1256 if adaptive_cfg.is_some() {
1257 maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
1258 }
1259 let chunk = &page.records[offset..offset + size];
1260 let t0 = std::time::Instant::now();
1261 let chunk_outcomes_result =
1267 with_retry_write!("sink_write", sink.write_batch_partial(chunk));
1268 let latency = t0.elapsed();
1269 let (mut chunk_outcomes, chunk_synthesized): (
1278 Vec<crate::RowOutcome>,
1279 bool,
1280 ) = match chunk_outcomes_result {
1281 Ok(o) => (o, false),
1282 Err(e) => match dlq_cfg.on_batch_error {
1283 OnBatchError::Propagate => return Err(e),
1284 OnBatchError::DlqAll => {
1285 outer_err_recovered = true;
1286 let msg = e.to_string();
1287 let synth = (0..chunk.len())
1288 .map(|_| Err(FaucetError::Sink(msg.clone())))
1289 .collect();
1290 (synth, true)
1291 }
1292 },
1293 };
1294
1295 if let Some(pp) = poison
1304 && !chunk_synthesized
1305 {
1306 let mut attempt = 1u32; while attempt < pp.max_row_attempts {
1308 let failing: Vec<usize> = chunk_outcomes
1309 .iter()
1310 .enumerate()
1311 .filter_map(|(j, o)| match o {
1312 Err(e)
1313 if retry_policy
1314 .as_ref()
1315 .map(|p| p.is_retriable(e))
1316 .unwrap_or(false) =>
1317 {
1318 Some(j)
1319 }
1320 _ => None,
1321 })
1322 .collect();
1323 if failing.is_empty() {
1324 break;
1325 }
1326 let subset: Vec<Value> =
1327 failing.iter().map(|&j| chunk[j].clone()).collect();
1328 let retried = sink.write_batch_partial(&subset).await?;
1336 let mut retried = retried.into_iter();
1340 for &j in failing.iter() {
1341 chunk_outcomes[j] = retried.next().unwrap_or(Ok(()));
1342 }
1343 attempt += 1;
1344 }
1345 }
1346
1347 let mut chunk_errors = 0usize;
1348 let mut poison_dlq = 0u64;
1354 let mut poison_drop = 0u64;
1355 for (j, outcome) in chunk_outcomes.iter().enumerate() {
1356 match outcome {
1357 Ok(()) => page_success += 1,
1358 Err(err) => {
1359 let action = poison
1364 .map(|pp| pp.action)
1365 .unwrap_or(crate::resilience::PoisonAction::Dlq);
1366 match action {
1367 crate::resilience::PoisonAction::Fail => {
1368 crate::observability::resilience::poison_rows(
1369 &pipeline_name,
1370 &row,
1371 "fail",
1372 1,
1373 );
1374 return Err(FaucetError::Sink(format!(
1375 "poison-pill row failed permanently: {err}"
1376 )));
1377 }
1378 crate::resilience::PoisonAction::Drop => {
1379 poison_drop += 1;
1382 if !warned_poison_drop {
1383 tracing::warn!(
1384 "poison-pill: dropping permanently-failing row(s) (action=drop); this warning fires once per run"
1385 );
1386 warned_poison_drop = true;
1387 }
1388 }
1389 crate::resilience::PoisonAction::Dlq => {
1390 poison_dlq += 1;
1391 chunk_errors += 1;
1392 if !chunk_synthesized {
1393 had_per_row_sink_failure = true;
1394 }
1395 envelopes.push(build_envelope(
1396 &chunk[j],
1397 err,
1398 DlqReason::Partial,
1399 sink_name,
1400 &pipeline_name,
1401 &row,
1402 offset + j,
1403 ));
1404 }
1405 }
1406 }
1407 }
1408 }
1409 if poison.is_some() {
1413 crate::observability::resilience::poison_rows(
1414 &pipeline_name,
1415 &row,
1416 "dlq",
1417 poison_dlq,
1418 );
1419 crate::observability::resilience::poison_rows(
1420 &pipeline_name,
1421 &row,
1422 "drop",
1423 poison_drop,
1424 );
1425 }
1426 if let Some(ctrl) = controller.as_mut() {
1427 let adj = ctrl.observe(crate::adaptive::Observation {
1428 batch_len: chunk.len(),
1429 errors: chunk_errors,
1430 latency,
1431 });
1432 emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
1433 }
1434 offset += size;
1435 }
1436 #[cfg(feature = "quality")]
1441 let quality_count = quality_envelopes.len();
1442 #[cfg(not(feature = "quality"))]
1443 let quality_count = 0usize;
1444 envelopes.splice(0..0, quality_envelopes);
1445 let page_failures = envelopes.len();
1446
1447 let mut budget_error: Option<FaucetError> = None;
1460 let mut circuit_error: Option<FaucetError> = None;
1469 if let Some((b, cooldown)) = breaker.as_mut() {
1470 if records_len > 0 && page_success == 0 {
1471 if b.record_failure() {
1472 crate::observability::resilience::circuit_opened(
1473 &pipeline_name,
1474 &row,
1475 );
1476 circuit_error = Some(FaucetError::CircuitOpen {
1477 failures: b.consecutive(),
1478 cooldown: *cooldown,
1479 });
1480 }
1481 } else if page_success > 0 {
1482 b.record_success();
1483 }
1484 }
1485 if let Some(limit) = dlq_cfg.max_failures_per_page
1486 && page_failures > limit
1487 {
1488 let mut lbl = metric_labels.clone();
1489 lbl.retain(|l| l.key() != "dlq_connector");
1490 lbl.push(Label::new("scope", SharedString::const_str("per_page")));
1491 counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
1492 budget_error = Some(FaucetError::Sink(format!(
1493 "DLQ per-page budget exceeded: {page_failures} > {limit}"
1494 )));
1495 }
1496 let new_total = dlq_stats.records_dlq + page_failures;
1497 if budget_error.is_none()
1498 && let Some(limit) = dlq_cfg.max_failures_total
1499 && new_total > limit
1500 {
1501 let mut lbl = metric_labels.clone();
1502 lbl.retain(|l| l.key() != "dlq_connector");
1503 lbl.push(Label::new("scope", SharedString::const_str("total")));
1504 counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
1505 budget_error = Some(FaucetError::Sink(format!(
1506 "DLQ total budget exceeded: {new_total} > {limit}"
1507 )));
1508 }
1509
1510 if !envelopes.is_empty() {
1512 let _dlq_write_timer = crate::observability::DurationGuard::new(
1513 "faucet_sink_dlq_write_duration_seconds",
1514 metric_labels.clone(),
1515 );
1516 dlq_cfg.sink.write_batch(&envelopes).await.map_err(|e| {
1517 let mut lbl = metric_labels.clone();
1518 lbl.push(Label::new(
1519 "kind",
1520 SharedString::const_str(
1521 crate::observability::decorator::error_kind(&e),
1522 ),
1523 ));
1524 counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1525 FaucetError::Sink(format!("DLQ sink write failed: {e}"))
1526 })?;
1527 dlq_stats.records_dlq += page_failures;
1528 dlq_stats.pages_with_failures += 1;
1529
1530 let reason_label = if had_per_row_sink_failure {
1549 DlqReason::Partial.as_str()
1550 } else if outer_err_recovered {
1551 DlqReason::DlqAll.as_str()
1552 } else if page_failures > quality_count {
1553 DlqReason::Partial.as_str()
1554 } else {
1555 DlqReason::Quality.as_str()
1556 };
1557 counter!("faucet_sink_dlq_records_total", metric_labels.clone())
1558 .increment(page_failures as u64);
1559 let mut page_labels = metric_labels.clone();
1560 page_labels
1561 .push(Label::new("reason", SharedString::const_str(reason_label)));
1562 counter!("faucet_sink_dlq_pages_total", page_labels).increment(1);
1563 }
1564
1565 records_written += page_success;
1566
1567 if let Some(bookmark) = page.bookmark {
1568 with_retry!("flush", sink.flush())?;
1574 let _dlq_flush_timer = crate::observability::DurationGuard::new(
1575 "faucet_sink_dlq_flush_duration_seconds",
1576 metric_labels.clone(),
1577 );
1578 with_retry!("flush", dlq_cfg.sink.flush()).map_err(|e| {
1579 let mut lbl = metric_labels.clone();
1580 lbl.push(Label::new(
1581 "kind",
1582 SharedString::const_str(
1583 crate::observability::decorator::error_kind(&e),
1584 ),
1585 ));
1586 counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1587 FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
1588 })?;
1589 let bm_labels =
1590 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1591 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1592 if let (Some(store), Some(key)) =
1593 (state_store.as_ref(), state_key.as_ref())
1594 {
1595 with_retry!("state_put", store.put(key, &bookmark))?;
1596 }
1597 last_bookmark = Some(bookmark);
1598 }
1599
1600 if let Some(e) = budget_error {
1607 return Err(e);
1608 }
1609 if let Some(e) = circuit_error {
1611 return Err(e);
1612 }
1613 if let Some(e) = drift_abort {
1618 return Err(e);
1619 }
1620 } else if exactly_once {
1621 if let Some(bookmark) = page.bookmark {
1625 next_seq += 1;
1626 let token = crate::idempotency::format_token_with_bookmark(
1633 next_seq,
1634 Some(&bookmark),
1635 );
1636 if next_seq <= committed_seq {
1637 use metrics::{Label, SharedString, counter};
1641 let skip_labels: Vec<Label> = vec![
1642 Label::new(
1643 "pipeline",
1644 SharedString::from(pipeline_name.clone()),
1645 ),
1646 Label::new("row", SharedString::from(row.clone())),
1647 ];
1648 counter!("faucet_pipeline_pages_skipped_total", skip_labels)
1649 .increment(1);
1650 } else {
1651 records_written += with_retry!(
1652 "sink_write",
1653 sink.write_batch_idempotent(&page.records, &scope, &token)
1654 )?;
1655 }
1656 with_retry!("flush", sink.flush())?;
1657 let bm_labels =
1658 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1659 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1660 if let (Some(store), Some(key)) =
1661 (state_store.as_ref(), state_key.as_ref())
1662 {
1663 let wrapped =
1664 crate::idempotency::wrap_state(Some(&bookmark), next_seq);
1665 with_retry!("state_put", store.put(key, &wrapped))?;
1666 }
1667 last_bookmark = Some(bookmark);
1668 } else if !page.records.is_empty() {
1669 records_written +=
1673 with_retry_write!("sink_write", sink.write_batch(&page.records))?;
1674 }
1675 } else {
1676 debug_assert!(
1678 quality_envelopes.is_empty(),
1679 "quality quarantine without DLQ should have been rejected at run start"
1680 );
1681 if !page.records.is_empty() {
1682 if let Some(cfg) = adaptive_cfg.as_ref() {
1683 let ctrl = controller.get_or_insert_with(|| {
1684 crate::adaptive::AimdController::new(cfg, page.records.len())
1685 });
1686 maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
1687 let mut offset = 0;
1688 while offset < page.records.len() {
1689 let size =
1690 ctrl.current().max(1).min(page.records.len() - offset);
1691 let chunk = &page.records[offset..offset + size];
1692 let t0 = std::time::Instant::now();
1693 let n = with_retry_write!("sink_write", sink.write_batch(chunk))?;
1694 let latency = t0.elapsed();
1695 records_written += n;
1696 offset += size;
1697 let adj = ctrl.observe(crate::adaptive::Observation {
1698 batch_len: chunk.len(),
1699 errors: 0,
1700 latency,
1701 });
1702 emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
1703 }
1704 } else {
1705 records_written +=
1706 with_retry_write!("sink_write", sink.write_batch(&page.records))?;
1707 }
1708 }
1709 if let Some(bookmark) = page.bookmark {
1710 with_retry!("flush", sink.flush())?;
1711 let bm_labels =
1712 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1713 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1714 if let (Some(store), Some(key)) =
1715 (state_store.as_ref(), state_key.as_ref())
1716 {
1717 with_retry!("state_put", store.put(key, &bookmark))?;
1718 }
1719 last_bookmark = Some(bookmark);
1720 }
1721 }
1722 }
1723 Some(Err(e)) => return Err(e),
1724 None => break,
1725 }
1726 }
1727 Ok(())
1728 }
1729 .await;
1730
1731 if let Err(e) = loop_result {
1738 if let Some(ref dlq_cfg) = dlq
1739 && let Err(flush_err) = dlq_cfg.sink.flush().await
1740 {
1741 tracing::warn!(
1742 error = %flush_err,
1743 "DLQ sink flush failed during error unwind; original error preserved"
1744 );
1745 }
1746 if let Err(flush_err) = sink.flush().await {
1747 tracing::warn!(
1748 error = %flush_err,
1749 "sink flush failed during error unwind; original error preserved"
1750 );
1751 }
1752 return Err(e);
1753 }
1754
1755 if let Some(ref dlq_cfg) = dlq {
1761 let final_metric_labels: Vec<metrics::Label> = vec![
1762 metrics::Label::new(
1763 "pipeline",
1764 metrics::SharedString::from(pipeline_name.clone()),
1765 ),
1766 metrics::Label::new("row", metrics::SharedString::from(row.clone())),
1767 metrics::Label::new(
1768 "connector",
1769 metrics::SharedString::from(sink_name.to_string()),
1770 ),
1771 metrics::Label::new(
1772 "dlq_connector",
1773 metrics::SharedString::from(dlq_sink_name.to_string()),
1774 ),
1775 ];
1776 let _final_dlq_flush_timer = crate::observability::DurationGuard::new(
1777 "faucet_sink_dlq_flush_duration_seconds",
1778 final_metric_labels.clone(),
1779 );
1780 dlq_cfg.sink.flush().await.map_err(|e| {
1781 let mut lbl = final_metric_labels.clone();
1782 lbl.push(metrics::Label::new(
1783 "kind",
1784 metrics::SharedString::const_str(crate::observability::decorator::error_kind(&e)),
1785 ));
1786 metrics::counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1787 FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
1788 })?;
1789 }
1790 sink.flush().await?;
1791
1792 if cancelled {
1793 tracing::info!(
1794 records_written,
1795 "pipeline run cancelled cooperatively; sink flushed (partial output is durable)"
1796 );
1797 }
1798
1799 tracing::info!(
1800 records_written,
1801 cancelled,
1802 has_bookmark = last_bookmark.is_some(),
1803 persisted = state_store.is_some() && state_key.is_some() && last_bookmark.is_some(),
1804 dlq_records = dlq_stats.records_dlq,
1805 "pipeline streaming run complete"
1806 );
1807
1808 Ok(PipelineResult {
1809 records_written,
1810 bookmark: last_bookmark,
1811 dlq: dlq.is_some().then_some(dlq_stats),
1812 })
1813}
1814
1815fn emit_adaptive_metrics(
1818 ctrl: &crate::adaptive::AimdController,
1819 adj: Option<crate::adaptive::Adjustment>,
1820 pipeline: &str,
1821 row: &str,
1822) {
1823 use metrics::{Label, SharedString, counter, gauge};
1824 let base = vec![
1825 Label::new("pipeline", SharedString::from(pipeline.to_string())),
1826 Label::new("row", SharedString::from(row.to_string())),
1827 ];
1828 gauge!("faucet_pipeline_adaptive_batch_size", base.clone()).set(ctrl.current() as f64);
1829 gauge!(
1830 "faucet_pipeline_adaptive_batch_cooldown_active",
1831 base.clone()
1832 )
1833 .set(if ctrl.cooldown_active() { 1.0 } else { 0.0 });
1834 if let Some(p50) = ctrl.p50_latency_ms() {
1835 gauge!(
1836 "faucet_pipeline_adaptive_batch_p50_latency_ms",
1837 base.clone()
1838 )
1839 .set(p50 as f64);
1840 }
1841 if let Some(a) = adj {
1842 let mut lbl = base;
1843 lbl.push(Label::new(
1844 "direction",
1845 SharedString::const_str(a.direction.as_str()),
1846 ));
1847 lbl.push(Label::new(
1848 "reason",
1849 SharedString::const_str(a.reason.as_str()),
1850 ));
1851 counter!("faucet_pipeline_adaptive_batch_adjustments_total", lbl).increment(1);
1852 }
1853}
1854
1855fn maybe_warn_noop_sink(sink_name: &str, warned: &mut bool) {
1858 if !*warned && matches!(sink_name, "jsonl" | "csv" | "stdout") {
1859 tracing::info!(
1860 sink = sink_name,
1861 "adaptive batch sizing is a no-op for this per-record sink"
1862 );
1863 *warned = true;
1864 }
1865}
1866
1867#[allow(clippy::too_many_arguments)]
1875async fn apply_drift_policy<Si: Sink + ?Sized>(
1876 policy: &crate::drift::SchemaDriftPolicy,
1877 diff: &crate::drift::SchemaDiff,
1878 dest: &Value,
1879 records: Vec<Value>,
1880 page_indices: &[usize],
1881 sink: &Si,
1882 sink_name: &str,
1883 pipeline_name: &str,
1884 row: &str,
1885 dest_schema_cache: &mut Option<Option<Value>>,
1886 drift_envelopes: &mut Vec<Value>,
1887) -> Result<(Vec<Value>, Option<FaucetError>), FaucetError> {
1888 use crate::drift::{OnDrift, OnIncompatible};
1889 use crate::observability::schema_drift as emit_drift;
1890
1891 let mode = match policy.on_drift {
1892 OnDrift::Warn => "warn",
1893 OnDrift::Ignore => "ignore",
1894 OnDrift::Quarantine => "quarantine",
1895 OnDrift::Fail => "fail",
1896 OnDrift::Evolve => "evolve",
1897 };
1898 emit_drift(
1899 pipeline_name,
1900 row,
1901 sink_name,
1902 mode,
1903 "added",
1904 diff.additions.len() as u64,
1905 );
1906 emit_drift(
1907 pipeline_name,
1908 row,
1909 sink_name,
1910 mode,
1911 "widened",
1912 diff.widenings.len() as u64,
1913 );
1914 emit_drift(
1915 pipeline_name,
1916 row,
1917 sink_name,
1918 mode,
1919 "narrowed",
1920 diff.incompatible.len() as u64,
1921 );
1922 emit_drift(
1923 pipeline_name,
1924 row,
1925 sink_name,
1926 mode,
1927 "dropped",
1928 diff.droppable_required.len() as u64,
1929 );
1930
1931 match policy.on_drift {
1932 OnDrift::Warn => {
1933 tracing::warn!(
1934 connector = sink_name,
1935 columns = ?diff.changed_columns(),
1936 "schema-drift detected (on_drift=warn); writing page unchanged"
1937 );
1938 Ok((records, None))
1939 }
1940 OnDrift::Fail => Ok((
1941 records,
1942 Some(FaucetError::SchemaDrift {
1943 columns: diff.changed_columns(),
1944 message: "schema drift detected (on_drift=fail)".to_string(),
1945 }),
1946 )),
1947 OnDrift::Ignore => {
1948 let allowed: std::collections::HashSet<String> = dest
1950 .get("properties")
1951 .and_then(|p| p.as_object())
1952 .map(|m| m.keys().cloned().collect())
1953 .unwrap_or_default();
1954 let trimmed = records
1955 .into_iter()
1956 .map(|r| match r {
1957 Value::Object(map) => Value::Object(
1958 map.into_iter()
1959 .filter(|(k, _)| allowed.contains(k))
1960 .collect(),
1961 ),
1962 other => other,
1963 })
1964 .collect();
1965 Ok((trimmed, None))
1966 }
1967 OnDrift::Quarantine => {
1968 let (kept, env) =
1969 quarantine_drift_rows(diff, records, page_indices, sink_name, pipeline_name, row);
1970 drift_envelopes.extend(env);
1971 Ok((kept, None))
1972 }
1973 OnDrift::Evolve => {
1974 let evolution = crate::drift::SchemaEvolution {
1975 additions: diff.additions.clone(),
1976 widenings: diff
1977 .widenings
1978 .iter()
1979 .filter(|c| {
1980 c.from
1981 .as_ref()
1982 .map(|f| crate::drift::base_widened(f, &c.to))
1983 .unwrap_or(false)
1984 })
1985 .cloned()
1986 .collect(),
1987 relax_nullability: diff
1988 .droppable_required
1989 .iter()
1990 .filter(|_| policy.relax_nullability_on_missing)
1995 .cloned()
1996 .chain(
1997 diff.widenings
1998 .iter()
1999 .filter(|c| {
2000 c.from
2001 .as_ref()
2002 .map(|f| crate::drift::adds_null(f, &c.to))
2003 .unwrap_or(false)
2004 })
2005 .map(|c| c.name.clone()),
2006 )
2007 .collect(),
2008 };
2009 if !evolution.is_empty() {
2010 sink.evolve_schema(&evolution).await?;
2011 *dest_schema_cache = Some(sink.current_schema().await?);
2014 }
2015 if diff.incompatible.is_empty() {
2017 Ok((records, None))
2018 } else {
2019 match policy.on_incompatible {
2020 OnIncompatible::Fail => Ok((
2021 records,
2022 Some(FaucetError::SchemaDrift {
2023 columns: diff.incompatible.iter().map(|c| c.name.clone()).collect(),
2024 message: "incompatible type change cannot be auto-evolved \
2025 (on_incompatible=fail)"
2026 .into(),
2027 }),
2028 )),
2029 OnIncompatible::Quarantine => {
2030 let incompat_only = crate::drift::SchemaDiff {
2032 incompatible: diff.incompatible.clone(),
2033 ..Default::default()
2034 };
2035 let (kept, env) = quarantine_drift_rows(
2036 &incompat_only,
2037 records,
2038 page_indices,
2039 sink_name,
2040 pipeline_name,
2041 row,
2042 );
2043 drift_envelopes.extend(env);
2044 Ok((kept, None))
2045 }
2046 }
2047 }
2048 }
2049 }
2050}
2051
2052fn quarantine_drift_rows(
2062 diff: &crate::drift::SchemaDiff,
2063 records: Vec<Value>,
2064 page_indices: &[usize],
2065 sink_name: &str,
2066 pipeline_name: &str,
2067 row: &str,
2068) -> (Vec<Value>, Vec<Value>) {
2069 use crate::dlq::{DlqReason, build_envelope};
2070 let present_cols: std::collections::HashSet<&str> = diff
2072 .additions
2073 .iter()
2074 .chain(&diff.widenings)
2075 .chain(&diff.incompatible)
2076 .map(|c| c.name.as_str())
2077 .collect();
2078 let required_cols: std::collections::HashSet<&str> =
2080 diff.droppable_required.iter().map(|s| s.as_str()).collect();
2081 let mut kept = Vec::new();
2082 let mut envelopes = Vec::new();
2083 for (idx, rec) in records.into_iter().enumerate() {
2084 let exhibits = rec
2085 .as_object()
2086 .map(|m| {
2087 m.keys().any(|k| present_cols.contains(k.as_str()))
2088 || required_cols.iter().any(|c| !m.contains_key(*c))
2089 })
2090 .unwrap_or(false);
2091 if exhibits {
2092 let err = FaucetError::SchemaDrift {
2093 columns: diff.changed_columns(),
2094 message: "row exhibits schema drift (on_drift=quarantine)".into(),
2095 };
2096 let page_index = page_indices.get(idx).copied().unwrap_or(idx);
2099 envelopes.push(build_envelope(
2100 &rec,
2101 &err,
2102 DlqReason::SchemaDrift,
2103 sink_name,
2104 pipeline_name,
2105 row,
2106 page_index,
2107 ));
2108 } else {
2109 kept.push(rec);
2110 }
2111 }
2112 (kept, envelopes)
2113}
2114
2115#[cfg(test)]
2116mod tests {
2117 use super::*;
2118 use async_trait::async_trait;
2119 use serde_json::json;
2120
2121 struct MockSource(Vec<Value>);
2124
2125 #[async_trait]
2126 impl Source for MockSource {
2127 async fn fetch_with_context(
2128 &self,
2129 _context: &std::collections::HashMap<String, Value>,
2130 ) -> Result<Vec<Value>, FaucetError> {
2131 Ok(self.0.clone())
2132 }
2133 }
2134
2135 struct IncrementalSource {
2136 records: Vec<Value>,
2137 bookmark: Value,
2138 }
2139
2140 #[async_trait]
2141 impl Source for IncrementalSource {
2142 async fn fetch_with_context(
2143 &self,
2144 _context: &std::collections::HashMap<String, Value>,
2145 ) -> Result<Vec<Value>, FaucetError> {
2146 Ok(self.records.clone())
2147 }
2148 async fn fetch_with_context_incremental(
2149 &self,
2150 _context: &std::collections::HashMap<String, Value>,
2151 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
2152 Ok((self.records.clone(), Some(self.bookmark.clone())))
2153 }
2154 }
2155
2156 struct FailingSource;
2157
2158 #[async_trait]
2159 impl Source for FailingSource {
2160 async fn fetch_with_context(
2161 &self,
2162 _context: &std::collections::HashMap<String, Value>,
2163 ) -> Result<Vec<Value>, FaucetError> {
2164 Err(FaucetError::Auth("no credentials".into()))
2165 }
2166 }
2167
2168 struct MockSink(std::sync::Mutex<Vec<Value>>);
2171
2172 impl MockSink {
2173 fn new() -> Self {
2174 Self(std::sync::Mutex::new(Vec::new()))
2175 }
2176 fn written(&self) -> Vec<Value> {
2177 self.0.lock().unwrap().clone()
2178 }
2179 }
2180
2181 #[async_trait]
2182 impl Sink for MockSink {
2183 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2184 self.0.lock().unwrap().extend(records.iter().cloned());
2185 Ok(records.len())
2186 }
2187 }
2188
2189 struct FailingSink;
2190
2191 #[async_trait]
2192 impl Sink for FailingSink {
2193 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
2194 Err(FaucetError::Sink("write failed".into()))
2195 }
2196 }
2197
2198 struct FlushTrackingSink {
2203 written: std::sync::Mutex<Vec<Value>>,
2204 flush_count: std::sync::atomic::AtomicUsize,
2205 }
2206
2207 impl FlushTrackingSink {
2208 fn new() -> Self {
2209 Self {
2210 written: std::sync::Mutex::new(Vec::new()),
2211 flush_count: std::sync::atomic::AtomicUsize::new(0),
2212 }
2213 }
2214 fn written(&self) -> Vec<Value> {
2215 self.written.lock().unwrap().clone()
2216 }
2217 fn flush_count(&self) -> usize {
2218 self.flush_count.load(std::sync::atomic::Ordering::SeqCst)
2219 }
2220 }
2221
2222 #[async_trait]
2223 impl Sink for FlushTrackingSink {
2224 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2225 self.written.lock().unwrap().extend(records.iter().cloned());
2226 Ok(records.len())
2227 }
2228 async fn flush(&self) -> Result<(), FaucetError> {
2229 self.flush_count
2230 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2231 Ok(())
2232 }
2233 }
2234
2235 struct IdempotentMockSink {
2239 rows: std::sync::Mutex<Vec<Value>>,
2240 tokens: std::sync::Mutex<std::collections::HashMap<String, String>>,
2241 }
2242 impl IdempotentMockSink {
2243 fn new() -> Self {
2244 Self {
2245 rows: std::sync::Mutex::new(Vec::new()),
2246 tokens: std::sync::Mutex::new(std::collections::HashMap::new()),
2247 }
2248 }
2249 fn rows(&self) -> Vec<Value> {
2250 self.rows.lock().unwrap().clone()
2251 }
2252 }
2253 #[async_trait]
2254 impl Sink for IdempotentMockSink {
2255 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2256 self.rows.lock().unwrap().extend(records.iter().cloned());
2257 Ok(records.len())
2258 }
2259 fn supports_idempotent_writes(&self) -> bool {
2260 true
2261 }
2262 async fn write_batch_idempotent(
2263 &self,
2264 records: &[Value],
2265 scope: &str,
2266 token: &str,
2267 ) -> Result<usize, FaucetError> {
2268 self.rows.lock().unwrap().extend(records.iter().cloned());
2269 self.tokens
2270 .lock()
2271 .unwrap()
2272 .insert(scope.to_string(), token.to_string());
2273 Ok(records.len())
2274 }
2275 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
2276 Ok(self.tokens.lock().unwrap().get(scope).cloned())
2277 }
2278 }
2279
2280 fn eo_opts(store: Arc<dyn StateStore>, key: &str, start_seq: u64) -> RunStreamOptions {
2281 RunStreamOptions::new()
2282 .with_state(store, key)
2283 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2284 .with_start_seq(start_seq)
2285 }
2286
2287 #[tokio::test]
2288 async fn exactly_once_writes_pages_and_persists_wrapped_state() {
2289 let pages = vec![
2290 Ok(StreamPage {
2291 records: vec![json!({"id": 1})],
2292 bookmark: Some(json!("b1")),
2293 }),
2294 Ok(StreamPage {
2295 records: vec![json!({"id": 2})],
2296 bookmark: Some(json!("b2")),
2297 }),
2298 ];
2299 let sink = IdempotentMockSink::new();
2300 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2301 let r = run_stream(
2302 futures::stream::iter(pages),
2303 &sink,
2304 eo_opts(store.clone(), "k", 0),
2305 )
2306 .await
2307 .unwrap();
2308 assert_eq!(r.records_written, 2);
2309 let (bm, seq) = crate::idempotency::unwrap_state(&store.get("k").await.unwrap().unwrap());
2310 assert_eq!(bm, Some(json!("b2")));
2311 assert_eq!(seq, 2);
2312 let token = sink.last_committed_token("k").await.unwrap().unwrap();
2315 assert_eq!(
2316 crate::idempotency::parse_token_parts(&token),
2317 Some((2, Some(json!("b2"))))
2318 );
2319 }
2320
2321 #[tokio::test]
2322 async fn exactly_once_skips_already_committed_pages_on_resume() {
2323 let sink = IdempotentMockSink::new();
2324 sink.write_batch_idempotent(
2326 &[json!({"id": 1})],
2327 "k",
2328 &crate::idempotency::format_token(1),
2329 )
2330 .await
2331 .unwrap();
2332 assert_eq!(sink.rows().len(), 1);
2333 let pages = vec![
2335 Ok(StreamPage {
2336 records: vec![json!({"id": 1})],
2337 bookmark: Some(json!("b1")),
2338 }),
2339 Ok(StreamPage {
2340 records: vec![json!({"id": 2})],
2341 bookmark: Some(json!("b2")),
2342 }),
2343 ];
2344 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2345 let r = run_stream(futures::stream::iter(pages), &sink, eo_opts(store, "k", 0))
2346 .await
2347 .unwrap();
2348 assert_eq!(r.records_written, 1);
2349 let rows = sink.rows();
2350 assert_eq!(
2351 rows.len(),
2352 2,
2353 "exactly one row per id — no duplicate of id=1"
2354 );
2355 assert_eq!(rows.iter().filter(|v| v["id"] == 1).count(), 1);
2356 }
2357
2358 #[tokio::test]
2359 async fn exactly_once_rejects_non_idempotent_sink() {
2360 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2361 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2362 let r = run_stream(
2363 futures::stream::iter(pages),
2364 &MockSink::new(),
2365 eo_opts(store, "k", 0),
2366 )
2367 .await;
2368 assert!(matches!(r, Err(FaucetError::Config(_))));
2369 }
2370
2371 #[tokio::test]
2372 async fn exactly_once_rejects_missing_state_store() {
2373 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2374 let opts =
2375 RunStreamOptions::new().with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce);
2376 let r = run_stream(
2377 futures::stream::iter(pages),
2378 &IdempotentMockSink::new(),
2379 opts,
2380 )
2381 .await;
2382 assert!(matches!(r, Err(FaucetError::Config(_))));
2383 }
2384
2385 struct KeyedMockSink(MockSink);
2388 #[async_trait]
2389 impl Sink for KeyedMockSink {
2390 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2391 self.0.write_batch(records).await
2392 }
2393 fn dedups_by_key(&self) -> bool {
2394 true
2395 }
2396 fn supported_write_modes(&self) -> &'static [crate::write_mode::WriteMode] {
2397 &[
2398 crate::write_mode::WriteMode::Append,
2399 crate::write_mode::WriteMode::Upsert,
2400 crate::write_mode::WriteMode::Delete,
2401 ]
2402 }
2403 }
2404
2405 #[tokio::test]
2406 async fn exactly_once_keyed_upsert_mechanism_uses_plain_write_path() {
2407 let sink = KeyedMockSink(MockSink::new());
2412 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2413 let pages = vec![Ok(StreamPage {
2414 records: vec![json!({"id": 1})],
2415 bookmark: Some(json!("b1")),
2416 })];
2417 let opts = eo_opts(store.clone(), "k", 0)
2418 .with_replay_guarantee(crate::idempotency::ReplayGuarantee::NonDeterministic);
2419 let r = run_stream(futures::stream::iter(pages), &sink, opts)
2420 .await
2421 .unwrap();
2422 assert_eq!(r.records_written, 1);
2423 assert_eq!(sink.0.written(), vec![json!({"id": 1})]);
2424 assert_eq!(store.get("k").await.unwrap(), Some(json!("b1")));
2426 }
2427
2428 #[tokio::test]
2429 async fn exactly_once_rejects_non_deterministic_source_without_keyed_dedup() {
2430 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2435 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2436 let opts = eo_opts(store, "k", 0)
2437 .with_replay_guarantee(crate::idempotency::ReplayGuarantee::NonDeterministic);
2438 let r = run_stream(
2439 futures::stream::iter(pages),
2440 &IdempotentMockSink::new(),
2441 opts,
2442 )
2443 .await;
2444 match r {
2445 Err(FaucetError::Config(msg)) => {
2446 assert!(msg.contains("write_mode: upsert"), "hint present: {msg}")
2447 }
2448 other => panic!("expected Config error, got {other:?}"),
2449 }
2450 }
2451
2452 #[tokio::test]
2453 async fn exactly_once_rejects_plain_sink_with_mechanism_hint() {
2454 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2455 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2456 let r = run_stream(
2457 futures::stream::iter(pages),
2458 &MockSink::new(),
2459 eo_opts(store, "k", 0),
2460 )
2461 .await;
2462 match r {
2463 Err(FaucetError::Config(msg)) => assert!(
2464 msg.contains("provides neither"),
2465 "names both mechanisms: {msg}"
2466 ),
2467 other => panic!("expected Config error, got {other:?}"),
2468 }
2469 }
2470
2471 struct AnchorRecordingSource {
2474 applied: std::sync::Mutex<Vec<Value>>,
2475 }
2476 #[async_trait]
2477 impl Source for AnchorRecordingSource {
2478 async fn fetch_with_context(
2479 &self,
2480 _context: &std::collections::HashMap<String, Value>,
2481 ) -> Result<Vec<Value>, FaucetError> {
2482 Ok(vec![json!({"id": 10})])
2483 }
2484 async fn fetch_with_context_incremental(
2485 &self,
2486 _context: &std::collections::HashMap<String, Value>,
2487 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
2488 Ok((vec![json!({"id": 10})], Some(json!("after"))))
2489 }
2490 fn state_key(&self) -> Option<String> {
2491 Some("anchor_key".to_string())
2492 }
2493 async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
2494 self.applied.lock().unwrap().push(bookmark);
2495 Ok(())
2496 }
2497 fn supports_exactly_once(&self) -> bool {
2498 true
2499 }
2500 }
2501
2502 #[tokio::test]
2503 async fn pipeline_run_anchors_resume_at_sink_embedded_bookmark() {
2504 let source = AnchorRecordingSource {
2511 applied: std::sync::Mutex::new(Vec::new()),
2512 };
2513 let sink = IdempotentMockSink::new();
2514 sink.tokens.lock().unwrap().insert(
2515 "anchor_key".to_string(),
2516 crate::idempotency::format_token_with_bookmark(5, Some(&json!("sink-pos"))),
2517 );
2518 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2519 store
2520 .put(
2521 "anchor_key",
2522 &crate::idempotency::wrap_state(Some(&json!("state-pos")), 4),
2523 )
2524 .await
2525 .unwrap();
2526
2527 let r = Pipeline::new(&source, &sink)
2528 .with_state_store(Arc::clone(&store))
2529 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2530 .run()
2531 .await
2532 .unwrap();
2533
2534 assert_eq!(
2535 *source.applied.lock().unwrap(),
2536 vec![json!("state-pos"), json!("sink-pos")],
2537 "state bookmark applied, then overridden by the sink anchor"
2538 );
2539 assert_eq!(r.records_written, 1);
2542 let token = sink
2543 .last_committed_token("anchor_key")
2544 .await
2545 .unwrap()
2546 .unwrap();
2547 assert_eq!(
2548 crate::idempotency::parse_token_parts(&token),
2549 Some((6, Some(json!("after"))))
2550 );
2551 let (bm, seq) =
2552 crate::idempotency::unwrap_state(&store.get("anchor_key").await.unwrap().unwrap());
2553 assert_eq!((bm, seq), (Some(json!("after")), 6));
2554 }
2555
2556 #[tokio::test]
2557 async fn pipeline_run_ignores_sink_token_behind_state_seq() {
2558 let source = AnchorRecordingSource {
2561 applied: std::sync::Mutex::new(Vec::new()),
2562 };
2563 let sink = IdempotentMockSink::new();
2564 sink.tokens.lock().unwrap().insert(
2565 "anchor_key".to_string(),
2566 crate::idempotency::format_token_with_bookmark(4, Some(&json!("sink-pos"))),
2567 );
2568 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2569 store
2570 .put(
2571 "anchor_key",
2572 &crate::idempotency::wrap_state(Some(&json!("state-pos")), 4),
2573 )
2574 .await
2575 .unwrap();
2576 Pipeline::new(&source, &sink)
2577 .with_state_store(Arc::clone(&store))
2578 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2579 .run()
2580 .await
2581 .unwrap();
2582 assert_eq!(*source.applied.lock().unwrap(), vec![json!("state-pos")]);
2583 }
2584
2585 #[tokio::test]
2586 async fn pipeline_run_legacy_bare_token_falls_back_to_skip_path() {
2587 let source = AnchorRecordingSource {
2591 applied: std::sync::Mutex::new(Vec::new()),
2592 };
2593 let sink = IdempotentMockSink::new();
2594 sink.tokens.lock().unwrap().insert(
2595 "anchor_key".to_string(),
2596 crate::idempotency::format_token(1),
2597 );
2598 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2599 let r = Pipeline::new(&source, &sink)
2600 .with_state_store(Arc::clone(&store))
2601 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
2602 .run()
2603 .await
2604 .unwrap();
2605 assert!(source.applied.lock().unwrap().is_empty());
2606 assert_eq!(r.records_written, 0, "page 1 already committed → skipped");
2607 assert!(sink.rows().is_empty());
2608 }
2609
2610 #[test]
2613 fn stream_page_constructs() {
2614 let page = StreamPage {
2615 records: vec![json!({"id": 1})],
2616 bookmark: Some(json!("2026-05-18")),
2617 };
2618 assert_eq!(page.records.len(), 1);
2619 assert_eq!(page.bookmark, Some(json!("2026-05-18")));
2620 }
2621
2622 #[test]
2623 fn validate_batch_size_accepts_zero_as_no_batching_sentinel() {
2624 assert_eq!(validate_batch_size(0).unwrap(), 0);
2626 }
2627
2628 #[test]
2629 fn validate_batch_size_rejects_too_large() {
2630 let err = validate_batch_size(MAX_BATCH_SIZE + 1).unwrap_err();
2631 assert!(matches!(err, FaucetError::Config(_)));
2632 }
2633
2634 #[test]
2635 fn validate_batch_size_accepts_one() {
2636 assert_eq!(validate_batch_size(1).unwrap(), 1);
2637 }
2638
2639 #[test]
2640 fn validate_batch_size_accepts_max() {
2641 assert_eq!(validate_batch_size(MAX_BATCH_SIZE).unwrap(), MAX_BATCH_SIZE);
2642 }
2643
2644 const _: () = {
2646 assert!(DEFAULT_BATCH_SIZE >= 1);
2647 assert!(DEFAULT_BATCH_SIZE <= MAX_BATCH_SIZE);
2648 };
2649
2650 #[tokio::test]
2653 async fn batch_pipeline_writes_all_records() {
2654 let source = MockSource(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
2655 let sink = MockSink::new();
2656
2657 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2658
2659 assert_eq!(result.records_written, 3);
2660 assert!(result.bookmark.is_none());
2661 assert_eq!(sink.written().len(), 3);
2662 }
2663
2664 #[tokio::test]
2665 async fn batch_pipeline_returns_bookmark() {
2666 let source = IncrementalSource {
2667 records: vec![json!({"id": 1, "ts": "2024-12-01"})],
2668 bookmark: json!("2024-12-01"),
2669 };
2670 let sink = MockSink::new();
2671
2672 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2673
2674 assert_eq!(result.records_written, 1);
2675 assert_eq!(result.bookmark, Some(json!("2024-12-01")));
2676 }
2677
2678 #[tokio::test]
2679 async fn batch_pipeline_empty_source() {
2680 let source = MockSource(vec![]);
2681 let sink = MockSink::new();
2682
2683 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2684
2685 assert_eq!(result.records_written, 0);
2686 assert!(sink.written().is_empty());
2687 }
2688
2689 #[tokio::test]
2690 async fn batch_pipeline_source_error_propagates() {
2691 let source = FailingSource;
2692 let sink = MockSink::new();
2693
2694 let result = Pipeline::new(&source, &sink).run().await;
2695 assert!(result.is_err());
2696 assert!(sink.written().is_empty());
2697 }
2698
2699 #[tokio::test]
2700 async fn batch_pipeline_sink_error_propagates() {
2701 let source = MockSource(vec![json!({"id": 1})]);
2702 let sink = FailingSink;
2703
2704 let result = Pipeline::new(&source, &sink).run().await;
2705 assert!(result.is_err());
2706 }
2707
2708 #[tokio::test]
2709 async fn batch_pipeline_with_trait_objects() {
2710 let source: Box<dyn Source> = Box::new(MockSource(vec![json!({"id": 1})]));
2711 let sink: Box<dyn Sink> = Box::new(MockSink::new());
2712
2713 let result = Pipeline::new(source.as_ref(), sink.as_ref())
2714 .run()
2715 .await
2716 .unwrap();
2717
2718 assert_eq!(result.records_written, 1);
2719 }
2720
2721 #[tokio::test]
2724 async fn stream_pipeline_writes_pages() {
2725 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2726 Ok(StreamPage {
2727 records: vec![json!({"id": 1}), json!({"id": 2})],
2728 bookmark: None,
2729 }),
2730 Ok(StreamPage {
2731 records: vec![json!({"id": 3})],
2732 bookmark: None,
2733 }),
2734 ];
2735 let stream = futures::stream::iter(pages);
2736 let sink = MockSink::new();
2737
2738 let result = run_stream(stream, &sink, RunStreamOptions::new())
2739 .await
2740 .unwrap();
2741
2742 assert_eq!(result.records_written, 3);
2743 assert!(result.bookmark.is_none());
2744 assert_eq!(sink.written().len(), 3);
2745 }
2746
2747 #[tokio::test]
2748 async fn stream_pipeline_flushes_sink_on_source_error() {
2749 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2753 Ok(StreamPage {
2754 records: vec![json!({"id": 1}), json!({"id": 2})],
2755 bookmark: None,
2756 }),
2757 Err(FaucetError::Source("transient blip mid-stream".into())),
2758 ];
2759 let stream = futures::stream::iter(pages);
2760 let sink = FlushTrackingSink::new();
2761
2762 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2763
2764 assert!(matches!(result, Err(FaucetError::Source(_))));
2766 assert_eq!(sink.written().len(), 2);
2768 assert!(
2770 sink.flush_count() >= 1,
2771 "sink must be flushed on the error path so partial output is durable"
2772 );
2773 }
2774
2775 #[tokio::test]
2776 async fn stream_pipeline_flushes_sink_on_cancel() {
2777 use tokio_util::sync::CancellationToken;
2782
2783 let stream = Box::pin(async_stream::stream! {
2785 yield Ok(StreamPage {
2786 records: vec![json!({"id": 1}), json!({"id": 2})],
2787 bookmark: None,
2788 });
2789 futures::future::pending::<()>().await;
2790 });
2791 let sink = FlushTrackingSink::new();
2792
2793 let token = CancellationToken::new();
2794 let canceller = token.clone();
2795 tokio::spawn(async move {
2796 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
2797 canceller.cancel();
2798 });
2799
2800 let result = run_stream(stream, &sink, RunStreamOptions::new().with_cancel(token))
2801 .await
2802 .expect("a cooperative cancel returns Ok with the partial result");
2803
2804 assert_eq!(result.records_written, 2);
2807 assert_eq!(sink.written().len(), 2);
2808 assert!(
2809 sink.flush_count() >= 1,
2810 "sink must be flushed on the cancel path so partial output is durable"
2811 );
2812 }
2813
2814 #[tokio::test]
2815 async fn stream_pipeline_empty() {
2816 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2817 let stream = futures::stream::iter(pages);
2818 let sink = MockSink::new();
2819
2820 let result = run_stream(stream, &sink, RunStreamOptions::new())
2821 .await
2822 .unwrap();
2823
2824 assert_eq!(result.records_written, 0);
2825 }
2826
2827 #[tokio::test]
2828 async fn stream_pipeline_skips_empty_pages() {
2829 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2830 Ok(StreamPage {
2831 records: vec![json!({"id": 1})],
2832 bookmark: None,
2833 }),
2834 Ok(StreamPage {
2835 records: vec![],
2836 bookmark: None,
2837 }),
2838 Ok(StreamPage {
2839 records: vec![json!({"id": 2})],
2840 bookmark: None,
2841 }),
2842 ];
2843 let stream = futures::stream::iter(pages);
2844 let sink = MockSink::new();
2845
2846 let result = run_stream(stream, &sink, RunStreamOptions::new())
2847 .await
2848 .unwrap();
2849
2850 assert_eq!(result.records_written, 2);
2851 }
2852
2853 #[tokio::test]
2854 async fn stream_pipeline_error_in_page_propagates() {
2855 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2856 Ok(StreamPage {
2857 records: vec![json!({"id": 1})],
2858 bookmark: None,
2859 }),
2860 Err(FaucetError::HttpStatus {
2861 status: 500,
2862 url: "https://example.com".into(),
2863 body: "Internal Server Error".into(),
2864 }),
2865 ];
2866 let stream = futures::stream::iter(pages);
2867 let sink = MockSink::new();
2868
2869 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2870 assert!(result.is_err());
2871 assert_eq!(sink.written().len(), 1);
2873 }
2874
2875 #[tokio::test]
2876 async fn stream_pipeline_sink_error_propagates() {
2877 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2878 records: vec![json!({"id": 1})],
2879 bookmark: None,
2880 })];
2881 let stream = futures::stream::iter(pages);
2882 let sink = FailingSink;
2883
2884 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2885 assert!(result.is_err());
2886 }
2887
2888 #[tokio::test]
2889 async fn stream_pipeline_with_trait_object_sink() {
2890 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2891 records: vec![json!({"id": 1})],
2892 bookmark: None,
2893 })];
2894 let stream = futures::stream::iter(pages);
2895 let sink: Box<dyn Sink> = Box::new(MockSink::new());
2896
2897 let result = run_stream(stream, sink.as_ref(), RunStreamOptions::new())
2898 .await
2899 .unwrap();
2900 assert_eq!(result.records_written, 1);
2901 }
2902
2903 #[tokio::test]
2904 async fn stream_pipeline_persists_bookmark_when_page_carries_one() {
2905 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2906 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2907 Ok(StreamPage {
2908 records: vec![json!({"id": 1})],
2909 bookmark: None,
2910 }),
2911 Ok(StreamPage {
2912 records: vec![json!({"id": 2})],
2913 bookmark: Some(json!("checkpoint-final")),
2914 }),
2915 ];
2916 let stream = futures::stream::iter(pages);
2917 let sink = MockSink::new();
2918
2919 let result = run_stream(
2920 stream,
2921 &sink,
2922 RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
2923 )
2924 .await
2925 .unwrap();
2926
2927 assert_eq!(result.records_written, 2);
2928 assert_eq!(result.bookmark, Some(json!("checkpoint-final")));
2929 assert_eq!(
2930 store.get("k").await.unwrap(),
2931 Some(json!("checkpoint-final"))
2932 );
2933 }
2934
2935 #[tokio::test]
2936 async fn stream_pipeline_persists_per_page_bookmarks() {
2937 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2938 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2939 Ok(StreamPage {
2940 records: vec![json!({"id": 1})],
2941 bookmark: Some(json!("tx-1")),
2942 }),
2943 Ok(StreamPage {
2944 records: vec![json!({"id": 2})],
2945 bookmark: Some(json!("tx-2")),
2946 }),
2947 ];
2948 let stream = futures::stream::iter(pages);
2949 let sink = MockSink::new();
2950
2951 run_stream(
2952 stream,
2953 &sink,
2954 RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
2955 )
2956 .await
2957 .unwrap();
2958
2959 assert_eq!(store.get("k").await.unwrap(), Some(json!("tx-2")));
2961 }
2962
2963 use crate::state::{FileStateStore, MemoryStateStore, StateStore};
2966 use std::sync::Arc;
2967 use tempfile::TempDir;
2968
2969 struct StatefulSource {
2973 key: String,
2974 records: Vec<Value>,
2975 new_bookmark: Value,
2976 seen_bookmark: std::sync::Mutex<Option<Value>>,
2977 }
2978
2979 impl StatefulSource {
2980 fn new(key: &str, records: Vec<Value>, new_bookmark: Value) -> Self {
2981 Self {
2982 key: key.into(),
2983 records,
2984 new_bookmark,
2985 seen_bookmark: std::sync::Mutex::new(None),
2986 }
2987 }
2988 fn observed_start(&self) -> Option<Value> {
2989 self.seen_bookmark.lock().unwrap().clone()
2990 }
2991 }
2992
2993 #[async_trait]
2994 impl Source for StatefulSource {
2995 async fn fetch_with_context(
2996 &self,
2997 _ctx: &std::collections::HashMap<String, Value>,
2998 ) -> Result<Vec<Value>, FaucetError> {
2999 Ok(self.records.clone())
3000 }
3001 async fn fetch_with_context_incremental(
3002 &self,
3003 _ctx: &std::collections::HashMap<String, Value>,
3004 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
3005 Ok((self.records.clone(), Some(self.new_bookmark.clone())))
3006 }
3007 fn state_key(&self) -> Option<String> {
3008 Some(self.key.clone())
3009 }
3010 async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
3011 *self.seen_bookmark.lock().unwrap() = Some(bookmark);
3012 Ok(())
3013 }
3014 }
3015
3016 #[tokio::test]
3017 async fn pipeline_with_state_store_persists_bookmark_after_sink() {
3018 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3019 let source = StatefulSource::new(
3020 "github_issues",
3021 vec![json!({"id": 1, "ts": "2026-05-01"})],
3022 json!("2026-05-01"),
3023 );
3024 let sink = MockSink::new();
3025 let result = Pipeline::new(&source, &sink)
3026 .with_state_store(Arc::clone(&store))
3027 .run()
3028 .await
3029 .unwrap();
3030
3031 assert_eq!(result.records_written, 1);
3032 assert_eq!(result.bookmark, Some(json!("2026-05-01")));
3033 let stored = store.get("github_issues").await.unwrap();
3035 assert_eq!(stored, Some(json!("2026-05-01")));
3036 }
3037
3038 #[tokio::test]
3039 async fn pipeline_with_state_store_resumes_from_stored_bookmark() {
3040 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3041 store
3042 .put("github_issues", &json!("2026-04-30"))
3043 .await
3044 .unwrap();
3045
3046 let source =
3047 StatefulSource::new("github_issues", vec![json!({"id": 2})], json!("2026-05-01"));
3048 let sink = MockSink::new();
3049 Pipeline::new(&source, &sink)
3050 .with_state_store(Arc::clone(&store))
3051 .run()
3052 .await
3053 .unwrap();
3054
3055 assert_eq!(source.observed_start(), Some(json!("2026-04-30")));
3057 assert_eq!(
3059 store.get("github_issues").await.unwrap(),
3060 Some(json!("2026-05-01"))
3061 );
3062 }
3063
3064 #[tokio::test]
3065 async fn pipeline_with_state_store_does_not_persist_when_sink_fails() {
3066 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3067 let source = StatefulSource::new("k", vec![json!({"id": 1})], json!("2026-05-01"));
3068 let sink = FailingSink;
3069
3070 let result = Pipeline::new(&source, &sink)
3071 .with_state_store(Arc::clone(&store))
3072 .run()
3073 .await;
3074 assert!(result.is_err());
3075 assert!(store.get("k").await.unwrap().is_none());
3076 }
3077
3078 #[tokio::test]
3079 async fn pipeline_with_state_store_no_state_key_means_no_persist() {
3080 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3081 let source = IncrementalSource {
3082 records: vec![json!({"id": 1})],
3083 bookmark: json!("ignored"),
3084 };
3085 let sink = MockSink::new();
3086 Pipeline::new(&source, &sink)
3087 .with_state_store(Arc::clone(&store))
3088 .run()
3089 .await
3090 .unwrap();
3091 assert!(store.get("anything").await.unwrap().is_none());
3094 }
3095
3096 #[tokio::test]
3097 async fn pipeline_with_state_store_skips_persist_when_bookmark_is_none() {
3098 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3099 struct NoBookmarkSource;
3100 #[async_trait]
3101 impl Source for NoBookmarkSource {
3102 async fn fetch_with_context(
3103 &self,
3104 _ctx: &std::collections::HashMap<String, Value>,
3105 ) -> Result<Vec<Value>, FaucetError> {
3106 Ok(vec![json!({"id": 1})])
3107 }
3108 fn state_key(&self) -> Option<String> {
3109 Some("k".into())
3110 }
3111 }
3112 let source = NoBookmarkSource;
3113 let sink = MockSink::new();
3114 Pipeline::new(&source, &sink)
3115 .with_state_store(Arc::clone(&store))
3116 .run()
3117 .await
3118 .unwrap();
3119 assert!(store.get("k").await.unwrap().is_none());
3120 }
3121
3122 struct PagedSource;
3127
3128 #[async_trait]
3129 impl Source for PagedSource {
3130 async fn fetch_with_context(
3131 &self,
3132 _ctx: &std::collections::HashMap<String, Value>,
3133 ) -> Result<Vec<Value>, FaucetError> {
3134 unreachable!("Pipeline::run must drive stream_pages, not fetch_with_context");
3136 }
3137 fn stream_pages<'a>(
3138 &'a self,
3139 _ctx: &'a std::collections::HashMap<String, Value>,
3140 _batch_size: usize,
3141 ) -> std::pin::Pin<
3142 Box<dyn futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>,
3143 > {
3144 Box::pin(async_stream::try_stream! {
3145 yield StreamPage { records: vec![json!({"i": 1})], bookmark: None };
3146 yield StreamPage { records: vec![json!({"i": 2})], bookmark: None };
3147 yield StreamPage { records: vec![json!({"i": 3})], bookmark: Some(json!("final")) };
3148 })
3149 }
3150 }
3151
3152 struct CountingSink {
3154 calls: std::sync::Mutex<Vec<usize>>,
3155 }
3156
3157 impl CountingSink {
3158 fn new() -> Self {
3159 Self {
3160 calls: std::sync::Mutex::new(Vec::new()),
3161 }
3162 }
3163 fn call_count(&self) -> usize {
3164 self.calls.lock().unwrap().len()
3165 }
3166 }
3167
3168 #[async_trait]
3169 impl Sink for CountingSink {
3170 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3171 self.calls.lock().unwrap().push(records.len());
3172 Ok(records.len())
3173 }
3174 }
3175
3176 #[tokio::test]
3177 async fn pipeline_run_drives_stream_pages() {
3178 let source = PagedSource;
3179 let sink = CountingSink::new();
3180
3181 let result = Pipeline::new(&source, &sink).run().await.unwrap();
3182
3183 assert_eq!(sink.call_count(), 3);
3185 assert_eq!(result.records_written, 3);
3186 assert_eq!(result.bookmark, Some(json!("final")));
3187 }
3188
3189 #[tokio::test]
3190 async fn pipeline_with_file_state_store_round_trips_across_runs() {
3191 let dir = TempDir::new().unwrap();
3192 let store: Arc<dyn StateStore> = Arc::new(FileStateStore::new(dir.path()));
3193
3194 let s1 = StatefulSource::new("k", vec![json!({"i": 1})], json!("v1"));
3196 let sink1 = MockSink::new();
3197 Pipeline::new(&s1, &sink1)
3198 .with_state_store(Arc::clone(&store))
3199 .run()
3200 .await
3201 .unwrap();
3202 assert_eq!(s1.observed_start(), None);
3203 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3204
3205 let s2 = StatefulSource::new("k", vec![json!({"i": 2})], json!("v2"));
3207 let sink2 = MockSink::new();
3208 Pipeline::new(&s2, &sink2)
3209 .with_state_store(Arc::clone(&store))
3210 .run()
3211 .await
3212 .unwrap();
3213 assert_eq!(s2.observed_start(), Some(json!("v1")));
3214 assert_eq!(store.get("k").await.unwrap(), Some(json!("v2")));
3215 }
3216
3217 #[tokio::test]
3218 #[allow(clippy::await_holding_lock)]
3219 async fn pipeline_run_increments_runs_total() {
3220 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3221 use metrics_util::debugging::DebugValue;
3222 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3223 let snap = snapshotter();
3224
3225 let source = MockSource(vec![json!({"i": 1})]);
3226 let sink = MockSink::new();
3227 let _ = Pipeline::new(&source, &sink)
3228 .with_name("test-pipeline")
3229 .with_row("rowA")
3230 .run()
3231 .await
3232 .unwrap();
3233
3234 let snapshot = snap.snapshot();
3235 let found = snapshot.into_vec().into_iter().any(
3236 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3237 key.key().name() == "faucet_pipeline_runs_total"
3238 && key.key().labels().any(|l: &metrics::Label| {
3239 l.key() == "pipeline" && l.value() == "test-pipeline"
3240 })
3241 && key
3242 .key()
3243 .labels()
3244 .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowA")
3245 && key
3246 .key()
3247 .labels()
3248 .any(|l: &metrics::Label| l.key() == "status" && l.value() == "ok")
3249 && matches!(v, DebugValue::Counter(c) if c >= 1)
3250 },
3251 );
3252 assert!(
3253 found,
3254 "expected faucet_pipeline_runs_total{{pipeline=test-pipeline, row=rowA, status=ok}}"
3255 );
3256 }
3257
3258 #[tokio::test]
3259 #[allow(clippy::await_holding_lock)]
3260 async fn pipeline_failure_attaches_kind_label_to_runs_total() {
3261 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3262 use metrics_util::debugging::DebugValue;
3263 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3264 let snap = snapshotter();
3265
3266 let source = FailingSource;
3267 let sink = MockSink::new();
3268 let _ = Pipeline::new(&source, &sink)
3269 .with_name("err-pipeline")
3270 .with_row("rowE")
3271 .run()
3272 .await;
3273
3274 let snapshot = snap.snapshot();
3275 let found = snapshot.into_vec().into_iter().any(
3276 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3277 key.key().name() == "faucet_pipeline_runs_total"
3278 && key.key().labels().any(|l: &metrics::Label| {
3279 l.key() == "pipeline" && l.value() == "err-pipeline"
3280 })
3281 && key
3282 .key()
3283 .labels()
3284 .any(|l: &metrics::Label| l.key() == "status" && l.value() == "err")
3285 && key
3286 .key()
3287 .labels()
3288 .any(|l: &metrics::Label| l.key() == "kind" && l.value() == "Auth")
3289 && matches!(v, DebugValue::Counter(c) if c >= 1)
3290 },
3291 );
3292 assert!(
3293 found,
3294 "expected faucet_pipeline_runs_total{{status=err, kind=Auth}} for failing source"
3295 );
3296 }
3297
3298 #[tokio::test]
3299 #[allow(clippy::await_holding_lock)]
3300 async fn pipeline_run_emits_start_time_gauge() {
3301 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3302 use metrics_util::debugging::DebugValue;
3303 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3304 let snap = snapshotter();
3305
3306 let source = MockSource(vec![json!({"i": 1})]);
3307 let sink = MockSink::new();
3308 let before = std::time::SystemTime::now()
3309 .duration_since(std::time::UNIX_EPOCH)
3310 .map(|d| d.as_secs_f64())
3311 .unwrap_or(0.0);
3312 let _ = Pipeline::new(&source, &sink)
3313 .with_name("start-time-pipeline")
3314 .with_row("rowS")
3315 .run()
3316 .await
3317 .unwrap();
3318
3319 let snapshot = snap.snapshot();
3320 let found = snapshot.into_vec().into_iter().any(
3321 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3322 if key.key().name() != "faucet_pipeline_start_time_unix_seconds" {
3323 return false;
3324 }
3325 let labels_match = key.key().labels().any(|l: &metrics::Label| {
3326 l.key() == "pipeline" && l.value() == "start-time-pipeline"
3327 }) && key
3328 .key()
3329 .labels()
3330 .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowS");
3331 if !labels_match {
3332 return false;
3333 }
3334 matches!(v, DebugValue::Gauge(g) if g.into_inner() >= before)
3335 },
3336 );
3337 assert!(
3338 found,
3339 "expected faucet_pipeline_start_time_unix_seconds gauge >= test-start timestamp"
3340 );
3341 }
3342
3343 #[tokio::test]
3344 #[allow(clippy::await_holding_lock)]
3345 async fn register_build_info_sets_version_gauge() {
3346 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3347 use metrics_util::debugging::DebugValue;
3348 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3349 let snap = snapshotter();
3350
3351 crate::observability::register_build_info();
3352
3353 let snapshot = snap.snapshot();
3354 let found = snapshot.into_vec().into_iter().any(
3355 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
3356 key.key().name() == "faucet_build_info"
3357 && key.key().labels().any(|l: &metrics::Label| {
3358 l.key() == "version" && l.value() == env!("CARGO_PKG_VERSION")
3359 })
3360 && matches!(v, DebugValue::Gauge(g) if (g.into_inner() - 1.0).abs() < f64::EPSILON)
3361 },
3362 );
3363 assert!(
3364 found,
3365 "expected faucet_build_info{{version=CARGO_PKG_VERSION}} = 1.0 after register_build_info()"
3366 );
3367 }
3368
3369 use crate::dlq::{DlqConfig, OnBatchError};
3372
3373 struct PartialSink {
3377 fail_indices: std::sync::Mutex<Vec<usize>>,
3378 committed: std::sync::Mutex<Vec<Value>>,
3379 }
3380
3381 impl PartialSink {
3382 fn new(fail_indices: Vec<usize>) -> Self {
3383 Self {
3384 fail_indices: std::sync::Mutex::new(fail_indices),
3385 committed: std::sync::Mutex::new(Vec::new()),
3386 }
3387 }
3388 }
3389
3390 #[async_trait]
3391 impl Sink for PartialSink {
3392 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
3393 unreachable!("PartialSink only overrides write_batch_partial");
3394 }
3395 async fn write_batch_partial(
3396 &self,
3397 records: &[Value],
3398 ) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
3399 let fails: std::collections::HashSet<usize> =
3400 self.fail_indices.lock().unwrap().iter().copied().collect();
3401 let mut outcomes = Vec::with_capacity(records.len());
3402 for (i, rec) in records.iter().enumerate() {
3403 if fails.contains(&i) {
3404 outcomes.push(Err(FaucetError::Sink(format!("row {i} rejected"))));
3405 } else {
3406 self.committed.lock().unwrap().push(rec.clone());
3407 outcomes.push(Ok(()));
3408 }
3409 }
3410 Ok(outcomes)
3411 }
3412 }
3413
3414 #[tokio::test]
3415 async fn dlq_routes_only_failed_rows_for_partial_success_sink() {
3416 let main = PartialSink::new(vec![1, 3]); let dlq = std::sync::Arc::new(MockSink::new());
3418 let dlq_cfg = DlqConfig::new(dlq.clone());
3419
3420 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3421 records: (0..4).map(|i| json!({"i": i})).collect(),
3422 bookmark: None,
3423 })];
3424 let stream = futures::stream::iter(pages);
3425 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg))
3426 .await
3427 .unwrap();
3428
3429 assert_eq!(result.records_written, 2); assert_eq!(main.committed.lock().unwrap().len(), 2);
3431 let envelopes = dlq.0.lock().unwrap();
3432 assert_eq!(envelopes.len(), 2);
3433 assert_eq!(envelopes[0]["payload"]["i"], 1);
3434 assert_eq!(envelopes[0]["record_index"], 1);
3435 assert_eq!(envelopes[1]["payload"]["i"], 3);
3436 assert_eq!(envelopes[1]["record_index"], 3);
3437 let stats = result.dlq.unwrap();
3438 assert_eq!(stats.records_dlq, 2);
3439 assert_eq!(stats.pages_with_failures, 1);
3440 }
3441
3442 #[cfg(feature = "masking")]
3443 #[tokio::test]
3444 async fn masking_runs_before_the_sink() {
3445 use crate::masking::{CompiledMasking, MaskingSpec};
3446 let sink = MockSink::new();
3447 let spec: MaskingSpec = serde_json::from_value(json!({
3448 "rules": [{ "match": { "value_detector": "email" },
3449 "action": { "type": "redact" } }]
3450 }))
3451 .unwrap();
3452 let m = std::sync::Arc::new(CompiledMasking::compile(&spec).unwrap());
3453 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3454 records: vec![json!({"email": "a@b.com", "name": "Al"})],
3455 bookmark: None,
3456 })];
3457 run_stream(
3458 futures::stream::iter(pages),
3459 &sink,
3460 RunStreamOptions::new().with_masking(m),
3461 )
3462 .await
3463 .unwrap();
3464 assert_eq!(sink.written()[0], json!({"email": "***", "name": "Al"}));
3465 }
3466
3467 #[cfg(feature = "masking")]
3468 #[tokio::test]
3469 async fn masking_applies_before_the_dlq_envelope() {
3470 use crate::masking::{CompiledMasking, MaskingSpec};
3474 let main = PartialSink::new(vec![0]); let dlq = std::sync::Arc::new(MockSink::new());
3476 let spec: MaskingSpec = serde_json::from_value(json!({
3477 "rules": [{ "match": { "value_detector": "email" },
3478 "action": { "type": "redact" } }]
3479 }))
3480 .unwrap();
3481 let m = std::sync::Arc::new(CompiledMasking::compile(&spec).unwrap());
3482 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3483 records: vec![
3484 json!({"email": "secret@x.com"}),
3485 json!({"email": "ok@y.com"}),
3486 ],
3487 bookmark: None,
3488 })];
3489 let opts = RunStreamOptions::new()
3490 .with_masking(m)
3491 .with_dlq(DlqConfig::new(dlq.clone()));
3492 run_stream(futures::stream::iter(pages), &main, opts)
3493 .await
3494 .unwrap();
3495 let env = dlq.0.lock().unwrap();
3497 assert_eq!(env.len(), 1);
3498 assert_eq!(
3499 env[0]["payload"]["email"], "***",
3500 "the DLQ payload must be masked, not raw PII"
3501 );
3502 assert_eq!(main.committed.lock().unwrap()[0]["email"], "***");
3503 }
3504
3505 #[tokio::test]
3506 async fn dlq_propagate_policy_bubbles_outer_err() {
3507 let main = FailingSink;
3508 let dlq = std::sync::Arc::new(MockSink::new());
3509 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3510 dlq_cfg.on_batch_error = OnBatchError::Propagate;
3511
3512 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3513 records: vec![json!({"i": 0}), json!({"i": 1})],
3514 bookmark: Some(json!("v1")),
3515 })];
3516 let stream = futures::stream::iter(pages);
3517 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3518 let result = run_stream(
3519 stream,
3520 &main,
3521 RunStreamOptions::new()
3522 .with_dlq(dlq_cfg)
3523 .with_state(std::sync::Arc::clone(&store), "k"),
3524 )
3525 .await;
3526 assert!(matches!(result, Err(FaucetError::Sink(_))));
3527 assert!(dlq.0.lock().unwrap().is_empty());
3528 assert!(store.get("k").await.unwrap().is_none());
3530 }
3531
3532 #[tokio::test]
3533 async fn dlq_dlq_all_policy_routes_every_row_on_outer_err() {
3534 let main = FailingSink;
3535 let dlq = std::sync::Arc::new(MockSink::new());
3536 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3537 dlq_cfg.on_batch_error = OnBatchError::DlqAll;
3538
3539 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3540 records: vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})],
3541 bookmark: Some(json!("v1")),
3542 })];
3543 let stream = futures::stream::iter(pages);
3544 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3545 let result = run_stream(
3546 stream,
3547 &main,
3548 RunStreamOptions::new()
3549 .with_dlq(dlq_cfg)
3550 .with_state(std::sync::Arc::clone(&store), "k"),
3551 )
3552 .await
3553 .unwrap();
3554 assert_eq!(result.records_written, 0);
3555 {
3556 let envelopes = dlq.0.lock().unwrap();
3557 assert_eq!(envelopes.len(), 3);
3558 for env in envelopes.iter() {
3560 let msg = env["error"]["message"].as_str().unwrap();
3561 assert!(msg.contains("write failed"), "got: {msg}");
3562 }
3563 }
3564 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3565 assert_eq!(result.dlq.unwrap().records_dlq, 3);
3566 }
3567
3568 #[tokio::test]
3569 async fn dlq_per_page_budget_exceeded_aborts() {
3570 let main = PartialSink::new(vec![0, 1, 2]);
3571 let dlq = std::sync::Arc::new(MockSink::new());
3572 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3573 dlq_cfg.max_failures_per_page = Some(2);
3574
3575 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3576 records: (0..3).map(|i| json!({"i": i})).collect(),
3577 bookmark: None,
3578 })];
3579 let stream = futures::stream::iter(pages);
3580 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3581 assert!(
3582 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
3583 "got: {result:?}"
3584 );
3585 }
3586
3587 #[tokio::test]
3588 async fn dlq_total_budget_exceeded_aborts_on_later_page() {
3589 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
3590 Ok(StreamPage {
3591 records: (0..3).map(|i| json!({"i": i})).collect(),
3592 bookmark: None,
3593 }),
3594 Ok(StreamPage {
3595 records: (3..6).map(|i| json!({"i": i})).collect(),
3596 bookmark: None,
3597 }),
3598 ];
3599 let main = PartialSink::new(vec![0, 1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3602 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3603 dlq_cfg.max_failures_total = Some(4);
3604
3605 let stream = futures::stream::iter(pages);
3606 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3607 assert!(
3608 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
3609 "got: {result:?}"
3610 );
3611 }
3612
3613 #[tokio::test]
3614 async fn dlq_per_page_budget_exceeded_commits_page_before_aborting() {
3615 let main = PartialSink::new(vec![1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3622 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3623 dlq_cfg.max_failures_per_page = Some(1); let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3626 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3627 records: (0..3).map(|i| json!({ "i": i })).collect(),
3628 bookmark: Some(json!("v1")),
3629 })];
3630 let stream = futures::stream::iter(pages);
3631 let result = run_stream(
3632 stream,
3633 &main,
3634 RunStreamOptions::new()
3635 .with_dlq(dlq_cfg)
3636 .with_state(std::sync::Arc::clone(&store), "k"),
3637 )
3638 .await;
3639
3640 assert!(
3642 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
3643 "got: {result:?}"
3644 );
3645 assert_eq!(main.committed.lock().unwrap().len(), 1);
3647 assert_eq!(dlq.0.lock().unwrap().len(), 2);
3649 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3651 }
3652
3653 #[tokio::test]
3654 async fn dlq_total_budget_exceeded_commits_tripping_page_before_aborting() {
3655 let main = PartialSink::new(vec![1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3660 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3661 dlq_cfg.max_failures_total = Some(1); let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3664 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3665 records: (0..3).map(|i| json!({ "i": i })).collect(),
3666 bookmark: Some(json!("v1")),
3667 })];
3668 let stream = futures::stream::iter(pages);
3669 let result = run_stream(
3670 stream,
3671 &main,
3672 RunStreamOptions::new()
3673 .with_dlq(dlq_cfg)
3674 .with_state(std::sync::Arc::clone(&store), "k"),
3675 )
3676 .await;
3677
3678 assert!(
3679 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
3680 "got: {result:?}"
3681 );
3682 assert_eq!(main.committed.lock().unwrap().len(), 1);
3683 assert_eq!(dlq.0.lock().unwrap().len(), 2);
3684 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3685 }
3686
3687 struct FailingDlqSink;
3690 #[async_trait]
3691 impl Sink for FailingDlqSink {
3692 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
3693 Err(FaucetError::Sink("dlq disk full".into()))
3694 }
3695 }
3696
3697 struct FailingFlushDlqSink {
3701 written: std::sync::Mutex<Vec<Value>>,
3702 }
3703 impl FailingFlushDlqSink {
3704 fn new() -> Self {
3705 Self {
3706 written: std::sync::Mutex::new(Vec::new()),
3707 }
3708 }
3709 }
3710 #[async_trait]
3711 impl Sink for FailingFlushDlqSink {
3712 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3713 self.written.lock().unwrap().extend(records.iter().cloned());
3714 Ok(records.len())
3715 }
3716 async fn flush(&self) -> Result<(), FaucetError> {
3717 Err(FaucetError::Sink("dlq flush failed".into()))
3718 }
3719 }
3720
3721 #[tokio::test]
3722 async fn dlq_sink_failure_is_fatal_no_recursion() {
3723 let main = PartialSink::new(vec![0]);
3724 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingDlqSink);
3725 let dlq_cfg = DlqConfig::new(dlq);
3726
3727 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3728 records: vec![json!({"i": 0}), json!({"i": 1})],
3729 bookmark: Some(json!("v1")),
3730 })];
3731 let stream = futures::stream::iter(pages);
3732 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3733 let result = run_stream(
3734 stream,
3735 &main,
3736 RunStreamOptions::new()
3737 .with_dlq(dlq_cfg)
3738 .with_state(std::sync::Arc::clone(&store), "k"),
3739 )
3740 .await;
3741 assert!(
3742 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink write failed")),
3743 "got: {result:?}"
3744 );
3745 assert!(store.get("k").await.unwrap().is_none());
3746 }
3747
3748 #[tokio::test]
3749 async fn dlq_bookmark_advances_only_after_both_flushes() {
3750 let main = PartialSink::new(vec![1]); let dlq = std::sync::Arc::new(MockSink::new());
3752 let dlq_cfg = DlqConfig::new(dlq.clone());
3753
3754 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3755 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3756 records: vec![json!({"i": 0}), json!({"i": 1})],
3757 bookmark: Some(json!("v1")),
3758 })];
3759 let stream = futures::stream::iter(pages);
3760 run_stream(
3761 stream,
3762 &main,
3763 RunStreamOptions::new()
3764 .with_dlq(dlq_cfg)
3765 .with_state(std::sync::Arc::clone(&store), "k"),
3766 )
3767 .await
3768 .unwrap();
3769 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3770 assert_eq!(dlq.0.lock().unwrap().len(), 1);
3771 assert_eq!(main.committed.lock().unwrap().len(), 1);
3772 }
3773
3774 #[tokio::test]
3775 async fn dlq_disabled_pipeline_behaves_identically_to_today() {
3776 let main = MockSink::new();
3778 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3779 records: vec![json!({"i": 0}), json!({"i": 1})],
3780 bookmark: None,
3781 })];
3782 let stream = futures::stream::iter(pages);
3783 let result = run_stream(stream, &main, RunStreamOptions::new())
3784 .await
3785 .unwrap();
3786 assert_eq!(result.records_written, 2);
3787 assert!(result.dlq.is_none());
3788 }
3789
3790 #[tokio::test]
3791 async fn dlq_per_page_flush_failure_is_fatal_and_blocks_bookmark() {
3792 let main = PartialSink::new(vec![1]);
3797 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
3798 let dlq_cfg = DlqConfig::new(dlq);
3799
3800 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3801 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3802 records: vec![json!({"i": 0}), json!({"i": 1})],
3803 bookmark: Some(json!("v1")),
3804 })];
3805 let stream = futures::stream::iter(pages);
3806 let result = run_stream(
3807 stream,
3808 &main,
3809 RunStreamOptions::new()
3810 .with_dlq(dlq_cfg)
3811 .with_state(std::sync::Arc::clone(&store), "k"),
3812 )
3813 .await;
3814 assert!(
3815 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
3816 "got: {result:?}"
3817 );
3818 assert!(store.get("k").await.unwrap().is_none());
3819 }
3820
3821 #[tokio::test]
3822 async fn dlq_end_of_stream_flush_failure_is_fatal() {
3823 let main = PartialSink::new(vec![1]);
3828 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
3829 let dlq_cfg = DlqConfig::new(dlq);
3830
3831 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3832 records: vec![json!({"i": 0}), json!({"i": 1})],
3833 bookmark: None,
3834 })];
3835 let stream = futures::stream::iter(pages);
3836 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3837 assert!(
3838 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
3839 "got: {result:?}"
3840 );
3841 }
3842
3843 #[tokio::test]
3844 #[allow(clippy::await_holding_lock)]
3845 async fn dlq_emits_records_total_and_pages_total() {
3846 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3847 use metrics_util::debugging::DebugValue;
3848
3849 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3850 let snap = snapshotter();
3851
3852 let source = MockSource(vec![json!({"i": 0}), json!({"i": 1})]);
3853 let main = PartialSink::new(vec![1]);
3854 let dlq = std::sync::Arc::new(MockSink::new());
3855 let _ = Pipeline::new(&source, &main)
3856 .with_name("p_dlq_metrics")
3857 .with_row("r1")
3858 .with_dlq(DlqConfig::new(dlq.clone()))
3859 .run()
3860 .await
3861 .unwrap();
3862
3863 let snapshot = snap.snapshot();
3864 let mut saw_records = false;
3865 let mut saw_pages = false;
3866 for (k, _u, _d, v) in snapshot.into_vec() {
3867 let key = k.key();
3868 let labels = key.labels().collect::<Vec<_>>();
3869 let has = |k: &str, v: &str| labels.iter().any(|l| l.key() == k && l.value() == v);
3870 if key.name() == "faucet_sink_dlq_records_total"
3871 && has("pipeline", "p_dlq_metrics")
3872 && has("row", "r1")
3873 && matches!(v, DebugValue::Counter(c) if c >= 1)
3874 {
3875 saw_records = true;
3876 }
3877 if key.name() == "faucet_sink_dlq_pages_total"
3878 && has("pipeline", "p_dlq_metrics")
3879 && matches!(v, DebugValue::Counter(c) if c >= 1)
3880 {
3881 saw_pages = true;
3882 }
3883 }
3884 assert!(saw_records, "faucet_sink_dlq_records_total not emitted");
3885 assert!(saw_pages, "faucet_sink_dlq_pages_total not emitted");
3886 }
3887
3888 #[tokio::test]
3889 #[allow(clippy::await_holding_lock)]
3890 async fn dlq_budget_exceeded_emits_counter() {
3891 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3892 use metrics_util::debugging::DebugValue;
3893
3894 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3895 let snap = snapshotter();
3896
3897 let source = MockSource((0..3).map(|i| json!({"i": i})).collect());
3898 let main = PartialSink::new(vec![0, 1, 2]);
3899 let dlq = std::sync::Arc::new(MockSink::new());
3900 let mut cfg = DlqConfig::new(dlq);
3901 cfg.max_failures_per_page = Some(1);
3902 let _ = Pipeline::new(&source, &main)
3903 .with_name("p_budget")
3904 .with_dlq(cfg)
3905 .run()
3906 .await;
3907
3908 let snapshot = snap.snapshot();
3909 let saw = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
3910 k.key().name() == "faucet_sink_dlq_budget_exceeded_total"
3911 && k.key()
3912 .labels()
3913 .any(|l| l.key() == "scope" && l.value() == "per_page")
3914 && matches!(v, DebugValue::Counter(c) if c >= 1)
3915 });
3916 assert!(saw, "faucet_sink_dlq_budget_exceeded_total not emitted");
3917 }
3918
3919 #[tokio::test]
3920 async fn pipeline_run_with_dlq_routes_partial_failures_end_to_end() {
3921 let source = MockSource(vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})]);
3923 let main = PartialSink::new(vec![1]);
3924 let dlq = std::sync::Arc::new(MockSink::new());
3925
3926 let result = Pipeline::new(&source, &main)
3927 .with_dlq(DlqConfig::new(dlq.clone()))
3928 .run()
3929 .await
3930 .unwrap();
3931
3932 assert_eq!(result.records_written, 2);
3933 let stats = result.dlq.unwrap();
3934 assert_eq!(stats.records_dlq, 1);
3935 {
3936 let dlq_records = dlq.0.lock().unwrap();
3937 assert_eq!(dlq_records.len(), 1);
3938 }
3939 }
3940
3941 #[cfg(feature = "quality")]
3944 #[tokio::test]
3945 async fn quality_quarantines_to_dlq_and_writes_survivors() {
3946 use crate::dlq::DlqConfig;
3947 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3948
3949 let main = Arc::new(MockSink::new());
3950 let dlq_sink = Arc::new(MockSink::new());
3951 let spec = QualitySpec {
3952 record: vec![RecordCheck::NotNull {
3953 field: "id".into(),
3954 treat_missing_as_null: true,
3955 on_failure: OnFailure::Quarantine,
3956 }],
3957 batch: vec![],
3958 };
3959 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3960 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3961 records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
3962 bookmark: None,
3963 })];
3964 let opts = RunStreamOptions::new()
3965 .with_dlq(DlqConfig::new(dlq_sink.clone()))
3966 .with_quality(quality);
3967 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
3968 .await
3969 .unwrap();
3970
3971 assert_eq!(result.records_written, 2); assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
3973 let dlq = dlq_sink.written();
3975 assert_eq!(dlq.len(), 1);
3976 assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
3977 assert_eq!(result.dlq.unwrap().records_dlq, 1);
3978 }
3979
3980 #[cfg(feature = "quality")]
3981 #[tokio::test]
3982 #[allow(clippy::await_holding_lock)]
3983 async fn quality_only_page_emits_quality_reason() {
3984 use crate::dlq::DlqConfig;
3988 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3989 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3990 use metrics_util::debugging::DebugValue;
3991
3992 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3993 let snap = snapshotter();
3994
3995 let main = Arc::new(MockSink::new());
3998 let dlq_sink = Arc::new(MockSink::new());
3999 let spec = QualitySpec {
4000 record: vec![RecordCheck::NotNull {
4001 field: "id".into(),
4002 treat_missing_as_null: true,
4003 on_failure: OnFailure::Quarantine,
4004 }],
4005 batch: vec![],
4006 };
4007 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4008 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4009 records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
4010 bookmark: None,
4011 })];
4012 let opts = RunStreamOptions::new()
4013 .with_name("p_quality_reason")
4014 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4015 .with_quality(quality);
4016 let _ = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4017 .await
4018 .unwrap();
4019
4020 let snapshot = snap.snapshot();
4021 let saw_quality_reason = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
4022 k.key().name() == "faucet_sink_dlq_pages_total"
4023 && k.key()
4024 .labels()
4025 .any(|l| l.key() == "pipeline" && l.value() == "p_quality_reason")
4026 && k.key()
4027 .labels()
4028 .any(|l| l.key() == "reason" && l.value() == "quality")
4029 && matches!(v, DebugValue::Counter(c) if c >= 1)
4030 });
4031 assert!(
4032 saw_quality_reason,
4033 "expected faucet_sink_dlq_pages_total with reason=\"quality\""
4034 );
4035 }
4036
4037 #[cfg(feature = "quality")]
4038 #[tokio::test]
4039 async fn quality_abort_fails_run() {
4040 use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
4041 let main = MockSink::new();
4042 let spec = QualitySpec {
4043 record: vec![],
4044 batch: vec![BatchCheck::RowCount {
4045 min: Some(5),
4046 max: None,
4047 on_failure: OnFailure::Abort,
4048 }],
4049 };
4050 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4051 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4052 records: vec![json!({"id": 1})],
4053 bookmark: None,
4054 })];
4055 let opts = RunStreamOptions::new().with_quality(quality);
4056 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4057 assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
4058 }
4059
4060 #[cfg(feature = "quality")]
4061 #[tokio::test]
4062 async fn quality_quarantine_without_dlq_is_rejected() {
4063 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
4064 let main = MockSink::new();
4065 let spec = QualitySpec {
4066 record: vec![RecordCheck::NotNull {
4067 field: "id".into(),
4068 treat_missing_as_null: true,
4069 on_failure: OnFailure::Quarantine,
4070 }],
4071 batch: vec![],
4072 };
4073 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4074 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4075 records: vec![json!({"id": null})],
4076 bookmark: None,
4077 })];
4078 let opts = RunStreamOptions::new().with_quality(quality);
4080 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4081 assert!(matches!(result, Err(FaucetError::Config(_))));
4082 }
4083
4084 #[cfg(feature = "contract")]
4085 fn compiled_contract(on_breach: &str) -> Arc<crate::contract::CompiledContract> {
4086 let spec: crate::contract::ContractSpec = serde_json::from_value(json!({
4087 "version": "1.0.0",
4088 "on_breach": on_breach,
4089 "fields": [{ "name": "id", "type": "integer" }]
4090 }))
4091 .unwrap();
4092 Arc::new(crate::contract::CompiledContract::compile(&spec).unwrap())
4093 }
4094
4095 #[cfg(feature = "contract")]
4096 #[tokio::test]
4097 async fn contract_quarantines_to_dlq_and_writes_survivors() {
4098 use crate::dlq::DlqConfig;
4099 let main = Arc::new(MockSink::new());
4100 let dlq_sink = Arc::new(MockSink::new());
4101 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4102 records: vec![json!({"id": 1}), json!({"id": "bad"}), json!({"id": 3})],
4103 bookmark: None,
4104 })];
4105 let opts = RunStreamOptions::new()
4106 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4107 .with_contract(compiled_contract("quarantine"));
4108 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4109 .await
4110 .unwrap();
4111
4112 assert_eq!(result.records_written, 2);
4113 assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
4114 let dlq = dlq_sink.written();
4115 assert_eq!(dlq.len(), 1);
4116 assert_eq!(dlq[0]["error"]["kind"], "ContractViolation");
4117 assert_eq!(dlq[0]["payload"], json!({"id": "bad"}));
4118 assert_eq!(dlq[0]["record_index"], 1);
4120 assert_eq!(result.dlq.unwrap().records_dlq, 1);
4121 }
4122
4123 #[cfg(feature = "contract")]
4124 #[tokio::test]
4125 async fn contract_fail_aborts_run_and_writes_nothing() {
4126 let main = MockSink::new();
4127 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4128 records: vec![json!({"id": 1}), json!({"id": "bad"})],
4129 bookmark: None,
4130 })];
4131 let opts = RunStreamOptions::new().with_contract(compiled_contract("fail"));
4132 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4133 match result {
4134 Err(FaucetError::ContractViolation { version, message }) => {
4135 assert_eq!(version, "1.0.0");
4136 assert!(message.contains("id"), "message: {message}");
4137 }
4138 other => panic!("expected ContractViolation, got {other:?}"),
4139 }
4140 assert!(
4141 main.written().is_empty(),
4142 "a contract fail must not commit any of the page's records"
4143 );
4144 }
4145
4146 #[cfg(feature = "contract")]
4147 #[tokio::test]
4148 async fn contract_warn_writes_everything() {
4149 let main = MockSink::new();
4150 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4151 records: vec![json!({"id": 1}), json!({"id": "bad"})],
4152 bookmark: None,
4153 })];
4154 let opts = RunStreamOptions::new().with_contract(compiled_contract("warn"));
4155 let result = run_stream(futures::stream::iter(pages), &main, opts)
4156 .await
4157 .unwrap();
4158 assert_eq!(result.records_written, 2);
4159 assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": "bad"})]);
4160 }
4161
4162 #[cfg(feature = "contract")]
4163 #[tokio::test]
4164 async fn contract_quarantine_without_dlq_is_rejected() {
4165 let main = MockSink::new();
4166 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4167 records: vec![json!({"id": 1})],
4168 bookmark: None,
4169 })];
4170 let opts = RunStreamOptions::new().with_contract(compiled_contract("quarantine"));
4172 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
4173 assert!(matches!(result, Err(FaucetError::Config(_))));
4174 }
4175
4176 #[cfg(all(feature = "contract", feature = "quality"))]
4177 #[tokio::test]
4178 async fn contract_runs_after_quality_and_shares_dlq() {
4179 use crate::dlq::DlqConfig;
4183 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
4184 let main = Arc::new(MockSink::new());
4185 let dlq_sink = Arc::new(MockSink::new());
4186 let quality = Arc::new(
4187 CompiledQuality::compile(&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 .unwrap(),
4196 );
4197 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4198 records: vec![json!({"id": null}), json!({"id": "bad"}), json!({"id": 3})],
4199 bookmark: None,
4200 })];
4201 let opts = RunStreamOptions::new()
4202 .with_dlq(DlqConfig::new(dlq_sink.clone()))
4203 .with_quality(quality)
4204 .with_contract(compiled_contract("quarantine"));
4205 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
4206 .await
4207 .unwrap();
4208
4209 assert_eq!(result.records_written, 1);
4210 assert_eq!(main.written(), vec![json!({"id": 3})]);
4211 let dlq = dlq_sink.written();
4212 assert_eq!(dlq.len(), 2);
4213 let kinds: Vec<&str> = dlq
4214 .iter()
4215 .map(|e| e["error"]["kind"].as_str().unwrap())
4216 .collect();
4217 assert!(kinds.contains(&"QualityFailure"), "kinds: {kinds:?}");
4218 assert!(kinds.contains(&"ContractViolation"), "kinds: {kinds:?}");
4219 assert_eq!(result.dlq.unwrap().records_dlq, 2);
4220
4221 let contract_env = dlq
4227 .iter()
4228 .find(|e| e["error"]["kind"] == "ContractViolation")
4229 .expect("a contract envelope");
4230 assert_eq!(
4231 contract_env["record_index"], 1,
4232 "contract record_index must be the original page position: {contract_env}"
4233 );
4234 let quality_env = dlq
4235 .iter()
4236 .find(|e| e["error"]["kind"] == "QualityFailure")
4237 .expect("a quality envelope");
4238 assert_eq!(
4239 quality_env["record_index"], 0,
4240 "null id was at page index 0"
4241 );
4242 }
4243
4244 struct FlakySink {
4247 every: usize,
4248 calls: std::sync::Mutex<Vec<usize>>,
4249 }
4250 impl FlakySink {
4251 fn new(every: usize) -> Self {
4252 Self {
4253 every,
4254 calls: std::sync::Mutex::new(Vec::new()),
4255 }
4256 }
4257 fn call_sizes(&self) -> Vec<usize> {
4258 self.calls.lock().unwrap().clone()
4259 }
4260 }
4261 #[async_trait]
4262 impl Sink for FlakySink {
4263 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4264 Ok(records.len())
4265 }
4266 async fn write_batch_partial(
4267 &self,
4268 records: &[Value],
4269 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4270 self.calls.lock().unwrap().push(records.len());
4271 Ok(records
4272 .iter()
4273 .enumerate()
4274 .map(|(i, _)| {
4275 if (i + 1) % self.every == 0 {
4276 Err(FaucetError::Sink("synthetic".into()))
4277 } else {
4278 Ok(())
4279 }
4280 })
4281 .collect())
4282 }
4283 }
4284
4285 #[tokio::test]
4286 async fn adaptive_shrinks_under_errors_on_dlq_path() {
4287 use crate::adaptive::AdaptiveBatchConfig;
4288 use crate::dlq::{DlqConfig, OnBatchError};
4289 let mk = || StreamPage {
4294 records: (0..400).map(|i| json!({"i": i})).collect(),
4295 bookmark: None,
4296 };
4297 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4298 let sink = FlakySink::new(4); let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
4300 let dlq = DlqConfig {
4301 sink: dlq_sink,
4302 on_batch_error: OnBatchError::Propagate,
4303 max_failures_per_page: None,
4304 max_failures_total: None,
4305 include_original_payload: true,
4306 };
4307 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4308 "enabled": true, "min": 50, "max": 400,
4309 "decrease_factor": 0.5, "cooldown_batches": 0, "error_threshold": 0.1
4310 }))
4311 .unwrap();
4312 let opts = RunStreamOptions::new().with_dlq(dlq).with_adaptive(cfg);
4313 let result = run_stream(stream, &sink, opts).await.unwrap();
4314 assert!(
4321 result.records_written >= 900,
4322 "expected ≥900 written, got {}",
4323 result.records_written
4324 );
4325 let sizes = sink.call_sizes();
4326 assert_eq!(sizes[0], 400, "first chunk is the full page");
4327 assert!(
4328 sizes.last().unwrap() < &400,
4329 "controller should shrink under errors: {sizes:?}"
4330 );
4331 assert!(
4332 result.dlq.unwrap().records_dlq >= 250,
4333 "expected ≥250 DLQ records"
4334 );
4335 }
4336
4337 struct RecordingSink {
4342 calls: std::sync::Mutex<Vec<usize>>,
4343 latency: std::time::Duration,
4344 }
4345 impl RecordingSink {
4346 fn new(latency_ms: u64) -> Self {
4347 Self {
4348 calls: std::sync::Mutex::new(Vec::new()),
4349 latency: std::time::Duration::from_millis(latency_ms),
4350 }
4351 }
4352 fn call_sizes(&self) -> Vec<usize> {
4353 self.calls.lock().unwrap().clone()
4354 }
4355 }
4356 #[async_trait]
4357 impl Sink for RecordingSink {
4358 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4359 tokio::time::sleep(self.latency).await;
4360 self.calls.lock().unwrap().push(records.len());
4361 Ok(records.len())
4362 }
4363 }
4364
4365 #[tokio::test]
4366 async fn adaptive_reslices_non_dlq_page_into_subbatches() {
4367 use crate::adaptive::AdaptiveBatchConfig;
4368 let page = StreamPage {
4369 records: (0..1000).map(|i| json!({ "i": i })).collect(),
4370 bookmark: None,
4371 };
4372 let stream = futures::stream::iter(vec![Ok(page)]);
4373 let sink = RecordingSink::new(0);
4374 let cfg: AdaptiveBatchConfig =
4375 serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
4376 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4377 .await
4378 .unwrap();
4379 assert_eq!(result.records_written, 1000);
4380 assert_eq!(sink.call_sizes(), vec![1000]);
4382 }
4383
4384 #[tokio::test]
4385 async fn adaptive_shrinks_under_latency_target_then_smaller_chunks() {
4386 use crate::adaptive::AdaptiveBatchConfig;
4387 let mk = || StreamPage {
4388 records: (0..400).map(|i| json!({"i": i})).collect(),
4389 bookmark: None,
4390 };
4391 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4392 let sink = RecordingSink::new(50);
4393 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4394 "enabled": true, "min": 50, "max": 400,
4395 "decrease_factor": 0.5, "cooldown_batches": 0,
4396 "target_latency_ms": 10, "latency_window": 1
4397 }))
4398 .unwrap();
4399 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4400 .await
4401 .unwrap();
4402 assert_eq!(result.records_written, 1200);
4403 let sizes = sink.call_sizes();
4404 assert_eq!(sizes[0], 400);
4405 assert!(
4406 sizes.last().unwrap() < &400,
4407 "controller should have shrunk: {sizes:?}"
4408 );
4409 }
4410
4411 #[tokio::test]
4412 #[allow(clippy::await_holding_lock)]
4413 async fn adaptive_emits_batch_size_and_adjustments_metrics() {
4414 use crate::adaptive::AdaptiveBatchConfig;
4417 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
4418 use metrics_util::debugging::DebugValue;
4419
4420 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
4421 let snap = snapshotter();
4422
4423 let mk = || StreamPage {
4429 records: (0..400).map(|i| json!({"i": i})).collect(),
4430 bookmark: None,
4431 };
4432 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
4433 let sink = RecordingSink::new(50);
4434 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
4435 "enabled": true, "min": 50, "max": 400,
4436 "decrease_factor": 0.5, "cooldown_batches": 0,
4437 "target_latency_ms": 10, "latency_window": 1
4438 }))
4439 .unwrap();
4440
4441 let _ = run_stream(
4442 stream,
4443 &sink,
4444 RunStreamOptions::new()
4445 .with_adaptive(cfg)
4446 .with_name("p")
4447 .with_row("r"),
4448 )
4449 .await
4450 .unwrap();
4451
4452 let snapshot = snap.snapshot();
4453 let mut saw_batch_size = false;
4454 let mut saw_adjustments = false;
4455 for (k, _u, _d, v) in snapshot.into_vec() {
4456 let key = k.key();
4457 let labels = key.labels().collect::<Vec<_>>();
4458 let has = |k: &str, val: &str| labels.iter().any(|l| l.key() == k && l.value() == val);
4459
4460 if key.name() == "faucet_pipeline_adaptive_batch_size"
4461 && has("pipeline", "p")
4462 && has("row", "r")
4463 && matches!(v, DebugValue::Gauge(_))
4464 {
4465 saw_batch_size = true;
4466 }
4467 if key.name() == "faucet_pipeline_adaptive_batch_adjustments_total"
4468 && has("pipeline", "p")
4469 && has("row", "r")
4470 && matches!(v, DebugValue::Counter(c) if c >= 1)
4471 {
4472 saw_adjustments = true;
4473 }
4474 }
4475 assert!(
4476 saw_batch_size,
4477 "expected faucet_pipeline_adaptive_batch_size gauge with pipeline=p, row=r"
4478 );
4479 assert!(
4480 saw_adjustments,
4481 "expected faucet_pipeline_adaptive_batch_adjustments_total counter with pipeline=p, row=r"
4482 );
4483 }
4484
4485 #[tokio::test]
4488 async fn exactly_once_rejects_dlq() {
4489 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4492 let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
4493 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
4494 let opts = eo_opts(store, "k", 0).with_dlq(DlqConfig::new(dlq_sink));
4495 let r = run_stream(
4496 futures::stream::iter(pages),
4497 &IdempotentMockSink::new(),
4498 opts,
4499 )
4500 .await;
4501 assert!(
4502 matches!(&r, Err(FaucetError::Config(m)) if m.contains("not compatible with a DLQ")),
4503 "got: {r:?}"
4504 );
4505 }
4506
4507 #[tokio::test]
4510 async fn exactly_once_writes_unbookmarked_page_at_least_once() {
4511 let sink = IdempotentMockSink::new();
4516 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4517 let pages = vec![Ok(StreamPage {
4518 records: vec![json!({"id": 1}), json!({"id": 2})],
4519 bookmark: None,
4520 })];
4521 let r = run_stream(
4522 futures::stream::iter(pages),
4523 &sink,
4524 eo_opts(store.clone(), "k", 0),
4525 )
4526 .await
4527 .unwrap();
4528 assert_eq!(r.records_written, 2);
4529 assert_eq!(r.bookmark, None);
4530 assert_eq!(sink.last_committed_token("k").await.unwrap(), None);
4533 assert!(store.get("k").await.unwrap().is_none());
4535 assert_eq!(sink.rows(), vec![json!({"id": 1}), json!({"id": 2})]);
4536 }
4537
4538 #[tokio::test]
4541 async fn dlq_with_bookmark_persists_after_routing_failures() {
4542 let main = PartialSink::new(vec![1]); let dlq = std::sync::Arc::new(MockSink::new());
4547 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4548 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
4549 records: vec![json!({"i": 0}), json!({"i": 1})],
4550 bookmark: Some(json!("ckpt")),
4551 })];
4552 let result = run_stream(
4553 futures::stream::iter(pages),
4554 &main,
4555 RunStreamOptions::new()
4556 .with_dlq(DlqConfig::new(dlq.clone()))
4557 .with_state(Arc::clone(&store), "k"),
4558 )
4559 .await
4560 .unwrap();
4561
4562 assert_eq!(result.records_written, 1); assert_eq!(result.bookmark, Some(json!("ckpt")));
4564 assert_eq!(store.get("k").await.unwrap(), Some(json!("ckpt")));
4566 let envelopes = dlq.0.lock().unwrap();
4568 assert_eq!(envelopes.len(), 1);
4569 assert_eq!(envelopes[0]["payload"]["i"], 1);
4570 }
4571
4572 #[tokio::test]
4575 async fn exactly_once_resumes_sequence_from_start_seq() {
4576 let sink = IdempotentMockSink::new();
4580 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4581 let pages = vec![Ok(StreamPage {
4582 records: vec![json!({"id": 9})],
4583 bookmark: Some(json!("bm-after-resume")),
4584 })];
4585 let r = run_stream(
4586 futures::stream::iter(pages),
4587 &sink,
4588 eo_opts(store.clone(), "eo_key", 7),
4589 )
4590 .await
4591 .unwrap();
4592
4593 assert_eq!(r.records_written, 1);
4594 let (bm, seq) =
4595 crate::idempotency::unwrap_state(&store.get("eo_key").await.unwrap().unwrap());
4596 assert_eq!(bm, Some(json!("bm-after-resume")));
4597 assert_eq!(seq, 8, "sequence resumes at start_seq + 1");
4598 let token = sink.last_committed_token("eo_key").await.unwrap().unwrap();
4599 assert_eq!(
4600 crate::idempotency::parse_token_parts(&token),
4601 Some((8, Some(json!("bm-after-resume"))))
4602 );
4603 }
4604
4605 #[cfg(feature = "quality")]
4608 #[tokio::test]
4609 async fn pipeline_run_with_quality_aborts_on_failed_batch_check() {
4610 use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
4611 let source = MockSource(vec![json!({"id": 1})]);
4612 let main = MockSink::new();
4613 let spec = QualitySpec {
4614 record: vec![],
4615 batch: vec![BatchCheck::RowCount {
4616 min: Some(5),
4617 max: None,
4618 on_failure: OnFailure::Abort,
4619 }],
4620 };
4621 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4622 let result = Pipeline::new(&source, &main)
4623 .with_quality(quality)
4624 .run()
4625 .await;
4626 assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
4627 assert!(main.written().is_empty());
4629 }
4630
4631 #[tokio::test]
4634 async fn pipeline_run_with_adaptive_reslices_page() {
4635 use crate::adaptive::AdaptiveBatchConfig;
4636 let source = MockSource((0..1000).map(|i| json!({ "i": i })).collect());
4637 let sink = MockSink::new();
4638 let cfg: AdaptiveBatchConfig =
4639 serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
4640 let result = Pipeline::new(&source, &sink)
4641 .with_adaptive(cfg)
4642 .run()
4643 .await
4644 .unwrap();
4645 assert_eq!(result.records_written, 1000);
4646 assert_eq!(sink.written().len(), 1000);
4647 }
4648
4649 #[tokio::test]
4652 async fn adaptive_noop_sink_name_is_handled() {
4653 use crate::adaptive::AdaptiveBatchConfig;
4654
4655 struct JsonlNamedSink(std::sync::Mutex<usize>);
4658 #[async_trait]
4659 impl Sink for JsonlNamedSink {
4660 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4661 *self.0.lock().unwrap() += records.len();
4662 Ok(records.len())
4663 }
4664 fn connector_name(&self) -> &'static str {
4665 "jsonl"
4666 }
4667 }
4668
4669 let page = StreamPage {
4670 records: (0..10).map(|i| json!({ "i": i })).collect(),
4671 bookmark: None,
4672 };
4673 let stream = futures::stream::iter(vec![Ok(page)]);
4674 let sink = JsonlNamedSink(std::sync::Mutex::new(0));
4675 let cfg: AdaptiveBatchConfig =
4676 serde_json::from_value(json!({"enabled": true, "min": 5, "max": 10})).unwrap();
4677 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4678 .await
4679 .unwrap();
4680 assert_eq!(result.records_written, 10);
4681 assert_eq!(*sink.0.lock().unwrap(), 10);
4682 }
4683
4684 #[tokio::test]
4685 async fn resilience_retries_transient_sink_write() {
4686 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4687 use std::sync::Arc;
4688 use std::sync::atomic::{AtomicU32, Ordering};
4689 use std::time::Duration;
4690
4691 struct TransientFlakySink {
4692 attempts: Arc<AtomicU32>,
4693 written: Arc<AtomicU32>,
4694 }
4695 #[async_trait::async_trait]
4696 impl Sink for TransientFlakySink {
4697 async fn write_batch(
4698 &self,
4699 records: &[serde_json::Value],
4700 ) -> Result<usize, FaucetError> {
4701 let n = self.attempts.fetch_add(1, Ordering::SeqCst);
4702 if n < 2 {
4703 return Err(FaucetError::HttpStatus {
4704 status: 503,
4705 url: "u".into(),
4706 body: "".into(),
4707 });
4708 }
4709 self.written
4710 .fetch_add(records.len() as u32, Ordering::SeqCst);
4711 Ok(records.len())
4712 }
4713 async fn flush(&self) -> Result<(), FaucetError> {
4714 Ok(())
4715 }
4716 fn supports_idempotent_writes(&self) -> bool {
4718 true
4719 }
4720 }
4721
4722 let attempts = Arc::new(AtomicU32::new(0));
4723 let written = Arc::new(AtomicU32::new(0));
4724 let sink = TransientFlakySink {
4725 attempts: attempts.clone(),
4726 written: written.clone(),
4727 };
4728 let pages = futures::stream::iter(vec![Ok(StreamPage {
4729 records: vec![serde_json::json!({"a": 1})],
4730 bookmark: None,
4731 })]);
4732 let policy = ResiliencePolicy {
4733 retry: RetryPolicy {
4734 max_attempts: 5,
4735 backoff: BackoffKind::None,
4736 base: Duration::ZERO,
4737 max: Duration::ZERO,
4738 jitter: false,
4739 ..RetryPolicy::default()
4740 },
4741 ..ResiliencePolicy::default()
4742 };
4743 let res = run_stream(
4744 pages,
4745 &sink,
4746 RunStreamOptions::new().with_resilience(policy),
4747 )
4748 .await
4749 .unwrap();
4750 assert_eq!(written.load(Ordering::SeqCst), 1);
4751 assert_eq!(attempts.load(Ordering::SeqCst), 3);
4752 assert_eq!(res.records_written, 1);
4753 }
4754
4755 #[tokio::test]
4756 async fn resilience_does_not_retry_non_idempotent_write_batch() {
4757 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4761 use std::sync::Arc;
4762 use std::sync::atomic::{AtomicU32, Ordering};
4763 use std::time::Duration;
4764
4765 struct NonIdempotentFlakySink {
4766 attempts: Arc<AtomicU32>,
4767 }
4768 #[async_trait::async_trait]
4769 impl Sink for NonIdempotentFlakySink {
4770 async fn write_batch(&self, _r: &[Value]) -> Result<usize, FaucetError> {
4771 self.attempts.fetch_add(1, Ordering::SeqCst);
4772 Err(FaucetError::HttpStatus {
4773 status: 503,
4774 url: "u".into(),
4775 body: "".into(),
4776 })
4777 }
4778 async fn flush(&self) -> Result<(), FaucetError> {
4779 Ok(())
4780 }
4781 }
4783
4784 let attempts = Arc::new(AtomicU32::new(0));
4785 let sink = NonIdempotentFlakySink {
4786 attempts: attempts.clone(),
4787 };
4788 let pages = futures::stream::iter(vec![Ok(StreamPage {
4789 records: vec![json!({"a": 1})],
4790 bookmark: None,
4791 })]);
4792 let policy = ResiliencePolicy {
4793 retry: RetryPolicy {
4794 max_attempts: 5,
4795 backoff: BackoffKind::None,
4796 base: Duration::ZERO,
4797 max: Duration::ZERO,
4798 jitter: false,
4799 ..RetryPolicy::default()
4800 },
4801 ..ResiliencePolicy::default()
4802 };
4803 let err = run_stream(
4804 pages,
4805 &sink,
4806 RunStreamOptions::new().with_resilience(policy),
4807 )
4808 .await
4809 .unwrap_err();
4810 assert!(matches!(err, FaucetError::HttpStatus { status: 503, .. }));
4811 assert_eq!(
4812 attempts.load(Ordering::SeqCst),
4813 1,
4814 "non-idempotent write_batch must be attempted exactly once (no retry)"
4815 );
4816 }
4817
4818 #[tokio::test]
4819 async fn resilience_circuit_opens_after_consecutive_failed_pages() {
4820 use crate::dlq::{DlqConfig, OnBatchError};
4821 use crate::resilience::{BackoffKind, CircuitBreakerConfig, ResiliencePolicy, RetryPolicy};
4822 use std::sync::Arc;
4823 use std::time::Duration;
4824
4825 struct DeadSink;
4827 #[async_trait]
4828 impl Sink for DeadSink {
4829 async fn write_batch(&self, _r: &[Value]) -> Result<usize, FaucetError> {
4830 Err(FaucetError::Sink("down".into()))
4831 }
4832 async fn flush(&self) -> Result<(), FaucetError> {
4833 Ok(())
4834 }
4835 }
4836 struct NullSink;
4838 #[async_trait]
4839 impl Sink for NullSink {
4840 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4841 Ok(r.len())
4842 }
4843 async fn flush(&self) -> Result<(), FaucetError> {
4844 Ok(())
4845 }
4846 }
4847
4848 let pages = futures::stream::iter((0..10).map(|i| {
4849 Ok(StreamPage {
4850 records: vec![json!({"i": i})],
4851 bookmark: None,
4852 })
4853 }));
4854 let dlq = DlqConfig {
4855 on_batch_error: OnBatchError::DlqAll,
4856 ..DlqConfig::new(Arc::new(NullSink))
4857 };
4858 let policy = ResiliencePolicy {
4859 retry: RetryPolicy {
4860 max_attempts: 1,
4861 backoff: BackoffKind::None,
4862 base: Duration::ZERO,
4863 max: Duration::ZERO,
4864 jitter: false,
4865 ..RetryPolicy::default()
4866 },
4867 circuit_breaker: Some(CircuitBreakerConfig {
4868 consecutive_failures: 3,
4869 cooldown: Duration::from_secs(60),
4870 }),
4871 poison: None,
4872 };
4873 let err = run_stream(
4874 pages,
4875 &DeadSink,
4876 RunStreamOptions::new()
4877 .with_dlq(dlq)
4878 .with_resilience(policy),
4879 )
4880 .await
4881 .unwrap_err();
4882 assert!(
4883 matches!(err, FaucetError::CircuitOpen { failures: 3, .. }),
4884 "got {err:?}"
4885 );
4886 }
4887
4888 #[tokio::test]
4889 async fn resilience_poison_retries_then_dlqs_failing_row() {
4890 use crate::dlq::DlqConfig;
4891 use crate::resilience::{
4892 BackoffKind, PoisonAction, PoisonPolicy, ResiliencePolicy, RetryPolicy,
4893 };
4894 use std::sync::{Arc, Mutex};
4895 use std::time::Duration;
4896
4897 struct PickySink {
4900 bad_attempts: Arc<Mutex<u32>>,
4901 }
4902 #[async_trait]
4903 impl Sink for PickySink {
4904 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4905 Ok(r.len())
4906 }
4907 async fn write_batch_partial(
4908 &self,
4909 records: &[Value],
4910 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4911 Ok(records
4912 .iter()
4913 .map(|rec| {
4914 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
4915 *self.bad_attempts.lock().unwrap() += 1;
4916 Err(FaucetError::HttpStatus {
4917 status: 503,
4918 url: "u".into(),
4919 body: "".into(),
4920 })
4921 } else {
4922 Ok(())
4923 }
4924 })
4925 .collect())
4926 }
4927 async fn flush(&self) -> Result<(), FaucetError> {
4928 Ok(())
4929 }
4930 }
4931 struct CaptureSink(Arc<Mutex<Vec<Value>>>);
4932 #[async_trait]
4933 impl Sink for CaptureSink {
4934 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4935 self.0.lock().unwrap().extend_from_slice(r);
4936 Ok(r.len())
4937 }
4938 async fn flush(&self) -> Result<(), FaucetError> {
4939 Ok(())
4940 }
4941 }
4942
4943 let captured = Arc::new(Mutex::new(Vec::new()));
4944 let bad_attempts = Arc::new(Mutex::new(0u32));
4945 let sink = PickySink {
4946 bad_attempts: bad_attempts.clone(),
4947 };
4948 let pages = futures::stream::iter(vec![Ok(StreamPage {
4949 records: vec![json!({"ok": 1}), json!({"bad": true})],
4950 bookmark: None,
4951 })]);
4952 let policy = ResiliencePolicy {
4953 retry: RetryPolicy {
4954 max_attempts: 1,
4955 backoff: BackoffKind::None,
4956 base: Duration::ZERO,
4957 max: Duration::ZERO,
4958 jitter: false,
4959 ..RetryPolicy::default()
4960 },
4961 circuit_breaker: None,
4962 poison: Some(PoisonPolicy {
4963 max_row_attempts: 3,
4964 action: PoisonAction::Dlq,
4965 }),
4966 };
4967 let res = run_stream(
4968 pages,
4969 &sink,
4970 RunStreamOptions::new()
4971 .with_dlq(DlqConfig::new(Arc::new(CaptureSink(captured.clone()))))
4972 .with_resilience(policy),
4973 )
4974 .await
4975 .unwrap();
4976
4977 assert_eq!(
4978 *bad_attempts.lock().unwrap(),
4979 3,
4980 "bad row tried max_row_attempts times"
4981 );
4982 assert_eq!(res.records_written, 1, "the ok row");
4983 let dlq = captured.lock().unwrap();
4984 assert_eq!(dlq.len(), 1, "one row to DLQ");
4985 assert_eq!(dlq[0]["payload"]["bad"], json!(true));
4986 }
4987
4988 #[tokio::test]
4989 async fn poison_loop_does_not_nest_resilience_retry_on_subset_resubmit() {
4990 use crate::dlq::DlqConfig;
4998 use crate::resilience::{
4999 BackoffKind, PoisonAction, PoisonPolicy, ResiliencePolicy, RetryPolicy,
5000 };
5001 use std::sync::{Arc, Mutex};
5002 use std::time::Duration;
5003
5004 struct OuterErrOnSubsetSink {
5005 subset_calls: Arc<Mutex<u32>>,
5006 }
5007 #[async_trait]
5008 impl Sink for OuterErrOnSubsetSink {
5009 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5010 Ok(r.len())
5011 }
5012 async fn write_batch_partial(
5013 &self,
5014 records: &[Value],
5015 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
5016 if records.len() == 1 {
5018 *self.subset_calls.lock().unwrap() += 1;
5019 return Err(FaucetError::HttpStatus {
5020 status: 503,
5021 url: "u".into(),
5022 body: "".into(),
5023 });
5024 }
5025 Ok(records
5026 .iter()
5027 .map(|rec| {
5028 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
5029 Err(FaucetError::HttpStatus {
5030 status: 503,
5031 url: "u".into(),
5032 body: "".into(),
5033 })
5034 } else {
5035 Ok(())
5036 }
5037 })
5038 .collect())
5039 }
5040 async fn flush(&self) -> Result<(), FaucetError> {
5041 Ok(())
5042 }
5043 fn supports_idempotent_writes(&self) -> bool {
5046 true
5047 }
5048 }
5049 struct CaptureSink(Arc<Mutex<Vec<Value>>>);
5050 #[async_trait]
5051 impl Sink for CaptureSink {
5052 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5053 self.0.lock().unwrap().extend_from_slice(r);
5054 Ok(r.len())
5055 }
5056 async fn flush(&self) -> Result<(), FaucetError> {
5057 Ok(())
5058 }
5059 }
5060
5061 let subset_calls = Arc::new(Mutex::new(0u32));
5062 let sink = OuterErrOnSubsetSink {
5063 subset_calls: subset_calls.clone(),
5064 };
5065 let pages = futures::stream::iter(vec![Ok(StreamPage {
5066 records: vec![json!({"ok": 1}), json!({"bad": true})],
5067 bookmark: None,
5068 })]);
5069 let policy = ResiliencePolicy {
5070 retry: RetryPolicy {
5071 max_attempts: 4, backoff: BackoffKind::None,
5073 base: Duration::ZERO,
5074 max: Duration::ZERO,
5075 jitter: false,
5076 ..RetryPolicy::default()
5077 },
5078 circuit_breaker: None,
5079 poison: Some(PoisonPolicy {
5080 max_row_attempts: 3,
5081 action: PoisonAction::Dlq,
5082 }),
5083 };
5084 let res = run_stream(
5085 pages,
5086 &sink,
5087 RunStreamOptions::new()
5088 .with_dlq(DlqConfig::new(Arc::new(CaptureSink(Arc::new(Mutex::new(
5089 Vec::new(),
5090 ))))))
5091 .with_resilience(policy),
5092 )
5093 .await;
5094
5095 assert!(matches!(
5097 res,
5098 Err(FaucetError::HttpStatus { status: 503, .. })
5099 ));
5100 assert_eq!(
5103 *subset_calls.lock().unwrap(),
5104 1,
5105 "poison subset resubmit must be bare (no nested resilience retry)"
5106 );
5107 }
5108
5109 #[tokio::test]
5110 async fn resilience_retries_transient_flush_and_state_on_dlq_path() {
5111 use crate::dlq::DlqConfig;
5118 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
5119 use crate::state::StateStore;
5120 use std::sync::Arc;
5121 use std::sync::atomic::{AtomicU32, Ordering};
5122 use std::time::Duration;
5123
5124 fn transient_503() -> FaucetError {
5125 FaucetError::HttpStatus {
5126 status: 503,
5127 url: "u".into(),
5128 body: "".into(),
5129 }
5130 }
5131
5132 struct FlakyFlushSink {
5135 flush_attempts: Arc<AtomicU32>,
5136 }
5137 #[async_trait]
5138 impl Sink for FlakyFlushSink {
5139 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5140 Ok(r.len())
5141 }
5142 async fn write_batch_partial(
5143 &self,
5144 records: &[Value],
5145 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
5146 Ok(records
5147 .iter()
5148 .map(|rec| {
5149 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
5150 Err(transient_503())
5151 } else {
5152 Ok(())
5153 }
5154 })
5155 .collect())
5156 }
5157 async fn flush(&self) -> Result<(), FaucetError> {
5158 let n = self.flush_attempts.fetch_add(1, Ordering::SeqCst);
5159 if n < 2 { Err(transient_503()) } else { Ok(()) }
5160 }
5161 }
5162 struct NullSink;
5163 #[async_trait]
5164 impl Sink for NullSink {
5165 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5166 Ok(r.len())
5167 }
5168 async fn flush(&self) -> Result<(), FaucetError> {
5169 Ok(())
5170 }
5171 }
5172
5173 struct FlakyStore {
5175 put_attempts: Arc<AtomicU32>,
5176 value: Arc<std::sync::Mutex<Option<Value>>>,
5177 }
5178 #[async_trait]
5179 impl StateStore for FlakyStore {
5180 async fn get(&self, _key: &str) -> Result<Option<Value>, FaucetError> {
5181 Ok(self.value.lock().unwrap().clone())
5182 }
5183 async fn put(&self, _key: &str, value: &Value) -> Result<(), FaucetError> {
5184 let n = self.put_attempts.fetch_add(1, Ordering::SeqCst);
5185 if n < 2 {
5186 return Err(transient_503());
5187 }
5188 *self.value.lock().unwrap() = Some(value.clone());
5189 Ok(())
5190 }
5191 async fn delete(&self, _key: &str) -> Result<(), FaucetError> {
5192 Ok(())
5193 }
5194 }
5195
5196 let flush_attempts = Arc::new(AtomicU32::new(0));
5197 let put_attempts = Arc::new(AtomicU32::new(0));
5198 let stored = Arc::new(std::sync::Mutex::new(None));
5199 let sink = FlakyFlushSink {
5200 flush_attempts: flush_attempts.clone(),
5201 };
5202 let store: Arc<dyn StateStore> = Arc::new(FlakyStore {
5203 put_attempts: put_attempts.clone(),
5204 value: stored.clone(),
5205 });
5206 let pages = futures::stream::iter(vec![Ok(StreamPage {
5207 records: vec![json!({"ok": 1}), json!({"bad": true})],
5208 bookmark: Some(json!({"cursor": 42})),
5209 })]);
5210 let policy = ResiliencePolicy {
5211 retry: RetryPolicy {
5212 max_attempts: 5,
5213 backoff: BackoffKind::None,
5214 base: Duration::ZERO,
5215 max: Duration::ZERO,
5216 jitter: false,
5217 ..RetryPolicy::default()
5218 },
5219 ..ResiliencePolicy::default()
5220 };
5221 let res = run_stream(
5222 pages,
5223 &sink,
5224 RunStreamOptions::new()
5225 .with_dlq(DlqConfig::new(Arc::new(NullSink)))
5226 .with_state(store, "k")
5227 .with_resilience(policy),
5228 )
5229 .await
5230 .unwrap();
5231
5232 assert_eq!(
5236 flush_attempts.load(Ordering::SeqCst),
5237 4,
5238 "page-gate flush retried past two transient failures (3) + 1 final flush"
5239 );
5240 assert_eq!(
5242 put_attempts.load(Ordering::SeqCst),
5243 3,
5244 "state put retried past two transient failures"
5245 );
5246 assert_eq!(*stored.lock().unwrap(), Some(json!({"cursor": 42})));
5247 assert_eq!(res.records_written, 1, "the ok row");
5248 }
5249
5250 #[tokio::test]
5251 #[allow(clippy::await_holding_lock)]
5252 async fn resilience_emits_retries_total_with_op_and_class_labels() {
5253 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
5257 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
5258 use metrics_util::debugging::DebugValue;
5259 use std::sync::Arc;
5260 use std::sync::atomic::{AtomicU32, Ordering};
5261 use std::time::Duration;
5262
5263 struct RetryProbeSink {
5264 attempts: Arc<AtomicU32>,
5265 }
5266 #[async_trait]
5267 impl Sink for RetryProbeSink {
5268 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
5269 let n = self.attempts.fetch_add(1, Ordering::SeqCst);
5270 if n < 2 {
5271 return Err(FaucetError::HttpStatus {
5272 status: 503,
5273 url: "u".into(),
5274 body: "".into(),
5275 });
5276 }
5277 Ok(r.len())
5278 }
5279 async fn flush(&self) -> Result<(), FaucetError> {
5280 Ok(())
5281 }
5282 fn connector_name(&self) -> &'static str {
5285 "retry-probe"
5286 }
5287 fn supports_idempotent_writes(&self) -> bool {
5289 true
5290 }
5291 }
5292
5293 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
5294 let snap = snapshotter();
5295
5296 let attempts = Arc::new(AtomicU32::new(0));
5297 let sink = RetryProbeSink {
5298 attempts: attempts.clone(),
5299 };
5300 let pages = futures::stream::iter(vec![Ok(StreamPage {
5301 records: vec![json!({"a": 1})],
5302 bookmark: None,
5303 })]);
5304 let policy = ResiliencePolicy {
5305 retry: RetryPolicy {
5306 max_attempts: 5,
5307 backoff: BackoffKind::None,
5308 base: Duration::ZERO,
5309 max: Duration::ZERO,
5310 jitter: false,
5311 ..RetryPolicy::default()
5312 },
5313 ..ResiliencePolicy::default()
5314 };
5315 run_stream(
5316 pages,
5317 &sink,
5318 RunStreamOptions::new()
5319 .with_name("retry-metrics-pipeline")
5320 .with_resilience(policy),
5321 )
5322 .await
5323 .unwrap();
5324 assert_eq!(attempts.load(Ordering::SeqCst), 3);
5325
5326 let snapshot = snap.snapshot();
5327 let retries: u64 = snapshot
5328 .into_vec()
5329 .into_iter()
5330 .filter_map(|(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
5331 if key.key().name() == "faucet_resilience_retries_total"
5332 && key.key().labels().any(|l: &metrics::Label| {
5333 l.key() == "pipeline" && l.value() == "retry-metrics-pipeline"
5334 })
5335 && key
5336 .key()
5337 .labels()
5338 .any(|l: &metrics::Label| l.key() == "op" && l.value() == "sink_write")
5339 && key
5340 .key()
5341 .labels()
5342 .any(|l: &metrics::Label| l.key() == "class" && l.value() == "http_5xx")
5343 && let DebugValue::Counter(c) = v
5344 {
5345 Some(c)
5346 } else {
5347 None
5348 }
5349 })
5350 .sum();
5351 assert_eq!(
5352 retries, 2,
5353 "expected 2 retries (2 transient 503s) counted with op=sink_write, class=http_5xx"
5354 );
5355 }
5356
5357 struct SchemaSink {
5361 schema: Value,
5362 written: std::sync::Mutex<Vec<Value>>,
5363 evolutions: std::sync::Mutex<Vec<crate::drift::SchemaEvolution>>,
5364 evolvable: bool,
5365 }
5366 impl SchemaSink {
5367 fn new(schema: Value, evolvable: bool) -> Self {
5368 Self {
5369 schema,
5370 written: std::sync::Mutex::new(Vec::new()),
5371 evolutions: std::sync::Mutex::new(Vec::new()),
5372 evolvable,
5373 }
5374 }
5375 fn written(&self) -> Vec<Value> {
5376 self.written.lock().unwrap().clone()
5377 }
5378 }
5379 #[async_trait]
5380 impl Sink for SchemaSink {
5381 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
5382 self.written.lock().unwrap().extend(records.iter().cloned());
5383 Ok(records.len())
5384 }
5385 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
5386 Ok(Some(self.schema.clone()))
5387 }
5388 fn supports_schema_evolution(&self) -> bool {
5389 self.evolvable
5390 }
5391 async fn evolve_schema(
5392 &self,
5393 evo: &crate::drift::SchemaEvolution,
5394 ) -> Result<(), FaucetError> {
5395 if !self.evolvable {
5396 return Err(FaucetError::Sink("not evolvable".into()));
5397 }
5398 self.evolutions.lock().unwrap().push(evo.clone());
5399 Ok(())
5400 }
5401 }
5402
5403 fn drift_opts(policy: crate::drift::SchemaDriftPolicy) -> RunStreamOptions {
5404 RunStreamOptions::new().with_schema_drift(policy)
5405 }
5406
5407 fn one_page(
5408 records: Vec<Value>,
5409 ) -> impl futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Unpin {
5410 Box::pin(futures::stream::iter(vec![Ok(StreamPage {
5411 records,
5412 bookmark: None,
5413 })]))
5414 }
5415
5416 #[tokio::test]
5417 async fn drift_warn_writes_unchanged() {
5418 let sink = SchemaSink::new(
5419 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5420 false,
5421 );
5422 let policy = crate::drift::SchemaDriftPolicy {
5423 on_drift: crate::drift::OnDrift::Warn,
5424 allow_widening: true,
5425 on_incompatible: crate::drift::OnIncompatible::Fail,
5426 relax_nullability_on_missing: false,
5427 };
5428 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5429 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5430 assert_eq!(res.records_written, 1);
5431 assert_eq!(sink.written()[0], json!({"id": 1, "email": "a@x.com"}));
5433 }
5434
5435 #[tokio::test]
5436 async fn drift_ignore_strips_unknown_fields() {
5437 let sink = SchemaSink::new(
5438 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5439 false,
5440 );
5441 let policy = crate::drift::SchemaDriftPolicy {
5442 on_drift: crate::drift::OnDrift::Ignore,
5443 allow_widening: true,
5444 on_incompatible: crate::drift::OnIncompatible::Fail,
5445 relax_nullability_on_missing: false,
5446 };
5447 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5448 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5449 assert_eq!(res.records_written, 1);
5450 assert_eq!(
5451 sink.written()[0],
5452 json!({"id": 1}),
5453 "email must be stripped"
5454 );
5455 }
5456
5457 #[tokio::test]
5458 async fn drift_fail_raises_schema_drift() {
5459 let sink = SchemaSink::new(
5460 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5461 false,
5462 );
5463 let policy = crate::drift::SchemaDriftPolicy {
5464 on_drift: crate::drift::OnDrift::Fail,
5465 allow_widening: true,
5466 on_incompatible: crate::drift::OnIncompatible::Fail,
5467 relax_nullability_on_missing: false,
5468 };
5469 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5470 let err = run_stream(pages, &sink, drift_opts(policy))
5471 .await
5472 .unwrap_err();
5473 assert!(matches!(err, FaucetError::SchemaDrift { .. }));
5474 }
5475
5476 #[tokio::test]
5477 async fn drift_evolve_calls_sink_then_writes() {
5478 let sink = SchemaSink::new(
5479 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5480 true,
5481 );
5482 let policy = crate::drift::SchemaDriftPolicy {
5483 on_drift: crate::drift::OnDrift::Evolve,
5484 allow_widening: true,
5485 on_incompatible: crate::drift::OnIncompatible::Fail,
5486 relax_nullability_on_missing: false,
5487 };
5488 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
5489 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5490 assert_eq!(res.records_written, 1);
5491 let evos = sink.evolutions.lock().unwrap();
5492 assert_eq!(evos.len(), 1);
5493 assert_eq!(evos[0].additions.len(), 1);
5494 assert_eq!(evos[0].additions[0].name, "email");
5495 assert_eq!(sink.written()[0], json!({"id": 1, "email": "a@x.com"}));
5497 }
5498
5499 #[tokio::test]
5500 async fn drift_evolve_does_not_relax_not_null_for_merely_absent_column() {
5501 let sink = SchemaSink::new(
5506 json!({"type":"object","properties":{
5507 "id":{"type":"integer"},
5508 "legacy":{"type":"string"}
5509 }}),
5510 true,
5511 );
5512 let policy = crate::drift::SchemaDriftPolicy {
5513 on_drift: crate::drift::OnDrift::Evolve,
5514 allow_widening: true,
5515 on_incompatible: crate::drift::OnIncompatible::Fail,
5516 relax_nullability_on_missing: false,
5517 };
5518 let pages = one_page(vec![json!({"id": 1})]);
5519 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5520 assert_eq!(res.records_written, 1);
5521 assert!(
5522 sink.evolutions.lock().unwrap().is_empty(),
5523 "NOT NULL must not be relaxed for a merely-absent column"
5524 );
5525 }
5526
5527 #[tokio::test]
5528 async fn drift_evolve_relaxes_absent_column_only_with_opt_in() {
5529 let sink = SchemaSink::new(
5531 json!({"type":"object","properties":{
5532 "id":{"type":"integer"},
5533 "legacy":{"type":"string"}
5534 }}),
5535 true,
5536 );
5537 let policy = crate::drift::SchemaDriftPolicy {
5538 on_drift: crate::drift::OnDrift::Evolve,
5539 allow_widening: true,
5540 on_incompatible: crate::drift::OnIncompatible::Fail,
5541 relax_nullability_on_missing: true,
5542 };
5543 let pages = one_page(vec![json!({"id": 1})]);
5544 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5545 assert_eq!(res.records_written, 1);
5546 let evos = sink.evolutions.lock().unwrap();
5547 assert_eq!(evos.len(), 1);
5548 assert_eq!(evos[0].relax_nullability, vec!["legacy".to_string()]);
5549 }
5550
5551 #[tokio::test]
5552 async fn drift_inert_when_sink_reports_no_schema() {
5553 let sink = MockSink::new();
5555 let policy = crate::drift::SchemaDriftPolicy {
5556 on_drift: crate::drift::OnDrift::Fail, allow_widening: true,
5558 on_incompatible: crate::drift::OnIncompatible::Fail,
5559 relax_nullability_on_missing: false,
5560 };
5561 let pages = one_page(vec![json!({"id": 1, "anything": true})]);
5562 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
5563 assert_eq!(res.records_written, 1);
5564 }
5565
5566 #[tokio::test]
5567 async fn drift_quarantine_routes_drift_rows_to_dlq() {
5568 let sink = SchemaSink::new(
5569 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5570 false,
5571 );
5572 let dlq_sink = std::sync::Arc::new(MockSink::new());
5573 let policy = crate::drift::SchemaDriftPolicy {
5574 on_drift: crate::drift::OnDrift::Quarantine,
5575 allow_widening: true,
5576 on_incompatible: crate::drift::OnIncompatible::Fail,
5577 relax_nullability_on_missing: false,
5578 };
5579 let pages = one_page(vec![
5580 json!({"id": 1}), json!({"id": 2, "email": "x"}), ]);
5583 let opts = RunStreamOptions::new()
5584 .with_schema_drift(policy)
5585 .with_dlq(crate::dlq::DlqConfig::new(dlq_sink.clone()));
5586 let res = run_stream(pages, &sink, opts).await.unwrap();
5587 assert_eq!(res.records_written, 1, "only the conforming row is written");
5588 assert_eq!(sink.written(), vec![json!({"id": 1})]);
5589 let dlq = dlq_sink.written();
5591 assert_eq!(dlq.len(), 1);
5592 assert_eq!(dlq[0]["payload"], json!({"id": 2, "email": "x"}));
5593 assert_eq!(dlq[0]["error"]["kind"], "SchemaDrift");
5594 }
5595
5596 #[test]
5597 fn quarantine_drift_rows_covers_widening_and_droppable_required() {
5598 use crate::drift::{ColumnChange, SchemaDiff};
5599 let diff = SchemaDiff {
5601 additions: vec![],
5602 widenings: vec![ColumnChange {
5603 name: "amount".into(),
5604 from: Some(json!({"type":"integer"})),
5605 to: json!({"type":"number"}),
5606 }],
5607 incompatible: vec![],
5608 droppable_required: vec!["legacy".into()],
5609 };
5610 let records = vec![
5611 json!({"id": 1, "amount": 1.5, "legacy": "x"}), json!({"id": 2, "amount": 7, "legacy": "y"}), json!({"id": 3, "legacy": "z"}), json!({"id": 4}), ];
5616 let page_indices = vec![10, 11, 12, 13];
5619 let (kept, env) = quarantine_drift_rows(&diff, records, &page_indices, "sink", "pl", "");
5620 assert_eq!(kept, vec![json!({"id": 3, "legacy": "z"})]);
5621 assert_eq!(
5622 env.len(),
5623 3,
5624 "widening rows + the missing-required row quarantined"
5625 );
5626 let indices: Vec<i64> = env
5628 .iter()
5629 .map(|e| e["record_index"].as_i64().unwrap())
5630 .collect();
5631 assert_eq!(indices, vec![10, 11, 13]);
5632 }
5633
5634 #[tokio::test]
5635 async fn drift_quarantine_without_dlq_is_rejected() {
5636 let sink = SchemaSink::new(json!({"type":"object","properties":{}}), false);
5637 let policy = crate::drift::SchemaDriftPolicy {
5638 on_drift: crate::drift::OnDrift::Quarantine,
5639 allow_widening: true,
5640 on_incompatible: crate::drift::OnIncompatible::Fail,
5641 relax_nullability_on_missing: false,
5642 };
5643 let pages = one_page(vec![json!({"id": 1})]);
5644 let err = run_stream(pages, &sink, drift_opts(policy))
5645 .await
5646 .unwrap_err();
5647 assert!(matches!(err, FaucetError::Config(_)));
5648 }
5649
5650 #[cfg(feature = "quality")]
5656 #[tokio::test]
5657 async fn drift_fail_with_dlq_still_routes_quality_quarantine() {
5658 use crate::dlq::DlqConfig;
5659 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
5660
5661 let sink = SchemaSink::new(
5664 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5665 false,
5666 );
5667 let dlq_sink = std::sync::Arc::new(MockSink::new());
5668 let policy = crate::drift::SchemaDriftPolicy {
5669 on_drift: crate::drift::OnDrift::Fail,
5670 allow_widening: true,
5671 on_incompatible: crate::drift::OnIncompatible::Fail,
5672 relax_nullability_on_missing: false,
5673 };
5674 let spec = QualitySpec {
5675 record: vec![RecordCheck::NotNull {
5676 field: "name".into(),
5677 treat_missing_as_null: true,
5678 on_failure: OnFailure::Quarantine,
5679 }],
5680 batch: vec![],
5681 };
5682 let quality = std::sync::Arc::new(CompiledQuality::compile(&spec).unwrap());
5683 let pages = one_page(vec![
5684 json!({"id": 1, "name": "ok", "email": "a@x"}), json!({"id": 2, "email": "b@x"}), ]);
5687 let opts = RunStreamOptions::new()
5688 .with_schema_drift(policy)
5689 .with_quality(quality)
5690 .with_dlq(DlqConfig::new(dlq_sink.clone()));
5691 let err = run_stream(pages, &sink, opts).await.unwrap_err();
5692 assert!(
5694 matches!(err, FaucetError::SchemaDrift { .. }),
5695 "got {err:?}"
5696 );
5697 let dlq = dlq_sink.written();
5699 assert_eq!(
5700 dlq.len(),
5701 1,
5702 "quarantined row must reach the DLQ before abort"
5703 );
5704 assert_eq!(dlq[0]["payload"], json!({"id": 2, "email": "b@x"}));
5705 assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
5706 assert_eq!(
5708 sink.written(),
5709 vec![json!({"id": 1, "name": "ok", "email": "a@x"})]
5710 );
5711 }
5712}
5713
5714#[cfg(test)]
5715mod cleanup_tests {
5716 use super::*;
5721 use crate::cleanup::{CleanupPolicy, CleanupTracker, SeenKeys};
5722 use async_trait::async_trait;
5723 use serde_json::json;
5724 use std::collections::BTreeMap;
5725 use std::sync::{Arc, Mutex};
5726 use tokio_util::sync::CancellationToken;
5727
5728 #[derive(Default)]
5731 struct CleanupSink {
5732 written: Mutex<Vec<Value>>,
5733 cleanup_calls: Mutex<Vec<Vec<String>>>,
5734 supports: bool,
5735 fail_writes: bool,
5736 }
5737
5738 impl CleanupSink {
5739 fn new() -> Self {
5740 Self {
5741 supports: true,
5742 ..Default::default()
5743 }
5744 }
5745 fn unsupported() -> Self {
5746 Self::default()
5747 }
5748 fn failing() -> Self {
5749 Self {
5750 supports: true,
5751 fail_writes: true,
5752 ..Default::default()
5753 }
5754 }
5755 fn calls(&self) -> Vec<Vec<String>> {
5756 self.cleanup_calls.lock().unwrap().clone()
5757 }
5758 }
5759
5760 #[async_trait]
5761 impl Sink for CleanupSink {
5762 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
5763 if self.fail_writes {
5764 return Err(FaucetError::Sink("boom".into()));
5765 }
5766 self.written.lock().unwrap().extend(records.iter().cloned());
5767 Ok(records.len())
5768 }
5769 fn supports_cleanup(&self) -> bool {
5770 self.supports
5771 }
5772 async fn cleanup_scope(
5773 &self,
5774 scope: &BTreeMap<String, Value>,
5775 seen: &SeenKeys,
5776 ) -> Result<u64, FaucetError> {
5777 assert!(
5778 !scope.is_empty(),
5779 "the pipeline must never pass an empty scope"
5780 );
5781 let mut ids: Vec<String> = seen
5782 .keys()
5783 .iter()
5784 .map(|k| {
5785 k.0.iter()
5786 .map(|(_, v)| v.to_string())
5787 .collect::<Vec<_>>()
5788 .join("|")
5789 })
5790 .collect();
5791 ids.sort();
5792 self.cleanup_calls.lock().unwrap().push(ids);
5793 Ok(3)
5794 }
5795 }
5796
5797 fn policy(max_keys: usize) -> Arc<CleanupPolicy> {
5798 Arc::new(
5799 CleanupPolicy::new(
5800 BTreeMap::from([("contact_id".to_string(), json!(7))]),
5801 vec!["id".to_string()],
5802 max_keys,
5803 )
5804 .unwrap(),
5805 )
5806 }
5807
5808 struct VecSource(Vec<Value>);
5810
5811 #[async_trait]
5812 impl Source for VecSource {
5813 async fn fetch_with_context(
5814 &self,
5815 _context: &std::collections::HashMap<String, Value>,
5816 ) -> Result<Vec<Value>, FaucetError> {
5817 Ok(self.0.clone())
5818 }
5819 }
5820
5821 fn src(records: Vec<Value>) -> VecSource {
5822 VecSource(records)
5823 }
5824
5825 #[tokio::test]
5828 async fn tracker_counts_rows_that_fail_into_the_dlq() {
5829 let sink = CleanupSink::new();
5833 let p = policy(100);
5834 let tracker = CleanupTracker::new(&sink, &p);
5835 let page = vec![json!({"id": 1}), json!({"id": 2})];
5836 tracker.write_batch_partial(&page).await.unwrap();
5838 assert_eq!(tracker.tracked(), 2);
5839 }
5840
5841 #[tokio::test]
5842 async fn tracker_refuses_to_delete_after_an_overflow() {
5843 let sink = CleanupSink::new();
5844 let p = policy(1);
5845 let tracker = CleanupTracker::new(&sink, &p);
5846 tracker
5847 .write_batch(&[json!({"id": 1}), json!({"id": 2})])
5848 .await
5849 .unwrap();
5850 let err = tracker
5851 .finish(&p)
5852 .await
5853 .expect_err("an overflowed tracker must refuse");
5854 assert!(err.to_string().contains("Nothing was deleted"), "{err}");
5855 assert!(
5856 sink.calls().is_empty(),
5857 "the sink must never be asked to delete"
5858 );
5859 }
5860
5861 #[tokio::test]
5864 async fn fires_once_at_a_natural_end_with_every_written_key() {
5865 let sink = CleanupSink::new();
5866 let source = src(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
5867 let result = Pipeline::new(&source, &sink)
5868 .with_cleanup(policy(100))
5869 .run()
5870 .await
5871 .expect("run succeeds");
5872 assert_eq!(result.records_written, 3);
5873 let calls = sink.calls();
5874 assert_eq!(calls.len(), 1, "exactly one cleanup per run");
5875 assert_eq!(calls[0], vec!["1", "2", "3"]);
5876 }
5877
5878 #[tokio::test]
5879 async fn empty_result_set_still_fires_with_no_keys() {
5880 let sink = CleanupSink::new();
5884 let source = src(vec![]);
5885 Pipeline::new(&source, &sink)
5886 .with_cleanup(policy(100))
5887 .run()
5888 .await
5889 .unwrap();
5890 let calls = sink.calls();
5891 assert_eq!(calls.len(), 1, "must still fire on an empty fetch");
5892 assert!(
5893 calls[0].is_empty(),
5894 "no keys written → delete the whole scope"
5895 );
5896 }
5897
5898 #[tokio::test]
5899 async fn skipped_when_the_run_was_cancelled() {
5900 let sink = CleanupSink::new();
5901 let source = src(vec![json!({"id": 1})]);
5902 let token = CancellationToken::new();
5903 token.cancel();
5904 Pipeline::new(&source, &sink)
5905 .with_cleanup(policy(100))
5906 .with_cancel(token)
5907 .run()
5908 .await
5909 .expect("a cancelled run still returns Ok with partial output flushed");
5910 assert!(
5911 sink.calls().is_empty(),
5912 "cleanup must not run after a cancellation"
5913 );
5914 }
5915
5916 #[tokio::test]
5917 async fn not_fired_when_the_run_failed() {
5918 let sink = CleanupSink::failing();
5919 let source = src(vec![json!({"id": 1})]);
5920 let err = Pipeline::new(&source, &sink)
5921 .with_cleanup(policy(100))
5922 .run()
5923 .await
5924 .expect_err("the write failure must propagate");
5925 assert!(err.to_string().contains("boom"), "{err}");
5926 assert!(
5927 sink.calls().is_empty(),
5928 "a failed run never wrote the authoritative set"
5929 );
5930 }
5931
5932 #[tokio::test]
5933 async fn refuses_and_deletes_nothing_when_the_key_ceiling_is_breached() {
5934 let sink = CleanupSink::new();
5935 let source = src(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
5936 let err = Pipeline::new(&source, &sink)
5937 .with_cleanup(policy(2))
5938 .run()
5939 .await
5940 .expect_err("must fail rather than partially delete");
5941 assert!(err.to_string().contains("Nothing was deleted"), "{err}");
5942 assert!(sink.calls().is_empty());
5943 }
5944
5945 #[tokio::test]
5946 async fn rejects_a_sink_that_does_not_support_cleanup() {
5947 let sink = CleanupSink::unsupported();
5948 let source = src(vec![json!({"id": 1})]);
5949 let err = Pipeline::new(&source, &sink)
5950 .with_cleanup(policy(100))
5951 .run()
5952 .await
5953 .expect_err("unsupported sink must be rejected");
5954 assert!(
5955 err.to_string().contains("does not support scoped cleanup"),
5956 "{err}"
5957 );
5958 assert!(
5959 sink.written.lock().unwrap().is_empty(),
5960 "must fail before writing anything"
5961 );
5962 }
5963
5964 #[tokio::test]
5965 async fn no_policy_leaves_the_write_path_untouched() {
5966 let sink = CleanupSink::new();
5967 let source = src(vec![json!({"id": 1})]);
5968 let result = Pipeline::new(&source, &sink).run().await.unwrap();
5969 assert_eq!(result.records_written, 1);
5970 assert!(sink.calls().is_empty(), "no policy → no cleanup");
5971 }
5972
5973 #[tokio::test]
5974 async fn rows_missing_the_key_are_not_tracked_but_do_not_break_the_pass() {
5975 let sink = CleanupSink::new();
5976 let source = src(vec![json!({"id": 1}), json!({"no_key": true})]);
5977 Pipeline::new(&source, &sink)
5978 .with_cleanup(policy(100))
5979 .run()
5980 .await
5981 .unwrap();
5982 assert_eq!(sink.calls()[0], vec!["1"]);
5983 }
5984}