1use mushroomdb::{ExportEdge, NodeInfo, RuleDef, Value};
31use std::collections::{BTreeMap, BTreeSet};
32use std::fmt::Write as _;
33use std::io::Write as _;
34use std::path::{Path, PathBuf};
35use std::sync::Arc;
36
37use arrow_array::{ArrayRef, BooleanArray, RecordBatch, StringArray};
38use arrow_schema::{DataType, Field, Schema};
39use parquet::arrow::ArrowWriter;
40use parquet::file::properties::WriterProperties;
41
42#[derive(Debug, Clone, PartialEq, Eq)]
44pub enum ExportFormat {
45 Jsonl,
46 JsonlSingle,
48 Csv,
49 Cypher,
50 Parquet,
51 Graphml,
52}
53
54impl ExportFormat {
55 pub fn parse(s: &str) -> Option<Self> {
56 match s {
57 "jsonl" => Some(Self::Jsonl),
58 "csv" => Some(Self::Csv),
59 "cypher" => Some(Self::Cypher),
60 "parquet" => Some(Self::Parquet),
61 "graphml" => Some(Self::Graphml),
62 _ => None,
63 }
64 }
65
66 pub fn name(&self) -> &'static str {
67 match self {
68 Self::Jsonl => "jsonl",
69 Self::JsonlSingle => "jsonl",
70 Self::Csv => "csv",
71 Self::Cypher => "cypher",
72 Self::Parquet => "parquet",
73 Self::Graphml => "graphml",
74 }
75 }
76}
77
78fn value_to_json(v: &Value) -> serde_json::Value {
83 match v {
84 Value::Int(i) => serde_json::Value::Number((*i).into()),
85 Value::Float(f) => serde_json::Number::from_f64(*f)
87 .map(serde_json::Value::Number)
88 .unwrap_or(serde_json::Value::Null),
89 Value::Str(s) => serde_json::Value::String(s.clone()),
90 Value::Bool(b) => serde_json::Value::Bool(*b),
91 Value::List(xs) => serde_json::Value::Array(xs.iter().map(value_to_json).collect()),
92 Value::Map(m) => serde_json::Value::Object(
93 m.iter()
94 .map(|(k, v)| (k.clone(), value_to_json(v)))
95 .collect(),
96 ),
97 Value::Date(_)
98 | Value::LocalDateTime(_)
99 | Value::DateTime { .. }
100 | Value::Duration { .. }
101 | Value::Binary(_) => mushroomdb::tagged_json(v)
104 .ok()
105 .flatten()
106 .unwrap_or(serde_json::Value::Null),
107 Value::Null => serde_json::Value::Null,
109 }
110}
111
112pub trait ExportSource {
117 fn for_each_node(
118 &self,
119 each: &mut dyn FnMut(NodeInfo) -> Result<(), crate::CliError>,
120 ) -> Result<(), crate::CliError>;
121 fn for_each_edge(
122 &self,
123 each: &mut dyn FnMut(ExportEdge) -> Result<(), crate::CliError>,
124 ) -> Result<(), crate::CliError>;
125 fn label_of(&self, key: &str) -> Option<String>;
127}
128
129impl ExportSource for mushroomdb::GraphDb<mushroomdb::RealFs> {
130 fn for_each_node(
131 &self,
132 each: &mut dyn FnMut(NodeInfo) -> Result<(), crate::CliError>,
133 ) -> Result<(), crate::CliError> {
134 self.for_each_node_for_export(each)
135 }
136 fn for_each_edge(
137 &self,
138 each: &mut dyn FnMut(ExportEdge) -> Result<(), crate::CliError>,
139 ) -> Result<(), crate::CliError> {
140 self.for_each_edge_for_export(each)
141 }
142 fn label_of(&self, key: &str) -> Option<String> {
143 self.node_ref(key).map(|n| n.label().to_string())
144 }
145}
146
147pub struct Slices<'a> {
149 pub nodes: &'a [NodeInfo],
150 pub edges: &'a [ExportEdge],
151}
152
153impl ExportSource for Slices<'_> {
154 fn for_each_node(
155 &self,
156 each: &mut dyn FnMut(NodeInfo) -> Result<(), crate::CliError>,
157 ) -> Result<(), crate::CliError> {
158 self.nodes.iter().try_for_each(|n| each(n.clone()))
159 }
160 fn for_each_edge(
161 &self,
162 each: &mut dyn FnMut(ExportEdge) -> Result<(), crate::CliError>,
163 ) -> Result<(), crate::CliError> {
164 self.edges.iter().try_for_each(|e| each(e.clone()))
165 }
166 fn label_of(&self, key: &str) -> Option<String> {
167 self.nodes
168 .iter()
169 .find(|n| n.key == key)
170 .map(|n| n.label.clone())
171 }
172}
173
174#[derive(Debug, Clone, Default, PartialEq, Eq)]
176pub struct ExportStats {
177 pub nodes: usize,
178 pub edges: usize,
180 pub derived_edges: usize,
181 pub non_finite_floats: usize,
184}
185
186fn count_non_finite(v: &Value) -> usize {
187 match v {
188 Value::Float(f) => usize::from(!f.is_finite()),
189 Value::List(xs) => xs.iter().map(count_non_finite).sum(),
190 Value::Map(m) => m.values().map(count_non_finite).sum(),
191 _ => 0,
192 }
193}
194
195fn json_line(w: &mut impl std::io::Write, v: &serde_json::Value) -> Result<(), crate::CliError> {
196 serde_json::to_writer(&mut *w, v)
197 .map_err(|e| crate::CliError(format!("json encode error: {e}")))?;
198 w.write_all(b"\n")?;
199 Ok(())
200}
201
202fn node_json(node: &NodeInfo, typed: bool) -> serde_json::Value {
208 let mut obj = serde_json::Map::new();
209 obj.insert("key".into(), serde_json::Value::String(node.key.clone()));
210 obj.insert(
211 "label".into(),
212 serde_json::Value::String(node.label.clone()),
213 );
214 let mut reserved = serde_json::Map::new();
215 for (k, v) in &node.props {
216 let own = matches!(k.as_str(), "key" | "label") || k == crate::import::JSONL_PROPS_FIELD;
217 if own || (typed && k == "type") {
218 reserved.insert(k.clone(), value_to_json(v));
219 } else {
220 obj.insert(k.clone(), value_to_json(v));
221 }
222 }
223 if !reserved.is_empty() {
224 obj.insert(
225 crate::import::JSONL_PROPS_FIELD.into(),
226 serde_json::Value::Object(reserved),
227 );
228 }
229 serde_json::Value::Object(obj)
230}
231
232fn edge_json(edge: &ExportEdge) -> serde_json::Value {
233 let mut obj = serde_json::Map::new();
234 obj.insert("edge_type".into(), edge.edge_type.clone().into());
235 obj.insert("src".into(), edge.src.clone().into());
236 obj.insert("dst".into(), edge.dst.clone().into());
237 obj.insert("derived".into(), edge.derived.into());
238 obj.insert(
239 "rule".into(),
240 edge.rule
241 .clone()
242 .map_or(serde_json::Value::Null, Into::into),
243 );
244 serde_json::Value::Object(obj)
245}
246
247fn typed_json_line(
250 w: &mut impl std::io::Write,
251 kind: &str,
252 value: &serde_json::Value,
253) -> Result<(), crate::CliError> {
254 let text = serde_json::to_string(value)
255 .map_err(|e| crate::CliError(format!("json encode error: {e}")))?;
256 let rest = text.strip_prefix('{').unwrap_or(&text);
257 let comma = if rest == "}" { "" } else { "," };
258 writeln!(w, "{{\"type\":\"{kind}\"{comma}{rest}")?;
259 Ok(())
260}
261
262fn create_file(path: PathBuf) -> Result<std::io::BufWriter<std::fs::File>, crate::CliError> {
263 Ok(std::io::BufWriter::with_capacity(
264 1 << 20,
265 std::fs::File::create(path)?,
266 ))
267}
268
269pub fn write_jsonl(
275 source: &dyn ExportSource,
276 rules: &[RuleDef],
277 dest: &Path,
278) -> Result<ExportStats, crate::CliError> {
279 std::fs::create_dir_all(dest)?;
280 let mut stats = ExportStats::default();
281 let mut f = create_file(dest.join("nodes.jsonl"))?;
282 source.for_each_node(&mut |node| {
283 stats.non_finite_floats += node.props.values().map(count_non_finite).sum::<usize>();
284 stats.nodes += 1;
285 json_line(&mut f, &node_json(&node, false))
286 })?;
287 f.flush()?;
288 let mut f = create_file(dest.join("edges.jsonl"))?;
289 source.for_each_edge(&mut |edge| {
290 stats.edges += 1;
291 stats.derived_edges += usize::from(edge.derived);
292 json_line(&mut f, &edge_json(&edge))
293 })?;
294 f.flush()?;
295 write_rules_jsonl(rules, dest)?;
296 Ok(stats)
297}
298
299pub fn write_jsonl_single(
312 source: &dyn ExportSource,
313 rules: &[RuleDef],
314 dest: &Path,
315) -> Result<ExportStats, crate::CliError> {
316 std::fs::create_dir_all(dest)?;
317 let mut stats = ExportStats::default();
318 let mut f = create_file(dest.join("graph.jsonl"))?;
319 source.for_each_node(&mut |node| {
320 stats.non_finite_floats += node.props.values().map(count_non_finite).sum::<usize>();
321 stats.nodes += 1;
322 typed_json_line(&mut f, "node", &node_json(&node, true))
323 })?;
324 source.for_each_edge(&mut |edge| {
325 stats.edges += 1;
326 stats.derived_edges += usize::from(edge.derived);
327 typed_json_line(&mut f, "edge", &edge_json(&edge))
328 })?;
329 for rule in rules {
330 let line = serde_json::json!({ "rule": rule });
331 typed_json_line(&mut f, "rule", &line)?;
332 }
333 f.flush()?;
334 Ok(stats)
335}
336
337fn write_rules_jsonl(rules: &[RuleDef], dest: &Path) -> Result<(), crate::CliError> {
340 let mut f = std::fs::File::create(dest.join("rules.jsonl"))?;
341 for rule in rules {
342 let line = serde_json::to_string(rule)
343 .map_err(|e| crate::CliError(format!("json encode rule: {e}")))?;
344 writeln!(f, "{line}")?;
345 }
346 Ok(())
347}
348
349#[derive(Debug, Clone, Copy, PartialEq, Eq)]
356enum CsvKind {
357 Str,
358 Int,
359 Float,
360 Bool,
361 Date,
362 LocalDateTime,
363 DateTime,
364 Duration,
365 Binary,
366 Json,
367}
368
369impl CsvKind {
370 fn suffix(self) -> &'static str {
371 match self {
372 Self::Str => "string",
373 Self::Int => "int",
374 Self::Float => "float",
375 Self::Bool => "boolean",
376 Self::Date => "date",
377 Self::LocalDateTime => "localdatetime",
378 Self::DateTime => "datetime",
379 Self::Duration => "duration",
380 Self::Binary => "binary",
381 Self::Json => "json",
382 }
383 }
384
385 fn of(v: &Value) -> Self {
387 match v {
388 Value::Str(s) if s.is_empty() => Self::Json,
390 Value::Str(_) => Self::Str,
391 Value::Int(_) => Self::Int,
392 Value::Float(f) if f.is_finite() => Self::Float,
394 Value::Bool(_) => Self::Bool,
395 Value::Date(_) => Self::Date,
396 Value::LocalDateTime(_) => Self::LocalDateTime,
397 Value::DateTime { .. } => Self::DateTime,
398 Value::Duration { .. } => Self::Duration,
399 Value::Binary(_) => Self::Binary,
400 Value::Float(_) | Value::List(_) | Value::Map(_) | Value::Null => Self::Json,
401 }
402 }
403}
404
405fn tagged_text_of(v: &Value) -> String {
407 match mushroomdb::tagged_json(v) {
408 Ok(Some(serde_json::Value::Object(obj))) => obj
409 .into_iter()
410 .next()
411 .and_then(|(_, text)| text.as_str().map(str::to_string))
412 .unwrap_or_default(),
413 _ => String::new(),
414 }
415}
416
417fn csv_cell(v: &Value, kind: CsvKind) -> String {
418 match (kind, v) {
419 (CsvKind::Str, Value::Str(s)) => s.clone(),
420 (CsvKind::Int, Value::Int(i)) => i.to_string(),
421 (CsvKind::Float, Value::Float(f)) => format!("{f:?}"),
424 (CsvKind::Bool, Value::Bool(b)) => b.to_string(),
425 (CsvKind::Json, v) => serde_json::to_string(&value_to_json(v)).unwrap_or_default(),
426 (_, v) => tagged_text_of(v),
427 }
428}
429
430fn csv_error(e: csv::Error) -> crate::CliError {
431 crate::CliError(format!("csv write error: {e}"))
432}
433
434pub fn write_csv(
449 source: &dyn ExportSource,
450 rules: &[RuleDef],
451 dest: &Path,
452) -> Result<ExportStats, crate::CliError> {
453 std::fs::create_dir_all(dest)?;
454 let mut stats = ExportStats::default();
455
456 let mut id_is_key = true;
458 let mut kinds: BTreeMap<String, CsvKind> = BTreeMap::new();
459 source.for_each_node(&mut |node| {
460 if node.label.contains(';') {
461 return Err(crate::CliError(format!(
462 "label `{}` of node `{}` cannot be written to CSV: the importer reads `;` as a multi-label",
463 node.label, node.key
464 )));
465 }
466 stats.nodes += 1;
467 id_is_key &= matches!(node.props.get("id"), Some(Value::Str(s)) if *s == node.key);
468 for (name, value) in &node.props {
469 if name.is_empty() {
470 return Err(crate::CliError(format!(
471 "node `{}` has a property with an empty name, which a CSV header cannot carry",
472 node.key
473 )));
474 }
475 stats.non_finite_floats += count_non_finite(value);
476 let kind = CsvKind::of(value);
477 match kinds.get_mut(name.as_str()) {
478 Some(k) if *k != kind => *k = CsvKind::Json,
479 Some(_) => {}
480 None => {
481 kinds.insert(name.clone(), kind);
482 }
483 }
484 }
485 Ok(())
486 })?;
487 let id_is_key = id_is_key && stats.nodes > 0;
488 if id_is_key {
489 kinds.remove("id");
490 }
491
492 let mut w = csv::Writer::from_path(dest.join("nodes.csv")).map_err(csv_error)?;
494 let mut header = vec![
495 if id_is_key { ":ID" } else { ":KEY" }.to_string(),
496 ":LABEL".to_string(),
497 ];
498 header.extend(kinds.iter().map(|(n, k)| format!("{n}:{}", k.suffix())));
501 w.write_record(&header).map_err(csv_error)?;
502 source.for_each_node(&mut |node| {
503 let mut row = vec![node.key.clone(), node.label.clone()];
504 for (name, kind) in &kinds {
505 row.push(
506 node.props
507 .get(name.as_str())
508 .map(|v| csv_cell(v, *kind))
509 .unwrap_or_default(),
510 );
511 }
512 w.write_record(&row).map_err(csv_error)
513 })?;
514 w.flush()?;
515
516 let mut w = csv::Writer::from_path(dest.join("relationships.csv")).map_err(csv_error)?;
518 w.write_record([":START_ID", ":END_ID", ":TYPE"])
519 .map_err(csv_error)?;
520 source.for_each_edge(&mut |edge| {
521 stats.edges += 1;
522 stats.derived_edges += usize::from(edge.derived);
523 if edge.derived {
524 return Ok(());
525 }
526 w.write_record([&edge.src, &edge.dst, &edge.edge_type])
527 .map_err(csv_error)
528 })?;
529 w.flush()?;
530
531 if !rules.is_empty() {
532 write_rules_jsonl(rules, dest)?;
533 }
534 Ok(stats)
535}
536
537#[derive(Debug, Clone, Default, PartialEq, Eq)]
542pub struct CypherLoss {
543 pub typed_properties: usize,
546 pub unwritable_numbers: usize,
550 pub differing_ids: usize,
553 pub ids_added: usize,
556}
557
558impl CypherLoss {
559 pub fn is_empty(&self) -> bool {
560 *self == Self::default()
561 }
562}
563
564fn cypher_label(label: &str) -> String {
567 if label.is_empty() {
568 String::new()
569 } else {
570 format!(":{}", cypher_ident(label))
571 }
572}
573
574fn cypher_ident(name: &str) -> String {
576 format!("`{}`", name.replace('`', "``"))
577}
578
579fn cypher_string(s: &str) -> String {
581 let mut out = String::with_capacity(s.len() + 2);
582 out.push('\'');
583 for c in s.chars() {
584 match c {
585 '\\' => out.push_str("\\\\"),
586 '\'' => out.push_str("\\'"),
587 '\n' => out.push_str("\\n"),
588 '\r' => out.push_str("\\r"),
589 '\t' => out.push_str("\\t"),
590 c if (c as u32) < 0x20 || c as u32 == 0x7f => {
591 let _ = write!(out, "\\u{:04x}", c as u32);
592 }
593 c => out.push(c),
594 }
595 }
596 out.push('\'');
597 out
598}
599
600fn cypher_literal(v: &Value) -> Result<String, CypherSkip> {
602 Ok(match v {
603 Value::Str(s) => cypher_string(s),
604 Value::Int(i64::MIN) => return Err(CypherSkip::Number),
606 Value::Int(i) => i.to_string(),
607 Value::Float(f) if !f.is_finite() => return Err(CypherSkip::Number),
608 Value::Float(f) => {
609 let text = format!("{f}");
611 if text.contains('.') {
612 text
613 } else {
614 format!("{text}.0")
615 }
616 }
617 Value::Bool(b) => b.to_string(),
618 Value::List(xs) => {
619 let items: Result<Vec<String>, _> = xs.iter().map(cypher_literal).collect();
620 format!("[{}]", items?.join(", "))
621 }
622 Value::Map(m) => {
623 let mut items = Vec::with_capacity(m.len());
624 for (k, v) in m {
625 if k.contains(['\n', '\r']) {
627 return Err(CypherSkip::Number);
628 }
629 items.push(format!("{}: {}", cypher_ident(k), cypher_literal(v)?));
630 }
631 format!("{{{}}}", items.join(", "))
632 }
633 Value::Date(_)
634 | Value::LocalDateTime(_)
635 | Value::DateTime { .. }
636 | Value::Duration { .. }
637 | Value::Binary(_) => return Err(CypherSkip::Typed),
638 Value::Null => return Err(CypherSkip::Number),
639 })
640}
641
642enum CypherSkip {
643 Typed,
644 Number,
645}
646
647pub fn write_cypher(
663 source: &dyn ExportSource,
664 dest: &Path,
665) -> Result<(ExportStats, CypherLoss), crate::CliError> {
666 std::fs::create_dir_all(dest)?;
667 let mut out = create_file(dest.join("graph.cypher"))?;
668 let mut loss = CypherLoss::default();
669 let mut stats = ExportStats::default();
670 source.for_each_node(&mut |node| {
671 stats.nodes += 1;
672 let mut sets = Vec::new();
673 let mut has_id = false;
674 for (name, value) in &node.props {
675 if name == "id" {
676 has_id = true;
678 if !matches!(value, Value::Str(s) if *s == node.key) {
679 loss.differing_ids += 1;
680 }
681 continue;
682 }
683 if name.contains(['\n', '\r']) {
684 loss.unwritable_numbers += 1;
685 continue;
686 }
687 match cypher_literal(value) {
688 Ok(lit) => sets.push(format!("n.{} = {lit}", cypher_ident(name))),
689 Err(CypherSkip::Typed) => loss.typed_properties += 1,
690 Err(CypherSkip::Number) => loss.unwritable_numbers += 1,
691 }
692 }
693 loss.ids_added += usize::from(!has_id);
694 write!(
695 out,
696 "MERGE (n{} {{id: {}}})",
697 cypher_label(&node.label),
698 cypher_string(&node.key)
699 )?;
700 if !sets.is_empty() {
701 write!(out, " ON CREATE SET {}", sets.join(", "))?;
702 }
703 writeln!(out)?;
704 Ok(())
705 })?;
706 source.for_each_edge(&mut |edge| {
707 stats.edges += 1;
708 stats.derived_edges += usize::from(edge.derived);
709 if edge.derived {
710 return Ok(());
711 }
712 let (Some(src_label), Some(dst_label)) =
713 (source.label_of(&edge.src), source.label_of(&edge.dst))
714 else {
715 return Ok(());
716 };
717 writeln!(
718 out,
719 "MERGE (a{} {{id: {}}})-[:{}]->(b{} {{id: {}}})",
720 cypher_label(&src_label),
721 cypher_string(&edge.src),
722 cypher_ident(&edge.edge_type),
723 cypher_label(&dst_label),
724 cypher_string(&edge.dst)
725 )?;
726 Ok(())
727 })?;
728 out.flush()?;
729 Ok((stats, loss))
730}
731
732pub fn write_parquet(
746 nodes: &[NodeInfo],
747 edges: &[ExportEdge],
748 rules: &[RuleDef],
749 dest: &Path,
750) -> Result<(), crate::CliError> {
751 std::fs::create_dir_all(dest)?;
752 let props = WriterProperties::builder().build();
753
754 {
756 let schema = Arc::new(Schema::new(vec![
757 Field::new("key", DataType::Utf8, false),
758 Field::new("label", DataType::Utf8, false),
759 Field::new("props", DataType::Utf8, false),
760 ]));
761 let mut keys = Vec::new();
762 let mut labels = Vec::new();
763 let mut props_json = Vec::new();
764 for node in nodes {
765 keys.push(node.key.clone());
766 labels.push(node.label.clone());
767 let p: serde_json::Map<String, serde_json::Value> = node
768 .props
769 .iter()
770 .map(|(k, v)| (k.clone(), value_to_json(v)))
771 .collect();
772 props_json.push(
773 serde_json::to_string(&serde_json::Value::Object(p))
774 .map_err(|e| crate::CliError(format!("json encode: {e}")))?,
775 );
776 }
777 let batch = RecordBatch::try_new(
778 schema.clone(),
779 vec![
780 Arc::new(StringArray::from(keys)) as ArrayRef,
781 Arc::new(StringArray::from(labels)) as ArrayRef,
782 Arc::new(StringArray::from(props_json)) as ArrayRef,
783 ],
784 )
785 .map_err(|e| crate::CliError(format!("arrow error: {e}")))?;
786
787 let file = std::fs::File::create(dest.join("nodes.parquet"))?;
788 let mut writer = ArrowWriter::try_new(file, schema, Some(props.clone()))
789 .map_err(|e| crate::CliError(format!("parquet writer: {e}")))?;
790 writer
791 .write(&batch)
792 .map_err(|e| crate::CliError(format!("parquet write: {e}")))?;
793 writer
794 .close()
795 .map_err(|e| crate::CliError(format!("parquet close: {e}")))?;
796 }
797
798 {
800 let schema = Arc::new(Schema::new(vec![
801 Field::new("edge_type", DataType::Utf8, false),
802 Field::new("src", DataType::Utf8, false),
803 Field::new("dst", DataType::Utf8, false),
804 Field::new("derived", DataType::Boolean, false),
805 Field::new("rule", DataType::Utf8, true),
806 ]));
807 let mut etypes = Vec::new();
808 let mut srcs = Vec::new();
809 let mut dsts = Vec::new();
810 let mut deriveds = Vec::new();
811 let mut rule_names: Vec<Option<String>> = Vec::new();
812 for edge in edges {
813 etypes.push(edge.edge_type.clone());
814 srcs.push(edge.src.clone());
815 dsts.push(edge.dst.clone());
816 deriveds.push(edge.derived);
817 rule_names.push(edge.rule.clone());
818 }
819 let batch = RecordBatch::try_new(
820 schema.clone(),
821 vec![
822 Arc::new(StringArray::from(etypes)) as ArrayRef,
823 Arc::new(StringArray::from(srcs)) as ArrayRef,
824 Arc::new(StringArray::from(dsts)) as ArrayRef,
825 Arc::new(BooleanArray::from(deriveds)) as ArrayRef,
826 Arc::new(StringArray::from(rule_names)) as ArrayRef,
827 ],
828 )
829 .map_err(|e| crate::CliError(format!("arrow error: {e}")))?;
830
831 let file = std::fs::File::create(dest.join("edges.parquet"))?;
832 let mut writer = ArrowWriter::try_new(file, schema, Some(props.clone()))
833 .map_err(|e| crate::CliError(format!("parquet writer: {e}")))?;
834 writer
835 .write(&batch)
836 .map_err(|e| crate::CliError(format!("parquet write: {e}")))?;
837 writer
838 .close()
839 .map_err(|e| crate::CliError(format!("parquet close: {e}")))?;
840 }
841
842 {
844 let schema = Arc::new(Schema::new(vec![
845 Field::new("name", DataType::Utf8, false),
846 Field::new("definition", DataType::Utf8, false),
847 ]));
848 let mut names = Vec::new();
849 let mut definitions = Vec::new();
850 for rule in rules {
851 names.push(rule.name.clone());
852 definitions.push(
853 serde_json::to_string(rule)
854 .map_err(|e| crate::CliError(format!("json encode: {e}")))?,
855 );
856 }
857 let batch = RecordBatch::try_new(
858 schema.clone(),
859 vec![
860 Arc::new(StringArray::from(names)) as ArrayRef,
861 Arc::new(StringArray::from(definitions)) as ArrayRef,
862 ],
863 )
864 .map_err(|e| crate::CliError(format!("arrow error: {e}")))?;
865
866 let file = std::fs::File::create(dest.join("rules.parquet"))?;
867 let mut writer = ArrowWriter::try_new(file, schema, Some(props))
868 .map_err(|e| crate::CliError(format!("parquet writer: {e}")))?;
869 writer
870 .write(&batch)
871 .map_err(|e| crate::CliError(format!("parquet write: {e}")))?;
872 writer
873 .close()
874 .map_err(|e| crate::CliError(format!("parquet close: {e}")))?;
875 }
876
877 Ok(())
878}
879
880#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
884enum GmlType {
885 Boolean,
886 Long,
891 Double,
892 String,
893}
894
895impl GmlType {
896 fn as_str(self) -> &'static str {
897 match self {
898 Self::Boolean => "boolean",
899 Self::Long => "long",
900 Self::Double => "double",
901 Self::String => "string",
902 }
903 }
904}
905
906fn value_gml_type(v: &Value) -> GmlType {
909 match v {
910 Value::Int(_) => GmlType::Long,
911 Value::Float(_) => GmlType::Double,
912 Value::Bool(_) => GmlType::Boolean,
913 Value::Str(_) | Value::List(_) | Value::Map(_) => GmlType::String,
914 Value::Date(_)
916 | Value::LocalDateTime(_)
917 | Value::DateTime { .. }
918 | Value::Duration { .. }
919 | Value::Binary(_) => GmlType::String,
920 Value::Null => GmlType::String,
921 }
922}
923
924fn value_gml_text(v: &Value) -> String {
932 match value_to_json(v) {
933 serde_json::Value::Null => String::new(),
934 serde_json::Value::Bool(b) => b.to_string(),
935 serde_json::Value::Number(n) => n.to_string(),
936 serde_json::Value::String(s) => s,
937 other => other.to_string(),
938 }
939}
940
941fn xml_escape(s: &str) -> String {
945 let mut out = String::with_capacity(s.len());
946 for c in s.chars() {
947 match c {
948 '&' => out.push_str("&"),
949 '<' => out.push_str("<"),
950 '>' => out.push_str(">"),
951 '"' => out.push_str("""),
952 '\'' => out.push_str("'"),
953 '\u{9}' | '\u{A}' | '\u{D}' | '\u{20}'..='\u{D7FF}' | '\u{E000}'..='\u{FFFD}' => {
954 out.push(c)
955 }
956 _ => {} }
958 }
959 out
960}
961
962fn sanitize_id_component(s: &str) -> String {
966 let mut out = String::with_capacity(s.len());
967 for (i, c) in s.chars().enumerate() {
968 let ok = if i == 0 {
969 c.is_ascii_alphabetic() || c == '_'
970 } else {
971 c.is_ascii_alphanumeric() || c == '_' || c == '.' || c == '-'
972 };
973 out.push(if ok { c } else { '_' });
974 }
975 if out.is_empty() {
976 out.push('_');
977 }
978 out
979}
980
981fn unique_id(base: String, used: &mut BTreeSet<String>) -> String {
985 if used.insert(base.clone()) {
986 return base;
987 }
988 let mut n = 2u32;
989 loop {
990 let candidate = format!("{base}_{n}");
991 if used.insert(candidate.clone()) {
992 return candidate;
993 }
994 n += 1;
995 }
996}
997
998fn resolve_graphml_dest(dest: &Path) -> Result<PathBuf, crate::CliError> {
1004 if dest.is_dir() {
1005 Ok(dest.join("graph.graphml"))
1006 } else {
1007 if let Some(parent) = dest.parent() {
1008 if !parent.as_os_str().is_empty() {
1009 std::fs::create_dir_all(parent)?;
1010 }
1011 }
1012 Ok(dest.to_path_buf())
1013 }
1014}
1015
1016pub fn write_graphml(
1046 nodes: &[NodeInfo],
1047 edges: &[ExportEdge],
1048 dest: &Path,
1049) -> Result<PathBuf, crate::CliError> {
1050 let file_path = resolve_graphml_dest(dest)?;
1051
1052 let mut node_key_types: BTreeMap<String, Option<GmlType>> = BTreeMap::new();
1056 node_key_types.insert("label".to_string(), Some(GmlType::String));
1057 for node in nodes {
1058 for (k, v) in &node.props {
1059 let t = value_gml_type(v);
1060 node_key_types
1061 .entry(k.clone())
1062 .and_modify(|existing| {
1063 if *existing != Some(t) {
1064 *existing = None;
1065 }
1066 })
1067 .or_insert(Some(t));
1068 }
1069 }
1070 let node_key_types: BTreeMap<String, GmlType> = node_key_types
1071 .into_iter()
1072 .map(|(k, t)| (k, t.unwrap_or(GmlType::String)))
1073 .collect();
1074
1075 let mut edge_key_types: Vec<(&str, GmlType)> = vec![
1077 ("type", GmlType::String),
1078 ("derived", GmlType::Boolean),
1079 ("rule", GmlType::String),
1080 ("weight", GmlType::Double),
1081 ];
1082 edge_key_types.sort_by(|a, b| a.0.cmp(b.0));
1083
1084 let mut used_ids: BTreeSet<String> = BTreeSet::new();
1087 let mut node_key_ids: BTreeMap<String, String> = BTreeMap::new();
1088 for name in node_key_types.keys() {
1089 let base = format!("n_{}", sanitize_id_component(name));
1090 node_key_ids.insert(name.clone(), unique_id(base, &mut used_ids));
1091 }
1092 let mut edge_key_ids: BTreeMap<&str, String> = BTreeMap::new();
1093 for (name, _) in &edge_key_types {
1094 let base = format!("e_{}", sanitize_id_component(name));
1095 edge_key_ids.insert(name, unique_id(base, &mut used_ids));
1096 }
1097
1098 let mut out = String::new();
1099 out.push_str("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
1100 out.push_str("<graphml xmlns=\"http://graphml.graphdrawing.org/xmlns\">\n");
1101
1102 for (name, gtype) in &node_key_types {
1103 let _ = writeln!(
1104 out,
1105 " <key id=\"{}\" for=\"node\" attr.name=\"{}\" attr.type=\"{}\"/>",
1106 xml_escape(&node_key_ids[name]),
1107 xml_escape(name),
1108 gtype.as_str()
1109 );
1110 }
1111 for (name, gtype) in &edge_key_types {
1112 let _ = writeln!(
1113 out,
1114 " <key id=\"{}\" for=\"edge\" attr.name=\"{}\" attr.type=\"{}\"/>",
1115 xml_escape(&edge_key_ids[name]),
1116 xml_escape(name),
1117 gtype.as_str()
1118 );
1119 }
1120
1121 out.push_str(" <graph id=\"G\" edgedefault=\"directed\">\n");
1122
1123 for node in nodes {
1124 let _ = writeln!(out, " <node id=\"{}\">", xml_escape(&node.key));
1125 let _ = writeln!(
1126 out,
1127 " <data key=\"{}\">{}</data>",
1128 xml_escape(&node_key_ids["label"]),
1129 xml_escape(&node.label)
1130 );
1131 for (k, v) in &node.props {
1132 let _ = writeln!(
1133 out,
1134 " <data key=\"{}\">{}</data>",
1135 xml_escape(&node_key_ids[k]),
1136 xml_escape(&value_gml_text(v))
1137 );
1138 }
1139 out.push_str(" </node>\n");
1140 }
1141
1142 for (i, edge) in edges.iter().enumerate() {
1143 let _ = writeln!(
1144 out,
1145 " <edge id=\"e{i}\" source=\"{}\" target=\"{}\">",
1146 xml_escape(&edge.src),
1147 xml_escape(&edge.dst)
1148 );
1149 let _ = writeln!(
1150 out,
1151 " <data key=\"{}\">{}</data>",
1152 xml_escape(&edge_key_ids["type"]),
1153 xml_escape(&edge.edge_type)
1154 );
1155 let _ = writeln!(
1156 out,
1157 " <data key=\"{}\">{}</data>",
1158 xml_escape(&edge_key_ids["derived"]),
1159 edge.derived
1160 );
1161 if let Some(rule) = &edge.rule {
1162 let _ = writeln!(
1163 out,
1164 " <data key=\"{}\">{}</data>",
1165 xml_escape(&edge_key_ids["rule"]),
1166 xml_escape(rule)
1167 );
1168 }
1169 if let Some(w) = edge.weight {
1170 let _ = writeln!(
1171 out,
1172 " <data key=\"{}\">{}</data>",
1173 xml_escape(&edge_key_ids["weight"]),
1174 xml_escape(&value_gml_text(&Value::Float(w)))
1175 );
1176 }
1177 out.push_str(" </edge>\n");
1178 }
1179
1180 out.push_str(" </graph>\n");
1181 out.push_str("</graphml>\n");
1182
1183 std::fs::write(&file_path, out)?;
1184 Ok(file_path)
1185}
1186
1187#[cfg(test)]
1188mod v11_value_tests {
1189 use super::{value_gml_text, value_gml_type, value_to_json, GmlType};
1190 use mushroomdb::{json_to_value, Value};
1191 use serde_json::json;
1192
1193 #[test]
1196 fn export_writes_the_tagged_form_that_reimports() {
1197 let date = Value::Date(19_737);
1198 assert_eq!(value_to_json(&date), json!({"$date": "2024-01-15"}));
1199 assert_eq!(json_to_value(value_to_json(&date)), Ok(Some(date.clone())));
1200 let blob = Value::Binary(vec![0xDE, 0xAD, 0xBE, 0xEF]);
1201 assert_eq!(json_to_value(value_to_json(&blob)), Ok(Some(blob.clone())));
1202 assert!(matches!(value_gml_type(&date), GmlType::String));
1204 assert_eq!(value_gml_text(&date), r#"{"$date":"2024-01-15"}"#);
1205 }
1206}
1207
1208#[cfg(test)]
1209mod round_trip_unit_tests {
1210 use super::{cypher_literal, cypher_string, write_csv, CypherSkip};
1211 use crate::import::{parse_node_header, parse_node_row};
1212 use mushroomdb::{NodeInfo, Value};
1213 use std::collections::BTreeMap;
1214
1215 fn typed_values() -> Vec<(&'static str, Value)> {
1216 vec![
1217 ("date", Value::Date(19_737)),
1218 ("ldt", Value::LocalDateTime(1_705_322_096_789_012)),
1219 (
1220 "dt",
1221 Value::DateTime {
1222 micros_utc: 1_705_302_296_789_012,
1223 offset_seconds: 19_800,
1224 },
1225 ),
1226 (
1227 "dur",
1228 Value::Duration {
1229 months: 14,
1230 days: 3,
1231 nanos: 1,
1232 },
1233 ),
1234 ("bin", Value::Binary(vec![0xDE, 0xAD, 0xBE, 0xEF])),
1235 (
1237 "nested",
1238 Value::List(vec![
1239 Value::Date(0),
1240 Value::Map(BTreeMap::from([("b".to_string(), Value::Binary(vec![1]))])),
1241 ]),
1242 ),
1243 ("whole", Value::Float(3.0)),
1244 ("empty", Value::Str(String::new())),
1245 ]
1246 }
1247
1248 #[test]
1252 fn csv_export_of_the_new_types_reads_back_typed() {
1253 let dir = std::env::temp_dir().join(format!("mushroomdb-csv-typed-{}", std::process::id()));
1254 let _ = std::fs::remove_dir_all(&dir);
1255 let node = NodeInfo {
1256 key: "n1".into(),
1257 label: "T".into(),
1258 props: typed_values()
1259 .into_iter()
1260 .map(|(k, v)| (k.to_string(), v))
1261 .collect(),
1262 };
1263 let source = super::Slices {
1264 nodes: std::slice::from_ref(&node),
1265 edges: &[],
1266 };
1267 write_csv(&source, &[], &dir).unwrap();
1268
1269 let mut reader = csv::Reader::from_path(dir.join("nodes.csv")).unwrap();
1270 let header = reader.headers().unwrap().clone();
1271 assert_eq!(&header[0], ":KEY");
1272 let schema = parse_node_header(&header).unwrap();
1273 assert!(schema.warnings.is_empty(), "{:?}", schema.warnings);
1274 let record = reader.records().next().unwrap().unwrap();
1275 let (label, props) = parse_node_row(&schema, &record, ';').unwrap();
1276 assert_eq!(label, "T");
1277 assert_eq!(props, node.props);
1278 let _ = std::fs::remove_dir_all(&dir);
1279 }
1280
1281 #[test]
1282 fn cypher_literals_have_no_exponent_and_refuse_what_has_no_literal() {
1283 let lit = |v: Value| cypher_literal(&v).ok();
1284 assert_eq!(lit(Value::Float(1.0)).as_deref(), Some("1.0"));
1285 assert_eq!(lit(Value::Float(-0.5)).as_deref(), Some("-0.5"));
1286 let huge = lit(Value::Float(1e300)).unwrap();
1287 assert!(!huge.contains('e') && huge.ends_with(".0"), "{huge}");
1288 assert_eq!(lit(Value::Int(-5)).as_deref(), Some("-5"));
1289 assert!(matches!(
1290 cypher_literal(&Value::Int(i64::MIN)),
1291 Err(CypherSkip::Number)
1292 ));
1293 assert!(matches!(
1294 cypher_literal(&Value::Float(f64::NAN)),
1295 Err(CypherSkip::Number)
1296 ));
1297 assert!(matches!(
1298 cypher_literal(&Value::List(vec![Value::Date(1)])),
1299 Err(CypherSkip::Typed)
1300 ));
1301 assert_eq!(
1302 lit(Value::List(vec![Value::Int(1), Value::Str("a".into())])).as_deref(),
1303 Some("[1, 'a']")
1304 );
1305 }
1306
1307 #[test]
1308 fn cypher_strings_stay_on_one_line() {
1309 assert_eq!(
1310 cypher_string("it's \\ \n\r\t\u{1}"),
1311 r"'it\'s \\ \n\r\t\u0001'"
1312 );
1313 assert!(!cypher_string("a\nb").contains('\n'));
1314 }
1315}