knut-bifrost 0.1.8

//! **Iceberg → table delivery**: the loop that reads an Iceberg table's change
//! windows and hands them to a [`TableSink`] (Doris), one snapshot at a time.
//!
//! Feature `skade`. The pieces it drives already existed — skade's
//! `read_changelog`, [`partition_change_rows`], [`deliver_window`] — and
//! nothing ran them end to end. This is that driver, and nothing more.
//!
//! # Where the table is
//!
//! [`IcebergLocator`] names the table three ways, one per kind of lake:
//!
//! | locator | how the table is found | feature |
//! |---|---|---|
//! | `Warehouse` | an embedded skade warehouse directory | `skade` |
//! | `S3` | its S3 root, **no catalog** — `version-hint.text`, else the newest `*.metadata.json` | `skade-s3` |
//! | `Rest` | an Iceberg REST catalog; MinIO AIStor's `/_iceberg` with SigV4 service `s3tables` | `skade-rest` |
//!
//! Every round loads the table again, so a new commit is seen on the next round
//! whichever way the table is found.
//!
//! # One snapshot per window
//!
//! A window is `(previous snapshot, next snapshot]`, never a wider range:
//!
//! * **memory** — `read_changelog` materialises the whole window; one commit at
//!   a time bounds it by the largest single commit, not by the backlog;
//! * **ordering** — inside one Iceberg commit an update is an equality delete
//!   plus an append of the same key; the sink ships deletes before upserts, so
//!   a one-commit window is always applied in the right order. A multi-commit
//!   window could hold insert-then-delete of one key, which the sink cannot
//!   order (`sink::doris` says so);
//! * **replay** — the label is derived from the snapshot pair, so a crash
//!   replays exactly the commit that was in flight.
//!
//! # Full reloads: allowed once, refused after
//!
//! An `overwrite`/`replace` commit, or a position-delete file, cannot be
//! expressed as a delta ([`skade::DeltaPlan::needs_full_reload`]).
//!
//! * **Before the first delivery** (no cursor) that is the initial load, and
//!   the right answer is the whole current snapshot as inserts — one window
//!   `(None, current]`. It is held in memory in full; say so in the log.
//! * **After** it, the window is refused and the cursor held
//!   ([`refuse_full_reload`]). The operator's resolution is explicit: empty the
//!   destination table, delete this table's cursor, start again.
//!
//! The cheap part of that check runs first, off the snapshot operations in the
//! metadata already loaded — no manifest or data file is read to refuse an
//! overwrite. A position-delete file is only visible in the manifests, so that
//! case is caught by the changelog's own plan, after the read.

use std::time::{Duration, Instant};

use arrow_array::RecordBatch;
use skade::iceberg::spec::Operation;
use skade::iceberg::table::Table as IceTable;

use crate::{
    error::{BifrostError, Result},
    sink::table::{SnapshotWindow, TableSink, TableWindow, deliver_window, refuse_full_reload},
    state::SyncState,
    transform::partition_change_rows,
};

/// Where an Iceberg table is. See the module docs.
#[derive(Debug, Clone)]
pub enum IcebergLocator {
    /// An embedded skade warehouse at `dir`, table `table` (`ns.name` or `name`).
    Warehouse { dir: std::path::PathBuf, table: String },
    /// A table on S3/MinIO by its root, no catalog.
    #[cfg(feature = "skade-s3")]
    S3(skade::connect::S3Table),
    /// A table behind an Iceberg REST catalog, `table` = `ns.name`.
    #[cfg(feature = "skade-rest")]
    Rest { catalog: skade::connect::IcebergRest, table: String },
}

impl IcebergLocator {
    /// A stable, secret-free name for the table — the cursor key's source half.
    pub fn source_id(&self) -> String {
        match self {
            IcebergLocator::Warehouse { dir, table } => format!("skade:{}#{table}", dir.display()),
            #[cfg(feature = "skade-s3")]
            IcebergLocator::S3(t) => format!("s3:{}", t.location.trim_end_matches('/')),
            #[cfg(feature = "skade-rest")]
            IcebergLocator::Rest { catalog, table } => format!(
                "rest:{}#{}#{table}",
                catalog.uri.trim_end_matches('/'),
                catalog.warehouse.as_deref().unwrap_or("-")
            ),
        }
    }

    /// Load the table at its current metadata.
    pub async fn load(&self) -> Result<IceTable> {
        let src = |e: skade::SkadeError| BifrostError::Source(format!("{}: {e}", self.source_id()));
        match self {
            IcebergLocator::Warehouse { dir, table } => {
                let wh = skade::Warehouse::open(dir).await.map_err(src)?;
                Ok(wh.table(table).await.map_err(src)?.inner().clone())
            }
            #[cfg(feature = "skade-s3")]
            IcebergLocator::S3(t) => t.load().await.map_err(src),
            #[cfg(feature = "skade-rest")]
            IcebergLocator::Rest { catalog, table } => {
                use skade::connect::CatalogConnector as _;
                let ident = skade::iceberg::TableIdent::from_strs(table.split('.')).map_err(|e| {
                    BifrostError::Source(format!("{}: bad table name {table:?}: {e}", self.source_id()))
                })?;
                let cat = catalog.open().await.map_err(src)?;
                cat.load_table(&ident)
                    .await
                    .map_err(|e| BifrostError::Source(format!("{}: {e}", self.source_id())))
            }
        }
    }
}

/// Knobs for [`IcebergDelivery`].
#[derive(Debug, Clone)]
pub struct DeliveryOptions {
    /// At most this many one-snapshot windows per [`IcebergDelivery::round`];
    /// the rest wait for the next round, so one round never runs unbounded.
    pub max_windows_per_round: usize,
}

impl Default for DeliveryOptions {
    fn default() -> Self {
        Self { max_windows_per_round: 64 }
    }
}

/// What one [`IcebergDelivery::round`] did.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RoundOutcome {
    /// Windows delivered (and cursor advances made) this round.
    pub windows: usize,
    pub rows_loaded: usize,
    pub rows_deleted: usize,
    /// Rows the store refused as an already-seen label: a replay, not an error.
    pub rows_duplicate: usize,
    /// The snapshot the cursor stands on after the round.
    pub at_snapshot: Option<i64>,
    /// Snapshots still to deliver after this round (0 = caught up).
    pub behind: usize,
    /// Wall time reading changelogs (table load + read_changelog + partition).
    pub read: Duration,
    /// Wall time inside the sink (encode + load + cursor commit).
    pub load: Duration,
}

/// The delivery loop: one Iceberg table into one table sink.
pub struct IcebergDelivery<S: TableSink> {
    locator: IcebergLocator,
    sink: S,
    state: SyncState,
    key: String,
    opts: DeliveryOptions,
    probed: bool,
}

impl<S: TableSink> IcebergDelivery<S> {
    /// `destination` names where the rows go (e.g. `doris:db.table`); it is
    /// part of the cursor key, so one source delivered to two destinations
    /// keeps two cursors.
    pub fn new(
        locator: IcebergLocator,
        destination: &str,
        sink: S,
        state: SyncState,
        opts: DeliveryOptions,
    ) -> Self {
        let key = format!("{} -> {destination}", locator.source_id());
        Self { locator, sink, state, key, opts, probed: false }
    }

    /// The cursor key this delivery reads and advances.
    pub fn cursor_key(&self) -> &str {
        &self.key
    }

    /// The snapshot the cursor stands on.
    pub fn cursor(&self) -> Option<i64> {
        self.state.get(&self.key).snapshot_id
    }

    /// Deliver what is waiting, up to [`DeliveryOptions::max_windows_per_round`]
    /// windows. Every delivered window has advanced the cursor before the next
    /// one is read; an error leaves the cursor on the last delivered window.
    pub async fn round(&mut self) -> Result<RoundOutcome> {
        let mut out = RoundOutcome::default();
        let t0 = Instant::now();
        let table = self.locator.load().await?;
        out.read += t0.elapsed();

        let from = self.cursor();
        out.at_snapshot = from;
        let Some(current) = table.metadata().current_snapshot().map(|s| s.snapshot_id()) else {
            return Ok(out); // no snapshot yet: nothing to deliver
        };
        if from == Some(current) {
            return Ok(out);
        }
        let chain = table
            .metadata()
            .snapshots_added_between(from, current)
            .map_err(|e| BifrostError::Source(format!("{}: snapshot lineage: {e}", self.key)))?;

        // The cheap refusal: an operation the delta cannot express, found in the
        // metadata already in hand.
        let rewrites = chain.iter().any(|id| {
            table
                .metadata()
                .snapshot_by_id(*id)
                .is_some_and(|s| !matches!(s.summary().operation, Operation::Append | Operation::Delete))
        });

        if rewrites && from.is_none() {
            tracing::warn!(
                target: "bifrost::delivery",
                event = "initial_full_snapshot",
                table = %self.key,
                to = current,
                snapshots = chain.len(),
                "history rewrites rows in place: the initial load is the whole current snapshot, held in memory in full",
            );
            self.deliver(&table, None, current, true, &mut out).await?;
            out.behind = 0;
            return Ok(out);
        }
        if rewrites {
            refuse_full_reload(&self.key, true)?;
        }

        let mut prev = from;
        for (i, &snap) in chain.iter().enumerate() {
            if i == self.opts.max_windows_per_round {
                out.behind = chain.len() - i;
                break;
            }
            self.deliver(&table, prev, snap, false, &mut out).await?;
            prev = Some(snap);
        }
        Ok(out)
    }

    /// Read one window, probe the sink once, deliver, account.
    async fn deliver(
        &mut self,
        table: &IceTable,
        from: Option<i64>,
        to: i64,
        initial_full: bool,
        out: &mut RoundOutcome,
    ) -> Result<()> {
        let src = |e: skade::SkadeError| BifrostError::Source(format!("{e}"));
        let t = Instant::now();
        if !self.probed {
            let schema = skade::arrow_schema_of(table).map_err(src)?;
            self.sink.probe(&schema).await?;
            self.probed = true;
        }
        let cl = skade::read_changelog(table, from, to).await.map_err(src)?;
        // A position delete is only visible here, in the plan of the read. Only
        // the initial load may take the whole-snapshot answer.
        let needs_full_reload = cl.plan.needs_full_reload && !(initial_full || from.is_none());
        let rows = partition_change_rows(&cl.rows)?;
        let deletes = project_to_keys(&rows.deletes, cl.key_columns())?;
        out.read += t.elapsed();

        let window = TableWindow {
            table: &self.key,
            snapshots: SnapshotWindow::new(from, to),
            upserts: &rows.upserts,
            deletes: &deletes,
            needs_full_reload,
        };
        let t = Instant::now();
        let ack = deliver_window(&mut self.sink, &mut self.state, &window).await?;
        out.load += t.elapsed();

        out.windows += 1;
        out.rows_loaded += ack.rows_loaded;
        out.rows_deleted += ack.rows_deleted;
        out.rows_duplicate += ack.rows_rejected_duplicate;
        out.at_snapshot = Some(ack.to_snapshot());
        Ok(())
    }

    /// Run rounds until `stop` resolves: `poll` between rounds while healthy,
    /// doubling up to `max_backoff` while rounds fail. A round in flight when
    /// `stop` fires is finished first — its cursor commit included — so a
    /// SIGTERM never leaves a loaded-but-uncommitted window behind.
    pub async fn run(
        &mut self,
        poll: Duration,
        max_backoff: Duration,
        stop: impl std::future::Future<Output = ()>,
    ) -> Result<()> {
        let mut backoff = crate::metrics::Backoff::new(poll, max_backoff);
        tokio::pin!(stop);
        loop {
            let result = self.round().await;
            let failed = match &result {
                Ok(o) => {
                    if o.windows > 0 {
                        tracing::info!(
                            target: "bifrost::delivery",
                            event = "round",
                            table = %self.key,
                            windows = o.windows,
                            loaded = o.rows_loaded,
                            deleted = o.rows_deleted,
                            duplicate = o.rows_duplicate,
                            at = ?o.at_snapshot,
                            behind = o.behind,
                            read_ms = o.read.as_millis() as u64,
                            load_ms = o.load.as_millis() as u64,
                            "delivered",
                        );
                    }
                    false
                }
                Err(e) => {
                    tracing::error!(target: "bifrost::delivery", event = "round_failed", table = %self.key, error = %e, "round failed, cursor held");
                    true
                }
            };
            let behind = matches!(&result, Ok(o) if o.behind > 0);
            let sleep = if behind { Duration::ZERO } else { backoff.next_sleep(failed) };
            tokio::select! {
                _ = &mut stop => return Ok(()),
                _ = tokio::time::sleep(sleep) => {}
            }
        }
    }
}

/// A delete names a row by its identity columns; everything else in the
/// widened changelog row is null. Ship only the identity, so a NOT NULL value
/// column in the destination cannot fail a delete. No key columns known (a
/// table without equality deletes) → the rows go as they are.
fn project_to_keys(deletes: &[RecordBatch], keys: &[String]) -> Result<Vec<RecordBatch>> {
    if keys.is_empty() {
        return Ok(deletes.to_vec());
    }
    deletes
        .iter()
        .map(|b| {
            let idx: Vec<usize> = keys
                .iter()
                .map(|k| {
                    b.schema().index_of(k).map_err(|_| {
                        BifrostError::Source(format!("delete rows have no identity column {k:?}"))
                    })
                })
                .collect::<Result<_>>()?;
            b.project(&idx).map_err(|e| BifrostError::Source(format!("projecting delete keys: {e}")))
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use arrow_array::{Int64Array, StringArray};
    use arrow_schema::{DataType, Field, Schema};
    use async_trait::async_trait;

    use super::*;
    use crate::sink::table::{AckToken, LoadAck};

    /// A sink that records every window it is handed.
    #[derive(Clone, Default)]
    struct Recording {
        windows: Arc<Mutex<Vec<(Option<i64>, i64, usize, usize)>>>,
        probes: Arc<Mutex<usize>>,
        fail_at: Option<i64>,
    }

    #[async_trait]
    impl TableSink for Recording {
        async fn probe(&mut self, _: &Schema) -> Result<()> {
            *self.probes.lock().unwrap() += 1;
            Ok(())
        }
        async fn apply(&mut self, w: &TableWindow<'_>, token: AckToken) -> Result<LoadAck> {
            if self.fail_at == Some(w.snapshots.to) {
                return Err(BifrostError::Other("sink down".into()));
            }
            self.windows.lock().unwrap().push((w.snapshots.from, w.snapshots.to, w.upsert_rows(), w.delete_rows()));
            Ok(token.ack(w.snapshots.to, vec![]).with_counts(w.upsert_rows(), w.delete_rows(), 0))
        }
    }

    fn people(ids: &[i64]) -> skade::arrow_array::RecordBatch {
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int64, false),
            Field::new("name", DataType::Utf8, true),
        ]));
        let names: Vec<String> = ids.iter().map(|i| format!("p{i}")).collect();
        RecordBatch::try_new(
            schema,
            vec![Arc::new(Int64Array::from(ids.to_vec())), Arc::new(StringArray::from(names))],
        )
        .unwrap()
    }

    async fn warehouse_with(commits: &[&[i64]]) -> (tempfile::TempDir, Vec<i64>) {
        let dir = tempfile::tempdir().unwrap();
        let wh = skade::Warehouse::open(dir.path()).await.unwrap();
        let mut t = wh.create_table("people", &people(&[]).schema()).await.unwrap();
        let mut snaps = Vec::new();
        for ids in commits {
            t.append(&[people(ids)]).await.unwrap();
            snaps.push(t.current_snapshot_id().unwrap());
        }
        (dir, snaps)
    }

    fn delivery(dir: &std::path::Path, sink: Recording, opts: DeliveryOptions) -> IcebergDelivery<Recording> {
        let state = SyncState::open(dir.join("cursor.json")).unwrap();
        let loc = IcebergLocator::Warehouse { dir: dir.to_path_buf(), table: "people".into() };
        IcebergDelivery::new(loc, "test:db.people", sink, state, opts)
    }

    #[tokio::test]
    async fn one_window_per_snapshot_and_the_cursor_follows_each() {
        let (dir, snaps) = warehouse_with(&[&[1, 2], &[3], &[4, 5, 6]]).await;
        let sink = Recording::default();
        let mut d = delivery(dir.path(), sink.clone(), DeliveryOptions::default());
        let o = d.round().await.unwrap();
        assert_eq!(o.windows, 3);
        assert_eq!(o.rows_loaded, 6);
        assert_eq!(o.at_snapshot, Some(snaps[2]));
        assert_eq!(d.cursor(), Some(snaps[2]));
        let w = sink.windows.lock().unwrap().clone();
        assert_eq!(
            w,
            vec![(None, snaps[0], 2, 0), (Some(snaps[0]), snaps[1], 1, 0), (Some(snaps[1]), snaps[2], 3, 0)]
        );
        assert_eq!(*sink.probes.lock().unwrap(), 1, "probe runs once, before the first window");

        // Caught up: nothing more, cursor unchanged.
        let o = d.round().await.unwrap();
        assert_eq!(o.windows, 0);
        assert_eq!(d.cursor(), Some(snaps[2]));
    }

    #[tokio::test]
    async fn a_round_is_bounded_and_says_how_far_behind_it_is() {
        let (dir, snaps) = warehouse_with(&[&[1], &[2], &[3]]).await;
        let mut d = delivery(dir.path(), Recording::default(), DeliveryOptions { max_windows_per_round: 2 });
        let o = d.round().await.unwrap();
        assert_eq!((o.windows, o.behind), (2, 1));
        assert_eq!(d.cursor(), Some(snaps[1]));
        let o = d.round().await.unwrap();
        assert_eq!((o.windows, o.behind), (1, 0));
        assert_eq!(d.cursor(), Some(snaps[2]));
    }

    #[tokio::test]
    async fn a_failed_window_holds_the_cursor_on_the_last_delivered_one() {
        let (dir, snaps) = warehouse_with(&[&[1], &[2], &[3]]).await;
        let sink = Recording { fail_at: Some(snaps[1]), ..Recording::default() };
        let mut d = delivery(dir.path(), sink, DeliveryOptions::default());
        assert!(d.round().await.is_err());
        assert_eq!(d.cursor(), Some(snaps[0]), "window 1 landed, window 2 failed: cursor on 1");

        // The cursor survives a restart (it is on disk), and delivery resumes at 2.
        let sink = Recording::default();
        let mut d = delivery(dir.path(), sink.clone(), DeliveryOptions::default());
        assert_eq!(d.cursor(), Some(snaps[0]));
        d.round().await.unwrap();
        let w = sink.windows.lock().unwrap().clone();
        assert_eq!(w.first().map(|w| (w.0, w.1)), Some((Some(snaps[0]), snaps[1])));
        assert_eq!(d.cursor(), Some(snaps[2]));
    }

    #[test]
    fn deletes_ship_only_their_identity_columns() {
        let b = people(&[7, 8]);
        let p = project_to_keys(&[b.clone()], &["id".to_string()]).unwrap();
        assert_eq!(p[0].num_columns(), 1);
        assert_eq!(p[0].schema().field(0).name(), "id");
        assert_eq!(project_to_keys(&[b.clone()], &[]).unwrap()[0].num_columns(), 2);
        assert!(project_to_keys(&[b], &["nope".to_string()]).is_err());
    }
}