Skip to main content

safe_migrate/_internal/engine/
config.rs

1use anyhow::{Result, bail};
2use serde::{Deserialize, Serialize};
3use std::collections::{BTreeSet, HashMap};
4use std::fs;
5use std::path::Path;
6
7#[derive(Debug, Clone, Serialize, Deserialize, Default)]
8#[serde(deny_unknown_fields)]
9pub struct RuleConfig {
10    pub disabled: Option<bool>,
11    pub tier1_threshold_rows: Option<u64>,
12    pub tier2_threshold_rows: Option<u64>,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16#[serde(default, deny_unknown_fields)]
17pub struct Config {
18    pub tier1_threshold_rows: u64,
19    pub tier2_threshold_rows: u64,
20    pub stale_stats_days: u64,
21    pub toast_width_threshold_bytes: i32,
22    pub default_rows: u64, // Fallback for offline/unanalyzed tables
23    pub auto_sync: bool,
24    pub cache_encryption: bool,
25    pub rules: HashMap<String, RuleConfig>, // Per-rule configuration
26    pub assume_pg_version: u32,
27    pub disabled_rules: Vec<String>,
28    pub schemas: Option<Vec<String>>,
29}
30
31impl Default for Config {
32    fn default() -> Self {
33        Self {
34            tier1_threshold_rows: 100_000,
35            tier2_threshold_rows: 10_000,
36            stale_stats_days: 7,
37            toast_width_threshold_bytes: 2048,
38            default_rows: 10_000,
39            auto_sync: false,
40            cache_encryption: false,
41            assume_pg_version: 100000,
42            disabled_rules: Vec::new(),
43            rules: HashMap::new(),
44            schemas: None,
45        }
46    }
47}
48
49impl Config {
50    pub fn load_from_file(path: &Path) -> Result<Self, anyhow::Error> {
51        match fs::read_to_string(path) {
52            Ok(contents) => Self::parse_file(path, &contents),
53            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()),
54            Err(error) => Err(error.into()),
55        }
56    }
57
58    pub fn load_required_from_file(path: &Path) -> Result<Self, anyhow::Error> {
59        let contents = fs::read_to_string(path)?;
60        Self::parse_file(path, &contents)
61    }
62
63    fn parse_file(path: &Path, contents: &str) -> Result<Self, anyhow::Error> {
64        toml::from_str(contents).map_err(|error| {
65            anyhow::anyhow!("Failed to parse config at {}: {}", path.display(), error)
66        })
67    }
68
69    /// Checks if a rule is completely disabled
70    pub fn is_rule_disabled(&self, rule_id: &str) -> bool {
71        if self.disabled_rules.contains(&rule_id.to_string()) {
72            return true;
73        }
74        self.rules
75            .get(rule_id)
76            .and_then(|r| r.disabled)
77            .unwrap_or(false)
78    }
79
80    /// Gets the Tier 1 threshold for a specific rule, falling back to the global default
81    pub fn rule_tier1_threshold(&self, rule_id: &str) -> u64 {
82        self.rules
83            .get(rule_id)
84            .and_then(|r| r.tier1_threshold_rows)
85            .unwrap_or(self.tier1_threshold_rows)
86    }
87
88    /// Gets the Tier 2 threshold for a specific rule, falling back to the global default
89    pub fn rule_tier2_threshold(&self, rule_id: &str) -> u64 {
90        self.rules
91            .get(rule_id)
92            .and_then(|r| r.tier2_threshold_rows)
93            .unwrap_or(self.tier2_threshold_rows)
94    }
95
96    /// Returns the schema filter for a direct sync. An explicit CLI value wins
97    /// over the team-wide configuration default.
98    pub fn sync_schemas<'a>(
99        &'a self,
100        cli_schemas: Option<&'a [String]>,
101    ) -> Result<Option<&'a [String]>> {
102        let schemas = cli_schemas.or(self.schemas.as_deref());
103        if schemas.is_some_and(|schemas| {
104            schemas.is_empty() || schemas.iter().any(|schema| schema.trim().is_empty())
105        }) {
106            bail!("schemas must not be empty and no schema name may be blank");
107        }
108        Ok(schemas)
109    }
110
111    /// Validates configured rule IDs against the primary rule registry.
112    pub fn validate_rule_ids<'a>(
113        &self,
114        primary_rule_ids: impl IntoIterator<Item = &'a str>,
115    ) -> Result<(), anyhow::Error> {
116        let valid: BTreeSet<String> = primary_rule_ids.into_iter().map(str::to_owned).collect();
117        let unknown: BTreeSet<&str> = self
118            .rules
119            .keys()
120            .map(String::as_str)
121            .chain(self.disabled_rules.iter().map(String::as_str))
122            .filter(|rule_id| !valid.contains(*rule_id))
123            .collect();
124
125        if unknown.is_empty() {
126            return Ok(());
127        }
128
129        Err(anyhow::anyhow!(
130            "Unknown primary rule ID(s): {}. Valid primary rule IDs: {}",
131            unknown.into_iter().collect::<Vec<_>>().join(", "),
132            valid.into_iter().collect::<Vec<_>>().join(", ")
133        ))
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use std::io::Write;
141    use tempfile::NamedTempFile;
142
143    #[test]
144    fn test_granular_rule_config() {
145        let mut file = NamedTempFile::new().expect("Failed to create temp file");
146        writeln!(
147            file,
148            r#"
149            tier1_threshold_rows = 500000
150
151            [rules.blocking-constraint]
152            tier1_threshold_rows = 5000
153
154            [rules.missing-idempotency]
155            disabled = true
156        "#
157        )
158        .expect("Failed to write temp config");
159
160        let config = Config::load_from_file(file.path()).expect("Failed to load valid config");
161
162        assert_eq!(config.tier1_threshold_rows, 500_000);
163
164        assert_eq!(config.rule_tier1_threshold("blocking-constraint"), 5000);
165        assert_eq!(config.rule_tier1_threshold("unspecified-rule"), 500_000);
166        assert!(!config.auto_sync);
167        assert!(!config.cache_encryption);
168
169        assert!(config.is_rule_disabled("missing-idempotency"));
170        assert!(!config.is_rule_disabled("blocking-constraint"));
171    }
172
173    #[test]
174    fn test_direct_sync_prefers_cli_schema_filter_over_configured_default() {
175        let config = Config {
176            schemas: Some(vec!["public".to_string()]),
177            ..Config::default()
178        };
179        let cli_schemas = vec!["auth".to_string()];
180
181        assert_eq!(
182            config.sync_schemas(None).unwrap(),
183            Some(["public".to_string()].as_slice())
184        );
185        assert_eq!(
186            config.sync_schemas(Some(&cli_schemas)).unwrap(),
187            Some(["auth".to_string()].as_slice())
188        );
189    }
190
191    #[test]
192    fn test_direct_sync_rejects_empty_schema_scope() {
193        let config = Config::default();
194        assert!(config.sync_schemas(Some(&[])).is_err());
195        assert!(config.sync_schemas(Some(&["".to_string()])).is_err());
196    }
197
198    #[test]
199    fn missing_optional_config_uses_defaults_but_required_config_fails() {
200        let directory = tempfile::tempdir().unwrap();
201        let missing = directory.path().join("missing.toml");
202
203        assert_eq!(
204            Config::load_from_file(&missing)
205                .unwrap()
206                .tier1_threshold_rows,
207            Config::default().tier1_threshold_rows
208        );
209        assert!(Config::load_required_from_file(&missing).is_err());
210    }
211
212    #[test]
213    fn rule_id_validation_rejects_unknown_rule_keys_and_disabled_ids() {
214        let mut config = Config::default();
215        config
216            .rules
217            .insert("typo-rule".to_string(), RuleConfig::default());
218        config.disabled_rules = vec!["known-rule".to_string(), "other-typo".to_string()];
219
220        let error = config
221            .validate_rule_ids(["known-rule"])
222            .expect_err("unknown rule IDs must fail validation")
223            .to_string();
224
225        assert!(error.contains("other-typo, typo-rule"));
226        assert!(error.contains("Valid primary rule IDs: known-rule"));
227    }
228
229    #[test]
230    fn rule_id_validation_accepts_known_rule_keys_and_disabled_ids() {
231        let mut config = Config::default();
232        config
233            .rules
234            .insert("known-rule".to_string(), RuleConfig::default());
235        config.disabled_rules = vec!["known-rule".to_string()];
236
237        config
238            .validate_rule_ids(["known-rule"])
239            .expect("known rule IDs must pass validation");
240    }
241
242    #[test]
243    fn config_rejects_unknown_top_level_setting() {
244        let error = toml::from_str::<Config>("auto_syn = true")
245            .expect_err("unknown top-level settings must fail")
246            .to_string();
247
248        assert!(error.contains("unknown field `auto_syn`"));
249        assert!(error.contains("auto_sync"));
250    }
251
252    #[test]
253    fn config_rejects_unknown_per_rule_setting() {
254        let error =
255            toml::from_str::<Config>("[rules.blocking-constraint]\ntier1_threshold_row = 1")
256                .expect_err("unknown per-rule settings must fail")
257                .to_string();
258
259        assert!(error.contains("unknown field `tier1_threshold_row`"));
260        assert!(error.contains("tier1_threshold_rows"));
261    }
262}