laminar_connectors/lakehouse/iceberg_reference/
mod.rs

1//! Iceberg startup snapshot source for reference tables.
2
3use arrow_array::RecordBatch;
4use arrow_schema::SchemaRef;
5use async_trait::async_trait;
6use futures_util::StreamExt;
7use iceberg::scan::ArrowRecordBatchStream;
8use iceberg::spec::SnapshotRef;
9use tracing::info;
10
11use crate::config::ConnectorConfig;
12use crate::error::ConnectorError;
13use crate::reference::ReferenceTableSource;
14
15use super::iceberg_config::IcebergReadMode;
16use super::iceberg_config::IcebergSourceConfig;
17use super::iceberg_scan::{
18    connector_scan_error, plan_files, preflight_snapshot, ManifestReadLimits,
19};
20use super::snapshot_schema::{conform_snapshot_batch, validate_snapshot_schema};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23enum Phase {
24    Ready,
25    Draining,
26    Done,
27    Failed,
28    Closed,
29}
30
31/// A finite snapshot of one Iceberg table.
32pub struct IcebergReferenceTableSource {
33    config: IcebergSourceConfig,
34    declared_schema: SchemaRef,
35    phase: Phase,
36    snapshot_stream: Option<ArrowRecordBatchStream>,
37    snapshot_id: Option<i64>,
38    emitted_rows: u64,
39    binding: Option<crate::schema::resolution::SchemaBinding>,
40}
41
42impl IcebergReferenceTableSource {
43    /// Creates a source from parsed configuration and the declared table schema.
44    ///
45    /// # Errors
46    ///
47    /// Returns an error when an explicit projection conflicts with the declared schema.
48    pub fn new(
49        config: IcebergSourceConfig,
50        declared_schema: SchemaRef,
51    ) -> Result<Self, ConnectorError> {
52        config.validate_read_limits()?;
53        if config.read_mode != IcebergReadMode::Snapshot {
54            return Err(ConnectorError::ConfigurationError(
55                "Iceberg reference tables require read.mode=snapshot".into(),
56            ));
57        }
58        let declared_columns = declared_schema
59            .fields()
60            .iter()
61            .map(|field| field.name().clone())
62            .collect::<Vec<_>>();
63        if !config.select_columns.is_empty() && config.select_columns != declared_columns {
64            return Err(ConnectorError::ConfigurationError(
65                "Iceberg select.columns must exactly match the declared reference-table columns"
66                    .into(),
67            ));
68        }
69        let config = {
70            let mut config = config;
71            config.select_columns = declared_columns;
72            config
73        };
74        Ok(Self {
75            config,
76            declared_schema,
77            phase: Phase::Ready,
78            snapshot_stream: None,
79            snapshot_id: None,
80            emitted_rows: 0,
81            binding: None,
82        })
83    }
84
85    /// Creates a source from SQL connector properties and the declared table schema.
86    ///
87    /// # Errors
88    ///
89    /// Returns an error when the Iceberg configuration or projection is invalid.
90    pub fn from_connector_config(
91        config: &ConnectorConfig,
92        declared_schema: SchemaRef,
93    ) -> Result<Self, ConnectorError> {
94        let mut source = Self::new(IcebergSourceConfig::from_config(config)?, declared_schema)?;
95        source.binding = config.schema_binding().cloned();
96        Ok(source)
97    }
98
99    async fn load_initial_snapshot(&mut self) -> Result<(), ConnectorError> {
100        let catalog =
101            super::iceberg_io::build_catalog(&self.config.catalog, &self.config.storage).await?;
102        let table = super::iceberg_io::load_table_with_timeout(
103            catalog.as_ref(),
104            &self.config.catalog.namespace,
105            &self.config.catalog.table_name,
106            self.config.catalog.request_timeout,
107        )
108        .await?;
109
110        super::schema_resolution::verify_identity(
111            self.binding.as_ref(),
112            &super::schema_resolution::iceberg_native(&table),
113        )?;
114        let snapshot = selected_snapshot(&table, &self.config)?;
115        let snapshot_schema = match &snapshot {
116            Some(snapshot) => snapshot
117                .schema(table.metadata())
118                .map_err(|error| connector_scan_error("resolve Iceberg snapshot schema", &error))?,
119            None => table.current_schema_ref(),
120        };
121        let predicate = super::iceberg_scan::parse_and_bind_filter(
122            self.config.filter.as_deref(),
123            snapshot_schema.clone(),
124        )?;
125        let physical_schema = iceberg::arrow::schema_to_arrow_schema(&snapshot_schema)
126            .map_err(|error| connector_scan_error("convert Iceberg snapshot schema", &error))?;
127        let projected_fields = self
128            .declared_schema
129            .fields()
130            .iter()
131            .map(|declared| {
132                physical_schema
133                    .index_of(declared.name())
134                    .map(|index| physical_schema.field(index).clone())
135                    .map_err(|_| {
136                        ConnectorError::ReadError(format!(
137                            "Iceberg snapshot is missing declared column '{}'",
138                            declared.name()
139                        ))
140                    })
141            })
142            .collect::<Result<Vec<_>, _>>()?;
143        let projected_schema = arrow_schema::Schema::new(projected_fields);
144        validate_snapshot_schema(&projected_schema, self.declared_schema.as_ref())?;
145        let Some(snapshot) = snapshot else {
146            return Ok(());
147        };
148
149        let mut builder = table
150            .scan()
151            .snapshot_id(snapshot.snapshot_id())
152            .with_batch_size(Some(8_192))
153            .with_concurrency_limit(self.config.scan_concurrency)
154            .select(self.config.select_columns.iter().map(String::as_str));
155        if let Some(predicate) = predicate {
156            builder = builder.with_filter(predicate);
157        }
158        let scan = builder
159            .build()
160            .map_err(|error| connector_scan_error("build Iceberg reference scan", &error))?;
161        let deadline = super::iceberg_io::checked_deadline(
162            self.config.storage.request_timeout,
163            "storage.request_timeout",
164        )?;
165        preflight_snapshot(
166            &table,
167            &snapshot,
168            ManifestReadLimits::from_source(&self.config),
169            deadline,
170        )
171        .await?;
172        let tasks = plan_files(&scan, self.config.max_planned_files, deadline).await?;
173        let reader = table
174            .reader_builder()
175            .with_batch_size(8_192)
176            .with_data_file_concurrency_limit(self.config.scan_concurrency)
177            .build()
178            .read(tasks)
179            .map_err(|error| connector_scan_error("create Iceberg reference reader", &error))?;
180        self.snapshot_id = Some(snapshot.snapshot_id());
181        self.snapshot_stream = Some(reader.stream());
182        Ok(())
183    }
184
185    async fn next_batch(&mut self) -> Result<Option<RecordBatch>, ConnectorError> {
186        loop {
187            let Some(stream) = self.snapshot_stream.as_mut() else {
188                return Ok(None);
189            };
190            let next = tokio::time::timeout(self.config.storage.request_timeout, stream.next())
191                .await
192                .map_err(|_| {
193                    ConnectorError::ReadError(
194                        "[LDB-ICEBERG-STORAGE-TIMEOUT] reference snapshot read made no progress"
195                            .into(),
196                    )
197                })?;
198            let Some(result) = next else {
199                self.snapshot_stream = None;
200                return Ok(None);
201            };
202            let batch = result.map_err(|error| {
203                connector_scan_error("Iceberg reference snapshot read failed", &error)
204            })?;
205            if batch.num_rows() == 0 {
206                continue;
207            }
208            let batch = conform_snapshot_batch(&batch, &self.declared_schema)?;
209            self.emitted_rows = self
210                .emitted_rows
211                .saturating_add(u64::try_from(batch.num_rows()).unwrap_or(u64::MAX));
212            return Ok(Some(batch));
213        }
214    }
215}
216
217fn selected_snapshot(
218    table: &iceberg::table::Table,
219    config: &IcebergSourceConfig,
220) -> Result<Option<SnapshotRef>, ConnectorError> {
221    if let Some(snapshot_id) = config.snapshot_id {
222        return table
223            .metadata()
224            .snapshot_by_id(snapshot_id)
225            .cloned()
226            .map(Some)
227            .ok_or_else(|| {
228                ConnectorError::ReadError(format!(
229                    "[LDB-ICEBERG-SNAPSHOT-MISSING] snapshot {snapshot_id} does not exist"
230                ))
231            });
232    }
233    if let Some(snapshot) = table.metadata().snapshot_for_ref(&config.table_ref) {
234        return Ok(Some(snapshot.clone()));
235    }
236    if config.table_ref == "main" && table.metadata().current_snapshot().is_none() {
237        return Ok(None);
238    }
239    Err(ConnectorError::ReadError(format!(
240        "[LDB-ICEBERG-REF-MISSING] table ref '{}' does not exist",
241        config.table_ref
242    )))
243}
244
245#[async_trait]
246impl ReferenceTableSource for IcebergReferenceTableSource {
247    async fn resolve_schema(
248        &mut self,
249        config: &ConnectorConfig,
250        explicit: Option<SchemaRef>,
251    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
252        super::schema_resolution::iceberg_source(config, explicit).await
253    }
254
255    async fn poll_snapshot(&mut self) -> Result<Option<RecordBatch>, ConnectorError> {
256        match self.phase {
257            Phase::Closed => {
258                return Err(ConnectorError::InvalidState {
259                    expected: "open reference snapshot source".into(),
260                    actual: "closed".into(),
261                });
262            }
263            Phase::Done => return Ok(None),
264            Phase::Failed => {
265                return Err(ConnectorError::InvalidState {
266                    expected: "readable reference snapshot source".into(),
267                    actual: "failed".into(),
268                });
269            }
270            Phase::Ready => {
271                if let Err(error) = self.load_initial_snapshot().await {
272                    self.phase = Phase::Failed;
273                    return Err(error);
274                }
275                self.phase = Phase::Draining;
276            }
277            Phase::Draining => {}
278        }
279
280        match self.next_batch().await {
281            Ok(Some(batch)) => Ok(Some(batch)),
282            Ok(None) => {
283                self.phase = Phase::Done;
284                info!(
285                    snapshot = ?self.snapshot_id,
286                    rows = self.emitted_rows,
287                    "Iceberg reference snapshot completed"
288                );
289                Ok(None)
290            }
291            Err(error) => {
292                self.snapshot_stream = None;
293                self.phase = Phase::Failed;
294                Err(error)
295            }
296        }
297    }
298
299    async fn close(&mut self) -> Result<(), ConnectorError> {
300        self.phase = Phase::Closed;
301        self.snapshot_stream = None;
302        Ok(())
303    }
304}
305
306#[cfg(test)]
307mod tests;
laminar_connectors/lakehouse/iceberg_reference/mod.rs

laminar_connectors/lakehouse/iceberg_reference/
mod.rs
