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
37type EdgeProps = BTreeMap<String, Value>;
39
40use arrow_array::{ArrayRef, BooleanArray, RecordBatch, StringArray};
41use arrow_schema::{DataType, Field, Schema};
42use parquet::arrow::ArrowWriter;
43use parquet::file::properties::WriterProperties;
44
45#[derive(Debug, Clone, PartialEq, Eq)]
47pub enum ExportFormat {
48 Jsonl,
49 JsonlSingle,
51 Csv,
52 Cypher,
53 Parquet,
54 Graphml,
55}
56
57impl ExportFormat {
58 pub fn parse(s: &str) -> Option<Self> {
59 match s {
60 "jsonl" => Some(Self::Jsonl),
61 "csv" => Some(Self::Csv),
62 "cypher" => Some(Self::Cypher),
63 "parquet" => Some(Self::Parquet),
64 "graphml" => Some(Self::Graphml),
65 _ => None,
66 }
67 }
68
69 pub fn name(&self) -> &'static str {
70 match self {
71 Self::Jsonl => "jsonl",
72 Self::JsonlSingle => "jsonl",
73 Self::Csv => "csv",
74 Self::Cypher => "cypher",
75 Self::Parquet => "parquet",
76 Self::Graphml => "graphml",
77 }
78 }
79}
80
81fn value_to_json(v: &Value) -> serde_json::Value {
86 match v {
87 Value::Int(i) => serde_json::Value::Number((*i).into()),
88 Value::Float(f) => serde_json::Number::from_f64(*f)
90 .map(serde_json::Value::Number)
91 .unwrap_or(serde_json::Value::Null),
92 Value::Str(s) => serde_json::Value::String(s.clone()),
93 Value::Bool(b) => serde_json::Value::Bool(*b),
94 Value::List(xs) => serde_json::Value::Array(xs.iter().map(value_to_json).collect()),
95 Value::Map(m) => serde_json::Value::Object(
96 m.iter()
97 .map(|(k, v)| (k.clone(), value_to_json(v)))
98 .collect(),
99 ),
100 Value::Date(_)
101 | Value::LocalDateTime(_)
102 | Value::DateTime { .. }
103 | Value::Duration { .. }
104 | Value::Binary(_) => mushroomdb::tagged_json(v)
107 .ok()
108 .flatten()
109 .unwrap_or(serde_json::Value::Null),
110 Value::Null => serde_json::Value::Null,
112 }
113}
114
115pub trait ExportSource {
120 fn for_each_node(
121 &self,
122 each: &mut dyn FnMut(NodeInfo) -> Result<(), crate::CliError>,
123 ) -> Result<(), crate::CliError>;
124 fn for_each_edge(
125 &self,
126 each: &mut dyn FnMut(ExportEdge) -> Result<(), crate::CliError>,
127 ) -> Result<(), crate::CliError>;
128 fn for_each_edge_with_props(
134 &self,
135 each: &mut dyn FnMut(ExportEdge, EdgeProps) -> Result<(), crate::CliError>,
136 ) -> Result<(), crate::CliError> {
137 self.for_each_edge(&mut |edge| each(edge, BTreeMap::new()))
138 }
139 fn label_of(&self, key: &str) -> Option<String>;
141}
142
143impl ExportSource for mushroomdb::GraphDb<mushroomdb::RealFs> {
144 fn for_each_node(
145 &self,
146 each: &mut dyn FnMut(NodeInfo) -> Result<(), crate::CliError>,
147 ) -> Result<(), crate::CliError> {
148 self.for_each_node_for_export(each)
149 }
150 fn for_each_edge(
151 &self,
152 each: &mut dyn FnMut(ExportEdge) -> Result<(), crate::CliError>,
153 ) -> Result<(), crate::CliError> {
154 self.for_each_edge_for_export(each)
155 }
156 fn for_each_edge_with_props(
157 &self,
158 each: &mut dyn FnMut(ExportEdge, EdgeProps) -> Result<(), crate::CliError>,
159 ) -> Result<(), crate::CliError> {
160 for (edge, props) in self.all_edges_with_props_for_export() {
161 each(edge, props)?;
162 }
163 Ok(())
164 }
165 fn label_of(&self, key: &str) -> Option<String> {
166 self.node_ref(key).map(|n| n.label().to_string())
167 }
168}
169
170pub struct Slices<'a> {
172 pub nodes: &'a [NodeInfo],
173 pub edges: &'a [ExportEdge],
174}
175
176impl ExportSource for Slices<'_> {
177 fn for_each_node(
178 &self,
179 each: &mut dyn FnMut(NodeInfo) -> Result<(), crate::CliError>,
180 ) -> Result<(), crate::CliError> {
181 self.nodes.iter().try_for_each(|n| each(n.clone()))
182 }
183 fn for_each_edge(
184 &self,
185 each: &mut dyn FnMut(ExportEdge) -> Result<(), crate::CliError>,
186 ) -> Result<(), crate::CliError> {
187 self.edges.iter().try_for_each(|e| each(e.clone()))
188 }
189 fn label_of(&self, key: &str) -> Option<String> {
190 self.nodes
191 .iter()
192 .find(|n| n.key == key)
193 .map(|n| n.label.clone())
194 }
195}
196
197#[derive(Debug, Clone, Default, PartialEq, Eq)]
199pub struct ExportStats {
200 pub nodes: usize,
201 pub edges: usize,
203 pub derived_edges: usize,
204 pub non_finite_floats: usize,
207}
208
209fn count_non_finite(v: &Value) -> usize {
210 match v {
211 Value::Float(f) => usize::from(!f.is_finite()),
212 Value::List(xs) => xs.iter().map(count_non_finite).sum(),
213 Value::Map(m) => m.values().map(count_non_finite).sum(),
214 _ => 0,
215 }
216}
217
218fn json_line(w: &mut impl std::io::Write, v: &serde_json::Value) -> Result<(), crate::CliError> {
219 serde_json::to_writer(&mut *w, v)
220 .map_err(|e| crate::CliError(format!("json encode error: {e}")))?;
221 w.write_all(b"\n")?;
222 Ok(())
223}
224
225fn node_json(node: &NodeInfo, typed: bool) -> serde_json::Value {
231 let mut obj = serde_json::Map::new();
232 obj.insert("key".into(), serde_json::Value::String(node.key.clone()));
233 obj.insert(
234 "label".into(),
235 serde_json::Value::String(node.label.clone()),
236 );
237 let mut reserved = serde_json::Map::new();
238 for (k, v) in &node.props {
239 let own = matches!(k.as_str(), "key" | "label") || k == crate::import::JSONL_PROPS_FIELD;
240 if own || (typed && k == "type") {
241 reserved.insert(k.clone(), value_to_json(v));
242 } else {
243 obj.insert(k.clone(), value_to_json(v));
244 }
245 }
246 if !reserved.is_empty() {
247 obj.insert(
248 crate::import::JSONL_PROPS_FIELD.into(),
249 serde_json::Value::Object(reserved),
250 );
251 }
252 serde_json::Value::Object(obj)
253}
254
255fn edge_json(edge: &ExportEdge, props: &EdgeProps) -> serde_json::Value {
256 let mut obj = serde_json::Map::new();
257 obj.insert("edge_type".into(), edge.edge_type.clone().into());
258 obj.insert("src".into(), edge.src.clone().into());
259 obj.insert("dst".into(), edge.dst.clone().into());
260 obj.insert("derived".into(), edge.derived.into());
261 obj.insert(
262 "rule".into(),
263 edge.rule
264 .clone()
265 .map_or(serde_json::Value::Null, Into::into),
266 );
267 if !props.is_empty() {
270 let props_obj: serde_json::Map<_, _> = props
271 .iter()
272 .map(|(k, v)| (k.clone(), value_to_json(v)))
273 .collect();
274 obj.insert("props".into(), serde_json::Value::Object(props_obj));
275 }
276 serde_json::Value::Object(obj)
277}
278
279fn typed_json_line(
282 w: &mut impl std::io::Write,
283 kind: &str,
284 value: &serde_json::Value,
285) -> Result<(), crate::CliError> {
286 let text = serde_json::to_string(value)
287 .map_err(|e| crate::CliError(format!("json encode error: {e}")))?;
288 let rest = text.strip_prefix('{').unwrap_or(&text);
289 let comma = if rest == "}" { "" } else { "," };
290 writeln!(w, "{{\"type\":\"{kind}\"{comma}{rest}")?;
291 Ok(())
292}
293
294fn create_file(path: PathBuf) -> Result<std::io::BufWriter<std::fs::File>, crate::CliError> {
295 Ok(std::io::BufWriter::with_capacity(
296 1 << 20,
297 std::fs::File::create(path)?,
298 ))
299}
300
301pub fn write_jsonl(
307 source: &dyn ExportSource,
308 rules: &[RuleDef],
309 dest: &Path,
310) -> Result<ExportStats, crate::CliError> {
311 std::fs::create_dir_all(dest)?;
312 let mut stats = ExportStats::default();
313 let mut f = create_file(dest.join("nodes.jsonl"))?;
314 source.for_each_node(&mut |node| {
315 stats.non_finite_floats += node.props.values().map(count_non_finite).sum::<usize>();
316 stats.nodes += 1;
317 json_line(&mut f, &node_json(&node, false))
318 })?;
319 f.flush()?;
320 let mut f = create_file(dest.join("edges.jsonl"))?;
321 source.for_each_edge_with_props(&mut |edge, props| {
322 stats.edges += 1;
323 stats.derived_edges += usize::from(edge.derived);
324 json_line(&mut f, &edge_json(&edge, &props))
325 })?;
326 f.flush()?;
327 write_rules_jsonl(rules, dest)?;
328 Ok(stats)
329}
330
331pub fn write_jsonl_single(
344 source: &dyn ExportSource,
345 rules: &[RuleDef],
346 dest: &Path,
347) -> Result<ExportStats, crate::CliError> {
348 std::fs::create_dir_all(dest)?;
349 let mut stats = ExportStats::default();
350 let mut f = create_file(dest.join("graph.jsonl"))?;
351 source.for_each_node(&mut |node| {
352 stats.non_finite_floats += node.props.values().map(count_non_finite).sum::<usize>();
353 stats.nodes += 1;
354 typed_json_line(&mut f, "node", &node_json(&node, true))
355 })?;
356 source.for_each_edge_with_props(&mut |edge, props| {
357 stats.edges += 1;
358 stats.derived_edges += usize::from(edge.derived);
359 typed_json_line(&mut f, "edge", &edge_json(&edge, &props))
360 })?;
361 for rule in rules {
362 let line = serde_json::json!({ "rule": rule });
363 typed_json_line(&mut f, "rule", &line)?;
364 }
365 f.flush()?;
366 Ok(stats)
367}
368
369fn write_rules_jsonl(rules: &[RuleDef], dest: &Path) -> Result<(), crate::CliError> {
372 let mut f = std::fs::File::create(dest.join("rules.jsonl"))?;
373 for rule in rules {
374 let line = serde_json::to_string(rule)
375 .map_err(|e| crate::CliError(format!("json encode rule: {e}")))?;
376 writeln!(f, "{line}")?;
377 }
378 Ok(())
379}
380
381#[derive(Debug, Clone, Copy, PartialEq, Eq)]
388enum CsvKind {
389 Str,
390 Int,
391 Float,
392 Bool,
393 Date,
394 LocalDateTime,
395 DateTime,
396 Duration,
397 Binary,
398 Json,
399}
400
401impl CsvKind {
402 fn suffix(self) -> &'static str {
403 match self {
404 Self::Str => "string",
405 Self::Int => "int",
406 Self::Float => "float",
407 Self::Bool => "boolean",
408 Self::Date => "date",
409 Self::LocalDateTime => "localdatetime",
410 Self::DateTime => "datetime",
411 Self::Duration => "duration",
412 Self::Binary => "binary",
413 Self::Json => "json",
414 }
415 }
416
417 fn of(v: &Value) -> Self {
419 match v {
420 Value::Str(s) if s.is_empty() => Self::Json,
422 Value::Str(_) => Self::Str,
423 Value::Int(_) => Self::Int,
424 Value::Float(f) if f.is_finite() => Self::Float,
426 Value::Bool(_) => Self::Bool,
427 Value::Date(_) => Self::Date,
428 Value::LocalDateTime(_) => Self::LocalDateTime,
429 Value::DateTime { .. } => Self::DateTime,
430 Value::Duration { .. } => Self::Duration,
431 Value::Binary(_) => Self::Binary,
432 Value::Float(_) | Value::List(_) | Value::Map(_) | Value::Null => Self::Json,
433 }
434 }
435}
436
437fn tagged_text_of(v: &Value) -> String {
439 match mushroomdb::tagged_json(v) {
440 Ok(Some(serde_json::Value::Object(obj))) => obj
441 .into_iter()
442 .next()
443 .and_then(|(_, text)| text.as_str().map(str::to_string))
444 .unwrap_or_default(),
445 _ => String::new(),
446 }
447}
448
449fn csv_cell(v: &Value, kind: CsvKind) -> String {
450 match (kind, v) {
451 (CsvKind::Str, Value::Str(s)) => s.clone(),
452 (CsvKind::Int, Value::Int(i)) => i.to_string(),
453 (CsvKind::Float, Value::Float(f)) => format!("{f:?}"),
456 (CsvKind::Bool, Value::Bool(b)) => b.to_string(),
457 (CsvKind::Json, v) => serde_json::to_string(&value_to_json(v)).unwrap_or_default(),
458 (_, v) => tagged_text_of(v),
459 }
460}
461
462fn csv_error(e: csv::Error) -> crate::CliError {
463 crate::CliError(format!("csv write error: {e}"))
464}
465
466pub fn write_csv(
481 source: &dyn ExportSource,
482 rules: &[RuleDef],
483 dest: &Path,
484) -> Result<ExportStats, crate::CliError> {
485 std::fs::create_dir_all(dest)?;
486 let mut stats = ExportStats::default();
487
488 let mut id_is_key = true;
490 let mut kinds: BTreeMap<String, CsvKind> = BTreeMap::new();
491 source.for_each_node(&mut |node| {
492 if node.label.contains(';') {
493 return Err(crate::CliError(format!(
494 "label `{}` of node `{}` cannot be written to CSV: the importer reads `;` as a multi-label",
495 node.label, node.key
496 )));
497 }
498 stats.nodes += 1;
499 id_is_key &= matches!(node.props.get("id"), Some(Value::Str(s)) if *s == node.key);
500 for (name, value) in &node.props {
501 if name.is_empty() {
502 return Err(crate::CliError(format!(
503 "node `{}` has a property with an empty name, which a CSV header cannot carry",
504 node.key
505 )));
506 }
507 stats.non_finite_floats += count_non_finite(value);
508 let kind = CsvKind::of(value);
509 match kinds.get_mut(name.as_str()) {
510 Some(k) if *k != kind => *k = CsvKind::Json,
511 Some(_) => {}
512 None => {
513 kinds.insert(name.clone(), kind);
514 }
515 }
516 }
517 Ok(())
518 })?;
519 let id_is_key = id_is_key && stats.nodes > 0;
520 if id_is_key {
521 kinds.remove("id");
522 }
523
524 let mut w = csv::Writer::from_path(dest.join("nodes.csv")).map_err(csv_error)?;
526 let mut header = vec![
527 if id_is_key { ":ID" } else { ":KEY" }.to_string(),
528 ":LABEL".to_string(),
529 ];
530 header.extend(kinds.iter().map(|(n, k)| format!("{n}:{}", k.suffix())));
533 w.write_record(&header).map_err(csv_error)?;
534 source.for_each_node(&mut |node| {
535 let mut row = vec![node.key.clone(), node.label.clone()];
536 for (name, kind) in &kinds {
537 row.push(
538 node.props
539 .get(name.as_str())
540 .map(|v| csv_cell(v, *kind))
541 .unwrap_or_default(),
542 );
543 }
544 w.write_record(&row).map_err(csv_error)
545 })?;
546 w.flush()?;
547
548 let mut all_edge_rows: Vec<(ExportEdge, EdgeProps)> = Vec::new();
553 source.for_each_edge_with_props(&mut |edge, props| {
554 all_edge_rows.push((edge, props));
555 Ok(())
556 })?;
557 let mut rel_kinds: BTreeMap<String, CsvKind> = BTreeMap::new();
558 for (edge, props) in &all_edge_rows {
559 if edge.derived {
560 continue;
561 }
562 for (name, value) in props {
563 stats.non_finite_floats += count_non_finite(value);
564 let kind = CsvKind::of(value);
565 match rel_kinds.get_mut(name.as_str()) {
566 Some(k) if *k != kind => *k = CsvKind::Json,
567 Some(_) => {}
568 None => {
569 rel_kinds.insert(name.clone(), kind);
570 }
571 }
572 }
573 }
574 let mut w = csv::Writer::from_path(dest.join("relationships.csv")).map_err(csv_error)?;
576 let mut rel_header = vec![
577 ":START_ID".to_string(),
578 ":END_ID".to_string(),
579 ":TYPE".to_string(),
580 ];
581 rel_header.extend(rel_kinds.iter().map(|(n, k)| format!("{n}:{}", k.suffix())));
582 w.write_record(&rel_header).map_err(csv_error)?;
583 for (edge, props) in all_edge_rows {
584 stats.edges += 1;
585 stats.derived_edges += usize::from(edge.derived);
586 if edge.derived {
587 continue;
588 }
589 let mut row = vec![edge.src.clone(), edge.dst.clone(), edge.edge_type.clone()];
590 for (name, kind) in &rel_kinds {
591 row.push(
592 props
593 .get(name.as_str())
594 .map(|v| csv_cell(v, *kind))
595 .unwrap_or_default(),
596 );
597 }
598 w.write_record(&row).map_err(csv_error)?;
599 }
600 w.flush()?;
601
602 if !rules.is_empty() {
603 write_rules_jsonl(rules, dest)?;
604 }
605 Ok(stats)
606}
607
608#[derive(Debug, Clone, Default, PartialEq, Eq)]
613pub struct CypherLoss {
614 pub typed_properties: usize,
617 pub unwritable_numbers: usize,
621 pub differing_ids: usize,
624 pub ids_added: usize,
627}
628
629impl CypherLoss {
630 pub fn is_empty(&self) -> bool {
631 *self == Self::default()
632 }
633}
634
635fn cypher_label(label: &str) -> String {
638 if label.is_empty() {
639 String::new()
640 } else {
641 format!(":{}", cypher_ident(label))
642 }
643}
644
645fn cypher_ident(name: &str) -> String {
647 format!("`{}`", name.replace('`', "``"))
648}
649
650fn cypher_string(s: &str) -> String {
652 let mut out = String::with_capacity(s.len() + 2);
653 out.push('\'');
654 for c in s.chars() {
655 match c {
656 '\\' => out.push_str("\\\\"),
657 '\'' => out.push_str("\\'"),
658 '\n' => out.push_str("\\n"),
659 '\r' => out.push_str("\\r"),
660 '\t' => out.push_str("\\t"),
661 c if (c as u32) < 0x20 || c as u32 == 0x7f => {
662 let _ = write!(out, "\\u{:04x}", c as u32);
663 }
664 c => out.push(c),
665 }
666 }
667 out.push('\'');
668 out
669}
670
671fn cypher_literal(v: &Value) -> Result<String, CypherSkip> {
673 Ok(match v {
674 Value::Str(s) => cypher_string(s),
675 Value::Int(i64::MIN) => return Err(CypherSkip::Number),
677 Value::Int(i) => i.to_string(),
678 Value::Float(f) if !f.is_finite() => return Err(CypherSkip::Number),
679 Value::Float(f) => {
680 let text = format!("{f}");
682 if text.contains('.') {
683 text
684 } else {
685 format!("{text}.0")
686 }
687 }
688 Value::Bool(b) => b.to_string(),
689 Value::List(xs) => {
690 let items: Result<Vec<String>, _> = xs.iter().map(cypher_literal).collect();
691 format!("[{}]", items?.join(", "))
692 }
693 Value::Map(m) => {
694 let mut items = Vec::with_capacity(m.len());
695 for (k, v) in m {
696 if k.contains(['\n', '\r']) {
698 return Err(CypherSkip::Number);
699 }
700 items.push(format!("{}: {}", cypher_ident(k), cypher_literal(v)?));
701 }
702 format!("{{{}}}", items.join(", "))
703 }
704 Value::Date(_)
705 | Value::LocalDateTime(_)
706 | Value::DateTime { .. }
707 | Value::Duration { .. }
708 | Value::Binary(_) => return Err(CypherSkip::Typed),
709 Value::Null => return Err(CypherSkip::Number),
710 })
711}
712
713enum CypherSkip {
714 Typed,
715 Number,
716}
717
718pub fn write_cypher(
734 source: &dyn ExportSource,
735 dest: &Path,
736) -> Result<(ExportStats, CypherLoss), crate::CliError> {
737 std::fs::create_dir_all(dest)?;
738 let mut out = create_file(dest.join("graph.cypher"))?;
739 let mut loss = CypherLoss::default();
740 let mut stats = ExportStats::default();
741 source.for_each_node(&mut |node| {
742 stats.nodes += 1;
743 let mut sets = Vec::new();
744 let mut has_id = false;
745 for (name, value) in &node.props {
746 if name == "id" {
747 has_id = true;
749 if !matches!(value, Value::Str(s) if *s == node.key) {
750 loss.differing_ids += 1;
751 }
752 continue;
753 }
754 if name.contains(['\n', '\r']) {
755 loss.unwritable_numbers += 1;
756 continue;
757 }
758 match cypher_literal(value) {
759 Ok(lit) => sets.push(format!("n.{} = {lit}", cypher_ident(name))),
760 Err(CypherSkip::Typed) => loss.typed_properties += 1,
761 Err(CypherSkip::Number) => loss.unwritable_numbers += 1,
762 }
763 }
764 loss.ids_added += usize::from(!has_id);
765 write!(
766 out,
767 "MERGE (n{} {{id: {}}})",
768 cypher_label(&node.label),
769 cypher_string(&node.key)
770 )?;
771 if !sets.is_empty() {
772 write!(out, " ON CREATE SET {}", sets.join(", "))?;
773 }
774 writeln!(out)?;
775 Ok(())
776 })?;
777 source.for_each_edge(&mut |edge| {
778 stats.edges += 1;
779 stats.derived_edges += usize::from(edge.derived);
780 if edge.derived {
781 return Ok(());
782 }
783 let (Some(src_label), Some(dst_label)) =
784 (source.label_of(&edge.src), source.label_of(&edge.dst))
785 else {
786 return Ok(());
787 };
788 writeln!(
789 out,
790 "MERGE (a{} {{id: {}}})-[:{}]->(b{} {{id: {}}})",
791 cypher_label(&src_label),
792 cypher_string(&edge.src),
793 cypher_ident(&edge.edge_type),
794 cypher_label(&dst_label),
795 cypher_string(&edge.dst)
796 )?;
797 Ok(())
798 })?;
799 out.flush()?;
800 Ok((stats, loss))
801}
802
803pub fn write_parquet(
817 nodes: &[NodeInfo],
818 edges: &[ExportEdge],
819 rules: &[RuleDef],
820 dest: &Path,
821) -> Result<(), crate::CliError> {
822 std::fs::create_dir_all(dest)?;
823 let props = WriterProperties::builder().build();
824
825 {
827 let schema = Arc::new(Schema::new(vec![
828 Field::new("key", DataType::Utf8, false),
829 Field::new("label", DataType::Utf8, false),
830 Field::new("props", DataType::Utf8, false),
831 ]));
832 let mut keys = Vec::new();
833 let mut labels = Vec::new();
834 let mut props_json = Vec::new();
835 for node in nodes {
836 keys.push(node.key.clone());
837 labels.push(node.label.clone());
838 let p: serde_json::Map<String, serde_json::Value> = node
839 .props
840 .iter()
841 .map(|(k, v)| (k.clone(), value_to_json(v)))
842 .collect();
843 props_json.push(
844 serde_json::to_string(&serde_json::Value::Object(p))
845 .map_err(|e| crate::CliError(format!("json encode: {e}")))?,
846 );
847 }
848 let batch = RecordBatch::try_new(
849 schema.clone(),
850 vec![
851 Arc::new(StringArray::from(keys)) as ArrayRef,
852 Arc::new(StringArray::from(labels)) as ArrayRef,
853 Arc::new(StringArray::from(props_json)) as ArrayRef,
854 ],
855 )
856 .map_err(|e| crate::CliError(format!("arrow error: {e}")))?;
857
858 let file = std::fs::File::create(dest.join("nodes.parquet"))?;
859 let mut writer = ArrowWriter::try_new(file, schema, Some(props.clone()))
860 .map_err(|e| crate::CliError(format!("parquet writer: {e}")))?;
861 writer
862 .write(&batch)
863 .map_err(|e| crate::CliError(format!("parquet write: {e}")))?;
864 writer
865 .close()
866 .map_err(|e| crate::CliError(format!("parquet close: {e}")))?;
867 }
868
869 {
871 let schema = Arc::new(Schema::new(vec![
872 Field::new("edge_type", DataType::Utf8, false),
873 Field::new("src", DataType::Utf8, false),
874 Field::new("dst", DataType::Utf8, false),
875 Field::new("derived", DataType::Boolean, false),
876 Field::new("rule", DataType::Utf8, true),
877 ]));
878 let mut etypes = Vec::new();
879 let mut srcs = Vec::new();
880 let mut dsts = Vec::new();
881 let mut deriveds = Vec::new();
882 let mut rule_names: Vec<Option<String>> = Vec::new();
883 for edge in edges {
884 etypes.push(edge.edge_type.clone());
885 srcs.push(edge.src.clone());
886 dsts.push(edge.dst.clone());
887 deriveds.push(edge.derived);
888 rule_names.push(edge.rule.clone());
889 }
890 let batch = RecordBatch::try_new(
891 schema.clone(),
892 vec![
893 Arc::new(StringArray::from(etypes)) as ArrayRef,
894 Arc::new(StringArray::from(srcs)) as ArrayRef,
895 Arc::new(StringArray::from(dsts)) as ArrayRef,
896 Arc::new(BooleanArray::from(deriveds)) as ArrayRef,
897 Arc::new(StringArray::from(rule_names)) as ArrayRef,
898 ],
899 )
900 .map_err(|e| crate::CliError(format!("arrow error: {e}")))?;
901
902 let file = std::fs::File::create(dest.join("edges.parquet"))?;
903 let mut writer = ArrowWriter::try_new(file, schema, Some(props.clone()))
904 .map_err(|e| crate::CliError(format!("parquet writer: {e}")))?;
905 writer
906 .write(&batch)
907 .map_err(|e| crate::CliError(format!("parquet write: {e}")))?;
908 writer
909 .close()
910 .map_err(|e| crate::CliError(format!("parquet close: {e}")))?;
911 }
912
913 {
915 let schema = Arc::new(Schema::new(vec![
916 Field::new("name", DataType::Utf8, false),
917 Field::new("definition", DataType::Utf8, false),
918 ]));
919 let mut names = Vec::new();
920 let mut definitions = Vec::new();
921 for rule in rules {
922 names.push(rule.name.clone());
923 definitions.push(
924 serde_json::to_string(rule)
925 .map_err(|e| crate::CliError(format!("json encode: {e}")))?,
926 );
927 }
928 let batch = RecordBatch::try_new(
929 schema.clone(),
930 vec![
931 Arc::new(StringArray::from(names)) as ArrayRef,
932 Arc::new(StringArray::from(definitions)) as ArrayRef,
933 ],
934 )
935 .map_err(|e| crate::CliError(format!("arrow error: {e}")))?;
936
937 let file = std::fs::File::create(dest.join("rules.parquet"))?;
938 let mut writer = ArrowWriter::try_new(file, schema, Some(props))
939 .map_err(|e| crate::CliError(format!("parquet writer: {e}")))?;
940 writer
941 .write(&batch)
942 .map_err(|e| crate::CliError(format!("parquet write: {e}")))?;
943 writer
944 .close()
945 .map_err(|e| crate::CliError(format!("parquet close: {e}")))?;
946 }
947
948 Ok(())
949}
950
951#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
955enum GmlType {
956 Boolean,
957 Long,
962 Double,
963 String,
964}
965
966impl GmlType {
967 fn as_str(self) -> &'static str {
968 match self {
969 Self::Boolean => "boolean",
970 Self::Long => "long",
971 Self::Double => "double",
972 Self::String => "string",
973 }
974 }
975}
976
977fn value_gml_type(v: &Value) -> GmlType {
980 match v {
981 Value::Int(_) => GmlType::Long,
982 Value::Float(_) => GmlType::Double,
983 Value::Bool(_) => GmlType::Boolean,
984 Value::Str(_) | Value::List(_) | Value::Map(_) => GmlType::String,
985 Value::Date(_)
987 | Value::LocalDateTime(_)
988 | Value::DateTime { .. }
989 | Value::Duration { .. }
990 | Value::Binary(_) => GmlType::String,
991 Value::Null => GmlType::String,
992 }
993}
994
995fn value_gml_text(v: &Value) -> String {
1003 match value_to_json(v) {
1004 serde_json::Value::Null => String::new(),
1005 serde_json::Value::Bool(b) => b.to_string(),
1006 serde_json::Value::Number(n) => n.to_string(),
1007 serde_json::Value::String(s) => s,
1008 other => other.to_string(),
1009 }
1010}
1011
1012fn xml_escape(s: &str) -> String {
1016 let mut out = String::with_capacity(s.len());
1017 for c in s.chars() {
1018 match c {
1019 '&' => out.push_str("&"),
1020 '<' => out.push_str("<"),
1021 '>' => out.push_str(">"),
1022 '"' => out.push_str("""),
1023 '\'' => out.push_str("'"),
1024 '\u{9}' | '\u{A}' | '\u{D}' | '\u{20}'..='\u{D7FF}' | '\u{E000}'..='\u{FFFD}' => {
1025 out.push(c)
1026 }
1027 _ => {} }
1029 }
1030 out
1031}
1032
1033fn sanitize_id_component(s: &str) -> String {
1037 let mut out = String::with_capacity(s.len());
1038 for (i, c) in s.chars().enumerate() {
1039 let ok = if i == 0 {
1040 c.is_ascii_alphabetic() || c == '_'
1041 } else {
1042 c.is_ascii_alphanumeric() || c == '_' || c == '.' || c == '-'
1043 };
1044 out.push(if ok { c } else { '_' });
1045 }
1046 if out.is_empty() {
1047 out.push('_');
1048 }
1049 out
1050}
1051
1052fn unique_id(base: String, used: &mut BTreeSet<String>) -> String {
1056 if used.insert(base.clone()) {
1057 return base;
1058 }
1059 let mut n = 2u32;
1060 loop {
1061 let candidate = format!("{base}_{n}");
1062 if used.insert(candidate.clone()) {
1063 return candidate;
1064 }
1065 n += 1;
1066 }
1067}
1068
1069fn resolve_graphml_dest(dest: &Path) -> Result<PathBuf, crate::CliError> {
1075 if dest.is_dir() {
1076 Ok(dest.join("graph.graphml"))
1077 } else {
1078 if let Some(parent) = dest.parent() {
1079 if !parent.as_os_str().is_empty() {
1080 std::fs::create_dir_all(parent)?;
1081 }
1082 }
1083 Ok(dest.to_path_buf())
1084 }
1085}
1086
1087pub fn write_graphml(
1117 nodes: &[NodeInfo],
1118 edges: &[ExportEdge],
1119 dest: &Path,
1120) -> Result<PathBuf, crate::CliError> {
1121 let file_path = resolve_graphml_dest(dest)?;
1122
1123 let mut node_key_types: BTreeMap<String, Option<GmlType>> = BTreeMap::new();
1127 node_key_types.insert("label".to_string(), Some(GmlType::String));
1128 for node in nodes {
1129 for (k, v) in &node.props {
1130 let t = value_gml_type(v);
1131 node_key_types
1132 .entry(k.clone())
1133 .and_modify(|existing| {
1134 if *existing != Some(t) {
1135 *existing = None;
1136 }
1137 })
1138 .or_insert(Some(t));
1139 }
1140 }
1141 let node_key_types: BTreeMap<String, GmlType> = node_key_types
1142 .into_iter()
1143 .map(|(k, t)| (k, t.unwrap_or(GmlType::String)))
1144 .collect();
1145
1146 let mut edge_key_types: Vec<(&str, GmlType)> = vec![
1148 ("type", GmlType::String),
1149 ("derived", GmlType::Boolean),
1150 ("rule", GmlType::String),
1151 ("weight", GmlType::Double),
1152 ];
1153 edge_key_types.sort_by(|a, b| a.0.cmp(b.0));
1154
1155 let mut used_ids: BTreeSet<String> = BTreeSet::new();
1158 let mut node_key_ids: BTreeMap<String, String> = BTreeMap::new();
1159 for name in node_key_types.keys() {
1160 let base = format!("n_{}", sanitize_id_component(name));
1161 node_key_ids.insert(name.clone(), unique_id(base, &mut used_ids));
1162 }
1163 let mut edge_key_ids: BTreeMap<&str, String> = BTreeMap::new();
1164 for (name, _) in &edge_key_types {
1165 let base = format!("e_{}", sanitize_id_component(name));
1166 edge_key_ids.insert(name, unique_id(base, &mut used_ids));
1167 }
1168
1169 let mut out = String::new();
1170 out.push_str("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
1171 out.push_str("<graphml xmlns=\"http://graphml.graphdrawing.org/xmlns\">\n");
1172
1173 for (name, gtype) in &node_key_types {
1174 let _ = writeln!(
1175 out,
1176 " <key id=\"{}\" for=\"node\" attr.name=\"{}\" attr.type=\"{}\"/>",
1177 xml_escape(&node_key_ids[name]),
1178 xml_escape(name),
1179 gtype.as_str()
1180 );
1181 }
1182 for (name, gtype) in &edge_key_types {
1183 let _ = writeln!(
1184 out,
1185 " <key id=\"{}\" for=\"edge\" attr.name=\"{}\" attr.type=\"{}\"/>",
1186 xml_escape(&edge_key_ids[name]),
1187 xml_escape(name),
1188 gtype.as_str()
1189 );
1190 }
1191
1192 out.push_str(" <graph id=\"G\" edgedefault=\"directed\">\n");
1193
1194 for node in nodes {
1195 let _ = writeln!(out, " <node id=\"{}\">", xml_escape(&node.key));
1196 let _ = writeln!(
1197 out,
1198 " <data key=\"{}\">{}</data>",
1199 xml_escape(&node_key_ids["label"]),
1200 xml_escape(&node.label)
1201 );
1202 for (k, v) in &node.props {
1203 let _ = writeln!(
1204 out,
1205 " <data key=\"{}\">{}</data>",
1206 xml_escape(&node_key_ids[k]),
1207 xml_escape(&value_gml_text(v))
1208 );
1209 }
1210 out.push_str(" </node>\n");
1211 }
1212
1213 for (i, edge) in edges.iter().enumerate() {
1214 let _ = writeln!(
1215 out,
1216 " <edge id=\"e{i}\" source=\"{}\" target=\"{}\">",
1217 xml_escape(&edge.src),
1218 xml_escape(&edge.dst)
1219 );
1220 let _ = writeln!(
1221 out,
1222 " <data key=\"{}\">{}</data>",
1223 xml_escape(&edge_key_ids["type"]),
1224 xml_escape(&edge.edge_type)
1225 );
1226 let _ = writeln!(
1227 out,
1228 " <data key=\"{}\">{}</data>",
1229 xml_escape(&edge_key_ids["derived"]),
1230 edge.derived
1231 );
1232 if let Some(rule) = &edge.rule {
1233 let _ = writeln!(
1234 out,
1235 " <data key=\"{}\">{}</data>",
1236 xml_escape(&edge_key_ids["rule"]),
1237 xml_escape(rule)
1238 );
1239 }
1240 if let Some(w) = edge.weight {
1241 let _ = writeln!(
1242 out,
1243 " <data key=\"{}\">{}</data>",
1244 xml_escape(&edge_key_ids["weight"]),
1245 xml_escape(&value_gml_text(&Value::Float(w)))
1246 );
1247 }
1248 out.push_str(" </edge>\n");
1249 }
1250
1251 out.push_str(" </graph>\n");
1252 out.push_str("</graphml>\n");
1253
1254 std::fs::write(&file_path, out)?;
1255 Ok(file_path)
1256}
1257
1258#[cfg(test)]
1259mod v11_value_tests {
1260 use super::{value_gml_text, value_gml_type, value_to_json, GmlType};
1261 use mushroomdb::{json_to_value, Value};
1262 use serde_json::json;
1263
1264 #[test]
1267 fn export_writes_the_tagged_form_that_reimports() {
1268 let date = Value::Date(19_737);
1269 assert_eq!(value_to_json(&date), json!({"$date": "2024-01-15"}));
1270 assert_eq!(json_to_value(value_to_json(&date)), Ok(Some(date.clone())));
1271 let blob = Value::Binary(vec![0xDE, 0xAD, 0xBE, 0xEF]);
1272 assert_eq!(json_to_value(value_to_json(&blob)), Ok(Some(blob.clone())));
1273 assert!(matches!(value_gml_type(&date), GmlType::String));
1275 assert_eq!(value_gml_text(&date), r#"{"$date":"2024-01-15"}"#);
1276 }
1277}
1278
1279#[cfg(test)]
1280mod round_trip_unit_tests {
1281 use super::{cypher_literal, cypher_string, write_csv, CypherSkip};
1282 use crate::import::{parse_node_header, parse_node_row};
1283 use mushroomdb::{NodeInfo, Value};
1284 use std::collections::BTreeMap;
1285
1286 fn typed_values() -> Vec<(&'static str, Value)> {
1287 vec![
1288 ("date", Value::Date(19_737)),
1289 ("ldt", Value::LocalDateTime(1_705_322_096_789_012)),
1290 (
1291 "dt",
1292 Value::DateTime {
1293 micros_utc: 1_705_302_296_789_012,
1294 offset_seconds: 19_800,
1295 },
1296 ),
1297 (
1298 "dur",
1299 Value::Duration {
1300 months: 14,
1301 days: 3,
1302 nanos: 1,
1303 },
1304 ),
1305 ("bin", Value::Binary(vec![0xDE, 0xAD, 0xBE, 0xEF])),
1306 (
1308 "nested",
1309 Value::List(vec![
1310 Value::Date(0),
1311 Value::Map(BTreeMap::from([("b".to_string(), Value::Binary(vec![1]))])),
1312 ]),
1313 ),
1314 ("whole", Value::Float(3.0)),
1315 ("empty", Value::Str(String::new())),
1316 ]
1317 }
1318
1319 #[test]
1323 fn csv_export_of_the_new_types_reads_back_typed() {
1324 let dir = std::env::temp_dir().join(format!("mushroomdb-csv-typed-{}", std::process::id()));
1325 let _ = std::fs::remove_dir_all(&dir);
1326 let node = NodeInfo {
1327 key: "n1".into(),
1328 label: "T".into(),
1329 props: typed_values()
1330 .into_iter()
1331 .map(|(k, v)| (k.to_string(), v))
1332 .collect(),
1333 };
1334 let source = super::Slices {
1335 nodes: std::slice::from_ref(&node),
1336 edges: &[],
1337 };
1338 write_csv(&source, &[], &dir).unwrap();
1339
1340 let mut reader = csv::Reader::from_path(dir.join("nodes.csv")).unwrap();
1341 let header = reader.headers().unwrap().clone();
1342 assert_eq!(&header[0], ":KEY");
1343 let schema = parse_node_header(&header).unwrap();
1344 assert!(schema.warnings.is_empty(), "{:?}", schema.warnings);
1345 let record = reader.records().next().unwrap().unwrap();
1346 let (label, props) = parse_node_row(&schema, &record, ';').unwrap();
1347 assert_eq!(label, "T");
1348 assert_eq!(props, node.props);
1349 let _ = std::fs::remove_dir_all(&dir);
1350 }
1351
1352 #[test]
1353 fn cypher_literals_have_no_exponent_and_refuse_what_has_no_literal() {
1354 let lit = |v: Value| cypher_literal(&v).ok();
1355 assert_eq!(lit(Value::Float(1.0)).as_deref(), Some("1.0"));
1356 assert_eq!(lit(Value::Float(-0.5)).as_deref(), Some("-0.5"));
1357 let huge = lit(Value::Float(1e300)).unwrap();
1358 assert!(!huge.contains('e') && huge.ends_with(".0"), "{huge}");
1359 assert_eq!(lit(Value::Int(-5)).as_deref(), Some("-5"));
1360 assert!(matches!(
1361 cypher_literal(&Value::Int(i64::MIN)),
1362 Err(CypherSkip::Number)
1363 ));
1364 assert!(matches!(
1365 cypher_literal(&Value::Float(f64::NAN)),
1366 Err(CypherSkip::Number)
1367 ));
1368 assert!(matches!(
1369 cypher_literal(&Value::List(vec![Value::Date(1)])),
1370 Err(CypherSkip::Typed)
1371 ));
1372 assert_eq!(
1373 lit(Value::List(vec![Value::Int(1), Value::Str("a".into())])).as_deref(),
1374 Some("[1, 'a']")
1375 );
1376 }
1377
1378 #[test]
1379 fn cypher_strings_stay_on_one_line() {
1380 assert_eq!(
1381 cypher_string("it's \\ \n\r\t\u{1}"),
1382 r"'it\'s \\ \n\r\t\u0001'"
1383 );
1384 assert!(!cypher_string("a\nb").contains('\n'));
1385 }
1386}