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) -> 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 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 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 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_eq!(config.tier1_threshold_rows, 500_000);
115
116 assert_eq!(config.rule_tier1_threshold("blocking-constraint"), 5000);
118 assert_eq!(config.rule_tier1_threshold("unspecified-rule"), 500_000);
119
120 assert!(config.is_rule_disabled("missing-idempotency"));
122 assert!(!config.is_rule_disabled("blocking-constraint"));
123 }
124}