Skip to main content

safe_migrate/engine/
config.rs

1// FILE: src/engine/config.rs
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fs;
5use std::path::Path;
6
7#[derive(Debug, Clone, Serialize, Deserialize, Default)]
8pub struct RuleConfig {
9    pub disabled: Option<bool>,
10    pub tier1_threshold_rows: Option<u64>,
11    pub tier2_threshold_rows: Option<u64>,
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15#[serde(default)] // Allows missing keys in TOML to fall back to Default::default()
16pub struct Config {
17    pub tier1_threshold_rows: u64,
18    pub tier2_threshold_rows: u64,
19    pub stale_stats_days: u64,
20    pub toast_width_threshold_bytes: i32,
21    pub default_rows: u64,           // Fallback for offline/unanalyzed tables
22    pub assume_pg_version: u32,      // Fallback for offline PG version
23    pub disabled_rules: Vec<String>, // Legacy global array
24    pub rules: HashMap<String, RuleConfig>, // NEW: Granular per-rule configs
25}
26
27impl Default for Config {
28    fn default() -> Self {
29        Self {
30            tier1_threshold_rows: 100_000,
31            tier2_threshold_rows: 10_000,
32            stale_stats_days: 7,
33            toast_width_threshold_bytes: 2048,
34            default_rows: 10_000,
35            assume_pg_version: 100000,
36            disabled_rules: Vec::new(),
37            rules: HashMap::new(),
38        }
39    }
40}
41
42impl Config {
43    pub fn load_from_file(path: &Path) -> Result<Self, anyhow::Error> {
44        if path.exists() {
45            let contents = fs::read_to_string(path)?;
46            match toml::from_str(&contents) {
47                Ok(config) => return Ok(config),
48                Err(e) => {
49                    return Err(anyhow::anyhow!(
50                        "Failed to parse config at {}: {}",
51                        path.display(),
52                        e
53                    ));
54                }
55            }
56        }
57        Ok(Self::default())
58    }
59
60    /// Checks if a rule is completely disabled
61    pub fn is_rule_disabled(&self, rule_id: &str) -> bool {
62        if self.disabled_rules.contains(&rule_id.to_string()) {
63            return true;
64        }
65        self.rules
66            .get(rule_id)
67            .and_then(|r| r.disabled)
68            .unwrap_or(false)
69    }
70
71    /// Gets the Tier 1 threshold for a specific rule, falling back to the global default
72    pub fn rule_tier1_threshold(&self, rule_id: &str) -> u64 {
73        self.rules
74            .get(rule_id)
75            .and_then(|r| r.tier1_threshold_rows)
76            .unwrap_or(self.tier1_threshold_rows)
77    }
78
79    /// Gets the Tier 2 threshold for a specific rule, falling back to the global default
80    pub fn rule_tier2_threshold(&self, rule_id: &str) -> u64 {
81        self.rules
82            .get(rule_id)
83            .and_then(|r| r.tier2_threshold_rows)
84            .unwrap_or(self.tier2_threshold_rows)
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use std::io::Write;
92    use tempfile::NamedTempFile;
93
94    #[test]
95    fn test_granular_rule_config() {
96        let mut file = NamedTempFile::new().expect("Failed to create temp file");
97        writeln!(
98            file,
99            r#"
100            tier1_threshold_rows = 500000
101
102            [rules.blocking-constraint]
103            tier1_threshold_rows = 5000
104
105            [rules.missing-idempotency]
106            disabled = true
107        "#
108        )
109        .expect("Failed to write temp config");
110
111        let config = Config::load_from_file(file.path()).expect("Failed to load valid config");
112
113        // Assert Global Overrides
114        assert_eq!(config.tier1_threshold_rows, 500_000);
115
116        // Assert Granular Fallbacks
117        assert_eq!(config.rule_tier1_threshold("blocking-constraint"), 5000);
118        assert_eq!(config.rule_tier1_threshold("unspecified-rule"), 500_000);
119
120        // Assert Rule Disabling
121        assert!(config.is_rule_disabled("missing-idempotency"));
122        assert!(!config.is_rule_disabled("blocking-constraint"));
123    }
124}