Skip to main content

stmo_cli/commands/
execute.rs

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 = &param.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) = &param.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(&param.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) = &param.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    serde_json::to_string_pretty(&rows)
176        .context("Failed to format results as JSON")
177}
178
179fn format_results_table(result: &crate::models::QueryResult, limit: Option<usize>) -> String {
180    use std::fmt::Write;
181
182    let mut output = String::new();
183
184    let _ = writeln!(output);
185    for col in &result.data.columns {
186        let _ = write!(output, "{:20} ", col.name);
187    }
188    let _ = writeln!(output);
189    let _ = writeln!(output, "{}", "-".repeat(result.data.columns.len() * 21));
190
191    let rows_to_show = limit.unwrap_or(result.data.rows.len()).min(result.data.rows.len());
192
193    for row in &result.data.rows[..rows_to_show] {
194        if let serde_json::Value::Object(obj) = row {
195            for col in &result.data.columns {
196                let value = obj.get(&col.name)
197                    .map(|v| match v {
198                        serde_json::Value::Null => "NULL".to_string(),
199                        serde_json::Value::String(s) => s.clone(),
200                        _ => v.to_string(),
201                    })
202                    .unwrap_or_default();
203
204                let truncated = if value.len() > 18 {
205                    format!("{}...", &value[..15])
206                } else {
207                    value
208                };
209
210                let _ = write!(output, "{truncated:20} ");
211            }
212            let _ = writeln!(output);
213        }
214    }
215
216    if rows_to_show < result.data.rows.len() {
217        let _ = write!(output, "\n... {} more rows (showing {} of {})\n",
218            result.data.rows.len() - rows_to_show,
219            rows_to_show,
220            result.data.rows.len()
221        );
222    }
223
224    let _ = write!(output, "\n✓ {} rows returned in {:.2}s\n",
225        result.data.rows.len(), result.runtime);
226
227    output
228}
229
230pub async fn execute(
231    client: &RedashClient,
232    query_id: u64,
233    param_args: Vec<String>,
234    format: OutputFormat,
235    interactive: bool,
236    timeout_secs: u64,
237    limit_rows: Option<usize>,
238) -> Result<()> {
239    let (metadata, _sql, yaml_path) = load_query_metadata_by_id(query_id)?;
240
241    eprintln!("Executing query: {} - {}", metadata.id, metadata.name);
242    eprintln!("Source: {yaml_path}\n");
243
244    let cli_params: Vec<(String, serde_json::Value)> = param_args
245        .iter()
246        .map(|arg| parse_parameter_arg(arg))
247        .collect::<Result<Vec<_>>>()?;
248
249    let parameters = build_parameter_map(&metadata, &cli_params, interactive)?;
250
251    if let Some(ref params) = parameters {
252        eprintln!("Parameters:");
253        for (name, value) in params {
254            eprintln!("  {name} = {value}");
255        }
256        eprintln!();
257    }
258
259    let result = client
260        .execute_query_with_polling(query_id, parameters, timeout_secs, 500)
261        .await?;
262
263    match format {
264        OutputFormat::Json => {
265            let json = format_results_json(&result, limit_rows)?;
266            println!("{json}");
267        }
268        OutputFormat::Table => {
269            let table = format_results_table(&result, limit_rows);
270            println!("{table}");
271        }
272    }
273
274    Ok(())
275}
276
277#[cfg(test)]
278#[allow(clippy::missing_errors_doc)]
279mod tests {
280    use super::*;
281    use crate::models::{QueryResult, QueryResultData, Column};
282
283    #[test]
284    fn test_parse_parameter_arg_string() {
285        let result = parse_parameter_arg("name=value").unwrap();
286        assert_eq!(result.0, "name");
287        assert_eq!(result.1, serde_json::Value::String("value".to_string()));
288    }
289
290    #[test]
291    fn test_parse_parameter_arg_json_array() {
292        let result = parse_parameter_arg("channels=[\"release\",\"beta\"]").unwrap();
293        assert_eq!(result.0, "channels");
294        assert_eq!(result.1, serde_json::json!(["release", "beta"]));
295    }
296
297    #[test]
298    fn test_parse_parameter_arg_number() {
299        let result = parse_parameter_arg("count=42").unwrap();
300        assert_eq!(result.0, "count");
301        assert_eq!(result.1, serde_json::json!(42));
302    }
303
304    #[test]
305    fn test_parse_parameter_arg_invalid() {
306        let result = parse_parameter_arg("invalid");
307        assert!(result.is_err());
308        assert!(result.unwrap_err().to_string().contains("Invalid parameter format"));
309    }
310
311    #[test]
312    fn test_format_results_json() {
313        let result = QueryResult {
314            id: 1,
315            data: QueryResultData {
316                columns: vec![
317                    Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
318                    Column { name: "col2".to_string(), type_name: "integer".to_string(), friendly_name: None },
319                ],
320                rows: vec![
321                    serde_json::json!({"col1": "value1", "col2": 123}),
322                    serde_json::json!({"col1": "value2", "col2": 456}),
323                ],
324            },
325            runtime: 1.5,
326            retrieved_at: "2026-01-21T10:00:00".to_string(),
327        };
328
329        let json = format_results_json(&result, None).unwrap();
330        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
331
332        let rows = parsed.as_array().unwrap();
333        assert_eq!(rows.len(), 2);
334        assert_eq!(rows[0]["col1"], "value1");
335        assert_eq!(rows[0]["col2"], 123);
336    }
337
338    #[test]
339    fn test_format_results_json_with_limit() {
340        let result = QueryResult {
341            id: 1,
342            data: QueryResultData {
343                columns: vec![
344                    Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
345                ],
346                rows: vec![
347                    serde_json::json!({"col1": "row1"}),
348                    serde_json::json!({"col1": "row2"}),
349                    serde_json::json!({"col1": "row3"}),
350                ],
351            },
352            runtime: 1.0,
353            retrieved_at: "2026-01-21T10:00:00".to_string(),
354        };
355
356        let json = format_results_json(&result, Some(2)).unwrap();
357        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
358
359        assert_eq!(parsed.as_array().unwrap().len(), 2);
360    }
361
362    #[test]
363    fn test_format_results_table() {
364        let result = QueryResult {
365            id: 1,
366            data: QueryResultData {
367                columns: vec![
368                    Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
369                    Column { name: "col2".to_string(), type_name: "integer".to_string(), friendly_name: None },
370                ],
371                rows: vec![
372                    serde_json::json!({"col1": "value1", "col2": 123}),
373                    serde_json::json!({"col1": "value2", "col2": 456}),
374                ],
375            },
376            runtime: 1.5,
377            retrieved_at: "2026-01-21T10:00:00".to_string(),
378        };
379
380        let table = format_results_table(&result, None);
381
382        assert!(table.contains("col1"));
383        assert!(table.contains("col2"));
384        assert!(table.contains("value1"));
385        assert!(table.contains("value2"));
386        assert!(table.contains("2 rows returned"));
387    }
388
389    #[test]
390    fn test_format_results_table_with_limit() {
391        let result = QueryResult {
392            id: 1,
393            data: QueryResultData {
394                columns: vec![
395                    Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
396                ],
397                rows: vec![
398                    serde_json::json!({"col1": "row1"}),
399                    serde_json::json!({"col1": "row2"}),
400                    serde_json::json!({"col1": "row3"}),
401                ],
402            },
403            runtime: 1.0,
404            retrieved_at: "2026-01-21T10:00:00".to_string(),
405        };
406
407        let table = format_results_table(&result, Some(2));
408
409        assert!(table.contains("row1"));
410        assert!(table.contains("row2"));
411        assert!(table.contains("... 1 more rows"));
412        assert!(table.contains("3 rows returned"));
413    }
414
415    #[test]
416    fn test_format_results_table_truncation() {
417        let result = QueryResult {
418            id: 1,
419            data: QueryResultData {
420                columns: vec![
421                    Column { name: "col1".to_string(), type_name: "string".to_string(), friendly_name: None },
422                ],
423                rows: vec![
424                    serde_json::json!({"col1": "this_is_a_very_long_value_that_should_be_truncated"}),
425                ],
426            },
427            runtime: 1.0,
428            retrieved_at: "2026-01-21T10:00:00".to_string(),
429        };
430
431        let table = format_results_table(&result, None);
432
433        assert!(table.contains("..."));
434    }
435
436    #[test]
437    fn test_output_format_from_str() {
438        assert!(matches!("json".parse::<OutputFormat>().unwrap(), OutputFormat::Json));
439        assert!(matches!("JSON".parse::<OutputFormat>().unwrap(), OutputFormat::Json));
440        assert!(matches!("table".parse::<OutputFormat>().unwrap(), OutputFormat::Table));
441        assert!(matches!("TABLE".parse::<OutputFormat>().unwrap(), OutputFormat::Table));
442    }
443
444    #[test]
445    fn test_output_format_from_str_invalid() {
446        let result = "csv".parse::<OutputFormat>();
447        assert!(result.is_err());
448        assert!(result.unwrap_err().to_string().contains("Invalid format"));
449    }
450}