knut_bifrost/

listener.rs

//! The sync loop: poll → transform → sink → advance state.

use std::sync::Arc;
use std::time::Duration;

use crate::{
    airflow::{AssetEmitter, round_to_events},
    error::{BifrostError, Result},
    mapping::{MappingSpec, RowFilter, SyncMode},
    metrics::{Backoff, Metrics},
    mutation::{GraphMutation, SinkReport},
    registry::TransformRegistry,
    sink::GraphSink,
    source::{TableSource, TableSpec},
    state::{SyncState, TableState},
    transform::{TableContext, Transform, deletes_from_change_rows, edge_deletes_from_change_rows, partition_change_rows},
};

/// Outcome of one [`Listener::sync_once`] pass over all tables.
#[derive(Debug, Default, Clone)]
pub struct RoundReport {
    /// (table, rows in, sink report) for each table that had news.
    pub synced: Vec<(String, usize, SinkReport)>,
    /// Tables that errored this round (left for the next round to retry).
    pub failed: Vec<(String, String)>,
}

impl RoundReport {
    /// Did this round make **no progress**? True when at least one table was
    /// attempted and *every* attempted table failed (a fully-failed round). An
    /// idle round — nothing synced and nothing failed — is **not** a failure
    /// (the source simply had no news), so idle pollers keep their cadence
    /// instead of backing off forever. Drives the poll-loop's [`Backoff`] and
    /// the `rounds_failed` metric.
    pub fn is_failure(&self) -> bool {
        self.synced.is_empty() && !self.failed.is_empty()
    }
}

/// One table's compiled, round-invariant work, prepared once at
/// [`Listener::new`] so a poll never re-clones the spec (H7) or recompiles the
/// pipeline (H6). The transforms are an `Arc`-shared slice so the whole list
/// can be cheaply cloned out of `&self` and iterated while `&mut self` runs.
struct CompiledTable {
    table: String,
    mode: SyncMode,
    /// `Ok` once-compiled transforms, or the compile error's message. A
    /// compile failure (empty pipeline / unknown custom transform) is
    /// deterministic; we surface it as a per-round `failed` entry every round,
    /// exactly as the old per-round `compile()` did.
    pipeline: std::result::Result<Arc<[Arc<dyn Transform>]>, String>,
    /// The **projection** this table's pipeline needs, resolved once at
    /// construction (it is a pure function of the spec). Empty = read every
    /// column, which is what a `custom` step forces — see
    /// [`TableMapping::projected_columns`](crate::mapping::TableMapping::projected_columns).
    columns: Vec<String>,
    /// The optional scan **predicate** from the mapping, pushed down every round.
    filter: Option<RowFilter>,
}

/// Watches every table in a [`MappingSpec`] and keeps a graph sink current.
pub struct Listener<Src, Snk> {
    /// Per-table work compiled once. `Arc`-shared so `sync_once` can clone the
    /// (cheap) handle and iterate it while holding `&mut self` for the sync.
    tables: Arc<[CompiledTable]>,
    source: Src,
    sink: Snk,
    state: SyncState,
    /// Optional AIP-82 asset emitter + URI prefix. When set, every successful
    /// round's `table_synced` results are announced as Airflow asset events,
    /// closing the loop: table ingested → asset event → downstream DAG.
    asset_emitter: Option<(Arc<dyn AssetEmitter>, String)>,
    /// Cumulative counters folded one round at a time (plan.md §7). Emitted as a
    /// structured `metrics` tracing event after every round and exposed via
    /// [`Listener::metrics`].
    metrics: Metrics,
}

impl<Src: TableSource, Snk: GraphSink> Listener<Src, Snk> {
    pub fn new(
        spec: MappingSpec,
        registry: TransformRegistry,
        source: Src,
        sink: Snk,
        state: SyncState,
    ) -> Self {
        // H6/H7: compile every pipeline once here and store the round-invariant
        // (name, mode, transforms) so polling never clones the spec or rebuilds
        // the transforms. Compile errors are kept as messages and re-surfaced
        // per round, preserving the old "table fails, retry next round" shape.
        let tables: Arc<[CompiledTable]> = spec
            .tables
            .iter()
            .map(|m| CompiledTable {
                table: m.table.clone(),
                mode: m.mode,
                pipeline: m
                    .compile(&registry)
                    .map(Arc::<[Arc<dyn Transform>]>::from)
                    .map_err(|e| e.to_string()),
                // Resolved ONCE: the projection is a pure function of the spec,
                // so a poll never recomputes it (same reason the pipeline is
                // compiled here).
                columns: m.projected_columns().unwrap_or_default(),
                filter: m.filter.clone(),
            })
            .collect();
        Self {
            tables,
            source,
            sink,
            state,
            asset_emitter: None,
            metrics: Metrics::default(),
        }
    }

    /// The cumulative run metrics so far (counters folded after every
    /// [`sync_once`](Self::sync_once) / poll round). See [`Metrics`] and
    /// [`crate::metrics::render_prometheus`].
    pub fn metrics(&self) -> &Metrics {
        &self.metrics
    }

    /// Attach an [`AssetEmitter`]: after every successful sync round the
    /// listener emits one Airflow asset event per ingested table, with URIs
    /// `{uri_prefix}{table}`. See [`crate::airflow`].
    pub fn with_asset_emitter(
        mut self,
        emitter: Arc<dyn AssetEmitter>,
        uri_prefix: impl Into<String>,
    ) -> Self {
        self.asset_emitter = Some((emitter, uri_prefix.into()));
        self
    }

    /// One pass over all mapped tables. A failing table does not stop the
    /// others; its state is not advanced, so the next pass retries it.
    pub async fn sync_once(&mut self) -> Result<RoundReport> {
        let mut report = RoundReport::default();
        // H7: clone the cheap `Arc` handle (not the spec) so we can iterate the
        // compiled tables while `sync_table` borrows `&mut self`.
        let tables = Arc::clone(&self.tables);
        for compiled in tables.iter() {
            match self.sync_table(compiled).await {
                Ok(Some((rows, sink_report))) => {
                    tracing::info!(
                        target: "bifrost::round",
                        event = "table_synced",
                        table = %compiled.table,
                        rows,
                        vertices = sink_report.vertices_written,
                        edges = sink_report.edges_written,
                        deleted = sink_report.deleted,
                        "table ingested",
                    );
                    report
                        .synced
                        .push((compiled.table.clone(), rows, sink_report));
                }
                Ok(None) => {}
                Err(e) => {
                    tracing::warn!(
                        target: "bifrost::round",
                        event = "table_failed",
                        table = %compiled.table,
                        error = %e,
                        "table sync failed (will retry next round)",
                    );
                    report.failed.push((compiled.table.clone(), e.to_string()));
                }
            }
        }
        // H8: persist the file-backed ledger once for the whole round (and only
        // if a table advanced). Each table's state was staged *after* its sink
        // succeeded, so a failed table contributed nothing to flush — the
        // crash-replay guarantee is preserved. (Valkey already wrote through
        // per advance; flush is a no-op there.)
        self.state.flush()?;
        tracing::info!(
            target: "bifrost::round",
            event = "round_complete",
            synced = report.synced.len(),
            failed = report.failed.len(),
            "sync round complete",
        );

        // Close the loop: announce each ingested table as an Airflow asset
        // event so downstream DAGs can schedule on it. Emission failures are
        // logged but never fail the round (the graph is already up to date).
        if let Some((emitter, prefix)) = &self.asset_emitter {
            let events = round_to_events(prefix, &report);
            if !events.is_empty() {
                match emitter.emit(&events).await {
                    Ok(()) => tracing::info!(
                        target: "bifrost::airflow",
                        event = "assets_emitted",
                        count = events.len(),
                        "emitted asset events",
                    ),
                    Err(e) => tracing::warn!(
                        target: "bifrost::airflow",
                        event = "asset_emit_failed",
                        error = %e,
                        "asset emission failed (round already applied)",
                    ),
                }
            }
        }

        // Fold the round into the cumulative counters and emit them as a
        // structured `metrics` event (target `metrics`). Purely additive — it
        // rides *after* the existing `round_complete`/`assets_emitted` events
        // and never touches them, so the log contract is unchanged.
        self.metrics.record_round(&report);
        let m = &self.metrics;
        tracing::info!(
            target: "metrics",
            event = "metrics",
            rounds_total = m.rounds_total,
            rounds_failed = m.rounds_failed,
            tables_synced = m.tables_synced,
            rows_ingested = m.rows_ingested,
            vertices_written = m.vertices_written,
            edges_written = m.edges_written,
            deletes = m.deletes,
            assets_emitted = m.assets_emitted,
            "cumulative metrics",
        );
        Ok(report)
    }

    async fn sync_table(
        &mut self,
        compiled: &CompiledTable,
    ) -> Result<Option<(usize, SinkReport)>> {
        // H6: transforms were compiled once at construction. A compile error
        // (empty pipeline / unknown custom step) is deterministic; reproduce it
        // every round so the table keeps landing in `report.failed`, as before.
        let transforms = compiled
            .pipeline
            .as_ref()
            .map_err(|msg| BifrostError::Other(msg.clone()))?;
        let spec = TableSpec {
            table: compiled.table.clone(),
            mode: compiled.mode,
            // The pushdown ask: exactly the columns this table's pipeline reads,
            // plus the mapping's optional row filter. A source that cannot prune
            // ignores both and returns everything, as before.
            columns: compiled.columns.clone(),
            filter: compiled.filter.clone(),
        };

        // M3: borrow the recorded state from the in-memory cache instead of
        // deep-cloning it (incl. a possibly-large `done_files`) every poll. An
        // idle round now does zero allocation here. `prev` only borrows
        // `self.state` immutably to build the source fetch below; its borrow
        // ends at the `fetch_new` call (NLL), well before the mutable
        // `advance_staged`. The never-synced case can't borrow the cache, so it
        // falls back to a stack-local default (cheap: an empty `done_files`).
        let default;
        let prev = match self.state.get_ref(&compiled.table) {
            Some(state) => state,
            None => {
                default = TableState::default();
                &default
            }
        };
        let Some((batches, mut new_state)) = self.source.fetch_new(&spec, prev).await? else {
            return Ok(None);
        };
        // Stamp WHICH CDC backend served this table onto the state the round is
        // about to persist, so `state_json` carries the decision (and any
        // fallback, with the engine's reason) — not just the log. Done for every
        // changelog-capable mode; `Full` has no CDC verdict to record.
        if matches!(compiled.mode, SyncMode::Incremental | SyncMode::Cdc) {
            new_state.cdc = self.source.cdc_capability(&compiled.table);
        }
        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();

        let ctx = TableContext {
            table: compiled.table.clone(),
        };
        let mut mutation = GraphMutation::default();
        // Changelog (CDC) mode: split the round's rows by `_change_type` into
        // upserts and deletes before transforming. `partition_change_rows`
        // strips the sentinel column, so upsert rows transform exactly like an
        // append-only batch and a non-CDC batch passes straight through.
        let change = partition_change_rows(&batches)?;
        for transform in transforms.iter() {
            mutation.merge(transform.apply(&ctx, &change.upserts)?);
        }
        mutation.deletes = deletes_from_change_rows(transforms, &ctx, &change.deletes)?;
        mutation.edge_deletes = edge_deletes_from_change_rows(transforms, &ctx, &change.deletes)?;

        let sink_report = if mutation.is_empty() {
            SinkReport::default()
        } else {
            self.sink.apply(&mutation).await?
        };

        // Only after the sink succeeded — crash here replays (MERGE dedupes).
        // Staged in memory; the round flushes the file once at the end (H8).
        self.state
            .advance_staged(&compiled.table, new_state)
            .await?;
        Ok(Some((rows, sink_report)))
    }

    /// Poll forever. The sleep between rounds is the base `interval` on the
    /// happy path, but a fully-failed (no-progress) round backs the next sleep
    /// off exponentially up to `interval * 10` (see [`Backoff`]); the first
    /// round that makes progress resets it. Errors are carried in the round
    /// reports handed to `on_round`; return `false` from it to stop.
    ///
    /// Signature-compatible with the original fixed-interval `run`; for a
    /// custom cap use [`run_with_backoff`](Self::run_with_backoff).
    pub async fn run(
        &mut self,
        interval: Duration,
        on_round: impl FnMut(RoundReport) -> bool,
    ) -> Result<()> {
        // Default cap: 10× the base interval (plan.md §7).
        self.run_with_backoff(
            Backoff::new(interval, interval.saturating_mul(10)),
            on_round,
        )
        .await
    }

    /// Poll forever with an explicit [`Backoff`] schedule. The happy-path
    /// cadence is `backoff.interval()`; consecutive no-progress rounds grow the
    /// inter-round sleep up to `backoff.max()`, resetting on the first round
    /// that makes progress. `run` is this with a `interval … interval*10`
    /// default. Return `false` from `on_round` to stop.
    pub async fn run_with_backoff(
        &mut self,
        mut backoff: Backoff,
        mut on_round: impl FnMut(RoundReport) -> bool,
    ) -> Result<()> {
        loop {
            let report = self.sync_once().await?;
            let failed = report.is_failure();
            if !on_round(report) {
                return Ok(());
            }
            let sleep = backoff.next_sleep(failed);
            if sleep > backoff.interval() {
                tracing::warn!(
                    target: "bifrost::round",
                    event = "backoff",
                    sleep_secs = sleep.as_secs(),
                    interval_secs = backoff.interval().as_secs(),
                    "round made no progress — backing off before retry",
                );
            }
            tokio::time::sleep(sleep).await;
        }
    }
}

/// The **streaming/delta** entry point: watch an Iceberg warehouse and inject
/// only the rows added since the last seen snapshot, directly into FalkorDB —
/// pure Rust, no Spark. A thin wrapper that wires an [`IcebergSource`] (in the
/// table's [`SyncMode`], so `incremental` ⇒ the manifest-level delta scan), a
/// [`FalkorDbSink`], and a file-backed [`SyncState`] cursor into a [`Listener`]
/// and polls forever at `interval` (with the default exponential backoff on
/// idle/failed rounds).
///
/// The `MappingSpec`'s tables should use `mode: incremental` to get the delta
/// path (`full` re-reads the whole table on every snapshot change). The cursor
/// (per-table last-ingested `snapshot_id`) is persisted to `state_path` after
/// each successful round, so a restart resumes from where it left off.
///
/// Return `false` from `on_round` to stop the loop.
///
/// ```no_run
/// # #[cfg(feature = "skade")]
/// # async fn demo() -> knut_bifrost::Result<()> {
/// use std::time::Duration;
/// use knut_bifrost::{MappingSpec, listener::watch_iceberg_delta};
///
/// let spec = MappingSpec::load("bifrost.yml")?; // tables: mode: incremental
/// watch_iceberg_delta(
///     spec,
///     "/var/lib/lake",            // skade Iceberg warehouse dir
///     "redis://127.0.0.1:6379",   // FalkorDB
///     "social",                   // graph name
///     "/var/lib/bifrost-state.json",
///     Duration::from_secs(15),
///     |round| { println!("{round:?}"); true },
/// ).await
/// # }
/// ```
#[cfg(feature = "skade")]
#[allow(clippy::too_many_arguments)]
pub async fn watch_iceberg_delta(
    spec: MappingSpec,
    warehouse_dir: impl Into<std::path::PathBuf>,
    falkordb_url: &str,
    graph: impl Into<String>,
    state_path: impl Into<std::path::PathBuf>,
    interval: Duration,
    on_round: impl FnMut(RoundReport) -> bool,
) -> Result<()> {
    watch_iceberg_delta_with_identity(spec, warehouse_dir, falkordb_url, None, graph, state_path, interval, on_round).await
}

/// [`watch_iceberg_delta`] with an optional mutual-TLS client identity, for a `rediss://`
/// FalkorDB running `tls-auth-clients yes`. `None` is the plain path, unchanged.
#[cfg(feature = "skade")]
#[allow(clippy::too_many_arguments)]
pub async fn watch_iceberg_delta_with_identity(
    spec: MappingSpec,
    warehouse_dir: impl Into<std::path::PathBuf>,
    falkordb_url: &str,
    tls: Option<graphar_falkordb::FalkorTls>,
    graph: impl Into<String>,
    state_path: impl Into<std::path::PathBuf>,
    interval: Duration,
    on_round: impl FnMut(RoundReport) -> bool,
) -> Result<()> {
    use crate::{sink::FalkorDbSink, source::IcebergSource};
    let mut listener = Listener::new(
        spec,
        TransformRegistry::new(),
        IcebergSource::new(warehouse_dir),
        FalkorDbSink::connect_with_identity(falkordb_url, graph, tls).await?,
        SyncState::open(state_path.into())?,
    );
    listener.run(interval, on_round).await
}