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