Skip to main content

faucet_core/
cleanup.rs

1//! Scoped cleanup — delete destination rows a run did not write (#478).
2//!
3//! An incremental sync into an upsert sink **cannot remove records deleted at the
4//! source**: a record that disappears simply stops appearing in the
5//! "updated since X" feed, so `write_mode: upsert` keeps it in the destination
6//! forever. The destination looks healthy, the run reports success, and stale
7//! rows accumulate indefinitely with nothing surfacing the divergence.
8//!
9//! Scoped cleanup closes that hole for the case where the source can make a
10//! **completeness claim**: "for scope S, these are *all* the records". The
11//! canonical shape is a parent/child incremental sync — a child row fetching one
12//! contact's associations is authoritative for `contact_id = <that contact>`.
13//!
14//! ## Why the claim comes from the source
15//!
16//! A sink cannot make it. It observes a page of records and cannot distinguish a
17//! complete set from page 1 of 3, or from a partial page preceding a failure. Two
18//! further reasons the scope is declared upstream rather than on the sink:
19//!
20//! 1. **Drift safety.** The predicate already exists in the source config
21//!    (`/contacts/${contacts.id}/associations`). Declaring it a second time on
22//!    the sink means two copies of one predicate, and when they diverge the
23//!    pipeline deletes the wrong rows.
24//! 2. **The empty-result case, which is decisive.** A contact that had five
25//!    associations and now has none produces a fetch returning **zero records**.
26//!    Any design inferring scopes from observed records never learns the scope
27//!    existed, so the five stale rows survive — i.e. it fails precisely the case
28//!    the feature exists to fix. A scope declared by the invocation comes from
29//!    the parent record and survives an empty result set.
30//!
31//! ## When it runs
32//!
33//! **Once per invocation, after a fully successful run — never per page.** Two
34//! ways to get this wrong, both silent data loss:
35//!
36//! - *Per page*: a scope's records spanning two pages → page 2's
37//!   "delete what I didn't see" wipes what page 1 just wrote.
38//! - *After a partial run*: the fetch dies at 40% → the delete removes the 60%
39//!   that had not yet arrived.
40//!
41//! So [`run_stream`](crate::run_stream) only invokes it when the stream reached
42//! its natural end **uncancelled**, and the CLI additionally attaches the policy
43//! only for real root invocations (never `--dry-run` / `--limit` / a shard).
44//!
45//! ## Key accumulation and the ceiling
46//!
47//! Deleting "what this run did not write" requires knowing what it wrote. The
48//! pipeline accumulates the **key tuples only** (not records) as pages are
49//! written — see [`SeenKeys`]. That is bounded by the scope's row count, which is
50//! small for the per-parent shape this targets, so a ceiling
51//! ([`CleanupPolicy::max_keys`]) guards the pathological case: on breach the
52//! cleanup **aborts with a typed error and deletes nothing**, rather than issuing
53//! a partial delete that would destroy rows it simply forgot about.
54//!
55//! ## Complete reads
56//!
57//! A source that sometimes reads its **whole dataset** in one run (a change
58//! stream that fell past its retention and re-reads the table) declares it
59//! through [`ReadCompleteness`], returned by `Source::read_completeness`. A
60//! policy built with [`CleanupPolicy::for_complete_read`] has an empty scope
61//! (the whole destination) and fires only when that run's read finished.
62//! Keys are tracked only once the read has begun, and the run's bookmarks are
63//! held back until the delete succeeds, so a crash or a failed delete makes
64//! the next run read the whole dataset again instead of skipping the delete.
65
66use crate::error::FaucetError;
67use crate::traits::Sink;
68use crate::write_mode::KeyTuple;
69use serde::{Deserialize, Serialize};
70use serde_json::Value;
71use std::collections::BTreeMap;
72
73/// Default ceiling on accumulated keys for a [complete-read](CleanupPolicy::for_complete_read)
74/// cleanup: every key of the dataset is tracked, so it is sized for tables
75/// rather than per-parent scopes.
76pub const DEFAULT_COMPLETE_READ_MAX_KEYS: usize = 1_000_000;
77
78/// Default ceiling on accumulated keys for one invocation's cleanup.
79///
80/// Sized for the shape this feature targets (a scope per parent record — tens to
81/// low thousands of rows). Above it the cleanup refuses rather than guessing.
82pub const DEFAULT_MAX_KEYS: usize = 100_000;
83
84/// What to do about destination rows inside the claimed scope that this run did
85/// not write.
86///
87/// Opted into on the **source**, through its `complete_for.on_missing: delete`
88/// claim; no sink config field selects it.
89#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
90#[serde(rename_all = "snake_case")]
91pub enum CleanupMode {
92    /// Delete rows in the scope whose key was not written by this invocation.
93    DeleteMissing,
94}
95
96/// When a [`CleanupPolicy`] deletes.
97#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
98#[non_exhaustive]
99pub enum CleanupTrigger {
100    /// After every successful, uncancelled run.
101    #[default]
102    Always,
103    /// Only after a run in which the source read its whole dataset
104    /// ([`ReadCompleteness`] reached [`ReadState::Complete`]).
105    CompleteRead,
106}
107
108/// How far a source got through reading its whole dataset in the current run.
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110#[non_exhaustive]
111pub enum ReadState {
112    /// The run is not a complete read.
113    Partial,
114    /// A complete read has started and not finished.
115    Reading,
116    /// The run read every record of the dataset.
117    Complete,
118}
119
120/// A shared signal through which a source reports a complete read of its
121/// dataset. Cheap to clone; the source keeps one and hands out clones.
122#[derive(Debug, Clone, Default)]
123pub struct ReadCompleteness(std::sync::Arc<std::sync::atomic::AtomicU8>);
124
125impl ReadCompleteness {
126    const PARTIAL: u8 = 0;
127    const READING: u8 = 1;
128    const COMPLETE: u8 = 2;
129
130    /// A signal in [`ReadState::Partial`].
131    pub fn new() -> Self {
132        Self::default()
133    }
134
135    /// Back to [`ReadState::Partial`] (start of a run).
136    pub fn reset(&self) {
137        self.set(Self::PARTIAL);
138    }
139
140    /// A complete read has started. Call before its first record is yielded.
141    pub fn begin(&self) {
142        self.set(Self::READING);
143    }
144
145    /// The complete read has yielded every record.
146    pub fn complete(&self) {
147        self.set(Self::COMPLETE);
148    }
149
150    /// The current state.
151    pub fn state(&self) -> ReadState {
152        match self.0.load(std::sync::atomic::Ordering::Acquire) {
153            Self::READING => ReadState::Reading,
154            Self::COMPLETE => ReadState::Complete,
155            _ => ReadState::Partial,
156        }
157    }
158
159    fn set(&self, v: u8) {
160        self.0.store(v, std::sync::atomic::Ordering::Release);
161    }
162}
163
164/// A compiled cleanup instruction for one invocation.
165#[derive(Debug, Clone)]
166pub struct CleanupPolicy {
167    /// The completeness claim, in **destination column** terms: the rows this
168    /// invocation is authoritative for. Every entry is an equality predicate,
169    /// AND-ed together.
170    ///
171    /// Destination terms rather than source terms because the `DELETE` executes
172    /// against destination columns, and a transform chain may rename fields
173    /// between the two.
174    pub scope: BTreeMap<String, Value>,
175    /// Key columns identifying a row — mirrors the sink's `key`.
176    pub key: Vec<String>,
177    /// Ceiling on accumulated keys before the cleanup refuses (see module docs).
178    pub max_keys: usize,
179    /// When the delete fires.
180    pub trigger: CleanupTrigger,
181}
182
183impl CleanupPolicy {
184    /// Build a policy, validating it is actionable.
185    pub fn new(
186        scope: BTreeMap<String, Value>,
187        key: Vec<String>,
188        max_keys: usize,
189    ) -> Result<Self, FaucetError> {
190        if scope.is_empty() {
191            // An empty scope is an every-row predicate. Refusing here is the
192            // difference between "delete this contact's stale rows" and
193            // "truncate the table".
194            return Err(FaucetError::Config(
195                "cleanup: the completeness claim (`complete_for`) must name at least one \
196                 column — an empty scope would match every row in the destination"
197                    .into(),
198            ));
199        }
200        if key.is_empty() {
201            return Err(FaucetError::Config(
202                "cleanup: requires a non-empty `key` so a written row can be told apart \
203                 from a stale one"
204                    .into(),
205            ));
206        }
207        if scope.values().any(Value::is_null) {
208            return Err(FaucetError::Config(
209                "cleanup: the completeness claim contains a null value — an unresolved \
210                 scope token would delete the wrong rows"
211                    .into(),
212            ));
213        }
214        Ok(Self {
215            scope,
216            key,
217            max_keys: max_keys.max(1),
218            trigger: CleanupTrigger::Always,
219        })
220    }
221
222    /// A policy over the **whole destination** that fires only after a run in
223    /// which the source read its whole dataset (see the module docs).
224    pub fn for_complete_read(key: Vec<String>, max_keys: usize) -> Result<Self, FaucetError> {
225        if key.is_empty() {
226            return Err(FaucetError::Config(
227                "cleanup: a complete-read cleanup requires a non-empty `key` so a written row \
228                 can be told apart from a deleted one"
229                    .into(),
230            ));
231        }
232        Ok(Self {
233            scope: BTreeMap::new(),
234            key,
235            max_keys: max_keys.max(1),
236            trigger: CleanupTrigger::CompleteRead,
237        })
238    }
239}
240
241/// The set of key tuples an invocation wrote, accumulated across pages.
242///
243/// Deliberately stores only keys, not records: memory is O(rows in scope × key
244/// width) rather than O(payload).
245#[derive(Debug, Default)]
246pub struct SeenKeys {
247    keys: Vec<KeyTuple>,
248    /// Set once the ceiling is breached. Sticky: a cleanup that lost track of
249    /// even one key must not run at all.
250    overflowed: bool,
251}
252
253impl SeenKeys {
254    pub fn new() -> Self {
255        Self::default()
256    }
257
258    /// Record the keys written for one page. Rows missing a key column, or
259    /// carrying a null there, are ignored — they cannot be matched by a keyed
260    /// delete anyway, and the write path already routes them to the DLQ or fails
261    /// the batch.
262    pub fn record_page(&mut self, page: &[Value], key: &[String], max_keys: usize) {
263        if self.overflowed {
264            return;
265        }
266        for rec in page {
267            let Some(obj) = rec.as_object() else { continue };
268            let mut tuple = Vec::with_capacity(key.len());
269            let mut complete = true;
270            for k in key {
271                match obj.get(k) {
272                    Some(v) if !v.is_null() => tuple.push((k.clone(), v.clone())),
273                    _ => {
274                        complete = false;
275                        break;
276                    }
277                }
278            }
279            if !complete {
280                continue;
281            }
282            if self.keys.len() >= max_keys {
283                self.overflowed = true;
284                self.keys.clear(); // free the buffer; the cleanup will refuse
285                return;
286            }
287            self.keys.push(KeyTuple(tuple));
288        }
289    }
290
291    /// Whether the ceiling was breached, which makes the cleanup unsafe to run.
292    pub fn overflowed(&self) -> bool {
293        self.overflowed
294    }
295
296    pub fn len(&self) -> usize {
297        self.keys.len()
298    }
299
300    pub fn is_empty(&self) -> bool {
301        self.keys.is_empty()
302    }
303
304    pub fn keys(&self) -> &[KeyTuple] {
305        &self.keys
306    }
307
308    /// The typed error a breached ceiling produces. Separate from the accumulator
309    /// so the caller decides whether to fail the run or log — but never to
310    /// delete.
311    pub fn overflow_error(&self, max_keys: usize) -> FaucetError {
312        FaucetError::Config(format!(
313            "cleanup: this invocation wrote more than {max_keys} rows in the claimed scope, \
314             so the set of written keys could not be tracked. Nothing was deleted — a \
315             partial delete would remove rows the run actually wrote. Narrow the scope \
316             (a smaller `complete_for`), or raise the ceiling if the destination can take \
317             a delete of this size"
318        ))
319    }
320}
321
322/// A sink wrapper that records the key tuples written through it (#478).
323///
324/// This is how scoped cleanup tracks "what this run wrote" **without** adding a
325/// field to [`RunStreamOptions`](crate::RunStreamOptions), which is an
326/// externally-constructible struct whose shape is part of the public API. It is
327/// also the more honest home for the bookkeeping: the thing that writes the rows
328/// is the thing that knows which rows were written, and it composes with the
329/// existing sink-decorator pattern (`InstrumentedSink`) rather than threading a
330/// second concern through the page loop.
331///
332/// Every write path is counted, including [`write_batch_partial`](Sink::write_batch_partial).
333/// That is deliberate: a row handed to the sink that fails and lands in the DLQ
334/// is still a record the source claimed present, so it must count as seen or the
335/// cleanup would delete its destination row.
336///
337/// Records that never reach the sink at all — quarantined by a quality, contract,
338/// or drift policy — are consequently *not* counted, which is why those
339/// combinations are rejected at config-load time rather than silently deleting
340/// the quarantined rows' destination counterparts.
341pub struct CleanupTracker<'a, S: Sink + ?Sized> {
342    inner: &'a S,
343    key: Vec<String>,
344    max_keys: usize,
345    seen: std::sync::Mutex<SeenKeys>,
346    gate: Option<ReadCompleteness>,
347}
348
349impl<'a, S: Sink + ?Sized> CleanupTracker<'a, S> {
350    pub fn new(inner: &'a S, policy: &CleanupPolicy) -> Self {
351        Self {
352            inner,
353            key: policy.key.clone(),
354            max_keys: policy.max_keys,
355            seen: std::sync::Mutex::new(SeenKeys::new()),
356            gate: None,
357        }
358    }
359
360    /// Record keys only while `signal` reports a complete read under way or
361    /// finished, so an ordinary run tracks nothing.
362    #[must_use]
363    pub fn gated_by(mut self, signal: ReadCompleteness) -> Self {
364        self.gate = Some(signal);
365        self
366    }
367
368    fn record(&self, records: &[Value]) {
369        if self
370            .gate
371            .as_ref()
372            .is_some_and(|g| g.state() == ReadState::Partial)
373        {
374            return;
375        }
376        if let Ok(mut seen) = self.seen.lock() {
377            seen.record_page(records, &self.key, self.max_keys);
378        }
379    }
380
381    /// Run the scoped delete against the wrapped sink. Call **only** after a
382    /// fully successful, uncancelled run — see the module docs.
383    #[allow(deprecated)] // drives the 1.x lifecycle hook until faucet-core 2.0 (RFC 0014)
384    pub async fn finish(&self, policy: &CleanupPolicy) -> Result<u64, FaucetError> {
385        // Take the set out and drop the guard *before* awaiting: holding a
386        // `MutexGuard` across an await makes the whole run future non-`Send`,
387        // which the executor's `JoinSet` requires.
388        let seen = {
389            let mut guard = self
390                .seen
391                .lock()
392                .map_err(|_| FaucetError::Sink("cleanup: key tracker poisoned".into()))?;
393            if guard.overflowed() {
394                // Refuse rather than partially delete: the tracked set is
395                // incomplete, so a delete would remove rows the run wrote.
396                return Err(guard.overflow_error(policy.max_keys));
397            }
398            std::mem::take(&mut *guard)
399        };
400        self.inner.cleanup_scope(&policy.scope, &seen).await
401    }
402
403    /// Whether more keys were written than the policy's ceiling, so the
404    /// cleanup was (or will be) refused rather than run.
405    pub fn overflowed(&self) -> bool {
406        self.seen.lock().map(|g| g.overflowed()).unwrap_or(true)
407    }
408
409    /// Number of keys tracked so far (for logging).
410    pub fn tracked(&self) -> usize {
411        self.seen.lock().map(|g| g.len()).unwrap_or(0)
412    }
413}
414
415/// Holds every bookmark write back while a complete read is under way, so
416/// the run's position becomes durable only after its cleanup succeeded: a
417/// crash, a cancel or a failed delete leaves the previous bookmark, and the
418/// next run reads the whole dataset again instead of skipping the delete.
419pub(crate) struct HoldingStateStore {
420    inner: std::sync::Arc<dyn crate::state::StateStore>,
421    signal: ReadCompleteness,
422    held: std::sync::Mutex<BTreeMap<String, Value>>,
423}
424
425impl HoldingStateStore {
426    pub(crate) fn new(
427        inner: std::sync::Arc<dyn crate::state::StateStore>,
428        signal: ReadCompleteness,
429    ) -> Self {
430        Self {
431            inner,
432            signal,
433            held: std::sync::Mutex::new(BTreeMap::new()),
434        }
435    }
436
437    fn held(&self) -> std::sync::MutexGuard<'_, BTreeMap<String, Value>> {
438        self.held.lock().unwrap_or_else(|e| e.into_inner())
439    }
440
441    /// Write the held bookmarks through. Call only after the cleanup succeeded.
442    pub(crate) async fn release(&self) -> Result<(), FaucetError> {
443        let held = std::mem::take(&mut *self.held());
444        for (key, value) in held {
445            self.inner.put(&key, &value).await?;
446        }
447        Ok(())
448    }
449}
450
451#[async_trait::async_trait]
452impl crate::state::StateStore for HoldingStateStore {
453    async fn get(&self, key: &str) -> Result<Option<Value>, FaucetError> {
454        if let Some(v) = self.held().get(key) {
455            return Ok(Some(v.clone()));
456        }
457        self.inner.get(key).await
458    }
459
460    async fn put(&self, key: &str, value: &Value) -> Result<(), FaucetError> {
461        if self.signal.state() == ReadState::Partial {
462            return self.inner.put(key, value).await;
463        }
464        self.held().insert(key.to_string(), value.clone());
465        Ok(())
466    }
467
468    async fn delete(&self, key: &str) -> Result<(), FaucetError> {
469        self.held().remove(key);
470        self.inner.delete(key).await
471    }
472}
473
474#[async_trait::async_trait]
475impl<S: Sink + ?Sized> Sink for CleanupTracker<'_, S> {
476    // The columnar and native paths carry no keys to record, so they stay at the
477    // trait default (off); everything else reaches the wrapped sink.
478    crate::forward_sink!(inner; except
479        write_batch, write_batch_reported, write_governed, write_batch_partial,
480        write_batch_idempotent, write_batch_idempotent_reported,
481        supports_columnar, write_batch_columnar, native_load_capabilities, load_native);
482
483    async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
484        let n = self.inner.write_batch(records).await?;
485        self.record(records);
486        Ok(n)
487    }
488
489    async fn write_batch_reported(
490        &self,
491        records: &[Value],
492    ) -> Result<crate::accounting::WriteReport, FaucetError> {
493        let r = self.inner.write_batch_reported(records).await?;
494        self.record(records);
495        Ok(r)
496    }
497
498    async fn write_governed(
499        &self,
500        page: &crate::governance::GovernedPage<'_>,
501    ) -> Result<crate::accounting::WriteReport, FaucetError> {
502        let r = self.inner.write_governed(page).await?;
503        self.record(page.records());
504        Ok(r)
505    }
506
507    async fn write_batch_partial(
508        &self,
509        records: &[Value],
510    ) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
511        let out = self.inner.write_batch_partial(records).await;
512        // Count every row handed to the sink, including the ones that failed —
513        // per row or as a whole batch under `on_batch_error: dlq_all` — and will
514        // be routed to the DLQ. See the type docs.
515        self.record(records);
516        out
517    }
518
519    #[allow(deprecated)] // drives the 1.x lifecycle hook until faucet-core 2.0 (RFC 0014)
520    async fn write_batch_idempotent(
521        &self,
522        records: &[Value],
523        scope: &str,
524        token: &str,
525    ) -> Result<usize, FaucetError> {
526        let n = self
527            .inner
528            .write_batch_idempotent(records, scope, token)
529            .await?;
530        self.record(records);
531        Ok(n)
532    }
533
534    #[allow(deprecated)] // drives the 1.x lifecycle hook until faucet-core 2.0 (RFC 0014)
535    async fn write_batch_idempotent_reported(
536        &self,
537        records: &[Value],
538        scope: &str,
539        token: &str,
540    ) -> Result<crate::accounting::WriteReport, FaucetError> {
541        let r = self
542            .inner
543            .write_batch_idempotent_reported(records, scope, token)
544            .await?;
545        self.record(records);
546        Ok(r)
547    }
548}
549
550#[allow(deprecated)] // the tests drive the 1.x lifecycle hooks directly (RFC 0014)
551#[cfg(test)]
552mod tests {
553    use super::*;
554
555    #[tokio::test]
556    async fn the_holding_store_holds_writes_only_during_a_complete_read() {
557        use crate::state::{MemoryStateStore, StateStore};
558        let inner: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
559        let signal = ReadCompleteness::new();
560        let store = HoldingStateStore::new(std::sync::Arc::clone(&inner), signal.clone());
561        store.put("k", &serde_json::json!(1)).await.unwrap();
562        assert_eq!(inner.get("k").await.unwrap(), Some(serde_json::json!(1)));
563        signal.begin();
564        assert_eq!(signal.state(), ReadState::Reading);
565        store.put("k", &serde_json::json!(2)).await.unwrap();
566        assert_eq!(store.get("k").await.unwrap(), Some(serde_json::json!(2)));
567        assert_eq!(inner.get("k").await.unwrap(), Some(serde_json::json!(1)));
568        signal.complete();
569        assert_eq!(signal.state(), ReadState::Complete);
570        store.release().await.unwrap();
571        assert_eq!(inner.get("k").await.unwrap(), Some(serde_json::json!(2)));
572        store.put("gone", &serde_json::json!(3)).await.unwrap();
573        store.delete("gone").await.unwrap();
574        store.release().await.unwrap();
575        assert_eq!(inner.get("gone").await.unwrap(), None);
576        signal.reset();
577        assert_eq!(signal.state(), ReadState::Partial);
578    }
579    use serde_json::json;
580
581    fn scope() -> BTreeMap<String, Value> {
582        BTreeMap::from([("contact_id".to_string(), json!(123))])
583    }
584
585    #[test]
586    fn policy_requires_a_non_empty_scope() {
587        // An empty scope is a truncate, not a cleanup.
588        let err = CleanupPolicy::new(BTreeMap::new(), vec!["id".into()], 10)
589            .expect_err("empty scope must be refused");
590        assert!(err.to_string().contains("at least one"), "{err}");
591    }
592
593    #[test]
594    fn policy_requires_a_key() {
595        let err = CleanupPolicy::new(scope(), vec![], 10).expect_err("no key must be refused");
596        assert!(err.to_string().contains("`key`"), "{err}");
597    }
598
599    #[test]
600    fn policy_refuses_a_null_scope_value() {
601        // An unresolved `${parent.id}` would land here as null and delete the
602        // wrong rows.
603        let s = BTreeMap::from([("contact_id".to_string(), Value::Null)]);
604        let err = CleanupPolicy::new(s, vec!["id".into()], 10).expect_err("null must be refused");
605        assert!(err.to_string().contains("null"), "{err}");
606    }
607
608    #[test]
609    fn policy_floors_max_keys_at_one() {
610        let p = CleanupPolicy::new(scope(), vec!["id".into()], 0).unwrap();
611        assert_eq!(p.max_keys, 1);
612    }
613
614    #[tokio::test]
615    async fn cleanup_tracker_forwards_overwrite_lifecycle() {
616        // The tracker must forward the overwrite lifecycle to the inner sink so
617        // an overwrite run still stages/swaps when scoped cleanup wraps it.
618        struct OvwSink {
619            log: std::sync::Mutex<Vec<&'static str>>,
620        }
621        #[async_trait::async_trait]
622        impl Sink for OvwSink {
623            async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
624                Ok(r.len())
625            }
626            fn is_overwrite(&self) -> bool {
627                true
628            }
629            async fn begin_overwrite(&self) -> Result<(), FaucetError> {
630                self.log.lock().unwrap().push("begin");
631                Ok(())
632            }
633            async fn commit_overwrite(&self) -> Result<(), FaucetError> {
634                self.log.lock().unwrap().push("commit");
635                Ok(())
636            }
637            async fn abort_overwrite(&self) -> Result<(), FaucetError> {
638                self.log.lock().unwrap().push("abort");
639                Ok(())
640            }
641        }
642        let inner = OvwSink {
643            log: std::sync::Mutex::new(Vec::new()),
644        };
645        let policy = CleanupPolicy::new(scope(), vec!["id".into()], 10).unwrap();
646        let tracker = CleanupTracker::new(&inner, &policy);
647        assert!(tracker.is_overwrite());
648        assert_eq!(
649            tracker.batch_atomicity(),
650            crate::dlq::BatchAtomicity::BestEffort
651        );
652        // A write through the tracker must forward to the inner sink too.
653        assert_eq!(tracker.write_batch(&[json!({"id": 1})]).await.unwrap(), 1);
654        tracker.begin_overwrite().await.unwrap();
655        tracker.commit_overwrite().await.unwrap();
656        tracker.abort_overwrite().await.unwrap();
657        tracker.complete_run().await.unwrap();
658        assert_eq!(*inner.log.lock().unwrap(), vec!["begin", "commit", "abort"]);
659    }
660
661    #[test]
662    fn accumulates_keys_across_pages() {
663        let mut seen = SeenKeys::new();
664        let key = vec!["id".to_string()];
665        seen.record_page(&[json!({"id": 1}), json!({"id": 2})], &key, 100);
666        seen.record_page(&[json!({"id": 3})], &key, 100);
667        assert_eq!(seen.len(), 3);
668        assert!(!seen.overflowed());
669    }
670
671    #[test]
672    fn accumulates_composite_keys_in_declared_order() {
673        let mut seen = SeenKeys::new();
674        let key = vec!["a".to_string(), "b".to_string()];
675        seen.record_page(&[json!({"b": 2, "a": 1})], &key, 100);
676        assert_eq!(seen.len(), 1);
677        let t = &seen.keys()[0].0;
678        assert_eq!(
679            t[0].0, "a",
680            "key order follows the declared `key`, not the record"
681        );
682        assert_eq!(t[1].0, "b");
683    }
684
685    #[test]
686    fn skips_rows_with_a_missing_or_null_key() {
687        let mut seen = SeenKeys::new();
688        let key = vec!["id".to_string()];
689        seen.record_page(
690            &[
691                json!({"id": 1}),
692                json!({"other": 9}), // missing key
693                json!({"id": null}), // null key
694                json!("not an object"),
695            ],
696            &key,
697            100,
698        );
699        assert_eq!(seen.len(), 1, "only the well-keyed row is tracked");
700    }
701
702    #[test]
703    fn overflow_is_sticky_and_frees_the_buffer() {
704        let mut seen = SeenKeys::new();
705        let key = vec!["id".to_string()];
706        let page: Vec<Value> = (0..5).map(|i| json!({"id": i})).collect();
707        seen.record_page(&page, &key, 3);
708        assert!(seen.overflowed(), "ceiling of 3 must trip on a 5-row page");
709        assert!(seen.is_empty(), "buffer is freed — the cleanup will refuse");
710        // Sticky: a later page cannot un-overflow it.
711        seen.record_page(&[json!({"id": 99})], &key, 3);
712        assert!(seen.overflowed());
713        assert!(seen.is_empty());
714    }
715
716    #[test]
717    fn overflow_error_explains_that_nothing_was_deleted() {
718        let seen = SeenKeys::new();
719        let msg = seen.overflow_error(50).to_string();
720        assert!(msg.contains("Nothing was deleted"), "{msg}");
721        assert!(msg.contains("50"), "{msg}");
722    }
723
724    #[tokio::test]
725    async fn cleanup_tracker_forwards_every_hook() {
726        let probe = crate::forwarding::HookSink::default();
727        let policy = CleanupPolicy::new(scope(), vec!["id".into()], 10).unwrap();
728        let t = CleanupTracker::new(&probe, &policy);
729        crate::forwarding::assert_sink_forwards_every_hook(
730            &t,
731            &probe,
732            &[
733                "supports_columnar",
734                "write_batch_columnar",
735                "native_load_capabilities",
736                "load_native",
737            ],
738        )
739        .await;
740    }
741
742    #[tokio::test]
743    async fn overflowed_tells_a_refusal_from_a_backend_failure() {
744        struct Fails;
745        #[async_trait::async_trait]
746        impl Sink for Fails {
747            async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
748                Ok(r.len())
749            }
750            fn supports_cleanup(&self) -> bool {
751                true
752            }
753            async fn cleanup_scope(
754                &self,
755                _: &BTreeMap<String, Value>,
756                _: &SeenKeys,
757            ) -> Result<u64, FaucetError> {
758                Err(FaucetError::Sink("db down".into()))
759            }
760        }
761        let policy = CleanupPolicy::new(scope(), vec!["id".into()], 1).unwrap();
762        let t = CleanupTracker::new(&Fails, &policy);
763        t.write_batch(&[json!({"id": 1})]).await.unwrap();
764        assert!(t.finish(&policy).await.is_err());
765        assert!(!t.overflowed());
766        let t = CleanupTracker::new(&Fails, &policy);
767        t.write_batch(&[json!({"id": 1}), json!({"id": 2})])
768            .await
769            .unwrap();
770        assert!(t.finish(&policy).await.is_err());
771        assert!(t.overflowed());
772    }
773}