1pub(crate) mod apply;
48mod page;
49pub(crate) mod write;
50
51#[cfg(feature = "arrow")]
52pub(crate) use page::GovernedBatch;
53pub(crate) use page::{Clearance, GovernedNative};
54pub use page::{GovernedEnvelopes, GovernedPage};
55
56use crate::budget::{BudgetSink, BudgetSpec, BudgetState, BudgetTimer};
57use crate::dlq::DlqConfig;
58use crate::error::FaucetError;
59use crate::metadata::{CompiledMetadata, MetadataContext, MetadataSink};
60use crate::profiling::{Profiler, ProfilingSink};
61use crate::traits::{RowOutcome, Sink};
62use crate::usage::UsageMeter;
63use async_trait::async_trait;
64use serde::Serialize;
65use serde_json::Value;
66use std::sync::atomic::{AtomicBool, Ordering};
67use std::sync::{Arc, Mutex};
68use tokio_util::sync::CancellationToken;
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)]
72#[serde(rename_all = "snake_case")]
73#[non_exhaustive]
74pub enum Pass {
75 Masking,
77 Quality,
79 Contract,
81 SchemaDrift,
83 Policy,
85 Budget,
87 Profiling,
89 Metadata,
91 Metering,
93 Resilience,
95 Cancel,
97 Dlq,
99}
100
101impl Pass {
102 pub const ALL: &'static [Pass] = &[
104 Pass::Masking,
105 Pass::Quality,
106 Pass::Contract,
107 Pass::SchemaDrift,
108 Pass::Policy,
109 Pass::Budget,
110 Pass::Profiling,
111 Pass::Metadata,
112 Pass::Metering,
113 Pass::Resilience,
114 Pass::Cancel,
115 Pass::Dlq,
116 ];
117
118 pub fn as_str(self) -> &'static str {
120 match self {
121 Pass::Masking => "masking",
122 Pass::Quality => "quality",
123 Pass::Contract => "contract",
124 Pass::SchemaDrift => "schema_drift",
125 Pass::Policy => "policy",
126 Pass::Budget => "budget",
127 Pass::Profiling => "profiling",
128 Pass::Metadata => "metadata",
129 Pass::Metering => "metering",
130 Pass::Resilience => "resilience",
131 Pass::Cancel => "cancel",
132 Pass::Dlq => "dlq",
133 }
134 }
135}
136
137impl std::fmt::Display for Pass {
138 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
139 f.write_str(self.as_str())
140 }
141}
142
143#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
145#[serde(rename_all = "snake_case")]
146#[non_exhaustive]
147pub enum SinkLayer {
148 Policy,
150 Profiling,
152 OuterTap,
154 Metadata,
156 InnerTap,
158 Budget,
160}
161
162pub const SINK_LAYERS: &[SinkLayer] = &[
165 SinkLayer::Policy,
166 SinkLayer::Profiling,
167 SinkLayer::OuterTap,
168 SinkLayer::Metadata,
169 SinkLayer::InnerTap,
170 SinkLayer::Budget,
171];
172
173#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
175pub struct OptOut {
176 pub pass: Pass,
178 pub reason: String,
180}
181
182#[cfg(feature = "policy")]
184#[derive(Clone)]
185#[non_exhaustive]
186pub struct PolicyGovernance {
187 pub policy: Arc<crate::policy::CompiledPolicy>,
189 pub facts: crate::policy::SinkFacts,
191 pub scope: crate::policy::PolicyScope,
193 pub masking: Option<Arc<crate::masking::CompiledMasking>>,
195}
196
197#[cfg(feature = "policy")]
198impl PolicyGovernance {
199 pub fn new(
201 policy: Arc<crate::policy::CompiledPolicy>,
202 facts: crate::policy::SinkFacts,
203 scope: crate::policy::PolicyScope,
204 ) -> Self {
205 Self {
206 policy,
207 facts,
208 scope,
209 masking: None,
210 }
211 }
212
213 pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
216 self.masking = Some(masking);
217 self
218 }
219
220 pub fn quarantines(&self) -> bool {
222 self.policy.quarantines()
223 }
224}
225
226#[derive(Clone)]
228#[non_exhaustive]
229pub struct DlqGovernance {
230 pub config: DlqConfig,
232 #[cfg(feature = "policy")]
234 pub facts: Option<crate::policy::SinkFacts>,
235 #[cfg(feature = "masking")]
238 pub masking: Option<Arc<crate::masking::CompiledMasking>>,
239}
240
241impl DlqGovernance {
242 pub fn new(config: DlqConfig) -> Self {
244 Self {
245 config,
246 #[cfg(feature = "policy")]
247 facts: None,
248 #[cfg(feature = "masking")]
249 masking: None,
250 }
251 }
252
253 #[cfg(feature = "policy")]
255 pub fn with_facts(mut self, facts: crate::policy::SinkFacts) -> Self {
256 self.facts = Some(facts);
257 self
258 }
259
260 #[cfg(feature = "masking")]
262 pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
263 if !masking.is_empty() {
264 self.masking = Some(masking);
265 }
266 self
267 }
268}
269
270#[derive(Clone, Default)]
275#[non_exhaustive]
276pub struct Governance {
277 #[cfg(feature = "masking")]
278 masking: Option<Arc<crate::masking::CompiledMasking>>,
279 #[cfg(feature = "quality")]
280 quality: Option<Arc<crate::quality::CompiledQuality>>,
281 #[cfg(feature = "contract")]
282 contract: Option<Arc<crate::contract::CompiledContract>>,
283 schema_drift: Option<crate::drift::SchemaDriftPolicy>,
284 #[cfg(feature = "policy")]
285 policy: Option<PolicyGovernance>,
286 budget: Option<BudgetSpec>,
287 budget_state: Option<Arc<BudgetState>>,
288 profiler: Option<Arc<Mutex<Profiler>>>,
289 metadata: Option<(CompiledMetadata, MetadataContext)>,
290 usage_meter: Option<Arc<UsageMeter>>,
291 resilience: Option<crate::resilience::ResiliencePolicy>,
292 cancel: Option<CancellationToken>,
293 dlq: Option<DlqGovernance>,
294 opt_outs: Vec<OptOut>,
295 sink_built: Arc<AtomicBool>,
296}
297
298impl std::fmt::Debug for Governance {
299 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
300 f.debug_struct("Governance")
301 .field("active", &self.active())
302 .field("opt_outs", &self.opt_outs)
303 .finish()
304 }
305}
306
307impl Governance {
308 pub fn new() -> Self {
310 Self::default()
311 }
312
313 #[cfg(feature = "masking")]
315 pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
316 self.masking = (!masking.is_empty()).then_some(masking);
317 self
318 }
319
320 #[cfg(feature = "quality")]
322 pub fn with_quality(mut self, quality: Arc<crate::quality::CompiledQuality>) -> Self {
323 self.quality = Some(quality);
324 self
325 }
326
327 #[cfg(feature = "contract")]
329 pub fn with_contract(mut self, contract: Arc<crate::contract::CompiledContract>) -> Self {
330 self.contract = Some(contract);
331 self
332 }
333
334 pub fn with_schema_drift(mut self, policy: crate::drift::SchemaDriftPolicy) -> Self {
336 self.schema_drift = Some(policy);
337 self
338 }
339
340 #[cfg(feature = "policy")]
342 pub fn with_policy(mut self, policy: PolicyGovernance) -> Self {
343 self.policy = Some(policy);
344 self
345 }
346
347 pub fn with_budget(mut self, spec: BudgetSpec) -> Self {
350 if spec.max_records.is_none()
351 && spec.max_bytes.is_none()
352 && spec.max_duration_secs.is_none()
353 {
354 self.budget = None;
355 self.budget_state = None;
356 } else {
357 self.budget = Some(spec);
358 self.budget_state = Some(Arc::new(BudgetState::default()));
359 }
360 self
361 }
362
363 pub fn with_profiler(mut self, profiler: Arc<Mutex<Profiler>>) -> Self {
365 self.profiler = Some(profiler);
366 self
367 }
368
369 pub fn with_metadata(mut self, meta: CompiledMetadata, ctx: MetadataContext) -> Self {
371 self.metadata = Some((meta, ctx));
372 self
373 }
374
375 pub fn with_usage_meter(mut self, meter: Arc<UsageMeter>) -> Self {
377 self.usage_meter = Some(meter);
378 self
379 }
380
381 pub fn with_resilience(mut self, policy: crate::resilience::ResiliencePolicy) -> Self {
383 self.resilience = Some(policy);
384 self
385 }
386
387 pub fn with_cancel(mut self, cancel: CancellationToken) -> Self {
390 self.cancel = Some(cancel);
391 self
392 }
393
394 pub fn with_dlq(mut self, dlq: DlqGovernance) -> Self {
396 self.dlq = Some(dlq);
397 self
398 }
399
400 pub fn without(mut self, pass: Pass, reason: impl Into<String>) -> Self {
403 let reason = reason.into();
404 match pass {
405 #[cfg(feature = "masking")]
406 Pass::Masking => self.masking = None,
407 #[cfg(feature = "quality")]
408 Pass::Quality => self.quality = None,
409 #[cfg(feature = "contract")]
410 Pass::Contract => self.contract = None,
411 Pass::SchemaDrift => self.schema_drift = None,
412 #[cfg(feature = "policy")]
413 Pass::Policy => self.policy = None,
414 Pass::Budget => {
415 self.budget = None;
416 self.budget_state = None;
417 }
418 Pass::Profiling => self.profiler = None,
419 Pass::Metadata => self.metadata = None,
420 Pass::Metering => self.usage_meter = None,
421 Pass::Resilience => self.resilience = None,
422 Pass::Cancel => self.cancel = None,
423 Pass::Dlq => self.dlq = None,
424 #[allow(unreachable_patterns)]
425 _ => {}
426 }
427 tracing::info!(pass = pass.as_str(), reason = %reason, "governance pass opted out");
428 self.opt_outs.retain(|o| o.pass != pass);
429 self.opt_outs.push(OptOut { pass, reason });
430 self
431 }
432
433 pub fn opt_outs(&self) -> &[OptOut] {
435 &self.opt_outs
436 }
437
438 pub fn opted_out(&self, pass: Pass) -> Option<&str> {
440 self.opt_outs
441 .iter()
442 .find(|o| o.pass == pass)
443 .map(|o| o.reason.as_str())
444 }
445
446 pub fn is_active(&self, pass: Pass) -> bool {
448 match pass {
449 #[cfg(feature = "masking")]
450 Pass::Masking => self.masking.is_some(),
451 #[cfg(feature = "quality")]
452 Pass::Quality => self.quality.is_some(),
453 #[cfg(feature = "contract")]
454 Pass::Contract => self.contract.is_some(),
455 Pass::SchemaDrift => self.schema_drift.is_some(),
456 #[cfg(feature = "policy")]
457 Pass::Policy => self.policy.is_some(),
458 Pass::Budget => self.budget.is_some(),
459 Pass::Profiling => self.profiler.is_some(),
460 Pass::Metadata => self.metadata.is_some(),
461 Pass::Metering => self.usage_meter.is_some(),
462 Pass::Resilience => self.resilience.is_some(),
463 Pass::Cancel => self.cancel.is_some(),
464 Pass::Dlq => self.dlq.is_some(),
465 #[allow(unreachable_patterns)]
466 _ => false,
467 }
468 }
469
470 pub fn active(&self) -> Vec<Pass> {
472 Pass::ALL
473 .iter()
474 .copied()
475 .filter(|p| self.is_active(*p))
476 .collect()
477 }
478
479 pub fn budget(&self) -> Option<&BudgetSpec> {
481 self.budget.as_ref()
482 }
483
484 pub fn budget_state(&self) -> Option<&Arc<BudgetState>> {
486 self.budget_state.as_ref()
487 }
488
489 pub fn cancel(&self) -> Option<&CancellationToken> {
491 self.cancel.as_ref()
492 }
493
494 pub fn usage_meter(&self) -> Option<&Arc<UsageMeter>> {
496 self.usage_meter.as_ref()
497 }
498
499 pub fn resilience(&self) -> Option<&crate::resilience::ResiliencePolicy> {
501 self.resilience.as_ref()
502 }
503
504 pub fn dlq(&self) -> Option<&DlqGovernance> {
506 self.dlq.as_ref()
507 }
508
509 #[cfg(feature = "masking")]
511 pub fn masking(&self) -> Option<&Arc<crate::masking::CompiledMasking>> {
512 self.masking.as_ref()
513 }
514
515 #[cfg(feature = "quality")]
517 pub fn quality(&self) -> Option<&Arc<crate::quality::CompiledQuality>> {
518 self.quality.as_ref()
519 }
520
521 #[cfg(feature = "contract")]
523 pub fn contract(&self) -> Option<&Arc<crate::contract::CompiledContract>> {
524 self.contract.as_ref()
525 }
526
527 pub fn schema_drift(&self) -> Option<crate::drift::SchemaDriftPolicy> {
529 self.schema_drift
530 }
531
532 #[cfg(feature = "policy")]
534 pub fn policy(&self) -> Option<&PolicyGovernance> {
535 self.policy.as_ref()
536 }
537
538 pub fn policy_quarantines(&self) -> bool {
541 #[cfg(feature = "policy")]
542 {
543 self.policy
544 .as_ref()
545 .is_some_and(PolicyGovernance::quarantines)
546 }
547 #[cfg(not(feature = "policy"))]
548 {
549 false
550 }
551 }
552
553 pub fn has_sink_layers(&self) -> bool {
556 self.is_active(Pass::Budget)
557 || self.is_active(Pass::Profiling)
558 || self.is_active(Pass::Metadata)
559 || self.policy_active()
560 }
561
562 fn policy_active(&self) -> bool {
563 #[cfg(feature = "policy")]
564 {
565 self.policy.is_some()
566 }
567 #[cfg(not(feature = "policy"))]
568 {
569 false
570 }
571 }
572
573 pub fn sink_built(&self) -> bool {
575 self.sink_built.load(Ordering::Relaxed)
576 }
577
578 pub fn sink_builder(&self) -> GovernedSinkBuilder<'_> {
580 GovernedSinkBuilder {
581 gov: self,
582 inner_tap: None,
583 outer_tap: None,
584 }
585 }
586
587 pub async fn govern<Si: Sink + ?Sized>(
598 &self,
599 records: Vec<Value>,
600 sink: &Si,
601 ) -> Result<GovernedOutput, FaucetError> {
602 let specs = apply::GovernanceSpecs {
603 #[cfg(feature = "masking")]
604 masking: self.masking.as_ref(),
605 #[cfg(feature = "quality")]
606 quality: self.quality.as_ref(),
607 #[cfg(feature = "contract")]
608 contract: self.contract.as_ref(),
609 schema_drift: self.schema_drift.as_ref(),
610 pipeline: "unnamed",
611 row: "",
612 run_id: "",
613 sink_name: sink.connector_name(),
614 has_dlq: false,
615 };
616 let mut state = apply::GovernanceState::default();
617 let out = apply::apply_governance(records, &specs, &mut state, sink).await?;
618 Ok(GovernedOutput {
619 page: out.records,
620 envelopes: out.envelopes,
621 })
622 }
623
624 pub(crate) fn compat_features(&self) -> Vec<crate::compat::Feature> {
626 use crate::compat::Feature as F;
627 [
628 (self.policy_quarantines(), F::PolicyQuarantine),
629 (self.policy_active(), F::Policy),
630 (self.is_active(Pass::Budget), F::Budget),
631 (self.is_active(Pass::Profiling), F::Profiling),
632 (self.is_active(Pass::Metadata), F::MetadataColumns),
633 ]
634 .into_iter()
635 .filter_map(|(on, f)| on.then_some(f))
636 .collect()
637 }
638
639 pub fn dlq_config(&self) -> Option<DlqConfig> {
644 self.governed_dlq()
645 }
646
647 pub(crate) fn governed_dlq(&self) -> Option<DlqConfig> {
650 let dlq = self.dlq.as_ref()?;
651 Some(govern_dlq(
652 dlq,
653 self.usage_meter.as_ref(),
654 self.budget.as_ref().zip(self.budget_state.as_ref()),
655 self.cancel.clone(),
656 #[cfg(feature = "policy")]
657 self.policy.as_ref(),
658 ))
659 }
660}
661
662#[derive(Debug, Clone)]
664#[non_exhaustive]
665pub struct GovernedOutput {
666 pub page: GovernedPage<'static>,
668 pub envelopes: GovernedEnvelopes,
670}
671
672pub(crate) fn govern_dlq(
675 dlq: &DlqGovernance,
676 meter: Option<&Arc<UsageMeter>>,
677 budget: Option<(&BudgetSpec, &Arc<BudgetState>)>,
678 cancel: Option<CancellationToken>,
679 #[cfg(feature = "policy")] policy: Option<&PolicyGovernance>,
680) -> DlqConfig {
681 let mut sink: Box<dyn Sink> = Box::new(SharedSink(Arc::clone(&dlq.config.sink)));
682 if let Some(m) = meter {
683 sink = Box::new(DlqMeteredSink {
684 inner: sink,
685 meter: Arc::clone(m),
686 });
687 }
688 if let Some((spec, state)) = budget {
689 sink = Box::new(BudgetSink::wrap_shared(
690 sink,
691 spec.clone(),
692 Arc::clone(state),
693 cancel.unwrap_or_default(),
694 ));
695 }
696 #[cfg(feature = "policy")]
697 if let (Some(p), Some(facts)) = (policy, dlq.facts.as_ref()) {
698 let mut ps = crate::policy::PolicySink::for_dlq(
699 sink,
700 Arc::clone(&p.policy),
701 facts.clone(),
702 p.scope.clone(),
703 );
704 if let Some(m) = &p.masking {
705 ps = ps.with_masking(Arc::clone(m));
706 }
707 sink = Box::new(ps);
708 }
709 #[cfg(feature = "masking")]
710 if let Some(m) = &dlq.masking {
711 sink = Box::new(DlqMaskingSink {
712 inner: sink,
713 masking: Arc::clone(m),
714 });
715 }
716 DlqConfig {
717 sink: Arc::from(sink),
718 ..dlq.config.clone()
719 }
720}
721
722type Tap = Box<dyn FnOnce(Box<dyn Sink>) -> Box<dyn Sink> + Send>;
723
724pub struct GovernedSinkBuilder<'g> {
727 gov: &'g Governance,
728 inner_tap: Option<Tap>,
729 outer_tap: Option<Tap>,
730}
731
732impl<'g> GovernedSinkBuilder<'g> {
733 pub fn inner_tap(
735 mut self,
736 tap: impl FnOnce(Box<dyn Sink>) -> Box<dyn Sink> + Send + 'static,
737 ) -> Self {
738 self.inner_tap = Some(Box::new(tap));
739 self
740 }
741
742 pub fn outer_tap(
744 mut self,
745 tap: impl FnOnce(Box<dyn Sink>) -> Box<dyn Sink> + Send + 'static,
746 ) -> Self {
747 self.outer_tap = Some(Box::new(tap));
748 self
749 }
750
751 pub fn build(mut self, destination: Box<dyn Sink>) -> GovernedSink {
754 let gov = self.gov;
755 let mut layers = Vec::new();
756 let mut sink = destination;
757 let mut timer = None;
758 for layer in SINK_LAYERS.iter().rev() {
759 match layer {
760 SinkLayer::Budget => {
761 if let (Some(spec), Some(state)) = (&gov.budget, &gov.budget_state) {
762 let cancel = gov.cancel.clone().unwrap_or_default();
763 timer = Some(BudgetTimer::start_for(
764 spec.max_duration_secs,
765 Arc::clone(state),
766 cancel.clone(),
767 ));
768 sink = Box::new(BudgetSink::wrap_shared(
769 sink,
770 spec.clone(),
771 Arc::clone(state),
772 cancel,
773 ));
774 layers.push(*layer);
775 }
776 }
777 SinkLayer::InnerTap => {
778 if let Some(tap) = self.inner_tap.take() {
779 sink = tap(sink);
780 layers.push(*layer);
781 }
782 }
783 SinkLayer::Metadata => {
784 if let Some((meta, ctx)) = &gov.metadata {
785 sink = Box::new(MetadataSink::new(sink, meta.clone(), ctx.clone()));
786 layers.push(*layer);
787 }
788 }
789 SinkLayer::OuterTap => {
790 if let Some(tap) = self.outer_tap.take() {
791 sink = tap(sink);
792 layers.push(*layer);
793 }
794 }
795 SinkLayer::Profiling => {
796 if let Some(p) = &gov.profiler {
797 sink = Box::new(ProfilingSink::new(sink, Arc::clone(p)));
798 layers.push(*layer);
799 }
800 }
801 SinkLayer::Policy =>
802 {
803 #[cfg(feature = "policy")]
804 if let Some(p) = &gov.policy {
805 let mut ps = crate::policy::PolicySink::new(
806 sink,
807 Arc::clone(&p.policy),
808 p.facts.clone(),
809 p.scope.clone(),
810 );
811 if let Some(facts) = gov.dlq.as_ref().and_then(|d| d.facts.clone()) {
812 ps = ps.with_dlq_destination(facts);
813 }
814 if let Some(m) = &p.masking {
815 ps = ps.with_masking(Arc::clone(m));
816 }
817 sink = Box::new(ps);
818 layers.push(*layer);
819 }
820 }
821 }
822 }
823 layers.reverse();
824 gov.sink_built.store(true, Ordering::Relaxed);
825 GovernedSink {
826 sink,
827 budget: gov.budget_state.clone(),
828 _timer: timer,
829 layers,
830 }
831 }
832}
833
834pub struct GovernedSink {
837 sink: Box<dyn Sink>,
838 budget: Option<Arc<BudgetState>>,
839 _timer: Option<BudgetTimer>,
840 layers: Vec<SinkLayer>,
841}
842
843impl GovernedSink {
844 pub fn layers(&self) -> &[SinkLayer] {
846 &self.layers
847 }
848
849 pub fn budget_state(&self) -> Option<&Arc<BudgetState>> {
851 self.budget.as_ref()
852 }
853
854 pub fn finish<T>(&self, result: Result<T, FaucetError>) -> Result<T, FaucetError> {
857 match (&result, self.budget.as_ref().and_then(|s| s.verdict())) {
858 (Ok(_), Some(verdict)) => Err(verdict.error()),
859 _ => result,
860 }
861 }
862
863 pub fn into_inner(self) -> Box<dyn Sink> {
865 self.sink
866 }
867}
868
869#[async_trait]
870impl Sink for GovernedSink {
871 crate::forward_sink!(sink);
872}
873
874pub(crate) struct SharedSink(pub(crate) Arc<dyn Sink>);
876
877#[async_trait]
878impl Sink for SharedSink {
879 crate::forward_sink!(0);
880}
881
882struct DlqMeteredSink {
884 inner: Box<dyn Sink>,
885 meter: Arc<UsageMeter>,
886}
887
888#[async_trait]
889impl Sink for DlqMeteredSink {
890 crate::forward_sink!(inner; except write_batch, write_batch_reported, write_batch_partial);
891
892 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
893 let n = self.inner.write_batch(records).await?;
894 self.meter
895 .add_dlq_written(n as u64, crate::usage::estimate_page_bytes(records));
896 Ok(n)
897 }
898
899 async fn write_batch_reported(
900 &self,
901 records: &[Value],
902 ) -> Result<crate::accounting::WriteReport, FaucetError> {
903 let r = self.inner.write_batch_reported(records).await?;
904 self.meter.add_dlq_written(
905 r.accepted as u64,
906 crate::usage::estimate_page_bytes(records),
907 );
908 Ok(r)
909 }
910
911 async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
912 let outcomes = self.inner.write_batch_partial(records).await?;
913 let ok = outcomes.iter().filter(|o| o.is_ok()).count();
914 self.meter
915 .add_dlq_written(ok as u64, crate::usage::estimate_page_bytes(records));
916 Ok(outcomes)
917 }
918}
919
920#[cfg(feature = "masking")]
922struct DlqMaskingSink {
923 inner: Box<dyn Sink>,
924 masking: Arc<crate::masking::CompiledMasking>,
925}
926
927#[cfg(feature = "masking")]
928impl DlqMaskingSink {
929 fn mask(&self, records: &[Value]) -> Vec<Value> {
930 let mut out = records.to_vec();
931 let payloads: Vec<Value> = out
932 .iter_mut()
933 .map(|e| match e.get_mut("payload") {
934 Some(p) => p.take(),
935 None => Value::Null,
936 })
937 .collect();
938 let masked = crate::masking::apply_masking(payloads, &self.masking).records;
939 for (envelope, payload) in out.iter_mut().zip(masked) {
940 if let Some(slot) = envelope.get_mut("payload") {
941 *slot = payload;
942 }
943 }
944 out
945 }
946}
947
948#[cfg(feature = "masking")]
949#[async_trait]
950impl Sink for DlqMaskingSink {
951 crate::forward_sink!(inner; except write_batch, write_batch_reported, write_batch_partial);
952
953 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
954 self.inner.write_batch(&self.mask(records)).await
955 }
956
957 async fn write_batch_reported(
958 &self,
959 records: &[Value],
960 ) -> Result<crate::accounting::WriteReport, FaucetError> {
961 self.inner.write_batch_reported(&self.mask(records)).await
962 }
963
964 async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
965 self.inner.write_batch_partial(&self.mask(records)).await
966 }
967}
968
969#[cfg(test)]
970mod tests {
971 use super::*;
972 use crate::traits::Source;
973 use serde_json::json;
974 use std::collections::HashMap;
975
976 type Seen = Arc<Mutex<Vec<(&'static str, Value)>>>;
977
978 struct Tap {
980 name: &'static str,
981 inner: Box<dyn Sink>,
982 seen: Seen,
983 }
984
985 #[async_trait]
986 impl Sink for Tap {
987 crate::forward_sink!(inner; except write_batch);
988 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
989 if let Some(r) = records.first() {
990 self.seen.lock().unwrap().push((self.name, r.clone()));
991 }
992 self.inner.write_batch(records).await
993 }
994 }
995
996 fn tap(
997 name: &'static str,
998 seen: &Seen,
999 ) -> impl FnOnce(Box<dyn Sink>) -> Box<dyn Sink> + Send + 'static {
1000 let seen = Arc::clone(seen);
1001 move |inner| Box::new(Tap { name, inner, seen })
1002 }
1003
1004 struct Dest(Seen);
1005
1006 #[async_trait]
1007 impl Sink for Dest {
1008 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1009 for r in records {
1010 self.0.lock().unwrap().push(("dest", r.clone()));
1011 }
1012 Ok(records.len())
1013 }
1014 }
1015
1016 fn profiler() -> Arc<Mutex<Profiler>> {
1017 Arc::new(Mutex::new(Profiler::new(
1018 serde_json::from_value(json!({})).unwrap(),
1019 )))
1020 }
1021
1022 fn metadata() -> (CompiledMetadata, MetadataContext) {
1023 (
1024 CompiledMetadata::compile(&crate::metadata::MetadataColumnsSpec::default())
1025 .unwrap()
1026 .unwrap(),
1027 MetadataContext {
1028 run_id: "r".into(),
1029 source: "s".into(),
1030 },
1031 )
1032 }
1033
1034 fn is_subsequence(layers: &[SinkLayer]) -> bool {
1035 let mut canon = SINK_LAYERS.iter();
1036 layers.iter().all(|l| canon.any(|c| c == l))
1037 }
1038
1039 #[tokio::test]
1043 async fn layers_wrap_the_destination_in_the_canonical_order() {
1044 let (meta, ctx) = metadata();
1045 let g = Governance::new()
1046 .with_budget(BudgetSpec {
1047 max_records: Some(10),
1048 ..Default::default()
1049 })
1050 .with_metadata(meta, ctx)
1051 .with_profiler(profiler());
1052 let seen: Seen = Arc::default();
1053 let governed = g
1054 .sink_builder()
1055 .inner_tap(tap("inner", &seen))
1056 .outer_tap(tap("outer", &seen))
1057 .build(Box::new(Dest(Arc::clone(&seen))));
1058 assert_eq!(
1059 governed.layers(),
1060 &[
1061 SinkLayer::Profiling,
1062 SinkLayer::OuterTap,
1063 SinkLayer::Metadata,
1064 SinkLayer::InnerTap,
1065 SinkLayer::Budget,
1066 ]
1067 );
1068 assert!(is_subsequence(governed.layers()));
1069 assert!(g.sink_built());
1070 governed.write_batch(&[json!({"a": 1})]).await.unwrap();
1071 let seen = seen.lock().unwrap().clone();
1072 let names: Vec<&str> = seen.iter().map(|(n, _)| *n).collect();
1073 assert_eq!(names, ["outer", "inner", "dest"]);
1074 assert!(
1075 seen[0].1.get("_faucet_run_id").is_none(),
1076 "the outer tap sees the source record"
1077 );
1078 assert_eq!(
1079 seen[1].1["_faucet_run_id"], "r",
1080 "metadata is stamped above the inner tap"
1081 );
1082 assert_eq!(governed.budget_state().unwrap().records(), 1);
1083 assert!(governed.finish(Ok::<_, FaucetError>(())).is_ok());
1084 let _ = governed.into_inner();
1085 }
1086
1087 #[tokio::test]
1088 async fn an_empty_governance_builds_the_bare_destination() {
1089 let g = Governance::new().with_budget(BudgetSpec::default());
1090 assert!(!g.has_sink_layers());
1091 assert!(g.budget().is_none());
1092 let governed = g.sink_builder().build(Box::new(Dest(Arc::default())));
1093 assert!(governed.layers().is_empty());
1094 assert!(governed.budget_state().is_none());
1095 assert_eq!(governed.write_batch(&[json!({})]).await.unwrap(), 1);
1096 }
1097
1098 #[tokio::test]
1099 async fn a_spent_duration_turns_the_partial_result_into_the_verdict() {
1100 let cancel = CancellationToken::new();
1101 let g = Governance::new()
1102 .with_cancel(cancel.clone())
1103 .with_budget(BudgetSpec {
1104 max_duration_secs: Some(1),
1105 ..Default::default()
1106 });
1107 let governed = g.sink_builder().build(Box::new(Dest(Arc::default())));
1108 tokio::time::timeout(std::time::Duration::from_secs(5), cancel.cancelled())
1109 .await
1110 .expect("the duration ceiling cancels the run");
1111 let err = governed.finish(Ok::<_, FaucetError>(())).unwrap_err();
1112 assert!(matches!(err, FaucetError::BudgetExceeded { .. }), "{err}");
1113 }
1114
1115 #[test]
1116 fn opt_outs_clear_the_pass_and_are_recorded() {
1117 let (meta, ctx) = metadata();
1118 let g = Governance::new()
1119 .with_budget(BudgetSpec {
1120 max_records: Some(1),
1121 ..Default::default()
1122 })
1123 .with_metadata(meta, ctx)
1124 .with_profiler(profiler())
1125 .with_cancel(CancellationToken::new())
1126 .with_usage_meter(Arc::new(UsageMeter::new()))
1127 .with_resilience(crate::resilience::ResiliencePolicy::default())
1128 .with_schema_drift(crate::drift::SchemaDriftPolicy::compile(
1129 &serde_json::from_value(json!({})).unwrap(),
1130 ));
1131 for p in [
1132 Pass::Budget,
1133 Pass::Metadata,
1134 Pass::Profiling,
1135 Pass::Cancel,
1136 Pass::Metering,
1137 Pass::Resilience,
1138 Pass::SchemaDrift,
1139 ] {
1140 assert!(g.is_active(p), "{p}");
1141 }
1142 assert!(g.has_sink_layers());
1143 let g = g
1144 .without(Pass::Budget, "why")
1145 .without(Pass::Profiling, "first")
1146 .without(Pass::Profiling, "second")
1147 .without(Pass::Dlq, "none");
1148 assert!(!g.is_active(Pass::Budget));
1149 assert!(g.budget_state().is_none());
1150 assert!(!g.is_active(Pass::Profiling));
1151 assert_eq!(g.opted_out(Pass::Profiling), Some("second"));
1152 assert_eq!(g.opted_out(Pass::Budget), Some("why"));
1153 assert_eq!(g.opted_out(Pass::Masking), None);
1154 assert_eq!(g.opt_outs().len(), 3);
1155 assert!(g.active().contains(&Pass::Metadata));
1156 assert!(format!("{g:?}").contains("opt_outs"));
1157 for p in Pass::ALL {
1158 assert_eq!(p.to_string(), p.as_str());
1159 }
1160 let g = g
1161 .without(Pass::Metadata, "m")
1162 .without(Pass::Cancel, "c")
1163 .without(Pass::Metering, "u")
1164 .without(Pass::Resilience, "r")
1165 .without(Pass::SchemaDrift, "d");
1166 assert!(!g.has_sink_layers());
1167 assert!(g.cancel().is_none() && g.usage_meter().is_none() && g.resilience().is_none());
1168 assert!(g.schema_drift().is_none() && g.dlq().is_none());
1169 }
1170
1171 struct Rows(Vec<Value>);
1172
1173 #[async_trait]
1174 impl Source for Rows {
1175 async fn fetch_with_context(
1176 &self,
1177 _ctx: &HashMap<String, Value>,
1178 ) -> Result<Vec<Value>, FaucetError> {
1179 Ok(self.0.clone())
1180 }
1181 }
1182
1183 #[tokio::test]
1186 async fn a_pipeline_refuses_an_unbuilt_destination() {
1187 let g = Governance::new().with_budget(BudgetSpec {
1188 max_records: Some(5),
1189 ..Default::default()
1190 });
1191 let source = Rows(vec![json!({"a": 1})]);
1192 let raw = Dest(Arc::default());
1193 let err = crate::Pipeline::new(&source, &raw)
1194 .with_governance(g.clone())
1195 .run()
1196 .await
1197 .unwrap_err();
1198 assert!(
1199 err.to_string().contains("Governance::sink_builder"),
1200 "{err}"
1201 );
1202 let governed = g.sink_builder().build(Box::new(raw));
1203 let r = crate::Pipeline::new(&source, &governed)
1204 .with_governance(g)
1205 .run()
1206 .await
1207 .unwrap();
1208 assert_eq!(r.records_written, 1);
1209 }
1210
1211 struct GovernedOnly {
1213 governed: std::sync::atomic::AtomicUsize,
1214 rows: Seen,
1215 }
1216
1217 #[async_trait]
1218 impl Sink for GovernedOnly {
1219 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1220 for r in records {
1221 self.rows.lock().unwrap().push(("plain", r.clone()));
1222 }
1223 Ok(records.len())
1224 }
1225 async fn write_governed(
1226 &self,
1227 page: &GovernedPage<'_>,
1228 ) -> Result<crate::accounting::WriteReport, FaucetError> {
1229 self.governed
1230 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1231 self.write_batch_reported(page.records()).await
1232 }
1233 }
1234
1235 #[tokio::test]
1237 async fn the_plain_write_path_writes_through_write_governed() {
1238 let source = Rows(vec![json!({"a": 1}), json!({"a": 2})]);
1239 let sink = GovernedOnly {
1240 governed: Default::default(),
1241 rows: Arc::default(),
1242 };
1243 let r = crate::Pipeline::new(&source, &sink).run().await.unwrap();
1244 assert_eq!(r.records_written, 2);
1245 assert_eq!(sink.governed.load(std::sync::atomic::Ordering::Relaxed), 1);
1246 assert_eq!(sink.rows.lock().unwrap().len(), 2);
1247 }
1248
1249 #[tokio::test]
1252 async fn a_wrapper_changing_write_batch_keeps_the_governed_write_its_own() {
1253 let seen: Seen = Arc::default();
1254 let wrapped = Tap {
1255 name: "tap",
1256 inner: Box::new(Dest(Arc::clone(&seen))),
1257 seen: Arc::clone(&seen),
1258 };
1259 let page = Governance::new()
1260 .govern(vec![json!({"k": 1})], &wrapped)
1261 .await
1262 .unwrap();
1263 assert!(page.envelopes.is_empty());
1264 let report = wrapped.write_governed(&page.page).await.unwrap();
1265 assert_eq!(report.accepted, 1);
1266 let names: Vec<&str> = seen.lock().unwrap().iter().map(|(n, _)| *n).collect();
1267 assert_eq!(names, ["tap", "dest"]);
1268 }
1269
1270 #[cfg(feature = "masking")]
1271 #[tokio::test]
1272 async fn govern_runs_the_passes_and_returns_a_governed_page() {
1273 let spec: crate::masking::MaskingSpec = serde_json::from_value(json!({
1274 "rules": [{"match": {"fields": ["email"]}, "action": {"type": "redact"}}]
1275 }))
1276 .unwrap();
1277 let g = Governance::new().with_masking(Arc::new(
1278 crate::masking::CompiledMasking::compile(&spec).unwrap(),
1279 ));
1280 let out = g
1281 .govern(vec![json!({"email": "a@b.co"})], &Dest(Arc::default()))
1282 .await
1283 .unwrap();
1284 assert_eq!(out.page.records()[0]["email"], "***");
1285 }
1286
1287 #[cfg(feature = "quality")]
1288 #[tokio::test]
1289 async fn govern_fails_on_an_aborting_pass() {
1290 let spec: crate::quality::QualitySpec = serde_json::from_value(json!({
1291 "record": [{"type": "not_null", "field": "id", "on_failure": "abort"}]
1292 }))
1293 .unwrap();
1294 let g = Governance::new().with_quality(Arc::new(
1295 crate::quality::CompiledQuality::compile(&spec).unwrap(),
1296 ));
1297 let err = g
1298 .govern(vec![json!({"id": null})], &Dest(Arc::default()))
1299 .await
1300 .unwrap_err();
1301 assert!(matches!(err, FaucetError::QualityFailure { .. }), "{err}");
1302 }
1303
1304 #[test]
1305 fn compat_features_follow_the_sink_layers() {
1306 use crate::compat::Feature as F;
1307 let (meta, ctx) = metadata();
1308 let g = Governance::new()
1309 .with_budget(BudgetSpec {
1310 max_bytes: Some(1),
1311 ..Default::default()
1312 })
1313 .with_metadata(meta, ctx)
1314 .with_profiler(profiler());
1315 let f = g.compat_features();
1316 assert!(f.contains(&F::Budget) && f.contains(&F::MetadataColumns));
1317 assert!(f.contains(&F::Profiling));
1318 assert!(!f.contains(&F::PolicyQuarantine));
1319 assert!(!g.policy_quarantines());
1320 assert!(!crate::compat::path_allowed(
1321 crate::compat::WritePath::Native,
1322 &f
1323 ));
1324 assert!(!crate::compat::path_allowed(
1325 crate::compat::WritePath::Columnar,
1326 &f
1327 ));
1328 assert!(crate::compat::path_allowed(
1329 crate::compat::WritePath::Row,
1330 &f
1331 ));
1332 }
1333
1334 #[cfg(feature = "masking")]
1338 #[tokio::test]
1339 async fn the_governed_dlq_masks_and_meters_every_write_kind() {
1340 let spec: crate::masking::MaskingSpec = serde_json::from_value(json!({
1341 "rules": [{"match": {"fields": ["email"]}, "action": {"type": "redact"}}]
1342 }))
1343 .unwrap();
1344 let masking = Arc::new(crate::masking::CompiledMasking::compile(&spec).unwrap());
1345 let seen: Seen = Arc::default();
1346 let meter = Arc::new(UsageMeter::new());
1347 let dlq = DlqGovernance::new(DlqConfig::new(Arc::new(Dest(Arc::clone(&seen)))))
1348 .with_masking(masking);
1349 let config = Governance::new()
1350 .with_usage_meter(Arc::clone(&meter))
1351 .with_dlq(dlq)
1352 .dlq_config()
1353 .unwrap();
1354 let envelope = || json!({"payload": {"email": "a@b.co"}});
1355 let report = config
1356 .sink
1357 .write_batch_reported(&[envelope()])
1358 .await
1359 .unwrap();
1360 assert_eq!(report.accepted, 1);
1361 let outcomes = config
1362 .sink
1363 .write_batch_partial(&[envelope(), envelope()])
1364 .await
1365 .unwrap();
1366 assert!(outcomes.iter().all(|o| o.is_ok()));
1367 let seen = seen.lock().unwrap().clone();
1368 assert_eq!(seen.len(), 3);
1369 assert!(seen.iter().all(|(_, r)| r["payload"]["email"] == "***"));
1370 assert_eq!(meter.snapshot().dlq_records_written, 3);
1371 }
1372
1373 #[cfg(all(feature = "quality", feature = "contract", feature = "policy"))]
1375 #[test]
1376 fn the_record_passes_can_be_opted_out() {
1377 let quality: crate::quality::QualitySpec = serde_json::from_value(json!({})).unwrap();
1378 let contract: crate::contract::ContractSpec = serde_json::from_value(json!({
1379 "version": "1", "fields": [{"name": "id", "type": "integer"}]
1380 }))
1381 .unwrap();
1382 let policy: crate::policy::PolicySpec = serde_json::from_value(json!({
1383 "classifications": [], "rules": []
1384 }))
1385 .unwrap();
1386 let g = Governance::new()
1387 .with_quality(Arc::new(
1388 crate::quality::CompiledQuality::compile(&quality).unwrap(),
1389 ))
1390 .with_contract(Arc::new(
1391 crate::contract::CompiledContract::compile(&contract).unwrap(),
1392 ))
1393 .with_policy(PolicyGovernance::new(
1394 Arc::new(crate::policy::CompiledPolicy::compile(&policy).unwrap()),
1395 crate::policy::SinkFacts {
1396 id: "w".into(),
1397 kind: "stdout".into(),
1398 attributes: Default::default(),
1399 },
1400 crate::policy::PolicyScope {
1401 pipeline: "p".into(),
1402 row: "w".into(),
1403 },
1404 ));
1405 assert!(g.quality().is_some() && g.contract().is_some() && g.policy().is_some());
1406 let g = g
1407 .without(Pass::Masking, "m")
1408 .without(Pass::Quality, "q")
1409 .without(Pass::Contract, "c")
1410 .without(Pass::Policy, "p");
1411 assert!(g.quality().is_none() && g.contract().is_none() && g.policy().is_none());
1412 assert_eq!(g.opt_outs().len(), 4);
1413 }
1414
1415 #[cfg(feature = "masking")]
1416 #[tokio::test]
1417 async fn dlq_only_masking_rewrites_the_envelope_payload() {
1418 let spec: crate::masking::MaskingSpec = serde_json::from_value(json!({
1419 "rules": [{"match": {"fields": ["email"]}, "action": {"type": "redact"}}]
1420 }))
1421 .unwrap();
1422 let masking = Arc::new(crate::masking::CompiledMasking::compile(&spec).unwrap());
1423 let seen: Seen = Arc::default();
1424 let dlq = DlqGovernance::new(DlqConfig::new(Arc::new(Dest(Arc::clone(&seen)))))
1425 .with_masking(Arc::clone(&masking));
1426 let config = Governance::new().with_dlq(dlq).dlq_config().unwrap();
1427 let envelope = json!({"payload": {"email": "a@b.co"}, "reason": "partial"});
1428 config
1429 .sink
1430 .write_batch(&[envelope, json!({"reason": "no payload"})])
1431 .await
1432 .unwrap();
1433 let seen = seen.lock().unwrap().clone();
1434 assert_eq!(seen[0].1["payload"]["email"], "***");
1435 assert!(seen[1].1.get("payload").is_none());
1436 let elsewhere: crate::masking::MaskingSpec = serde_json::from_value(json!({
1437 "rules": [{"match": {"fields": ["x"]}, "action": {"type": "redact"},
1438 "applies_to": ["elsewhere"]}]
1439 }))
1440 .unwrap();
1441 let empty = DlqGovernance::new(DlqConfig::new(Arc::new(Dest(Arc::default()))))
1442 .with_masking(Arc::new(
1443 crate::masking::CompiledMasking::compile_for_sink(&elsewhere, &["dlq"]).unwrap(),
1444 ));
1445 assert!(empty.masking.is_none(), "an empty policy is ignored");
1446 assert!(Governance::new().dlq_config().is_none());
1447 }
1448}