kglite 0.16.22

//! `from_records` — build a graph from an inline JSON records spec.
//!
//! A JSON-native sibling to [`crate::graph::blueprint::build`]: instead of a
//! blueprint pointing at CSV files on disk, the caller passes node and
//! connection records inline. Agent-authored graphs are JSON-native, so this
//! is the natural ingestion path for them.
//!
//! The spec shape (all `records` are arrays of flat JSON objects):
//!
//! ```json
//! {
//!   "nodes": [
//!     { "type": "Person", "id_field": "id", "title_field": "name",
//!       "conflict_handling": "update",
//!       "records": [ {"id": 1, "name": "Alice", "aliases": ["a", "b"]} ] }
//!   ],
//!   "connections": [
//!     { "type": "KNOWS", "source_type": "Person", "source_id_field": "from",
//!       "target_type": "Person", "target_id_field": "to",
//!       "records": [ {"from": 1, "to": 2, "since": 2020} ] }
//!   ],
//!   // `target_type` also takes a list, for a relationship whose range is a
//!   // union of types; `target_type_column` then names the field holding
//!   // each record's target type, and without it the listed types are
//!   // probed for the record's target id.
//!   "on_missing_endpoint": "vivify"
//! }
//! ```
//!
//! Those key sets are closed: an unknown key at the top level or in any node
//! or connection spec is refused (with a "did you mean" suggestion when one is
//! close), because a key this parser does not read is silently dropped — a
//! spec written with `"relationships"` would build zero edges and report
//! success.
//!
//! Column types are **inferred** from the record values (across all rows, via
//! [`DataFrame::from_cypher_rows`]), so a JSON array becomes a native list
//! property, an integer an `Int64`, etc. All graph mutation reuses the
//! existing engine — [`maintain::add_nodes`] and [`maintain::add_connections`]
//! (the latter's Pass-A/B/C endpoint vivification fills in missing endpoints) —
//! so there is no duplicated mutation logic.

use crate::datatypes::values::{DataFrame, Value};
use crate::graph::mutation::maintain;
use crate::graph::DirGraph;
use serde_json::Value as Json;

/// Summary of a `from_records` build.
#[derive(Debug, Default, Clone)]
pub struct RecordsReport {
    pub nodes_added: usize,
    pub edges_added: usize,
    /// Connection records omitted by `on_missing_endpoint: "drop"` because
    /// at least one endpoint was null or absent from the graph.
    pub edges_dropped_missing_endpoint: usize,
    pub node_types: Vec<String>,
    pub connection_types: Vec<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MissingEndpointPolicy {
    Vivify,
    Drop,
    Error,
}

/// Build (or extend) `graph` from an inline JSON records spec. See the module
/// docs for the spec shape. Returns a per-build summary.
pub fn from_records(graph: &mut DirGraph, spec: &Json) -> Result<RecordsReport, String> {
    let obj = spec
        .as_object()
        .ok_or_else(|| "from_records: top-level JSON must be an object".to_string())?;

    reject_unknown_keys(obj, TOP_LEVEL_KEYS, "from_records")?;

    let endpoint_policy = parse_endpoint_policy(obj.get("on_missing_endpoint"))?;
    if endpoint_policy == MissingEndpointPolicy::Error {
        // `error` is an all-or-nothing ingestion mode. Build against a
        // backend-aware transaction fork and publish it only after every node
        // and connection spec has succeeded.
        let mut working = graph.fork_transaction();
        let report = load_records(&mut working, obj, endpoint_policy)?;
        *graph = working;
        return Ok(report);
    }

    load_records(graph, obj, endpoint_policy)
}

fn load_records(
    graph: &mut DirGraph,
    obj: &serde_json::Map<String, Json>,
    endpoint_policy: MissingEndpointPolicy,
) -> Result<RecordsReport, String> {
    let mut report = RecordsReport::default();
    let mut declared_labels: Vec<(String, Vec<String>)> = Vec::new();

    // ── Nodes ────────────────────────────────────────────────────────────
    if let Some(nodes) = obj.get("nodes") {
        let arr = nodes
            .as_array()
            .ok_or_else(|| "from_records: 'nodes' must be an array".to_string())?;
        for (i, node_spec) in arr.iter().enumerate() {
            let labels = load_node_spec(graph, node_spec, i, &mut report)?;
            declared_labels.extend(labels);
        }
    }

    // ── Connections ──────────────────────────────────────────────────────
    if let Some(conns) = obj.get("connections") {
        let arr = conns
            .as_array()
            .ok_or_else(|| "from_records: 'connections' must be an array".to_string())?;
        for (i, conn_spec) in arr.iter().enumerate() {
            load_connection_spec(graph, conn_spec, i, endpoint_policy, &mut report)?;
        }
    }

    // After the connections, not between the two phases: `vivify` mints a stub
    // for an endpoint no record supplied, and a spec's labels cover every node
    // of the type it declares — otherwise `MATCH (:Place)` misses exactly the
    // nodes that arrived as an endpoint rather than as a record.
    stamp_declared_labels(graph, &declared_labels)?;

    Ok(report)
}

/// Stamp each node spec's declared `labels` on every node of its type.
///
/// One bulk call per (type, label) — see
/// [`DirGraph::add_node_labels_bulk`](crate::graph::schema::DirGraph::add_node_labels_bulk)
/// for why the per-node loop is quadratic. Interning is validated for the
/// whole set first, so a name over the interner's ceiling refuses the load
/// rather than leaving the type half-labelled.
fn stamp_declared_labels(
    graph: &mut DirGraph,
    declared: &[(String, Vec<String>)],
) -> Result<(), String> {
    for (node_type, labels) in declared {
        graph
            .interner
            .validate_names(labels.iter().map(String::as_str))
            .map_err(|e| format!("from_records: node '{node_type}': labels: {e}"))?;
        let Some(nodes) = graph.type_indices.get(node_type) else {
            continue;
        };
        let indices: Vec<petgraph::graph::NodeIndex> = nodes.iter().collect();
        for label in labels {
            let key = graph.interner.get_or_intern(label);
            graph.add_node_labels_bulk(&indices, key);
        }
    }
    Ok(())
}

/// Returns the spec's `(type, labels)` when it declares any, for the stamping
/// pass that runs once every spec — nodes *and* connections — has loaded.
fn load_node_spec(
    graph: &mut DirGraph,
    spec: &Json,
    idx: usize,
    report: &mut RecordsReport,
) -> Result<Option<(String, Vec<String>)>, String> {
    let ctx = || format!("from_records: nodes[{}]", idx);
    // A non-object spec falls through to the required-field error below.
    if let Some(obj) = spec.as_object() {
        reject_unknown_keys(obj, NODE_SPEC_KEYS, &ctx())?;
    }
    let node_type = required_str(spec, "type", &ctx)?;
    let id_field = required_str(spec, "id_field", &ctx)?;
    let title_field = optional_str(spec, "title_field");
    let conflict_handling = optional_str(spec, "conflict_handling");

    let labels = optional_str_list(spec, "labels", &ctx)?;

    let records = records_array(spec, &ctx)?;
    if records.is_empty() {
        // No rows to load, but the declaration still stands: every node of
        // this type carries its labels, and those nodes can all arrive as
        // vivified edge endpoints. Returning before the labels dropped them
        // in silence — the blueprint builder keeps every spec through its
        // own stamping phase for exactly this reason.
        return Ok(labels.map(|labels| (node_type, labels)));
    }

    // The id field always leads the column order so a record missing it
    // still produces a (null id) row that add_nodes' validity check catches.
    let borrowed: Vec<&Json> = records.iter().collect();
    let (columns, rows) = records_to_columns_rows(&borrowed, &[&id_field], &ctx)?;
    let df = DataFrame::from_cypher_rows(columns, rows).map_err(|e| format!("{}: {}", ctx(), e))?;

    let rep = maintain::add_nodes(
        graph,
        df,
        node_type.clone(),
        id_field,
        title_field,
        conflict_handling,
    )
    .map_err(|e| format!("{}: {}", ctx(), e))?;

    report.nodes_added += rep.nodes_created + rep.nodes_updated;
    report.node_types.push(node_type.clone());
    Ok(labels.map(|labels| (node_type, labels)))
}

/// Read an optional array-of-strings field. Refuses a non-array or a
/// non-string element rather than skipping it: a spec that says `"labels":
/// "Human"` means to label something, and dropping it silently is the failure
/// this whole key exists to remove.
fn optional_str_list(
    spec: &Json,
    field: &str,
    ctx: &impl Fn() -> String,
) -> Result<Option<Vec<String>>, String> {
    let Some(value) = spec.get(field) else {
        return Ok(None);
    };
    if value.is_null() {
        return Ok(None);
    }
    let arr = value
        .as_array()
        .ok_or_else(|| format!("{}: '{field}' must be an array of strings", ctx()))?;
    let mut out = Vec::with_capacity(arr.len());
    for item in arr {
        let s = item
            .as_str()
            .ok_or_else(|| format!("{}: '{field}' must be an array of strings", ctx()))?;
        out.push(s.to_string());
    }
    Ok(Some(out))
}

fn load_connection_spec(
    graph: &mut DirGraph,
    spec: &Json,
    idx: usize,
    endpoint_policy: MissingEndpointPolicy,
    report: &mut RecordsReport,
) -> Result<(), String> {
    let ctx = || format!("from_records: connections[{}]", idx);
    // A non-object spec falls through to the required-field error below.
    if let Some(obj) = spec.as_object() {
        reject_unknown_keys(obj, CONNECTION_SPEC_KEYS, &ctx())?;
    }
    let connection_type = required_str(spec, "type", &ctx)?;
    let source_type = required_str(spec, "source_type", &ctx)?;
    let source_id_field = required_str(spec, "source_id_field", &ctx)?;
    let target_types = required_type_list(spec, "target_type", &ctx)?;
    let target_id_field = required_str(spec, "target_id_field", &ctx)?;
    let type_field = optional_str(spec, "target_type_column");

    let records = records_array(spec, &ctx)?;
    if records.is_empty() {
        return Ok(());
    }

    let groups = group_records_by_target(
        graph,
        records,
        &target_types,
        type_field.as_deref(),
        &target_id_field,
        &ctx,
    )?;

    for (target_type, group) in groups {
        let (columns, rows) =
            records_to_columns_rows(&group, &[&source_id_field, &target_id_field], &ctx)?;
        let df =
            DataFrame::from_cypher_rows(columns, rows).map_err(|e| format!("{}: {}", ctx(), e))?;

        match endpoint_policy {
            MissingEndpointPolicy::Vivify => {
                let rep = maintain::add_connections(
                    graph,
                    df,
                    connection_type.clone(),
                    source_type.clone(),
                    source_id_field.clone(),
                    target_type.to_string(),
                    target_id_field.clone(),
                    None,
                    None,
                    None,
                )
                .map_err(|e| format!("{}: {}", ctx(), e))?;
                report.edges_added += rep.connections_created;
            }
            MissingEndpointPolicy::Drop | MissingEndpointPolicy::Error => {
                let edge_context = EdgeFrameContext {
                    connection_type: &connection_type,
                    source_type: &source_type,
                    source_id_field: &source_id_field,
                    target_type,
                    target_id_field: &target_id_field,
                    connection_idx: idx,
                    endpoint_policy,
                };
                let edge_specs = edge_specs_from_frame(graph, &df, &edge_context)?;
                let rep = maintain::add_edges_from_specs(graph, edge_specs)
                    .map_err(|e| format!("{}: {}", ctx(), e))?;
                report.edges_added += rep.connections_created;
                report.edges_dropped_missing_endpoint += rep.skipped_missing_endpoint;
            }
        }
    }
    report.connection_types.push(connection_type);
    Ok(())
}

/// One connection spec's records grouped by the target type each routes to,
/// in declaration order. The blueprint junction loader's union-target rules,
/// on records instead of CSV rows: `target_type_column` names each record's
/// type, and without it the declared types are probed for the record's target
/// id, an id none has taking the first.
///
/// The one difference is what an unknown type value does. `from_records`
/// key sets are closed and it raises where the blueprint loader warns, so a
/// record naming a type outside the declaration is an error here.
fn group_records_by_target<'a>(
    graph: &DirGraph,
    records: &'a [Json],
    target_types: &'a [String],
    type_field: Option<&str>,
    target_id_field: &str,
    ctx: &impl Fn() -> String,
) -> Result<Vec<(&'a str, Vec<&'a Json>)>, String> {
    if target_types.len() == 1 && type_field.is_none() {
        return Ok(vec![(
            target_types[0].as_str(),
            records.iter().collect::<Vec<_>>(),
        )]);
    }
    let mut groups: Vec<Vec<&Json>> = vec![Vec::new(); target_types.len()];
    for rec in records {
        let owner = match type_field {
            Some(field) => {
                let declared = rec.get(field).and_then(|v| v.as_str()).ok_or_else(|| {
                    format!(
                        "{}: record has no string '{field}' — 'target_type_column' names the \
                         field holding each record's target type",
                        ctx()
                    )
                })?;
                target_types
                    .iter()
                    .position(|t| t == declared)
                    .ok_or_else(|| {
                        format!(
                            "{}: record names target type '{declared}', which is not in \
                             'target_type' ({})",
                            ctx(),
                            target_types.join(", ")
                        )
                    })?
            }
            None => {
                let id = rec.get(target_id_field).map(json_to_value);
                id.as_ref()
                    .and_then(|v| {
                        target_types
                            .iter()
                            .position(|t| graph.id_indices.lookup(t, v).is_some())
                    })
                    .unwrap_or(0)
            }
        };
        groups[owner].push(rec);
    }
    Ok(target_types
        .iter()
        .zip(groups)
        .filter(|(_, group)| !group.is_empty())
        .map(|(t, group)| (t.as_str(), group))
        .collect())
}

/// A required `"Type"` or `["TypeA", "TypeB"]` field.
fn required_type_list(
    spec: &Json,
    key: &str,
    ctx: &impl Fn() -> String,
) -> Result<Vec<String>, String> {
    match spec.get(key) {
        Some(Json::String(s)) => Ok(vec![s.clone()]),
        Some(Json::Array(items)) if !items.is_empty() => items
            .iter()
            .map(|item| {
                item.as_str().map(str::to_string).ok_or_else(|| {
                    format!("{}: '{key}' must be a string or an array of strings", ctx())
                })
            })
            .collect(),
        Some(Json::Array(_)) => Err(format!("{}: '{key}' names no node type", ctx())),
        _ => Err(format!(
            "{}: missing required string field '{}'",
            ctx(),
            key
        )),
    }
}

// ── Helpers ──────────────────────────────────────────────────────────────

/// The only keys a records spec may carry at its top level.
///
/// `on_missing_endpoint` is accepted here as well as taken as a binding
/// argument: the Python shim injects the argument into the spec object before
/// calling in (`kglite-py` `from_records_rust`).
const TOP_LEVEL_KEYS: &[&str] = &["nodes", "connections", "on_missing_endpoint"];

/// The only keys one entry of `nodes` may carry.
const NODE_SPEC_KEYS: &[&str] = &[
    "type",
    "id_field",
    "title_field",
    "labels",
    "conflict_handling",
    "records",
];

/// The only keys one entry of `connections` may carry.
const CONNECTION_SPEC_KEYS: &[&str] = &[
    "type",
    "source_type",
    "source_id_field",
    "target_type",
    "target_type_column",
    "target_id_field",
    "records",
];

/// Refuse the first key of `map` that `accepted` does not contain.
///
/// A key this parser does not read is otherwise dropped in silence and the
/// build reports success on a graph the caller did not describe. The near-miss
/// hint comes from the same
/// [`did_you_mean`](crate::graph::mutation::validation::did_you_mean) the
/// ontology and schema parsers use, and there is no escape-hatch prefix, for
/// the same reason they have none.
fn reject_unknown_keys(
    map: &serde_json::Map<String, Json>,
    accepted: &[&str],
    ctx: &str,
) -> Result<(), String> {
    for key in map.keys() {
        if accepted.contains(&key.as_str()) {
            continue;
        }
        let suggestion = crate::graph::mutation::validation::did_you_mean(key, accepted);
        if !suggestion.is_empty() {
            return Err(format!("{ctx}: unknown key '{key}'.{suggestion}"));
        }
        let list = accepted
            .iter()
            .map(|k| format!("'{k}'"))
            .collect::<Vec<_>>()
            .join(", ");
        return Err(format!(
            "{ctx}: unknown key '{key}'. Accepted keys: {list}."
        ));
    }
    Ok(())
}

fn parse_endpoint_policy(value: Option<&Json>) -> Result<MissingEndpointPolicy, String> {
    match value {
        None => Ok(MissingEndpointPolicy::Vivify),
        Some(Json::String(s)) if s == "vivify" => Ok(MissingEndpointPolicy::Vivify),
        Some(Json::String(s)) if s == "drop" => Ok(MissingEndpointPolicy::Drop),
        Some(Json::String(s)) if s == "error" => Ok(MissingEndpointPolicy::Error),
        Some(Json::String(s)) => Err(format!(
            "from_records: unknown on_missing_endpoint mode '{s}'; expected 'vivify', 'drop', or 'error'"
        )),
        Some(_) => Err(
            "from_records: 'on_missing_endpoint' must be 'vivify', 'drop', or 'error'"
                .to_string(),
        ),
    }
}

struct EdgeFrameContext<'a> {
    connection_type: &'a str,
    source_type: &'a str,
    source_id_field: &'a str,
    target_type: &'a str,
    target_id_field: &'a str,
    connection_idx: usize,
    endpoint_policy: MissingEndpointPolicy,
}

fn edge_specs_from_frame(
    graph: &DirGraph,
    frame: &DataFrame,
    context: &EdgeFrameContext<'_>,
) -> Result<Vec<maintain::EdgeSpec>, String> {
    let property_columns: Vec<String> = frame
        .get_column_names()
        .into_iter()
        .filter(|name| name != context.source_id_field && name != context.target_id_field)
        .collect();
    let mut specs = Vec::with_capacity(frame.row_count());

    for row_idx in 0..frame.row_count() {
        let source_id = frame
            .get_value(row_idx, context.source_id_field)
            .unwrap_or(Value::Null);
        let target_id = frame
            .get_value(row_idx, context.target_id_field)
            .unwrap_or(Value::Null);

        if context.endpoint_policy == MissingEndpointPolicy::Error {
            validate_endpoint(
                graph,
                context.source_type,
                context.source_id_field,
                &source_id,
                context.connection_idx,
                row_idx,
                "source",
            )?;
            validate_endpoint(
                graph,
                context.target_type,
                context.target_id_field,
                &target_id,
                context.connection_idx,
                row_idx,
                "target",
            )?;
        }

        let properties = property_columns
            .iter()
            .filter_map(|name| {
                frame
                    .get_value(row_idx, name)
                    .filter(|value| !matches!(value, Value::Null))
                    .map(|value| (name.clone(), value))
            })
            .collect();
        specs.push(maintain::EdgeSpec {
            source_type: context.source_type.to_string(),
            source_id,
            target_type: context.target_type.to_string(),
            target_id,
            edge_type: context.connection_type.to_string(),
            properties,
        });
    }

    Ok(specs)
}

fn validate_endpoint(
    graph: &DirGraph,
    node_type: &str,
    id_field: &str,
    id: &Value,
    connection_idx: usize,
    row_idx: usize,
    endpoint_name: &str,
) -> Result<(), String> {
    let ctx = format!("from_records: connections[{connection_idx}].records[{row_idx}]");
    if matches!(id, Value::Null) {
        return Err(format!(
            "{ctx}: {endpoint_name} endpoint id field '{id_field}' is null"
        ));
    }
    if graph.lookup_by_id_readonly(node_type, id).is_none() {
        return Err(format!(
            "{ctx}: {endpoint_name} endpoint {node_type}({id}) does not exist"
        ));
    }
    Ok(())
}

fn required_str(spec: &Json, key: &str, ctx: &impl Fn() -> String) -> Result<String, String> {
    spec.get(key)
        .and_then(|v| v.as_str())
        .map(|s| s.to_string())
        .ok_or_else(|| format!("{}: missing required string field '{}'", ctx(), key))
}

fn optional_str(spec: &Json, key: &str) -> Option<String> {
    spec.get(key)
        .and_then(|v| v.as_str())
        .map(|s| s.to_string())
}

fn records_array<'a>(spec: &'a Json, ctx: &impl Fn() -> String) -> Result<&'a Vec<Json>, String> {
    spec.get("records")
        .and_then(|v| v.as_array())
        .ok_or_else(|| format!("{}: 'records' must be an array", ctx()))
}

/// Flatten an array of JSON objects into `(column_names, rows)`. Column order
/// is `required` fields first (in the given order), then every other key in
/// first-seen order. A record missing a column yields `Value::Null` there.
fn records_to_columns_rows(
    records: &[&Json],
    required: &[&str],
    ctx: &impl Fn() -> String,
) -> Result<(Vec<String>, Vec<Vec<Value>>), String> {
    let mut columns: Vec<String> = required.iter().map(|s| s.to_string()).collect();
    let mut seen: std::collections::HashSet<String> = columns.iter().cloned().collect();
    for rec in records {
        let obj = rec
            .as_object()
            .ok_or_else(|| format!("{}: every record must be a JSON object", ctx()))?;
        for key in obj.keys() {
            if seen.insert(key.clone()) {
                columns.push(key.clone());
            }
        }
    }

    let rows = records
        .iter()
        .map(|rec| {
            let obj = rec.as_object().expect("validated above");
            columns
                .iter()
                .map(|col| obj.get(col).map(json_to_value).unwrap_or(Value::Null))
                .collect()
        })
        .collect();

    Ok((columns, rows))
}

/// Recursive JSON → [`Value`]. Arrays become native `Value::List`, objects
/// become `Value::Map`; scalars map to their natural typed variant.
fn json_to_value(j: &Json) -> Value {
    match j {
        Json::Null => Value::Null,
        Json::Bool(b) => Value::Boolean(*b),
        Json::Number(n) => n
            .as_i64()
            .map(Value::Int64)
            .or_else(|| n.as_f64().map(Value::Float64))
            .unwrap_or(Value::Null),
        Json::String(s) => Value::String(s.clone()),
        Json::Array(items) => Value::List(items.iter().map(json_to_value).collect()),
        Json::Object(map) => Value::Map(
            map.iter()
                .map(|(k, v)| (k.clone(), json_to_value(v)))
                .collect(),
        ),
    }
}

#[cfg(test)]
#[path = "json_records_tests.rs"]
mod tests;