knut-bifrost 0.1.3

//! Graph sinks. The FalkorDB sink upserts with `MERGE`, keyed on the
//! mutation's id columns, so replaying a round after a crash is idempotent.

use std::collections::HashSet;

use arrow_array::{Array, RecordBatch};
use async_trait::async_trait;
use graphar_falkordb::{FalkorDbLoader, cypher::arrow_value_to_cypher};

use crate::{
    error::{BifrostError, Result},
    mutation::{EdgeSet, GraphMutation, SinkReport, VertexDelete, VertexSet},
};

#[async_trait]
pub trait GraphSink: Send {
    async fn apply(&mut self, mutation: &GraphMutation) -> Result<SinkReport>;
}

/// Upserts a [`GraphMutation`] into FalkorDB.
pub struct FalkorDbSink {
    loader: FalkorDbLoader,
    batch_size: usize,
}

impl FalkorDbSink {
    pub async fn connect(url: &str, graph: impl Into<String>) -> Result<Self> {
        Ok(Self {
            loader: FalkorDbLoader::connect(url, graph).await?,
            batch_size: graphar_falkordb::DEFAULT_BATCH_SIZE,
        })
    }

    pub fn with_batch_size(mut self, n: usize) -> Self {
        self.batch_size = n.max(1);
        self
    }

    async fn upsert_vertices(&mut self, set: &VertexSet) -> Result<usize> {
        let rows = rows_as_cypher_maps(&set.batch)?;
        let mut written = 0;
        for chunk in rows.chunks(self.batch_size) {
            // MERGE on the id property, then overwrite all properties.
            let query = format!(
                "UNWIND [{rows}] AS r MERGE (n:`{label}` {{`{id}`: r.`{id}`}}) SET n += r",
                rows = chunk.join(", "),
                label = set.label,
                id = set.id_column,
            );
            self.loader.run_query(&query).await?;
            written += chunk.len();
        }
        Ok(written)
    }

    async fn upsert_edges(&mut self, set: &EdgeSet) -> Result<usize> {
        let rows = rows_as_cypher_maps(&set.batch)?;
        let mut written = 0;
        for chunk in rows.chunks(self.batch_size) {
            // MERGE (not MATCH) on the endpoints: tables sync in arbitrary
            // order, so an edge may arrive before its endpoint's table.
            // The placeholder node is enriched later by the vertex MERGE
            // on the same id key.
            let query = format!(
                "UNWIND [{rows}] AS r \
                 MERGE (a:`{src_label}` {{`{src_id}`: r.src}}) \
                 MERGE (b:`{dst_label}` {{`{dst_id}`: r.dst}}) \
                 MERGE (a)-[e:`{etype}`]->(b) \
                 SET e += r.props",
                rows = chunk.join(", "),
                src_label = set.src_label,
                src_id = set.src_id_column,
                dst_label = set.dst_label,
                dst_id = set.dst_id_column,
                etype = set.edge_type,
            );
            self.loader.run_query(&query).await?;
            written += chunk.len();
        }
        Ok(written)
    }

    /// `DETACH DELETE` every node named by a [`VertexDelete`]. Deleting the
    /// node drops its incident edges too, so a changelog `DELETE` on a vertex
    /// table cleans up its relationships for free. Idempotent: a `MATCH` that
    /// finds nothing (replay after a crash, or a never-ingested id) is a no-op.
    async fn delete_vertices(
        &mut self,
        set: &VertexDelete,
        upserted: &HashSet<(String, String)>,
    ) -> Result<usize> {
        let ids = surviving_delete_ids(set, upserted)?;
        let mut deleted = 0;
        for chunk in ids.chunks(self.batch_size) {
            self.loader
                .run_query(&detach_delete_query(&set.label, &set.id_column, chunk))
                .await?;
            deleted += chunk.len();
        }
        Ok(deleted)
    }
}

#[async_trait]
impl GraphSink for FalkorDbSink {
    async fn apply(&mut self, mutation: &GraphMutation) -> Result<SinkReport> {
        let mut report = SinkReport::default();
        // Keys upserted this round; a delete of the same (label, id) must be
        // suppressed so an in-place UPDATE (UPDATE_BEFORE + UPDATE_AFTER on a
        // stable key) doesn't destroy the node it just rewrote.
        let upserted = upserted_vertex_keys(mutation)?;
        // Vertices strictly before edges: edges MATCH their endpoints.
        for set in &mutation.vertices {
            let n = self.upsert_vertices(set).await?;
            tracing::debug!(
                target: "bifrost::sink",
                event = "vertices_upserted",
                label = %set.label,
                count = n,
                "MERGE vertices",
            );
            report.vertices_written += n;
        }
        for set in &mutation.edges {
            let n = self.upsert_edges(set).await?;
            tracing::debug!(
                target: "bifrost::sink",
                event = "edges_upserted",
                edge_type = %set.edge_type,
                count = n,
                "MERGE edges",
            );
            report.edges_written += n;
        }
        // Deletes last: an upsert + delete of the same key in one changelog
        // round must resolve to "gone" (see GraphMutation phase ordering).
        for set in &mutation.deletes {
            let n = self.delete_vertices(set, &upserted).await?;
            tracing::debug!(
                target: "bifrost::sink",
                event = "vertices_deleted",
                label = %set.label,
                count = n,
                "DETACH DELETE vertices",
            );
            report.deleted += n;
        }
        tracing::info!(
            target: "bifrost::sink",
            event = "sink_applied",
            vertices = report.vertices_written,
            edges = report.edges_written,
            deleted = report.deleted,
            "graph mutation applied",
        );
        Ok(report)
    }
}

/// Build the `DETACH DELETE` query for one chunk of pre-rendered id literals.
/// Pure (no IO) so the query shape is unit-testable without a database.
fn detach_delete_query(label: &str, id_column: &str, ids: &[String]) -> String {
    format!(
        "UNWIND [{ids}] AS r \
         MATCH (n:`{label}` {{`{id}`: r}}) DETACH DELETE n",
        ids = ids.join(", "),
        id = id_column,
    )
}

/// Render the id column of a [`VertexDelete`] batch as a list of Cypher scalar
/// literals — one per node to delete. Only `id_column` is read; other columns
/// (e.g. a `DELETE` row's stale property values) are ignored.
fn delete_ids_as_cypher(set: &VertexDelete) -> Result<Vec<String>> {
    let schema = set.batch.schema();
    let idx = schema
        .index_of(&set.id_column)
        .map_err(|_| BifrostError::ColumnNotFound(set.id_column.clone()))?;
    let col = set.batch.column(idx).as_ref();
    let mut ids = Vec::with_capacity(set.batch.num_rows());
    for row in 0..set.batch.num_rows() {
        ids.push(arrow_value_to_cypher(col, row).map_err(BifrostError::Falkor)?);
    }
    Ok(ids)
}

/// Collect the `(label, id-literal)` of every vertex this round upserts.
///
/// An in-place UPDATE on a stable primary key emits `UPDATE_BEFORE(id=X)` +
/// `UPDATE_AFTER(id=X)`; CDC routing sends the BEFORE row to `deletes` and the
/// AFTER row to `vertices`. Applying all deletes last would then DETACH-DELETE
/// the node the upsert just wrote. By recording upserted keys we can let the
/// upsert win over a same-key delete in the same round. (The same window can
/// arise in the pure-Rust skade path: an equality-delete of the old row + an
/// append of the new row for one id.)
fn upserted_vertex_keys(mutation: &GraphMutation) -> Result<HashSet<(String, String)>> {
    let mut keys = HashSet::new();
    for set in &mutation.vertices {
        let schema = set.batch.schema();
        let Ok(idx) = schema.index_of(&set.id_column) else {
            continue;
        };
        let col = set.batch.column(idx).as_ref();
        for row in 0..set.batch.num_rows() {
            let lit = arrow_value_to_cypher(col, row).map_err(BifrostError::Falkor)?;
            keys.insert((set.label.clone(), lit));
        }
    }
    Ok(keys)
}

/// Ids of a [`VertexDelete`] that should actually be deleted: every id *except*
/// those whose `(label, id)` was upserted this round (the upsert wins). Pure
/// (no IO) so the net-effect rule is unit-testable without a database.
fn surviving_delete_ids(
    set: &VertexDelete,
    upserted: &HashSet<(String, String)>,
) -> Result<Vec<String>> {
    Ok(delete_ids_as_cypher(set)?
        .into_iter()
        .filter(|id| !upserted.contains(&(set.label.clone(), id.clone())))
        .collect())
}

/// Render each row as a Cypher map literal.
///
/// Vertex batches become `{col: val, …}`. Edge batches become
/// `{src: …, dst: …, props: {…}}` so the property overwrite can't clobber
/// the endpoint keys.
pub(crate) fn rows_as_cypher_maps(batch: &RecordBatch) -> Result<Vec<String>> {
    let schema = batch.schema();
    let is_edge = schema.index_of("src").is_ok() && schema.index_of("dst").is_ok();
    let cols: Vec<(&str, &dyn Array)> = schema
        .fields()
        .iter()
        .zip(batch.columns())
        .map(|(f, c)| (f.name().as_str(), c.as_ref()))
        .collect();

    let mut rows = Vec::with_capacity(batch.num_rows());
    for row in 0..batch.num_rows() {
        if is_edge {
            let mut src = None;
            let mut dst = None;
            let mut props = Vec::new();
            for (name, array) in &cols {
                let val = arrow_value_to_cypher(*array, row).map_err(BifrostError::Falkor)?;
                match *name {
                    "src" => src = Some(val),
                    "dst" => dst = Some(val),
                    _ => props.push(format!("`{name}`: {val}")),
                }
            }
            rows.push(format!(
                "{{src: {}, dst: {}, props: {{{}}}}}",
                src.expect("src checked above"),
                dst.expect("dst checked above"),
                props.join(", "),
            ));
        } else {
            let parts: Vec<String> = cols
                .iter()
                .map(|(name, array)| {
                    let val = arrow_value_to_cypher(*array, row).map_err(BifrostError::Falkor)?;
                    Ok(format!("`{name}`: {val}"))
                })
                .collect::<Result<_>>()?;
            rows.push(format!("{{{}}}", parts.join(", ")));
        }
    }
    Ok(rows)
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_array::{Int64Array, StringArray};
    use arrow_schema::{DataType, Field, Schema};
    use std::sync::Arc;

    #[test]
    fn vertex_rows_render_as_flat_maps() {
        let batch = RecordBatch::try_new(
            Arc::new(Schema::new(vec![
                Field::new("user_id", DataType::Int64, false),
                Field::new("name", DataType::Utf8, false),
            ])),
            vec![
                Arc::new(Int64Array::from(vec![7])),
                Arc::new(StringArray::from(vec!["alice"])),
            ],
        )
        .unwrap();
        let rows = rows_as_cypher_maps(&batch).unwrap();
        assert_eq!(rows, vec!["{`user_id`: 7, `name`: 'alice'}"]);
        crate::testmatrix::emit(
            crate::testmatrix::COMPONENT,
            "sink::vertex_rows_render_as_flat_maps",
            rows == vec!["{`user_id`: 7, `name`: 'alice'}"],
            &format!("cypher={}", rows[0]),
        );
    }

    #[test]
    fn edge_rows_separate_endpoints_from_props() {
        let batch = RecordBatch::try_new(
            Arc::new(Schema::new(vec![
                Field::new("src", DataType::Int64, false),
                Field::new("dst", DataType::Int64, false),
                Field::new("since", DataType::Utf8, false),
            ])),
            vec![
                Arc::new(Int64Array::from(vec![1])),
                Arc::new(Int64Array::from(vec![2])),
                Arc::new(StringArray::from(vec!["2024"])),
            ],
        )
        .unwrap();
        let rows = rows_as_cypher_maps(&batch).unwrap();
        assert_eq!(rows, vec!["{src: 1, dst: 2, props: {`since`: '2024'}}"]);
        crate::testmatrix::emit(
            crate::testmatrix::COMPONENT,
            "sink::edge_rows_separate_endpoints_from_props",
            rows == vec!["{src: 1, dst: 2, props: {`since`: '2024'}}"],
            &format!("cypher={}", rows[0]),
        );
    }

    #[test]
    fn detach_delete_query_matches_on_id_and_detaches() {
        let q = detach_delete_query("Person", "user_id", &["7".into(), "8".into()]);
        assert_eq!(
            q,
            "UNWIND [7, 8] AS r MATCH (n:`Person` {`user_id`: r}) DETACH DELETE n"
        );
        crate::testmatrix::emit(
            crate::testmatrix::COMPONENT,
            "sink::detach_delete_query_matches_on_id_and_detaches",
            q == "UNWIND [7, 8] AS r MATCH (n:`Person` {`user_id`: r}) DETACH DELETE n",
            &format!("query={q}"),
        );
    }

    #[test]
    fn delete_ids_read_only_the_id_column() {
        // A DELETE/UPDATE_BEFORE row carries stale props; only the id is used.
        let set = VertexDelete {
            label: "Person".into(),
            id_column: "user_id".into(),
            batch: RecordBatch::try_new(
                Arc::new(Schema::new(vec![
                    Field::new("user_id", DataType::Int64, false),
                    Field::new("name", DataType::Utf8, false),
                ])),
                vec![
                    Arc::new(Int64Array::from(vec![7, 9])),
                    Arc::new(StringArray::from(vec!["stale", "gone"])),
                ],
            )
            .unwrap(),
        };
        let ids = delete_ids_as_cypher(&set).unwrap();
        assert_eq!(ids, vec!["7", "9"]);
        crate::testmatrix::emit(
            crate::testmatrix::COMPONENT,
            "sink::delete_ids_read_only_the_id_column",
            ids == vec!["7", "9"],
            &format!("ids={ids:?} (stale props ignored)"),
        );
    }

    #[test]
    fn in_place_update_does_not_delete_the_upserted_node() {
        // An in-place UPDATE on a stable key: UPDATE_AFTER(id=7) lands in
        // `vertices`, UPDATE_BEFORE(id=7) lands in `deletes`. Without net-effect
        // resolution the sink would DETACH-DELETE id=7 last, destroying the node
        // it just upserted. id=9 is a *true* delete (no matching upsert) and
        // must still be removed.
        let upsert = VertexSet {
            label: "Person".into(),
            id_column: "user_id".into(),
            batch: RecordBatch::try_new(
                Arc::new(Schema::new(vec![
                    Field::new("user_id", DataType::Int64, false),
                    Field::new("name", DataType::Utf8, false),
                ])),
                vec![
                    Arc::new(Int64Array::from(vec![7])),
                    Arc::new(StringArray::from(vec!["new-name"])),
                ],
            )
            .unwrap(),
        };
        let delete = VertexDelete {
            label: "Person".into(),
            id_column: "user_id".into(),
            batch: RecordBatch::try_new(
                Arc::new(Schema::new(vec![Field::new(
                    "user_id",
                    DataType::Int64,
                    false,
                )])),
                vec![Arc::new(Int64Array::from(vec![7, 9]))],
            )
            .unwrap(),
        };
        let mutation = GraphMutation {
            vertices: vec![upsert],
            edges: vec![],
            deletes: vec![delete.clone()],
        };

        let upserted = upserted_vertex_keys(&mutation).unwrap();
        let surviving = surviving_delete_ids(&delete, &upserted).unwrap();

        // id 7 was upserted → suppressed; id 9 is a real delete → survives.
        let pass = surviving == vec!["9"];
        assert!(pass, "expected only id 9 to be deleted, got {surviving:?}");
        crate::testmatrix::emit(
            crate::testmatrix::COMPONENT,
            "sink::in_place_update_does_not_delete_the_upserted_node",
            pass,
            &format!("surviving_deletes={surviving:?} (id 7 upsert wins over delete)"),
        );
    }
}