knut-bifrost 0.3.0

//! **PuppyGraph `schema.json`** as a bifröst mapping.
//!
//! PuppyGraph maps graph-first: a `node` list and an `edge` list, each element
//! naming the table it comes from. bifröst maps table-first: per table, a
//! pipeline of steps. [`mapping_from_puppygraph`] turns the one into the other —
//! each node becomes a `vertex` step and each edge an `edge` step of its table,
//! tables in first-seen order — so the listener, the changelog path and the
//! sinks run it unchanged. It is the format of PuppyGraph's current docs:
//!
//! ```json
//! { "catalog": [ … ],
//!   "node": [ { "label": "person",
//!               "dataSourceGroup": { "externalDataSource": {
//!                   "catalog": "lake", "schema": "modern", "table": "person",
//!                   "mappedField": [ { "sourceFieldName": "id", "targetFieldName": "id" }, … ] } },
//!               "id": [ { "name": "id", "type": "STRING" } ],
//!               "attribute": [ { "name": "name", "type": "STRING" } ] } ],
//!   "edge": [ { "label": "knows", "fromNodeLabel": "person", "toNodeLabel": "person",
//!               "dataSourceGroup": { … "table": "knows" … },
//!               "id": [ { "name": "id" } ], "fromKey": [ { "name": "from_id" } ],
//!               "toKey": [ { "name": "to_id" } ], "attribute": [ … ] } ] }
//! ```
//!
//! How each part lands:
//!
//! * the table is `schema.table` — `ns.name` in a REST catalog, `<root>/ns/name`
//!   on S3; the `catalog` section is not read (the lake comes from the command
//!   line: `--rest-uri`, `--s3-root`, `--warehouse`);
//! * `mappedField` is the source-column → property rename; a field it does not
//!   name keeps its name;
//! * a node's `id` is its MERGE key; an edge's `fromKey`/`toKey` are matched
//!   against its end nodes' `id` properties, and its `id` is the edge key, so two
//!   rows between the same nodes are two (parallel) edges;
//! * `attribute`s are the properties; `type`s are not read — the table's own
//!   Arrow types are written.
//!
//! **Not yet:** composite keys (an `id`, `fromKey` or `toKey` of more than one
//! field) and the older `oneToOne`/`tableSource` layout — both are refused by
//! name rather than half-mapped.

use std::collections::HashMap;

use serde::Deserialize;

use crate::error::{BifrostError, Result};
use crate::mapping::{MappingSpec, PipelineStep, SyncMode, TableMapping};
use crate::transform::{EdgeTransform, EndpointSpec, VertexTransform};

#[derive(Debug, Deserialize)]
struct Schema {
    #[serde(default)]
    node: Vec<Node>,
    #[serde(default)]
    edge: Vec<Edge>,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct Node {
    label: String,
    data_source_group: DataSourceGroup,
    id: Vec<FieldRef>,
    #[serde(default)]
    attribute: Vec<FieldRef>,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct Edge {
    label: String,
    from_node_label: String,
    to_node_label: String,
    data_source_group: DataSourceGroup,
    #[serde(default)]
    id: Vec<FieldRef>,
    from_key: Vec<FieldRef>,
    to_key: Vec<FieldRef>,
    #[serde(default)]
    attribute: Vec<FieldRef>,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DataSourceGroup {
    external_data_source: ExternalDataSource,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ExternalDataSource {
    #[serde(default = "yes")]
    enabled: bool,
    #[serde(default)]
    schema: Option<String>,
    table: String,
    #[serde(default)]
    mapped_field: Vec<MappedField>,
}

fn yes() -> bool {
    true
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct MappedField {
    source_field_name: String,
    target_field_name: String,
}

#[derive(Debug, Deserialize)]
struct FieldRef {
    name: String,
}

impl ExternalDataSource {
    fn table(&self) -> String {
        match &self.schema {
            Some(s) if !s.is_empty() => format!("{s}.{}", self.table),
            _ => self.table.clone(),
        }
    }

    /// Property → source column.
    fn sources(&self) -> HashMap<&str, &str> {
        self.mapped_field.iter().map(|m| (m.target_field_name.as_str(), m.source_field_name.as_str())).collect()
    }
}

/// Source column → property, for the fields whose names differ.
fn renames(src: &ExternalDataSource) -> HashMap<String, String> {
    src.mapped_field
        .iter()
        .filter(|m| m.source_field_name != m.target_field_name)
        .map(|m| (m.source_field_name.clone(), m.target_field_name.clone()))
        .collect()
}

/// The one field of a key, or a refusal naming what is not supported yet.
fn single<'a>(what: &str, label: &str, key: &'a [FieldRef]) -> Result<&'a str> {
    match key {
        [one] => Ok(one.name.as_str()),
        [] => Err(BifrostError::Other(format!("puppygraph: {what} of `{label}` is empty"))),
        _ => Err(BifrostError::Other(format!(
            "puppygraph: {what} of `{label}` has {} fields - composite keys are not supported yet",
            key.len()
        ))),
    }
}

/// A PuppyGraph `schema.json` as a [`MappingSpec`], every table in `mode`.
pub fn mapping_from_puppygraph(json: &str, mode: SyncMode) -> Result<MappingSpec> {
    let raw: serde_json::Value = serde_json::from_str(json).map_err(|e| BifrostError::Other(format!("puppygraph: {e}")))?;
    if raw.get("graph").is_some() || json.contains("\"oneToOne\"") {
        return Err(BifrostError::Other(
            "puppygraph: this is the older `graph`/`oneToOne`/`tableSource` layout - only the current \
             `node`/`edge` + `dataSourceGroup` layout is read"
                .into(),
        ));
    }
    let schema: Schema = serde_json::from_value(raw).map_err(|e| BifrostError::Other(format!("puppygraph: {e}")))?;

    let mut tables: Vec<TableMapping> = Vec::new();
    let mut push = |table: String, step: PipelineStep| match tables.iter_mut().find(|t| t.table == table) {
        Some(t) => t.pipeline.push(step),
        None => tables.push(TableMapping { table, mode, pipeline: vec![step], columns: Vec::new(), filter: None }),
    };

    // A node's id PROPERTY, by label - what an edge's endpoint is matched on.
    let mut node_id: HashMap<&str, &str> = HashMap::new();
    for n in &schema.node {
        let src = &n.data_source_group.external_data_source;
        let id = single("id", &n.label, &n.id)?;
        node_id.insert(n.label.as_str(), id);
        if !src.enabled {
            continue;
        }
        let sources = src.sources();
        let col = |prop: &str| sources.get(prop).copied().unwrap_or(prop).to_string();
        push(
            src.table(),
            PipelineStep::Vertex {
                vertex: VertexTransform {
                    label: n.label.clone(),
                    id_column: col(id),
                    properties: n.attribute.iter().filter(|a| a.name != id).map(|a| col(&a.name)).collect(),
                    rename: renames(src),
                },
            },
        );
    }

    for e in &schema.edge {
        let src = &e.data_source_group.external_data_source;
        if !src.enabled {
            continue;
        }
        let end = |label: &str| {
            node_id.get(label).map(|s| s.to_string()).ok_or_else(|| {
                BifrostError::Other(format!("puppygraph: edge `{}` names node `{label}`, which the schema does not define", e.label))
            })
        };
        let sources = src.sources();
        let col = |prop: &str| sources.get(prop).copied().unwrap_or(prop).to_string();
        let from = col(single("fromKey", &e.label, &e.from_key)?);
        let to = col(single("toKey", &e.label, &e.to_key)?);
        let key = match e.id.as_slice() {
            [] => Vec::new(),
            _ => vec![col(single("id", &e.label, &e.id)?)],
        };
        // The endpoint columns are the edge's `src`/`dst`; as attributes they are
        // also properties, as PuppyGraph writes them.
        let properties: Vec<String> =
            e.attribute.iter().map(|a| col(&a.name)).filter(|c| !key.contains(c)).collect();
        push(
            src.table(),
            PipelineStep::Edge {
                edge: EdgeTransform {
                    edge_type: e.label.clone(),
                    src: EndpointSpec { label: e.from_node_label.clone(), column: from, id_column: Some(end(&e.from_node_label)?) },
                    dst: EndpointSpec { label: e.to_node_label.clone(), column: to, id_column: Some(end(&e.to_node_label)?) },
                    properties,
                    rename: renames(src),
                    key,
                    // An edge table: each row is its own edge.
                    replace: Some(false),
                },
            },
        );
    }
    if tables.is_empty() {
        return Err(BifrostError::Other("puppygraph: the schema maps no enabled node or edge".into()));
    }
    Ok(MappingSpec { version: 1, tables })
}

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

    /// PuppyGraph's own "modern" example, as in their Iceberg getting-started.
    pub(crate) const MODERN: &str = r#"{
      "catalog": [ { "name": "iceberg_data", "type": "iceberg",
                     "metastore": { "type": "rest", "uri": "http://iceberg-rest:8181" } } ],
      "node": [
        { "label": "software",
          "dataSourceGroup": { "externalDataSource": { "enabled": true, "catalog": "iceberg_data",
              "schema": "modern", "table": "software",
              "mappedField": [ { "sourceFieldName": "id", "targetFieldName": "id" },
                               { "sourceFieldName": "name", "targetFieldName": "name" },
                               { "sourceFieldName": "lang", "targetFieldName": "lang" } ] } },
          "id": [ { "name": "id", "type": "STRING" } ],
          "attribute": [ { "name": "name", "type": "STRING" }, { "name": "lang", "type": "STRING" } ] },
        { "label": "person",
          "dataSourceGroup": { "externalDataSource": { "enabled": true, "catalog": "iceberg_data",
              "schema": "modern", "table": "person",
              "mappedField": [ { "sourceFieldName": "id", "targetFieldName": "id" },
                               { "sourceFieldName": "name", "targetFieldName": "name" },
                               { "sourceFieldName": "age", "targetFieldName": "age" } ] } },
          "id": [ { "name": "id", "type": "STRING" } ],
          "attribute": [ { "name": "name", "type": "STRING" }, { "name": "age", "type": "INT" } ] } ],
      "edge": [
        { "label": "created", "fromNodeLabel": "person", "toNodeLabel": "software",
          "dataSourceGroup": { "externalDataSource": { "enabled": true, "catalog": "iceberg_data",
              "schema": "modern", "table": "created",
              "mappedField": [ { "sourceFieldName": "id", "targetFieldName": "id" },
                               { "sourceFieldName": "from_id", "targetFieldName": "from_id" },
                               { "sourceFieldName": "to_id", "targetFieldName": "to_id" },
                               { "sourceFieldName": "weight", "targetFieldName": "weight" } ] } },
          "id": [ { "name": "id", "type": "STRING" } ],
          "fromKey": [ { "name": "from_id", "type": "STRING" } ],
          "toKey": [ { "name": "to_id", "type": "STRING" } ],
          "attribute": [ { "name": "from_id", "type": "STRING" }, { "name": "to_id", "type": "STRING" },
                         { "name": "weight", "type": "DOUBLE" } ] },
        { "label": "knows", "fromNodeLabel": "person", "toNodeLabel": "person",
          "dataSourceGroup": { "externalDataSource": { "enabled": true, "catalog": "iceberg_data",
              "schema": "modern", "table": "knows",
              "mappedField": [ { "sourceFieldName": "id", "targetFieldName": "id" },
                               { "sourceFieldName": "from_id", "targetFieldName": "from_id" },
                               { "sourceFieldName": "to_id", "targetFieldName": "to_id" },
                               { "sourceFieldName": "weight", "targetFieldName": "weight" } ] } },
          "id": [ { "name": "id", "type": "STRING" } ],
          "fromKey": [ { "name": "from_id", "type": "STRING" } ],
          "toKey": [ { "name": "to_id", "type": "STRING" } ],
          "attribute": [ { "name": "from_id", "type": "STRING" }, { "name": "to_id", "type": "STRING" },
                         { "name": "weight", "type": "DOUBLE" } ] } ]
    }"#;

    #[test]
    fn the_modern_example_maps_to_four_tables() {
        let m = mapping_from_puppygraph(MODERN, SyncMode::Cdc).unwrap();
        let tables: Vec<&str> = m.tables.iter().map(|t| t.table.as_str()).collect();
        assert_eq!(tables, vec!["modern.software", "modern.person", "modern.created", "modern.knows"]);
        let PipelineStep::Vertex { vertex } = &m.tables[1].pipeline[0] else { panic!("person is a vertex") };
        assert_eq!((vertex.label.as_str(), vertex.id_column.as_str()), ("person", "id"));
        assert_eq!(vertex.properties, vec!["name", "age"]);
        let PipelineStep::Edge { edge } = &m.tables[3].pipeline[0] else { panic!("knows is an edge") };
        assert_eq!((edge.src.label.as_str(), edge.src.column.as_str(), edge.src.id_column.as_deref()), ("person", "from_id", Some("id")));
        assert_eq!(edge.key, vec!["id"]);
        assert_eq!(edge.properties, vec!["from_id", "to_id", "weight"]);
        assert!(m.tables.iter().all(|t| t.mode == SyncMode::Cdc));
    }

    #[test]
    fn mapped_fields_rename_source_columns() {
        let json = MODERN.replace(r#""sourceFieldName": "age", "targetFieldName": "age""#, r#""sourceFieldName": "age_years", "targetFieldName": "age""#);
        let m = mapping_from_puppygraph(&json, SyncMode::Full).unwrap();
        let PipelineStep::Vertex { vertex } = &m.tables[1].pipeline[0] else { panic!() };
        assert_eq!(vertex.properties, vec!["name", "age_years"]);
        assert_eq!(vertex.rename.get("age_years").map(String::as_str), Some("age"));
    }

    #[test]
    fn composite_keys_and_the_old_layout_are_refused_by_name() {
        let composite = MODERN.replacen(r#""id": [ { "name": "id", "type": "STRING" } ],"#, r#""id": [ { "name": "id" }, { "name": "name" } ],"#, 1);
        let e = mapping_from_puppygraph(&composite, SyncMode::Cdc).unwrap_err().to_string();
        assert!(e.contains("composite keys"), "{e}");
        let old = r#"{ "catalogs": [], "graph": { "vertices": [ { "label": "p", "oneToOne": {} } ] } }"#;
        let e = mapping_from_puppygraph(old, SyncMode::Cdc).unwrap_err().to_string();
        assert!(e.contains("older"), "{e}");
    }

    #[test]
    fn an_edge_to_an_undefined_node_is_refused() {
        let json = MODERN.replace(r#""toNodeLabel": "software""#, r#""toNodeLabel": "robot""#);
        let e = mapping_from_puppygraph(&json, SyncMode::Cdc).unwrap_err().to_string();
        assert!(e.contains("robot"), "{e}");
    }
}