Skip to main content

stmo_cli/commands/
execute.rs

1#![allow(clippy::missing_errors_doc)]
2#![allow(clippy::missing_panics_doc)]
3
4use super::OutputFormat;
5use crate::api::RedashClient;
6use crate::models::{Parameter, QueryMetadata};
7use anyhow::{Context, Result, bail};
8use std::collections::HashMap;
9use std::fs;
10use std::io::IsTerminal;
11use std::path::Path;
12
13fn parse_parameter_arg(arg: &str) -> Result<(String, serde_json::Value)> {
14    let parts: Vec<&str> = arg.splitn(2, '=').collect();
15    if parts.len() != 2 {
16        bail!("Invalid parameter format. Use: --param name=value");
17    }
18
19    let name = parts[0].to_string();
20    let value_str = parts[1];
21
22    let value = if let Ok(json_value) = serde_json::from_str(value_str) {
23        json_value
24    } else {
25        serde_json::Value::String(value_str.to_string())
26    };
27
28    Ok((name, value))
29}
30
31fn load_query_metadata_by_id(query_id: u64) -> Result<(QueryMetadata, String, String)> {
32    let queries_dir = Path::new("queries");
33
34    for entry in fs::read_dir(queries_dir).context("Failed to read queries directory")? {
35        let entry = entry.context("Failed to read directory entry")?;
36        let path = entry.path();
37
38        if path.extension().is_some_and(|ext| ext == "yaml")
39            && let Some(filename) = path.file_name().and_then(|f| f.to_str())
40            && let Some(id_str) = filename.split('-').next()
41            && let Ok(id) = id_str.parse::<u64>()
42            && id == query_id
43        {
44            let yaml_content =
45                fs::read_to_string(&path).context(format!("Failed to read {}", path.display()))?;
46
47            let metadata: QueryMetadata = serde_yaml::from_str(&yaml_content)
48                .context(format!("Failed to parse {}", path.display()))?;
49
50            let yaml_path = path.display().to_string();
51            let sql_path = yaml_path.replace(".yaml", ".sql");
52
53            if !Path::new(&sql_path).exists() {
54                bail!("SQL file not found: {sql_path}");
55            }
56
57            let sql =
58                fs::read_to_string(&sql_path).context(format!("Failed to read {sql_path}"))?;
59
60            return Ok((metadata, sql, yaml_path));
61        }
62    }
63
64    bail!(
65        "Query {query_id} not found in queries/ directory. Run 'stmo-cli fetch {query_id}' first."
66    );
67}
68
69fn prompt_for_parameter(param: &Parameter) -> Result<serde_json::Value> {
70    use dialoguer::{Input, Select};
71
72    let title = &param.title;
73
74    match param.param_type.as_str() {
75        "date" => {
76            let input: String = Input::new()
77                .with_prompt(format!("{title} (YYYY-MM-DD)"))
78                .interact_text()?;
79            Ok(serde_json::Value::String(input))
80        }
81        "enum" => {
82            if let Some(enum_options) = &param.enum_options {
83                let options: Vec<&str> = enum_options.lines().collect();
84
85                if param.multi_values_options.is_some() {
86                    use dialoguer::MultiSelect;
87                    let selections = MultiSelect::new()
88                        .with_prompt(title)
89                        .items(&options)
90                        .interact()?;
91
92                    let selected: Vec<String> =
93                        selections.iter().map(|&i| options[i].to_string()).collect();
94
95                    Ok(serde_json::Value::Array(
96                        selected
97                            .into_iter()
98                            .map(serde_json::Value::String)
99                            .collect(),
100                    ))
101                } else {
102                    let selection = Select::new()
103                        .with_prompt(title)
104                        .items(&options)
105                        .default(0)
106                        .interact()?;
107
108                    Ok(serde_json::Value::String(options[selection].to_string()))
109                }
110            } else {
111                let input: String = Input::new().with_prompt(title).interact_text()?;
112                Ok(serde_json::Value::String(input))
113            }
114        }
115        "number" => {
116            let input: f64 = Input::new().with_prompt(title).interact_text()?;
117            Ok(serde_json::json!(input))
118        }
119        _ => {
120            let input: String = Input::new().with_prompt(title).interact_text()?;
121            Ok(serde_json::Value::String(input))
122        }
123    }
124}
125
126fn coerce_for_type(value: &serde_json::Value, param_type: &str) -> serde_json::Value {
127    match param_type {
128        "text" | "date" => match value {
129            serde_json::Value::String(_) => value.clone(),
130            other => serde_json::Value::String(other.to_string()),
131        },
132        _ => value.clone(),
133    }
134}
135
136fn build_parameter_map(
137    metadata: &QueryMetadata,
138    cli_params: &[(String, serde_json::Value)],
139    interactive: bool,
140) -> Result<Option<HashMap<String, serde_json::Value>>> {
141    if metadata.options.parameters.is_empty() {
142        return Ok(None);
143    }
144
145    let mut param_map = HashMap::new();
146
147    for (name, value) in cli_params {
148        let coerced = metadata
149            .options
150            .parameters
151            .iter()
152            .find(|p| p.name == *name)
153            .map_or_else(|| value.clone(), |p| coerce_for_type(value, &p.param_type));
154        param_map.insert(name.clone(), coerced);
155    }
156
157    let has_tty = std::io::stdin().is_terminal();
158
159    for param in &metadata.options.parameters {
160        if !param_map.contains_key(&param.name) {
161            if interactive && has_tty {
162                eprintln!("\nParameter '{}' required:", param.title);
163                let value = prompt_for_parameter(param)?;
164                param_map.insert(param.name.clone(), value);
165            } else if interactive && !has_tty {
166                if let Some(default_value) = &param.value {
167                    param_map.insert(param.name.clone(), default_value.clone());
168                } else {
169                    bail!(
170                        "No TTY available for interactive prompt. \
171                         Supply parameter explicitly: --param {}=<value>",
172                        param.name
173                    );
174                }
175            } else if let Some(default_value) = &param.value {
176                param_map.insert(param.name.clone(), default_value.clone());
177            } else {
178                bail!(
179                    "Missing required parameter: '{}' ({}). Use --param {}=value or --interactive",
180                    param.name,
181                    param.title,
182                    param.name
183                );
184            }
185        }
186    }
187
188    Ok(if param_map.is_empty() {
189        None
190    } else {
191        Some(param_map)
192    })
193}
194
195fn format_results_json(
196    result: &crate::models::QueryResult,
197    limit: Option<usize>,
198) -> Result<String> {
199    let rows = if let Some(limit) = limit {
200        result
201            .data
202            .rows
203            .iter()
204            .take(limit)
205            .cloned()
206            .collect::<Vec<_>>()
207    } else {
208        result.data.rows.clone()
209    };
210
211    serde_json::to_string_pretty(&rows).context("Failed to format results as JSON")
212}
213
214fn format_results_table(result: &crate::models::QueryResult, limit: Option<usize>) -> String {
215    use std::fmt::Write;
216
217    let mut output = String::new();
218
219    let _ = writeln!(output);
220    for col in &result.data.columns {
221        let _ = write!(output, "{:20} ", col.name);
222    }
223    let _ = writeln!(output);
224    let _ = writeln!(output, "{}", "-".repeat(result.data.columns.len() * 21));
225
226    let rows_to_show = limit
227        .unwrap_or(result.data.rows.len())
228        .min(result.data.rows.len());
229
230    for row in &result.data.rows[..rows_to_show] {
231        if let serde_json::Value::Object(obj) = row {
232            for col in &result.data.columns {
233                let value = obj
234                    .get(&col.name)
235                    .map(|v| match v {
236                        serde_json::Value::Null => "NULL".to_string(),
237                        serde_json::Value::String(s) => s.clone(),
238                        _ => v.to_string(),
239                    })
240                    .unwrap_or_default();
241
242                let truncated = if value.len() > 18 {
243                    format!("{}...", &value[..15])
244                } else {
245                    value
246                };
247
248                let _ = write!(output, "{truncated:20} ");
249            }
250            let _ = writeln!(output);
251        }
252    }
253
254    if rows_to_show < result.data.rows.len() {
255        let _ = write!(
256            output,
257            "\n... {} more rows (showing {} of {})\n",
258            result.data.rows.len() - rows_to_show,
259            rows_to_show,
260            result.data.rows.len()
261        );
262    }
263
264    let _ = write!(
265        output,
266        "\n✓ {} rows returned in {:.2}s\n",
267        result.data.rows.len(),
268        result.runtime
269    );
270
271    output
272}
273
274pub async fn execute(
275    client: &RedashClient,
276    query_id: u64,
277    param_args: Vec<String>,
278    format: OutputFormat,
279    interactive: bool,
280    timeout_secs: u64,
281    limit_rows: Option<usize>,
282) -> Result<()> {
283    let (metadata, _sql, yaml_path) = load_query_metadata_by_id(query_id)?;
284
285    eprintln!("Executing query: {} - {}", metadata.id, metadata.name);
286    eprintln!("Source: {yaml_path}\n");
287
288    let cli_params: Vec<(String, serde_json::Value)> = param_args
289        .iter()
290        .map(|arg| parse_parameter_arg(arg))
291        .collect::<Result<Vec<_>>>()?;
292
293    let parameters = build_parameter_map(&metadata, &cli_params, interactive)?;
294
295    if let Some(ref params) = parameters {
296        eprintln!("Parameters:");
297        for (name, value) in params {
298            eprintln!("  {name} = {value}");
299        }
300        eprintln!();
301    }
302
303    let result = client
304        .execute_query_with_polling(query_id, parameters, timeout_secs, 500)
305        .await?;
306
307    match format {
308        OutputFormat::Json => {
309            let json = format_results_json(&result, limit_rows)?;
310            println!("{json}");
311        }
312        OutputFormat::Table => {
313            let table = format_results_table(&result, limit_rows);
314            println!("{table}");
315        }
316    }
317
318    Ok(())
319}
320
321#[cfg(test)]
322#[allow(clippy::missing_errors_doc)]
323mod tests {
324    use super::*;
325    use crate::models::{Column, QueryResult, QueryResultData};
326
327    #[test]
328    fn test_parse_parameter_arg_string() {
329        let result = parse_parameter_arg("name=value").unwrap();
330        assert_eq!(result.0, "name");
331        assert_eq!(result.1, serde_json::Value::String("value".to_string()));
332    }
333
334    #[test]
335    fn test_parse_parameter_arg_json_array() {
336        let result = parse_parameter_arg("channels=[\"release\",\"beta\"]").unwrap();
337        assert_eq!(result.0, "channels");
338        assert_eq!(result.1, serde_json::json!(["release", "beta"]));
339    }
340
341    #[test]
342    fn test_parse_parameter_arg_number() {
343        let result = parse_parameter_arg("count=42").unwrap();
344        assert_eq!(result.0, "count");
345        assert_eq!(result.1, serde_json::json!(42));
346    }
347
348    #[test]
349    fn test_parse_parameter_arg_invalid() {
350        let result = parse_parameter_arg("invalid");
351        assert!(result.is_err());
352        assert!(
353            result
354                .unwrap_err()
355                .to_string()
356                .contains("Invalid parameter format")
357        );
358    }
359
360    #[test]
361    fn test_format_results_json() {
362        let result = QueryResult {
363            id: 1,
364            data: QueryResultData {
365                columns: vec![
366                    Column {
367                        name: "col1".to_string(),
368                        type_name: "string".to_string(),
369                        friendly_name: None,
370                    },
371                    Column {
372                        name: "col2".to_string(),
373                        type_name: "integer".to_string(),
374                        friendly_name: None,
375                    },
376                ],
377                rows: vec![
378                    serde_json::json!({"col1": "value1", "col2": 123}),
379                    serde_json::json!({"col1": "value2", "col2": 456}),
380                ],
381            },
382            runtime: 1.5,
383            retrieved_at: "2026-01-21T10:00:00".to_string(),
384        };
385
386        let json = format_results_json(&result, None).unwrap();
387        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
388
389        let rows = parsed.as_array().unwrap();
390        assert_eq!(rows.len(), 2);
391        assert_eq!(rows[0]["col1"], "value1");
392        assert_eq!(rows[0]["col2"], 123);
393    }
394
395    #[test]
396    fn test_format_results_json_with_limit() {
397        let result = QueryResult {
398            id: 1,
399            data: QueryResultData {
400                columns: vec![Column {
401                    name: "col1".to_string(),
402                    type_name: "string".to_string(),
403                    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 json = format_results_json(&result, Some(2)).unwrap();
416        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
417
418        assert_eq!(parsed.as_array().unwrap().len(), 2);
419    }
420
421    #[test]
422    fn test_format_results_table() {
423        let result = QueryResult {
424            id: 1,
425            data: QueryResultData {
426                columns: vec![
427                    Column {
428                        name: "col1".to_string(),
429                        type_name: "string".to_string(),
430                        friendly_name: None,
431                    },
432                    Column {
433                        name: "col2".to_string(),
434                        type_name: "integer".to_string(),
435                        friendly_name: None,
436                    },
437                ],
438                rows: vec![
439                    serde_json::json!({"col1": "value1", "col2": 123}),
440                    serde_json::json!({"col1": "value2", "col2": 456}),
441                ],
442            },
443            runtime: 1.5,
444            retrieved_at: "2026-01-21T10:00:00".to_string(),
445        };
446
447        let table = format_results_table(&result, None);
448
449        assert!(table.contains("col1"));
450        assert!(table.contains("col2"));
451        assert!(table.contains("value1"));
452        assert!(table.contains("value2"));
453        assert!(table.contains("2 rows returned"));
454    }
455
456    #[test]
457    fn test_format_results_table_with_limit() {
458        let result = QueryResult {
459            id: 1,
460            data: QueryResultData {
461                columns: vec![Column {
462                    name: "col1".to_string(),
463                    type_name: "string".to_string(),
464                    friendly_name: None,
465                }],
466                rows: vec![
467                    serde_json::json!({"col1": "row1"}),
468                    serde_json::json!({"col1": "row2"}),
469                    serde_json::json!({"col1": "row3"}),
470                ],
471            },
472            runtime: 1.0,
473            retrieved_at: "2026-01-21T10:00:00".to_string(),
474        };
475
476        let table = format_results_table(&result, Some(2));
477
478        assert!(table.contains("row1"));
479        assert!(table.contains("row2"));
480        assert!(table.contains("... 1 more rows"));
481        assert!(table.contains("3 rows returned"));
482    }
483
484    #[test]
485    fn test_format_results_table_truncation() {
486        let result = QueryResult {
487            id: 1,
488            data: QueryResultData {
489                columns: vec![Column {
490                    name: "col1".to_string(),
491                    type_name: "string".to_string(),
492                    friendly_name: None,
493                }],
494                rows: vec![
495                    serde_json::json!({"col1": "this_is_a_very_long_value_that_should_be_truncated"}),
496                ],
497            },
498            runtime: 1.0,
499            retrieved_at: "2026-01-21T10:00:00".to_string(),
500        };
501
502        let table = format_results_table(&result, None);
503
504        assert!(table.contains("..."));
505    }
506
507    #[test]
508    fn test_output_format_from_str() {
509        assert!(matches!(
510            "json".parse::<OutputFormat>().unwrap(),
511            OutputFormat::Json
512        ));
513        assert!(matches!(
514            "JSON".parse::<OutputFormat>().unwrap(),
515            OutputFormat::Json
516        ));
517        assert!(matches!(
518            "table".parse::<OutputFormat>().unwrap(),
519            OutputFormat::Table
520        ));
521        assert!(matches!(
522            "TABLE".parse::<OutputFormat>().unwrap(),
523            OutputFormat::Table
524        ));
525    }
526
527    #[test]
528    fn test_output_format_from_str_invalid() {
529        let result = "csv".parse::<OutputFormat>();
530        assert!(result.is_err());
531        assert!(result.unwrap_err().to_string().contains("Invalid format"));
532    }
533
534    #[test]
535    fn test_coerce_for_type_text_coerces_number_to_string() {
536        let result = coerce_for_type(&serde_json::json!(90), "text");
537        assert_eq!(result, serde_json::Value::String("90".to_string()));
538    }
539
540    #[test]
541    fn test_coerce_for_type_text_leaves_string_unchanged() {
542        let result = coerce_for_type(&serde_json::json!("90"), "text");
543        assert_eq!(result, serde_json::Value::String("90".to_string()));
544    }
545
546    #[test]
547    fn test_coerce_for_type_date_coerces_to_string() {
548        let result = coerce_for_type(&serde_json::json!(20_260_507), "date");
549        assert_eq!(result, serde_json::Value::String("20260507".to_string()));
550    }
551
552    #[test]
553    fn test_coerce_for_type_number_leaves_number_unchanged() {
554        let result = coerce_for_type(&serde_json::json!(42), "number");
555        assert_eq!(result, serde_json::json!(42));
556    }
557
558    fn make_metadata_with_param(name: &str, default: Option<serde_json::Value>) -> QueryMetadata {
559        use crate::models::{Parameter, QueryOptions};
560        QueryMetadata {
561            id: 1,
562            name: "test".to_string(),
563            description: None,
564            data_source_id: 1,
565            user_id: None,
566            schedule: None,
567            options: QueryOptions {
568                parameters: vec![Parameter {
569                    name: name.to_string(),
570                    title: name.to_string(),
571                    param_type: "text".to_string(),
572                    value: default,
573                    enum_options: None,
574                    query_id: None,
575                    multi_values_options: None,
576                }],
577            },
578            visualizations: vec![],
579            tags: None,
580        }
581    }
582
583    #[test]
584    fn test_build_parameter_map_coerces_text_param() {
585        let metadata = make_metadata_with_param("days", None);
586        let cli_params = vec![("days".to_string(), serde_json::json!(90))];
587        let result = build_parameter_map(&metadata, &cli_params, false)
588            .unwrap()
589            .unwrap();
590        assert_eq!(result["days"], serde_json::Value::String("90".to_string()));
591    }
592
593    #[test]
594    fn test_build_parameter_map_interactive_no_tty_uses_default() {
595        let metadata = make_metadata_with_param("p", Some(serde_json::json!("default_val")));
596        let result = build_parameter_map(&metadata, &[], true).unwrap();
597        let map = result.unwrap();
598        assert_eq!(map["p"], serde_json::json!("default_val"));
599    }
600
601    #[test]
602    fn test_build_parameter_map_interactive_no_tty_no_default_bails() {
603        let metadata = make_metadata_with_param("p", None);
604        let err = build_parameter_map(&metadata, &[], true).unwrap_err();
605        assert!(err.to_string().contains("--param p="));
606    }
607
608    #[test]
609    fn test_build_parameter_map_interactive_no_tty_cli_param_overrides() {
610        let metadata = make_metadata_with_param("p", None);
611        let cli_params = vec![("p".to_string(), serde_json::json!("provided"))];
612        let result = build_parameter_map(&metadata, &cli_params, true).unwrap();
613        let map = result.unwrap();
614        assert_eq!(map["p"], serde_json::json!("provided"));
615    }
616}