laminar_connectors/lakehouse/
iceberg_source.rs

1//! Apache Iceberg source connector implementation.
2//!
3//! [`IcebergSource`] implements [`SourceConnector`] for polling Iceberg
4//! snapshots. It is a **reference/lookup table** source — Iceberg tables
5//! are snapshot-based with no push mechanism.
6//!
7//! On each poll cycle the source checks for a newer snapshot. If one exists,
8//! only the newly added data files are read via manifest-level diff
9//! (see `iceberg_incremental::scan_incremental`). The first read is a full scan.
10
11use std::collections::VecDeque;
12use std::sync::Arc;
13use std::time::Instant;
14
15use arrow_array::RecordBatch;
16use arrow_schema::SchemaRef;
17use async_trait::async_trait;
18#[cfg(feature = "iceberg")]
19use tracing::debug;
20#[cfg(feature = "iceberg")]
21use tracing::info;
22
23use crate::checkpoint::SourceCheckpoint;
24use crate::config::{ConnectorConfig, ConnectorState};
25use crate::connector::{
26    SourceBatch, SourceConnector, SourceConsistency, SourceContract, SourceTopology,
27};
28use crate::connector::{SourcePosition, SourceStart};
29use crate::error::ConnectorError;
30
31use super::iceberg_config::IcebergSourceConfig;
32
33/// Apache Iceberg source connector.
34///
35/// Polls for new snapshots on a configurable interval and emits
36/// `RecordBatch` data. Supports pinning to a specific snapshot.
37#[allow(dead_code)] // Fields used by feature-gated I/O methods.
38pub struct IcebergSource {
39    /// Source configuration — reparsed from `ConnectorConfig` in `start()`.
40    config: IcebergSourceConfig,
41    /// Discovered Arrow schema.
42    schema: Option<SchemaRef>,
43    /// Connector lifecycle state.
44    state: ConnectorState,
45    /// Buffered batches from the most recent scan.
46    buffer: VecDeque<RecordBatch>,
47    /// Last fully-ingested snapshot ID.
48    last_snapshot_id: Option<i64>,
49    /// Time of last snapshot poll.
50    last_poll_time: Option<Instant>,
51    /// Cached catalog connection.
52    #[cfg(feature = "iceberg")]
53    catalog: Option<Arc<dyn iceberg::Catalog>>,
54    /// Cached table handle.
55    #[cfg(feature = "iceberg")]
56    table: Option<iceberg::table::Table>,
57}
58
59impl IcebergSource {
60    /// Creates a new Iceberg source with the given configuration.
61    #[must_use]
62    pub fn new(config: IcebergSourceConfig, _registry: Option<&prometheus::Registry>) -> Self {
63        Self {
64            config,
65            schema: None,
66            state: ConnectorState::Created,
67            buffer: VecDeque::new(),
68            last_snapshot_id: None,
69            last_poll_time: None,
70            #[cfg(feature = "iceberg")]
71            catalog: None,
72            #[cfg(feature = "iceberg")]
73            table: None,
74        }
75    }
76
77    /// Checks for a new snapshot and loads data if available.
78    #[cfg(feature = "iceberg")]
79    async fn refresh(&mut self) -> Result<(), ConnectorError> {
80        if let Some(last) = self.last_poll_time {
81            if last.elapsed() < self.config.poll_interval {
82                return Ok(());
83            }
84        }
85        self.last_poll_time = Some(Instant::now());
86
87        // Reload table metadata to see new snapshots.
88        let catalog = self
89            .catalog
90            .as_ref()
91            .ok_or_else(|| ConnectorError::InvalidState {
92                expected: "started".into(),
93                actual: "catalog not initialized".into(),
94            })?;
95        let table = super::iceberg_io::load_table(
96            catalog.as_ref(),
97            &self.config.catalog.namespace,
98            &self.config.catalog.table_name,
99        )
100        .await?;
101        self.table = Some(table);
102
103        let table = self.table.as_ref().unwrap();
104        let current_snap = super::iceberg_io::current_snapshot_id(table);
105
106        // If pinned to a specific snapshot, only load once.
107        if let Some(pinned) = self.config.snapshot_id {
108            if self.last_snapshot_id.is_some() {
109                return Ok(());
110            }
111            let batches =
112                super::iceberg_io::scan_table(table, Some(pinned), &self.config.select_columns)
113                    .await?;
114            self.buffer.extend(batches);
115            self.last_snapshot_id = Some(pinned);
116            return Ok(());
117        }
118
119        if current_snap == self.last_snapshot_id {
120            return Ok(());
121        }
122
123        // Incremental read if we have a previous snapshot, full scan otherwise.
124        let batches = if let (Some(old), Some(new)) = (self.last_snapshot_id, current_snap) {
125            super::iceberg_incremental::scan_incremental(
126                table,
127                old,
128                new,
129                &self.config.select_columns,
130            )
131            .await?
132        } else {
133            super::iceberg_io::scan_table(table, current_snap, &self.config.select_columns).await?
134        };
135
136        if self.schema.is_none() {
137            if let Some(first) = batches.first() {
138                self.schema = Some(first.schema());
139            }
140        }
141
142        let rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
143        debug!(snapshot = ?current_snap, rows, "iceberg source loaded snapshot");
144
145        self.buffer.extend(batches);
146        self.last_snapshot_id = current_snap;
147
148        Ok(())
149    }
150}
151
152#[async_trait]
153impl SourceConnector for IcebergSource {
154    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
155        let (config, position, _) = request.into_parts();
156        if let SourcePosition::Resume { attempt, .. } = position {
157            return Err(ConnectorError::ConfigurationError(format!(
158                "Iceberg is an ephemeral source and cannot resume checkpoint attempt {attempt:?}"
159            )));
160        }
161        let config = &config;
162        // Re-parse config from runtime ConnectorConfig (not factory defaults).
163        if !config.properties().is_empty() {
164            self.config = IcebergSourceConfig::from_config(config)?;
165        }
166
167        #[cfg(feature = "iceberg")]
168        {
169            let catalog = super::iceberg_io::build_catalog(&self.config.catalog).await?;
170            let table = super::iceberg_io::load_table(
171                catalog.as_ref(),
172                &self.config.catalog.namespace,
173                &self.config.catalog.table_name,
174            )
175            .await?;
176
177            let iceberg_schema = table.current_schema_ref();
178            let arrow_schema =
179                iceberg::arrow::schema_to_arrow_schema(&iceberg_schema).map_err(|e| {
180                    ConnectorError::SchemaMismatch(format!("iceberg→arrow schema: {e}"))
181                })?;
182            self.schema = Some(Arc::new(arrow_schema));
183
184            info!(
185                table = self.config.catalog.table_name,
186                namespace = self.config.catalog.namespace,
187                "iceberg source connected"
188            );
189
190            self.catalog = Some(catalog);
191            self.table = Some(table);
192            self.state = ConnectorState::Running;
193            return Ok(());
194        }
195
196        #[cfg(not(feature = "iceberg"))]
197        {
198            self.state = ConnectorState::Failed;
199            Err(ConnectorError::ConfigurationError(
200                "Apache Iceberg requires the 'iceberg' feature".into(),
201            ))
202        }
203    }
204
205    async fn poll_batch(
206        &mut self,
207        max_records: usize,
208    ) -> Result<Option<SourceBatch>, ConnectorError> {
209        if let Some(batch) = self.buffer.pop_front() {
210            if batch.num_rows() <= max_records {
211                return Ok(Some(SourceBatch::new(batch)));
212            }
213            let take = batch.slice(0, max_records);
214            let remainder = batch.slice(max_records, batch.num_rows() - max_records);
215            self.buffer.push_front(remainder);
216            return Ok(Some(SourceBatch::new(take)));
217        }
218
219        #[cfg(feature = "iceberg")]
220        self.refresh().await?;
221
222        Ok(self.buffer.pop_front().map(SourceBatch::new))
223    }
224
225    fn schema(&self) -> SchemaRef {
226        self.schema
227            .clone()
228            .unwrap_or_else(|| Arc::new(arrow_schema::Schema::empty()))
229    }
230
231    fn checkpoint(&self) -> SourceCheckpoint {
232        let mut cp = SourceCheckpoint::new();
233        if let Some(sid) = self.last_snapshot_id {
234            cp.set_offset("snapshot_id", sid.to_string());
235        }
236        cp.set_metadata("connector_type", "iceberg");
237        cp
238    }
239
240    fn contract(&self, _config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
241        Ok(SourceContract::new(
242            SourceConsistency::Ephemeral,
243            SourceTopology::Singleton,
244        ))
245    }
246
247    async fn close(&mut self) -> Result<(), ConnectorError> {
248        #[cfg(feature = "iceberg")]
249        {
250            self.catalog = None;
251            self.table = None;
252        }
253        self.state = ConnectorState::Closed;
254        Ok(())
255    }
256}
257
258#[cfg(test)]
259mod tests {
260    use super::*;
261    use crate::config::ConnectorConfig;
262
263    fn test_source_config() -> IcebergSourceConfig {
264        let mut config = ConnectorConfig::new("iceberg");
265        config.set("catalog.uri", "http://localhost:8181");
266        config.set("warehouse", "s3://test/wh");
267        config.set("namespace", "test");
268        config.set("table.name", "dim_customers");
269        IcebergSourceConfig::from_config(&config).unwrap()
270    }
271
272    #[test]
273    fn test_new_source() {
274        let source = IcebergSource::new(test_source_config(), None);
275        assert!(source.schema.is_none());
276        assert!(source.last_snapshot_id.is_none());
277        assert!(source.buffer.is_empty());
278    }
279
280    #[test]
281    fn test_checkpoint_round_trip() {
282        let mut source = IcebergSource::new(test_source_config(), None);
283        source.last_snapshot_id = Some(42);
284
285        let cp = source.checkpoint();
286        assert_eq!(cp.get_offset("snapshot_id"), Some("42"));
287        assert_eq!(cp.get_metadata("connector_type"), Some("iceberg"));
288    }
289
290    #[tokio::test]
291    async fn resume_fails_before_opening_the_ephemeral_source() {
292        let mut source = IcebergSource::new(test_source_config(), None);
293        let error = source
294            .start(
295                SourceStart::new(
296                    ConnectorConfig::new("iceberg"),
297                    SourcePosition::Resume {
298                        attempt: laminar_core::state::CheckpointAttempt::canonical(11),
299                        checkpoint: SourceCheckpoint::new(),
300                    },
301                    crate::connector::DeliveryGuarantee::BestEffort,
302                )
303                .unwrap(),
304            )
305            .await
306            .expect_err("ephemeral Iceberg source must reject recovery");
307        assert!(error.to_string().contains("ephemeral"));
308        assert_eq!(source.state, ConnectorState::Created);
309    }
310
311    #[test]
312    fn test_source_contract_is_ephemeral_singleton() {
313        let source = IcebergSource::new(test_source_config(), None);
314        let contract = source.contract(&ConnectorConfig::new("iceberg")).unwrap();
315        assert_eq!(contract.consistency, SourceConsistency::Ephemeral);
316        assert_eq!(contract.topology, SourceTopology::Singleton);
317    }
318}
laminar_connectors/lakehouse/iceberg_source.rs

laminar_connectors/lakehouse/
iceberg_source.rs
