laminar_connectors/lakehouse/

iceberg_io.rs

//! Feature-gated I/O operations for Apache Iceberg.
//!
//! Contains catalog construction, table loading, scanning, and writing
//! functions. All code requires the `iceberg` feature.
#![allow(clippy::disallowed_types)] // cold path: lakehouse I/O
#![cfg(feature = "iceberg")]

use std::collections::HashMap;
use std::sync::Arc;

use arrow_array::RecordBatch;
use iceberg::table::Table;
use iceberg::transaction::{ApplyTransactionAction, Transaction};
use iceberg::{Catalog, CatalogBuilder, TableIdent};
use iceberg_catalog_rest::RestCatalogBuilder;
use iceberg_storage_opendal::OpenDalStorageFactory;
use tokio_stream::StreamExt;

use super::iceberg_config::{IcebergCatalogConfig, IcebergCatalogType};
use crate::error::ConnectorError;

/// Selects the `OpenDalStorageFactory` for the table-data URLs the catalog
/// will return. Explicit `storage.type` wins; otherwise inferred from the
/// `s3://` / `s3a://` / `file://` warehouse URL.
fn storage_factory(
    warehouse: &str,
    storage_type: Option<&str>,
) -> Result<Arc<dyn iceberg::io::StorageFactory>, ConnectorError> {
    let scheme = storage_type
        .map(str::to_lowercase)
        .or_else(|| {
            if warehouse.starts_with("s3a://") {
                Some("s3a".to_string())
            } else if warehouse.starts_with("s3://") {
                Some("s3".to_string())
            } else if warehouse.starts_with("file://") {
                Some("fs".to_string())
            } else {
                None
            }
        })
        .ok_or_else(|| {
            ConnectorError::ConfigurationError(format!(
                "[LDB-5100] cannot infer storage backend from warehouse '{warehouse}'; \
                 set storage.type = 's3' | 's3a' | 'fs'"
            ))
        })?;

    // configured_scheme is the bare scheme ("s3"), NOT "s3://" —
    // iceberg-storage-opendal formats the prefix as `{scheme}://{bucket}/`.
    let factory: Arc<dyn iceberg::io::StorageFactory> = match scheme.as_str() {
        "s3" | "s3a" => Arc::new(OpenDalStorageFactory::S3 {
            configured_scheme: scheme,
            customized_credential_load: None,
        }),
        "fs" => Arc::new(OpenDalStorageFactory::Fs),
        other => {
            return Err(ConnectorError::ConfigurationError(format!(
                "[LDB-5101] unsupported storage.type '{other}'; expected s3 | s3a | fs"
            )));
        }
    };
    Ok(factory)
}

/// Builds a REST catalog from configuration.
///
/// # Errors
///
/// Returns `ConnectorError::ConnectionFailed` if catalog initialization fails.
pub async fn build_catalog(
    config: &IcebergCatalogConfig,
) -> Result<Arc<dyn Catalog>, ConnectorError> {
    match config.catalog_type {
        IcebergCatalogType::Rest => build_rest_catalog(config).await,
    }
}

async fn build_rest_catalog(
    config: &IcebergCatalogConfig,
) -> Result<Arc<dyn Catalog>, ConnectorError> {
    let storage_factory = storage_factory(&config.warehouse, config.storage_type.as_deref())?;

    let mut props = HashMap::new();
    props.insert("uri".to_string(), config.catalog_uri.clone());
    props.insert("warehouse".to_string(), config.warehouse.clone());

    for (k, v) in &config.properties {
        props.insert(k.clone(), v.clone());
    }

    let catalog = RestCatalogBuilder::default()
        .with_storage_factory(storage_factory)
        .load("laminardb", props)
        .await
        .map_err(|e| ConnectorError::ConnectionFailed(format!("iceberg catalog: {e}")))?;

    Ok(Arc::new(catalog))
}

/// Loads an Iceberg table from the catalog.
///
/// # Errors
///
/// Returns `ConnectorError::ReadError` if the table cannot be loaded.
pub async fn load_table(
    catalog: &dyn Catalog,
    namespace: &str,
    table_name: &str,
) -> Result<Table, ConnectorError> {
    let ns = iceberg::NamespaceIdent::from_strs(namespace.split('.').collect::<Vec<_>>())
        .map_err(|e| ConnectorError::ConfigurationError(format!("invalid namespace: {e}")))?;

    let ident = TableIdent::new(ns, table_name.to_string());

    catalog
        .load_table(&ident)
        .await
        .map_err(|e| ConnectorError::ReadError(format!("load table '{table_name}': {e}")))
}

/// Scans a table and returns all record batches for the current snapshot.
///
/// # Errors
///
/// Returns `ConnectorError::ReadError` on scan failure.
pub async fn scan_table(
    table: &Table,
    snapshot_id: Option<i64>,
    select_columns: &[String],
) -> Result<Vec<RecordBatch>, ConnectorError> {
    let mut scan_builder = table.scan();

    if let Some(sid) = snapshot_id {
        scan_builder = scan_builder.snapshot_id(sid);
    }

    if select_columns.is_empty() {
        scan_builder = scan_builder.select_all();
    } else {
        scan_builder = scan_builder.select(select_columns.iter().map(String::as_str));
    }

    let scan = scan_builder
        .build()
        .map_err(|e| ConnectorError::ReadError(format!("build scan: {e}")))?;

    let stream = scan
        .to_arrow()
        .await
        .map_err(|e| ConnectorError::ReadError(format!("scan to arrow: {e}")))?;

    let mut batches = Vec::new();
    let mut stream = std::pin::pin!(stream);
    while let Some(result) = stream.next().await {
        let batch = result.map_err(|e| ConnectorError::ReadError(format!("read batch: {e}")))?;
        batches.push(batch);
    }

    Ok(batches)
}

/// Returns the current snapshot ID of a table, if any.
#[must_use]
pub fn current_snapshot_id(table: &Table) -> Option<i64> {
    table.metadata().current_snapshot().map(|s| s.snapshot_id())
}

/// Commits data files to an Iceberg table via a fast-append transaction.
///
/// When `epoch_metadata` is provided, also stores `writer_id` and `epoch`
/// in table properties for exactly-once recovery. Both the data append
/// and property update are committed atomically in a single transaction.
///
/// Returns the updated table with the new snapshot.
///
/// # Errors
///
/// Returns `ConnectorError::TransactionError` on commit failure.
pub async fn commit_data_files(
    table: &Table,
    catalog: &dyn Catalog,
    data_files: Vec<iceberg::spec::DataFile>,
    epoch_metadata: Option<(&str, u64)>,
) -> Result<Table, ConnectorError> {
    let tx = Transaction::new(table);
    let tx: Transaction = tx
        .fast_append()
        .add_data_files(data_files)
        .apply(tx)
        .map_err(|e| ConnectorError::TransactionError(format!("apply fast_append: {e}")))?;

    // Chain epoch metadata into the same atomic transaction.
    let tx = if let Some((writer_id, epoch)) = epoch_metadata {
        let key = format!("laminardb.writer.{writer_id}.last_epoch");
        tx.update_table_properties()
            .set(key, epoch.to_string())
            .apply(tx)
            .map_err(|e| {
                ConnectorError::TransactionError(format!("apply update_table_properties: {e}"))
            })?
    } else {
        tx
    };

    tx.commit(catalog)
        .await
        .map_err(|e| ConnectorError::TransactionError(format!("commit: {e}")))
}

/// Reads the last committed epoch for a writer from Iceberg table properties.
///
/// Returns `None` if the writer has never committed to this table.
#[must_use]
pub fn get_last_committed_epoch(table: &Table, writer_id: &str) -> Option<u64> {
    let key = format!("laminardb.writer.{writer_id}.last_epoch");
    table.metadata().properties().get(&key)?.parse().ok()
}

/// Creates an Iceberg table (and namespace) if it does not already exist.
///
/// # Errors
///
/// Returns `ConnectorError` on creation failure.
pub async fn ensure_table_exists(
    catalog: &dyn Catalog,
    namespace: &str,
    table_name: &str,
    arrow_schema: &arrow_schema::SchemaRef,
) -> Result<(), ConnectorError> {
    let ns = iceberg::NamespaceIdent::from_strs(namespace.split('.').collect::<Vec<_>>())
        .map_err(|e| ConnectorError::ConfigurationError(format!("invalid namespace: {e}")))?;

    let ident = TableIdent::new(ns.clone(), table_name.to_string());

    if catalog
        .table_exists(&ident)
        .await
        .map_err(|e| ConnectorError::ReadError(format!("table_exists: {e}")))?
    {
        return Ok(());
    }

    // Pipeline-derived Arrow schemas don't carry `PARQUET:field_id`
    // metadata; let iceberg-rust assign sequential IDs.
    let iceberg_schema = iceberg::arrow::arrow_schema_to_schema_auto_assign_ids(arrow_schema)
        .map_err(|e| {
            ConnectorError::SchemaMismatch(format!("arrow→iceberg schema conversion: {e}"))
        })?;

    if !catalog
        .namespace_exists(&ns)
        .await
        .map_err(|e| ConnectorError::ReadError(format!("namespace_exists: {e}")))?
    {
        catalog
            .create_namespace(&ns, HashMap::new())
            .await
            .map_err(|e| ConnectorError::WriteError(format!("create namespace: {e}")))?;
    }

    let creation = iceberg::TableCreation::builder()
        .name(table_name.to_string())
        .schema(iceberg_schema)
        .build();

    catalog
        .create_table(&ns, creation)
        .await
        .map_err(|e| ConnectorError::WriteError(format!("create table: {e}")))?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_storage_factory_infers_s3_from_warehouse_url() {
        let f = storage_factory("s3://bucket/warehouse", None).unwrap();
        assert!(format!("{f:?}").contains("S3"));
    }

    #[test]
    fn test_storage_factory_infers_s3a_from_warehouse_url() {
        let f = storage_factory("s3a://bucket/warehouse", None).unwrap();
        assert!(format!("{f:?}").contains("S3"));
    }

    #[test]
    fn test_storage_factory_infers_fs_from_file_url() {
        let f = storage_factory("file:///tmp/warehouse", None).unwrap();
        assert!(format!("{f:?}").contains("Fs"));
    }

    #[test]
    fn test_storage_factory_bare_path_requires_explicit_storage_type() {
        // Trimmed `/` and `./` inference: REST catalogs use logical names
        // and we don't want a silent default to local fs.
        let err = storage_factory("/tmp/warehouse", None)
            .unwrap_err()
            .to_string();
        assert!(err.contains("LDB-5100"), "got: {err}");
    }

    #[test]
    fn test_storage_factory_explicit_overrides_inference() {
        // Lakekeeper-style: warehouse is a name, storage backend is S3.
        let f = storage_factory("demo", Some("s3")).unwrap();
        assert!(format!("{f:?}").contains("S3"));
    }

    #[test]
    fn test_storage_factory_unknown_warehouse_without_storage_type_errors() {
        let err = storage_factory("demo", None).unwrap_err().to_string();
        assert!(err.contains("LDB-5100"), "got: {err}");
    }

    #[test]
    fn test_storage_factory_rejects_unknown_storage_type() {
        let err = storage_factory("demo", Some("hdfs"))
            .unwrap_err()
            .to_string();
        assert!(err.contains("LDB-5101"), "got: {err}");
    }

    #[test]
    fn test_current_snapshot_id_empty_table() {
        let schema = iceberg::spec::Schema::builder()
            .with_fields(vec![])
            .build()
            .unwrap();

        let creation = iceberg::TableCreation::builder()
            .name("test_table".to_string())
            .schema(schema)
            .location("s3://test/location".to_string())
            .build();

        let metadata = iceberg::spec::TableMetadataBuilder::from_table_creation(creation)
            .unwrap()
            .build()
            .unwrap();

        let table = Table::builder()
            .metadata(metadata.metadata)
            .identifier(TableIdent::new(
                iceberg::NamespaceIdent::new("test".to_string()),
                "t".to_string(),
            ))
            .file_io(iceberg::io::FileIO::new_with_memory())
            .build()
            .unwrap();

        assert!(current_snapshot_id(&table).is_none());
    }
}