1#![allow(clippy::missing_errors_doc)]
2#![allow(clippy::missing_panics_doc)]
3
4use anyhow::{Context, Result, bail};
5use std::collections::HashMap;
6use std::fs;
7use std::path::Path;
8use crate::api::RedashClient;
9use crate::models::{QueryMetadata, Parameter};
10use super::OutputFormat;
11
12fn parse_parameter_arg(arg: &str) -> Result<(String, serde_json::Value)> {
13 let parts: Vec<&str> = arg.splitn(2, '=').collect();
14 if parts.len() != 2 {
15 bail!("Invalid parameter format. Use: --param name=value");
16 }
17
18 let name = parts[0].to_string();
19 let value_str = parts[1];
20
21 let value = if let Ok(json_value) = serde_json::from_str(value_str) {
22 json_value
23 } else {
24 serde_json::Value::String(value_str.to_string())
25 };
26
27 Ok((name, value))
28}
29
30fn load_query_metadata_by_id(query_id: u64) -> Result<(QueryMetadata, String, String)> {
31 let queries_dir = Path::new("queries");
32
33 for entry in fs::read_dir(queries_dir).context("Failed to read queries directory")? {
34 let entry = entry.context("Failed to read directory entry")?;
35 let path = entry.path();
36
37 if path.extension().is_some_and(|ext| ext == "yaml")
38 && let Some(filename) = path.file_name().and_then(|f| f.to_str())
39 && let Some(id_str) = filename.split('-').next()
40 && let Ok(id) = id_str.parse::<u64>()
41 && id == query_id
42 {
43 let yaml_content = fs::read_to_string(&path)
44 .context(format!("Failed to read {}", path.display()))?;
45
46 let metadata: QueryMetadata = serde_yaml::from_str(&yaml_content)
47 .context(format!("Failed to parse {}", path.display()))?;
48
49 let yaml_path = path.display().to_string();
50 let sql_path = yaml_path.replace(".yaml", ".sql");
51
52 if !Path::new(&sql_path).exists() {
53 bail!("SQL file not found: {sql_path}");
54 }
55
56 let sql = fs::read_to_string(&sql_path)
57 .context(format!("Failed to read {sql_path}"))?;
58
59 return Ok((metadata, sql, yaml_path));
60 }
61 }
62
63 bail!("Query {query_id} not found in queries/ directory. Run 'stmo-cli fetch {query_id}' first.");
64}
65
66fn prompt_for_parameter(param: &Parameter) -> Result<serde_json::Value> {
67 use dialoguer::{Input, Select};
68
69 let title = ¶m.title;
70
71 match param.param_type.as_str() {
72 "date" => {
73 let input: String = Input::new()
74 .with_prompt(format!("{title} (YYYY-MM-DD)"))
75 .interact_text()?;
76 Ok(serde_json::Value::String(input))
77 }
78 "enum" => {
79 if let Some(enum_options) = ¶m.enum_options {
80 let options: Vec<&str> = enum_options.lines().collect();
81
82 if param.multi_values_options.is_some() {
83 use dialoguer::MultiSelect;
84 let selections = MultiSelect::new()
85 .with_prompt(title)
86 .items(&options)
87 .interact()?;
88
89 let selected: Vec<String> = selections.iter()
90 .map(|&i| options[i].to_string())
91 .collect();
92
93 Ok(serde_json::Value::Array(
94 selected.into_iter().map(serde_json::Value::String).collect()
95 ))
96 } else {
97 let selection = Select::new()
98 .with_prompt(title)
99 .items(&options)
100 .default(0)
101 .interact()?;
102
103 Ok(serde_json::Value::String(options[selection].to_string()))
104 }
105 } else {
106 let input: String = Input::new()
107 .with_prompt(title)
108 .interact_text()?;
109 Ok(serde_json::Value::String(input))
110 }
111 }
112 "number" => {
113 let input: f64 = Input::new()
114 .with_prompt(title)
115 .interact_text()?;
116 Ok(serde_json::json!(input))
117 }
118 "text" => {
119 let input: String = Input::new()
120 .with_prompt(title)
121 .interact_text()?;
122 Ok(serde_json::Value::String(input))
123 }
124 _ => {
125 let input: String = Input::new()
126 .with_prompt(title)
127 .interact_text()?;
128 Ok(serde_json::Value::String(input))
129 }
130 }
131}
132
133fn build_parameter_map(
134 metadata: &QueryMetadata,
135 cli_params: &[(String, serde_json::Value)],
136 interactive: bool,
137) -> Result<Option<HashMap<String, serde_json::Value>>> {
138 if metadata.options.parameters.is_empty() {
139 return Ok(None);
140 }
141
142 let mut param_map = HashMap::new();
143
144 for (name, value) in cli_params {
145 param_map.insert(name.clone(), value.clone());
146 }
147
148 for param in &metadata.options.parameters {
149 if !param_map.contains_key(¶m.name) {
150 if interactive {
151 eprintln!("\nParameter '{}' required:", param.title);
152 let value = prompt_for_parameter(param)?;
153 param_map.insert(param.name.clone(), value);
154 } else if let Some(default_value) = ¶m.value {
155 param_map.insert(param.name.clone(), default_value.clone());
156 } else {
157 bail!(
158 "Missing required parameter: '{}' ({}). Use --param {}=value or --interactive",
159 param.name, param.title, param.name
160 );
161 }
162 }
163 }
164
165 Ok(if param_map.is_empty() { None } else { Some(param_map) })
166}
167
168fn format_results_json(result: &crate::models::QueryResult, limit: Option<usize>) -> Result<String> {
169 let rows = if let Some(limit) = limit {
170 result.data.rows.iter().take(limit).cloned().collect::<Vec<_>>()
171 } else {
172 result.data.rows.clone()
173 };
174
175 let formatted = serde_json::json!({
176 "columns": result.data.columns.iter().map(|c| &c.name).collect::<Vec<_>>(),
177 "rows": rows,
178 "row_count": result.data.rows.len(),
179 "runtime_seconds": result.runtime,
180 });
181
182 serde_json::to_string_pretty(&formatted)
183 .context("Failed to format results as JSON")
184}
185
186fn format_results_table(result: &crate::models::QueryResult, limit: Option<usize>) -> String {
187 use std::fmt::Write;
188
189 let mut output = String::new();
190
191 let _ = writeln!(output);
192 for col in &result.data.columns {
193 let _ = write!(output, "{:20} ", col.name);
194 }
195 let _ = writeln!(output);
196 let _ = writeln!(output, "{}", "-".repeat(result.data.columns.len() * 21));
197
198 let rows_to_show = limit.unwrap_or(result.data.rows.len()).min(result.data.rows.len());
199
200 for row in &result.data.rows[..rows_to_show] {
201 if let serde_json::Value::Object(obj) = row {
202 for col in &result.data.columns {
203 let value = obj.get(&col.name)
204 .map(|v| match v {
205 serde_json::Value::Null => "NULL".to_string(),
206 serde_json::Value::String(s) => s.clone(),
207 _ => v.to_string(),
208 })
209 .unwrap_or_default();
210
211 let truncated = if value.len() > 18 {
212 format!("{}...", &value[..15])
213 } else {
214 value
215 };
216
217 let _ = write!(output, "{truncated:20} ");
218 }
219 let _ = writeln!(output);
220 }
221 }
222
223 if rows_to_show < result.data.rows.len() {
224 let _ = write!(output, "\n... {} more rows (showing {} of {})\n",
225 result.data.rows.len() - rows_to_show,
226 rows_to_show,
227 result.data.rows.len()
228 );
229 }
230
231 let _ = write!(output, "\n✓ {} rows returned in {:.2}s\n",
232 result.data.rows.len(), result.runtime);
233
234 output
235}
236
237pub async fn execute(
238 client: &RedashClient,
239 query_id: u64,
240 param_args: Vec<String>,
241 format: OutputFormat,
242 interactive: bool,
243 timeout_secs: u64,
244 limit_rows: Option<usize>,
245) -> Result<()> {
246 let (metadata, _sql, yaml_path) = load_query_metadata_by_id(query_id)?;
247
248 println!("Executing query: {} - {}", metadata.id, metadata.name);
249 println!("Source: {yaml_path}\n");
250
251 let cli_params: Vec<(String, serde_json::Value)> = param_args
252 .iter()
253 .map(|arg| parse_parameter_arg(arg))
254 .collect::<Result<Vec<_>>>()?;
255
256 let parameters = build_parameter_map(&metadata, &cli_params, interactive)?;
257
258 if let Some(ref params) = parameters {
259 eprintln!("Parameters:");
260 for (name, value) in params {
261 eprintln!(" {name} = {value}");
262 }
263 eprintln!();
264 }
265
266 let result = client
267 .execute_query_with_polling(query_id, parameters, timeout_secs, 500)
268 .await?;
269
270 match format {
271 OutputFormat::Json => {
272 let json = format_results_json(&result, limit_rows)?;
273 println!("\n{json}");
274 }
275 OutputFormat::Table => {
276 let table = format_results_table(&result, limit_rows);
277 println!("{table}");
278 }
279 }
280
281 Ok(())
282}
283
284#[cfg(test)]
285#[allow(clippy::missing_errors_doc)]
286mod tests {
287 use super::*;
288 use crate::models::{QueryResult, QueryResultData, Column};
289
290 #[test]
291 fn test_parse_parameter_arg_string() {
292 let result = parse_parameter_arg("name=value").unwrap();
293 assert_eq!(result.0, "name");
294 assert_eq!(result.1, serde_json::Value::String("value".to_string()));
295 }
296
297 #[test]
298 fn test_parse_parameter_arg_json_array() {
299 let result = parse_parameter_arg("channels=[\"release\",\"beta\"]").unwrap();
300 assert_eq!(result.0, "channels");
301 assert_eq!(result.1, serde_json::json!(["release", "beta"]));
302 }
303
304 #[test]
305 fn test_parse_parameter_arg_number() {
306 let result = parse_parameter_arg("count=42").unwrap();
307 assert_eq!(result.0, "count");
308 assert_eq!(result.1, serde_json::json!(42));
309 }
310
311 #[test]
312 fn test_parse_parameter_arg_invalid() {
313 let result = parse_parameter_arg("invalid");
314 assert!(result.is_err());
315 assert!(result.unwrap_err().to_string().contains("Invalid parameter format"));
316 }
317
318 #[test]
319 fn test_format_results_json() {
320 let result = QueryResult {
321 id: 1,
322 data: QueryResultData {
323 columns: vec![
324 Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
325 Column { name: "col2".to_string(), type_name: "integer".to_string(), friendly_name: None },
326 ],
327 rows: vec![
328 serde_json::json!({"col1": "value1", "col2": 123}),
329 serde_json::json!({"col1": "value2", "col2": 456}),
330 ],
331 },
332 runtime: 1.5,
333 retrieved_at: "2026-01-21T10:00:00".to_string(),
334 };
335
336 let json = format_results_json(&result, None).unwrap();
337 let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
338
339 assert_eq!(parsed["row_count"], 2);
340 assert_eq!(parsed["runtime_seconds"], 1.5);
341 assert_eq!(parsed["columns"], serde_json::json!(["col1", "col2"]));
342 assert_eq!(parsed["rows"][0]["col1"], "value1");
343 }
344
345 #[test]
346 fn test_format_results_json_with_limit() {
347 let result = QueryResult {
348 id: 1,
349 data: QueryResultData {
350 columns: vec![
351 Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
352 ],
353 rows: vec![
354 serde_json::json!({"col1": "row1"}),
355 serde_json::json!({"col1": "row2"}),
356 serde_json::json!({"col1": "row3"}),
357 ],
358 },
359 runtime: 1.0,
360 retrieved_at: "2026-01-21T10:00:00".to_string(),
361 };
362
363 let json = format_results_json(&result, Some(2)).unwrap();
364 let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
365
366 assert_eq!(parsed["row_count"], 3);
367 assert_eq!(parsed["rows"].as_array().unwrap().len(), 2);
368 }
369
370 #[test]
371 fn test_format_results_table() {
372 let result = QueryResult {
373 id: 1,
374 data: QueryResultData {
375 columns: vec![
376 Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
377 Column { name: "col2".to_string(), type_name: "integer".to_string(), friendly_name: None },
378 ],
379 rows: vec![
380 serde_json::json!({"col1": "value1", "col2": 123}),
381 serde_json::json!({"col1": "value2", "col2": 456}),
382 ],
383 },
384 runtime: 1.5,
385 retrieved_at: "2026-01-21T10:00:00".to_string(),
386 };
387
388 let table = format_results_table(&result, None);
389
390 assert!(table.contains("col1"));
391 assert!(table.contains("col2"));
392 assert!(table.contains("value1"));
393 assert!(table.contains("value2"));
394 assert!(table.contains("2 rows returned"));
395 }
396
397 #[test]
398 fn test_format_results_table_with_limit() {
399 let result = QueryResult {
400 id: 1,
401 data: QueryResultData {
402 columns: vec![
403 Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
404 ],
405 rows: vec![
406 serde_json::json!({"col1": "row1"}),
407 serde_json::json!({"col1": "row2"}),
408 serde_json::json!({"col1": "row3"}),
409 ],
410 },
411 runtime: 1.0,
412 retrieved_at: "2026-01-21T10:00:00".to_string(),
413 };
414
415 let table = format_results_table(&result, Some(2));
416
417 assert!(table.contains("row1"));
418 assert!(table.contains("row2"));
419 assert!(table.contains("... 1 more rows"));
420 assert!(table.contains("3 rows returned"));
421 }
422
423 #[test]
424 fn test_format_results_table_truncation() {
425 let result = QueryResult {
426 id: 1,
427 data: QueryResultData {
428 columns: vec![
429 Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
430 ],
431 rows: vec![
432 serde_json::json!({"col1": "this_is_a_very_long_value_that_should_be_truncated"}),
433 ],
434 },
435 runtime: 1.0,
436 retrieved_at: "2026-01-21T10:00:00".to_string(),
437 };
438
439 let table = format_results_table(&result, None);
440
441 assert!(table.contains("..."));
442 }
443
444 #[test]
445 fn test_output_format_from_str() {
446 assert!(matches!("json".parse::<OutputFormat>().unwrap(), OutputFormat::Json));
447 assert!(matches!("JSON".parse::<OutputFormat>().unwrap(), OutputFormat::Json));
448 assert!(matches!("table".parse::<OutputFormat>().unwrap(), OutputFormat::Table));
449 assert!(matches!("TABLE".parse::<OutputFormat>().unwrap(), OutputFormat::Table));
450 }
451
452 #[test]
453 fn test_output_format_from_str_invalid() {
454 let result = "csv".parse::<OutputFormat>();
455 assert!(result.is_err());
456 assert!(result.unwrap_err().to_string().contains("Invalid format"));
457 }
458}