1use crate::error::FaucetError;
56use crate::traits::Sink;
57use crate::write_mode::KeyTuple;
58use serde::{Deserialize, Serialize};
59use serde_json::Value;
60use std::collections::BTreeMap;
61
62pub const DEFAULT_MAX_KEYS: usize = 100_000;
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
74#[serde(rename_all = "snake_case")]
75pub enum CleanupMode {
76 DeleteMissing,
78}
79
80#[derive(Debug, Clone)]
82pub struct CleanupPolicy {
83 pub scope: BTreeMap<String, Value>,
91 pub key: Vec<String>,
93 pub max_keys: usize,
95}
96
97impl CleanupPolicy {
98 pub fn new(
100 scope: BTreeMap<String, Value>,
101 key: Vec<String>,
102 max_keys: usize,
103 ) -> Result<Self, FaucetError> {
104 if scope.is_empty() {
105 return Err(FaucetError::Config(
109 "cleanup: the completeness claim (`complete_for`) must name at least one \
110 column — an empty scope would match every row in the destination"
111 .into(),
112 ));
113 }
114 if key.is_empty() {
115 return Err(FaucetError::Config(
116 "cleanup: requires a non-empty `key` so a written row can be told apart \
117 from a stale one"
118 .into(),
119 ));
120 }
121 if scope.values().any(Value::is_null) {
122 return Err(FaucetError::Config(
123 "cleanup: the completeness claim contains a null value — an unresolved \
124 scope token would delete the wrong rows"
125 .into(),
126 ));
127 }
128 Ok(Self {
129 scope,
130 key,
131 max_keys: max_keys.max(1),
132 })
133 }
134}
135
136#[derive(Debug, Default)]
141pub struct SeenKeys {
142 keys: Vec<KeyTuple>,
143 overflowed: bool,
146}
147
148impl SeenKeys {
149 pub fn new() -> Self {
150 Self::default()
151 }
152
153 pub fn record_page(&mut self, page: &[Value], key: &[String], max_keys: usize) {
158 if self.overflowed {
159 return;
160 }
161 for rec in page {
162 let Some(obj) = rec.as_object() else { continue };
163 let mut tuple = Vec::with_capacity(key.len());
164 let mut complete = true;
165 for k in key {
166 match obj.get(k) {
167 Some(v) if !v.is_null() => tuple.push((k.clone(), v.clone())),
168 _ => {
169 complete = false;
170 break;
171 }
172 }
173 }
174 if !complete {
175 continue;
176 }
177 if self.keys.len() >= max_keys {
178 self.overflowed = true;
179 self.keys.clear(); return;
181 }
182 self.keys.push(KeyTuple(tuple));
183 }
184 }
185
186 pub fn overflowed(&self) -> bool {
188 self.overflowed
189 }
190
191 pub fn len(&self) -> usize {
192 self.keys.len()
193 }
194
195 pub fn is_empty(&self) -> bool {
196 self.keys.is_empty()
197 }
198
199 pub fn keys(&self) -> &[KeyTuple] {
200 &self.keys
201 }
202
203 pub fn overflow_error(&self, max_keys: usize) -> FaucetError {
207 FaucetError::Config(format!(
208 "cleanup: this invocation wrote more than {max_keys} rows in the claimed scope, \
209 so the set of written keys could not be tracked. Nothing was deleted — a \
210 partial delete would remove rows the run actually wrote. Narrow the scope \
211 (a smaller `complete_for`), or raise the ceiling if the destination can take \
212 a delete of this size"
213 ))
214 }
215}
216
217pub struct CleanupTracker<'a, S: Sink + ?Sized> {
237 inner: &'a S,
238 key: Vec<String>,
239 max_keys: usize,
240 seen: std::sync::Mutex<SeenKeys>,
241}
242
243impl<'a, S: Sink + ?Sized> CleanupTracker<'a, S> {
244 pub fn new(inner: &'a S, policy: &CleanupPolicy) -> Self {
245 Self {
246 inner,
247 key: policy.key.clone(),
248 max_keys: policy.max_keys,
249 seen: std::sync::Mutex::new(SeenKeys::new()),
250 }
251 }
252
253 fn record(&self, records: &[Value]) {
254 if let Ok(mut seen) = self.seen.lock() {
255 seen.record_page(records, &self.key, self.max_keys);
256 }
257 }
258
259 #[allow(deprecated)] pub async fn finish(&self, policy: &CleanupPolicy) -> Result<u64, FaucetError> {
263 let seen = {
267 let mut guard = self
268 .seen
269 .lock()
270 .map_err(|_| FaucetError::Sink("cleanup: key tracker poisoned".into()))?;
271 if guard.overflowed() {
272 return Err(guard.overflow_error(policy.max_keys));
275 }
276 std::mem::take(&mut *guard)
277 };
278 self.inner.cleanup_scope(&policy.scope, &seen).await
279 }
280
281 pub fn overflowed(&self) -> bool {
284 self.seen.lock().map(|g| g.overflowed()).unwrap_or(true)
285 }
286
287 pub fn tracked(&self) -> usize {
289 self.seen.lock().map(|g| g.len()).unwrap_or(0)
290 }
291}
292
293#[async_trait::async_trait]
294impl<S: Sink + ?Sized> Sink for CleanupTracker<'_, S> {
295 crate::forward_sink!(inner; except
298 write_batch, write_batch_reported, write_governed, write_batch_partial,
299 write_batch_idempotent, write_batch_idempotent_reported,
300 supports_columnar, write_batch_columnar, native_load_capabilities, load_native);
301
302 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
303 let n = self.inner.write_batch(records).await?;
304 self.record(records);
305 Ok(n)
306 }
307
308 async fn write_batch_reported(
309 &self,
310 records: &[Value],
311 ) -> Result<crate::accounting::WriteReport, FaucetError> {
312 let r = self.inner.write_batch_reported(records).await?;
313 self.record(records);
314 Ok(r)
315 }
316
317 async fn write_governed(
318 &self,
319 page: &crate::governance::GovernedPage<'_>,
320 ) -> Result<crate::accounting::WriteReport, FaucetError> {
321 let r = self.inner.write_governed(page).await?;
322 self.record(page.records());
323 Ok(r)
324 }
325
326 async fn write_batch_partial(
327 &self,
328 records: &[Value],
329 ) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
330 let out = self.inner.write_batch_partial(records).await;
331 self.record(records);
335 out
336 }
337
338 #[allow(deprecated)] async fn write_batch_idempotent(
340 &self,
341 records: &[Value],
342 scope: &str,
343 token: &str,
344 ) -> Result<usize, FaucetError> {
345 let n = self
346 .inner
347 .write_batch_idempotent(records, scope, token)
348 .await?;
349 self.record(records);
350 Ok(n)
351 }
352
353 #[allow(deprecated)] async fn write_batch_idempotent_reported(
355 &self,
356 records: &[Value],
357 scope: &str,
358 token: &str,
359 ) -> Result<crate::accounting::WriteReport, FaucetError> {
360 let r = self
361 .inner
362 .write_batch_idempotent_reported(records, scope, token)
363 .await?;
364 self.record(records);
365 Ok(r)
366 }
367}
368
369#[allow(deprecated)] #[cfg(test)]
371mod tests {
372 use super::*;
373 use serde_json::json;
374
375 fn scope() -> BTreeMap<String, Value> {
376 BTreeMap::from([("contact_id".to_string(), json!(123))])
377 }
378
379 #[test]
380 fn policy_requires_a_non_empty_scope() {
381 let err = CleanupPolicy::new(BTreeMap::new(), vec!["id".into()], 10)
383 .expect_err("empty scope must be refused");
384 assert!(err.to_string().contains("at least one"), "{err}");
385 }
386
387 #[test]
388 fn policy_requires_a_key() {
389 let err = CleanupPolicy::new(scope(), vec![], 10).expect_err("no key must be refused");
390 assert!(err.to_string().contains("`key`"), "{err}");
391 }
392
393 #[test]
394 fn policy_refuses_a_null_scope_value() {
395 let s = BTreeMap::from([("contact_id".to_string(), Value::Null)]);
398 let err = CleanupPolicy::new(s, vec!["id".into()], 10).expect_err("null must be refused");
399 assert!(err.to_string().contains("null"), "{err}");
400 }
401
402 #[test]
403 fn policy_floors_max_keys_at_one() {
404 let p = CleanupPolicy::new(scope(), vec!["id".into()], 0).unwrap();
405 assert_eq!(p.max_keys, 1);
406 }
407
408 #[tokio::test]
409 async fn cleanup_tracker_forwards_overwrite_lifecycle() {
410 struct OvwSink {
413 log: std::sync::Mutex<Vec<&'static str>>,
414 }
415 #[async_trait::async_trait]
416 impl Sink for OvwSink {
417 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
418 Ok(r.len())
419 }
420 fn is_overwrite(&self) -> bool {
421 true
422 }
423 async fn begin_overwrite(&self) -> Result<(), FaucetError> {
424 self.log.lock().unwrap().push("begin");
425 Ok(())
426 }
427 async fn commit_overwrite(&self) -> Result<(), FaucetError> {
428 self.log.lock().unwrap().push("commit");
429 Ok(())
430 }
431 async fn abort_overwrite(&self) -> Result<(), FaucetError> {
432 self.log.lock().unwrap().push("abort");
433 Ok(())
434 }
435 }
436 let inner = OvwSink {
437 log: std::sync::Mutex::new(Vec::new()),
438 };
439 let policy = CleanupPolicy::new(scope(), vec!["id".into()], 10).unwrap();
440 let tracker = CleanupTracker::new(&inner, &policy);
441 assert!(tracker.is_overwrite());
442 assert_eq!(
443 tracker.batch_atomicity(),
444 crate::dlq::BatchAtomicity::BestEffort
445 );
446 assert_eq!(tracker.write_batch(&[json!({"id": 1})]).await.unwrap(), 1);
448 tracker.begin_overwrite().await.unwrap();
449 tracker.commit_overwrite().await.unwrap();
450 tracker.abort_overwrite().await.unwrap();
451 tracker.complete_run().await.unwrap();
452 assert_eq!(*inner.log.lock().unwrap(), vec!["begin", "commit", "abort"]);
453 }
454
455 #[test]
456 fn accumulates_keys_across_pages() {
457 let mut seen = SeenKeys::new();
458 let key = vec!["id".to_string()];
459 seen.record_page(&[json!({"id": 1}), json!({"id": 2})], &key, 100);
460 seen.record_page(&[json!({"id": 3})], &key, 100);
461 assert_eq!(seen.len(), 3);
462 assert!(!seen.overflowed());
463 }
464
465 #[test]
466 fn accumulates_composite_keys_in_declared_order() {
467 let mut seen = SeenKeys::new();
468 let key = vec!["a".to_string(), "b".to_string()];
469 seen.record_page(&[json!({"b": 2, "a": 1})], &key, 100);
470 assert_eq!(seen.len(), 1);
471 let t = &seen.keys()[0].0;
472 assert_eq!(
473 t[0].0, "a",
474 "key order follows the declared `key`, not the record"
475 );
476 assert_eq!(t[1].0, "b");
477 }
478
479 #[test]
480 fn skips_rows_with_a_missing_or_null_key() {
481 let mut seen = SeenKeys::new();
482 let key = vec!["id".to_string()];
483 seen.record_page(
484 &[
485 json!({"id": 1}),
486 json!({"other": 9}), json!({"id": null}), json!("not an object"),
489 ],
490 &key,
491 100,
492 );
493 assert_eq!(seen.len(), 1, "only the well-keyed row is tracked");
494 }
495
496 #[test]
497 fn overflow_is_sticky_and_frees_the_buffer() {
498 let mut seen = SeenKeys::new();
499 let key = vec!["id".to_string()];
500 let page: Vec<Value> = (0..5).map(|i| json!({"id": i})).collect();
501 seen.record_page(&page, &key, 3);
502 assert!(seen.overflowed(), "ceiling of 3 must trip on a 5-row page");
503 assert!(seen.is_empty(), "buffer is freed — the cleanup will refuse");
504 seen.record_page(&[json!({"id": 99})], &key, 3);
506 assert!(seen.overflowed());
507 assert!(seen.is_empty());
508 }
509
510 #[test]
511 fn overflow_error_explains_that_nothing_was_deleted() {
512 let seen = SeenKeys::new();
513 let msg = seen.overflow_error(50).to_string();
514 assert!(msg.contains("Nothing was deleted"), "{msg}");
515 assert!(msg.contains("50"), "{msg}");
516 }
517
518 #[tokio::test]
519 async fn cleanup_tracker_forwards_every_hook() {
520 let probe = crate::forwarding::HookSink::default();
521 let policy = CleanupPolicy::new(scope(), vec!["id".into()], 10).unwrap();
522 let t = CleanupTracker::new(&probe, &policy);
523 crate::forwarding::assert_sink_forwards_every_hook(
524 &t,
525 &probe,
526 &[
527 "supports_columnar",
528 "write_batch_columnar",
529 "native_load_capabilities",
530 "load_native",
531 ],
532 )
533 .await;
534 }
535
536 #[tokio::test]
537 async fn overflowed_tells_a_refusal_from_a_backend_failure() {
538 struct Fails;
539 #[async_trait::async_trait]
540 impl Sink for Fails {
541 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
542 Ok(r.len())
543 }
544 fn supports_cleanup(&self) -> bool {
545 true
546 }
547 async fn cleanup_scope(
548 &self,
549 _: &BTreeMap<String, Value>,
550 _: &SeenKeys,
551 ) -> Result<u64, FaucetError> {
552 Err(FaucetError::Sink("db down".into()))
553 }
554 }
555 let policy = CleanupPolicy::new(scope(), vec!["id".into()], 1).unwrap();
556 let t = CleanupTracker::new(&Fails, &policy);
557 t.write_batch(&[json!({"id": 1})]).await.unwrap();
558 assert!(t.finish(&policy).await.is_err());
559 assert!(!t.overflowed());
560 let t = CleanupTracker::new(&Fails, &policy);
561 t.write_batch(&[json!({"id": 1}), json!({"id": 2})])
562 .await
563 .unwrap();
564 assert!(t.finish(&policy).await.is_err());
565 assert!(t.overflowed());
566 }
567}