laminar_connectors/
reference.rs

1//! Startup snapshot sources for reference tables.
2
3#[cfg(any(test, feature = "testing"))]
4use std::collections::VecDeque;
5
6use arrow_array::RecordBatch;
7
8use crate::error::ConnectorError;
9
10/// A finite source used to hydrate a reference table before processing starts.
11#[async_trait::async_trait]
12pub trait ReferenceTableSource: Send {
13    /// Returns the next snapshot batch, or `None` after the complete snapshot was delivered.
14    async fn poll_snapshot(&mut self) -> Result<Option<RecordBatch>, ConnectorError>;
15
16    /// Releases source resources.
17    async fn close(&mut self) -> Result<(), ConnectorError>;
18}
19
20/// In-memory finite snapshot source for tests.
21#[cfg(any(test, feature = "testing"))]
22pub struct MockReferenceTableSource {
23    snapshot_batches: VecDeque<RecordBatch>,
24    /// Whether [`ReferenceTableSource::close`] has been called.
25    pub closed: bool,
26}
27
28#[cfg(any(test, feature = "testing"))]
29impl MockReferenceTableSource {
30    /// Creates a source that drains the supplied snapshot batches in order.
31    #[must_use]
32    pub fn new(snapshot_batches: Vec<RecordBatch>) -> Self {
33        Self {
34            snapshot_batches: VecDeque::from(snapshot_batches),
35            closed: false,
36        }
37    }
38
39    /// Creates a source with an empty snapshot.
40    #[must_use]
41    pub fn empty() -> Self {
42        Self::new(Vec::new())
43    }
44}
45
46#[cfg(any(test, feature = "testing"))]
47#[async_trait::async_trait]
48impl ReferenceTableSource for MockReferenceTableSource {
49    async fn poll_snapshot(&mut self) -> Result<Option<RecordBatch>, ConnectorError> {
50        if self.closed {
51            return Err(ConnectorError::InvalidState {
52                expected: "open reference snapshot source".into(),
53                actual: "closed".into(),
54            });
55        }
56        Ok(self.snapshot_batches.pop_front())
57    }
58
59    async fn close(&mut self) -> Result<(), ConnectorError> {
60        self.closed = true;
61        Ok(())
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use std::sync::Arc;
68
69    use arrow_array::Int32Array;
70    use arrow_schema::{DataType, Field, Schema};
71
72    use super::*;
73
74    fn test_batch(values: &[i32]) -> RecordBatch {
75        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
76        RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values.to_vec()))]).unwrap()
77    }
78
79    #[tokio::test]
80    async fn snapshot_exhaustion_and_close_are_stable() {
81        let mut source = MockReferenceTableSource::new(vec![test_batch(&[1, 2]), test_batch(&[3])]);
82
83        assert_eq!(source.poll_snapshot().await.unwrap().unwrap().num_rows(), 2);
84        assert_eq!(source.poll_snapshot().await.unwrap().unwrap().num_rows(), 1);
85        assert!(source.poll_snapshot().await.unwrap().is_none());
86        assert!(source.poll_snapshot().await.unwrap().is_none());
87        source.close().await.unwrap();
88        source.close().await.unwrap();
89        assert!(source.closed);
90        assert!(source.poll_snapshot().await.is_err());
91    }
92}
laminar_connectors/reference.rs

laminar_connectors/
reference.rs
