1use crate::{CodeGraph, Result};
6use std::collections::HashSet;
7use std::fs::File;
8use std::io::Write;
9use std::path::Path;
10
11pub fn export_csv_nodes(graph: &CodeGraph, path: &Path) -> Result<()> {
13 let mut file = File::create(path).map_err(|e| crate::GraphError::Storage {
14 message: format!("Failed to create CSV file: {}", path.display()),
15 source: Some(Box::new(e)),
16 })?;
17
18 let mut all_keys = HashSet::new();
20 for node_id in 0..graph.node_count() as u64 {
21 if let Ok(node) = graph.get_node(node_id) {
22 for (key, _) in node.properties.iter() {
23 all_keys.insert(key.clone());
24 }
25 }
26 }
27
28 let mut keys_vec: Vec<String> = all_keys.into_iter().collect();
29 keys_vec.sort();
30
31 write!(file, "id,type").map_err(|e| crate::GraphError::Storage {
33 message: "Failed to write CSV header".to_string(),
34 source: Some(Box::new(e)),
35 })?;
36 for key in &keys_vec {
37 write!(file, ",{key}").map_err(|e| crate::GraphError::Storage {
38 message: "Failed to write CSV header".to_string(),
39 source: Some(Box::new(e)),
40 })?;
41 }
42 writeln!(file).map_err(|e| crate::GraphError::Storage {
43 message: "Failed to write CSV header".to_string(),
44 source: Some(Box::new(e)),
45 })?;
46
47 for node_id in 0..graph.node_count() as u64 {
49 if let Ok(node) = graph.get_node(node_id) {
50 write!(file, "{},{:?}", node_id, node.node_type).map_err(|e| {
51 crate::GraphError::Storage {
52 message: "Failed to write CSV row".to_string(),
53 source: Some(Box::new(e)),
54 }
55 })?;
56
57 for key in &keys_vec {
58 write!(file, ",").map_err(|e| crate::GraphError::Storage {
59 message: "Failed to write CSV row".to_string(),
60 source: Some(Box::new(e)),
61 })?;
62 if let Some(value) = node.properties.get(key) {
63 write!(file, "{}", escape_csv(&format_property_value(value))).map_err(|e| {
64 crate::GraphError::Storage {
65 message: "Failed to write CSV row".to_string(),
66 source: Some(Box::new(e)),
67 }
68 })?;
69 }
70 }
71 writeln!(file).map_err(|e| crate::GraphError::Storage {
72 message: "Failed to write CSV row".to_string(),
73 source: Some(Box::new(e)),
74 })?;
75 }
76 }
77
78 Ok(())
79}
80
81pub fn export_csv_edges(graph: &CodeGraph, path: &Path) -> Result<()> {
83 let mut file = File::create(path).map_err(|e| crate::GraphError::Storage {
84 message: format!("Failed to create CSV file: {}", path.display()),
85 source: Some(Box::new(e)),
86 })?;
87
88 let mut all_keys = HashSet::new();
90 for edge_id in 0..graph.edge_count() as u64 {
91 if let Ok(edge) = graph.get_edge(edge_id) {
92 for (key, _) in edge.properties.iter() {
93 all_keys.insert(key.clone());
94 }
95 }
96 }
97
98 let mut keys_vec: Vec<String> = all_keys.into_iter().collect();
99 keys_vec.sort();
100
101 write!(file, "id,source,target,type").map_err(|e| crate::GraphError::Storage {
103 message: "Failed to write CSV header".to_string(),
104 source: Some(Box::new(e)),
105 })?;
106 for key in &keys_vec {
107 write!(file, ",{key}").map_err(|e| crate::GraphError::Storage {
108 message: "Failed to write CSV header".to_string(),
109 source: Some(Box::new(e)),
110 })?;
111 }
112 writeln!(file).map_err(|e| crate::GraphError::Storage {
113 message: "Failed to write CSV header".to_string(),
114 source: Some(Box::new(e)),
115 })?;
116
117 for edge_id in 0..graph.edge_count() as u64 {
119 if let Ok(edge) = graph.get_edge(edge_id) {
120 write!(
121 file,
122 "{},{},{},{:?}",
123 edge_id, edge.source_id, edge.target_id, edge.edge_type
124 )
125 .map_err(|e| crate::GraphError::Storage {
126 message: "Failed to write CSV row".to_string(),
127 source: Some(Box::new(e)),
128 })?;
129
130 for key in &keys_vec {
131 write!(file, ",").map_err(|e| crate::GraphError::Storage {
132 message: "Failed to write CSV row".to_string(),
133 source: Some(Box::new(e)),
134 })?;
135 if let Some(value) = edge.properties.get(key) {
136 write!(file, "{}", escape_csv(&format_property_value(value))).map_err(|e| {
137 crate::GraphError::Storage {
138 message: "Failed to write CSV row".to_string(),
139 source: Some(Box::new(e)),
140 }
141 })?;
142 }
143 }
144 writeln!(file).map_err(|e| crate::GraphError::Storage {
145 message: "Failed to write CSV row".to_string(),
146 source: Some(Box::new(e)),
147 })?;
148 }
149 }
150
151 Ok(())
152}
153
154pub fn export_csv(graph: &CodeGraph, nodes_path: &Path, edges_path: &Path) -> Result<()> {
156 export_csv_nodes(graph, nodes_path)?;
157 export_csv_edges(graph, edges_path)?;
158 Ok(())
159}
160
161fn format_property_value(value: &crate::PropertyValue) -> String {
163 match value {
164 crate::PropertyValue::String(s) => s.clone(),
165 crate::PropertyValue::Int(i) => i.to_string(),
166 crate::PropertyValue::Float(f) => f.to_string(),
167 crate::PropertyValue::Bool(b) => b.to_string(),
168 crate::PropertyValue::StringList(v) => v.join(";"),
169 crate::PropertyValue::IntList(v) => v
170 .iter()
171 .map(|i| i.to_string())
172 .collect::<Vec<_>>()
173 .join(";"),
174 crate::PropertyValue::Null => String::new(),
175 }
176}
177
178fn escape_csv(s: &str) -> String {
180 if s.contains(',') || s.contains('"') || s.contains('\n') {
181 format!("\"{}\"", s.replace('"', "\"\""))
182 } else {
183 s.to_string()
184 }
185}
186
187#[cfg(test)]
188mod tests {
189 use super::*;
190
191 #[test]
192 fn test_escape_csv() {
193 assert_eq!(escape_csv("hello"), "hello");
194 assert_eq!(escape_csv("hello,world"), "\"hello,world\"");
195 assert_eq!(escape_csv("say \"hi\""), "\"say \"\"hi\"\"\"");
196 }
197}