use crate::db::GraphDb;
use core_rules::{default_max_edges, Predicate, RuleDef};
use core_storage::fs::Fs;
use core_storage::{GraphError, Result, Value};
use serde::Serialize;
use std::collections::{BTreeMap, BTreeSet};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IngestOptions {
pub key_field: String,
pub auto_fk: AutoFk,
}
impl Default for IngestOptions {
fn default() -> Self {
Self {
key_field: "id".into(),
auto_fk: AutoFk::default(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AutoFk {
Auto { suffix: String },
Off,
}
impl Default for AutoFk {
fn default() -> Self {
AutoFk::Auto {
suffix: "_id".into(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct FkSkip {
pub field: String,
pub reason: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct IngestReport {
pub inserted: usize,
pub row_errors: Vec<(usize, String)>,
pub rules_created: Vec<String>,
pub skipped_fk_fields: Vec<FkSkip>,
pub edges_inserted: usize,
}
pub fn json_to_value(v: serde_json::Value) -> Option<Value> {
match v {
serde_json::Value::Null => None,
serde_json::Value::Bool(b) => Some(Value::Bool(b)),
serde_json::Value::Number(n) => number_to_value(&n),
serde_json::Value::String(s) => Some(Value::Str(s)),
serde_json::Value::Array(items) => {
let mut out = Vec::with_capacity(items.len());
for item in items {
out.push(json_to_value(item)?);
}
Some(Value::List(out))
}
serde_json::Value::Object(obj) => {
let mut map = std::collections::BTreeMap::new();
for (k, v) in obj {
if let Some(val) = json_to_value(v) {
map.insert(k, val);
}
}
Some(Value::Map(map))
}
}
}
fn number_to_value(n: &serde_json::Number) -> Option<Value> {
if let Some(i) = n.as_i64() {
return Some(Value::Int(i));
}
let f = n.as_f64()?;
if f.is_finite() && f.fract() == 0.0 && f >= i64::MIN as f64 && f <= i64::MAX as f64 {
Some(Value::Int(f as i64))
} else {
Some(Value::Float(f))
}
}
fn object_to_row(
obj: &serde_json::Map<String, serde_json::Value>,
) -> std::result::Result<BTreeMap<String, Value>, String> {
let mut row = BTreeMap::new();
for (k, v) in obj {
if let Some(val) = json_to_value(v.clone()) {
row.insert(k.clone(), val);
}
}
Ok(row)
}
pub struct JsonRows {
pub rows: Vec<BTreeMap<String, Value>>,
kept_indices: Vec<usize>,
shape_errors: Vec<(usize, String)>,
}
impl JsonRows {
pub fn into_report(self, mut report: IngestReport) -> IngestReport {
for (idx, _) in &mut report.row_errors {
*idx = self.kept_indices[*idx];
}
report.row_errors.extend(self.shape_errors);
report.row_errors.sort_by_key(|(i, _)| *i);
report
}
}
pub fn json_to_rows(value: &serde_json::Value) -> Result<JsonRows> {
let arr = value.as_array().ok_or_else(|| GraphError::IngestError {
detail: "top-level JSON must be an array of objects".into(),
})?;
if !arr.iter().all(|v| v.is_object()) {
return Err(GraphError::IngestError {
detail: "top-level JSON must be an array of objects".into(),
});
}
let mut rows = Vec::new();
let mut shape_errors = Vec::new();
let mut kept_indices = Vec::new();
for (i, item) in arr.iter().enumerate() {
let obj = item
.as_object()
.expect("top-level checked as array of objects");
match object_to_row(obj) {
Ok(row) => {
kept_indices.push(i);
rows.push(row);
}
Err(msg) => shape_errors.push((i, msg)),
}
}
Ok(JsonRows {
rows,
kept_indices,
shape_errors,
})
}
pub(crate) fn run_json<F: Fs>(
db: &mut GraphDb<F>,
label: &str,
json: &str,
opts: &IngestOptions,
) -> Result<IngestReport> {
let parsed: serde_json::Value =
serde_json::from_str(json).map_err(|e| GraphError::IngestError {
detail: e.to_string(),
})?;
let mut converted = json_to_rows(&parsed)?;
let rows = std::mem::take(&mut converted.rows);
let report = run(db, label, rows, opts, &[])?;
Ok(converted.into_report(report))
}
type PropMap = BTreeMap<String, Value>;
struct Classified {
accepted: Vec<(String, PropMap)>,
row_errors: Vec<(usize, String)>,
}
pub(crate) fn run<F: Fs>(
db: &mut GraphDb<F>,
label: &str,
rows: Vec<BTreeMap<String, Value>>,
opts: &IngestOptions,
edges: &[(String, String, String)],
) -> Result<IngestReport> {
let Classified {
accepted,
row_errors,
} = classify_rows(db, rows, &opts.key_field);
let (new_rules, skipped_fk_fields) = match &opts.auto_fk {
AutoFk::Off => (Vec::new(), Vec::new()),
AutoFk::Auto { suffix } => infer_auto_fk(db, label, suffix, &opts.key_field, &accepted),
};
let rules_created: Vec<String> = new_rules.iter().map(|r| r.name.clone()).collect();
let mut batch = db.batch();
for def in new_rules {
batch.create_rule(def);
}
for (key, props) in &accepted {
let prop_vec: Vec<(String, Value)> =
props.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
batch.insert_node(label, key, prop_vec);
}
for (etype, src, dst) in edges {
batch.insert_edge(etype, src, dst);
}
let (_, edges_inserted) = batch.commit_ingest(label, accepted.len())?;
Ok(IngestReport {
inserted: accepted.len(),
row_errors,
rules_created,
skipped_fk_fields,
edges_inserted,
})
}
fn classify_rows<F: Fs>(db: &GraphDb<F>, rows: Vec<PropMap>, key_field: &str) -> Classified {
let mut accepted = Vec::new();
let mut row_errors = Vec::new();
let mut seen: BTreeSet<String> = BTreeSet::new();
for (i, row) in rows.into_iter().enumerate() {
match row.get(key_field) {
None => row_errors.push((i, format!("missing key field {key_field}"))),
Some(Value::Str(key)) => {
if db.has_node(key) || seen.contains(key) {
row_errors.push((i, format!("duplicate key {key}")));
} else {
seen.insert(key.clone());
accepted.push((key.clone(), row));
}
}
Some(_) => row_errors.push((i, format!("key field {key_field} is not a string"))),
}
}
Classified {
accepted,
row_errors,
}
}
fn infer_auto_fk<F: Fs>(
db: &GraphDb<F>,
src_label: &str,
suffix: &str,
key_field: &str,
accepted: &[(String, PropMap)],
) -> (Vec<RuleDef>, Vec<FkSkip>) {
let existing_rule_names: BTreeSet<String> = db.rules().into_iter().map(|r| r.name).collect();
let accepted_keys: BTreeSet<&str> = accepted.iter().map(|(k, _)| k.as_str()).collect();
let mut fields: BTreeSet<String> = BTreeSet::new();
for (_, row) in accepted {
for field in row.keys() {
if field != key_field && field.ends_with(suffix) && field.len() > suffix.len() {
fields.insert(field.clone());
}
}
}
let mut new_rules = Vec::new();
let mut skipped = Vec::new();
for field in fields {
let mut values: BTreeSet<&str> = BTreeSet::new();
for (_, row) in accepted {
if let Some(Value::Str(s)) = row.get(&field) {
values.insert(s.as_str());
}
}
let mut labels: BTreeSet<String> = BTreeSet::new();
for value in values {
if let Some(n) = db.node_ref(value) {
labels.insert(n.label().to_string());
}
if accepted_keys.contains(value) {
labels.insert(src_label.to_string());
}
}
match labels.len() {
0 => skipped.push(FkSkip {
field,
reason: "no matching target keys".into(),
}),
1 => {
let dst_label = labels.into_iter().next().expect("len == 1");
let name = format!("auto_fk_{}_{field}", src_label.to_lowercase());
if existing_rule_names.contains(&name) {
continue;
}
let remainder = &field[..field.len() - suffix.len()];
let predicate = Predicate::KeyMatch {
field: field.clone(),
};
let max_edges = Some(default_max_edges(&predicate));
new_rules.push(RuleDef {
name,
src_label: src_label.to_string(),
dst_label,
predicate,
edge_type: remainder.to_uppercase(),
weight_prop: None,
max_edges,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
namespace: None,
});
}
_ => {
let listed = labels.into_iter().collect::<Vec<_>>().join(", ");
skipped.push(FkSkip {
field,
reason: format!("ambiguous target labels: {listed}"),
});
}
}
}
(new_rules, skipped)
}