pub(crate) mod apply;
mod page;
pub(crate) mod write;
#[cfg(feature = "arrow")]
pub(crate) use page::GovernedBatch;
pub(crate) use page::{Clearance, GovernedNative};
pub use page::{GovernedEnvelopes, GovernedPage};
use crate::budget::{BudgetSink, BudgetSpec, BudgetState, BudgetTimer};
use crate::dlq::DlqConfig;
use crate::error::FaucetError;
use crate::metadata::{CompiledMetadata, MetadataContext, MetadataSink};
use crate::profiling::{Profiler, ProfilingSink};
use crate::traits::{RowOutcome, Sink};
use crate::usage::UsageMeter;
use async_trait::async_trait;
use serde::Serialize;
use serde_json::Value;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tokio_util::sync::CancellationToken;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum Pass {
Masking,
Quality,
Contract,
SchemaDrift,
Policy,
Budget,
Profiling,
Metadata,
Metering,
Resilience,
Cancel,
Dlq,
}
impl Pass {
pub const ALL: &'static [Pass] = &[
Pass::Masking,
Pass::Quality,
Pass::Contract,
Pass::SchemaDrift,
Pass::Policy,
Pass::Budget,
Pass::Profiling,
Pass::Metadata,
Pass::Metering,
Pass::Resilience,
Pass::Cancel,
Pass::Dlq,
];
pub fn as_str(self) -> &'static str {
match self {
Pass::Masking => "masking",
Pass::Quality => "quality",
Pass::Contract => "contract",
Pass::SchemaDrift => "schema_drift",
Pass::Policy => "policy",
Pass::Budget => "budget",
Pass::Profiling => "profiling",
Pass::Metadata => "metadata",
Pass::Metering => "metering",
Pass::Resilience => "resilience",
Pass::Cancel => "cancel",
Pass::Dlq => "dlq",
}
}
}
impl std::fmt::Display for Pass {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum SinkLayer {
Policy,
Profiling,
OuterTap,
Metadata,
InnerTap,
Budget,
}
pub const SINK_LAYERS: &[SinkLayer] = &[
SinkLayer::Policy,
SinkLayer::Profiling,
SinkLayer::OuterTap,
SinkLayer::Metadata,
SinkLayer::InnerTap,
SinkLayer::Budget,
];
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct OptOut {
pub pass: Pass,
pub reason: String,
}
#[cfg(feature = "policy")]
#[derive(Clone)]
#[non_exhaustive]
pub struct PolicyGovernance {
pub policy: Arc<crate::policy::CompiledPolicy>,
pub facts: crate::policy::SinkFacts,
pub scope: crate::policy::PolicyScope,
pub masking: Option<Arc<crate::masking::CompiledMasking>>,
}
#[cfg(feature = "policy")]
impl PolicyGovernance {
pub fn new(
policy: Arc<crate::policy::CompiledPolicy>,
facts: crate::policy::SinkFacts,
scope: crate::policy::PolicyScope,
) -> Self {
Self {
policy,
facts,
scope,
masking: None,
}
}
pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
self.masking = Some(masking);
self
}
pub fn quarantines(&self) -> bool {
self.policy.quarantines()
}
}
#[derive(Clone)]
#[non_exhaustive]
pub struct DlqGovernance {
pub config: DlqConfig,
#[cfg(feature = "policy")]
pub facts: Option<crate::policy::SinkFacts>,
#[cfg(feature = "masking")]
pub masking: Option<Arc<crate::masking::CompiledMasking>>,
}
impl DlqGovernance {
pub fn new(config: DlqConfig) -> Self {
Self {
config,
#[cfg(feature = "policy")]
facts: None,
#[cfg(feature = "masking")]
masking: None,
}
}
#[cfg(feature = "policy")]
pub fn with_facts(mut self, facts: crate::policy::SinkFacts) -> Self {
self.facts = Some(facts);
self
}
#[cfg(feature = "masking")]
pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
if !masking.is_empty() {
self.masking = Some(masking);
}
self
}
}
#[derive(Clone, Default)]
#[non_exhaustive]
pub struct Governance {
#[cfg(feature = "masking")]
masking: Option<Arc<crate::masking::CompiledMasking>>,
#[cfg(feature = "quality")]
quality: Option<Arc<crate::quality::CompiledQuality>>,
#[cfg(feature = "contract")]
contract: Option<Arc<crate::contract::CompiledContract>>,
schema_drift: Option<crate::drift::SchemaDriftPolicy>,
#[cfg(feature = "policy")]
policy: Option<PolicyGovernance>,
budget: Option<BudgetSpec>,
budget_state: Option<Arc<BudgetState>>,
profiler: Option<Arc<Mutex<Profiler>>>,
metadata: Option<(CompiledMetadata, MetadataContext)>,
usage_meter: Option<Arc<UsageMeter>>,
resilience: Option<crate::resilience::ResiliencePolicy>,
cancel: Option<CancellationToken>,
dlq: Option<DlqGovernance>,
opt_outs: Vec<OptOut>,
sink_built: Arc<AtomicBool>,
}
impl std::fmt::Debug for Governance {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Governance")
.field("active", &self.active())
.field("opt_outs", &self.opt_outs)
.finish()
}
}
impl Governance {
pub fn new() -> Self {
Self::default()
}
#[cfg(feature = "masking")]
pub fn with_masking(mut self, masking: Arc<crate::masking::CompiledMasking>) -> Self {
self.masking = (!masking.is_empty()).then_some(masking);
self
}
#[cfg(feature = "quality")]
pub fn with_quality(mut self, quality: Arc<crate::quality::CompiledQuality>) -> Self {
self.quality = Some(quality);
self
}
#[cfg(feature = "contract")]
pub fn with_contract(mut self, contract: Arc<crate::contract::CompiledContract>) -> Self {
self.contract = Some(contract);
self
}
pub fn with_schema_drift(mut self, policy: crate::drift::SchemaDriftPolicy) -> Self {
self.schema_drift = Some(policy);
self
}
#[cfg(feature = "policy")]
pub fn with_policy(mut self, policy: PolicyGovernance) -> Self {
self.policy = Some(policy);
self
}
pub fn with_budget(mut self, spec: BudgetSpec) -> Self {
if spec.max_records.is_none()
&& spec.max_bytes.is_none()
&& spec.max_duration_secs.is_none()
{
self.budget = None;
self.budget_state = None;
} else {
self.budget = Some(spec);
self.budget_state = Some(Arc::new(BudgetState::default()));
}
self
}
pub fn with_profiler(mut self, profiler: Arc<Mutex<Profiler>>) -> Self {
self.profiler = Some(profiler);
self
}
pub fn with_metadata(mut self, meta: CompiledMetadata, ctx: MetadataContext) -> Self {
self.metadata = Some((meta, ctx));
self
}
pub fn with_usage_meter(mut self, meter: Arc<UsageMeter>) -> Self {
self.usage_meter = Some(meter);
self
}
pub fn with_resilience(mut self, policy: crate::resilience::ResiliencePolicy) -> Self {
self.resilience = Some(policy);
self
}
pub fn with_cancel(mut self, cancel: CancellationToken) -> Self {
self.cancel = Some(cancel);
self
}
pub fn with_dlq(mut self, dlq: DlqGovernance) -> Self {
self.dlq = Some(dlq);
self
}
pub fn without(mut self, pass: Pass, reason: impl Into<String>) -> Self {
let reason = reason.into();
match pass {
#[cfg(feature = "masking")]
Pass::Masking => self.masking = None,
#[cfg(feature = "quality")]
Pass::Quality => self.quality = None,
#[cfg(feature = "contract")]
Pass::Contract => self.contract = None,
Pass::SchemaDrift => self.schema_drift = None,
#[cfg(feature = "policy")]
Pass::Policy => self.policy = None,
Pass::Budget => {
self.budget = None;
self.budget_state = None;
}
Pass::Profiling => self.profiler = None,
Pass::Metadata => self.metadata = None,
Pass::Metering => self.usage_meter = None,
Pass::Resilience => self.resilience = None,
Pass::Cancel => self.cancel = None,
Pass::Dlq => self.dlq = None,
#[allow(unreachable_patterns)]
_ => {}
}
tracing::info!(pass = pass.as_str(), reason = %reason, "governance pass opted out");
self.opt_outs.retain(|o| o.pass != pass);
self.opt_outs.push(OptOut { pass, reason });
self
}
pub fn opt_outs(&self) -> &[OptOut] {
&self.opt_outs
}
pub fn opted_out(&self, pass: Pass) -> Option<&str> {
self.opt_outs
.iter()
.find(|o| o.pass == pass)
.map(|o| o.reason.as_str())
}
pub fn is_active(&self, pass: Pass) -> bool {
match pass {
#[cfg(feature = "masking")]
Pass::Masking => self.masking.is_some(),
#[cfg(feature = "quality")]
Pass::Quality => self.quality.is_some(),
#[cfg(feature = "contract")]
Pass::Contract => self.contract.is_some(),
Pass::SchemaDrift => self.schema_drift.is_some(),
#[cfg(feature = "policy")]
Pass::Policy => self.policy.is_some(),
Pass::Budget => self.budget.is_some(),
Pass::Profiling => self.profiler.is_some(),
Pass::Metadata => self.metadata.is_some(),
Pass::Metering => self.usage_meter.is_some(),
Pass::Resilience => self.resilience.is_some(),
Pass::Cancel => self.cancel.is_some(),
Pass::Dlq => self.dlq.is_some(),
#[allow(unreachable_patterns)]
_ => false,
}
}
pub fn active(&self) -> Vec<Pass> {
Pass::ALL
.iter()
.copied()
.filter(|p| self.is_active(*p))
.collect()
}
pub fn budget(&self) -> Option<&BudgetSpec> {
self.budget.as_ref()
}
pub fn budget_state(&self) -> Option<&Arc<BudgetState>> {
self.budget_state.as_ref()
}
pub fn cancel(&self) -> Option<&CancellationToken> {
self.cancel.as_ref()
}
pub fn usage_meter(&self) -> Option<&Arc<UsageMeter>> {
self.usage_meter.as_ref()
}
pub fn resilience(&self) -> Option<&crate::resilience::ResiliencePolicy> {
self.resilience.as_ref()
}
pub fn dlq(&self) -> Option<&DlqGovernance> {
self.dlq.as_ref()
}
#[cfg(feature = "masking")]
pub fn masking(&self) -> Option<&Arc<crate::masking::CompiledMasking>> {
self.masking.as_ref()
}
#[cfg(feature = "quality")]
pub fn quality(&self) -> Option<&Arc<crate::quality::CompiledQuality>> {
self.quality.as_ref()
}
#[cfg(feature = "contract")]
pub fn contract(&self) -> Option<&Arc<crate::contract::CompiledContract>> {
self.contract.as_ref()
}
pub fn schema_drift(&self) -> Option<crate::drift::SchemaDriftPolicy> {
self.schema_drift
}
#[cfg(feature = "policy")]
pub fn policy(&self) -> Option<&PolicyGovernance> {
self.policy.as_ref()
}
pub fn policy_quarantines(&self) -> bool {
#[cfg(feature = "policy")]
{
self.policy
.as_ref()
.is_some_and(PolicyGovernance::quarantines)
}
#[cfg(not(feature = "policy"))]
{
false
}
}
pub fn has_sink_layers(&self) -> bool {
self.is_active(Pass::Budget)
|| self.is_active(Pass::Profiling)
|| self.is_active(Pass::Metadata)
|| self.policy_active()
}
fn policy_active(&self) -> bool {
#[cfg(feature = "policy")]
{
self.policy.is_some()
}
#[cfg(not(feature = "policy"))]
{
false
}
}
pub fn sink_built(&self) -> bool {
self.sink_built.load(Ordering::Relaxed)
}
pub fn sink_builder(&self) -> GovernedSinkBuilder<'_> {
GovernedSinkBuilder {
gov: self,
inner_tap: None,
outer_tap: None,
}
}
pub async fn govern<Si: Sink + ?Sized>(
&self,
records: Vec<Value>,
sink: &Si,
) -> Result<GovernedOutput, FaucetError> {
let specs = apply::GovernanceSpecs {
#[cfg(feature = "masking")]
masking: self.masking.as_ref(),
#[cfg(feature = "quality")]
quality: self.quality.as_ref(),
#[cfg(feature = "contract")]
contract: self.contract.as_ref(),
schema_drift: self.schema_drift.as_ref(),
pipeline: "unnamed",
row: "",
run_id: "",
sink_name: sink.connector_name(),
has_dlq: false,
};
let mut state = apply::GovernanceState::default();
let out = apply::apply_governance(records, &specs, &mut state, sink).await?;
Ok(GovernedOutput {
page: out.records,
envelopes: out.envelopes,
})
}
pub(crate) fn compat_features(&self) -> Vec<crate::compat::Feature> {
use crate::compat::Feature as F;
[
(self.policy_quarantines(), F::PolicyQuarantine),
(self.policy_active(), F::Policy),
(self.is_active(Pass::Budget), F::Budget),
(self.is_active(Pass::Profiling), F::Profiling),
(self.is_active(Pass::Metadata), F::MetadataColumns),
]
.into_iter()
.filter_map(|(on, f)| on.then_some(f))
.collect()
}
pub fn dlq_config(&self) -> Option<DlqConfig> {
self.governed_dlq()
}
pub(crate) fn governed_dlq(&self) -> Option<DlqConfig> {
let dlq = self.dlq.as_ref()?;
Some(govern_dlq(
dlq,
self.usage_meter.as_ref(),
self.budget.as_ref().zip(self.budget_state.as_ref()),
self.cancel.clone(),
#[cfg(feature = "policy")]
self.policy.as_ref(),
))
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct GovernedOutput {
pub page: GovernedPage<'static>,
pub envelopes: GovernedEnvelopes,
}
pub(crate) fn govern_dlq(
dlq: &DlqGovernance,
meter: Option<&Arc<UsageMeter>>,
budget: Option<(&BudgetSpec, &Arc<BudgetState>)>,
cancel: Option<CancellationToken>,
#[cfg(feature = "policy")] policy: Option<&PolicyGovernance>,
) -> DlqConfig {
let mut sink: Box<dyn Sink> = Box::new(SharedSink(Arc::clone(&dlq.config.sink)));
if let Some(m) = meter {
sink = Box::new(DlqMeteredSink {
inner: sink,
meter: Arc::clone(m),
});
}
if let Some((spec, state)) = budget {
sink = Box::new(BudgetSink::wrap_shared(
sink,
spec.clone(),
Arc::clone(state),
cancel.unwrap_or_default(),
));
}
#[cfg(feature = "policy")]
if let (Some(p), Some(facts)) = (policy, dlq.facts.as_ref()) {
let mut ps = crate::policy::PolicySink::for_dlq(
sink,
Arc::clone(&p.policy),
facts.clone(),
p.scope.clone(),
);
if let Some(m) = &p.masking {
ps = ps.with_masking(Arc::clone(m));
}
sink = Box::new(ps);
}
#[cfg(feature = "masking")]
if let Some(m) = &dlq.masking {
sink = Box::new(DlqMaskingSink {
inner: sink,
masking: Arc::clone(m),
});
}
DlqConfig {
sink: Arc::from(sink),
..dlq.config.clone()
}
}
type Tap = Box<dyn FnOnce(Box<dyn Sink>) -> Box<dyn Sink> + Send>;
pub struct GovernedSinkBuilder<'g> {
gov: &'g Governance,
inner_tap: Option<Tap>,
outer_tap: Option<Tap>,
}
impl<'g> GovernedSinkBuilder<'g> {
pub fn inner_tap(
mut self,
tap: impl FnOnce(Box<dyn Sink>) -> Box<dyn Sink> + Send + 'static,
) -> Self {
self.inner_tap = Some(Box::new(tap));
self
}
pub fn outer_tap(
mut self,
tap: impl FnOnce(Box<dyn Sink>) -> Box<dyn Sink> + Send + 'static,
) -> Self {
self.outer_tap = Some(Box::new(tap));
self
}
pub fn build(mut self, destination: Box<dyn Sink>) -> GovernedSink {
let gov = self.gov;
let mut layers = Vec::new();
let mut sink = destination;
let mut timer = None;
for layer in SINK_LAYERS.iter().rev() {
match layer {
SinkLayer::Budget => {
if let (Some(spec), Some(state)) = (&gov.budget, &gov.budget_state) {
let cancel = gov.cancel.clone().unwrap_or_default();
timer = Some(BudgetTimer::start_for(
spec.max_duration_secs,
Arc::clone(state),
cancel.clone(),
));
sink = Box::new(BudgetSink::wrap_shared(
sink,
spec.clone(),
Arc::clone(state),
cancel,
));
layers.push(*layer);
}
}
SinkLayer::InnerTap => {
if let Some(tap) = self.inner_tap.take() {
sink = tap(sink);
layers.push(*layer);
}
}
SinkLayer::Metadata => {
if let Some((meta, ctx)) = &gov.metadata {
sink = Box::new(MetadataSink::new(sink, meta.clone(), ctx.clone()));
layers.push(*layer);
}
}
SinkLayer::OuterTap => {
if let Some(tap) = self.outer_tap.take() {
sink = tap(sink);
layers.push(*layer);
}
}
SinkLayer::Profiling => {
if let Some(p) = &gov.profiler {
sink = Box::new(ProfilingSink::new(sink, Arc::clone(p)));
layers.push(*layer);
}
}
SinkLayer::Policy =>
{
#[cfg(feature = "policy")]
if let Some(p) = &gov.policy {
let mut ps = crate::policy::PolicySink::new(
sink,
Arc::clone(&p.policy),
p.facts.clone(),
p.scope.clone(),
);
if let Some(facts) = gov.dlq.as_ref().and_then(|d| d.facts.clone()) {
ps = ps.with_dlq_destination(facts);
}
if let Some(m) = &p.masking {
ps = ps.with_masking(Arc::clone(m));
}
sink = Box::new(ps);
layers.push(*layer);
}
}
}
}
layers.reverse();
gov.sink_built.store(true, Ordering::Relaxed);
GovernedSink {
sink,
budget: gov.budget_state.clone(),
_timer: timer,
layers,
}
}
}
pub struct GovernedSink {
sink: Box<dyn Sink>,
budget: Option<Arc<BudgetState>>,
_timer: Option<BudgetTimer>,
layers: Vec<SinkLayer>,
}
impl GovernedSink {
pub fn layers(&self) -> &[SinkLayer] {
&self.layers
}
pub fn budget_state(&self) -> Option<&Arc<BudgetState>> {
self.budget.as_ref()
}
pub fn finish<T>(&self, result: Result<T, FaucetError>) -> Result<T, FaucetError> {
match (&result, self.budget.as_ref().and_then(|s| s.verdict())) {
(Ok(_), Some(verdict)) => Err(verdict.error()),
_ => result,
}
}
pub fn into_inner(self) -> Box<dyn Sink> {
self.sink
}
}
#[async_trait]
impl Sink for GovernedSink {
crate::forward_sink!(sink);
}
pub(crate) struct SharedSink(pub(crate) Arc<dyn Sink>);
#[async_trait]
impl Sink for SharedSink {
crate::forward_sink!(0);
}
struct DlqMeteredSink {
inner: Box<dyn Sink>,
meter: Arc<UsageMeter>,
}
#[async_trait]
impl Sink for DlqMeteredSink {
crate::forward_sink!(inner; except write_batch, write_batch_reported, write_batch_partial);
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
let n = self.inner.write_batch(records).await?;
self.meter
.add_dlq_written(n as u64, crate::usage::estimate_page_bytes(records));
Ok(n)
}
async fn write_batch_reported(
&self,
records: &[Value],
) -> Result<crate::accounting::WriteReport, FaucetError> {
let r = self.inner.write_batch_reported(records).await?;
self.meter.add_dlq_written(
r.accepted as u64,
crate::usage::estimate_page_bytes(records),
);
Ok(r)
}
async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
let outcomes = self.inner.write_batch_partial(records).await?;
let ok = outcomes.iter().filter(|o| o.is_ok()).count();
self.meter
.add_dlq_written(ok as u64, crate::usage::estimate_page_bytes(records));
Ok(outcomes)
}
}
#[cfg(feature = "masking")]
struct DlqMaskingSink {
inner: Box<dyn Sink>,
masking: Arc<crate::masking::CompiledMasking>,
}
#[cfg(feature = "masking")]
impl DlqMaskingSink {
fn mask(&self, records: &[Value]) -> Vec<Value> {
let mut out = records.to_vec();
let payloads: Vec<Value> = out
.iter_mut()
.map(|e| match e.get_mut("payload") {
Some(p) => p.take(),
None => Value::Null,
})
.collect();
let masked = crate::masking::apply_masking(payloads, &self.masking).records;
for (envelope, payload) in out.iter_mut().zip(masked) {
if let Some(slot) = envelope.get_mut("payload") {
*slot = payload;
}
}
out
}
}
#[cfg(feature = "masking")]
#[async_trait]
impl Sink for DlqMaskingSink {
crate::forward_sink!(inner; except write_batch, write_batch_reported, write_batch_partial);
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
self.inner.write_batch(&self.mask(records)).await
}
async fn write_batch_reported(
&self,
records: &[Value],
) -> Result<crate::accounting::WriteReport, FaucetError> {
self.inner.write_batch_reported(&self.mask(records)).await
}
async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
self.inner.write_batch_partial(&self.mask(records)).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::traits::Source;
use serde_json::json;
use std::collections::HashMap;
type Seen = Arc<Mutex<Vec<(&'static str, Value)>>>;
struct Tap {
name: &'static str,
inner: Box<dyn Sink>,
seen: Seen,
}
#[async_trait]
impl Sink for Tap {
crate::forward_sink!(inner; except write_batch);
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
if let Some(r) = records.first() {
self.seen.lock().unwrap().push((self.name, r.clone()));
}
self.inner.write_batch(records).await
}
}
fn tap(
name: &'static str,
seen: &Seen,
) -> impl FnOnce(Box<dyn Sink>) -> Box<dyn Sink> + Send + 'static {
let seen = Arc::clone(seen);
move |inner| Box::new(Tap { name, inner, seen })
}
struct Dest(Seen);
#[async_trait]
impl Sink for Dest {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
for r in records {
self.0.lock().unwrap().push(("dest", r.clone()));
}
Ok(records.len())
}
}
fn profiler() -> Arc<Mutex<Profiler>> {
Arc::new(Mutex::new(Profiler::new(
serde_json::from_value(json!({})).unwrap(),
)))
}
fn metadata() -> (CompiledMetadata, MetadataContext) {
(
CompiledMetadata::compile(&crate::metadata::MetadataColumnsSpec::default())
.unwrap()
.unwrap(),
MetadataContext {
run_id: "r".into(),
source: "s".into(),
},
)
}
fn is_subsequence(layers: &[SinkLayer]) -> bool {
let mut canon = SINK_LAYERS.iter();
layers.iter().all(|l| canon.any(|c| c == l))
}
#[tokio::test]
async fn layers_wrap_the_destination_in_the_canonical_order() {
let (meta, ctx) = metadata();
let g = Governance::new()
.with_budget(BudgetSpec {
max_records: Some(10),
..Default::default()
})
.with_metadata(meta, ctx)
.with_profiler(profiler());
let seen: Seen = Arc::default();
let governed = g
.sink_builder()
.inner_tap(tap("inner", &seen))
.outer_tap(tap("outer", &seen))
.build(Box::new(Dest(Arc::clone(&seen))));
assert_eq!(
governed.layers(),
&[
SinkLayer::Profiling,
SinkLayer::OuterTap,
SinkLayer::Metadata,
SinkLayer::InnerTap,
SinkLayer::Budget,
]
);
assert!(is_subsequence(governed.layers()));
assert!(g.sink_built());
governed.write_batch(&[json!({"a": 1})]).await.unwrap();
let seen = seen.lock().unwrap().clone();
let names: Vec<&str> = seen.iter().map(|(n, _)| *n).collect();
assert_eq!(names, ["outer", "inner", "dest"]);
assert!(
seen[0].1.get("_faucet_run_id").is_none(),
"the outer tap sees the source record"
);
assert_eq!(
seen[1].1["_faucet_run_id"], "r",
"metadata is stamped above the inner tap"
);
assert_eq!(governed.budget_state().unwrap().records(), 1);
assert!(governed.finish(Ok::<_, FaucetError>(())).is_ok());
let _ = governed.into_inner();
}
#[tokio::test]
async fn an_empty_governance_builds_the_bare_destination() {
let g = Governance::new().with_budget(BudgetSpec::default());
assert!(!g.has_sink_layers());
assert!(g.budget().is_none());
let governed = g.sink_builder().build(Box::new(Dest(Arc::default())));
assert!(governed.layers().is_empty());
assert!(governed.budget_state().is_none());
assert_eq!(governed.write_batch(&[json!({})]).await.unwrap(), 1);
}
#[tokio::test]
async fn a_spent_duration_turns_the_partial_result_into_the_verdict() {
let cancel = CancellationToken::new();
let g = Governance::new()
.with_cancel(cancel.clone())
.with_budget(BudgetSpec {
max_duration_secs: Some(1),
..Default::default()
});
let governed = g.sink_builder().build(Box::new(Dest(Arc::default())));
tokio::time::timeout(std::time::Duration::from_secs(5), cancel.cancelled())
.await
.expect("the duration ceiling cancels the run");
let err = governed.finish(Ok::<_, FaucetError>(())).unwrap_err();
assert!(matches!(err, FaucetError::BudgetExceeded { .. }), "{err}");
}
#[test]
fn opt_outs_clear_the_pass_and_are_recorded() {
let (meta, ctx) = metadata();
let g = Governance::new()
.with_budget(BudgetSpec {
max_records: Some(1),
..Default::default()
})
.with_metadata(meta, ctx)
.with_profiler(profiler())
.with_cancel(CancellationToken::new())
.with_usage_meter(Arc::new(UsageMeter::new()))
.with_resilience(crate::resilience::ResiliencePolicy::default())
.with_schema_drift(crate::drift::SchemaDriftPolicy::compile(
&serde_json::from_value(json!({})).unwrap(),
));
for p in [
Pass::Budget,
Pass::Metadata,
Pass::Profiling,
Pass::Cancel,
Pass::Metering,
Pass::Resilience,
Pass::SchemaDrift,
] {
assert!(g.is_active(p), "{p}");
}
assert!(g.has_sink_layers());
let g = g
.without(Pass::Budget, "why")
.without(Pass::Profiling, "first")
.without(Pass::Profiling, "second")
.without(Pass::Dlq, "none");
assert!(!g.is_active(Pass::Budget));
assert!(g.budget_state().is_none());
assert!(!g.is_active(Pass::Profiling));
assert_eq!(g.opted_out(Pass::Profiling), Some("second"));
assert_eq!(g.opted_out(Pass::Budget), Some("why"));
assert_eq!(g.opted_out(Pass::Masking), None);
assert_eq!(g.opt_outs().len(), 3);
assert!(g.active().contains(&Pass::Metadata));
assert!(format!("{g:?}").contains("opt_outs"));
for p in Pass::ALL {
assert_eq!(p.to_string(), p.as_str());
}
let g = g
.without(Pass::Metadata, "m")
.without(Pass::Cancel, "c")
.without(Pass::Metering, "u")
.without(Pass::Resilience, "r")
.without(Pass::SchemaDrift, "d");
assert!(!g.has_sink_layers());
assert!(g.cancel().is_none() && g.usage_meter().is_none() && g.resilience().is_none());
assert!(g.schema_drift().is_none() && g.dlq().is_none());
}
struct Rows(Vec<Value>);
#[async_trait]
impl Source for Rows {
async fn fetch_with_context(
&self,
_ctx: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
Ok(self.0.clone())
}
}
#[tokio::test]
async fn a_pipeline_refuses_an_unbuilt_destination() {
let g = Governance::new().with_budget(BudgetSpec {
max_records: Some(5),
..Default::default()
});
let source = Rows(vec![json!({"a": 1})]);
let raw = Dest(Arc::default());
let err = crate::Pipeline::new(&source, &raw)
.with_governance(g.clone())
.run()
.await
.unwrap_err();
assert!(
err.to_string().contains("Governance::sink_builder"),
"{err}"
);
let governed = g.sink_builder().build(Box::new(raw));
let r = crate::Pipeline::new(&source, &governed)
.with_governance(g)
.run()
.await
.unwrap();
assert_eq!(r.records_written, 1);
}
struct GovernedOnly {
governed: std::sync::atomic::AtomicUsize,
rows: Seen,
}
#[async_trait]
impl Sink for GovernedOnly {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
for r in records {
self.rows.lock().unwrap().push(("plain", r.clone()));
}
Ok(records.len())
}
async fn write_governed(
&self,
page: &GovernedPage<'_>,
) -> Result<crate::accounting::WriteReport, FaucetError> {
self.governed
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.write_batch_reported(page.records()).await
}
}
#[tokio::test]
async fn the_plain_write_path_writes_through_write_governed() {
let source = Rows(vec![json!({"a": 1}), json!({"a": 2})]);
let sink = GovernedOnly {
governed: Default::default(),
rows: Arc::default(),
};
let r = crate::Pipeline::new(&source, &sink).run().await.unwrap();
assert_eq!(r.records_written, 2);
assert_eq!(sink.governed.load(std::sync::atomic::Ordering::Relaxed), 1);
assert_eq!(sink.rows.lock().unwrap().len(), 2);
}
#[tokio::test]
async fn a_wrapper_changing_write_batch_keeps_the_governed_write_its_own() {
let seen: Seen = Arc::default();
let wrapped = Tap {
name: "tap",
inner: Box::new(Dest(Arc::clone(&seen))),
seen: Arc::clone(&seen),
};
let page = Governance::new()
.govern(vec![json!({"k": 1})], &wrapped)
.await
.unwrap();
assert!(page.envelopes.is_empty());
let report = wrapped.write_governed(&page.page).await.unwrap();
assert_eq!(report.accepted, 1);
let names: Vec<&str> = seen.lock().unwrap().iter().map(|(n, _)| *n).collect();
assert_eq!(names, ["tap", "dest"]);
}
#[cfg(feature = "masking")]
#[tokio::test]
async fn govern_runs_the_passes_and_returns_a_governed_page() {
let spec: crate::masking::MaskingSpec = serde_json::from_value(json!({
"rules": [{"match": {"fields": ["email"]}, "action": {"type": "redact"}}]
}))
.unwrap();
let g = Governance::new().with_masking(Arc::new(
crate::masking::CompiledMasking::compile(&spec).unwrap(),
));
let out = g
.govern(vec![json!({"email": "a@b.co"})], &Dest(Arc::default()))
.await
.unwrap();
assert_eq!(out.page.records()[0]["email"], "***");
}
#[cfg(feature = "quality")]
#[tokio::test]
async fn govern_fails_on_an_aborting_pass() {
let spec: crate::quality::QualitySpec = serde_json::from_value(json!({
"record": [{"type": "not_null", "field": "id", "on_failure": "abort"}]
}))
.unwrap();
let g = Governance::new().with_quality(Arc::new(
crate::quality::CompiledQuality::compile(&spec).unwrap(),
));
let err = g
.govern(vec![json!({"id": null})], &Dest(Arc::default()))
.await
.unwrap_err();
assert!(matches!(err, FaucetError::QualityFailure { .. }), "{err}");
}
#[test]
fn compat_features_follow_the_sink_layers() {
use crate::compat::Feature as F;
let (meta, ctx) = metadata();
let g = Governance::new()
.with_budget(BudgetSpec {
max_bytes: Some(1),
..Default::default()
})
.with_metadata(meta, ctx)
.with_profiler(profiler());
let f = g.compat_features();
assert!(f.contains(&F::Budget) && f.contains(&F::MetadataColumns));
assert!(f.contains(&F::Profiling));
assert!(!f.contains(&F::PolicyQuarantine));
assert!(!g.policy_quarantines());
assert!(!crate::compat::path_allowed(
crate::compat::WritePath::Native,
&f
));
assert!(!crate::compat::path_allowed(
crate::compat::WritePath::Columnar,
&f
));
assert!(crate::compat::path_allowed(
crate::compat::WritePath::Row,
&f
));
}
#[cfg(feature = "masking")]
#[tokio::test]
async fn the_governed_dlq_masks_and_meters_every_write_kind() {
let spec: crate::masking::MaskingSpec = serde_json::from_value(json!({
"rules": [{"match": {"fields": ["email"]}, "action": {"type": "redact"}}]
}))
.unwrap();
let masking = Arc::new(crate::masking::CompiledMasking::compile(&spec).unwrap());
let seen: Seen = Arc::default();
let meter = Arc::new(UsageMeter::new());
let dlq = DlqGovernance::new(DlqConfig::new(Arc::new(Dest(Arc::clone(&seen)))))
.with_masking(masking);
let config = Governance::new()
.with_usage_meter(Arc::clone(&meter))
.with_dlq(dlq)
.dlq_config()
.unwrap();
let envelope = || json!({"payload": {"email": "a@b.co"}});
let report = config
.sink
.write_batch_reported(&[envelope()])
.await
.unwrap();
assert_eq!(report.accepted, 1);
let outcomes = config
.sink
.write_batch_partial(&[envelope(), envelope()])
.await
.unwrap();
assert!(outcomes.iter().all(|o| o.is_ok()));
let seen = seen.lock().unwrap().clone();
assert_eq!(seen.len(), 3);
assert!(seen.iter().all(|(_, r)| r["payload"]["email"] == "***"));
assert_eq!(meter.snapshot().dlq_records_written, 3);
}
#[cfg(all(feature = "quality", feature = "contract", feature = "policy"))]
#[test]
fn the_record_passes_can_be_opted_out() {
let quality: crate::quality::QualitySpec = serde_json::from_value(json!({})).unwrap();
let contract: crate::contract::ContractSpec = serde_json::from_value(json!({
"version": "1", "fields": [{"name": "id", "type": "integer"}]
}))
.unwrap();
let policy: crate::policy::PolicySpec = serde_json::from_value(json!({
"classifications": [], "rules": []
}))
.unwrap();
let g = Governance::new()
.with_quality(Arc::new(
crate::quality::CompiledQuality::compile(&quality).unwrap(),
))
.with_contract(Arc::new(
crate::contract::CompiledContract::compile(&contract).unwrap(),
))
.with_policy(PolicyGovernance::new(
Arc::new(crate::policy::CompiledPolicy::compile(&policy).unwrap()),
crate::policy::SinkFacts {
id: "w".into(),
kind: "stdout".into(),
attributes: Default::default(),
},
crate::policy::PolicyScope {
pipeline: "p".into(),
row: "w".into(),
},
));
assert!(g.quality().is_some() && g.contract().is_some() && g.policy().is_some());
let g = g
.without(Pass::Masking, "m")
.without(Pass::Quality, "q")
.without(Pass::Contract, "c")
.without(Pass::Policy, "p");
assert!(g.quality().is_none() && g.contract().is_none() && g.policy().is_none());
assert_eq!(g.opt_outs().len(), 4);
}
#[cfg(feature = "masking")]
#[tokio::test]
async fn dlq_only_masking_rewrites_the_envelope_payload() {
let spec: crate::masking::MaskingSpec = serde_json::from_value(json!({
"rules": [{"match": {"fields": ["email"]}, "action": {"type": "redact"}}]
}))
.unwrap();
let masking = Arc::new(crate::masking::CompiledMasking::compile(&spec).unwrap());
let seen: Seen = Arc::default();
let dlq = DlqGovernance::new(DlqConfig::new(Arc::new(Dest(Arc::clone(&seen)))))
.with_masking(Arc::clone(&masking));
let config = Governance::new().with_dlq(dlq).dlq_config().unwrap();
let envelope = json!({"payload": {"email": "a@b.co"}, "reason": "partial"});
config
.sink
.write_batch(&[envelope, json!({"reason": "no payload"})])
.await
.unwrap();
let seen = seen.lock().unwrap().clone();
assert_eq!(seen[0].1["payload"]["email"], "***");
assert!(seen[1].1.get("payload").is_none());
let elsewhere: crate::masking::MaskingSpec = serde_json::from_value(json!({
"rules": [{"match": {"fields": ["x"]}, "action": {"type": "redact"},
"applies_to": ["elsewhere"]}]
}))
.unwrap();
let empty = DlqGovernance::new(DlqConfig::new(Arc::new(Dest(Arc::default()))))
.with_masking(Arc::new(
crate::masking::CompiledMasking::compile_for_sink(&elsewhere, &["dlq"]).unwrap(),
));
assert!(empty.masking.is_none(), "an empty policy is ignored");
assert!(Governance::new().dlq_config().is_none());
}
}