1use 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
9pub 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
26fn 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
56fn 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
93fn 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(®istry_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
126fn format_query_results(ast: &ZqlQuery, rows: Vec<HashMap<String, String>>) -> Result<String> {
128 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 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 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 let columns: Vec<&str> = if ast.select.len() == 1 && ast.select[0] == "*" {
196 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 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}