laminar_connectors/reference/
mod.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;
laminar_connectors/reference/mod.rs

laminar_connectors/reference/
mod.rs
