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);
}
}
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 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> {
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_partial(
&self,
records: &[Value],
) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
let out = self.inner.write_batch_partial(records).await?;
self.record(records);
Ok(out)
}
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)
}
async fn flush(&self) -> Result<(), FaucetError> {
self.inner.flush().await
}
fn supports_cleanup(&self) -> bool {
self.inner.supports_cleanup()
}
async fn cleanup_scope(
&self,
scope: &BTreeMap<String, Value>,
seen: &SeenKeys,
) -> Result<u64, FaucetError> {
self.inner.cleanup_scope(scope, seen).await
}
fn supports_idempotent_writes(&self) -> bool {
self.inner.supports_idempotent_writes()
}
async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
self.inner.last_committed_token(scope).await
}
fn supported_write_modes(&self) -> &'static [crate::write_mode::WriteMode] {
self.inner.supported_write_modes()
}
fn dedups_by_key(&self) -> bool {
self.inner.dedups_by_key()
}
fn sink_guarantee(&self) -> crate::idempotency::SinkGuarantee {
self.inner.sink_guarantee()
}
async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
self.inner.current_schema().await
}
fn supports_schema_evolution(&self) -> bool {
self.inner.supports_schema_evolution()
}
async fn evolve_schema(
&self,
evolution: &crate::drift::SchemaEvolution,
) -> Result<(), FaucetError> {
self.inner.evolve_schema(evolution).await
}
fn config_schema(&self) -> Value {
self.inner.config_schema()
}
fn connector_name(&self) -> &'static str {
self.inner.connector_name()
}
fn dataset_uri(&self) -> String {
self.inner.dataset_uri()
}
fn is_overwrite(&self) -> bool {
self.inner.is_overwrite()
}
async fn begin_overwrite(&self) -> Result<(), FaucetError> {
self.inner.begin_overwrite().await
}
async fn commit_overwrite(&self) -> Result<(), FaucetError> {
self.inner.commit_overwrite().await
}
async fn abort_overwrite(&self) -> Result<(), FaucetError> {
self.inner.abort_overwrite().await
}
}
#[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.write_batch(&[json!({"id": 1})]).await.unwrap(), 1);
tracker.begin_overwrite().await.unwrap();
tracker.commit_overwrite().await.unwrap();
tracker.abort_overwrite().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}");
}
}