Skip to main content

stmo_cli/commands/
schedule.rs

1#![allow(clippy::missing_errors_doc)]
2
3use anyhow::{Context, Result, bail};
4use std::fs;
5use std::path::{Path, PathBuf};
6
7use crate::models::{QueryMetadata, Schedule};
8
9fn find_yaml_path_in(dir: &Path, query_id: u64) -> Result<Option<PathBuf>> {
10    if !dir.exists() {
11        return Ok(None);
12    }
13    for entry in fs::read_dir(dir).context("Failed to read queries directory")? {
14        let entry = entry.context("Failed to read directory entry")?;
15        let path = entry.path();
16        if path.extension().is_some_and(|ext| ext == "yaml")
17            && let Some(filename) = path.file_name().and_then(|f| f.to_str())
18            && let Some(id_str) = filename.split('-').next()
19            && let Ok(id) = id_str.parse::<u64>()
20            && id == query_id
21        {
22            return Ok(Some(path));
23        }
24    }
25    Ok(None)
26}
27
28fn update_yaml_schedule(yaml_path: &Path, schedule: Option<Schedule>) -> Result<String> {
29    let content =
30        fs::read_to_string(yaml_path).context(format!("Failed to read {}", yaml_path.display()))?;
31    let mut metadata: QueryMetadata = serde_yaml::from_str(&content)
32        .context(format!("Failed to parse {}", yaml_path.display()))?;
33
34    let name = metadata.name.clone();
35    metadata.schedule = schedule;
36
37    let yaml_content =
38        serde_yaml::to_string(&metadata).context("Failed to serialize query metadata")?;
39    fs::write(yaml_path, yaml_content)
40        .context(format!("Failed to write {}", yaml_path.display()))?;
41
42    Ok(name)
43}
44
45pub fn schedule(
46    query_ids: &[u64],
47    interval: Option<u64>,
48    time: Option<&str>,
49    day_of_week: Option<&str>,
50    clear: bool,
51) -> Result<()> {
52    if !clear && interval.is_none() {
53        bail!(
54            "Either --interval or --clear must be specified.\n\nExamples:\n  stmo-cli schedule 123 456 --interval 86400 --time 07:15\n  stmo-cli schedule 123 --clear"
55        );
56    }
57
58    let new_schedule = if clear {
59        None
60    } else {
61        Some(Schedule {
62            interval,
63            time: time.map(str::to_owned),
64            day_of_week: day_of_week.map(str::to_owned),
65            until: None,
66        })
67    };
68
69    let queries_dir = Path::new("queries");
70    let mut errors: Vec<(u64, anyhow::Error)> = Vec::new();
71    let mut updated_count = 0;
72
73    for &query_id in query_ids {
74        match find_yaml_path_in(queries_dir, query_id) {
75            Err(e) => {
76                eprintln!("  ✗ Failed to search for query {query_id}: {e}");
77                errors.push((query_id, e));
78            }
79            Ok(None) => {
80                let e = anyhow::anyhow!(
81                    "No local file found for query {query_id}. Run 'stmo-cli fetch {query_id}' first."
82                );
83                eprintln!("  ✗ {e}");
84                errors.push((query_id, e));
85            }
86            Ok(Some(yaml_path)) => match update_yaml_schedule(&yaml_path, new_schedule.clone()) {
87                Ok(name) => {
88                    let action = if clear {
89                        "Cleared schedule from"
90                    } else {
91                        "Set schedule on"
92                    };
93                    println!("  ✓ {action} query {query_id} - {name}");
94                    updated_count += 1;
95                }
96                Err(e) => {
97                    eprintln!("  ✗ Failed to update query {query_id}: {e}");
98                    errors.push((query_id, e));
99                }
100            },
101        }
102    }
103
104    println!("\n✓ Updated {updated_count}/{} queries", query_ids.len());
105    println!("Run 'stmo-cli deploy' to push the schedule changes to Redash.");
106
107    if !errors.is_empty() {
108        bail!("Failed to update {} queries", errors.len());
109    }
110    Ok(())
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use std::fs;
117    use tempfile::TempDir;
118
119    const SAMPLE_YAML: &str = "id: 121795\nname: test query\ndescription: null\ndata_source_id: 63\nschedule: null\noptions:\n  parameters: []\nvisualizations: []\ntags: null\n";
120
121    fn create_test_query(dir: &Path, id: u64, yaml_content: &str) {
122        let yaml_path = dir.join(format!("{id}-test-query.yaml"));
123        fs::write(yaml_path, yaml_content).unwrap();
124        let sql_path = dir.join(format!("{id}-test-query.sql"));
125        fs::write(sql_path, "SELECT 1").unwrap();
126    }
127
128    #[test]
129    fn set_schedule_writes_yaml() {
130        let temp_dir = TempDir::new().unwrap();
131        let queries_dir = temp_dir.path();
132        create_test_query(queries_dir, 121_795, SAMPLE_YAML);
133
134        let yaml_path = queries_dir.join("121795-test-query.yaml");
135        update_yaml_schedule(
136            &yaml_path,
137            Some(Schedule {
138                interval: Some(86400),
139                time: Some("07:15".to_owned()),
140                day_of_week: None,
141                until: None,
142            }),
143        )
144        .unwrap();
145
146        let content = fs::read_to_string(&yaml_path).unwrap();
147        let metadata: QueryMetadata = serde_yaml::from_str(&content).unwrap();
148        let s = metadata.schedule.unwrap();
149        assert_eq!(s.interval, Some(86_400));
150        assert_eq!(s.time, Some("07:15".to_owned()));
151        assert_eq!(s.day_of_week, None);
152        assert_eq!(s.until, None);
153    }
154
155    #[test]
156    fn clear_schedule_writes_null() {
157        let temp_dir = TempDir::new().unwrap();
158        let queries_dir = temp_dir.path();
159        let yaml_with_schedule = "id: 121795\nname: test query\ndescription: null\ndata_source_id: 63\nschedule:\n  interval: 86400\n  time: '07:15'\n  day_of_week: null\n  until: null\noptions:\n  parameters: []\nvisualizations: []\ntags: null\n";
160        create_test_query(queries_dir, 121_795, yaml_with_schedule);
161
162        let yaml_path = queries_dir.join("121795-test-query.yaml");
163        update_yaml_schedule(&yaml_path, None).unwrap();
164
165        let content = fs::read_to_string(&yaml_path).unwrap();
166        let metadata: QueryMetadata = serde_yaml::from_str(&content).unwrap();
167        assert!(metadata.schedule.is_none());
168    }
169
170    #[test]
171    fn find_yaml_finds_by_id_prefix() {
172        let temp_dir = TempDir::new().unwrap();
173        let queries_dir = temp_dir.path();
174        create_test_query(queries_dir, 121_795, SAMPLE_YAML);
175
176        let found = find_yaml_path_in(queries_dir, 121_795).unwrap();
177        assert!(found.is_some());
178
179        let not_found = find_yaml_path_in(queries_dir, 99_999).unwrap();
180        assert!(not_found.is_none());
181    }
182
183    #[test]
184    fn find_yaml_returns_none_for_missing_dir() {
185        let path = find_yaml_path_in(Path::new("/nonexistent/path"), 123).unwrap();
186        assert!(path.is_none());
187    }
188}