Skip to main content

kyu_visualizer/graph/
loader.rs

1//! Load graph data from KyuGraph catalog and query results.
2
3use 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
10/// Load schema metadata from the catalog for the sidebar browser.
11///
12/// Queries actual row counts since the catalog `num_rows` may not reflect
13/// rows inserted via individual CREATE statements.
14pub 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
70/// Query actual row count for a node table.
71fn 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
76/// Query actual row count for a relationship table.
77fn 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
98/// Load all nodes from a specific node table into the graph.
99pub 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        // Skip if already loaded.
129        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
155/// Load all relationships from a specific rel table into the graph.
156pub 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    // Find source and target table names.
172    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
221/// Load all node tables and rel tables from the database.
222pub fn load_full_graph(conn: &Connection, db: &Database) -> KyuResult<GraphData> {
223    let mut graph = GraphData::new();
224    let schema = load_schema(db);
225
226    // Load all node tables first.
227    for table in &schema.node_tables {
228        let _ = load_node_table(conn, db, &table.name, &mut graph);
229    }
230
231    // Then load all relationship tables.
232    for table in &schema.rel_tables {
233        let _ = load_rel_table(conn, db, &table.name, &mut graph);
234    }
235
236    Ok(graph)
237}
238
239/// Format a TypedValue to a display string.
240pub 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
252/// Generate a random initial position spread around the origin.
253fn random_position() -> Vec2 {
254    use std::sync::atomic::{AtomicU32, Ordering};
255    static SEED: AtomicU32 = AtomicU32::new(42);
256
257    // Advance state with a simple LCG.
258    let s = SEED.fetch_add(1, Ordering::Relaxed);
259    let hash = s.wrapping_mul(2654435761); // Knuth's multiplicative hash
260    let x_bits = (hash & 0xFFFF) as f32 / 65535.0; // [0, 1]
261    let y_bits = ((hash >> 16) & 0xFFFF) as f32 / 65535.0;
262
263    // Spread within [-200, 200] — layout will refine positions.
264    Vec2::new(x_bits * 400.0 - 200.0, y_bits * 400.0 - 200.0)
265}