1use crate::model::LockTier;
2use anyhow::{Context, Result, anyhow};
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::fs;
6use std::path::Path;
7
8#[derive(Deserialize, Serialize, Debug, Clone)]
9pub struct RuleConfig {
10 pub tier: LockTier,
11 pub threshold: Option<u64>,
12}
13
14#[derive(Deserialize, Debug)]
15struct PartialRuleConfig {
16 tier: Option<LockTier>,
17 threshold: Option<u64>,
18}
19
20#[derive(Deserialize, Debug)]
21struct PartialConfig {
22 default_threshold: Option<u64>,
23 rules: Option<HashMap<String, PartialRuleConfig>>,
24}
25
26#[derive(Deserialize, Serialize, Debug, Clone)]
27pub struct Config {
28 pub default_threshold: u64,
29 pub rules: HashMap<String, RuleConfig>,
30}
31
32impl Config {
33 pub fn default_config() -> Self {
34 let mut rules = HashMap::new();
35
36 let tier1_rules = [
37 "adding-field-with-default",
38 "changing-column-type",
39 "adding-not-nullable-field",
40 "adding-serial-primary-key-field",
41 "adding-required-field",
42 "renaming-column",
43 "renaming-table",
44 "disallowed-unique-constraint",
45 "ban-drop-table",
46 "ban-drop-column",
47 "executing-unclassified-statement", ];
49
50 for r in tier1_rules {
51 rules.insert(
52 r.to_string(),
53 RuleConfig {
54 tier: LockTier::Tier1,
55 threshold: None,
56 },
57 );
58 }
59
60 let tier2_rules = [
61 "require-concurrent-index-creation",
62 "require-concurrent-index-deletion",
63 "adding-foreign-key-constraint",
64 "constraint-missing-not-valid",
65 ];
66
67 for r in tier2_rules {
68 rules.insert(
69 r.to_string(),
70 RuleConfig {
71 tier: LockTier::Tier2,
72 threshold: None,
73 },
74 );
75 }
76
77 Config {
78 default_threshold: 100_000,
79 rules,
80 }
81 }
82
83 pub fn load(path: &str) -> Result<Self> {
84 let mut config = Self::default_config();
85 let path_obj = Path::new(path);
86
87 if path_obj.exists() {
88 let contents = fs::read_to_string(path_obj)
89 .with_context(|| format!("Failed to read config file at {}", path))?;
90
91 let partial: PartialConfig = toml::from_str(&contents)
92 .context("Malformed safe-migrate.toml. Ensure it is valid TOML.")?;
93
94 if let Some(dt) = partial.default_threshold {
95 config.default_threshold = dt;
96 }
97 if let Some(user_rules) = partial.rules {
98 for (k, v) in user_rules {
99 let existing = config.rules.get_mut(&k).ok_or_else(|| {
100 anyhow!(
101 "Unknown rule '{}' found in config. Please check for typos.",
102 k
103 )
104 })?;
105
106 if let Some(t) = v.tier {
107 existing.tier = t;
108 }
109 if let Some(th) = v.threshold {
110 existing.threshold = Some(th);
111 }
112 }
113 }
114 }
115 Ok(config)
116 }
117}
118
119pub fn get_recipe(rule: &str) -> &'static str {
120 match rule {
121 "require-concurrent-index-creation" => {
122 "CREATE INDEX blocks all writes (INSERT/UPDATE/DELETE) to the table.\n\
123 Safe Migration:\n\
124 1. Use 'CREATE INDEX CONCURRENTLY'.\n\
125 2. Remove any surrounding BEGIN/COMMIT blocks, as CONCURRENTLY cannot run inside a transaction."
126 }
127 "require-concurrent-index-deletion" => {
128 "DROP INDEX requires an ACCESS EXCLUSIVE lock, blocking all reads and writes.\n\
129 Safe Migration:\n\
130 1. Use 'DROP INDEX CONCURRENTLY'.\n\
131 2. Remove any surrounding BEGIN/COMMIT blocks."
132 }
133 "adding-field-with-default" => {
134 "Since PostgreSQL 11, adding a column with a constant default is instant. However, volatile defaults (e.g., gen_random_uuid()) will rewrite the entire table.\n\
135 Safe Migration (Expand/Contract):\n\
136 1. Add the column as nullable.\n\
137 2. Backfill existing rows in small batches to avoid long lock durations.\n\
138 3. Add a check constraint using NOT VALID to avoid a table scan.\n\
139 4. In a separate migration, validate the constraint."
140 }
141 "adding-not-nullable-field" | "adding-required-field" => {
142 "Adding a NOT NULL column without a default will fail or rewrite the table.\n\
143 Safe Migration:\n\
144 1. Add the column as nullable.\n\
145 2. Backfill existing rows in small batches.\n\
146 3. Add a check constraint using NOT VALID.\n\
147 4. Run VALIDATE CONSTRAINT in a separate migration."
148 }
149 "changing-column-type" => {
150 "Changing a column type rewrites the entire table on disk, holding an ACCESS EXCLUSIVE lock.\n\
151 Safe Migration:\n\
152 1. Create a new column with the desired type.\n\
153 2. Add a trigger to keep the new column in sync with the old one during writes.\n\
154 3. Backfill existing data in batches.\n\
155 4. Deploy app changes to read/write to the new column.\n\
156 5. Drop the old column and rename the new one."
157 }
158 "renaming-column" => {
159 "Renaming a column is instant but breaks concurrent queries using the old name.\n\
160 Safe Migration:\n\
161 1. Add a new column with the new name.\n\
162 2. Add a trigger to sync writes between the old and new columns.\n\
163 3. Backfill data in batches.\n\
164 4. Deploy your application to use the new column.\n\
165 5. Drop the old column."
166 }
167 "renaming-table" => {
168 "Renaming a table is instant but breaks concurrent application queries using the old name.\n\
169 Safe Migration:\n\
170 1. Rename the table.\n\
171 2. Immediately create a VIEW with the old table name pointing to the new table.\n\
172 3. Update your application to use the new name.\n\
173 4. Drop the view when ready."
174 }
175 "adding-foreign-key-constraint" | "constraint-missing-not-valid" => {
176 "Adding a standard constraint acquires a ShareRowExclusiveLock and triggers a full table scan, blocking concurrent writes.\n\
177 Safe Migration:\n\
178 1. Add the constraint using the NOT VALID option. This skips the initial integrity check and commits immediately.\n\
179 2. In a separate migration, run VALIDATE CONSTRAINT. This checks pre-existing rows without locking out concurrent updates."
180 }
181 "disallowed-unique-constraint" => {
182 "Adding a UNIQUE constraint builds an index using an ACCESS EXCLUSIVE lock, blocking reads and writes.\n\
183 Safe Migration:\n\
184 1. Create a unique index concurrently: CREATE UNIQUE INDEX CONCURRENTLY idx_name ON table_name (column_name);\n\
185 2. Add the unique constraint using the existing index: ALTER TABLE table_name ADD CONSTRAINT const_name UNIQUE USING INDEX idx_name;"
186 }
187 "adding-serial-primary-key-field" => {
188 "Adding a SERIAL PRIMARY KEY rewrites the entire table to generate sequence values and build the index, locking it exclusively.\n\
189 Safe Migration:\n\
190 1. Add a nullable integer column.\n\
191 2. Create a sequence and set the column default to the sequence.\n\
192 3. Backfill existing rows in batches.\n\
193 4. Create a UNIQUE INDEX CONCURRENTLY.\n\
194 5. Add the PRIMARY KEY constraint USING INDEX."
195 }
196 "ban-drop-table" | "ban-drop-column" => {
197 "Dropping data structures requires an ACCESS EXCLUSIVE lock.\n\
198 Safe Migration:\n\
199 1. Ensure application code completely ignores this object.\n\
200 2. Always precede this command with 'SET lock_timeout = '2s';' so busy databases fail gracefully instead of taking down production.\n\
201 3. Let the pipeline retry later."
202 }
203 "benign-statement" => {
204 "Standard transactional, session, or DML block. No blocking schema lock required."
205 }
206 "create-table" => "Creating a new table does not block existing application queries.",
207 _ => {
208 "This statement triggers an unclassified heavy lock.\n\
209 General DB Safety Rules:\n\
210 1. Always run DDL with a short timeout (e.g., SET lock_timeout = '2s';).\n\
211 2. Avoid running DDL in long-running transactions.\n\
212 3. If a table rewrite is unavoidable, schedule downtime or use zero-downtime tools like pg_repack."
213 }
214 }
215}
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220 use std::io::Write;
221 use tempfile::NamedTempFile;
222
223 #[test]
224 fn test_default_config() {
225 let cfg = Config::default_config();
226 assert_eq!(cfg.default_threshold, 100_000);
227 assert_eq!(
228 cfg.rules.get("adding-field-with-default").unwrap().tier,
229 LockTier::Tier1
230 );
231 }
232
233 #[test]
234 fn test_valid_toml_override() {
235 let mut file = NamedTempFile::new().unwrap();
236 let toml = r#"
237 default_threshold = 500
238
239 [rules.adding-field-with-default]
240 tier = "Tier2"
241 threshold = 1000
242 "#;
243 file.write_all(toml.as_bytes()).unwrap();
244
245 let cfg = Config::load(file.path().to_str().unwrap()).unwrap();
246
247 assert_eq!(cfg.default_threshold, 500);
248
249 let rule_cfg = cfg.rules.get("adding-field-with-default").unwrap();
250 assert_eq!(rule_cfg.tier, LockTier::Tier2);
251 assert_eq!(rule_cfg.threshold, Some(1000));
252
253 let untouched = cfg.rules.get("ban-drop-column").unwrap();
255 assert_eq!(untouched.tier, LockTier::Tier1);
256 assert_eq!(untouched.threshold, None);
257 }
258
259 #[test]
260 fn test_malformed_config_fails() {
261 let mut file = NamedTempFile::new().unwrap();
262 let toml = r#"
263 default_threshold = 5000
264
265 [rules.typo-fake-rule]
266 tier = "Tier3"
267 "#;
268 file.write_all(toml.as_bytes()).unwrap();
269
270 let result = Config::load(file.path().to_str().unwrap());
271 assert!(result.is_err());
272 assert!(result.unwrap_err().to_string().contains("Unknown rule"));
273 }
274}