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) -> Self {
44        if path.exists()
45            && let Ok(contents) = fs::read_to_string(path)
46        {
47            if let Ok(config) = toml::from_str(&contents) {
48                return config;
49            } else {
50                eprintln!(
51                    "[WARN] Failed to parse config at {}. Using defaults.",
52                    path.display()
53                );
54            }
55        }
56        Self::default()
57    }
58
59    /// Checks if a rule is completely disabled
60    pub fn is_rule_disabled(&self, rule_id: &str) -> bool {
61        if self.disabled_rules.contains(&rule_id.to_string()) {
62            return true;
63        }
64        self.rules
65            .get(rule_id)
66            .and_then(|r| r.disabled)
67            .unwrap_or(false)
68    }
69
70    /// Gets the Tier 1 threshold for a specific rule, falling back to the global default
71    pub fn rule_tier1_threshold(&self, rule_id: &str) -> u64 {
72        self.rules
73            .get(rule_id)
74            .and_then(|r| r.tier1_threshold_rows)
75            .unwrap_or(self.tier1_threshold_rows)
76    }
77
78    /// Gets the Tier 2 threshold for a specific rule, falling back to the global default
79    pub fn rule_tier2_threshold(&self, rule_id: &str) -> u64 {
80        self.rules
81            .get(rule_id)
82            .and_then(|r| r.tier2_threshold_rows)
83            .unwrap_or(self.tier2_threshold_rows)
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use std::io::Write;
91    use tempfile::NamedTempFile;
92
93    #[test]
94    fn test_granular_rule_config() {
95        let mut file = NamedTempFile::new().expect("Failed to create temp file");
96        writeln!(
97            file,
98            r#"
99            tier1_threshold_rows = 500000
100
101            [rules.blocking-constraint]
102            tier1_threshold_rows = 5000
103
104            [rules.missing-idempotency]
105            disabled = true
106        "#
107        )
108        .expect("Failed to write temp config");
109
110        let config = Config::load_from_file(file.path());
111
112        // Assert Global Overrides
113        assert_eq!(config.tier1_threshold_rows, 500_000);
114
115        // Assert Granular Fallbacks
116        assert_eq!(config.rule_tier1_threshold("blocking-constraint"), 5000);
117        assert_eq!(config.rule_tier1_threshold("unspecified-rule"), 500_000);
118
119        // Assert Rule Disabling
120        assert!(config.is_rule_disabled("missing-idempotency"));
121        assert!(!config.is_rule_disabled("blocking-constraint"));
122    }
123}