Skip to main content

dbkit/
initialization.rs

1use crate::DbkitError;
2use crate::base_handler::{BaseHandler, FetchMode, WriteOp};
3use crate::config::Backend;
4use sqlx::{AnyPool, Row};
5use std::collections::hash_map::DefaultHasher;
6use std::hash::{Hash, Hasher};
7use tracing::{error, info};
8
9/// Batch DDL migration executor with tracking.
10///
11/// Maintains a `_dbkit_migrations` table so already-applied migrations
12/// are skipped on subsequent runs. Each migration is identified by a
13/// user-provided name and a content hash.
14///
15/// The tracking table DDL and the internal queries are backend-aware
16/// (Postgres / MySQL / SQLite), so the migration runner works on any backend.
17/// The *user-supplied* migration SQL is still backend-native — write it for the
18/// database you connected to.
19pub struct InitializationHandler {
20    handler: BaseHandler,
21    backend: Backend,
22}
23
24impl InitializationHandler {
25    /// Create a handler for the given pool and backend.
26    ///
27    /// The backend is typically obtained from
28    /// [`ConnectionManager::backend`](crate::ConnectionManager::backend).
29    pub fn new(pool: AnyPool, backend: Backend) -> Self {
30        Self {
31            handler: BaseHandler::new(pool),
32            backend,
33        }
34    }
35
36    /// The placeholder for the `n`-th bind parameter (1-based) for this backend.
37    fn placeholder(&self, n: usize) -> String {
38        match self.backend {
39            Backend::Postgres => format!("${n}"),
40            Backend::MySql | Backend::Sqlite => "?".to_string(),
41        }
42    }
43
44    /// Backend-specific DDL for the migrations tracking table.
45    ///
46    /// `name` is the primary key (no auto-increment column needed), which keeps
47    /// the schema portable. `applied_at` is a timestamp defaulting to "now".
48    fn tracking_table_ddl(&self) -> &'static str {
49        match self.backend {
50            Backend::Postgres => {
51                "CREATE TABLE IF NOT EXISTS _dbkit_migrations (
52                    name TEXT PRIMARY KEY,
53                    hash TEXT NOT NULL,
54                    applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
55                )"
56            }
57            Backend::MySql => {
58                // TEXT can't be a primary key without a prefix length in MySQL.
59                "CREATE TABLE IF NOT EXISTS _dbkit_migrations (
60                    name VARCHAR(255) PRIMARY KEY,
61                    hash TEXT NOT NULL,
62                    applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
63                )"
64            }
65            Backend::Sqlite => {
66                "CREATE TABLE IF NOT EXISTS _dbkit_migrations (
67                    name TEXT PRIMARY KEY,
68                    hash TEXT NOT NULL,
69                    applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
70                )"
71            }
72        }
73    }
74
75    /// A SQL expression that renders `applied_at` as text, so it can be read
76    /// through sqlx's `Any` driver (which can't decode native timestamp types).
77    fn applied_at_as_text(&self) -> &'static str {
78        match self.backend {
79            // MySQL casts to CHAR; Postgres and SQLite cast to TEXT.
80            Backend::MySql => "CAST(applied_at AS CHAR)",
81            Backend::Postgres | Backend::Sqlite => "CAST(applied_at AS TEXT)",
82        }
83    }
84
85    /// Ensure the migrations tracking table exists.
86    async fn ensure_tracking_table(&self) -> Result<(), DbkitError> {
87        self.handler
88            .execute_write(WriteOp::BatchDDL {
89                queries: &[self.tracking_table_ddl()],
90            })
91            .await?;
92        Ok(())
93    }
94
95    /// Compute a hash of the SQL content for change detection.
96    ///
97    /// FNV-1a (64-bit), implemented inline: the algorithm is fixed forever, so
98    /// hashes persisted in `_dbkit_migrations` stay comparable across dbkit and
99    /// Rust releases. (`DefaultHasher` — used before 0.5 — explicitly does not
100    /// guarantee a stable algorithm between Rust releases, which would have
101    /// made every already-applied migration report "content has changed" after
102    /// a toolchain upgrade.)
103    fn hash_sql(sql: &str) -> String {
104        const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
105        const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
106        let mut h = FNV_OFFSET;
107        for b in sql.as_bytes() {
108            h ^= u64::from(*b);
109            h = h.wrapping_mul(FNV_PRIME);
110        }
111        format!("{h:016x}")
112    }
113
114    /// The pre-0.5 hash (`DefaultHasher`), kept only to recognize rows recorded
115    /// by older dbkit versions so they can be upgraded in place rather than
116    /// misreported as changed content.
117    fn legacy_hash_sql(sql: &str) -> String {
118        let mut hasher = DefaultHasher::new();
119        sql.hash(&mut hasher);
120        format!("{:016x}", hasher.finish())
121    }
122
123    /// Run a named migration. Skips if already applied with the same content hash.
124    ///
125    /// If the migration name exists but the hash differs, it returns an error
126    /// (content changed after being applied).
127    pub async fn run_named_migration(&self, name: &str, sql: &str) -> Result<(), DbkitError> {
128        self.ensure_tracking_table().await?;
129
130        let hash = Self::hash_sql(sql);
131
132        // Check if already applied
133        let select = format!(
134            "SELECT hash FROM _dbkit_migrations WHERE name = {}",
135            self.placeholder(1)
136        );
137        let result = self
138            .handler
139            .execute_write(WriteOp::Single {
140                query: select.as_str(),
141                params: vec![name.into()],
142                mode: FetchMode::Optional,
143            })
144            .await?;
145
146        if let Some(row) = result.optional()? {
147            let existing_hash: String = row.get(0);
148            if existing_hash == hash {
149                info!("migration '{}' already applied, skipping", name);
150                return Ok(());
151            }
152            // A row recorded by dbkit < 0.5 holds a `DefaultHasher` hash. If it
153            // matches the same SQL under the legacy algorithm, the content is
154            // unchanged — upgrade the stored hash in place and skip.
155            if existing_hash == Self::legacy_hash_sql(sql) {
156                info!(
157                    "migration '{}' already applied (pre-0.5 hash), upgrading stored hash",
158                    name
159                );
160                let update = format!(
161                    "UPDATE _dbkit_migrations SET hash = {} WHERE name = {}",
162                    self.placeholder(1),
163                    self.placeholder(2)
164                );
165                self.handler
166                    .execute_write(WriteOp::Single {
167                        query: update.as_str(),
168                        params: vec![hash.into(), name.into()],
169                        mode: FetchMode::None,
170                    })
171                    .await?;
172                return Ok(());
173            }
174            return Err(DbkitError::Migration(format!(
175                "migration '{}' was already applied but content has changed (hash {} → {})",
176                name, existing_hash, hash
177            )));
178        }
179
180        // Run the migration
181        info!("applying migration '{}'...", name);
182        let queries: Vec<String> = sql
183            .split(';')
184            .map(|s| s.trim().to_string())
185            .filter(|s| !s.is_empty())
186            .collect();
187
188        let query_refs: Vec<&str> = queries.iter().map(|s| s.as_str()).collect();
189
190        match self
191            .handler
192            .execute_write(WriteOp::BatchDDL {
193                queries: &query_refs,
194            })
195            .await
196        {
197            Ok(_) => {
198                info!(
199                    "migration '{}': {} DDL statements executed",
200                    name,
201                    query_refs.len()
202                );
203            }
204            Err(e) => {
205                error!("migration '{}' failed: {:?}", name, e);
206                return Err(DbkitError::Migration(e.to_string()));
207            }
208        }
209
210        // Record the migration
211        let insert = format!(
212            "INSERT INTO _dbkit_migrations (name, hash) VALUES ({}, {})",
213            self.placeholder(1),
214            self.placeholder(2)
215        );
216        self.handler
217            .execute_write(WriteOp::Single {
218                query: insert.as_str(),
219                params: vec![name.into(), hash.clone().into()],
220                mode: FetchMode::None,
221            })
222            .await?;
223
224        info!("migration '{}' recorded", name);
225        Ok(())
226    }
227
228    /// Run migrations from a SQL string (semicolon-separated DDL statements).
229    ///
230    /// Statements are split on `;`, so this does not support SQL that embeds
231    /// semicolons inside a single statement (string literals, PL/pgSQL
232    /// function bodies, triggers). Run such migrations as a single statement
233    /// through [`WriteOp::BatchDDL`] directly.
234    ///
235    /// This is the simple/legacy API — it runs all statements unconditionally
236    /// without tracking. Use [`run_named_migration`](Self::run_named_migration)
237    /// for tracked migrations.
238    pub async fn run_migrations(&self, sql: &str) -> Result<(), DbkitError> {
239        info!("running database migrations...");
240
241        let queries: Vec<String> = sql
242            .split(';')
243            .map(|s| s.trim().to_string())
244            .filter(|s| !s.is_empty())
245            .collect();
246
247        let query_refs: Vec<&str> = queries.iter().map(|s| s.as_str()).collect();
248
249        match self
250            .handler
251            .execute_write(WriteOp::BatchDDL {
252                queries: &query_refs,
253            })
254            .await
255        {
256            Ok(_) => {
257                info!("{} DDL statements executed", query_refs.len());
258            }
259            Err(e) => {
260                error!("migration failed: {:?}", e);
261                return Err(DbkitError::Migration(e.to_string()));
262            }
263        }
264
265        Ok(())
266    }
267
268    /// List all applied migrations (name, hash, applied_at) in application order.
269    pub async fn applied_migrations(&self) -> Result<Vec<(String, String, String)>, DbkitError> {
270        self.ensure_tracking_table().await?;
271
272        let select = format!(
273            "SELECT name, hash, {} FROM _dbkit_migrations ORDER BY applied_at, name",
274            self.applied_at_as_text()
275        );
276        let result = self
277            .handler
278            .execute_write(WriteOp::Single {
279                query: select.as_str(),
280                params: vec![],
281                mode: FetchMode::All,
282            })
283            .await?;
284
285        let rows = result.all()?;
286        Ok(rows
287            .iter()
288            .map(|row| {
289                let name: String = row.get(0);
290                let hash: String = row.get(1);
291                let applied_at: String = row.get(2);
292                (name, hash, applied_at)
293            })
294            .collect())
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    /// The stored migration hashes must never change meaning, so pin the
303    /// algorithm to the published FNV-1a 64-bit reference vectors.
304    #[test]
305    fn hash_sql_is_stable_fnv1a() {
306        assert_eq!(InitializationHandler::hash_sql(""), "cbf29ce484222325");
307        assert_eq!(InitializationHandler::hash_sql("a"), "af63dc4c8601ec8c");
308        assert_eq!(InitializationHandler::hash_sql("foobar"), "85944171f73967e8");
309    }
310}