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

safe_migrate/engine/
config.rs
