faucet_source_singer/
assemble.rs

1//! Pure page-assembly state machine.
2//!
3//! Extracted from the stream so the RECORD/STATE/EOF → [`StreamPage`] logic is
4//! unit-testable without spawning a process. The stream feeds parsed messages
5//! in; the process/error paths (idle timeout, non-zero exit, malformed-fail)
6//! stay in `stream.rs`.
7
8use faucet_core::{StreamPage, Value};
9
10/// Accumulates RECORDs into pages and attaches the latest STATE as the page
11/// bookmark, per the v0 single-stream rules.
12pub struct PageAssembler {
13    target_stream: String,
14    batch_size: usize,
15    flush_on_state: bool,
16    buffer: Vec<Value>,
17    /// Latest STATE `value` seen but not yet attached to a flushed page.
18    pending_state: Option<Value>,
19}
20
21impl PageAssembler {
22    /// `batch_size == 0` disables size-based flushing (flush only on STATE/EOF).
23    pub fn new(target_stream: impl Into<String>, batch_size: usize, flush_on_state: bool) -> Self {
24        Self {
25            target_stream: target_stream.into(),
26            batch_size,
27            flush_on_state,
28            buffer: Vec::new(),
29            pending_state: None,
30        }
31    }
32
33    /// Feed a RECORD. Records for a different stream are ignored (single-stream
34    /// v0). Returns a page when the `batch_size` threshold is reached.
35    pub fn on_record(&mut self, stream: &str, record: Value) -> Option<StreamPage> {
36        if stream != self.target_stream {
37            return None;
38        }
39        self.buffer.push(record);
40        if self.batch_size != 0 && self.buffer.len() >= self.batch_size {
41            Some(self.flush())
42        } else {
43            None
44        }
45    }
46
47    /// Feed a STATE. Always becomes the pending checkpoint; flushes immediately
48    /// when `flush_on_state` (yielding a page — possibly empty — that carries
49    /// the STATE as its bookmark).
50    pub fn on_state(&mut self, value: Value) -> Option<StreamPage> {
51        self.pending_state = Some(value);
52        if self.flush_on_state {
53            Some(self.flush())
54        } else {
55            None
56        }
57    }
58
59    /// Flush any trailing buffer + pending checkpoint at end-of-stream. Returns
60    /// `None` when there is nothing to emit.
61    pub fn on_eof(&mut self) -> Option<StreamPage> {
62        if self.buffer.is_empty() && self.pending_state.is_none() {
63            None
64        } else {
65            Some(self.flush())
66        }
67    }
68
69    /// Drain the buffer + pending checkpoint into a page.
70    fn flush(&mut self) -> StreamPage {
71        StreamPage {
72            records: std::mem::take(&mut self.buffer),
73            bookmark: self.pending_state.take(),
74        }
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use serde_json::json;
82
83    fn rec(id: i64) -> Value {
84        json!({ "id": id })
85    }
86
87    #[test]
88    fn flushes_at_batch_size_with_no_bookmark() {
89        let mut a = PageAssembler::new("s", 2, true);
90        assert!(a.on_record("s", rec(1)).is_none());
91        let page = a.on_record("s", rec(2)).expect("flush at size 2");
92        assert_eq!(page.records, vec![rec(1), rec(2)]);
93        assert!(
94            page.bookmark.is_none(),
95            "no STATE covers a size-based flush"
96        );
97        // third record buffered until EOF
98        assert!(a.on_record("s", rec(3)).is_none());
99        let tail = a.on_eof().unwrap();
100        assert_eq!(tail.records, vec![rec(3)]);
101    }
102
103    #[test]
104    fn state_flush_attaches_bookmark() {
105        let mut a = PageAssembler::new("s", 1000, true);
106        assert!(a.on_record("s", rec(1)).is_none());
107        assert!(a.on_record("s", rec(2)).is_none());
108        let page = a.on_state(json!({"last_id": 2})).expect("flush on state");
109        assert_eq!(page.records, vec![rec(1), rec(2)]);
110        assert_eq!(page.bookmark, Some(json!({"last_id": 2})));
111        // nothing left
112        assert!(a.on_eof().is_none());
113    }
114
115    #[test]
116    fn empty_run_with_trailing_state_yields_empty_page_with_bookmark() {
117        let mut a = PageAssembler::new("s", 1000, true);
118        let page = a.on_state(json!({"last_id": 0})).unwrap();
119        assert!(page.records.is_empty());
120        assert_eq!(page.bookmark, Some(json!({"last_id": 0})));
121    }
122
123    #[test]
124    fn records_for_other_streams_are_ignored() {
125        let mut a = PageAssembler::new("wanted", 1, true);
126        assert!(a.on_record("other", rec(1)).is_none());
127        assert!(a.on_record("other", rec(2)).is_none());
128        // only the wanted stream flushes
129        let page = a.on_record("wanted", rec(3)).unwrap();
130        assert_eq!(page.records, vec![rec(3)]);
131    }
132
133    #[test]
134    fn no_flush_on_state_defers_bookmark_to_next_flush() {
135        let mut a = PageAssembler::new("s", 2, false);
136        // STATE arrives first; not flushed because flush_on_state == false
137        assert!(a.on_state(json!({"last_id": 0})).is_none());
138        assert!(a.on_record("s", rec(1)).is_none());
139        // batch flush picks up the pending bookmark
140        let page = a.on_record("s", rec(2)).unwrap();
141        assert_eq!(page.records, vec![rec(1), rec(2)]);
142        assert_eq!(page.bookmark, Some(json!({"last_id": 0})));
143    }
144
145    #[test]
146    fn eof_with_empty_buffer_and_no_state_yields_nothing() {
147        let mut a = PageAssembler::new("s", 1000, true);
148        assert!(a.on_eof().is_none());
149    }
150}
faucet_source_singer/assemble.rs

faucet_source_singer/
assemble.rs
