zernel/zql/
executor.rs

1// Copyright (C) 2026 Dyber, Inc. — Proprietary
2
3use crate::experiments::store::ExperimentStore;
4use crate::experiments::tracker;
5use crate::zql::parser::{self, CompareOp, Direction, Value, ZqlQuery};
6use anyhow::Result;
7use std::collections::HashMap;
8
9/// Execute a ZQL query and return formatted results.
10pub fn execute(query: &str) -> Result<String> {
11    let ast = match parser::parse(query) {
12        Ok(ast) => ast,
13        Err(e) => return Ok(format!("Parse error: {e}")),
14    };
15
16    match ast.from.as_str() {
17        "experiments" => execute_experiments(&ast),
18        "jobs" => execute_jobs(&ast),
19        "models" => execute_models(&ast),
20        other => Ok(format!(
21            "Unknown table: '{other}'. Available: experiments, jobs, models"
22        )),
23    }
24}
25
26/// Execute against the experiments SQLite table.
27fn execute_experiments(ast: &ZqlQuery) -> Result<String> {
28    let db_path = tracker::experiments_db_path();
29    let store = ExperimentStore::open(&db_path)?;
30    let all = store.list(10000)?;
31
32    let rows: Vec<HashMap<String, String>> = all
33        .into_iter()
34        .map(|exp| {
35            let mut row = HashMap::new();
36            row.insert("id".into(), exp.id.clone());
37            row.insert("name".into(), exp.name.clone());
38            row.insert("status".into(), exp.status.to_string());
39            row.insert("script".into(), exp.script.clone().unwrap_or_default());
40            row.insert(
41                "duration".into(),
42                exp.duration_secs
43                    .map(|d| format!("{d:.1}s"))
44                    .unwrap_or_default(),
45            );
46            for (k, v) in &exp.metrics {
47                row.insert(k.clone(), format!("{v:.4}"));
48            }
49            row
50        })
51        .collect();
52
53    format_query_results(ast, rows)
54}
55
56/// Execute against the jobs SQLite table.
57fn execute_jobs(ast: &ZqlQuery) -> Result<String> {
58    let db_path = tracker::zernel_dir().join("jobs").join("jobs.db");
59    if !db_path.exists() {
60        return Ok("No jobs database found. Submit a job with: zernel job submit <script>".into());
61    }
62
63    let conn = rusqlite::Connection::open(&db_path)?;
64    let mut stmt = conn.prepare(
65        "SELECT id, script, status, gpus_per_node, nodes, framework, backend, exit_code, submitted_at FROM jobs ORDER BY submitted_at DESC",
66    )?;
67
68    let rows: Vec<HashMap<String, String>> = stmt
69        .query_map([], |row| {
70            let mut map = HashMap::new();
71            map.insert("id".into(), row.get::<_, String>(0)?);
72            map.insert("script".into(), row.get::<_, String>(1)?);
73            map.insert("status".into(), row.get::<_, String>(2)?);
74            map.insert("gpus_per_node".into(), row.get::<_, u32>(3)?.to_string());
75            map.insert("nodes".into(), row.get::<_, u32>(4)?.to_string());
76            map.insert("framework".into(), row.get::<_, String>(5)?);
77            map.insert("backend".into(), row.get::<_, String>(6)?);
78            map.insert(
79                "exit_code".into(),
80                row.get::<_, Option<i32>>(7)?
81                    .map(|e| e.to_string())
82                    .unwrap_or_default(),
83            );
84            map.insert("submitted_at".into(), row.get::<_, String>(8)?);
85            Ok(map)
86        })?
87        .filter_map(|r| r.ok())
88        .collect();
89
90    format_query_results(ast, rows)
91}
92
93/// Execute against the models JSON registry.
94fn execute_models(ast: &ZqlQuery) -> Result<String> {
95    let registry_path = tracker::zernel_dir().join("models").join("registry.json");
96    if !registry_path.exists() {
97        return Ok("No model registry found. Save a model with: zernel model save <path>".into());
98    }
99
100    let data = std::fs::read_to_string(&registry_path)?;
101    let entries: Vec<serde_json::Value> = serde_json::from_str(&data).unwrap_or_default();
102
103    let rows: Vec<HashMap<String, String>> = entries
104        .into_iter()
105        .map(|entry| {
106            let mut map = HashMap::new();
107            if let Some(obj) = entry.as_object() {
108                for (k, v) in obj {
109                    map.insert(
110                        k.clone(),
111                        match v {
112                            serde_json::Value::String(s) => s.clone(),
113                            serde_json::Value::Number(n) => n.to_string(),
114                            other => other.to_string(),
115                        },
116                    );
117                }
118            }
119            map
120        })
121        .collect();
122
123    format_query_results(ast, rows)
124}
125
126/// Generic query executor: applies WHERE, ORDER BY, LIMIT, then formats output.
127fn format_query_results(ast: &ZqlQuery, rows: Vec<HashMap<String, String>>) -> Result<String> {
128    // Apply WHERE filter
129    let mut filtered: Vec<_> = rows
130        .into_iter()
131        .filter(|row| {
132            if let Some(ref wc) = ast.where_clause {
133                wc.conditions.iter().all(|cond| {
134                    let field_val = row.get(&cond.field);
135                    let Some(fv) = field_val else {
136                        return false;
137                    };
138
139                    match &cond.value {
140                        Value::Number(target) => {
141                            let Ok(v) = fv.parse::<f64>() else {
142                                return false;
143                            };
144                            match cond.op {
145                                CompareOp::Eq => (v - target).abs() < f64::EPSILON,
146                                CompareOp::NotEq => (v - target).abs() > f64::EPSILON,
147                                CompareOp::Lt => v < *target,
148                                CompareOp::Gt => v > *target,
149                                CompareOp::Lte => v <= *target,
150                                CompareOp::Gte => v >= *target,
151                            }
152                        }
153                        Value::Text(target) => match cond.op {
154                            CompareOp::Eq => fv == target,
155                            CompareOp::NotEq => fv != target,
156                            _ => false,
157                        },
158                    }
159                })
160            } else {
161                true
162            }
163        })
164        .collect();
165
166    // Apply ORDER BY
167    if let Some(ref ob) = ast.order_by {
168        filtered.sort_by(|a, b| {
169            let va = a
170                .get(&ob.field)
171                .and_then(|s| s.parse::<f64>().ok())
172                .unwrap_or(f64::MAX);
173            let vb = b
174                .get(&ob.field)
175                .and_then(|s| s.parse::<f64>().ok())
176                .unwrap_or(f64::MAX);
177            let cmp = va.partial_cmp(&vb).unwrap_or(std::cmp::Ordering::Equal);
178            match ob.direction {
179                Direction::Asc => cmp,
180                Direction::Desc => cmp.reverse(),
181            }
182        });
183    }
184
185    // Apply LIMIT
186    if let Some(limit) = ast.limit {
187        filtered.truncate(limit);
188    }
189
190    if filtered.is_empty() {
191        return Ok("No results.".into());
192    }
193
194    // Determine columns
195    let columns: Vec<&str> = if ast.select.len() == 1 && ast.select[0] == "*" {
196        // Auto-detect columns from first row
197        let mut cols: Vec<&str> = filtered[0].keys().map(|s| s.as_str()).collect();
198        cols.sort();
199        cols
200    } else {
201        ast.select.iter().map(|s| s.as_str()).collect()
202    };
203
204    // Format output
205    let mut out = String::new();
206    for col in &columns {
207        out.push_str(&format!("{:<24}", col));
208    }
209    out.push('\n');
210    out.push_str(&"-".repeat(columns.len() * 24));
211    out.push('\n');
212
213    for row in &filtered {
214        for col in &columns {
215            let val = row.get(*col).map(|s| s.as_str()).unwrap_or("-");
216            out.push_str(&format!("{:<24}", val));
217        }
218        out.push('\n');
219    }
220
221    out.push_str(&format!("\n{} row(s)", filtered.len()));
222
223    Ok(out)
224}