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 build_parameter_map(
126    metadata: &QueryMetadata,
127    cli_params: &[(String, serde_json::Value)],
128    interactive: bool,
129) -> Result<Option<HashMap<String, serde_json::Value>>> {
130    if metadata.options.parameters.is_empty() {
131        return Ok(None);
132    }
133
134    let mut param_map = HashMap::new();
135
136    for (name, value) in cli_params {
137        param_map.insert(name.clone(), value.clone());
138    }
139
140    for param in &metadata.options.parameters {
141        if !param_map.contains_key(&param.name) {
142            if interactive {
143                eprintln!("\nParameter '{}' required:", param.title);
144                let value = prompt_for_parameter(param)?;
145                param_map.insert(param.name.clone(), value);
146            } else if let Some(default_value) = &param.value {
147                param_map.insert(param.name.clone(), default_value.clone());
148            } else {
149                bail!(
150                    "Missing required parameter: '{}' ({}). Use --param {}=value or --interactive",
151                    param.name,
152                    param.title,
153                    param.name
154                );
155            }
156        }
157    }
158
159    Ok(if param_map.is_empty() {
160        None
161    } else {
162        Some(param_map)
163    })
164}
165
166fn format_results_json(
167    result: &crate::models::QueryResult,
168    limit: Option<usize>,
169) -> Result<String> {
170    let rows = if let Some(limit) = limit {
171        result
172            .data
173            .rows
174            .iter()
175            .take(limit)
176            .cloned()
177            .collect::<Vec<_>>()
178    } else {
179        result.data.rows.clone()
180    };
181
182    serde_json::to_string_pretty(&rows).context("Failed to format results as JSON")
183}
184
185fn format_results_table(result: &crate::models::QueryResult, limit: Option<usize>) -> String {
186    use std::fmt::Write;
187
188    let mut output = String::new();
189
190    let _ = writeln!(output);
191    for col in &result.data.columns {
192        let _ = write!(output, "{:20} ", col.name);
193    }
194    let _ = writeln!(output);
195    let _ = writeln!(output, "{}", "-".repeat(result.data.columns.len() * 21));
196
197    let rows_to_show = limit
198        .unwrap_or(result.data.rows.len())
199        .min(result.data.rows.len());
200
201    for row in &result.data.rows[..rows_to_show] {
202        if let serde_json::Value::Object(obj) = row {
203            for col in &result.data.columns {
204                let value = obj
205                    .get(&col.name)
206                    .map(|v| match v {
207                        serde_json::Value::Null => "NULL".to_string(),
208                        serde_json::Value::String(s) => s.clone(),
209                        _ => v.to_string(),
210                    })
211                    .unwrap_or_default();
212
213                let truncated = if value.len() > 18 {
214                    format!("{}...", &value[..15])
215                } else {
216                    value
217                };
218
219                let _ = write!(output, "{truncated:20} ");
220            }
221            let _ = writeln!(output);
222        }
223    }
224
225    if rows_to_show < result.data.rows.len() {
226        let _ = write!(
227            output,
228            "\n... {} more rows (showing {} of {})\n",
229            result.data.rows.len() - rows_to_show,
230            rows_to_show,
231            result.data.rows.len()
232        );
233    }
234
235    let _ = write!(
236        output,
237        "\n✓ {} rows returned in {:.2}s\n",
238        result.data.rows.len(),
239        result.runtime
240    );
241
242    output
243}
244
245pub async fn execute(
246    client: &RedashClient,
247    query_id: u64,
248    param_args: Vec<String>,
249    format: OutputFormat,
250    interactive: bool,
251    timeout_secs: u64,
252    limit_rows: Option<usize>,
253) -> Result<()> {
254    let (metadata, _sql, yaml_path) = load_query_metadata_by_id(query_id)?;
255
256    eprintln!("Executing query: {} - {}", metadata.id, metadata.name);
257    eprintln!("Source: {yaml_path}\n");
258
259    let cli_params: Vec<(String, serde_json::Value)> = param_args
260        .iter()
261        .map(|arg| parse_parameter_arg(arg))
262        .collect::<Result<Vec<_>>>()?;
263
264    let parameters = build_parameter_map(&metadata, &cli_params, interactive)?;
265
266    if let Some(ref params) = parameters {
267        eprintln!("Parameters:");
268        for (name, value) in params {
269            eprintln!("  {name} = {value}");
270        }
271        eprintln!();
272    }
273
274    let result = client
275        .execute_query_with_polling(query_id, parameters, timeout_secs, 500)
276        .await?;
277
278    match format {
279        OutputFormat::Json => {
280            let json = format_results_json(&result, limit_rows)?;
281            println!("{json}");
282        }
283        OutputFormat::Table => {
284            let table = format_results_table(&result, limit_rows);
285            println!("{table}");
286        }
287    }
288
289    Ok(())
290}
291
292#[cfg(test)]
293#[allow(clippy::missing_errors_doc)]
294mod tests {
295    use super::*;
296    use crate::models::{Column, QueryResult, QueryResultData};
297
298    #[test]
299    fn test_parse_parameter_arg_string() {
300        let result = parse_parameter_arg("name=value").unwrap();
301        assert_eq!(result.0, "name");
302        assert_eq!(result.1, serde_json::Value::String("value".to_string()));
303    }
304
305    #[test]
306    fn test_parse_parameter_arg_json_array() {
307        let result = parse_parameter_arg("channels=[\"release\",\"beta\"]").unwrap();
308        assert_eq!(result.0, "channels");
309        assert_eq!(result.1, serde_json::json!(["release", "beta"]));
310    }
311
312    #[test]
313    fn test_parse_parameter_arg_number() {
314        let result = parse_parameter_arg("count=42").unwrap();
315        assert_eq!(result.0, "count");
316        assert_eq!(result.1, serde_json::json!(42));
317    }
318
319    #[test]
320    fn test_parse_parameter_arg_invalid() {
321        let result = parse_parameter_arg("invalid");
322        assert!(result.is_err());
323        assert!(
324            result
325                .unwrap_err()
326                .to_string()
327                .contains("Invalid parameter format")
328        );
329    }
330
331    #[test]
332    fn test_format_results_json() {
333        let result = QueryResult {
334            id: 1,
335            data: QueryResultData {
336                columns: vec![
337                    Column {
338                        name: "col1".to_string(),
339                        type_name: "string".to_string(),
340                        friendly_name: None,
341                    },
342                    Column {
343                        name: "col2".to_string(),
344                        type_name: "integer".to_string(),
345                        friendly_name: None,
346                    },
347                ],
348                rows: vec![
349                    serde_json::json!({"col1": "value1", "col2": 123}),
350                    serde_json::json!({"col1": "value2", "col2": 456}),
351                ],
352            },
353            runtime: 1.5,
354            retrieved_at: "2026-01-21T10:00:00".to_string(),
355        };
356
357        let json = format_results_json(&result, None).unwrap();
358        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
359
360        let rows = parsed.as_array().unwrap();
361        assert_eq!(rows.len(), 2);
362        assert_eq!(rows[0]["col1"], "value1");
363        assert_eq!(rows[0]["col2"], 123);
364    }
365
366    #[test]
367    fn test_format_results_json_with_limit() {
368        let result = QueryResult {
369            id: 1,
370            data: QueryResultData {
371                columns: vec![Column {
372                    name: "col1".to_string(),
373                    type_name: "string".to_string(),
374                    friendly_name: None,
375                }],
376                rows: vec![
377                    serde_json::json!({"col1": "row1"}),
378                    serde_json::json!({"col1": "row2"}),
379                    serde_json::json!({"col1": "row3"}),
380                ],
381            },
382            runtime: 1.0,
383            retrieved_at: "2026-01-21T10:00:00".to_string(),
384        };
385
386        let json = format_results_json(&result, Some(2)).unwrap();
387        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
388
389        assert_eq!(parsed.as_array().unwrap().len(), 2);
390    }
391
392    #[test]
393    fn test_format_results_table() {
394        let result = QueryResult {
395            id: 1,
396            data: QueryResultData {
397                columns: vec![
398                    Column {
399                        name: "col1".to_string(),
400                        type_name: "string".to_string(),
401                        friendly_name: None,
402                    },
403                    Column {
404                        name: "col2".to_string(),
405                        type_name: "integer".to_string(),
406                        friendly_name: None,
407                    },
408                ],
409                rows: vec![
410                    serde_json::json!({"col1": "value1", "col2": 123}),
411                    serde_json::json!({"col1": "value2", "col2": 456}),
412                ],
413            },
414            runtime: 1.5,
415            retrieved_at: "2026-01-21T10:00:00".to_string(),
416        };
417
418        let table = format_results_table(&result, None);
419
420        assert!(table.contains("col1"));
421        assert!(table.contains("col2"));
422        assert!(table.contains("value1"));
423        assert!(table.contains("value2"));
424        assert!(table.contains("2 rows returned"));
425    }
426
427    #[test]
428    fn test_format_results_table_with_limit() {
429        let result = QueryResult {
430            id: 1,
431            data: QueryResultData {
432                columns: vec![Column {
433                    name: "col1".to_string(),
434                    type_name: "string".to_string(),
435                    friendly_name: None,
436                }],
437                rows: vec![
438                    serde_json::json!({"col1": "row1"}),
439                    serde_json::json!({"col1": "row2"}),
440                    serde_json::json!({"col1": "row3"}),
441                ],
442            },
443            runtime: 1.0,
444            retrieved_at: "2026-01-21T10:00:00".to_string(),
445        };
446
447        let table = format_results_table(&result, Some(2));
448
449        assert!(table.contains("row1"));
450        assert!(table.contains("row2"));
451        assert!(table.contains("... 1 more rows"));
452        assert!(table.contains("3 rows returned"));
453    }
454
455    #[test]
456    fn test_format_results_table_truncation() {
457        let result = QueryResult {
458            id: 1,
459            data: QueryResultData {
460                columns: vec![Column {
461                    name: "col1".to_string(),
462                    type_name: "string".to_string(),
463                    friendly_name: None,
464                }],
465                rows: vec![
466                    serde_json::json!({"col1": "this_is_a_very_long_value_that_should_be_truncated"}),
467                ],
468            },
469            runtime: 1.0,
470            retrieved_at: "2026-01-21T10:00:00".to_string(),
471        };
472
473        let table = format_results_table(&result, None);
474
475        assert!(table.contains("..."));
476    }
477
478    #[test]
479    fn test_output_format_from_str() {
480        assert!(matches!(
481            "json".parse::<OutputFormat>().unwrap(),
482            OutputFormat::Json
483        ));
484        assert!(matches!(
485            "JSON".parse::<OutputFormat>().unwrap(),
486            OutputFormat::Json
487        ));
488        assert!(matches!(
489            "table".parse::<OutputFormat>().unwrap(),
490            OutputFormat::Table
491        ));
492        assert!(matches!(
493            "TABLE".parse::<OutputFormat>().unwrap(),
494            OutputFormat::Table
495        ));
496    }
497
498    #[test]
499    fn test_output_format_from_str_invalid() {
500        let result = "csv".parse::<OutputFormat>();
501        assert!(result.is_err());
502        assert!(result.unwrap_err().to_string().contains("Invalid format"));
503    }
504}