use crate::error::FaucetError;
use crate::traits::Sink;
use crate::write_mode::KeyTuple;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::BTreeMap;
pub const DEFAULT_MAX_KEYS: usize = 100_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum CleanupMode {
DeleteMissing,
}
#[derive(Debug, Clone)]
pub struct CleanupPolicy {
pub scope: BTreeMap<String, Value>,
pub key: Vec<String>,
pub max_keys: usize,
}
impl CleanupPolicy {
pub fn new(
scope: BTreeMap<String, Value>,
key: Vec<String>,
max_keys: usize,
) -> Result<Self, FaucetError> {
if scope.is_empty() {
return Err(FaucetError::Config(
"cleanup: the completeness claim (`complete_for`) must name at least one \
column — an empty scope would match every row in the destination"
.into(),
));
}
if key.is_empty() {
return Err(FaucetError::Config(
"cleanup: requires a non-empty `key` so a written row can be told apart \
from a stale one"
.into(),
));
}
if scope.values().any(Value::is_null) {
return Err(FaucetError::Config(
"cleanup: the completeness claim contains a null value — an unresolved \
scope token would delete the wrong rows"
.into(),
));
}
Ok(Self {
scope,
key,
max_keys: max_keys.max(1),
})
}
}
#[derive(Debug, Default)]
pub struct SeenKeys {
keys: Vec<KeyTuple>,
overflowed: bool,
}
impl SeenKeys {
pub fn new() -> Self {
Self::default()
}
pub fn record_page(&mut self, page: &[Value], key: &[String], max_keys: usize) {
if self.overflowed {
return;
}
for rec in page {
let Some(obj) = rec.as_object() else { continue };
let mut tuple = Vec::with_capacity(key.len());
let mut complete = true;
for k in key {
match obj.get(k) {
Some(v) if !v.is_null() => tuple.push((k.clone(), v.clone())),
_ => {
complete = false;
break;
}
}
}
if !complete {
continue;
}
if self.keys.len() >= max_keys {
self.overflowed = true;
self.keys.clear(); return;
}
self.keys.push(KeyTuple(tuple));
}
}
pub fn overflowed(&self) -> bool {
self.overflowed
}
pub fn len(&self) -> usize {
self.keys.len()
}
pub fn is_empty(&self) -> bool {
self.keys.is_empty()
}
pub fn keys(&self) -> &[KeyTuple] {
&self.keys
}
pub fn overflow_error(&self, max_keys: usize) -> FaucetError {
FaucetError::Config(format!(
"cleanup: this invocation wrote more than {max_keys} rows in the claimed scope, \
so the set of written keys could not be tracked. Nothing was deleted — a \
partial delete would remove rows the run actually wrote. Narrow the scope \
(a smaller `complete_for`), or raise the ceiling if the destination can take \
a delete of this size"
))
}
}
pub struct CleanupTracker<'a, S: Sink + ?Sized> {
inner: &'a S,
key: Vec<String>,
max_keys: usize,
seen: std::sync::Mutex<SeenKeys>,
}
impl<'a, S: Sink + ?Sized> CleanupTracker<'a, S> {
pub fn new(inner: &'a S, policy: &CleanupPolicy) -> Self {
Self {
inner,
key: policy.key.clone(),
max_keys: policy.max_keys,
seen: std::sync::Mutex::new(SeenKeys::new()),
}
}
fn record(&self, records: &[Value]) {
if let Ok(mut seen) = self.seen.lock() {
seen.record_page(records, &self.key, self.max_keys);
}
}
#[allow(deprecated)] pub async fn finish(&self, policy: &CleanupPolicy) -> Result<u64, FaucetError> {
let seen = {
let mut guard = self
.seen
.lock()
.map_err(|_| FaucetError::Sink("cleanup: key tracker poisoned".into()))?;
if guard.overflowed() {
return Err(guard.overflow_error(policy.max_keys));
}
std::mem::take(&mut *guard)
};
self.inner.cleanup_scope(&policy.scope, &seen).await
}
pub fn overflowed(&self) -> bool {
self.seen.lock().map(|g| g.overflowed()).unwrap_or(true)
}
pub fn tracked(&self) -> usize {
self.seen.lock().map(|g| g.len()).unwrap_or(0)
}
}
#[async_trait::async_trait]
impl<S: Sink + ?Sized> Sink for CleanupTracker<'_, S> {
crate::forward_sink!(inner; except
write_batch, write_batch_reported, write_governed, write_batch_partial,
write_batch_idempotent, write_batch_idempotent_reported,
supports_columnar, write_batch_columnar, native_load_capabilities, load_native);
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
let n = self.inner.write_batch(records).await?;
self.record(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.record(records);
Ok(r)
}
async fn write_governed(
&self,
page: &crate::governance::GovernedPage<'_>,
) -> Result<crate::accounting::WriteReport, FaucetError> {
let r = self.inner.write_governed(page).await?;
self.record(page.records());
Ok(r)
}
async fn write_batch_partial(
&self,
records: &[Value],
) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
let out = self.inner.write_batch_partial(records).await;
self.record(records);
out
}
#[allow(deprecated)] async fn write_batch_idempotent(
&self,
records: &[Value],
scope: &str,
token: &str,
) -> Result<usize, FaucetError> {
let n = self
.inner
.write_batch_idempotent(records, scope, token)
.await?;
self.record(records);
Ok(n)
}
#[allow(deprecated)] async fn write_batch_idempotent_reported(
&self,
records: &[Value],
scope: &str,
token: &str,
) -> Result<crate::accounting::WriteReport, FaucetError> {
let r = self
.inner
.write_batch_idempotent_reported(records, scope, token)
.await?;
self.record(records);
Ok(r)
}
}
#[allow(deprecated)] #[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn scope() -> BTreeMap<String, Value> {
BTreeMap::from([("contact_id".to_string(), json!(123))])
}
#[test]
fn policy_requires_a_non_empty_scope() {
let err = CleanupPolicy::new(BTreeMap::new(), vec!["id".into()], 10)
.expect_err("empty scope must be refused");
assert!(err.to_string().contains("at least one"), "{err}");
}
#[test]
fn policy_requires_a_key() {
let err = CleanupPolicy::new(scope(), vec![], 10).expect_err("no key must be refused");
assert!(err.to_string().contains("`key`"), "{err}");
}
#[test]
fn policy_refuses_a_null_scope_value() {
let s = BTreeMap::from([("contact_id".to_string(), Value::Null)]);
let err = CleanupPolicy::new(s, vec!["id".into()], 10).expect_err("null must be refused");
assert!(err.to_string().contains("null"), "{err}");
}
#[test]
fn policy_floors_max_keys_at_one() {
let p = CleanupPolicy::new(scope(), vec!["id".into()], 0).unwrap();
assert_eq!(p.max_keys, 1);
}
#[tokio::test]
async fn cleanup_tracker_forwards_overwrite_lifecycle() {
struct OvwSink {
log: std::sync::Mutex<Vec<&'static str>>,
}
#[async_trait::async_trait]
impl Sink for OvwSink {
async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
Ok(r.len())
}
fn is_overwrite(&self) -> bool {
true
}
async fn begin_overwrite(&self) -> Result<(), FaucetError> {
self.log.lock().unwrap().push("begin");
Ok(())
}
async fn commit_overwrite(&self) -> Result<(), FaucetError> {
self.log.lock().unwrap().push("commit");
Ok(())
}
async fn abort_overwrite(&self) -> Result<(), FaucetError> {
self.log.lock().unwrap().push("abort");
Ok(())
}
}
let inner = OvwSink {
log: std::sync::Mutex::new(Vec::new()),
};
let policy = CleanupPolicy::new(scope(), vec!["id".into()], 10).unwrap();
let tracker = CleanupTracker::new(&inner, &policy);
assert!(tracker.is_overwrite());
assert_eq!(
tracker.batch_atomicity(),
crate::dlq::BatchAtomicity::BestEffort
);
assert_eq!(tracker.write_batch(&[json!({"id": 1})]).await.unwrap(), 1);
tracker.begin_overwrite().await.unwrap();
tracker.commit_overwrite().await.unwrap();
tracker.abort_overwrite().await.unwrap();
tracker.complete_run().await.unwrap();
assert_eq!(*inner.log.lock().unwrap(), vec!["begin", "commit", "abort"]);
}
#[test]
fn accumulates_keys_across_pages() {
let mut seen = SeenKeys::new();
let key = vec!["id".to_string()];
seen.record_page(&[json!({"id": 1}), json!({"id": 2})], &key, 100);
seen.record_page(&[json!({"id": 3})], &key, 100);
assert_eq!(seen.len(), 3);
assert!(!seen.overflowed());
}
#[test]
fn accumulates_composite_keys_in_declared_order() {
let mut seen = SeenKeys::new();
let key = vec!["a".to_string(), "b".to_string()];
seen.record_page(&[json!({"b": 2, "a": 1})], &key, 100);
assert_eq!(seen.len(), 1);
let t = &seen.keys()[0].0;
assert_eq!(
t[0].0, "a",
"key order follows the declared `key`, not the record"
);
assert_eq!(t[1].0, "b");
}
#[test]
fn skips_rows_with_a_missing_or_null_key() {
let mut seen = SeenKeys::new();
let key = vec!["id".to_string()];
seen.record_page(
&[
json!({"id": 1}),
json!({"other": 9}), json!({"id": null}), json!("not an object"),
],
&key,
100,
);
assert_eq!(seen.len(), 1, "only the well-keyed row is tracked");
}
#[test]
fn overflow_is_sticky_and_frees_the_buffer() {
let mut seen = SeenKeys::new();
let key = vec!["id".to_string()];
let page: Vec<Value> = (0..5).map(|i| json!({"id": i})).collect();
seen.record_page(&page, &key, 3);
assert!(seen.overflowed(), "ceiling of 3 must trip on a 5-row page");
assert!(seen.is_empty(), "buffer is freed — the cleanup will refuse");
seen.record_page(&[json!({"id": 99})], &key, 3);
assert!(seen.overflowed());
assert!(seen.is_empty());
}
#[test]
fn overflow_error_explains_that_nothing_was_deleted() {
let seen = SeenKeys::new();
let msg = seen.overflow_error(50).to_string();
assert!(msg.contains("Nothing was deleted"), "{msg}");
assert!(msg.contains("50"), "{msg}");
}
#[tokio::test]
async fn cleanup_tracker_forwards_every_hook() {
let probe = crate::forwarding::HookSink::default();
let policy = CleanupPolicy::new(scope(), vec!["id".into()], 10).unwrap();
let t = CleanupTracker::new(&probe, &policy);
crate::forwarding::assert_sink_forwards_every_hook(
&t,
&probe,
&[
"supports_columnar",
"write_batch_columnar",
"native_load_capabilities",
"load_native",
],
)
.await;
}
#[tokio::test]
async fn overflowed_tells_a_refusal_from_a_backend_failure() {
struct Fails;
#[async_trait::async_trait]
impl Sink for Fails {
async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
Ok(r.len())
}
fn supports_cleanup(&self) -> bool {
true
}
async fn cleanup_scope(
&self,
_: &BTreeMap<String, Value>,
_: &SeenKeys,
) -> Result<u64, FaucetError> {
Err(FaucetError::Sink("db down".into()))
}
}
let policy = CleanupPolicy::new(scope(), vec!["id".into()], 1).unwrap();
let t = CleanupTracker::new(&Fails, &policy);
t.write_batch(&[json!({"id": 1})]).await.unwrap();
assert!(t.finish(&policy).await.is_err());
assert!(!t.overflowed());
let t = CleanupTracker::new(&Fails, &policy);
t.write_batch(&[json!({"id": 1}), json!({"id": 2})])
.await
.unwrap();
assert!(t.finish(&policy).await.is_err());
assert!(t.overflowed());
}
}