1use kyu_api::{Connection, Database};
4use kyu_common::KyuResult;
5use kyu_types::TypedValue;
6
7use crate::state::{EdgeVisual, GraphData, NodeVisual, SchemaData, SchemaTable, Vec2};
8use crate::theme;
9
10pub fn load_schema(db: &Database) -> SchemaData {
15 let catalog = db.catalog();
16 let cat = catalog.read();
17 let conn = db.connect();
18
19 let node_tables = cat
20 .node_tables()
21 .iter()
22 .map(|entry| {
23 let count = count_rows(&conn, &entry.name);
24 SchemaTable {
25 name: entry.name.to_string(),
26 num_rows: count,
27 properties: entry
28 .properties
29 .iter()
30 .map(|p| (p.name.to_string(), format!("{:?}", p.data_type)))
31 .collect(),
32 }
33 })
34 .collect();
35
36 let node_tables_cat = cat.node_tables();
37 let rel_tables = cat
38 .rel_tables()
39 .iter()
40 .map(|entry| {
41 let from_name = node_tables_cat
42 .iter()
43 .find(|n| n.table_id == entry.from_table_id)
44 .map(|n| n.name.as_str())
45 .unwrap_or("_");
46 let to_name = node_tables_cat
47 .iter()
48 .find(|n| n.table_id == entry.to_table_id)
49 .map(|n| n.name.as_str())
50 .unwrap_or("_");
51 let count = count_rel_rows(&conn, from_name, &entry.name, to_name);
52 SchemaTable {
53 name: entry.name.to_string(),
54 num_rows: count,
55 properties: entry
56 .properties
57 .iter()
58 .map(|p| (p.name.to_string(), format!("{:?}", p.data_type)))
59 .collect(),
60 }
61 })
62 .collect();
63
64 SchemaData {
65 node_tables,
66 rel_tables,
67 }
68}
69
70fn count_rows(conn: &Connection, table_name: &str) -> u64 {
72 let query = format!("MATCH (n:{table_name}) RETURN count(*) AS cnt");
73 extract_count(conn, &query)
74}
75
76fn count_rel_rows(conn: &Connection, from: &str, rel: &str, to: &str) -> u64 {
78 let query = format!("MATCH (:{from})-[:{rel}]->(:{to}) RETURN count(*) AS cnt");
79 extract_count(conn, &query)
80}
81
82fn extract_count(conn: &Connection, query: &str) -> u64 {
83 match conn.query(query) {
84 Ok(result) => {
85 if let Some(row) = result.iter_rows().next() {
86 match &row[0] {
87 TypedValue::Int64(n) => *n as u64,
88 _ => 0,
89 }
90 } else {
91 0
92 }
93 }
94 Err(_) => 0,
95 }
96}
97
98pub fn load_node_table(
100 conn: &Connection,
101 db: &Database,
102 table_name: &str,
103 graph: &mut GraphData,
104) -> KyuResult<()> {
105 let catalog = db.catalog();
106 let cat = catalog.read();
107 let node_tables = cat.node_tables();
108 let entry = match node_tables.iter().find(|e| e.name == table_name) {
109 Some(e) => e,
110 None => return Ok(()),
111 };
112
113 let prop_names: Vec<&str> = entry.properties.iter().map(|p| p.name.as_str()).collect();
114 let return_clause = prop_names
115 .iter()
116 .map(|p| format!("n.{p}"))
117 .collect::<Vec<_>>()
118 .join(", ");
119
120 let query = format!("MATCH (n:{table_name}) RETURN {return_clause} LIMIT 500");
121 let result = conn.query(&query)?;
122
123 let num_cols = result.column_names.len();
124 for row in result.iter_rows() {
125 let pk_val = format_typed_value(&row[entry.primary_key_idx]);
126 let id = format!("{table_name}:{pk_val}");
127
128 if graph.node_index.contains_key(&id) {
130 continue;
131 }
132
133 let properties: Vec<(String, String)> = (0..num_cols)
134 .map(|i| (prop_names[i].to_string(), format_typed_value(&row[i])))
135 .collect();
136
137 let node = NodeVisual {
138 id: id.clone(),
139 label: table_name.to_string(),
140 properties,
141 pos: random_position(),
142 vel: Vec2::default(),
143 pinned: false,
144 color_idx: theme::label_color_idx(table_name),
145 };
146
147 let idx = graph.nodes.len();
148 graph.node_index.insert(id, idx);
149 graph.nodes.push(node);
150 }
151
152 Ok(())
153}
154
155pub fn load_rel_table(
157 conn: &Connection,
158 db: &Database,
159 rel_name: &str,
160 graph: &mut GraphData,
161) -> KyuResult<()> {
162 let catalog = db.catalog();
163 let cat = catalog.read();
164
165 let rel_tables = cat.rel_tables();
166 let rel_entry = match rel_tables.iter().find(|e| e.name == rel_name) {
167 Some(e) => e,
168 None => return Ok(()),
169 };
170
171 let node_tables = cat.node_tables();
173 let from_table = node_tables
174 .iter()
175 .find(|e| e.table_id == rel_entry.from_table_id);
176 let to_table = node_tables
177 .iter()
178 .find(|e| e.table_id == rel_entry.to_table_id);
179
180 let (from_name, from_pk) = match from_table {
181 Some(e) => (
182 e.name.as_str(),
183 e.properties[e.primary_key_idx].name.as_str(),
184 ),
185 None => return Ok(()),
186 };
187 let (to_name, to_pk) = match to_table {
188 Some(e) => (
189 e.name.as_str(),
190 e.properties[e.primary_key_idx].name.as_str(),
191 ),
192 None => return Ok(()),
193 };
194
195 let query = format!(
196 "MATCH (a:{from_name})-[:{rel_name}]->(b:{to_name}) RETURN a.{from_pk}, b.{to_pk} LIMIT 1000"
197 );
198 let result = conn.query(&query)?;
199
200 for row in result.iter_rows() {
201 let src_pk = format_typed_value(&row[0]);
202 let dst_pk = format_typed_value(&row[1]);
203 let src_id = format!("{from_name}:{src_pk}");
204 let dst_id = format!("{to_name}:{dst_pk}");
205
206 if let (Some(&src_idx), Some(&dst_idx)) =
207 (graph.node_index.get(&src_id), graph.node_index.get(&dst_id))
208 {
209 graph.edges.push(EdgeVisual {
210 src: src_idx,
211 dst: dst_idx,
212 rel_type: rel_name.to_string(),
213 properties: Vec::new(),
214 });
215 }
216 }
217
218 Ok(())
219}
220
221pub fn load_full_graph(conn: &Connection, db: &Database) -> KyuResult<GraphData> {
223 let mut graph = GraphData::new();
224 let schema = load_schema(db);
225
226 for table in &schema.node_tables {
228 let _ = load_node_table(conn, db, &table.name, &mut graph);
229 }
230
231 for table in &schema.rel_tables {
233 let _ = load_rel_table(conn, db, &table.name, &mut graph);
234 }
235
236 Ok(graph)
237}
238
239pub fn format_typed_value(val: &TypedValue) -> String {
241 match val {
242 TypedValue::Null => "null".to_string(),
243 TypedValue::Bool(b) => b.to_string(),
244 TypedValue::Int64(i) => i.to_string(),
245 TypedValue::Double(f) => format!("{f:.4}"),
246 TypedValue::String(s) => s.to_string(),
247 TypedValue::InternalId(id) => format!("{id}"),
248 other => format!("{other:?}"),
249 }
250}
251
252fn random_position() -> Vec2 {
254 use std::sync::atomic::{AtomicU32, Ordering};
255 static SEED: AtomicU32 = AtomicU32::new(42);
256
257 let s = SEED.fetch_add(1, Ordering::Relaxed);
259 let hash = s.wrapping_mul(2654435761); let x_bits = (hash & 0xFFFF) as f32 / 65535.0; let y_bits = ((hash >> 16) & 0xFFFF) as f32 / 65535.0;
262
263 Vec2::new(x_bits * 400.0 - 200.0, y_bits * 400.0 - 200.0)
265}