Skip to main content

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    /// Resolve a snapshot reader before hydration, without consuming snapshot rows.
14    ///
15    /// # Errors
16    /// The default extension contract requires explicit fields. Metadata-capable
17    /// implementations override this hook and retain native identity at hydration.
18    async fn resolve_schema(
19        &mut self,
20        config: &crate::config::ConnectorConfig,
21        explicit: Option<arrow_schema::SchemaRef>,
22    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
23        let schema = explicit.ok_or_else(|| ConnectorError::FeatureUnsupported(
24            "reference-table connector requires explicit fields; no metadata resolver is implemented".into()))?;
25        crate::schema::resolution::logical_binding(
26            config,
27            crate::schema::resolution::SchemaDirection::Source,
28            crate::schema::resolution::SchemaOrigin::Explicit,
29            &schema,
30        )
31    }
32
33    /// Returns the next snapshot batch, or `None` after the complete snapshot was delivered.
34    async fn poll_snapshot(&mut self) -> Result<Option<RecordBatch>, ConnectorError>;
35
36    /// Releases source resources.
37    async fn close(&mut self) -> Result<(), ConnectorError>;
38}
39
40/// In-memory finite snapshot source for tests.
41#[cfg(any(test, feature = "testing"))]
42pub struct MockReferenceTableSource {
43    snapshot_batches: VecDeque<RecordBatch>,
44    /// Whether [`ReferenceTableSource::close`] has been called.
45    pub closed: bool,
46}
47
48#[cfg(any(test, feature = "testing"))]
49impl MockReferenceTableSource {
50    /// Creates a source that drains the supplied snapshot batches in order.
51    #[must_use]
52    pub fn new(snapshot_batches: Vec<RecordBatch>) -> Self {
53        Self {
54            snapshot_batches: VecDeque::from(snapshot_batches),
55            closed: false,
56        }
57    }
58
59    /// Creates a source with an empty snapshot.
60    #[must_use]
61    pub fn empty() -> Self {
62        Self::new(Vec::new())
63    }
64}
65
66#[cfg(any(test, feature = "testing"))]
67#[async_trait::async_trait]
68impl ReferenceTableSource for MockReferenceTableSource {
69    async fn poll_snapshot(&mut self) -> Result<Option<RecordBatch>, ConnectorError> {
70        if self.closed {
71            return Err(ConnectorError::InvalidState {
72                expected: "open reference snapshot source".into(),
73                actual: "closed".into(),
74            });
75        }
76        Ok(self.snapshot_batches.pop_front())
77    }
78
79    async fn close(&mut self) -> Result<(), ConnectorError> {
80        self.closed = true;
81        Ok(())
82    }
83}
84
85#[cfg(test)]
86mod tests;