1use crate::dlq::{DlqConfig, DlqStats};
43use crate::error::FaucetError;
44use crate::observability::RunStreamOptions;
45use crate::state::{StateStore, validate_state_key};
46use crate::traits::{Sink, Source};
47use futures_core::Stream;
48use serde_json::Value;
49use std::pin::Pin;
50use std::sync::Arc;
51
52pub const DEFAULT_BATCH_SIZE: usize = 1000;
58
59pub const MAX_BATCH_SIZE: usize = 1_000_000;
64
65pub fn validate_batch_size(batch_size: usize) -> Result<usize, FaucetError> {
77 if batch_size > MAX_BATCH_SIZE {
78 return Err(FaucetError::Config(format!(
79 "batch_size {batch_size} exceeds maximum {MAX_BATCH_SIZE} \
80 (use 0 to opt out of batching entirely)"
81 )));
82 }
83 Ok(batch_size)
84}
85
86#[derive(Debug, Clone, Default)]
96pub struct StreamPage {
97 pub records: Vec<Value>,
99 pub bookmark: Option<Value>,
101}
102
103#[derive(Debug, Clone)]
105pub struct PipelineResult {
106 pub records_written: usize,
108 pub bookmark: Option<Value>,
116 pub dlq: Option<DlqStats>,
118}
119
120pub struct Pipeline<'a, So: Source + ?Sized, Si: Sink + ?Sized> {
125 source: &'a So,
126 sink: &'a Si,
127 state_store: Option<Arc<dyn StateStore>>,
128 name: Option<String>,
129 row: Option<String>,
130 run_id: Option<String>,
131 dlq: Option<DlqConfig>,
132 #[cfg(feature = "quality")]
133 quality: Option<Arc<crate::quality::CompiledQuality>>,
134 #[cfg(feature = "contract")]
135 contract: Option<Arc<crate::contract::CompiledContract>>,
136 #[cfg(feature = "masking")]
137 masking: Option<Arc<crate::masking::CompiledMasking>>,
138 adaptive: Option<crate::adaptive::AdaptiveBatchConfig>,
139 cancel: Option<tokio_util::sync::CancellationToken>,
140 delivery: crate::idempotency::DeliveryMode,
141 resilience: Option<crate::resilience::ResiliencePolicy>,
142 schema_drift: Option<crate::drift::SchemaDriftPolicy>,
143}
144
145impl<'a, So: Source + ?Sized, Si: Sink + ?Sized> Pipeline<'a, So, Si> {
146 pub fn new(source: &'a So, sink: &'a Si) -> Self {
148 Self {
149 source,
150 sink,
151 state_store: None,
152 name: None,
153 row: None,
154 run_id: None,
155 dlq: None,
156 #[cfg(feature = "quality")]
157 quality: None,
158 #[cfg(feature = "contract")]
159 contract: None,
160 #[cfg(feature = "masking")]
161 masking: None,
162 adaptive: None,
163 cancel: None,
164 delivery: crate::idempotency::DeliveryMode::AtLeastOnce,
165 resilience: None,
166 schema_drift: None,
167 }
168 }
169
170 pub fn with_state_store(mut self, store: Arc<dyn StateStore>) -> Self {
184 self.state_store = Some(store);
185 self
186 }
187
188 pub fn with_name(mut self, name: impl Into<String>) -> Self {
191 self.name = Some(name.into());
192 self
193 }
194
195 pub fn with_row(mut self, row: impl Into<String>) -> Self {
198 self.row = Some(row.into());
199 self
200 }
201
202 pub fn with_run_id(mut self, run_id: impl Into<String>) -> Self {
206 self.run_id = Some(run_id.into());
207 self
208 }
209
210 pub fn with_dlq(mut self, dlq: DlqConfig) -> Self {
212 self.dlq = Some(dlq);
213 self
214 }
215
216 #[cfg(feature = "quality")]
219 pub fn with_quality(mut self, quality: Arc<crate::quality::CompiledQuality>) -> Self {
220 self.quality = Some(quality);
221 self
222 }
223
224 #[cfg(feature = "contract")]
227 pub fn with_contract(mut self, contract: Arc<crate::contract::CompiledContract>) -> Self {
228 self.contract = Some(contract);
229 self
230 }
231
232 #[cfg(feature = "masking")]
236 pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
237 self.masking = Some(masking);
238 self
239 }
240
241 pub fn with_adaptive(mut self, cfg: crate::adaptive::AdaptiveBatchConfig) -> Self {
245 self.adaptive = Some(cfg);
246 self
247 }
248
249 pub fn with_cancel(mut self, cancel: tokio_util::sync::CancellationToken) -> Self {
254 self.cancel = Some(cancel);
255 self
256 }
257
258 pub fn with_delivery(mut self, mode: crate::idempotency::DeliveryMode) -> Self {
263 self.delivery = mode;
264 self
265 }
266
267 pub fn with_resilience(mut self, policy: crate::resilience::ResiliencePolicy) -> Self {
269 self.resilience = Some(policy);
270 self
271 }
272
273 pub fn with_schema_drift(mut self, policy: crate::drift::SchemaDriftPolicy) -> Self {
276 self.schema_drift = Some(policy);
277 self
278 }
279
280 pub async fn run(&self) -> Result<PipelineResult, FaucetError> {
294 use crate::observability::{
295 DurationGuard, InstrumentedSink, InstrumentedSource, InstrumentedStateStore, Labels,
296 };
297 use metrics::{Label, SharedString, counter, gauge};
298 use tracing::Instrument;
299
300 let name = self.name.clone().unwrap_or_else(|| "unnamed".to_string());
302 let row = self.row.clone().unwrap_or_default();
303 let run_id = self
304 .run_id
305 .clone()
306 .unwrap_or_else(|| uuid::Uuid::now_v7().to_string());
307 let obs_labels = Labels::new(name.clone(), row.clone(), run_id.clone());
308
309 let wrapped_source = InstrumentedSource::new(self.source, obs_labels.clone());
311 let wrapped_sink = InstrumentedSink::new(self.sink, obs_labels.clone());
312 let wrapped_state_store: Option<Arc<dyn StateStore>> = self.state_store.as_ref().map(|s| {
313 Arc::new(InstrumentedStateStore::new(
314 Arc::clone(s),
315 obs_labels.clone(),
316 )) as Arc<dyn StateStore>
317 });
318
319 let span = tracing::info_span!(
322 "faucet.pipeline.run",
323 pipeline = %name,
324 row = %row,
325 run_id = %run_id,
326 source = %wrapped_source.connector_name(),
327 sink = %wrapped_sink.connector_name(),
328 );
329
330 let base_labels: Vec<Label> = vec![
332 Label::new("pipeline", SharedString::from(name.clone())),
333 Label::new("row", SharedString::from(row.clone())),
334 ];
335 let run_labels: Vec<Label> = {
336 let mut v = base_labels.clone();
337 v.push(Label::new(
338 "source",
339 SharedString::from(wrapped_source.connector_name().to_string()),
340 ));
341 v.push(Label::new(
342 "sink",
343 SharedString::from(wrapped_sink.connector_name().to_string()),
344 ));
345 v
346 };
347
348 struct InFlightGuard(Vec<Label>);
350 impl Drop for InFlightGuard {
351 fn drop(&mut self) {
352 gauge!("faucet_pipeline_in_flight", self.0.clone()).decrement(1.0);
353 }
354 }
355 gauge!("faucet_pipeline_in_flight", base_labels.clone()).increment(1.0);
356 let _in_flight = InFlightGuard(base_labels.clone());
357
358 let start_unix = std::time::SystemTime::now()
361 .duration_since(std::time::UNIX_EPOCH)
362 .map(|d| d.as_secs_f64())
363 .unwrap_or(0.0);
364 gauge!(
365 "faucet_pipeline_start_time_unix_seconds",
366 base_labels.clone()
367 )
368 .set(start_unix);
369
370 let _run_timer =
372 DurationGuard::new("faucet_pipeline_run_duration_seconds", run_labels.clone());
373
374 let result = async {
376 let state_key = self.source.state_key();
379 let mut start_seq = 0u64;
380 if let (Some(store), Some(key)) = (wrapped_state_store.as_ref(), state_key.as_ref()) {
381 validate_state_key(key)?;
382 if let Some(prior) = store.get(key).await? {
383 if self.delivery == crate::idempotency::DeliveryMode::ExactlyOnce {
384 let (bookmark, seq) = crate::idempotency::unwrap_state(&prior);
385 start_seq = seq;
386 if let Some(bm) = bookmark {
387 wrapped_source.apply_start_bookmark(bm).await?;
388 }
389 } else {
390 wrapped_source.apply_start_bookmark(prior).await?;
391 }
392 }
393 }
394
395 let ctx = std::collections::HashMap::new();
396 let pages = wrapped_source.stream_pages(&ctx, DEFAULT_BATCH_SIZE);
397
398 let mut opts = RunStreamOptions::new()
399 .with_name(name.clone())
400 .with_row(row.clone())
401 .with_run_id(run_id.clone());
402 if let (Some(store), Some(key)) = (wrapped_state_store.clone(), state_key) {
403 opts = opts.with_state(store, key);
404 }
405 if let Some(dlq) = self.dlq.clone() {
406 opts = opts.with_dlq(dlq);
407 }
408 #[cfg(feature = "quality")]
409 if let Some(q) = self.quality.clone() {
410 opts = opts.with_quality(q);
411 }
412 #[cfg(feature = "contract")]
413 if let Some(c) = self.contract.clone() {
414 opts = opts.with_contract(c);
415 }
416 #[cfg(feature = "masking")]
417 if let Some(m) = self.masking.clone() {
418 opts = opts.with_masking(m);
419 }
420 if let Some(ad) = self.adaptive.clone() {
421 opts = opts.with_adaptive(ad);
422 }
423 if let Some(cancel) = self.cancel.clone() {
424 opts = opts.with_cancel(cancel);
425 }
426 if let Some(policy) = self.resilience.clone() {
427 opts = opts.with_resilience(policy);
428 }
429 if let Some(p) = self.schema_drift {
430 opts = opts.with_schema_drift(p);
431 }
432 opts = opts.with_delivery(self.delivery).with_start_seq(start_seq);
433
434 run_stream(pages, &wrapped_sink, opts).await
435 }
436 .instrument(span)
437 .await;
438
439 let status = if result.is_ok() { "ok" } else { "err" };
443 let mut final_labels = run_labels;
444 final_labels.push(Label::new("status", SharedString::const_str(status)));
445 if let Err(ref e) = result {
446 final_labels.push(Label::new(
447 "kind",
448 SharedString::const_str(crate::observability::decorator::error_kind(e)),
449 ));
450 }
451 counter!("faucet_pipeline_runs_total", final_labels).increment(1);
452
453 result
454 }
455}
456
457pub async fn run_stream<S, Si>(
475 mut pages: S,
476 sink: &Si,
477 options: RunStreamOptions,
478) -> Result<PipelineResult, FaucetError>
479where
480 S: Stream<Item = Result<StreamPage, FaucetError>> + Unpin,
481 Si: Sink + ?Sized,
482{
483 use crate::dlq::{DlqReason, DlqStats, OnBatchError, build_envelope};
484
485 let state_store = options.state_store.clone();
486 let state_key = options.state_key.clone();
487 let pipeline_name = options.pipeline_name.unwrap_or_else(|| "unnamed".into());
488 let row = options.row.unwrap_or_default();
489 let run_id = options.run_id.unwrap_or_default();
490 let dlq = options.dlq.clone();
491 let cancel = options.cancel.clone();
492
493 #[cfg(feature = "quality")]
494 let quality = options.quality.clone();
495
496 #[cfg(feature = "quality")]
498 if let Some(q) = quality.as_ref()
499 && q.requires_dlq()
500 && dlq.is_none()
501 {
502 return Err(FaucetError::Config(
503 "quality: on_failure 'quarantine'/'quarantine_batch' requires a DLQ sink".into(),
504 ));
505 }
506
507 #[cfg(feature = "contract")]
508 let contract = options.contract.clone();
509 #[cfg(feature = "contract")]
511 if let Some(c) = contract.as_ref()
512 && c.requires_dlq()
513 && dlq.is_none()
514 {
515 return Err(FaucetError::Config(
516 "contract: on_breach 'quarantine' requires a DLQ sink".into(),
517 ));
518 }
519 #[cfg(feature = "contract")]
521 let mut warned_contract_breach = false;
522
523 #[cfg(feature = "masking")]
527 let masking = options.masking.clone();
528
529 let schema_drift = options.schema_drift;
531 if let Some(p) = schema_drift.as_ref()
533 && p.requires_dlq()
534 && dlq.is_none()
535 {
536 return Err(FaucetError::Config(
537 "schema: on_drift 'quarantine' (or on_incompatible 'quarantine') requires a DLQ sink"
538 .into(),
539 ));
540 }
541 let mut dest_schema_cache: Option<Option<Value>> = None;
545 let mut warned_drift_inert = false;
546
547 if let Some(key) = state_key.as_ref() {
548 validate_state_key(key)?;
549 }
550
551 let exactly_once = options.delivery == crate::idempotency::DeliveryMode::ExactlyOnce;
553 if exactly_once {
554 if !sink.supports_idempotent_writes() {
555 return Err(FaucetError::Config(format!(
556 "delivery: exactly_once requires an idempotent sink, but '{}' does not support it",
557 sink.connector_name()
558 )));
559 }
560 if state_store.is_none() || state_key.is_none() {
561 return Err(FaucetError::Config(
562 "delivery: exactly_once requires a state store".into(),
563 ));
564 }
565 if dlq.is_some() {
566 return Err(FaucetError::Config(
567 "delivery: exactly_once is not compatible with a DLQ in this version".into(),
568 ));
569 }
570 }
571 let scope = state_key.clone().unwrap_or_default();
572 let mut next_seq = options.start_seq;
573 let committed_seq = if exactly_once {
574 sink.last_committed_token(&scope)
575 .await?
576 .and_then(|t| crate::idempotency::parse_token(&t))
577 .unwrap_or(0)
578 } else {
579 0
580 };
581
582 let mut records_written = 0usize;
583 let mut last_bookmark: Option<Value> = None;
584 let mut dlq_stats = DlqStats::default();
585
586 let adaptive_cfg = options.adaptive.clone().filter(|c| c.enabled);
587 if let Some(cfg) = adaptive_cfg.as_ref() {
591 cfg.validate()?;
592 }
593 let mut controller: Option<crate::adaptive::AimdController> = None;
594 let mut warned_noop_sink = false;
595 let mut warned_poison_drop = false;
597
598 let sink_name = sink.connector_name();
599 let dlq_sink_name = dlq.as_ref().map(|d| d.sink.connector_name()).unwrap_or("");
600
601 let mut cancelled = false;
612
613 let resilience = options.resilience.clone();
620 let retry_policy = resilience.as_ref().map(|r| r.retry.clone());
621 let mut breaker = resilience
622 .as_ref()
623 .and_then(|r| r.circuit_breaker)
624 .map(|cb| {
625 (
626 crate::resilience::CircuitBreaker::new(cb.consecutive_failures),
627 cb.cooldown,
628 )
629 });
630 let poison = resilience.as_ref().and_then(|r| r.poison);
632
633 macro_rules! with_retry {
646 ($op_label:literal, $op:expr) => {
647 match &retry_policy {
648 Some(p) => {
649 let m = crate::resilience::RetryMetrics {
650 pipeline: pipeline_name.to_string(),
651 row: row.to_string(),
652 op: $op_label,
653 };
654 crate::resilience::execute_with_policy_metered(p, cancel.as_ref(), &m, || $op)
655 .await
656 }
657 None => $op.await,
658 }
659 };
660 }
661
662 macro_rules! with_retry_write {
673 ($op_label:literal, $op:expr) => {
674 if retry_policy.is_some() && sink.supports_idempotent_writes() {
675 with_retry!($op_label, $op)
676 } else {
677 $op.await
678 }
679 };
680 }
681
682 let loop_result: Result<(), FaucetError> = async {
683 loop {
684 let page = match &cancel {
688 Some(token) => tokio::select! {
689 biased;
690 _ = token.cancelled() => {
691 cancelled = true;
692 break;
693 }
694 p = std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)) => p,
695 },
696 None => std::future::poll_fn(|cx| Pin::new(&mut pages).poll_next(cx)).await,
697 };
698 match page {
699 Some(Ok(page)) => {
700 if page.records.is_empty() && page.bookmark.is_none() {
701 continue;
702 }
703
704 #[cfg(feature = "masking")]
710 let page = if let Some(m) = masking.as_ref() {
711 let labels =
712 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
713 let outcome = crate::observability::instrumented_apply_masking(
714 page.records,
715 m,
716 &labels,
717 );
718 StreamPage {
719 records: outcome.records,
720 bookmark: page.bookmark,
721 }
722 } else {
723 page
724 };
725
726 #[cfg(feature = "quality")]
728 let (records, quality_envelopes): (Vec<Value>, Vec<Value>) =
729 if let Some(q) = quality.as_ref() {
730 let labels =
731 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
732 let outcome = crate::observability::instrumented_apply_quality(
733 page.records,
734 q,
735 &labels,
736 )?;
737 let envelopes: Vec<Value> = outcome
738 .quarantined
739 .iter()
740 .map(|qr| {
741 let err = FaucetError::QualityFailure {
742 check: qr.check.to_string(),
743 message: qr.message.clone(),
744 };
745 build_envelope(
749 &qr.record,
750 &err,
751 DlqReason::Quality,
752 sink_name,
753 &pipeline_name,
754 &row,
755 qr.page_index,
756 )
757 })
758 .collect();
759 (outcome.survivors, envelopes)
760 } else {
761 (page.records, Vec::new())
762 };
763 #[cfg(not(feature = "quality"))]
764 let (records, quality_envelopes): (Vec<Value>, Vec<Value>) =
765 (page.records, Vec::new());
766
767 #[cfg(feature = "contract")]
774 let (records, contract_envelopes): (Vec<Value>, Vec<Value>) =
775 if let Some(c) = contract.as_ref() {
776 let labels =
777 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
778 let outcome = crate::observability::instrumented_apply_contract(
779 records, c, &labels,
780 )?;
781 if !outcome.warned.is_empty() && !warned_contract_breach {
782 tracing::warn!(
783 version = %c.version,
784 breaches = outcome.warned.len(),
785 first = %outcome.warned[0].describe(),
786 "contract: breaching records written unchanged \
787 (on_breach=warn); this warning fires once per run"
788 );
789 warned_contract_breach = true;
790 }
791 let envelopes: Vec<Value> = outcome
792 .quarantined
793 .iter()
794 .map(|vr| {
795 let err = FaucetError::ContractViolation {
796 version: c.version.clone(),
797 message: vr.violation.describe(),
798 };
799 build_envelope(
802 &vr.record,
803 &err,
804 DlqReason::Contract,
805 sink_name,
806 &pipeline_name,
807 &row,
808 vr.violation.page_index,
809 )
810 })
811 .collect();
812 (outcome.survivors, envelopes)
813 } else {
814 (records, Vec::new())
815 };
816
817 let mut drift_envelopes: Vec<Value> = Vec::new();
819 let (records, drift_abort): (Vec<Value>, Option<FaucetError>) =
820 if let Some(policy) = schema_drift.as_ref().filter(|_| !records.is_empty()) {
821 if dest_schema_cache.is_none() {
823 dest_schema_cache = Some(sink.current_schema().await?);
824 }
825 let dest = dest_schema_cache.as_ref().and_then(|o| o.as_ref());
826 match dest {
827 None => {
828 if !warned_drift_inert {
829 tracing::info!(
830 connector = sink_name,
831 "schema-drift: sink reports no destination schema; \
832 drift handling is inert this run"
833 );
834 warned_drift_inert = true;
835 }
836 (records, None)
837 }
838 Some(dest) => {
839 let inferred = crate::schema::infer_schema(&records);
840 let diff = crate::drift::diff_schema(
841 dest,
842 &inferred,
843 policy.allow_widening,
844 );
845 if diff.is_empty() {
846 (records, None)
847 } else {
848 let dest_owned = dest.clone();
852 apply_drift_policy(
853 policy,
854 &diff,
855 &dest_owned,
856 records,
857 sink,
858 sink_name,
859 &pipeline_name,
860 &row,
861 &mut dest_schema_cache,
862 &mut drift_envelopes,
863 )
864 .await?
865 }
866 }
867 }
868 } else {
869 (records, None)
870 };
871 let quality_envelopes = {
875 let mut q = quality_envelopes;
876 #[cfg(feature = "contract")]
877 q.extend(contract_envelopes);
878 q.append(&mut drift_envelopes);
879 q
880 };
881 let mut drift_abort = drift_abort;
891 if dlq.is_none()
892 && let Some(e) = drift_abort.take()
893 {
894 return Err(e);
895 }
896
897 let page = StreamPage {
898 records,
899 bookmark: page.bookmark,
900 };
901
902 if let Some(ref dlq_cfg) = dlq {
903 use metrics::{Label, SharedString, counter};
905 let metric_labels: Vec<Label> = vec![
906 Label::new("pipeline", SharedString::from(pipeline_name.clone())),
907 Label::new("row", SharedString::from(row.clone())),
908 Label::new("connector", SharedString::from(sink_name.to_string())),
909 Label::new(
910 "dlq_connector",
911 SharedString::from(dlq_sink_name.to_string()),
912 ),
913 ];
914 let span = tracing::info_span!(
915 "faucet.dlq.route",
916 pipeline = %pipeline_name,
917 row = %row,
918 run_id = %run_id,
919 connector = %sink_name,
920 dlq_connector = %dlq_sink_name,
921 );
922 let _enter = span.enter();
923
924 let mut envelopes: Vec<Value> = Vec::new();
928 let mut page_success = 0usize;
929 let mut outer_err_recovered = false;
930 let mut had_per_row_sink_failure = false;
935 let records_len = page.records.len();
936 let mut offset = 0usize;
937 while offset < records_len {
938 let size = match adaptive_cfg.as_ref() {
939 Some(cfg) => {
940 let ctrl = controller.get_or_insert_with(|| {
941 crate::adaptive::AimdController::new(cfg, records_len)
942 });
943 ctrl.current().max(1).min(records_len - offset)
944 }
945 None => records_len - offset, };
947 if adaptive_cfg.is_some() {
948 maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
949 }
950 let chunk = &page.records[offset..offset + size];
951 let t0 = std::time::Instant::now();
952 let chunk_outcomes_result =
958 with_retry_write!("sink_write", sink.write_batch_partial(chunk));
959 let latency = t0.elapsed();
960 let (mut chunk_outcomes, chunk_synthesized): (
969 Vec<crate::RowOutcome>,
970 bool,
971 ) = match chunk_outcomes_result {
972 Ok(o) => (o, false),
973 Err(e) => match dlq_cfg.on_batch_error {
974 OnBatchError::Propagate => return Err(e),
975 OnBatchError::DlqAll => {
976 outer_err_recovered = true;
977 let msg = e.to_string();
978 let synth = (0..chunk.len())
979 .map(|_| Err(FaucetError::Sink(msg.clone())))
980 .collect();
981 (synth, true)
982 }
983 },
984 };
985
986 if let Some(pp) = poison
995 && !chunk_synthesized
996 {
997 let mut attempt = 1u32; while attempt < pp.max_row_attempts {
999 let failing: Vec<usize> = chunk_outcomes
1000 .iter()
1001 .enumerate()
1002 .filter_map(|(j, o)| match o {
1003 Err(e)
1004 if retry_policy
1005 .as_ref()
1006 .map(|p| p.is_retriable(e))
1007 .unwrap_or(false) =>
1008 {
1009 Some(j)
1010 }
1011 _ => None,
1012 })
1013 .collect();
1014 if failing.is_empty() {
1015 break;
1016 }
1017 let subset: Vec<Value> =
1018 failing.iter().map(|&j| chunk[j].clone()).collect();
1019 let retried = sink.write_batch_partial(&subset).await?;
1027 let mut retried = retried.into_iter();
1031 for &j in failing.iter() {
1032 chunk_outcomes[j] = retried.next().unwrap_or(Ok(()));
1033 }
1034 attempt += 1;
1035 }
1036 }
1037
1038 let mut chunk_errors = 0usize;
1039 let mut poison_dlq = 0u64;
1045 let mut poison_drop = 0u64;
1046 for (j, outcome) in chunk_outcomes.iter().enumerate() {
1047 match outcome {
1048 Ok(()) => page_success += 1,
1049 Err(err) => {
1050 let action = poison
1055 .map(|pp| pp.action)
1056 .unwrap_or(crate::resilience::PoisonAction::Dlq);
1057 match action {
1058 crate::resilience::PoisonAction::Fail => {
1059 crate::observability::resilience::poison_rows(
1060 &pipeline_name,
1061 &row,
1062 "fail",
1063 1,
1064 );
1065 return Err(FaucetError::Sink(format!(
1066 "poison-pill row failed permanently: {err}"
1067 )));
1068 }
1069 crate::resilience::PoisonAction::Drop => {
1070 poison_drop += 1;
1073 if !warned_poison_drop {
1074 tracing::warn!(
1075 "poison-pill: dropping permanently-failing row(s) (action=drop); this warning fires once per run"
1076 );
1077 warned_poison_drop = true;
1078 }
1079 }
1080 crate::resilience::PoisonAction::Dlq => {
1081 poison_dlq += 1;
1082 chunk_errors += 1;
1083 if !chunk_synthesized {
1084 had_per_row_sink_failure = true;
1085 }
1086 envelopes.push(build_envelope(
1087 &chunk[j],
1088 err,
1089 DlqReason::Partial,
1090 sink_name,
1091 &pipeline_name,
1092 &row,
1093 offset + j,
1094 ));
1095 }
1096 }
1097 }
1098 }
1099 }
1100 if poison.is_some() {
1104 crate::observability::resilience::poison_rows(
1105 &pipeline_name,
1106 &row,
1107 "dlq",
1108 poison_dlq,
1109 );
1110 crate::observability::resilience::poison_rows(
1111 &pipeline_name,
1112 &row,
1113 "drop",
1114 poison_drop,
1115 );
1116 }
1117 if let Some(ctrl) = controller.as_mut() {
1118 let adj = ctrl.observe(crate::adaptive::Observation {
1119 batch_len: chunk.len(),
1120 errors: chunk_errors,
1121 latency,
1122 });
1123 emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
1124 }
1125 offset += size;
1126 }
1127 #[cfg(feature = "quality")]
1132 let quality_count = quality_envelopes.len();
1133 #[cfg(not(feature = "quality"))]
1134 let quality_count = 0usize;
1135 envelopes.splice(0..0, quality_envelopes);
1136 let page_failures = envelopes.len();
1137
1138 let mut budget_error: Option<FaucetError> = None;
1151 let mut circuit_error: Option<FaucetError> = None;
1160 if let Some((b, cooldown)) = breaker.as_mut() {
1161 if records_len > 0 && page_success == 0 {
1162 if b.record_failure() {
1163 crate::observability::resilience::circuit_opened(
1164 &pipeline_name,
1165 &row,
1166 );
1167 circuit_error = Some(FaucetError::CircuitOpen {
1168 failures: b.consecutive(),
1169 cooldown: *cooldown,
1170 });
1171 }
1172 } else if page_success > 0 {
1173 b.record_success();
1174 }
1175 }
1176 if let Some(limit) = dlq_cfg.max_failures_per_page
1177 && page_failures > limit
1178 {
1179 let mut lbl = metric_labels.clone();
1180 lbl.retain(|l| l.key() != "dlq_connector");
1181 lbl.push(Label::new("scope", SharedString::const_str("per_page")));
1182 counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
1183 budget_error = Some(FaucetError::Sink(format!(
1184 "DLQ per-page budget exceeded: {page_failures} > {limit}"
1185 )));
1186 }
1187 let new_total = dlq_stats.records_dlq + page_failures;
1188 if budget_error.is_none()
1189 && let Some(limit) = dlq_cfg.max_failures_total
1190 && new_total > limit
1191 {
1192 let mut lbl = metric_labels.clone();
1193 lbl.retain(|l| l.key() != "dlq_connector");
1194 lbl.push(Label::new("scope", SharedString::const_str("total")));
1195 counter!("faucet_sink_dlq_budget_exceeded_total", lbl).increment(1);
1196 budget_error = Some(FaucetError::Sink(format!(
1197 "DLQ total budget exceeded: {new_total} > {limit}"
1198 )));
1199 }
1200
1201 if !envelopes.is_empty() {
1203 let _dlq_write_timer = crate::observability::DurationGuard::new(
1204 "faucet_sink_dlq_write_duration_seconds",
1205 metric_labels.clone(),
1206 );
1207 dlq_cfg.sink.write_batch(&envelopes).await.map_err(|e| {
1208 let mut lbl = metric_labels.clone();
1209 lbl.push(Label::new(
1210 "kind",
1211 SharedString::const_str(
1212 crate::observability::decorator::error_kind(&e),
1213 ),
1214 ));
1215 counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1216 FaucetError::Sink(format!("DLQ sink write failed: {e}"))
1217 })?;
1218 dlq_stats.records_dlq += page_failures;
1219 dlq_stats.pages_with_failures += 1;
1220
1221 let reason_label = if had_per_row_sink_failure {
1240 DlqReason::Partial.as_str()
1241 } else if outer_err_recovered {
1242 DlqReason::DlqAll.as_str()
1243 } else if page_failures > quality_count {
1244 DlqReason::Partial.as_str()
1245 } else {
1246 DlqReason::Quality.as_str()
1247 };
1248 counter!("faucet_sink_dlq_records_total", metric_labels.clone())
1249 .increment(page_failures as u64);
1250 let mut page_labels = metric_labels.clone();
1251 page_labels
1252 .push(Label::new("reason", SharedString::const_str(reason_label)));
1253 counter!("faucet_sink_dlq_pages_total", page_labels).increment(1);
1254 }
1255
1256 records_written += page_success;
1257
1258 if let Some(bookmark) = page.bookmark {
1259 with_retry!("flush", sink.flush())?;
1265 let _dlq_flush_timer = crate::observability::DurationGuard::new(
1266 "faucet_sink_dlq_flush_duration_seconds",
1267 metric_labels.clone(),
1268 );
1269 with_retry!("flush", dlq_cfg.sink.flush()).map_err(|e| {
1270 let mut lbl = metric_labels.clone();
1271 lbl.push(Label::new(
1272 "kind",
1273 SharedString::const_str(
1274 crate::observability::decorator::error_kind(&e),
1275 ),
1276 ));
1277 counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1278 FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
1279 })?;
1280 let bm_labels =
1281 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1282 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1283 if let (Some(store), Some(key)) =
1284 (state_store.as_ref(), state_key.as_ref())
1285 {
1286 with_retry!("state_put", store.put(key, &bookmark))?;
1287 }
1288 last_bookmark = Some(bookmark);
1289 }
1290
1291 if let Some(e) = budget_error {
1298 return Err(e);
1299 }
1300 if let Some(e) = circuit_error {
1302 return Err(e);
1303 }
1304 if let Some(e) = drift_abort {
1309 return Err(e);
1310 }
1311 } else if exactly_once {
1312 if let Some(bookmark) = page.bookmark {
1316 next_seq += 1;
1317 let token = crate::idempotency::format_token(next_seq);
1318 if next_seq <= committed_seq {
1319 use metrics::{Label, SharedString, counter};
1323 let skip_labels: Vec<Label> = vec![
1324 Label::new(
1325 "pipeline",
1326 SharedString::from(pipeline_name.clone()),
1327 ),
1328 Label::new("row", SharedString::from(row.clone())),
1329 ];
1330 counter!("faucet_pipeline_pages_skipped_total", skip_labels)
1331 .increment(1);
1332 } else {
1333 records_written += with_retry!(
1334 "sink_write",
1335 sink.write_batch_idempotent(&page.records, &scope, &token)
1336 )?;
1337 }
1338 with_retry!("flush", sink.flush())?;
1339 let bm_labels =
1340 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1341 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1342 if let (Some(store), Some(key)) =
1343 (state_store.as_ref(), state_key.as_ref())
1344 {
1345 let wrapped =
1346 crate::idempotency::wrap_state(Some(&bookmark), next_seq);
1347 with_retry!("state_put", store.put(key, &wrapped))?;
1348 }
1349 last_bookmark = Some(bookmark);
1350 } else if !page.records.is_empty() {
1351 records_written +=
1355 with_retry_write!("sink_write", sink.write_batch(&page.records))?;
1356 }
1357 } else {
1358 debug_assert!(
1360 quality_envelopes.is_empty(),
1361 "quality quarantine without DLQ should have been rejected at run start"
1362 );
1363 if !page.records.is_empty() {
1364 if let Some(cfg) = adaptive_cfg.as_ref() {
1365 let ctrl = controller.get_or_insert_with(|| {
1366 crate::adaptive::AimdController::new(cfg, page.records.len())
1367 });
1368 maybe_warn_noop_sink(sink_name, &mut warned_noop_sink);
1369 let mut offset = 0;
1370 while offset < page.records.len() {
1371 let size =
1372 ctrl.current().max(1).min(page.records.len() - offset);
1373 let chunk = &page.records[offset..offset + size];
1374 let t0 = std::time::Instant::now();
1375 let n = with_retry_write!("sink_write", sink.write_batch(chunk))?;
1376 let latency = t0.elapsed();
1377 records_written += n;
1378 offset += size;
1379 let adj = ctrl.observe(crate::adaptive::Observation {
1380 batch_len: chunk.len(),
1381 errors: 0,
1382 latency,
1383 });
1384 emit_adaptive_metrics(ctrl, adj, &pipeline_name, &row);
1385 }
1386 } else {
1387 records_written +=
1388 with_retry_write!("sink_write", sink.write_batch(&page.records))?;
1389 }
1390 }
1391 if let Some(bookmark) = page.bookmark {
1392 with_retry!("flush", sink.flush())?;
1393 let bm_labels =
1394 crate::observability::Labels::new(&*pipeline_name, &*row, &*run_id);
1395 crate::observability::update_bookmark_lag(&bookmark, &bm_labels);
1396 if let (Some(store), Some(key)) =
1397 (state_store.as_ref(), state_key.as_ref())
1398 {
1399 with_retry!("state_put", store.put(key, &bookmark))?;
1400 }
1401 last_bookmark = Some(bookmark);
1402 }
1403 }
1404 }
1405 Some(Err(e)) => return Err(e),
1406 None => break,
1407 }
1408 }
1409 Ok(())
1410 }
1411 .await;
1412
1413 if let Err(e) = loop_result {
1420 if let Some(ref dlq_cfg) = dlq
1421 && let Err(flush_err) = dlq_cfg.sink.flush().await
1422 {
1423 tracing::warn!(
1424 error = %flush_err,
1425 "DLQ sink flush failed during error unwind; original error preserved"
1426 );
1427 }
1428 if let Err(flush_err) = sink.flush().await {
1429 tracing::warn!(
1430 error = %flush_err,
1431 "sink flush failed during error unwind; original error preserved"
1432 );
1433 }
1434 return Err(e);
1435 }
1436
1437 if let Some(ref dlq_cfg) = dlq {
1443 let final_metric_labels: Vec<metrics::Label> = vec![
1444 metrics::Label::new(
1445 "pipeline",
1446 metrics::SharedString::from(pipeline_name.clone()),
1447 ),
1448 metrics::Label::new("row", metrics::SharedString::from(row.clone())),
1449 metrics::Label::new(
1450 "connector",
1451 metrics::SharedString::from(sink_name.to_string()),
1452 ),
1453 metrics::Label::new(
1454 "dlq_connector",
1455 metrics::SharedString::from(dlq_sink_name.to_string()),
1456 ),
1457 ];
1458 let _final_dlq_flush_timer = crate::observability::DurationGuard::new(
1459 "faucet_sink_dlq_flush_duration_seconds",
1460 final_metric_labels.clone(),
1461 );
1462 dlq_cfg.sink.flush().await.map_err(|e| {
1463 let mut lbl = final_metric_labels.clone();
1464 lbl.push(metrics::Label::new(
1465 "kind",
1466 metrics::SharedString::const_str(crate::observability::decorator::error_kind(&e)),
1467 ));
1468 metrics::counter!("faucet_sink_dlq_errors_total", lbl).increment(1);
1469 FaucetError::Sink(format!("DLQ sink flush failed: {e}"))
1470 })?;
1471 }
1472 sink.flush().await?;
1473
1474 if cancelled {
1475 tracing::info!(
1476 records_written,
1477 "pipeline run cancelled cooperatively; sink flushed (partial output is durable)"
1478 );
1479 }
1480
1481 tracing::info!(
1482 records_written,
1483 cancelled,
1484 has_bookmark = last_bookmark.is_some(),
1485 persisted = state_store.is_some() && state_key.is_some() && last_bookmark.is_some(),
1486 dlq_records = dlq_stats.records_dlq,
1487 "pipeline streaming run complete"
1488 );
1489
1490 Ok(PipelineResult {
1491 records_written,
1492 bookmark: last_bookmark,
1493 dlq: dlq.is_some().then_some(dlq_stats),
1494 })
1495}
1496
1497fn emit_adaptive_metrics(
1500 ctrl: &crate::adaptive::AimdController,
1501 adj: Option<crate::adaptive::Adjustment>,
1502 pipeline: &str,
1503 row: &str,
1504) {
1505 use metrics::{Label, SharedString, counter, gauge};
1506 let base = vec![
1507 Label::new("pipeline", SharedString::from(pipeline.to_string())),
1508 Label::new("row", SharedString::from(row.to_string())),
1509 ];
1510 gauge!("faucet_pipeline_adaptive_batch_size", base.clone()).set(ctrl.current() as f64);
1511 gauge!(
1512 "faucet_pipeline_adaptive_batch_cooldown_active",
1513 base.clone()
1514 )
1515 .set(if ctrl.cooldown_active() { 1.0 } else { 0.0 });
1516 if let Some(p50) = ctrl.p50_latency_ms() {
1517 gauge!(
1518 "faucet_pipeline_adaptive_batch_p50_latency_ms",
1519 base.clone()
1520 )
1521 .set(p50 as f64);
1522 }
1523 if let Some(a) = adj {
1524 let mut lbl = base;
1525 lbl.push(Label::new(
1526 "direction",
1527 SharedString::const_str(a.direction.as_str()),
1528 ));
1529 lbl.push(Label::new(
1530 "reason",
1531 SharedString::const_str(a.reason.as_str()),
1532 ));
1533 counter!("faucet_pipeline_adaptive_batch_adjustments_total", lbl).increment(1);
1534 }
1535}
1536
1537fn maybe_warn_noop_sink(sink_name: &str, warned: &mut bool) {
1540 if !*warned && matches!(sink_name, "jsonl" | "csv" | "stdout") {
1541 tracing::info!(
1542 sink = sink_name,
1543 "adaptive batch sizing is a no-op for this per-record sink"
1544 );
1545 *warned = true;
1546 }
1547}
1548
1549#[allow(clippy::too_many_arguments)]
1557async fn apply_drift_policy<Si: Sink + ?Sized>(
1558 policy: &crate::drift::SchemaDriftPolicy,
1559 diff: &crate::drift::SchemaDiff,
1560 dest: &Value,
1561 records: Vec<Value>,
1562 sink: &Si,
1563 sink_name: &str,
1564 pipeline_name: &str,
1565 row: &str,
1566 dest_schema_cache: &mut Option<Option<Value>>,
1567 drift_envelopes: &mut Vec<Value>,
1568) -> Result<(Vec<Value>, Option<FaucetError>), FaucetError> {
1569 use crate::drift::{OnDrift, OnIncompatible};
1570 use crate::observability::schema_drift as emit_drift;
1571
1572 let mode = match policy.on_drift {
1573 OnDrift::Warn => "warn",
1574 OnDrift::Ignore => "ignore",
1575 OnDrift::Quarantine => "quarantine",
1576 OnDrift::Fail => "fail",
1577 OnDrift::Evolve => "evolve",
1578 };
1579 emit_drift(
1580 pipeline_name,
1581 row,
1582 sink_name,
1583 mode,
1584 "added",
1585 diff.additions.len() as u64,
1586 );
1587 emit_drift(
1588 pipeline_name,
1589 row,
1590 sink_name,
1591 mode,
1592 "widened",
1593 diff.widenings.len() as u64,
1594 );
1595 emit_drift(
1596 pipeline_name,
1597 row,
1598 sink_name,
1599 mode,
1600 "narrowed",
1601 diff.incompatible.len() as u64,
1602 );
1603 emit_drift(
1604 pipeline_name,
1605 row,
1606 sink_name,
1607 mode,
1608 "dropped",
1609 diff.droppable_required.len() as u64,
1610 );
1611
1612 match policy.on_drift {
1613 OnDrift::Warn => {
1614 tracing::warn!(
1615 connector = sink_name,
1616 columns = ?diff.changed_columns(),
1617 "schema-drift detected (on_drift=warn); writing page unchanged"
1618 );
1619 Ok((records, None))
1620 }
1621 OnDrift::Fail => Ok((
1622 records,
1623 Some(FaucetError::SchemaDrift {
1624 columns: diff.changed_columns(),
1625 message: "schema drift detected (on_drift=fail)".to_string(),
1626 }),
1627 )),
1628 OnDrift::Ignore => {
1629 let allowed: std::collections::HashSet<String> = dest
1631 .get("properties")
1632 .and_then(|p| p.as_object())
1633 .map(|m| m.keys().cloned().collect())
1634 .unwrap_or_default();
1635 let trimmed = records
1636 .into_iter()
1637 .map(|r| match r {
1638 Value::Object(map) => Value::Object(
1639 map.into_iter()
1640 .filter(|(k, _)| allowed.contains(k))
1641 .collect(),
1642 ),
1643 other => other,
1644 })
1645 .collect();
1646 Ok((trimmed, None))
1647 }
1648 OnDrift::Quarantine => {
1649 let (kept, env) = quarantine_drift_rows(diff, records, sink_name, pipeline_name, row);
1650 drift_envelopes.extend(env);
1651 Ok((kept, None))
1652 }
1653 OnDrift::Evolve => {
1654 let evolution = crate::drift::SchemaEvolution {
1655 additions: diff.additions.clone(),
1656 widenings: diff
1657 .widenings
1658 .iter()
1659 .filter(|c| {
1660 c.from
1661 .as_ref()
1662 .map(|f| crate::drift::base_widened(f, &c.to))
1663 .unwrap_or(false)
1664 })
1665 .cloned()
1666 .collect(),
1667 relax_nullability: diff
1668 .droppable_required
1669 .iter()
1670 .filter(|_| policy.relax_nullability_on_missing)
1675 .cloned()
1676 .chain(
1677 diff.widenings
1678 .iter()
1679 .filter(|c| {
1680 c.from
1681 .as_ref()
1682 .map(|f| crate::drift::adds_null(f, &c.to))
1683 .unwrap_or(false)
1684 })
1685 .map(|c| c.name.clone()),
1686 )
1687 .collect(),
1688 };
1689 if !evolution.is_empty() {
1690 sink.evolve_schema(&evolution).await?;
1691 *dest_schema_cache = Some(sink.current_schema().await?);
1694 }
1695 if diff.incompatible.is_empty() {
1697 Ok((records, None))
1698 } else {
1699 match policy.on_incompatible {
1700 OnIncompatible::Fail => Ok((
1701 records,
1702 Some(FaucetError::SchemaDrift {
1703 columns: diff.incompatible.iter().map(|c| c.name.clone()).collect(),
1704 message: "incompatible type change cannot be auto-evolved \
1705 (on_incompatible=fail)"
1706 .into(),
1707 }),
1708 )),
1709 OnIncompatible::Quarantine => {
1710 let incompat_only = crate::drift::SchemaDiff {
1712 incompatible: diff.incompatible.clone(),
1713 ..Default::default()
1714 };
1715 let (kept, env) = quarantine_drift_rows(
1716 &incompat_only,
1717 records,
1718 sink_name,
1719 pipeline_name,
1720 row,
1721 );
1722 drift_envelopes.extend(env);
1723 Ok((kept, None))
1724 }
1725 }
1726 }
1727 }
1728 }
1729}
1730
1731fn quarantine_drift_rows(
1741 diff: &crate::drift::SchemaDiff,
1742 records: Vec<Value>,
1743 sink_name: &str,
1744 pipeline_name: &str,
1745 row: &str,
1746) -> (Vec<Value>, Vec<Value>) {
1747 use crate::dlq::{DlqReason, build_envelope};
1748 let present_cols: std::collections::HashSet<&str> = diff
1750 .additions
1751 .iter()
1752 .chain(&diff.widenings)
1753 .chain(&diff.incompatible)
1754 .map(|c| c.name.as_str())
1755 .collect();
1756 let required_cols: std::collections::HashSet<&str> =
1758 diff.droppable_required.iter().map(|s| s.as_str()).collect();
1759 let mut kept = Vec::new();
1760 let mut envelopes = Vec::new();
1761 for (idx, rec) in records.into_iter().enumerate() {
1762 let exhibits = rec
1763 .as_object()
1764 .map(|m| {
1765 m.keys().any(|k| present_cols.contains(k.as_str()))
1766 || required_cols.iter().any(|c| !m.contains_key(*c))
1767 })
1768 .unwrap_or(false);
1769 if exhibits {
1770 let err = FaucetError::SchemaDrift {
1771 columns: diff.changed_columns(),
1772 message: "row exhibits schema drift (on_drift=quarantine)".into(),
1773 };
1774 envelopes.push(build_envelope(
1775 &rec,
1776 &err,
1777 DlqReason::SchemaDrift,
1778 sink_name,
1779 pipeline_name,
1780 row,
1781 idx,
1782 ));
1783 } else {
1784 kept.push(rec);
1785 }
1786 }
1787 (kept, envelopes)
1788}
1789
1790#[cfg(test)]
1791mod tests {
1792 use super::*;
1793 use async_trait::async_trait;
1794 use serde_json::json;
1795
1796 struct MockSource(Vec<Value>);
1799
1800 #[async_trait]
1801 impl Source for MockSource {
1802 async fn fetch_with_context(
1803 &self,
1804 _context: &std::collections::HashMap<String, Value>,
1805 ) -> Result<Vec<Value>, FaucetError> {
1806 Ok(self.0.clone())
1807 }
1808 }
1809
1810 struct IncrementalSource {
1811 records: Vec<Value>,
1812 bookmark: Value,
1813 }
1814
1815 #[async_trait]
1816 impl Source for IncrementalSource {
1817 async fn fetch_with_context(
1818 &self,
1819 _context: &std::collections::HashMap<String, Value>,
1820 ) -> Result<Vec<Value>, FaucetError> {
1821 Ok(self.records.clone())
1822 }
1823 async fn fetch_with_context_incremental(
1824 &self,
1825 _context: &std::collections::HashMap<String, Value>,
1826 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
1827 Ok((self.records.clone(), Some(self.bookmark.clone())))
1828 }
1829 }
1830
1831 struct FailingSource;
1832
1833 #[async_trait]
1834 impl Source for FailingSource {
1835 async fn fetch_with_context(
1836 &self,
1837 _context: &std::collections::HashMap<String, Value>,
1838 ) -> Result<Vec<Value>, FaucetError> {
1839 Err(FaucetError::Auth("no credentials".into()))
1840 }
1841 }
1842
1843 struct MockSink(std::sync::Mutex<Vec<Value>>);
1846
1847 impl MockSink {
1848 fn new() -> Self {
1849 Self(std::sync::Mutex::new(Vec::new()))
1850 }
1851 fn written(&self) -> Vec<Value> {
1852 self.0.lock().unwrap().clone()
1853 }
1854 }
1855
1856 #[async_trait]
1857 impl Sink for MockSink {
1858 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1859 self.0.lock().unwrap().extend(records.iter().cloned());
1860 Ok(records.len())
1861 }
1862 }
1863
1864 struct FailingSink;
1865
1866 #[async_trait]
1867 impl Sink for FailingSink {
1868 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
1869 Err(FaucetError::Sink("write failed".into()))
1870 }
1871 }
1872
1873 struct FlushTrackingSink {
1878 written: std::sync::Mutex<Vec<Value>>,
1879 flush_count: std::sync::atomic::AtomicUsize,
1880 }
1881
1882 impl FlushTrackingSink {
1883 fn new() -> Self {
1884 Self {
1885 written: std::sync::Mutex::new(Vec::new()),
1886 flush_count: std::sync::atomic::AtomicUsize::new(0),
1887 }
1888 }
1889 fn written(&self) -> Vec<Value> {
1890 self.written.lock().unwrap().clone()
1891 }
1892 fn flush_count(&self) -> usize {
1893 self.flush_count.load(std::sync::atomic::Ordering::SeqCst)
1894 }
1895 }
1896
1897 #[async_trait]
1898 impl Sink for FlushTrackingSink {
1899 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1900 self.written.lock().unwrap().extend(records.iter().cloned());
1901 Ok(records.len())
1902 }
1903 async fn flush(&self) -> Result<(), FaucetError> {
1904 self.flush_count
1905 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1906 Ok(())
1907 }
1908 }
1909
1910 struct IdempotentMockSink {
1914 rows: std::sync::Mutex<Vec<Value>>,
1915 tokens: std::sync::Mutex<std::collections::HashMap<String, String>>,
1916 }
1917 impl IdempotentMockSink {
1918 fn new() -> Self {
1919 Self {
1920 rows: std::sync::Mutex::new(Vec::new()),
1921 tokens: std::sync::Mutex::new(std::collections::HashMap::new()),
1922 }
1923 }
1924 fn rows(&self) -> Vec<Value> {
1925 self.rows.lock().unwrap().clone()
1926 }
1927 }
1928 #[async_trait]
1929 impl Sink for IdempotentMockSink {
1930 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1931 self.rows.lock().unwrap().extend(records.iter().cloned());
1932 Ok(records.len())
1933 }
1934 fn supports_idempotent_writes(&self) -> bool {
1935 true
1936 }
1937 async fn write_batch_idempotent(
1938 &self,
1939 records: &[Value],
1940 scope: &str,
1941 token: &str,
1942 ) -> Result<usize, FaucetError> {
1943 self.rows.lock().unwrap().extend(records.iter().cloned());
1944 self.tokens
1945 .lock()
1946 .unwrap()
1947 .insert(scope.to_string(), token.to_string());
1948 Ok(records.len())
1949 }
1950 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
1951 Ok(self.tokens.lock().unwrap().get(scope).cloned())
1952 }
1953 }
1954
1955 fn eo_opts(store: Arc<dyn StateStore>, key: &str, start_seq: u64) -> RunStreamOptions {
1956 RunStreamOptions::new()
1957 .with_state(store, key)
1958 .with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce)
1959 .with_start_seq(start_seq)
1960 }
1961
1962 #[tokio::test]
1963 async fn exactly_once_writes_pages_and_persists_wrapped_state() {
1964 let pages = vec![
1965 Ok(StreamPage {
1966 records: vec![json!({"id": 1})],
1967 bookmark: Some(json!("b1")),
1968 }),
1969 Ok(StreamPage {
1970 records: vec![json!({"id": 2})],
1971 bookmark: Some(json!("b2")),
1972 }),
1973 ];
1974 let sink = IdempotentMockSink::new();
1975 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
1976 let r = run_stream(
1977 futures::stream::iter(pages),
1978 &sink,
1979 eo_opts(store.clone(), "k", 0),
1980 )
1981 .await
1982 .unwrap();
1983 assert_eq!(r.records_written, 2);
1984 let (bm, seq) = crate::idempotency::unwrap_state(&store.get("k").await.unwrap().unwrap());
1985 assert_eq!(bm, Some(json!("b2")));
1986 assert_eq!(seq, 2);
1987 assert_eq!(
1988 sink.last_committed_token("k").await.unwrap(),
1989 Some(crate::idempotency::format_token(2))
1990 );
1991 }
1992
1993 #[tokio::test]
1994 async fn exactly_once_skips_already_committed_pages_on_resume() {
1995 let sink = IdempotentMockSink::new();
1996 sink.write_batch_idempotent(
1998 &[json!({"id": 1})],
1999 "k",
2000 &crate::idempotency::format_token(1),
2001 )
2002 .await
2003 .unwrap();
2004 assert_eq!(sink.rows().len(), 1);
2005 let pages = vec![
2007 Ok(StreamPage {
2008 records: vec![json!({"id": 1})],
2009 bookmark: Some(json!("b1")),
2010 }),
2011 Ok(StreamPage {
2012 records: vec![json!({"id": 2})],
2013 bookmark: Some(json!("b2")),
2014 }),
2015 ];
2016 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2017 let r = run_stream(futures::stream::iter(pages), &sink, eo_opts(store, "k", 0))
2018 .await
2019 .unwrap();
2020 assert_eq!(r.records_written, 1);
2021 let rows = sink.rows();
2022 assert_eq!(
2023 rows.len(),
2024 2,
2025 "exactly one row per id — no duplicate of id=1"
2026 );
2027 assert_eq!(rows.iter().filter(|v| v["id"] == 1).count(), 1);
2028 }
2029
2030 #[tokio::test]
2031 async fn exactly_once_rejects_non_idempotent_sink() {
2032 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2033 let store: Arc<dyn StateStore> = Arc::new(crate::state::MemoryStateStore::new());
2034 let r = run_stream(
2035 futures::stream::iter(pages),
2036 &MockSink::new(),
2037 eo_opts(store, "k", 0),
2038 )
2039 .await;
2040 assert!(matches!(r, Err(FaucetError::Config(_))));
2041 }
2042
2043 #[tokio::test]
2044 async fn exactly_once_rejects_missing_state_store() {
2045 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2046 let opts =
2047 RunStreamOptions::new().with_delivery(crate::idempotency::DeliveryMode::ExactlyOnce);
2048 let r = run_stream(
2049 futures::stream::iter(pages),
2050 &IdempotentMockSink::new(),
2051 opts,
2052 )
2053 .await;
2054 assert!(matches!(r, Err(FaucetError::Config(_))));
2055 }
2056
2057 #[test]
2060 fn stream_page_constructs() {
2061 let page = StreamPage {
2062 records: vec![json!({"id": 1})],
2063 bookmark: Some(json!("2026-05-18")),
2064 };
2065 assert_eq!(page.records.len(), 1);
2066 assert_eq!(page.bookmark, Some(json!("2026-05-18")));
2067 }
2068
2069 #[test]
2070 fn validate_batch_size_accepts_zero_as_no_batching_sentinel() {
2071 assert_eq!(validate_batch_size(0).unwrap(), 0);
2073 }
2074
2075 #[test]
2076 fn validate_batch_size_rejects_too_large() {
2077 let err = validate_batch_size(MAX_BATCH_SIZE + 1).unwrap_err();
2078 assert!(matches!(err, FaucetError::Config(_)));
2079 }
2080
2081 #[test]
2082 fn validate_batch_size_accepts_one() {
2083 assert_eq!(validate_batch_size(1).unwrap(), 1);
2084 }
2085
2086 #[test]
2087 fn validate_batch_size_accepts_max() {
2088 assert_eq!(validate_batch_size(MAX_BATCH_SIZE).unwrap(), MAX_BATCH_SIZE);
2089 }
2090
2091 const _: () = {
2093 assert!(DEFAULT_BATCH_SIZE >= 1);
2094 assert!(DEFAULT_BATCH_SIZE <= MAX_BATCH_SIZE);
2095 };
2096
2097 #[tokio::test]
2100 async fn batch_pipeline_writes_all_records() {
2101 let source = MockSource(vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})]);
2102 let sink = MockSink::new();
2103
2104 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2105
2106 assert_eq!(result.records_written, 3);
2107 assert!(result.bookmark.is_none());
2108 assert_eq!(sink.written().len(), 3);
2109 }
2110
2111 #[tokio::test]
2112 async fn batch_pipeline_returns_bookmark() {
2113 let source = IncrementalSource {
2114 records: vec![json!({"id": 1, "ts": "2024-12-01"})],
2115 bookmark: json!("2024-12-01"),
2116 };
2117 let sink = MockSink::new();
2118
2119 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2120
2121 assert_eq!(result.records_written, 1);
2122 assert_eq!(result.bookmark, Some(json!("2024-12-01")));
2123 }
2124
2125 #[tokio::test]
2126 async fn batch_pipeline_empty_source() {
2127 let source = MockSource(vec![]);
2128 let sink = MockSink::new();
2129
2130 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2131
2132 assert_eq!(result.records_written, 0);
2133 assert!(sink.written().is_empty());
2134 }
2135
2136 #[tokio::test]
2137 async fn batch_pipeline_source_error_propagates() {
2138 let source = FailingSource;
2139 let sink = MockSink::new();
2140
2141 let result = Pipeline::new(&source, &sink).run().await;
2142 assert!(result.is_err());
2143 assert!(sink.written().is_empty());
2144 }
2145
2146 #[tokio::test]
2147 async fn batch_pipeline_sink_error_propagates() {
2148 let source = MockSource(vec![json!({"id": 1})]);
2149 let sink = FailingSink;
2150
2151 let result = Pipeline::new(&source, &sink).run().await;
2152 assert!(result.is_err());
2153 }
2154
2155 #[tokio::test]
2156 async fn batch_pipeline_with_trait_objects() {
2157 let source: Box<dyn Source> = Box::new(MockSource(vec![json!({"id": 1})]));
2158 let sink: Box<dyn Sink> = Box::new(MockSink::new());
2159
2160 let result = Pipeline::new(source.as_ref(), sink.as_ref())
2161 .run()
2162 .await
2163 .unwrap();
2164
2165 assert_eq!(result.records_written, 1);
2166 }
2167
2168 #[tokio::test]
2171 async fn stream_pipeline_writes_pages() {
2172 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2173 Ok(StreamPage {
2174 records: vec![json!({"id": 1}), json!({"id": 2})],
2175 bookmark: None,
2176 }),
2177 Ok(StreamPage {
2178 records: vec![json!({"id": 3})],
2179 bookmark: None,
2180 }),
2181 ];
2182 let stream = futures::stream::iter(pages);
2183 let sink = MockSink::new();
2184
2185 let result = run_stream(stream, &sink, RunStreamOptions::new())
2186 .await
2187 .unwrap();
2188
2189 assert_eq!(result.records_written, 3);
2190 assert!(result.bookmark.is_none());
2191 assert_eq!(sink.written().len(), 3);
2192 }
2193
2194 #[tokio::test]
2195 async fn stream_pipeline_flushes_sink_on_source_error() {
2196 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2200 Ok(StreamPage {
2201 records: vec![json!({"id": 1}), json!({"id": 2})],
2202 bookmark: None,
2203 }),
2204 Err(FaucetError::Source("transient blip mid-stream".into())),
2205 ];
2206 let stream = futures::stream::iter(pages);
2207 let sink = FlushTrackingSink::new();
2208
2209 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2210
2211 assert!(matches!(result, Err(FaucetError::Source(_))));
2213 assert_eq!(sink.written().len(), 2);
2215 assert!(
2217 sink.flush_count() >= 1,
2218 "sink must be flushed on the error path so partial output is durable"
2219 );
2220 }
2221
2222 #[tokio::test]
2223 async fn stream_pipeline_flushes_sink_on_cancel() {
2224 use tokio_util::sync::CancellationToken;
2229
2230 let stream = Box::pin(async_stream::stream! {
2232 yield Ok(StreamPage {
2233 records: vec![json!({"id": 1}), json!({"id": 2})],
2234 bookmark: None,
2235 });
2236 futures::future::pending::<()>().await;
2237 });
2238 let sink = FlushTrackingSink::new();
2239
2240 let token = CancellationToken::new();
2241 let canceller = token.clone();
2242 tokio::spawn(async move {
2243 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
2244 canceller.cancel();
2245 });
2246
2247 let result = run_stream(stream, &sink, RunStreamOptions::new().with_cancel(token))
2248 .await
2249 .expect("a cooperative cancel returns Ok with the partial result");
2250
2251 assert_eq!(result.records_written, 2);
2254 assert_eq!(sink.written().len(), 2);
2255 assert!(
2256 sink.flush_count() >= 1,
2257 "sink must be flushed on the cancel path so partial output is durable"
2258 );
2259 }
2260
2261 #[tokio::test]
2262 async fn stream_pipeline_empty() {
2263 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
2264 let stream = futures::stream::iter(pages);
2265 let sink = MockSink::new();
2266
2267 let result = run_stream(stream, &sink, RunStreamOptions::new())
2268 .await
2269 .unwrap();
2270
2271 assert_eq!(result.records_written, 0);
2272 }
2273
2274 #[tokio::test]
2275 async fn stream_pipeline_skips_empty_pages() {
2276 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2277 Ok(StreamPage {
2278 records: vec![json!({"id": 1})],
2279 bookmark: None,
2280 }),
2281 Ok(StreamPage {
2282 records: vec![],
2283 bookmark: None,
2284 }),
2285 Ok(StreamPage {
2286 records: vec![json!({"id": 2})],
2287 bookmark: None,
2288 }),
2289 ];
2290 let stream = futures::stream::iter(pages);
2291 let sink = MockSink::new();
2292
2293 let result = run_stream(stream, &sink, RunStreamOptions::new())
2294 .await
2295 .unwrap();
2296
2297 assert_eq!(result.records_written, 2);
2298 }
2299
2300 #[tokio::test]
2301 async fn stream_pipeline_error_in_page_propagates() {
2302 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2303 Ok(StreamPage {
2304 records: vec![json!({"id": 1})],
2305 bookmark: None,
2306 }),
2307 Err(FaucetError::HttpStatus {
2308 status: 500,
2309 url: "https://example.com".into(),
2310 body: "Internal Server Error".into(),
2311 }),
2312 ];
2313 let stream = futures::stream::iter(pages);
2314 let sink = MockSink::new();
2315
2316 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2317 assert!(result.is_err());
2318 assert_eq!(sink.written().len(), 1);
2320 }
2321
2322 #[tokio::test]
2323 async fn stream_pipeline_sink_error_propagates() {
2324 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2325 records: vec![json!({"id": 1})],
2326 bookmark: None,
2327 })];
2328 let stream = futures::stream::iter(pages);
2329 let sink = FailingSink;
2330
2331 let result = run_stream(stream, &sink, RunStreamOptions::new()).await;
2332 assert!(result.is_err());
2333 }
2334
2335 #[tokio::test]
2336 async fn stream_pipeline_with_trait_object_sink() {
2337 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2338 records: vec![json!({"id": 1})],
2339 bookmark: None,
2340 })];
2341 let stream = futures::stream::iter(pages);
2342 let sink: Box<dyn Sink> = Box::new(MockSink::new());
2343
2344 let result = run_stream(stream, sink.as_ref(), RunStreamOptions::new())
2345 .await
2346 .unwrap();
2347 assert_eq!(result.records_written, 1);
2348 }
2349
2350 #[tokio::test]
2351 async fn stream_pipeline_persists_bookmark_when_page_carries_one() {
2352 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2353 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2354 Ok(StreamPage {
2355 records: vec![json!({"id": 1})],
2356 bookmark: None,
2357 }),
2358 Ok(StreamPage {
2359 records: vec![json!({"id": 2})],
2360 bookmark: Some(json!("checkpoint-final")),
2361 }),
2362 ];
2363 let stream = futures::stream::iter(pages);
2364 let sink = MockSink::new();
2365
2366 let result = run_stream(
2367 stream,
2368 &sink,
2369 RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
2370 )
2371 .await
2372 .unwrap();
2373
2374 assert_eq!(result.records_written, 2);
2375 assert_eq!(result.bookmark, Some(json!("checkpoint-final")));
2376 assert_eq!(
2377 store.get("k").await.unwrap(),
2378 Some(json!("checkpoint-final"))
2379 );
2380 }
2381
2382 #[tokio::test]
2383 async fn stream_pipeline_persists_per_page_bookmarks() {
2384 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2385 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
2386 Ok(StreamPage {
2387 records: vec![json!({"id": 1})],
2388 bookmark: Some(json!("tx-1")),
2389 }),
2390 Ok(StreamPage {
2391 records: vec![json!({"id": 2})],
2392 bookmark: Some(json!("tx-2")),
2393 }),
2394 ];
2395 let stream = futures::stream::iter(pages);
2396 let sink = MockSink::new();
2397
2398 run_stream(
2399 stream,
2400 &sink,
2401 RunStreamOptions::new().with_state(Arc::clone(&store), "k"),
2402 )
2403 .await
2404 .unwrap();
2405
2406 assert_eq!(store.get("k").await.unwrap(), Some(json!("tx-2")));
2408 }
2409
2410 use crate::state::{FileStateStore, MemoryStateStore, StateStore};
2413 use std::sync::Arc;
2414 use tempfile::TempDir;
2415
2416 struct StatefulSource {
2420 key: String,
2421 records: Vec<Value>,
2422 new_bookmark: Value,
2423 seen_bookmark: std::sync::Mutex<Option<Value>>,
2424 }
2425
2426 impl StatefulSource {
2427 fn new(key: &str, records: Vec<Value>, new_bookmark: Value) -> Self {
2428 Self {
2429 key: key.into(),
2430 records,
2431 new_bookmark,
2432 seen_bookmark: std::sync::Mutex::new(None),
2433 }
2434 }
2435 fn observed_start(&self) -> Option<Value> {
2436 self.seen_bookmark.lock().unwrap().clone()
2437 }
2438 }
2439
2440 #[async_trait]
2441 impl Source for StatefulSource {
2442 async fn fetch_with_context(
2443 &self,
2444 _ctx: &std::collections::HashMap<String, Value>,
2445 ) -> Result<Vec<Value>, FaucetError> {
2446 Ok(self.records.clone())
2447 }
2448 async fn fetch_with_context_incremental(
2449 &self,
2450 _ctx: &std::collections::HashMap<String, Value>,
2451 ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
2452 Ok((self.records.clone(), Some(self.new_bookmark.clone())))
2453 }
2454 fn state_key(&self) -> Option<String> {
2455 Some(self.key.clone())
2456 }
2457 async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
2458 *self.seen_bookmark.lock().unwrap() = Some(bookmark);
2459 Ok(())
2460 }
2461 }
2462
2463 #[tokio::test]
2464 async fn pipeline_with_state_store_persists_bookmark_after_sink() {
2465 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2466 let source = StatefulSource::new(
2467 "github_issues",
2468 vec![json!({"id": 1, "ts": "2026-05-01"})],
2469 json!("2026-05-01"),
2470 );
2471 let sink = MockSink::new();
2472 let result = Pipeline::new(&source, &sink)
2473 .with_state_store(Arc::clone(&store))
2474 .run()
2475 .await
2476 .unwrap();
2477
2478 assert_eq!(result.records_written, 1);
2479 assert_eq!(result.bookmark, Some(json!("2026-05-01")));
2480 let stored = store.get("github_issues").await.unwrap();
2482 assert_eq!(stored, Some(json!("2026-05-01")));
2483 }
2484
2485 #[tokio::test]
2486 async fn pipeline_with_state_store_resumes_from_stored_bookmark() {
2487 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2488 store
2489 .put("github_issues", &json!("2026-04-30"))
2490 .await
2491 .unwrap();
2492
2493 let source =
2494 StatefulSource::new("github_issues", vec![json!({"id": 2})], json!("2026-05-01"));
2495 let sink = MockSink::new();
2496 Pipeline::new(&source, &sink)
2497 .with_state_store(Arc::clone(&store))
2498 .run()
2499 .await
2500 .unwrap();
2501
2502 assert_eq!(source.observed_start(), Some(json!("2026-04-30")));
2504 assert_eq!(
2506 store.get("github_issues").await.unwrap(),
2507 Some(json!("2026-05-01"))
2508 );
2509 }
2510
2511 #[tokio::test]
2512 async fn pipeline_with_state_store_does_not_persist_when_sink_fails() {
2513 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2514 let source = StatefulSource::new("k", vec![json!({"id": 1})], json!("2026-05-01"));
2515 let sink = FailingSink;
2516
2517 let result = Pipeline::new(&source, &sink)
2518 .with_state_store(Arc::clone(&store))
2519 .run()
2520 .await;
2521 assert!(result.is_err());
2522 assert!(store.get("k").await.unwrap().is_none());
2523 }
2524
2525 #[tokio::test]
2526 async fn pipeline_with_state_store_no_state_key_means_no_persist() {
2527 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2528 let source = IncrementalSource {
2529 records: vec![json!({"id": 1})],
2530 bookmark: json!("ignored"),
2531 };
2532 let sink = MockSink::new();
2533 Pipeline::new(&source, &sink)
2534 .with_state_store(Arc::clone(&store))
2535 .run()
2536 .await
2537 .unwrap();
2538 assert!(store.get("anything").await.unwrap().is_none());
2541 }
2542
2543 #[tokio::test]
2544 async fn pipeline_with_state_store_skips_persist_when_bookmark_is_none() {
2545 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
2546 struct NoBookmarkSource;
2547 #[async_trait]
2548 impl Source for NoBookmarkSource {
2549 async fn fetch_with_context(
2550 &self,
2551 _ctx: &std::collections::HashMap<String, Value>,
2552 ) -> Result<Vec<Value>, FaucetError> {
2553 Ok(vec![json!({"id": 1})])
2554 }
2555 fn state_key(&self) -> Option<String> {
2556 Some("k".into())
2557 }
2558 }
2559 let source = NoBookmarkSource;
2560 let sink = MockSink::new();
2561 Pipeline::new(&source, &sink)
2562 .with_state_store(Arc::clone(&store))
2563 .run()
2564 .await
2565 .unwrap();
2566 assert!(store.get("k").await.unwrap().is_none());
2567 }
2568
2569 struct PagedSource;
2574
2575 #[async_trait]
2576 impl Source for PagedSource {
2577 async fn fetch_with_context(
2578 &self,
2579 _ctx: &std::collections::HashMap<String, Value>,
2580 ) -> Result<Vec<Value>, FaucetError> {
2581 unreachable!("Pipeline::run must drive stream_pages, not fetch_with_context");
2583 }
2584 fn stream_pages<'a>(
2585 &'a self,
2586 _ctx: &'a std::collections::HashMap<String, Value>,
2587 _batch_size: usize,
2588 ) -> std::pin::Pin<
2589 Box<dyn futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>,
2590 > {
2591 Box::pin(async_stream::try_stream! {
2592 yield StreamPage { records: vec![json!({"i": 1})], bookmark: None };
2593 yield StreamPage { records: vec![json!({"i": 2})], bookmark: None };
2594 yield StreamPage { records: vec![json!({"i": 3})], bookmark: Some(json!("final")) };
2595 })
2596 }
2597 }
2598
2599 struct CountingSink {
2601 calls: std::sync::Mutex<Vec<usize>>,
2602 }
2603
2604 impl CountingSink {
2605 fn new() -> Self {
2606 Self {
2607 calls: std::sync::Mutex::new(Vec::new()),
2608 }
2609 }
2610 fn call_count(&self) -> usize {
2611 self.calls.lock().unwrap().len()
2612 }
2613 }
2614
2615 #[async_trait]
2616 impl Sink for CountingSink {
2617 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
2618 self.calls.lock().unwrap().push(records.len());
2619 Ok(records.len())
2620 }
2621 }
2622
2623 #[tokio::test]
2624 async fn pipeline_run_drives_stream_pages() {
2625 let source = PagedSource;
2626 let sink = CountingSink::new();
2627
2628 let result = Pipeline::new(&source, &sink).run().await.unwrap();
2629
2630 assert_eq!(sink.call_count(), 3);
2632 assert_eq!(result.records_written, 3);
2633 assert_eq!(result.bookmark, Some(json!("final")));
2634 }
2635
2636 #[tokio::test]
2637 async fn pipeline_with_file_state_store_round_trips_across_runs() {
2638 let dir = TempDir::new().unwrap();
2639 let store: Arc<dyn StateStore> = Arc::new(FileStateStore::new(dir.path()));
2640
2641 let s1 = StatefulSource::new("k", vec![json!({"i": 1})], json!("v1"));
2643 let sink1 = MockSink::new();
2644 Pipeline::new(&s1, &sink1)
2645 .with_state_store(Arc::clone(&store))
2646 .run()
2647 .await
2648 .unwrap();
2649 assert_eq!(s1.observed_start(), None);
2650 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
2651
2652 let s2 = StatefulSource::new("k", vec![json!({"i": 2})], json!("v2"));
2654 let sink2 = MockSink::new();
2655 Pipeline::new(&s2, &sink2)
2656 .with_state_store(Arc::clone(&store))
2657 .run()
2658 .await
2659 .unwrap();
2660 assert_eq!(s2.observed_start(), Some(json!("v1")));
2661 assert_eq!(store.get("k").await.unwrap(), Some(json!("v2")));
2662 }
2663
2664 #[tokio::test]
2665 #[allow(clippy::await_holding_lock)]
2666 async fn pipeline_run_increments_runs_total() {
2667 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2668 use metrics_util::debugging::DebugValue;
2669 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2670 let snap = snapshotter();
2671
2672 let source = MockSource(vec![json!({"i": 1})]);
2673 let sink = MockSink::new();
2674 let _ = Pipeline::new(&source, &sink)
2675 .with_name("test-pipeline")
2676 .with_row("rowA")
2677 .run()
2678 .await
2679 .unwrap();
2680
2681 let snapshot = snap.snapshot();
2682 let found = snapshot.into_vec().into_iter().any(
2683 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
2684 key.key().name() == "faucet_pipeline_runs_total"
2685 && key.key().labels().any(|l: &metrics::Label| {
2686 l.key() == "pipeline" && l.value() == "test-pipeline"
2687 })
2688 && key
2689 .key()
2690 .labels()
2691 .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowA")
2692 && key
2693 .key()
2694 .labels()
2695 .any(|l: &metrics::Label| l.key() == "status" && l.value() == "ok")
2696 && matches!(v, DebugValue::Counter(c) if c >= 1)
2697 },
2698 );
2699 assert!(
2700 found,
2701 "expected faucet_pipeline_runs_total{{pipeline=test-pipeline, row=rowA, status=ok}}"
2702 );
2703 }
2704
2705 #[tokio::test]
2706 #[allow(clippy::await_holding_lock)]
2707 async fn pipeline_failure_attaches_kind_label_to_runs_total() {
2708 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2709 use metrics_util::debugging::DebugValue;
2710 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2711 let snap = snapshotter();
2712
2713 let source = FailingSource;
2714 let sink = MockSink::new();
2715 let _ = Pipeline::new(&source, &sink)
2716 .with_name("err-pipeline")
2717 .with_row("rowE")
2718 .run()
2719 .await;
2720
2721 let snapshot = snap.snapshot();
2722 let found = snapshot.into_vec().into_iter().any(
2723 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
2724 key.key().name() == "faucet_pipeline_runs_total"
2725 && key.key().labels().any(|l: &metrics::Label| {
2726 l.key() == "pipeline" && l.value() == "err-pipeline"
2727 })
2728 && key
2729 .key()
2730 .labels()
2731 .any(|l: &metrics::Label| l.key() == "status" && l.value() == "err")
2732 && key
2733 .key()
2734 .labels()
2735 .any(|l: &metrics::Label| l.key() == "kind" && l.value() == "Auth")
2736 && matches!(v, DebugValue::Counter(c) if c >= 1)
2737 },
2738 );
2739 assert!(
2740 found,
2741 "expected faucet_pipeline_runs_total{{status=err, kind=Auth}} for failing source"
2742 );
2743 }
2744
2745 #[tokio::test]
2746 #[allow(clippy::await_holding_lock)]
2747 async fn pipeline_run_emits_start_time_gauge() {
2748 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2749 use metrics_util::debugging::DebugValue;
2750 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2751 let snap = snapshotter();
2752
2753 let source = MockSource(vec![json!({"i": 1})]);
2754 let sink = MockSink::new();
2755 let before = std::time::SystemTime::now()
2756 .duration_since(std::time::UNIX_EPOCH)
2757 .map(|d| d.as_secs_f64())
2758 .unwrap_or(0.0);
2759 let _ = Pipeline::new(&source, &sink)
2760 .with_name("start-time-pipeline")
2761 .with_row("rowS")
2762 .run()
2763 .await
2764 .unwrap();
2765
2766 let snapshot = snap.snapshot();
2767 let found = snapshot.into_vec().into_iter().any(
2768 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
2769 if key.key().name() != "faucet_pipeline_start_time_unix_seconds" {
2770 return false;
2771 }
2772 let labels_match = key.key().labels().any(|l: &metrics::Label| {
2773 l.key() == "pipeline" && l.value() == "start-time-pipeline"
2774 }) && key
2775 .key()
2776 .labels()
2777 .any(|l: &metrics::Label| l.key() == "row" && l.value() == "rowS");
2778 if !labels_match {
2779 return false;
2780 }
2781 matches!(v, DebugValue::Gauge(g) if g.into_inner() >= before)
2782 },
2783 );
2784 assert!(
2785 found,
2786 "expected faucet_pipeline_start_time_unix_seconds gauge >= test-start timestamp"
2787 );
2788 }
2789
2790 #[tokio::test]
2791 #[allow(clippy::await_holding_lock)]
2792 async fn register_build_info_sets_version_gauge() {
2793 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
2794 use metrics_util::debugging::DebugValue;
2795 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
2796 let snap = snapshotter();
2797
2798 crate::observability::register_build_info();
2799
2800 let snapshot = snap.snapshot();
2801 let found = snapshot.into_vec().into_iter().any(
2802 |(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
2803 key.key().name() == "faucet_build_info"
2804 && key.key().labels().any(|l: &metrics::Label| {
2805 l.key() == "version" && l.value() == env!("CARGO_PKG_VERSION")
2806 })
2807 && matches!(v, DebugValue::Gauge(g) if (g.into_inner() - 1.0).abs() < f64::EPSILON)
2808 },
2809 );
2810 assert!(
2811 found,
2812 "expected faucet_build_info{{version=CARGO_PKG_VERSION}} = 1.0 after register_build_info()"
2813 );
2814 }
2815
2816 use crate::dlq::{DlqConfig, OnBatchError};
2819
2820 struct PartialSink {
2824 fail_indices: std::sync::Mutex<Vec<usize>>,
2825 committed: std::sync::Mutex<Vec<Value>>,
2826 }
2827
2828 impl PartialSink {
2829 fn new(fail_indices: Vec<usize>) -> Self {
2830 Self {
2831 fail_indices: std::sync::Mutex::new(fail_indices),
2832 committed: std::sync::Mutex::new(Vec::new()),
2833 }
2834 }
2835 }
2836
2837 #[async_trait]
2838 impl Sink for PartialSink {
2839 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
2840 unreachable!("PartialSink only overrides write_batch_partial");
2841 }
2842 async fn write_batch_partial(
2843 &self,
2844 records: &[Value],
2845 ) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
2846 let fails: std::collections::HashSet<usize> =
2847 self.fail_indices.lock().unwrap().iter().copied().collect();
2848 let mut outcomes = Vec::with_capacity(records.len());
2849 for (i, rec) in records.iter().enumerate() {
2850 if fails.contains(&i) {
2851 outcomes.push(Err(FaucetError::Sink(format!("row {i} rejected"))));
2852 } else {
2853 self.committed.lock().unwrap().push(rec.clone());
2854 outcomes.push(Ok(()));
2855 }
2856 }
2857 Ok(outcomes)
2858 }
2859 }
2860
2861 #[tokio::test]
2862 async fn dlq_routes_only_failed_rows_for_partial_success_sink() {
2863 let main = PartialSink::new(vec![1, 3]); let dlq = std::sync::Arc::new(MockSink::new());
2865 let dlq_cfg = DlqConfig::new(dlq.clone());
2866
2867 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2868 records: (0..4).map(|i| json!({"i": i})).collect(),
2869 bookmark: None,
2870 })];
2871 let stream = futures::stream::iter(pages);
2872 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg))
2873 .await
2874 .unwrap();
2875
2876 assert_eq!(result.records_written, 2); assert_eq!(main.committed.lock().unwrap().len(), 2);
2878 let envelopes = dlq.0.lock().unwrap();
2879 assert_eq!(envelopes.len(), 2);
2880 assert_eq!(envelopes[0]["payload"]["i"], 1);
2881 assert_eq!(envelopes[0]["record_index"], 1);
2882 assert_eq!(envelopes[1]["payload"]["i"], 3);
2883 assert_eq!(envelopes[1]["record_index"], 3);
2884 let stats = result.dlq.unwrap();
2885 assert_eq!(stats.records_dlq, 2);
2886 assert_eq!(stats.pages_with_failures, 1);
2887 }
2888
2889 #[cfg(feature = "masking")]
2890 #[tokio::test]
2891 async fn masking_runs_before_the_sink() {
2892 use crate::masking::{CompiledMasking, MaskingSpec};
2893 let sink = MockSink::new();
2894 let spec: MaskingSpec = serde_json::from_value(json!({
2895 "rules": [{ "match": { "value_detector": "email" },
2896 "action": { "type": "redact" } }]
2897 }))
2898 .unwrap();
2899 let m = std::sync::Arc::new(CompiledMasking::compile(&spec).unwrap());
2900 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2901 records: vec![json!({"email": "a@b.com", "name": "Al"})],
2902 bookmark: None,
2903 })];
2904 run_stream(
2905 futures::stream::iter(pages),
2906 &sink,
2907 RunStreamOptions::new().with_masking(m),
2908 )
2909 .await
2910 .unwrap();
2911 assert_eq!(sink.written()[0], json!({"email": "***", "name": "Al"}));
2912 }
2913
2914 #[cfg(feature = "masking")]
2915 #[tokio::test]
2916 async fn masking_applies_before_the_dlq_envelope() {
2917 use crate::masking::{CompiledMasking, MaskingSpec};
2921 let main = PartialSink::new(vec![0]); let dlq = std::sync::Arc::new(MockSink::new());
2923 let spec: MaskingSpec = serde_json::from_value(json!({
2924 "rules": [{ "match": { "value_detector": "email" },
2925 "action": { "type": "redact" } }]
2926 }))
2927 .unwrap();
2928 let m = std::sync::Arc::new(CompiledMasking::compile(&spec).unwrap());
2929 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2930 records: vec![
2931 json!({"email": "secret@x.com"}),
2932 json!({"email": "ok@y.com"}),
2933 ],
2934 bookmark: None,
2935 })];
2936 let opts = RunStreamOptions::new()
2937 .with_masking(m)
2938 .with_dlq(DlqConfig::new(dlq.clone()));
2939 run_stream(futures::stream::iter(pages), &main, opts)
2940 .await
2941 .unwrap();
2942 let env = dlq.0.lock().unwrap();
2944 assert_eq!(env.len(), 1);
2945 assert_eq!(
2946 env[0]["payload"]["email"], "***",
2947 "the DLQ payload must be masked, not raw PII"
2948 );
2949 assert_eq!(main.committed.lock().unwrap()[0]["email"], "***");
2950 }
2951
2952 #[tokio::test]
2953 async fn dlq_propagate_policy_bubbles_outer_err() {
2954 let main = FailingSink;
2955 let dlq = std::sync::Arc::new(MockSink::new());
2956 let mut dlq_cfg = DlqConfig::new(dlq.clone());
2957 dlq_cfg.on_batch_error = OnBatchError::Propagate;
2958
2959 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2960 records: vec![json!({"i": 0}), json!({"i": 1})],
2961 bookmark: Some(json!("v1")),
2962 })];
2963 let stream = futures::stream::iter(pages);
2964 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2965 let result = run_stream(
2966 stream,
2967 &main,
2968 RunStreamOptions::new()
2969 .with_dlq(dlq_cfg)
2970 .with_state(std::sync::Arc::clone(&store), "k"),
2971 )
2972 .await;
2973 assert!(matches!(result, Err(FaucetError::Sink(_))));
2974 assert!(dlq.0.lock().unwrap().is_empty());
2975 assert!(store.get("k").await.unwrap().is_none());
2977 }
2978
2979 #[tokio::test]
2980 async fn dlq_dlq_all_policy_routes_every_row_on_outer_err() {
2981 let main = FailingSink;
2982 let dlq = std::sync::Arc::new(MockSink::new());
2983 let mut dlq_cfg = DlqConfig::new(dlq.clone());
2984 dlq_cfg.on_batch_error = OnBatchError::DlqAll;
2985
2986 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
2987 records: vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})],
2988 bookmark: Some(json!("v1")),
2989 })];
2990 let stream = futures::stream::iter(pages);
2991 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
2992 let result = run_stream(
2993 stream,
2994 &main,
2995 RunStreamOptions::new()
2996 .with_dlq(dlq_cfg)
2997 .with_state(std::sync::Arc::clone(&store), "k"),
2998 )
2999 .await
3000 .unwrap();
3001 assert_eq!(result.records_written, 0);
3002 {
3003 let envelopes = dlq.0.lock().unwrap();
3004 assert_eq!(envelopes.len(), 3);
3005 for env in envelopes.iter() {
3007 let msg = env["error"]["message"].as_str().unwrap();
3008 assert!(msg.contains("write failed"), "got: {msg}");
3009 }
3010 }
3011 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3012 assert_eq!(result.dlq.unwrap().records_dlq, 3);
3013 }
3014
3015 #[tokio::test]
3016 async fn dlq_per_page_budget_exceeded_aborts() {
3017 let main = PartialSink::new(vec![0, 1, 2]);
3018 let dlq = std::sync::Arc::new(MockSink::new());
3019 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3020 dlq_cfg.max_failures_per_page = Some(2);
3021
3022 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3023 records: (0..3).map(|i| json!({"i": i})).collect(),
3024 bookmark: None,
3025 })];
3026 let stream = futures::stream::iter(pages);
3027 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3028 assert!(
3029 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
3030 "got: {result:?}"
3031 );
3032 }
3033
3034 #[tokio::test]
3035 async fn dlq_total_budget_exceeded_aborts_on_later_page() {
3036 let pages: Vec<Result<StreamPage, FaucetError>> = vec![
3037 Ok(StreamPage {
3038 records: (0..3).map(|i| json!({"i": i})).collect(),
3039 bookmark: None,
3040 }),
3041 Ok(StreamPage {
3042 records: (3..6).map(|i| json!({"i": i})).collect(),
3043 bookmark: None,
3044 }),
3045 ];
3046 let main = PartialSink::new(vec![0, 1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3049 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3050 dlq_cfg.max_failures_total = Some(4);
3051
3052 let stream = futures::stream::iter(pages);
3053 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3054 assert!(
3055 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
3056 "got: {result:?}"
3057 );
3058 }
3059
3060 #[tokio::test]
3061 async fn dlq_per_page_budget_exceeded_commits_page_before_aborting() {
3062 let main = PartialSink::new(vec![1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3069 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3070 dlq_cfg.max_failures_per_page = Some(1); let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3073 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3074 records: (0..3).map(|i| json!({ "i": i })).collect(),
3075 bookmark: Some(json!("v1")),
3076 })];
3077 let stream = futures::stream::iter(pages);
3078 let result = run_stream(
3079 stream,
3080 &main,
3081 RunStreamOptions::new()
3082 .with_dlq(dlq_cfg)
3083 .with_state(std::sync::Arc::clone(&store), "k"),
3084 )
3085 .await;
3086
3087 assert!(
3089 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("per-page budget exceeded")),
3090 "got: {result:?}"
3091 );
3092 assert_eq!(main.committed.lock().unwrap().len(), 1);
3094 assert_eq!(dlq.0.lock().unwrap().len(), 2);
3096 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3098 }
3099
3100 #[tokio::test]
3101 async fn dlq_total_budget_exceeded_commits_tripping_page_before_aborting() {
3102 let main = PartialSink::new(vec![1, 2]); let dlq = std::sync::Arc::new(MockSink::new());
3107 let mut dlq_cfg = DlqConfig::new(dlq.clone());
3108 dlq_cfg.max_failures_total = Some(1); let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3111 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3112 records: (0..3).map(|i| json!({ "i": i })).collect(),
3113 bookmark: Some(json!("v1")),
3114 })];
3115 let stream = futures::stream::iter(pages);
3116 let result = run_stream(
3117 stream,
3118 &main,
3119 RunStreamOptions::new()
3120 .with_dlq(dlq_cfg)
3121 .with_state(std::sync::Arc::clone(&store), "k"),
3122 )
3123 .await;
3124
3125 assert!(
3126 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("total budget exceeded")),
3127 "got: {result:?}"
3128 );
3129 assert_eq!(main.committed.lock().unwrap().len(), 1);
3130 assert_eq!(dlq.0.lock().unwrap().len(), 2);
3131 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3132 }
3133
3134 struct FailingDlqSink;
3137 #[async_trait]
3138 impl Sink for FailingDlqSink {
3139 async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
3140 Err(FaucetError::Sink("dlq disk full".into()))
3141 }
3142 }
3143
3144 struct FailingFlushDlqSink {
3148 written: std::sync::Mutex<Vec<Value>>,
3149 }
3150 impl FailingFlushDlqSink {
3151 fn new() -> Self {
3152 Self {
3153 written: std::sync::Mutex::new(Vec::new()),
3154 }
3155 }
3156 }
3157 #[async_trait]
3158 impl Sink for FailingFlushDlqSink {
3159 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3160 self.written.lock().unwrap().extend(records.iter().cloned());
3161 Ok(records.len())
3162 }
3163 async fn flush(&self) -> Result<(), FaucetError> {
3164 Err(FaucetError::Sink("dlq flush failed".into()))
3165 }
3166 }
3167
3168 #[tokio::test]
3169 async fn dlq_sink_failure_is_fatal_no_recursion() {
3170 let main = PartialSink::new(vec![0]);
3171 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingDlqSink);
3172 let dlq_cfg = DlqConfig::new(dlq);
3173
3174 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3175 records: vec![json!({"i": 0}), json!({"i": 1})],
3176 bookmark: Some(json!("v1")),
3177 })];
3178 let stream = futures::stream::iter(pages);
3179 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3180 let result = run_stream(
3181 stream,
3182 &main,
3183 RunStreamOptions::new()
3184 .with_dlq(dlq_cfg)
3185 .with_state(std::sync::Arc::clone(&store), "k"),
3186 )
3187 .await;
3188 assert!(
3189 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink write failed")),
3190 "got: {result:?}"
3191 );
3192 assert!(store.get("k").await.unwrap().is_none());
3193 }
3194
3195 #[tokio::test]
3196 async fn dlq_bookmark_advances_only_after_both_flushes() {
3197 let main = PartialSink::new(vec![1]); let dlq = std::sync::Arc::new(MockSink::new());
3199 let dlq_cfg = DlqConfig::new(dlq.clone());
3200
3201 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3202 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3203 records: vec![json!({"i": 0}), json!({"i": 1})],
3204 bookmark: Some(json!("v1")),
3205 })];
3206 let stream = futures::stream::iter(pages);
3207 run_stream(
3208 stream,
3209 &main,
3210 RunStreamOptions::new()
3211 .with_dlq(dlq_cfg)
3212 .with_state(std::sync::Arc::clone(&store), "k"),
3213 )
3214 .await
3215 .unwrap();
3216 assert_eq!(store.get("k").await.unwrap(), Some(json!("v1")));
3217 assert_eq!(dlq.0.lock().unwrap().len(), 1);
3218 assert_eq!(main.committed.lock().unwrap().len(), 1);
3219 }
3220
3221 #[tokio::test]
3222 async fn dlq_disabled_pipeline_behaves_identically_to_today() {
3223 let main = MockSink::new();
3225 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3226 records: vec![json!({"i": 0}), json!({"i": 1})],
3227 bookmark: None,
3228 })];
3229 let stream = futures::stream::iter(pages);
3230 let result = run_stream(stream, &main, RunStreamOptions::new())
3231 .await
3232 .unwrap();
3233 assert_eq!(result.records_written, 2);
3234 assert!(result.dlq.is_none());
3235 }
3236
3237 #[tokio::test]
3238 async fn dlq_per_page_flush_failure_is_fatal_and_blocks_bookmark() {
3239 let main = PartialSink::new(vec![1]);
3244 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
3245 let dlq_cfg = DlqConfig::new(dlq);
3246
3247 let store: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
3248 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3249 records: vec![json!({"i": 0}), json!({"i": 1})],
3250 bookmark: Some(json!("v1")),
3251 })];
3252 let stream = futures::stream::iter(pages);
3253 let result = run_stream(
3254 stream,
3255 &main,
3256 RunStreamOptions::new()
3257 .with_dlq(dlq_cfg)
3258 .with_state(std::sync::Arc::clone(&store), "k"),
3259 )
3260 .await;
3261 assert!(
3262 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
3263 "got: {result:?}"
3264 );
3265 assert!(store.get("k").await.unwrap().is_none());
3266 }
3267
3268 #[tokio::test]
3269 async fn dlq_end_of_stream_flush_failure_is_fatal() {
3270 let main = PartialSink::new(vec![1]);
3275 let dlq: std::sync::Arc<dyn Sink> = std::sync::Arc::new(FailingFlushDlqSink::new());
3276 let dlq_cfg = DlqConfig::new(dlq);
3277
3278 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3279 records: vec![json!({"i": 0}), json!({"i": 1})],
3280 bookmark: None,
3281 })];
3282 let stream = futures::stream::iter(pages);
3283 let result = run_stream(stream, &main, RunStreamOptions::new().with_dlq(dlq_cfg)).await;
3284 assert!(
3285 matches!(&result, Err(FaucetError::Sink(m)) if m.contains("DLQ sink flush failed")),
3286 "got: {result:?}"
3287 );
3288 }
3289
3290 #[tokio::test]
3291 #[allow(clippy::await_holding_lock)]
3292 async fn dlq_emits_records_total_and_pages_total() {
3293 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3294 use metrics_util::debugging::DebugValue;
3295
3296 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3297 let snap = snapshotter();
3298
3299 let source = MockSource(vec![json!({"i": 0}), json!({"i": 1})]);
3300 let main = PartialSink::new(vec![1]);
3301 let dlq = std::sync::Arc::new(MockSink::new());
3302 let _ = Pipeline::new(&source, &main)
3303 .with_name("p_dlq_metrics")
3304 .with_row("r1")
3305 .with_dlq(DlqConfig::new(dlq.clone()))
3306 .run()
3307 .await
3308 .unwrap();
3309
3310 let snapshot = snap.snapshot();
3311 let mut saw_records = false;
3312 let mut saw_pages = false;
3313 for (k, _u, _d, v) in snapshot.into_vec() {
3314 let key = k.key();
3315 let labels = key.labels().collect::<Vec<_>>();
3316 let has = |k: &str, v: &str| labels.iter().any(|l| l.key() == k && l.value() == v);
3317 if key.name() == "faucet_sink_dlq_records_total"
3318 && has("pipeline", "p_dlq_metrics")
3319 && has("row", "r1")
3320 && matches!(v, DebugValue::Counter(c) if c >= 1)
3321 {
3322 saw_records = true;
3323 }
3324 if key.name() == "faucet_sink_dlq_pages_total"
3325 && has("pipeline", "p_dlq_metrics")
3326 && matches!(v, DebugValue::Counter(c) if c >= 1)
3327 {
3328 saw_pages = true;
3329 }
3330 }
3331 assert!(saw_records, "faucet_sink_dlq_records_total not emitted");
3332 assert!(saw_pages, "faucet_sink_dlq_pages_total not emitted");
3333 }
3334
3335 #[tokio::test]
3336 #[allow(clippy::await_holding_lock)]
3337 async fn dlq_budget_exceeded_emits_counter() {
3338 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3339 use metrics_util::debugging::DebugValue;
3340
3341 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3342 let snap = snapshotter();
3343
3344 let source = MockSource((0..3).map(|i| json!({"i": i})).collect());
3345 let main = PartialSink::new(vec![0, 1, 2]);
3346 let dlq = std::sync::Arc::new(MockSink::new());
3347 let mut cfg = DlqConfig::new(dlq);
3348 cfg.max_failures_per_page = Some(1);
3349 let _ = Pipeline::new(&source, &main)
3350 .with_name("p_budget")
3351 .with_dlq(cfg)
3352 .run()
3353 .await;
3354
3355 let snapshot = snap.snapshot();
3356 let saw = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
3357 k.key().name() == "faucet_sink_dlq_budget_exceeded_total"
3358 && k.key()
3359 .labels()
3360 .any(|l| l.key() == "scope" && l.value() == "per_page")
3361 && matches!(v, DebugValue::Counter(c) if c >= 1)
3362 });
3363 assert!(saw, "faucet_sink_dlq_budget_exceeded_total not emitted");
3364 }
3365
3366 #[tokio::test]
3367 async fn pipeline_run_with_dlq_routes_partial_failures_end_to_end() {
3368 let source = MockSource(vec![json!({"i": 0}), json!({"i": 1}), json!({"i": 2})]);
3370 let main = PartialSink::new(vec![1]);
3371 let dlq = std::sync::Arc::new(MockSink::new());
3372
3373 let result = Pipeline::new(&source, &main)
3374 .with_dlq(DlqConfig::new(dlq.clone()))
3375 .run()
3376 .await
3377 .unwrap();
3378
3379 assert_eq!(result.records_written, 2);
3380 let stats = result.dlq.unwrap();
3381 assert_eq!(stats.records_dlq, 1);
3382 {
3383 let dlq_records = dlq.0.lock().unwrap();
3384 assert_eq!(dlq_records.len(), 1);
3385 }
3386 }
3387
3388 #[cfg(feature = "quality")]
3391 #[tokio::test]
3392 async fn quality_quarantines_to_dlq_and_writes_survivors() {
3393 use crate::dlq::DlqConfig;
3394 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3395
3396 let main = Arc::new(MockSink::new());
3397 let dlq_sink = Arc::new(MockSink::new());
3398 let spec = QualitySpec {
3399 record: vec![RecordCheck::NotNull {
3400 field: "id".into(),
3401 treat_missing_as_null: true,
3402 on_failure: OnFailure::Quarantine,
3403 }],
3404 batch: vec![],
3405 };
3406 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3407 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3408 records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
3409 bookmark: None,
3410 })];
3411 let opts = RunStreamOptions::new()
3412 .with_dlq(DlqConfig::new(dlq_sink.clone()))
3413 .with_quality(quality);
3414 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
3415 .await
3416 .unwrap();
3417
3418 assert_eq!(result.records_written, 2); assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
3420 let dlq = dlq_sink.written();
3422 assert_eq!(dlq.len(), 1);
3423 assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
3424 assert_eq!(result.dlq.unwrap().records_dlq, 1);
3425 }
3426
3427 #[cfg(feature = "quality")]
3428 #[tokio::test]
3429 #[allow(clippy::await_holding_lock)]
3430 async fn quality_only_page_emits_quality_reason() {
3431 use crate::dlq::DlqConfig;
3435 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3436 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3437 use metrics_util::debugging::DebugValue;
3438
3439 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3440 let snap = snapshotter();
3441
3442 let main = Arc::new(MockSink::new());
3445 let dlq_sink = Arc::new(MockSink::new());
3446 let spec = QualitySpec {
3447 record: vec![RecordCheck::NotNull {
3448 field: "id".into(),
3449 treat_missing_as_null: true,
3450 on_failure: OnFailure::Quarantine,
3451 }],
3452 batch: vec![],
3453 };
3454 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3455 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3456 records: vec![json!({"id": 1}), json!({"id": null}), json!({"id": 3})],
3457 bookmark: None,
3458 })];
3459 let opts = RunStreamOptions::new()
3460 .with_name("p_quality_reason")
3461 .with_dlq(DlqConfig::new(dlq_sink.clone()))
3462 .with_quality(quality);
3463 let _ = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
3464 .await
3465 .unwrap();
3466
3467 let snapshot = snap.snapshot();
3468 let saw_quality_reason = snapshot.into_vec().into_iter().any(|(k, _, _, v)| {
3469 k.key().name() == "faucet_sink_dlq_pages_total"
3470 && k.key()
3471 .labels()
3472 .any(|l| l.key() == "pipeline" && l.value() == "p_quality_reason")
3473 && k.key()
3474 .labels()
3475 .any(|l| l.key() == "reason" && l.value() == "quality")
3476 && matches!(v, DebugValue::Counter(c) if c >= 1)
3477 });
3478 assert!(
3479 saw_quality_reason,
3480 "expected faucet_sink_dlq_pages_total with reason=\"quality\""
3481 );
3482 }
3483
3484 #[cfg(feature = "quality")]
3485 #[tokio::test]
3486 async fn quality_abort_fails_run() {
3487 use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
3488 let main = MockSink::new();
3489 let spec = QualitySpec {
3490 record: vec![],
3491 batch: vec![BatchCheck::RowCount {
3492 min: Some(5),
3493 max: None,
3494 on_failure: OnFailure::Abort,
3495 }],
3496 };
3497 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3498 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3499 records: vec![json!({"id": 1})],
3500 bookmark: None,
3501 })];
3502 let opts = RunStreamOptions::new().with_quality(quality);
3503 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
3504 assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
3505 }
3506
3507 #[cfg(feature = "quality")]
3508 #[tokio::test]
3509 async fn quality_quarantine_without_dlq_is_rejected() {
3510 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3511 let main = MockSink::new();
3512 let spec = QualitySpec {
3513 record: vec![RecordCheck::NotNull {
3514 field: "id".into(),
3515 treat_missing_as_null: true,
3516 on_failure: OnFailure::Quarantine,
3517 }],
3518 batch: vec![],
3519 };
3520 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
3521 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3522 records: vec![json!({"id": null})],
3523 bookmark: None,
3524 })];
3525 let opts = RunStreamOptions::new().with_quality(quality);
3527 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
3528 assert!(matches!(result, Err(FaucetError::Config(_))));
3529 }
3530
3531 #[cfg(feature = "contract")]
3532 fn compiled_contract(on_breach: &str) -> Arc<crate::contract::CompiledContract> {
3533 let spec: crate::contract::ContractSpec = serde_json::from_value(json!({
3534 "version": "1.0.0",
3535 "on_breach": on_breach,
3536 "fields": [{ "name": "id", "type": "integer" }]
3537 }))
3538 .unwrap();
3539 Arc::new(crate::contract::CompiledContract::compile(&spec).unwrap())
3540 }
3541
3542 #[cfg(feature = "contract")]
3543 #[tokio::test]
3544 async fn contract_quarantines_to_dlq_and_writes_survivors() {
3545 use crate::dlq::DlqConfig;
3546 let main = Arc::new(MockSink::new());
3547 let dlq_sink = Arc::new(MockSink::new());
3548 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3549 records: vec![json!({"id": 1}), json!({"id": "bad"}), json!({"id": 3})],
3550 bookmark: None,
3551 })];
3552 let opts = RunStreamOptions::new()
3553 .with_dlq(DlqConfig::new(dlq_sink.clone()))
3554 .with_contract(compiled_contract("quarantine"));
3555 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
3556 .await
3557 .unwrap();
3558
3559 assert_eq!(result.records_written, 2);
3560 assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": 3})]);
3561 let dlq = dlq_sink.written();
3562 assert_eq!(dlq.len(), 1);
3563 assert_eq!(dlq[0]["error"]["kind"], "ContractViolation");
3564 assert_eq!(dlq[0]["payload"], json!({"id": "bad"}));
3565 assert_eq!(dlq[0]["record_index"], 1);
3567 assert_eq!(result.dlq.unwrap().records_dlq, 1);
3568 }
3569
3570 #[cfg(feature = "contract")]
3571 #[tokio::test]
3572 async fn contract_fail_aborts_run_and_writes_nothing() {
3573 let main = MockSink::new();
3574 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3575 records: vec![json!({"id": 1}), json!({"id": "bad"})],
3576 bookmark: None,
3577 })];
3578 let opts = RunStreamOptions::new().with_contract(compiled_contract("fail"));
3579 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
3580 match result {
3581 Err(FaucetError::ContractViolation { version, message }) => {
3582 assert_eq!(version, "1.0.0");
3583 assert!(message.contains("id"), "message: {message}");
3584 }
3585 other => panic!("expected ContractViolation, got {other:?}"),
3586 }
3587 assert!(
3588 main.written().is_empty(),
3589 "a contract fail must not commit any of the page's records"
3590 );
3591 }
3592
3593 #[cfg(feature = "contract")]
3594 #[tokio::test]
3595 async fn contract_warn_writes_everything() {
3596 let main = MockSink::new();
3597 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3598 records: vec![json!({"id": 1}), json!({"id": "bad"})],
3599 bookmark: None,
3600 })];
3601 let opts = RunStreamOptions::new().with_contract(compiled_contract("warn"));
3602 let result = run_stream(futures::stream::iter(pages), &main, opts)
3603 .await
3604 .unwrap();
3605 assert_eq!(result.records_written, 2);
3606 assert_eq!(main.written(), vec![json!({"id": 1}), json!({"id": "bad"})]);
3607 }
3608
3609 #[cfg(feature = "contract")]
3610 #[tokio::test]
3611 async fn contract_quarantine_without_dlq_is_rejected() {
3612 let main = MockSink::new();
3613 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3614 records: vec![json!({"id": 1})],
3615 bookmark: None,
3616 })];
3617 let opts = RunStreamOptions::new().with_contract(compiled_contract("quarantine"));
3619 let result = run_stream(futures::stream::iter(pages), &main, opts).await;
3620 assert!(matches!(result, Err(FaucetError::Config(_))));
3621 }
3622
3623 #[cfg(all(feature = "contract", feature = "quality"))]
3624 #[tokio::test]
3625 async fn contract_runs_after_quality_and_shares_dlq() {
3626 use crate::dlq::DlqConfig;
3630 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
3631 let main = Arc::new(MockSink::new());
3632 let dlq_sink = Arc::new(MockSink::new());
3633 let quality = Arc::new(
3634 CompiledQuality::compile(&QualitySpec {
3635 record: vec![RecordCheck::NotNull {
3636 field: "id".into(),
3637 treat_missing_as_null: true,
3638 on_failure: OnFailure::Quarantine,
3639 }],
3640 batch: vec![],
3641 })
3642 .unwrap(),
3643 );
3644 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3645 records: vec![json!({"id": null}), json!({"id": "bad"}), json!({"id": 3})],
3646 bookmark: None,
3647 })];
3648 let opts = RunStreamOptions::new()
3649 .with_dlq(DlqConfig::new(dlq_sink.clone()))
3650 .with_quality(quality)
3651 .with_contract(compiled_contract("quarantine"));
3652 let result = run_stream(futures::stream::iter(pages), main.as_ref(), opts)
3653 .await
3654 .unwrap();
3655
3656 assert_eq!(result.records_written, 1);
3657 assert_eq!(main.written(), vec![json!({"id": 3})]);
3658 let dlq = dlq_sink.written();
3659 assert_eq!(dlq.len(), 2);
3660 let kinds: Vec<&str> = dlq
3661 .iter()
3662 .map(|e| e["error"]["kind"].as_str().unwrap())
3663 .collect();
3664 assert!(kinds.contains(&"QualityFailure"), "kinds: {kinds:?}");
3665 assert!(kinds.contains(&"ContractViolation"), "kinds: {kinds:?}");
3666 assert_eq!(result.dlq.unwrap().records_dlq, 2);
3667 }
3668
3669 struct FlakySink {
3672 every: usize,
3673 calls: std::sync::Mutex<Vec<usize>>,
3674 }
3675 impl FlakySink {
3676 fn new(every: usize) -> Self {
3677 Self {
3678 every,
3679 calls: std::sync::Mutex::new(Vec::new()),
3680 }
3681 }
3682 fn call_sizes(&self) -> Vec<usize> {
3683 self.calls.lock().unwrap().clone()
3684 }
3685 }
3686 #[async_trait]
3687 impl Sink for FlakySink {
3688 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3689 Ok(records.len())
3690 }
3691 async fn write_batch_partial(
3692 &self,
3693 records: &[Value],
3694 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
3695 self.calls.lock().unwrap().push(records.len());
3696 Ok(records
3697 .iter()
3698 .enumerate()
3699 .map(|(i, _)| {
3700 if (i + 1) % self.every == 0 {
3701 Err(FaucetError::Sink("synthetic".into()))
3702 } else {
3703 Ok(())
3704 }
3705 })
3706 .collect())
3707 }
3708 }
3709
3710 #[tokio::test]
3711 async fn adaptive_shrinks_under_errors_on_dlq_path() {
3712 use crate::adaptive::AdaptiveBatchConfig;
3713 use crate::dlq::{DlqConfig, OnBatchError};
3714 let mk = || StreamPage {
3719 records: (0..400).map(|i| json!({"i": i})).collect(),
3720 bookmark: None,
3721 };
3722 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
3723 let sink = FlakySink::new(4); let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
3725 let dlq = DlqConfig {
3726 sink: dlq_sink,
3727 on_batch_error: OnBatchError::Propagate,
3728 max_failures_per_page: None,
3729 max_failures_total: None,
3730 include_original_payload: true,
3731 };
3732 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
3733 "enabled": true, "min": 50, "max": 400,
3734 "decrease_factor": 0.5, "cooldown_batches": 0, "error_threshold": 0.1
3735 }))
3736 .unwrap();
3737 let opts = RunStreamOptions::new().with_dlq(dlq).with_adaptive(cfg);
3738 let result = run_stream(stream, &sink, opts).await.unwrap();
3739 assert!(
3746 result.records_written >= 900,
3747 "expected ≥900 written, got {}",
3748 result.records_written
3749 );
3750 let sizes = sink.call_sizes();
3751 assert_eq!(sizes[0], 400, "first chunk is the full page");
3752 assert!(
3753 sizes.last().unwrap() < &400,
3754 "controller should shrink under errors: {sizes:?}"
3755 );
3756 assert!(
3757 result.dlq.unwrap().records_dlq >= 250,
3758 "expected ≥250 DLQ records"
3759 );
3760 }
3761
3762 struct RecordingSink {
3767 calls: std::sync::Mutex<Vec<usize>>,
3768 latency: std::time::Duration,
3769 }
3770 impl RecordingSink {
3771 fn new(latency_ms: u64) -> Self {
3772 Self {
3773 calls: std::sync::Mutex::new(Vec::new()),
3774 latency: std::time::Duration::from_millis(latency_ms),
3775 }
3776 }
3777 fn call_sizes(&self) -> Vec<usize> {
3778 self.calls.lock().unwrap().clone()
3779 }
3780 }
3781 #[async_trait]
3782 impl Sink for RecordingSink {
3783 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
3784 tokio::time::sleep(self.latency).await;
3785 self.calls.lock().unwrap().push(records.len());
3786 Ok(records.len())
3787 }
3788 }
3789
3790 #[tokio::test]
3791 async fn adaptive_reslices_non_dlq_page_into_subbatches() {
3792 use crate::adaptive::AdaptiveBatchConfig;
3793 let page = StreamPage {
3794 records: (0..1000).map(|i| json!({ "i": i })).collect(),
3795 bookmark: None,
3796 };
3797 let stream = futures::stream::iter(vec![Ok(page)]);
3798 let sink = RecordingSink::new(0);
3799 let cfg: AdaptiveBatchConfig =
3800 serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
3801 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
3802 .await
3803 .unwrap();
3804 assert_eq!(result.records_written, 1000);
3805 assert_eq!(sink.call_sizes(), vec![1000]);
3807 }
3808
3809 #[tokio::test]
3810 async fn adaptive_shrinks_under_latency_target_then_smaller_chunks() {
3811 use crate::adaptive::AdaptiveBatchConfig;
3812 let mk = || StreamPage {
3813 records: (0..400).map(|i| json!({"i": i})).collect(),
3814 bookmark: None,
3815 };
3816 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
3817 let sink = RecordingSink::new(50);
3818 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
3819 "enabled": true, "min": 50, "max": 400,
3820 "decrease_factor": 0.5, "cooldown_batches": 0,
3821 "target_latency_ms": 10, "latency_window": 1
3822 }))
3823 .unwrap();
3824 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
3825 .await
3826 .unwrap();
3827 assert_eq!(result.records_written, 1200);
3828 let sizes = sink.call_sizes();
3829 assert_eq!(sizes[0], 400);
3830 assert!(
3831 sizes.last().unwrap() < &400,
3832 "controller should have shrunk: {sizes:?}"
3833 );
3834 }
3835
3836 #[tokio::test]
3837 #[allow(clippy::await_holding_lock)]
3838 async fn adaptive_emits_batch_size_and_adjustments_metrics() {
3839 use crate::adaptive::AdaptiveBatchConfig;
3842 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
3843 use metrics_util::debugging::DebugValue;
3844
3845 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
3846 let snap = snapshotter();
3847
3848 let mk = || StreamPage {
3854 records: (0..400).map(|i| json!({"i": i})).collect(),
3855 bookmark: None,
3856 };
3857 let stream = futures::stream::iter(vec![Ok(mk()), Ok(mk()), Ok(mk())]);
3858 let sink = RecordingSink::new(50);
3859 let cfg: AdaptiveBatchConfig = serde_json::from_value(json!({
3860 "enabled": true, "min": 50, "max": 400,
3861 "decrease_factor": 0.5, "cooldown_batches": 0,
3862 "target_latency_ms": 10, "latency_window": 1
3863 }))
3864 .unwrap();
3865
3866 let _ = run_stream(
3867 stream,
3868 &sink,
3869 RunStreamOptions::new()
3870 .with_adaptive(cfg)
3871 .with_name("p")
3872 .with_row("r"),
3873 )
3874 .await
3875 .unwrap();
3876
3877 let snapshot = snap.snapshot();
3878 let mut saw_batch_size = false;
3879 let mut saw_adjustments = false;
3880 for (k, _u, _d, v) in snapshot.into_vec() {
3881 let key = k.key();
3882 let labels = key.labels().collect::<Vec<_>>();
3883 let has = |k: &str, val: &str| labels.iter().any(|l| l.key() == k && l.value() == val);
3884
3885 if key.name() == "faucet_pipeline_adaptive_batch_size"
3886 && has("pipeline", "p")
3887 && has("row", "r")
3888 && matches!(v, DebugValue::Gauge(_))
3889 {
3890 saw_batch_size = true;
3891 }
3892 if key.name() == "faucet_pipeline_adaptive_batch_adjustments_total"
3893 && has("pipeline", "p")
3894 && has("row", "r")
3895 && matches!(v, DebugValue::Counter(c) if c >= 1)
3896 {
3897 saw_adjustments = true;
3898 }
3899 }
3900 assert!(
3901 saw_batch_size,
3902 "expected faucet_pipeline_adaptive_batch_size gauge with pipeline=p, row=r"
3903 );
3904 assert!(
3905 saw_adjustments,
3906 "expected faucet_pipeline_adaptive_batch_adjustments_total counter with pipeline=p, row=r"
3907 );
3908 }
3909
3910 #[tokio::test]
3913 async fn exactly_once_rejects_dlq() {
3914 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3917 let dlq_sink: Arc<dyn Sink> = Arc::new(MockSink::new());
3918 let pages: Vec<Result<StreamPage, FaucetError>> = vec![];
3919 let opts = eo_opts(store, "k", 0).with_dlq(DlqConfig::new(dlq_sink));
3920 let r = run_stream(
3921 futures::stream::iter(pages),
3922 &IdempotentMockSink::new(),
3923 opts,
3924 )
3925 .await;
3926 assert!(
3927 matches!(&r, Err(FaucetError::Config(m)) if m.contains("not compatible with a DLQ")),
3928 "got: {r:?}"
3929 );
3930 }
3931
3932 #[tokio::test]
3935 async fn exactly_once_writes_unbookmarked_page_at_least_once() {
3936 let sink = IdempotentMockSink::new();
3941 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3942 let pages = vec![Ok(StreamPage {
3943 records: vec![json!({"id": 1}), json!({"id": 2})],
3944 bookmark: None,
3945 })];
3946 let r = run_stream(
3947 futures::stream::iter(pages),
3948 &sink,
3949 eo_opts(store.clone(), "k", 0),
3950 )
3951 .await
3952 .unwrap();
3953 assert_eq!(r.records_written, 2);
3954 assert_eq!(r.bookmark, None);
3955 assert_eq!(sink.last_committed_token("k").await.unwrap(), None);
3958 assert!(store.get("k").await.unwrap().is_none());
3960 assert_eq!(sink.rows(), vec![json!({"id": 1}), json!({"id": 2})]);
3961 }
3962
3963 #[tokio::test]
3966 async fn dlq_with_bookmark_persists_after_routing_failures() {
3967 let main = PartialSink::new(vec![1]); let dlq = std::sync::Arc::new(MockSink::new());
3972 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
3973 let pages: Vec<Result<StreamPage, FaucetError>> = vec![Ok(StreamPage {
3974 records: vec![json!({"i": 0}), json!({"i": 1})],
3975 bookmark: Some(json!("ckpt")),
3976 })];
3977 let result = run_stream(
3978 futures::stream::iter(pages),
3979 &main,
3980 RunStreamOptions::new()
3981 .with_dlq(DlqConfig::new(dlq.clone()))
3982 .with_state(Arc::clone(&store), "k"),
3983 )
3984 .await
3985 .unwrap();
3986
3987 assert_eq!(result.records_written, 1); assert_eq!(result.bookmark, Some(json!("ckpt")));
3989 assert_eq!(store.get("k").await.unwrap(), Some(json!("ckpt")));
3991 let envelopes = dlq.0.lock().unwrap();
3993 assert_eq!(envelopes.len(), 1);
3994 assert_eq!(envelopes[0]["payload"]["i"], 1);
3995 }
3996
3997 #[tokio::test]
4000 async fn exactly_once_resumes_sequence_from_start_seq() {
4001 let sink = IdempotentMockSink::new();
4005 let store: Arc<dyn StateStore> = Arc::new(MemoryStateStore::new());
4006 let pages = vec![Ok(StreamPage {
4007 records: vec![json!({"id": 9})],
4008 bookmark: Some(json!("bm-after-resume")),
4009 })];
4010 let r = run_stream(
4011 futures::stream::iter(pages),
4012 &sink,
4013 eo_opts(store.clone(), "eo_key", 7),
4014 )
4015 .await
4016 .unwrap();
4017
4018 assert_eq!(r.records_written, 1);
4019 let (bm, seq) =
4020 crate::idempotency::unwrap_state(&store.get("eo_key").await.unwrap().unwrap());
4021 assert_eq!(bm, Some(json!("bm-after-resume")));
4022 assert_eq!(seq, 8, "sequence resumes at start_seq + 1");
4023 assert_eq!(
4024 sink.last_committed_token("eo_key").await.unwrap(),
4025 Some(crate::idempotency::format_token(8))
4026 );
4027 }
4028
4029 #[cfg(feature = "quality")]
4032 #[tokio::test]
4033 async fn pipeline_run_with_quality_aborts_on_failed_batch_check() {
4034 use crate::quality::{BatchCheck, CompiledQuality, OnFailure, QualitySpec};
4035 let source = MockSource(vec![json!({"id": 1})]);
4036 let main = MockSink::new();
4037 let spec = QualitySpec {
4038 record: vec![],
4039 batch: vec![BatchCheck::RowCount {
4040 min: Some(5),
4041 max: None,
4042 on_failure: OnFailure::Abort,
4043 }],
4044 };
4045 let quality = Arc::new(CompiledQuality::compile(&spec).unwrap());
4046 let result = Pipeline::new(&source, &main)
4047 .with_quality(quality)
4048 .run()
4049 .await;
4050 assert!(matches!(result, Err(FaucetError::QualityFailure { .. })));
4051 assert!(main.written().is_empty());
4053 }
4054
4055 #[tokio::test]
4058 async fn pipeline_run_with_adaptive_reslices_page() {
4059 use crate::adaptive::AdaptiveBatchConfig;
4060 let source = MockSource((0..1000).map(|i| json!({ "i": i })).collect());
4061 let sink = MockSink::new();
4062 let cfg: AdaptiveBatchConfig =
4063 serde_json::from_value(json!({"enabled": true, "min": 100, "max": 1000})).unwrap();
4064 let result = Pipeline::new(&source, &sink)
4065 .with_adaptive(cfg)
4066 .run()
4067 .await
4068 .unwrap();
4069 assert_eq!(result.records_written, 1000);
4070 assert_eq!(sink.written().len(), 1000);
4071 }
4072
4073 #[tokio::test]
4076 async fn adaptive_noop_sink_name_is_handled() {
4077 use crate::adaptive::AdaptiveBatchConfig;
4078
4079 struct JsonlNamedSink(std::sync::Mutex<usize>);
4082 #[async_trait]
4083 impl Sink for JsonlNamedSink {
4084 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4085 *self.0.lock().unwrap() += records.len();
4086 Ok(records.len())
4087 }
4088 fn connector_name(&self) -> &'static str {
4089 "jsonl"
4090 }
4091 }
4092
4093 let page = StreamPage {
4094 records: (0..10).map(|i| json!({ "i": i })).collect(),
4095 bookmark: None,
4096 };
4097 let stream = futures::stream::iter(vec![Ok(page)]);
4098 let sink = JsonlNamedSink(std::sync::Mutex::new(0));
4099 let cfg: AdaptiveBatchConfig =
4100 serde_json::from_value(json!({"enabled": true, "min": 5, "max": 10})).unwrap();
4101 let result = run_stream(stream, &sink, RunStreamOptions::new().with_adaptive(cfg))
4102 .await
4103 .unwrap();
4104 assert_eq!(result.records_written, 10);
4105 assert_eq!(*sink.0.lock().unwrap(), 10);
4106 }
4107
4108 #[tokio::test]
4109 async fn resilience_retries_transient_sink_write() {
4110 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4111 use std::sync::Arc;
4112 use std::sync::atomic::{AtomicU32, Ordering};
4113 use std::time::Duration;
4114
4115 struct TransientFlakySink {
4116 attempts: Arc<AtomicU32>,
4117 written: Arc<AtomicU32>,
4118 }
4119 #[async_trait::async_trait]
4120 impl Sink for TransientFlakySink {
4121 async fn write_batch(
4122 &self,
4123 records: &[serde_json::Value],
4124 ) -> Result<usize, FaucetError> {
4125 let n = self.attempts.fetch_add(1, Ordering::SeqCst);
4126 if n < 2 {
4127 return Err(FaucetError::HttpStatus {
4128 status: 503,
4129 url: "u".into(),
4130 body: "".into(),
4131 });
4132 }
4133 self.written
4134 .fetch_add(records.len() as u32, Ordering::SeqCst);
4135 Ok(records.len())
4136 }
4137 async fn flush(&self) -> Result<(), FaucetError> {
4138 Ok(())
4139 }
4140 fn supports_idempotent_writes(&self) -> bool {
4142 true
4143 }
4144 }
4145
4146 let attempts = Arc::new(AtomicU32::new(0));
4147 let written = Arc::new(AtomicU32::new(0));
4148 let sink = TransientFlakySink {
4149 attempts: attempts.clone(),
4150 written: written.clone(),
4151 };
4152 let pages = futures::stream::iter(vec![Ok(StreamPage {
4153 records: vec![serde_json::json!({"a": 1})],
4154 bookmark: None,
4155 })]);
4156 let policy = ResiliencePolicy {
4157 retry: RetryPolicy {
4158 max_attempts: 5,
4159 backoff: BackoffKind::None,
4160 base: Duration::ZERO,
4161 max: Duration::ZERO,
4162 jitter: false,
4163 ..RetryPolicy::default()
4164 },
4165 ..ResiliencePolicy::default()
4166 };
4167 let res = run_stream(
4168 pages,
4169 &sink,
4170 RunStreamOptions::new().with_resilience(policy),
4171 )
4172 .await
4173 .unwrap();
4174 assert_eq!(written.load(Ordering::SeqCst), 1);
4175 assert_eq!(attempts.load(Ordering::SeqCst), 3);
4176 assert_eq!(res.records_written, 1);
4177 }
4178
4179 #[tokio::test]
4180 async fn resilience_does_not_retry_non_idempotent_write_batch() {
4181 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4185 use std::sync::Arc;
4186 use std::sync::atomic::{AtomicU32, Ordering};
4187 use std::time::Duration;
4188
4189 struct NonIdempotentFlakySink {
4190 attempts: Arc<AtomicU32>,
4191 }
4192 #[async_trait::async_trait]
4193 impl Sink for NonIdempotentFlakySink {
4194 async fn write_batch(&self, _r: &[Value]) -> Result<usize, FaucetError> {
4195 self.attempts.fetch_add(1, Ordering::SeqCst);
4196 Err(FaucetError::HttpStatus {
4197 status: 503,
4198 url: "u".into(),
4199 body: "".into(),
4200 })
4201 }
4202 async fn flush(&self) -> Result<(), FaucetError> {
4203 Ok(())
4204 }
4205 }
4207
4208 let attempts = Arc::new(AtomicU32::new(0));
4209 let sink = NonIdempotentFlakySink {
4210 attempts: attempts.clone(),
4211 };
4212 let pages = futures::stream::iter(vec![Ok(StreamPage {
4213 records: vec![json!({"a": 1})],
4214 bookmark: None,
4215 })]);
4216 let policy = ResiliencePolicy {
4217 retry: RetryPolicy {
4218 max_attempts: 5,
4219 backoff: BackoffKind::None,
4220 base: Duration::ZERO,
4221 max: Duration::ZERO,
4222 jitter: false,
4223 ..RetryPolicy::default()
4224 },
4225 ..ResiliencePolicy::default()
4226 };
4227 let err = run_stream(
4228 pages,
4229 &sink,
4230 RunStreamOptions::new().with_resilience(policy),
4231 )
4232 .await
4233 .unwrap_err();
4234 assert!(matches!(err, FaucetError::HttpStatus { status: 503, .. }));
4235 assert_eq!(
4236 attempts.load(Ordering::SeqCst),
4237 1,
4238 "non-idempotent write_batch must be attempted exactly once (no retry)"
4239 );
4240 }
4241
4242 #[tokio::test]
4243 async fn resilience_circuit_opens_after_consecutive_failed_pages() {
4244 use crate::dlq::{DlqConfig, OnBatchError};
4245 use crate::resilience::{BackoffKind, CircuitBreakerConfig, ResiliencePolicy, RetryPolicy};
4246 use std::sync::Arc;
4247 use std::time::Duration;
4248
4249 struct DeadSink;
4251 #[async_trait]
4252 impl Sink for DeadSink {
4253 async fn write_batch(&self, _r: &[Value]) -> Result<usize, FaucetError> {
4254 Err(FaucetError::Sink("down".into()))
4255 }
4256 async fn flush(&self) -> Result<(), FaucetError> {
4257 Ok(())
4258 }
4259 }
4260 struct NullSink;
4262 #[async_trait]
4263 impl Sink for NullSink {
4264 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4265 Ok(r.len())
4266 }
4267 async fn flush(&self) -> Result<(), FaucetError> {
4268 Ok(())
4269 }
4270 }
4271
4272 let pages = futures::stream::iter((0..10).map(|i| {
4273 Ok(StreamPage {
4274 records: vec![json!({"i": i})],
4275 bookmark: None,
4276 })
4277 }));
4278 let dlq = DlqConfig {
4279 on_batch_error: OnBatchError::DlqAll,
4280 ..DlqConfig::new(Arc::new(NullSink))
4281 };
4282 let policy = ResiliencePolicy {
4283 retry: RetryPolicy {
4284 max_attempts: 1,
4285 backoff: BackoffKind::None,
4286 base: Duration::ZERO,
4287 max: Duration::ZERO,
4288 jitter: false,
4289 ..RetryPolicy::default()
4290 },
4291 circuit_breaker: Some(CircuitBreakerConfig {
4292 consecutive_failures: 3,
4293 cooldown: Duration::from_secs(60),
4294 }),
4295 poison: None,
4296 };
4297 let err = run_stream(
4298 pages,
4299 &DeadSink,
4300 RunStreamOptions::new()
4301 .with_dlq(dlq)
4302 .with_resilience(policy),
4303 )
4304 .await
4305 .unwrap_err();
4306 assert!(
4307 matches!(err, FaucetError::CircuitOpen { failures: 3, .. }),
4308 "got {err:?}"
4309 );
4310 }
4311
4312 #[tokio::test]
4313 async fn resilience_poison_retries_then_dlqs_failing_row() {
4314 use crate::dlq::DlqConfig;
4315 use crate::resilience::{
4316 BackoffKind, PoisonAction, PoisonPolicy, ResiliencePolicy, RetryPolicy,
4317 };
4318 use std::sync::{Arc, Mutex};
4319 use std::time::Duration;
4320
4321 struct PickySink {
4324 bad_attempts: Arc<Mutex<u32>>,
4325 }
4326 #[async_trait]
4327 impl Sink for PickySink {
4328 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4329 Ok(r.len())
4330 }
4331 async fn write_batch_partial(
4332 &self,
4333 records: &[Value],
4334 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4335 Ok(records
4336 .iter()
4337 .map(|rec| {
4338 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
4339 *self.bad_attempts.lock().unwrap() += 1;
4340 Err(FaucetError::HttpStatus {
4341 status: 503,
4342 url: "u".into(),
4343 body: "".into(),
4344 })
4345 } else {
4346 Ok(())
4347 }
4348 })
4349 .collect())
4350 }
4351 async fn flush(&self) -> Result<(), FaucetError> {
4352 Ok(())
4353 }
4354 }
4355 struct CaptureSink(Arc<Mutex<Vec<Value>>>);
4356 #[async_trait]
4357 impl Sink for CaptureSink {
4358 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4359 self.0.lock().unwrap().extend_from_slice(r);
4360 Ok(r.len())
4361 }
4362 async fn flush(&self) -> Result<(), FaucetError> {
4363 Ok(())
4364 }
4365 }
4366
4367 let captured = Arc::new(Mutex::new(Vec::new()));
4368 let bad_attempts = Arc::new(Mutex::new(0u32));
4369 let sink = PickySink {
4370 bad_attempts: bad_attempts.clone(),
4371 };
4372 let pages = futures::stream::iter(vec![Ok(StreamPage {
4373 records: vec![json!({"ok": 1}), json!({"bad": true})],
4374 bookmark: None,
4375 })]);
4376 let policy = ResiliencePolicy {
4377 retry: RetryPolicy {
4378 max_attempts: 1,
4379 backoff: BackoffKind::None,
4380 base: Duration::ZERO,
4381 max: Duration::ZERO,
4382 jitter: false,
4383 ..RetryPolicy::default()
4384 },
4385 circuit_breaker: None,
4386 poison: Some(PoisonPolicy {
4387 max_row_attempts: 3,
4388 action: PoisonAction::Dlq,
4389 }),
4390 };
4391 let res = run_stream(
4392 pages,
4393 &sink,
4394 RunStreamOptions::new()
4395 .with_dlq(DlqConfig::new(Arc::new(CaptureSink(captured.clone()))))
4396 .with_resilience(policy),
4397 )
4398 .await
4399 .unwrap();
4400
4401 assert_eq!(
4402 *bad_attempts.lock().unwrap(),
4403 3,
4404 "bad row tried max_row_attempts times"
4405 );
4406 assert_eq!(res.records_written, 1, "the ok row");
4407 let dlq = captured.lock().unwrap();
4408 assert_eq!(dlq.len(), 1, "one row to DLQ");
4409 assert_eq!(dlq[0]["payload"]["bad"], json!(true));
4410 }
4411
4412 #[tokio::test]
4413 async fn poison_loop_does_not_nest_resilience_retry_on_subset_resubmit() {
4414 use crate::dlq::DlqConfig;
4422 use crate::resilience::{
4423 BackoffKind, PoisonAction, PoisonPolicy, ResiliencePolicy, RetryPolicy,
4424 };
4425 use std::sync::{Arc, Mutex};
4426 use std::time::Duration;
4427
4428 struct OuterErrOnSubsetSink {
4429 subset_calls: Arc<Mutex<u32>>,
4430 }
4431 #[async_trait]
4432 impl Sink for OuterErrOnSubsetSink {
4433 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4434 Ok(r.len())
4435 }
4436 async fn write_batch_partial(
4437 &self,
4438 records: &[Value],
4439 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4440 if records.len() == 1 {
4442 *self.subset_calls.lock().unwrap() += 1;
4443 return Err(FaucetError::HttpStatus {
4444 status: 503,
4445 url: "u".into(),
4446 body: "".into(),
4447 });
4448 }
4449 Ok(records
4450 .iter()
4451 .map(|rec| {
4452 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
4453 Err(FaucetError::HttpStatus {
4454 status: 503,
4455 url: "u".into(),
4456 body: "".into(),
4457 })
4458 } else {
4459 Ok(())
4460 }
4461 })
4462 .collect())
4463 }
4464 async fn flush(&self) -> Result<(), FaucetError> {
4465 Ok(())
4466 }
4467 fn supports_idempotent_writes(&self) -> bool {
4470 true
4471 }
4472 }
4473 struct CaptureSink(Arc<Mutex<Vec<Value>>>);
4474 #[async_trait]
4475 impl Sink for CaptureSink {
4476 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4477 self.0.lock().unwrap().extend_from_slice(r);
4478 Ok(r.len())
4479 }
4480 async fn flush(&self) -> Result<(), FaucetError> {
4481 Ok(())
4482 }
4483 }
4484
4485 let subset_calls = Arc::new(Mutex::new(0u32));
4486 let sink = OuterErrOnSubsetSink {
4487 subset_calls: subset_calls.clone(),
4488 };
4489 let pages = futures::stream::iter(vec![Ok(StreamPage {
4490 records: vec![json!({"ok": 1}), json!({"bad": true})],
4491 bookmark: None,
4492 })]);
4493 let policy = ResiliencePolicy {
4494 retry: RetryPolicy {
4495 max_attempts: 4, backoff: BackoffKind::None,
4497 base: Duration::ZERO,
4498 max: Duration::ZERO,
4499 jitter: false,
4500 ..RetryPolicy::default()
4501 },
4502 circuit_breaker: None,
4503 poison: Some(PoisonPolicy {
4504 max_row_attempts: 3,
4505 action: PoisonAction::Dlq,
4506 }),
4507 };
4508 let res = run_stream(
4509 pages,
4510 &sink,
4511 RunStreamOptions::new()
4512 .with_dlq(DlqConfig::new(Arc::new(CaptureSink(Arc::new(Mutex::new(
4513 Vec::new(),
4514 ))))))
4515 .with_resilience(policy),
4516 )
4517 .await;
4518
4519 assert!(matches!(
4521 res,
4522 Err(FaucetError::HttpStatus { status: 503, .. })
4523 ));
4524 assert_eq!(
4527 *subset_calls.lock().unwrap(),
4528 1,
4529 "poison subset resubmit must be bare (no nested resilience retry)"
4530 );
4531 }
4532
4533 #[tokio::test]
4534 async fn resilience_retries_transient_flush_and_state_on_dlq_path() {
4535 use crate::dlq::DlqConfig;
4542 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4543 use crate::state::StateStore;
4544 use std::sync::Arc;
4545 use std::sync::atomic::{AtomicU32, Ordering};
4546 use std::time::Duration;
4547
4548 fn transient_503() -> FaucetError {
4549 FaucetError::HttpStatus {
4550 status: 503,
4551 url: "u".into(),
4552 body: "".into(),
4553 }
4554 }
4555
4556 struct FlakyFlushSink {
4559 flush_attempts: Arc<AtomicU32>,
4560 }
4561 #[async_trait]
4562 impl Sink for FlakyFlushSink {
4563 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4564 Ok(r.len())
4565 }
4566 async fn write_batch_partial(
4567 &self,
4568 records: &[Value],
4569 ) -> Result<Vec<crate::RowOutcome>, FaucetError> {
4570 Ok(records
4571 .iter()
4572 .map(|rec| {
4573 if rec.get("bad").and_then(|v| v.as_bool()).unwrap_or(false) {
4574 Err(transient_503())
4575 } else {
4576 Ok(())
4577 }
4578 })
4579 .collect())
4580 }
4581 async fn flush(&self) -> Result<(), FaucetError> {
4582 let n = self.flush_attempts.fetch_add(1, Ordering::SeqCst);
4583 if n < 2 { Err(transient_503()) } else { Ok(()) }
4584 }
4585 }
4586 struct NullSink;
4587 #[async_trait]
4588 impl Sink for NullSink {
4589 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4590 Ok(r.len())
4591 }
4592 async fn flush(&self) -> Result<(), FaucetError> {
4593 Ok(())
4594 }
4595 }
4596
4597 struct FlakyStore {
4599 put_attempts: Arc<AtomicU32>,
4600 value: Arc<std::sync::Mutex<Option<Value>>>,
4601 }
4602 #[async_trait]
4603 impl StateStore for FlakyStore {
4604 async fn get(&self, _key: &str) -> Result<Option<Value>, FaucetError> {
4605 Ok(self.value.lock().unwrap().clone())
4606 }
4607 async fn put(&self, _key: &str, value: &Value) -> Result<(), FaucetError> {
4608 let n = self.put_attempts.fetch_add(1, Ordering::SeqCst);
4609 if n < 2 {
4610 return Err(transient_503());
4611 }
4612 *self.value.lock().unwrap() = Some(value.clone());
4613 Ok(())
4614 }
4615 async fn delete(&self, _key: &str) -> Result<(), FaucetError> {
4616 Ok(())
4617 }
4618 }
4619
4620 let flush_attempts = Arc::new(AtomicU32::new(0));
4621 let put_attempts = Arc::new(AtomicU32::new(0));
4622 let stored = Arc::new(std::sync::Mutex::new(None));
4623 let sink = FlakyFlushSink {
4624 flush_attempts: flush_attempts.clone(),
4625 };
4626 let store: Arc<dyn StateStore> = Arc::new(FlakyStore {
4627 put_attempts: put_attempts.clone(),
4628 value: stored.clone(),
4629 });
4630 let pages = futures::stream::iter(vec![Ok(StreamPage {
4631 records: vec![json!({"ok": 1}), json!({"bad": true})],
4632 bookmark: Some(json!({"cursor": 42})),
4633 })]);
4634 let policy = ResiliencePolicy {
4635 retry: RetryPolicy {
4636 max_attempts: 5,
4637 backoff: BackoffKind::None,
4638 base: Duration::ZERO,
4639 max: Duration::ZERO,
4640 jitter: false,
4641 ..RetryPolicy::default()
4642 },
4643 ..ResiliencePolicy::default()
4644 };
4645 let res = run_stream(
4646 pages,
4647 &sink,
4648 RunStreamOptions::new()
4649 .with_dlq(DlqConfig::new(Arc::new(NullSink)))
4650 .with_state(store, "k")
4651 .with_resilience(policy),
4652 )
4653 .await
4654 .unwrap();
4655
4656 assert_eq!(
4660 flush_attempts.load(Ordering::SeqCst),
4661 4,
4662 "page-gate flush retried past two transient failures (3) + 1 final flush"
4663 );
4664 assert_eq!(
4666 put_attempts.load(Ordering::SeqCst),
4667 3,
4668 "state put retried past two transient failures"
4669 );
4670 assert_eq!(*stored.lock().unwrap(), Some(json!({"cursor": 42})));
4671 assert_eq!(res.records_written, 1, "the ok row");
4672 }
4673
4674 #[tokio::test]
4675 #[allow(clippy::await_holding_lock)]
4676 async fn resilience_emits_retries_total_with_op_and_class_labels() {
4677 use crate::observability::decorator::source_tests::{LOCK, snapshotter};
4681 use crate::resilience::{BackoffKind, ResiliencePolicy, RetryPolicy};
4682 use metrics_util::debugging::DebugValue;
4683 use std::sync::Arc;
4684 use std::sync::atomic::{AtomicU32, Ordering};
4685 use std::time::Duration;
4686
4687 struct RetryProbeSink {
4688 attempts: Arc<AtomicU32>,
4689 }
4690 #[async_trait]
4691 impl Sink for RetryProbeSink {
4692 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
4693 let n = self.attempts.fetch_add(1, Ordering::SeqCst);
4694 if n < 2 {
4695 return Err(FaucetError::HttpStatus {
4696 status: 503,
4697 url: "u".into(),
4698 body: "".into(),
4699 });
4700 }
4701 Ok(r.len())
4702 }
4703 async fn flush(&self) -> Result<(), FaucetError> {
4704 Ok(())
4705 }
4706 fn connector_name(&self) -> &'static str {
4709 "retry-probe"
4710 }
4711 fn supports_idempotent_writes(&self) -> bool {
4713 true
4714 }
4715 }
4716
4717 let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
4718 let snap = snapshotter();
4719
4720 let attempts = Arc::new(AtomicU32::new(0));
4721 let sink = RetryProbeSink {
4722 attempts: attempts.clone(),
4723 };
4724 let pages = futures::stream::iter(vec![Ok(StreamPage {
4725 records: vec![json!({"a": 1})],
4726 bookmark: None,
4727 })]);
4728 let policy = ResiliencePolicy {
4729 retry: RetryPolicy {
4730 max_attempts: 5,
4731 backoff: BackoffKind::None,
4732 base: Duration::ZERO,
4733 max: Duration::ZERO,
4734 jitter: false,
4735 ..RetryPolicy::default()
4736 },
4737 ..ResiliencePolicy::default()
4738 };
4739 run_stream(
4740 pages,
4741 &sink,
4742 RunStreamOptions::new()
4743 .with_name("retry-metrics-pipeline")
4744 .with_resilience(policy),
4745 )
4746 .await
4747 .unwrap();
4748 assert_eq!(attempts.load(Ordering::SeqCst), 3);
4749
4750 let snapshot = snap.snapshot();
4751 let retries: u64 = snapshot
4752 .into_vec()
4753 .into_iter()
4754 .filter_map(|(key, _u, _d, v): (metrics_util::CompositeKey, _, _, _)| {
4755 if key.key().name() == "faucet_resilience_retries_total"
4756 && key.key().labels().any(|l: &metrics::Label| {
4757 l.key() == "pipeline" && l.value() == "retry-metrics-pipeline"
4758 })
4759 && key
4760 .key()
4761 .labels()
4762 .any(|l: &metrics::Label| l.key() == "op" && l.value() == "sink_write")
4763 && key
4764 .key()
4765 .labels()
4766 .any(|l: &metrics::Label| l.key() == "class" && l.value() == "http_5xx")
4767 && let DebugValue::Counter(c) = v
4768 {
4769 Some(c)
4770 } else {
4771 None
4772 }
4773 })
4774 .sum();
4775 assert_eq!(
4776 retries, 2,
4777 "expected 2 retries (2 transient 503s) counted with op=sink_write, class=http_5xx"
4778 );
4779 }
4780
4781 struct SchemaSink {
4785 schema: Value,
4786 written: std::sync::Mutex<Vec<Value>>,
4787 evolutions: std::sync::Mutex<Vec<crate::drift::SchemaEvolution>>,
4788 evolvable: bool,
4789 }
4790 impl SchemaSink {
4791 fn new(schema: Value, evolvable: bool) -> Self {
4792 Self {
4793 schema,
4794 written: std::sync::Mutex::new(Vec::new()),
4795 evolutions: std::sync::Mutex::new(Vec::new()),
4796 evolvable,
4797 }
4798 }
4799 fn written(&self) -> Vec<Value> {
4800 self.written.lock().unwrap().clone()
4801 }
4802 }
4803 #[async_trait]
4804 impl Sink for SchemaSink {
4805 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
4806 self.written.lock().unwrap().extend(records.iter().cloned());
4807 Ok(records.len())
4808 }
4809 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
4810 Ok(Some(self.schema.clone()))
4811 }
4812 fn supports_schema_evolution(&self) -> bool {
4813 self.evolvable
4814 }
4815 async fn evolve_schema(
4816 &self,
4817 evo: &crate::drift::SchemaEvolution,
4818 ) -> Result<(), FaucetError> {
4819 if !self.evolvable {
4820 return Err(FaucetError::Sink("not evolvable".into()));
4821 }
4822 self.evolutions.lock().unwrap().push(evo.clone());
4823 Ok(())
4824 }
4825 }
4826
4827 fn drift_opts(policy: crate::drift::SchemaDriftPolicy) -> RunStreamOptions {
4828 RunStreamOptions::new().with_schema_drift(policy)
4829 }
4830
4831 fn one_page(
4832 records: Vec<Value>,
4833 ) -> impl futures_core::Stream<Item = Result<StreamPage, FaucetError>> + Unpin {
4834 Box::pin(futures::stream::iter(vec![Ok(StreamPage {
4835 records,
4836 bookmark: None,
4837 })]))
4838 }
4839
4840 #[tokio::test]
4841 async fn drift_warn_writes_unchanged() {
4842 let sink = SchemaSink::new(
4843 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
4844 false,
4845 );
4846 let policy = crate::drift::SchemaDriftPolicy {
4847 on_drift: crate::drift::OnDrift::Warn,
4848 allow_widening: true,
4849 on_incompatible: crate::drift::OnIncompatible::Fail,
4850 relax_nullability_on_missing: false,
4851 };
4852 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
4853 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
4854 assert_eq!(res.records_written, 1);
4855 assert_eq!(sink.written()[0], json!({"id": 1, "email": "a@x.com"}));
4857 }
4858
4859 #[tokio::test]
4860 async fn drift_ignore_strips_unknown_fields() {
4861 let sink = SchemaSink::new(
4862 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
4863 false,
4864 );
4865 let policy = crate::drift::SchemaDriftPolicy {
4866 on_drift: crate::drift::OnDrift::Ignore,
4867 allow_widening: true,
4868 on_incompatible: crate::drift::OnIncompatible::Fail,
4869 relax_nullability_on_missing: false,
4870 };
4871 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
4872 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
4873 assert_eq!(res.records_written, 1);
4874 assert_eq!(
4875 sink.written()[0],
4876 json!({"id": 1}),
4877 "email must be stripped"
4878 );
4879 }
4880
4881 #[tokio::test]
4882 async fn drift_fail_raises_schema_drift() {
4883 let sink = SchemaSink::new(
4884 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
4885 false,
4886 );
4887 let policy = crate::drift::SchemaDriftPolicy {
4888 on_drift: crate::drift::OnDrift::Fail,
4889 allow_widening: true,
4890 on_incompatible: crate::drift::OnIncompatible::Fail,
4891 relax_nullability_on_missing: false,
4892 };
4893 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
4894 let err = run_stream(pages, &sink, drift_opts(policy))
4895 .await
4896 .unwrap_err();
4897 assert!(matches!(err, FaucetError::SchemaDrift { .. }));
4898 }
4899
4900 #[tokio::test]
4901 async fn drift_evolve_calls_sink_then_writes() {
4902 let sink = SchemaSink::new(
4903 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
4904 true,
4905 );
4906 let policy = crate::drift::SchemaDriftPolicy {
4907 on_drift: crate::drift::OnDrift::Evolve,
4908 allow_widening: true,
4909 on_incompatible: crate::drift::OnIncompatible::Fail,
4910 relax_nullability_on_missing: false,
4911 };
4912 let pages = one_page(vec![json!({"id": 1, "email": "a@x.com"})]);
4913 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
4914 assert_eq!(res.records_written, 1);
4915 let evos = sink.evolutions.lock().unwrap();
4916 assert_eq!(evos.len(), 1);
4917 assert_eq!(evos[0].additions.len(), 1);
4918 assert_eq!(evos[0].additions[0].name, "email");
4919 assert_eq!(sink.written()[0], json!({"id": 1, "email": "a@x.com"}));
4921 }
4922
4923 #[tokio::test]
4924 async fn drift_evolve_does_not_relax_not_null_for_merely_absent_column() {
4925 let sink = SchemaSink::new(
4930 json!({"type":"object","properties":{
4931 "id":{"type":"integer"},
4932 "legacy":{"type":"string"}
4933 }}),
4934 true,
4935 );
4936 let policy = crate::drift::SchemaDriftPolicy {
4937 on_drift: crate::drift::OnDrift::Evolve,
4938 allow_widening: true,
4939 on_incompatible: crate::drift::OnIncompatible::Fail,
4940 relax_nullability_on_missing: false,
4941 };
4942 let pages = one_page(vec![json!({"id": 1})]);
4943 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
4944 assert_eq!(res.records_written, 1);
4945 assert!(
4946 sink.evolutions.lock().unwrap().is_empty(),
4947 "NOT NULL must not be relaxed for a merely-absent column"
4948 );
4949 }
4950
4951 #[tokio::test]
4952 async fn drift_evolve_relaxes_absent_column_only_with_opt_in() {
4953 let sink = SchemaSink::new(
4955 json!({"type":"object","properties":{
4956 "id":{"type":"integer"},
4957 "legacy":{"type":"string"}
4958 }}),
4959 true,
4960 );
4961 let policy = crate::drift::SchemaDriftPolicy {
4962 on_drift: crate::drift::OnDrift::Evolve,
4963 allow_widening: true,
4964 on_incompatible: crate::drift::OnIncompatible::Fail,
4965 relax_nullability_on_missing: true,
4966 };
4967 let pages = one_page(vec![json!({"id": 1})]);
4968 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
4969 assert_eq!(res.records_written, 1);
4970 let evos = sink.evolutions.lock().unwrap();
4971 assert_eq!(evos.len(), 1);
4972 assert_eq!(evos[0].relax_nullability, vec!["legacy".to_string()]);
4973 }
4974
4975 #[tokio::test]
4976 async fn drift_inert_when_sink_reports_no_schema() {
4977 let sink = MockSink::new();
4979 let policy = crate::drift::SchemaDriftPolicy {
4980 on_drift: crate::drift::OnDrift::Fail, allow_widening: true,
4982 on_incompatible: crate::drift::OnIncompatible::Fail,
4983 relax_nullability_on_missing: false,
4984 };
4985 let pages = one_page(vec![json!({"id": 1, "anything": true})]);
4986 let res = run_stream(pages, &sink, drift_opts(policy)).await.unwrap();
4987 assert_eq!(res.records_written, 1);
4988 }
4989
4990 #[tokio::test]
4991 async fn drift_quarantine_routes_drift_rows_to_dlq() {
4992 let sink = SchemaSink::new(
4993 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
4994 false,
4995 );
4996 let dlq_sink = std::sync::Arc::new(MockSink::new());
4997 let policy = crate::drift::SchemaDriftPolicy {
4998 on_drift: crate::drift::OnDrift::Quarantine,
4999 allow_widening: true,
5000 on_incompatible: crate::drift::OnIncompatible::Fail,
5001 relax_nullability_on_missing: false,
5002 };
5003 let pages = one_page(vec![
5004 json!({"id": 1}), json!({"id": 2, "email": "x"}), ]);
5007 let opts = RunStreamOptions::new()
5008 .with_schema_drift(policy)
5009 .with_dlq(crate::dlq::DlqConfig::new(dlq_sink.clone()));
5010 let res = run_stream(pages, &sink, opts).await.unwrap();
5011 assert_eq!(res.records_written, 1, "only the conforming row is written");
5012 assert_eq!(sink.written(), vec![json!({"id": 1})]);
5013 let dlq = dlq_sink.written();
5015 assert_eq!(dlq.len(), 1);
5016 assert_eq!(dlq[0]["payload"], json!({"id": 2, "email": "x"}));
5017 assert_eq!(dlq[0]["error"]["kind"], "SchemaDrift");
5018 }
5019
5020 #[test]
5021 fn quarantine_drift_rows_covers_widening_and_droppable_required() {
5022 use crate::drift::{ColumnChange, SchemaDiff};
5023 let diff = SchemaDiff {
5025 additions: vec![],
5026 widenings: vec![ColumnChange {
5027 name: "amount".into(),
5028 from: Some(json!({"type":"integer"})),
5029 to: json!({"type":"number"}),
5030 }],
5031 incompatible: vec![],
5032 droppable_required: vec!["legacy".into()],
5033 };
5034 let records = vec![
5035 json!({"id": 1, "amount": 1.5, "legacy": "x"}), json!({"id": 2, "amount": 7, "legacy": "y"}), json!({"id": 3, "legacy": "z"}), json!({"id": 4}), ];
5040 let (kept, env) = quarantine_drift_rows(&diff, records, "sink", "pl", "");
5041 assert_eq!(kept, vec![json!({"id": 3, "legacy": "z"})]);
5042 assert_eq!(
5043 env.len(),
5044 3,
5045 "widening rows + the missing-required row quarantined"
5046 );
5047 }
5048
5049 #[tokio::test]
5050 async fn drift_quarantine_without_dlq_is_rejected() {
5051 let sink = SchemaSink::new(json!({"type":"object","properties":{}}), false);
5052 let policy = crate::drift::SchemaDriftPolicy {
5053 on_drift: crate::drift::OnDrift::Quarantine,
5054 allow_widening: true,
5055 on_incompatible: crate::drift::OnIncompatible::Fail,
5056 relax_nullability_on_missing: false,
5057 };
5058 let pages = one_page(vec![json!({"id": 1})]);
5059 let err = run_stream(pages, &sink, drift_opts(policy))
5060 .await
5061 .unwrap_err();
5062 assert!(matches!(err, FaucetError::Config(_)));
5063 }
5064
5065 #[cfg(feature = "quality")]
5071 #[tokio::test]
5072 async fn drift_fail_with_dlq_still_routes_quality_quarantine() {
5073 use crate::dlq::DlqConfig;
5074 use crate::quality::{CompiledQuality, OnFailure, QualitySpec, RecordCheck};
5075
5076 let sink = SchemaSink::new(
5079 json!({"type":"object","properties":{"id":{"type":"integer"}}}),
5080 false,
5081 );
5082 let dlq_sink = std::sync::Arc::new(MockSink::new());
5083 let policy = crate::drift::SchemaDriftPolicy {
5084 on_drift: crate::drift::OnDrift::Fail,
5085 allow_widening: true,
5086 on_incompatible: crate::drift::OnIncompatible::Fail,
5087 relax_nullability_on_missing: false,
5088 };
5089 let spec = QualitySpec {
5090 record: vec![RecordCheck::NotNull {
5091 field: "name".into(),
5092 treat_missing_as_null: true,
5093 on_failure: OnFailure::Quarantine,
5094 }],
5095 batch: vec![],
5096 };
5097 let quality = std::sync::Arc::new(CompiledQuality::compile(&spec).unwrap());
5098 let pages = one_page(vec![
5099 json!({"id": 1, "name": "ok", "email": "a@x"}), json!({"id": 2, "email": "b@x"}), ]);
5102 let opts = RunStreamOptions::new()
5103 .with_schema_drift(policy)
5104 .with_quality(quality)
5105 .with_dlq(DlqConfig::new(dlq_sink.clone()));
5106 let err = run_stream(pages, &sink, opts).await.unwrap_err();
5107 assert!(
5109 matches!(err, FaucetError::SchemaDrift { .. }),
5110 "got {err:?}"
5111 );
5112 let dlq = dlq_sink.written();
5114 assert_eq!(
5115 dlq.len(),
5116 1,
5117 "quarantined row must reach the DLQ before abort"
5118 );
5119 assert_eq!(dlq[0]["payload"], json!({"id": 2, "email": "b@x"}));
5120 assert_eq!(dlq[0]["error"]["kind"], "QualityFailure");
5121 assert_eq!(
5123 sink.written(),
5124 vec![json!({"id": 1, "name": "ok", "email": "a@x"})]
5125 );
5126 }
5127}