laminar_connectors/lakehouse/
iceberg_reference.rs

1//! Iceberg startup snapshot source for reference tables.
2
3use std::collections::VecDeque;
4
5use arrow_array::RecordBatch;
6use arrow_schema::SchemaRef;
7use async_trait::async_trait;
8use tracing::info;
9
10use crate::config::ConnectorConfig;
11use crate::error::ConnectorError;
12use crate::reference::ReferenceTableSource;
13
14use super::iceberg_config::IcebergSourceConfig;
15use super::snapshot_schema::{conform_snapshot_batch, validate_snapshot_schema};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18enum Phase {
19    Ready,
20    Draining,
21    Done,
22    Closed,
23}
24
25/// A finite snapshot of one Iceberg table.
26pub struct IcebergReferenceTableSource {
27    config: IcebergSourceConfig,
28    declared_schema: SchemaRef,
29    phase: Phase,
30    snapshot_batches: VecDeque<RecordBatch>,
31}
32
33impl IcebergReferenceTableSource {
34    /// Creates a source from parsed configuration and the declared table schema.
35    ///
36    /// # Errors
37    ///
38    /// Returns an error when an explicit projection conflicts with the declared schema.
39    pub fn new(
40        config: IcebergSourceConfig,
41        declared_schema: SchemaRef,
42    ) -> Result<Self, ConnectorError> {
43        let declared_columns = declared_schema
44            .fields()
45            .iter()
46            .map(|field| field.name().clone())
47            .collect::<Vec<_>>();
48        if !config.select_columns.is_empty() && config.select_columns != declared_columns {
49            return Err(ConnectorError::ConfigurationError(
50                "Iceberg select.columns must exactly match the declared reference-table columns"
51                    .into(),
52            ));
53        }
54        let config = {
55            let mut config = config;
56            config.select_columns = declared_columns;
57            config
58        };
59        Ok(Self {
60            config,
61            declared_schema,
62            phase: Phase::Ready,
63            snapshot_batches: VecDeque::new(),
64        })
65    }
66
67    /// Creates a source from SQL connector properties and the declared table schema.
68    ///
69    /// # Errors
70    ///
71    /// Returns an error when the Iceberg configuration or projection is invalid.
72    pub fn from_connector_config(
73        config: &ConnectorConfig,
74        declared_schema: SchemaRef,
75    ) -> Result<Self, ConnectorError> {
76        Self::new(IcebergSourceConfig::from_config(config)?, declared_schema)
77    }
78
79    async fn load_initial_snapshot(&mut self) -> Result<(), ConnectorError> {
80        let catalog = super::iceberg_io::build_catalog(&self.config.catalog).await?;
81        let table = super::iceberg_io::load_table(
82            catalog.as_ref(),
83            &self.config.catalog.namespace,
84            &self.config.catalog.table_name,
85        )
86        .await?;
87
88        let physical_schema =
89            match iceberg::arrow::schema_to_arrow_schema(&table.current_schema_ref()) {
90                Ok(schema) => schema,
91                Err(error) => {
92                    return Err(ConnectorError::SchemaMismatch(format!(
93                        "Iceberg to Arrow schema: {error}"
94                    )));
95                }
96            };
97        let projected_fields = self
98            .declared_schema
99            .fields()
100            .iter()
101            .map(|declared| {
102                physical_schema
103                    .index_of(declared.name())
104                    .map(|index| physical_schema.field(index).clone())
105                    .map_err(|_| {
106                        ConnectorError::ReadError(format!(
107                            "Iceberg snapshot is missing declared column '{}'",
108                            declared.name()
109                        ))
110                    })
111            })
112            .collect::<Result<Vec<_>, _>>()?;
113        let projected_schema = arrow_schema::Schema::new(projected_fields);
114        validate_snapshot_schema(&projected_schema, self.declared_schema.as_ref())?;
115
116        let snapshot_id = self
117            .config
118            .snapshot_id
119            .or_else(|| super::iceberg_io::current_snapshot_id(&table));
120        let batches =
121            super::iceberg_io::scan_table(&table, snapshot_id, &self.config.select_columns)
122                .await?
123                .into_iter()
124                .map(|batch| conform_snapshot_batch(&batch, &self.declared_schema))
125                .collect::<Result<Vec<_>, _>>()?;
126        let rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
127        info!(snapshot = ?snapshot_id, rows, "Iceberg reference snapshot loaded");
128        self.snapshot_batches = batches.into();
129        Ok(())
130    }
131}
132
133#[async_trait]
134impl ReferenceTableSource for IcebergReferenceTableSource {
135    async fn poll_snapshot(&mut self) -> Result<Option<RecordBatch>, ConnectorError> {
136        match self.phase {
137            Phase::Closed => {
138                return Err(ConnectorError::InvalidState {
139                    expected: "open reference snapshot source".into(),
140                    actual: "closed".into(),
141                });
142            }
143            Phase::Done => return Ok(None),
144            Phase::Ready => {
145                self.load_initial_snapshot().await?;
146                self.phase = Phase::Draining;
147            }
148            Phase::Draining => {}
149        }
150
151        if let Some(batch) = self.snapshot_batches.pop_front() {
152            return Ok(Some(batch));
153        }
154        self.phase = Phase::Done;
155        Ok(None)
156    }
157
158    async fn close(&mut self) -> Result<(), ConnectorError> {
159        self.phase = Phase::Closed;
160        self.snapshot_batches.clear();
161        Ok(())
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use std::sync::Arc;
168
169    use arrow_schema::{DataType, Field, Schema};
170
171    use super::*;
172
173    fn test_source_config() -> IcebergSourceConfig {
174        let mut config = ConnectorConfig::new("iceberg");
175        config.set("catalog.uri", "http://localhost:8181");
176        config.set("warehouse", "s3://test/wh");
177        config.set("namespace", "test");
178        config.set("table.name", "dim_customers");
179        IcebergSourceConfig::from_config(&config).unwrap()
180    }
181
182    fn declared_schema() -> SchemaRef {
183        Arc::new(Schema::new(vec![
184            Field::new("id", DataType::Int64, false),
185            Field::new("name", DataType::Utf8, true),
186        ]))
187    }
188
189    #[test]
190    fn construction_carries_declared_non_null_key_schema() {
191        let source =
192            IcebergReferenceTableSource::new(test_source_config(), declared_schema()).unwrap();
193        assert_eq!(source.config.select_columns, vec!["id", "name"]);
194        assert_eq!(source.declared_schema.field(0).name(), "id");
195        assert!(!source.declared_schema.field(0).is_nullable());
196    }
197
198    #[test]
199    fn conflicting_explicit_projection_is_rejected() {
200        let mut config = test_source_config();
201        config.select_columns = vec!["name".into()];
202        assert!(IcebergReferenceTableSource::new(config, declared_schema()).is_err());
203    }
204
205    #[tokio::test]
206    async fn exhaustion_and_close_are_stable_without_external_io() {
207        let mut source =
208            IcebergReferenceTableSource::new(test_source_config(), declared_schema()).unwrap();
209        source.phase = Phase::Draining;
210        assert!(source.poll_snapshot().await.unwrap().is_none());
211        assert!(source.poll_snapshot().await.unwrap().is_none());
212        source.close().await.unwrap();
213        source.close().await.unwrap();
214        assert!(source.poll_snapshot().await.is_err());
215    }
216}
laminar_connectors/lakehouse/iceberg_reference.rs

laminar_connectors/lakehouse/
iceberg_reference.rs
