waypoint_core/
hooks.rs

1//! SQL callback hooks that run before/after migrations (Flyway-compatible).
2
3use std::collections::HashMap;
4use std::fmt;
5use std::path::PathBuf;
6
7#[cfg(feature = "postgres")]
8use tokio_postgres::Client;
9
10use crate::config::HooksConfig;
11#[cfg(feature = "postgres")]
12use crate::db;
13use crate::db::DbClient;
14use crate::error::{Result, WaypointError};
15use crate::placeholder::replace_placeholders;
16
17/// The phase at which a hook runs.
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub enum HookType {
20    /// Runs once before the entire migration run begins.
21    BeforeMigrate,
22    /// Runs once after the entire migration run completes.
23    AfterMigrate,
24    /// Runs before each individual migration is applied.
25    BeforeEachMigrate,
26    /// Runs after each individual migration is applied.
27    AfterEachMigrate,
28}
29
30impl fmt::Display for HookType {
31    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32        match self {
33            HookType::BeforeMigrate => write!(f, "beforeMigrate"),
34            HookType::AfterMigrate => write!(f, "afterMigrate"),
35            HookType::BeforeEachMigrate => write!(f, "beforeEachMigrate"),
36            HookType::AfterEachMigrate => write!(f, "afterEachMigrate"),
37        }
38    }
39}
40
41/// A hook SQL script discovered on disk or specified in config.
42#[derive(Debug, Clone)]
43pub struct ResolvedHook {
44    /// The phase at which this hook should be executed.
45    pub hook_type: HookType,
46    /// Filename of the hook SQL script.
47    pub script_name: String,
48    /// Raw SQL content of the hook file.
49    pub sql: String,
50}
51
52/// File prefixes that indicate hook callback files (Flyway-compatible).
53type HookPrefixEntry = (&'static str, fn() -> HookType);
54const HOOK_PREFIXES: &[HookPrefixEntry] = &[
55    ("beforeEachMigrate", || HookType::BeforeEachMigrate),
56    ("afterEachMigrate", || HookType::AfterEachMigrate),
57    ("beforeMigrate", || HookType::BeforeMigrate),
58    ("afterMigrate", || HookType::AfterMigrate),
59];
60
61/// Check if a filename is a hook callback file (not a migration).
62pub fn is_hook_file(filename: &str) -> bool {
63    HOOK_PREFIXES
64        .iter()
65        .any(|(prefix, _)| filename.starts_with(prefix) && filename.ends_with(".sql"))
66}
67
68/// Scan migration locations for SQL callback hook files.
69///
70/// Recognizes:
71///   - `beforeMigrate.sql` / `beforeMigrate__*.sql`
72///   - `afterMigrate.sql` / `afterMigrate__*.sql`
73///   - `beforeEachMigrate.sql` / `beforeEachMigrate__*.sql`
74///   - `afterEachMigrate.sql` / `afterEachMigrate__*.sql`
75///
76/// Multiple files per hook type are sorted alphabetically.
77pub fn scan_hooks(locations: &[PathBuf]) -> Result<Vec<ResolvedHook>> {
78    let mut hooks = Vec::new();
79
80    for location in locations {
81        if !location.exists() {
82            continue;
83        }
84
85        let entries = std::fs::read_dir(location).map_err(|e| {
86            WaypointError::IoError(std::io::Error::new(
87                e.kind(),
88                format!(
89                    "Failed to read hook directory '{}': {}",
90                    location.display(),
91                    e
92                ),
93            ))
94        })?;
95
96        let mut files: Vec<_> = entries
97            .filter_map(|e| e.ok())
98            .filter(|e| e.path().is_file())
99            .collect();
100
101        // Sort alphabetically for deterministic ordering
102        files.sort_by_key(|e| e.file_name());
103
104        for entry in files {
105            let path = entry.path();
106            let filename = match path.file_name().and_then(|n| n.to_str()) {
107                Some(name) => name.to_string(),
108                None => continue,
109            };
110
111            if !filename.ends_with(".sql") {
112                continue;
113            }
114
115            // Check each hook prefix
116            for (prefix, type_fn) in HOOK_PREFIXES {
117                if filename.starts_with(prefix) {
118                    // Must be exactly `prefix.sql` or `prefix__*.sql`
119                    let rest = &filename[prefix.len()..filename.len() - 4]; // strip prefix and .sql
120                    if rest.is_empty() || rest.starts_with("__") {
121                        let sql = std::fs::read_to_string(&path)?;
122                        hooks.push(ResolvedHook {
123                            hook_type: type_fn(),
124                            script_name: filename.clone(),
125                            sql,
126                        });
127                        break;
128                    }
129                }
130            }
131        }
132    }
133
134    // Sort within each hook type alphabetically by script name
135    hooks.sort_by(|a, b| {
136        a.hook_type
137            .to_string()
138            .cmp(&b.hook_type.to_string())
139            .then_with(|| a.script_name.cmp(&b.script_name))
140    });
141
142    Ok(hooks)
143}
144
145/// Load hook SQL files specified in the TOML `[hooks]` config section.
146pub fn load_config_hooks(config: &HooksConfig) -> Result<Vec<ResolvedHook>> {
147    let mut hooks = Vec::new();
148
149    let sections: &[(HookType, &[PathBuf])] = &[
150        (HookType::BeforeMigrate, &config.before_migrate),
151        (HookType::AfterMigrate, &config.after_migrate),
152        (HookType::BeforeEachMigrate, &config.before_each_migrate),
153        (HookType::AfterEachMigrate, &config.after_each_migrate),
154    ];
155
156    for (hook_type, paths) in sections {
157        for path in *paths {
158            let sql = std::fs::read_to_string(path).map_err(|e| {
159                WaypointError::IoError(std::io::Error::new(
160                    e.kind(),
161                    format!("Failed to read hook file '{}': {}", path.display(), e),
162                ))
163            })?;
164
165            let script_name = path
166                .file_name()
167                .and_then(|n| n.to_str())
168                .unwrap_or_else(|| path.to_str().unwrap_or("unknown"))
169                .to_string();
170
171            hooks.push(ResolvedHook {
172                hook_type: hook_type.clone(),
173                script_name,
174                sql,
175            });
176        }
177    }
178
179    Ok(hooks)
180}
181
182/// Run all hooks of a given type.
183///
184/// Returns total execution time in milliseconds.
185#[cfg(feature = "postgres")]
186pub async fn run_hooks(
187    client: &Client,
188    hooks: &[ResolvedHook],
189    phase: &HookType,
190    placeholders: &HashMap<String, String>,
191) -> Result<(usize, i32)> {
192    let mut total_ms = 0;
193    let mut count = 0;
194
195    for hook in hooks.iter().filter(|h| &h.hook_type == phase) {
196        log::info!("Running {} hook: {}", phase, hook.script_name);
197
198        let sql = replace_placeholders(&hook.sql, placeholders)?;
199
200        match db::execute_in_transaction(client, &sql).await {
201            Ok(exec_time) => {
202                total_ms += exec_time;
203                count += 1;
204            }
205            Err(e) => {
206                let reason = match &e {
207                    WaypointError::DatabaseError(db_err) => crate::error::format_db_error(db_err),
208                    other => other.to_string(),
209                };
210                return Err(WaypointError::HookFailed {
211                    phase: phase.to_string(),
212                    script: hook.script_name.clone(),
213                    reason,
214                });
215            }
216        }
217    }
218
219    Ok((count, total_ms))
220}
221
222/// Run all hooks of a given phase (dialect-aware entry).
223///
224/// On PostgreSQL each hook is wrapped in a transaction (matching the legacy
225/// `run_hooks` PG entry). On MySQL hooks execute via `execute_raw` — MySQL DDL
226/// auto-commits, so a transaction wrapper would buy nothing for DDL hooks.
227/// Returns `(hook_count, total_ms)`.
228pub async fn run_hooks_db(
229    client: &DbClient,
230    hooks: &[ResolvedHook],
231    phase: &HookType,
232    placeholders: &HashMap<String, String>,
233) -> Result<(usize, i32)> {
234    let mut total_ms = 0;
235    let mut count = 0;
236
237    for hook in hooks.iter().filter(|h| &h.hook_type == phase) {
238        log::info!("Running {} hook: {}", phase, hook.script_name);
239
240        let sql = replace_placeholders(&hook.sql, placeholders)?;
241
242        let exec_result = match client.dialect_kind() {
243            crate::dialect::DialectKind::Postgres => client.execute_in_transaction(&sql).await,
244            crate::dialect::DialectKind::Mysql => client.execute_raw(&sql).await,
245        };
246
247        match exec_result {
248            Ok(exec_time) => {
249                total_ms += exec_time;
250                count += 1;
251            }
252            Err(e) => {
253                // Match the legacy `run_hooks` error format: when the cause is
254                // a tokio_postgres::Error, surface the inner DbError detail/hint
255                // (`format_db_error`) without the "Database error: " prefix
256                // that WaypointError::Display would prepend.
257                #[cfg(feature = "postgres")]
258                let reason = match &e {
259                    WaypointError::DatabaseError(db_err) => crate::error::format_db_error(db_err),
260                    other => other.to_string(),
261                };
262                #[cfg(not(feature = "postgres"))]
263                let reason = e.to_string();
264                return Err(WaypointError::HookFailed {
265                    phase: phase.to_string(),
266                    script: hook.script_name.clone(),
267                    reason,
268                });
269            }
270        }
271    }
272
273    Ok((count, total_ms))
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279    use std::fs;
280
281    fn create_temp_dir(name: &str) -> PathBuf {
282        let dir = std::env::temp_dir().join(format!("waypoint_hooks_test_{}", name));
283        let _ = fs::remove_dir_all(&dir);
284        fs::create_dir_all(&dir).unwrap();
285        dir
286    }
287
288    #[test]
289    fn test_is_hook_file() {
290        assert!(is_hook_file("beforeMigrate.sql"));
291        assert!(is_hook_file("afterMigrate.sql"));
292        assert!(is_hook_file("beforeEachMigrate.sql"));
293        assert!(is_hook_file("afterEachMigrate.sql"));
294        assert!(is_hook_file("beforeMigrate__Disable_triggers.sql"));
295        assert!(is_hook_file("afterMigrate__Refresh_views.sql"));
296
297        assert!(!is_hook_file("V1__Create_table.sql"));
298        assert!(!is_hook_file("R__Create_view.sql"));
299        assert!(!is_hook_file("beforeMigrate.txt"));
300        assert!(!is_hook_file("random.sql"));
301    }
302
303    #[test]
304    fn test_scan_hooks_finds_callback_files() {
305        let dir = create_temp_dir("scan");
306        fs::write(dir.join("beforeMigrate.sql"), "SELECT 1;").unwrap();
307        fs::write(dir.join("afterMigrate__Refresh_views.sql"), "SELECT 2;").unwrap();
308        fs::write(dir.join("V1__Create_table.sql"), "CREATE TABLE t(id INT);").unwrap();
309        fs::write(dir.join("R__Create_view.sql"), "CREATE VIEW v AS SELECT 1;").unwrap();
310
311        let hooks = scan_hooks(std::slice::from_ref(&dir)).unwrap();
312
313        assert_eq!(hooks.len(), 2);
314
315        let before: Vec<_> = hooks
316            .iter()
317            .filter(|h| h.hook_type == HookType::BeforeMigrate)
318            .collect();
319        let after: Vec<_> = hooks
320            .iter()
321            .filter(|h| h.hook_type == HookType::AfterMigrate)
322            .collect();
323        assert_eq!(before.len(), 1);
324        assert_eq!(before[0].script_name, "beforeMigrate.sql");
325        assert_eq!(after.len(), 1);
326        assert_eq!(after[0].script_name, "afterMigrate__Refresh_views.sql");
327
328        let _ = fs::remove_dir_all(&dir);
329    }
330
331    #[test]
332    fn test_scan_hooks_multiple_sorted_alphabetically() {
333        let dir = create_temp_dir("multi");
334        fs::write(dir.join("beforeMigrate__B_second.sql"), "SELECT 2;").unwrap();
335        fs::write(dir.join("beforeMigrate__A_first.sql"), "SELECT 1;").unwrap();
336        fs::write(dir.join("beforeMigrate.sql"), "SELECT 0;").unwrap();
337
338        let hooks = scan_hooks(std::slice::from_ref(&dir)).unwrap();
339
340        assert_eq!(hooks.len(), 3);
341        assert_eq!(hooks[0].script_name, "beforeMigrate.sql");
342        assert_eq!(hooks[1].script_name, "beforeMigrate__A_first.sql");
343        assert_eq!(hooks[2].script_name, "beforeMigrate__B_second.sql");
344
345        let _ = fs::remove_dir_all(&dir);
346    }
347
348    #[test]
349    fn test_load_config_hooks() {
350        let dir = create_temp_dir("config");
351        let hook_file = dir.join("pre.sql");
352        fs::write(&hook_file, "SET work_mem = '256MB';").unwrap();
353
354        let config = HooksConfig {
355            before_migrate: vec![hook_file],
356            after_migrate: vec![],
357            before_each_migrate: vec![],
358            after_each_migrate: vec![],
359        };
360
361        let hooks = load_config_hooks(&config).unwrap();
362        assert_eq!(hooks.len(), 1);
363        assert_eq!(hooks[0].hook_type, HookType::BeforeMigrate);
364        assert_eq!(hooks[0].sql, "SET work_mem = '256MB';");
365
366        let _ = fs::remove_dir_all(&dir);
367    }
368
369    #[test]
370    fn test_load_config_hooks_missing_file() {
371        let config = HooksConfig {
372            before_migrate: vec![PathBuf::from("/nonexistent/hook.sql")],
373            after_migrate: vec![],
374            before_each_migrate: vec![],
375            after_each_migrate: vec![],
376        };
377
378        assert!(load_config_hooks(&config).is_err());
379    }
380}
waypoint_core/hooks.rs

waypoint_core/
hooks.rs
