aster_forge_db_migration/
coordination.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use sea_orm_migration::MigratorTrait;
6use sea_orm_migration::sea_orm::{
7    ConnectionTrait, DatabaseConnection, DatabaseExecutor, DatabaseTransaction, DbBackend, DbErr,
8    RuntimeErr, Statement, TransactionTrait,
9};
10
11const DEFAULT_MYSQL_LOCK_TIMEOUT_SECONDS: u64 = 300;
12const MYSQL_LOCK_NAME_MAX_BYTES: usize = 64;
13
14/// Boxed migration callback future tied to the coordinated database connection.
15pub type MigrationFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, DbErr>> + Send + 'a>>;
16
17/// Stable cross-process migration lock configuration.
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct MigrationLockOptions {
20    namespace: String,
21    postgres_advisory_key: i64,
22    mysql_timeout_seconds: u64,
23}
24
25impl MigrationLockOptions {
26    /// Creates options using a deterministic `PostgreSQL` advisory key derived from `namespace`.
27    ///
28    /// Products migrating from an existing lock implementation should call
29    /// [`Self::with_postgres_advisory_key`] to preserve the old key during rolling upgrades.
30    pub fn new(namespace: impl Into<String>) -> Self {
31        let namespace = namespace.into();
32        Self {
33            postgres_advisory_key: stable_advisory_key(&namespace),
34            namespace,
35            mysql_timeout_seconds: DEFAULT_MYSQL_LOCK_TIMEOUT_SECONDS,
36        }
37    }
38
39    /// Overrides the `PostgreSQL` advisory key while preserving the shared namespace for `MySQL`.
40    #[must_use]
41    pub const fn with_postgres_advisory_key(mut self, key: i64) -> Self {
42        self.postgres_advisory_key = key;
43        self
44    }
45
46    /// Overrides the `MySQL` named-lock wait timeout in whole seconds.
47    #[must_use]
48    pub const fn with_mysql_timeout_seconds(mut self, seconds: u64) -> Self {
49        self.mysql_timeout_seconds = seconds;
50        self
51    }
52
53    /// Returns the stable product namespace used for `MySQL` named locks.
54    #[must_use]
55    pub fn namespace(&self) -> &str {
56        &self.namespace
57    }
58
59    /// Returns the `PostgreSQL` transaction-scoped advisory-lock key.
60    #[must_use]
61    pub const fn postgres_advisory_key(&self) -> i64 {
62        self.postgres_advisory_key
63    }
64
65    /// Returns the `MySQL` named-lock wait timeout in seconds.
66    #[must_use]
67    pub const fn mysql_timeout_seconds(&self) -> u64 {
68        self.mysql_timeout_seconds
69    }
70
71    fn validate(&self) -> Result<(), DbErr> {
72        if self.namespace.is_empty() || self.namespace.len() > MYSQL_LOCK_NAME_MAX_BYTES {
73            return Err(DbErr::Custom(format!(
74                "migration lock namespace must contain 1-{MYSQL_LOCK_NAME_MAX_BYTES} bytes"
75            )));
76        }
77        if self.namespace.contains('\0') {
78            return Err(DbErr::Custom(
79                "migration lock namespace must not contain NUL bytes".to_string(),
80            ));
81        }
82        i64::try_from(self.mysql_timeout_seconds).map_err(|_| {
83            DbErr::Custom("MySQL migration lock timeout exceeds signed 64-bit range".to_string())
84        })?;
85        Ok(())
86    }
87}
88
89/// Runs a product migration callback while holding the backend's process-wide migration lock.
90///
91/// `PostgreSQL` uses a transaction-scoped advisory lock. `MySQL` uses a dedicated single-connection
92/// pool and a connection-bound named lock without wrapping DDL in a transaction. `SQLite` runs the
93/// callback in a transaction without an additional external lock.
94///
95/// # Errors
96///
97/// Returns an error when lock options are invalid, the backend is unsupported, connection,
98/// transaction, or lock operations fail, the callback fails, or transactional finalization fails.
99pub async fn with_migration_lock<T, F>(
100    database: &DatabaseConnection,
101    options: &MigrationLockOptions,
102    operation: F,
103) -> Result<T, DbErr>
104where
105    F: for<'a> FnOnce(DatabaseExecutor<'a>) -> MigrationFuture<'a, T>,
106{
107    options.validate()?;
108    match database.get_database_backend() {
109        DbBackend::Postgres => {
110            run_transactional_migration(database, options, operation, true).await
111        }
112        DbBackend::Sqlite => run_transactional_migration(database, options, operation, false).await,
113        DbBackend::MySql => run_mysql_migration(database, options, operation).await,
114        _ => Err(DbErr::Custom(
115            "unsupported database backend for migration coordination".to_string(),
116        )),
117    }
118}
119
120async fn run_transactional_migration<T, F>(
121    database: &DatabaseConnection,
122    options: &MigrationLockOptions,
123    operation: F,
124    acquire_postgres_advisory_lock: bool,
125) -> Result<T, DbErr>
126where
127    F: for<'a> FnOnce(DatabaseExecutor<'a>) -> MigrationFuture<'a, T>,
128{
129    let transaction = database.begin().await?;
130    if acquire_postgres_advisory_lock {
131        acquire_postgres_lock(&transaction, options).await?;
132    }
133    let operation_result = operation((&transaction).into()).await;
134
135    match operation_result {
136        Ok(value) => {
137            transaction.commit().await?;
138            Ok(value)
139        }
140        Err(error) => rollback_preserving_error(transaction, error).await,
141    }
142}
143
144async fn run_mysql_migration<T, F>(
145    database: &DatabaseConnection,
146    options: &MigrationLockOptions,
147    operation: F,
148) -> Result<T, DbErr>
149where
150    F: for<'a> FnOnce(DatabaseExecutor<'a>) -> MigrationFuture<'a, T>,
151{
152    let dedicated_database = create_mysql_migration_connection(database).await?;
153    acquire_mysql_lock(&dedicated_database, options).await?;
154    let operation_result = operation((&dedicated_database).into()).await;
155    let release_result = release_mysql_lock(&dedicated_database, options).await;
156
157    let result = match (operation_result, release_result) {
158        (Ok(value), Ok(())) => Ok(value),
159        (Err(error), Ok(())) => Err(error),
160        (Ok(_), Err(release_error)) => Err(release_error),
161        (Err(operation_error), Err(release_error)) => Err(DbErr::Custom(format!(
162            "migration operation failed: {operation_error}; additionally failed to release \
163             the MySQL migration lock: {release_error}"
164        ))),
165    };
166
167    if let Err(close_error) = dedicated_database.close().await {
168        tracing::warn!(%close_error, "failed to close dedicated MySQL migration connection");
169    }
170    result
171}
172
173async fn create_mysql_migration_connection(
174    database: &DatabaseConnection,
175) -> Result<DatabaseConnection, DbErr> {
176    let source_pool = database.get_mysql_connection_pool();
177    let connect_options = source_pool.connect_options();
178    let dedicated_pool = source_pool
179        .options()
180        .clone()
181        .max_connections(1)
182        .min_connections(1)
183        .idle_timeout(None)
184        .max_lifetime(None)
185        .test_before_acquire(false)
186        .before_acquire(|_, _| Box::pin(async { Ok(true) }))
187        .after_release(|_, _| Box::pin(async { Ok(true) }))
188        .connect_with((*connect_options).clone())
189        .await
190        .map_err(|error| DbErr::Conn(RuntimeErr::SqlxError(Arc::new(error))))?;
191    Ok(dedicated_pool.into())
192}
193
194/// Runs a standard `SeaORM` migrator while holding the backend migration lock.
195///
196/// # Errors
197///
198/// Returns any migration coordination, transaction, lock, or [`MigratorTrait::up`] failure.
199pub async fn run_migrator_with_lock<M>(
200    database: &DatabaseConnection,
201    options: &MigrationLockOptions,
202    steps: Option<u32>,
203) -> Result<(), DbErr>
204where
205    M: MigratorTrait + 'static,
206{
207    with_migration_lock(database, options, |connection| {
208        Box::pin(M::up(connection, steps))
209    })
210    .await
211}
212
213async fn acquire_postgres_lock(
214    transaction: &DatabaseTransaction,
215    options: &MigrationLockOptions,
216) -> Result<(), DbErr> {
217    transaction
218        .query_one_raw(Statement::from_sql_and_values(
219            DbBackend::Postgres,
220            "SELECT pg_advisory_xact_lock($1)",
221            [options.postgres_advisory_key.into()],
222        ))
223        .await?;
224    Ok(())
225}
226
227async fn acquire_mysql_lock(
228    connection: &DatabaseConnection,
229    options: &MigrationLockOptions,
230) -> Result<(), DbErr> {
231    let timeout = i64::try_from(options.mysql_timeout_seconds).map_err(|_| {
232        DbErr::Custom("MySQL migration lock timeout exceeds signed 64-bit range".to_string())
233    })?;
234    let acquired = mysql_lock_query_result(
235        connection,
236        "SELECT GET_LOCK(?, ?)",
237        [options.namespace.clone().into(), timeout.into()],
238        "acquire",
239    )
240    .await?;
241    if acquired {
242        Ok(())
243    } else {
244        Err(DbErr::Custom(format!(
245            "timed out after {} seconds waiting for MySQL migration lock '{}'",
246            options.mysql_timeout_seconds, options.namespace
247        )))
248    }
249}
250
251async fn release_mysql_lock(
252    connection: &DatabaseConnection,
253    options: &MigrationLockOptions,
254) -> Result<(), DbErr> {
255    let released = mysql_lock_query_result(
256        connection,
257        "SELECT RELEASE_LOCK(?)",
258        [options.namespace.clone().into()],
259        "release",
260    )
261    .await?;
262    if released {
263        Ok(())
264    } else {
265        Err(DbErr::Custom(format!(
266            "MySQL migration lock '{}' was not owned by the migration connection",
267            options.namespace
268        )))
269    }
270}
271
272async fn mysql_lock_query_result<C, const N: usize>(
273    connection: &C,
274    sql: &str,
275    values: [sea_orm_migration::sea_orm::Value; N],
276    operation: &str,
277) -> Result<bool, DbErr>
278where
279    C: ConnectionTrait,
280{
281    let row = connection
282        .query_one_raw(Statement::from_sql_and_values(
283            DbBackend::MySql,
284            sql,
285            values,
286        ))
287        .await?
288        .ok_or_else(|| {
289            DbErr::Custom(format!(
290                "MySQL migration lock {operation} query returned no rows"
291            ))
292        })?;
293
294    if let Ok(value) = row.try_get_by_index::<Option<i64>>(0) {
295        return value.map(|value| value == 1).ok_or_else(|| {
296            DbErr::Custom(format!(
297                "MySQL migration lock {operation} query returned NULL"
298            ))
299        });
300    }
301    if let Ok(value) = row.try_get_by_index::<Option<i32>>(0) {
302        return value.map(|value| value == 1).ok_or_else(|| {
303            DbErr::Custom(format!(
304                "MySQL migration lock {operation} query returned NULL"
305            ))
306        });
307    }
308
309    Err(DbErr::Custom(format!(
310        "failed to decode MySQL migration lock {operation} result"
311    )))
312}
313
314async fn rollback_preserving_error<T>(
315    transaction: DatabaseTransaction,
316    error: DbErr,
317) -> Result<T, DbErr> {
318    if let Err(rollback_error) = transaction.rollback().await {
319        tracing::warn!(%rollback_error, "failed to rollback migration transaction after callback error");
320    }
321    Err(error)
322}
323
324fn stable_advisory_key(namespace: &str) -> i64 {
325    const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
326    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
327
328    let hash = namespace.bytes().fold(FNV_OFFSET_BASIS, |hash, byte| {
329        (hash ^ u64::from(byte)).wrapping_mul(FNV_PRIME)
330    });
331    i64::from_be_bytes(hash.to_be_bytes())
332}
333
334#[cfg(test)]
335mod tests {
336    use sea_orm_migration::sea_orm::{ConnectionTrait, Database, DbErr};
337
338    use super::{MigrationLockOptions, stable_advisory_key, with_migration_lock};
339
340    #[test]
341    fn advisory_key_is_stable_and_namespace_sensitive() {
342        assert_eq!(
343            stable_advisory_key("aster_drive:database_migrations"),
344            stable_advisory_key("aster_drive:database_migrations")
345        );
346        assert_ne!(
347            stable_advisory_key("aster_drive:database_migrations"),
348            stable_advisory_key("aster_yggdrasil:database_migrations")
349        );
350    }
351
352    #[test]
353    fn migration_lock_options_validate_namespace_boundaries() {
354        assert!(MigrationLockOptions::new("a").validate().is_ok());
355        assert!(MigrationLockOptions::new("a".repeat(64)).validate().is_ok());
356        assert!(MigrationLockOptions::new("").validate().is_err());
357        assert!(
358            MigrationLockOptions::new("a".repeat(65))
359                .validate()
360                .is_err()
361        );
362        assert!(MigrationLockOptions::new("bad\0name").validate().is_err());
363    }
364
365    #[tokio::test]
366    async fn sqlite_callback_commits_on_success() {
367        let database = Database::connect("sqlite::memory:").await.unwrap();
368        with_migration_lock(
369            &database,
370            &MigrationLockOptions::new("sqlite-success"),
371            |transaction| {
372                Box::pin(async move {
373                    transaction
374                        .execute_unprepared("CREATE TABLE example (id INTEGER PRIMARY KEY)")
375                        .await?;
376                    Ok(())
377                })
378            },
379        )
380        .await
381        .unwrap();
382
383        database
384            .execute_unprepared("INSERT INTO example (id) VALUES (1)")
385            .await
386            .unwrap();
387    }
388
389    #[tokio::test]
390    async fn sqlite_callback_error_rolls_back_and_is_preserved() {
391        let database = Database::connect("sqlite::memory:").await.unwrap();
392        let error = with_migration_lock(
393            &database,
394            &MigrationLockOptions::new("sqlite-error"),
395            |transaction| {
396                Box::pin(async move {
397                    transaction
398                        .execute_unprepared("CREATE TABLE rolled_back (id INTEGER PRIMARY KEY)")
399                        .await?;
400                    Err::<(), _>(DbErr::Custom("product migration failed".to_string()))
401                })
402            },
403        )
404        .await
405        .unwrap_err();
406
407        assert_eq!(error.to_string(), "Custom Error: product migration failed");
408        assert!(
409            database
410                .execute_unprepared("INSERT INTO rolled_back (id) VALUES (1)")
411                .await
412                .is_err()
413        );
414    }
415}
aster_forge_db_migration/coordination.rs

aster_forge_db_migration/
coordination.rs
