safe_migrate/engine/
config.rs1use 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)] pub 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, pub assume_pg_version: u32, pub disabled_rules: Vec<String>, pub rules: HashMap<String, RuleConfig>, }
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 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 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 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_eq!(config.tier1_threshold_rows, 500_000);
114
115 assert_eq!(config.rule_tier1_threshold("blocking-constraint"), 5000);
117 assert_eq!(config.rule_tier1_threshold("unspecified-rule"), 500_000);
118
119 assert!(config.is_rule_disabled("missing-idempotency"));
121 assert!(!config.is_rule_disabled("blocking-constraint"));
122 }
123}