Add per-migration transaction control (#2980)

* Add per-migration transaction control for SeaORM migrations

+ Add MigrationTrait::use_transaction() (None/Some(true)/Some(false))
  and SchemaManager::begin()/commit() for manual control
+ Change up()/down() from single batch transaction to per-migration
+ Add DatabaseExecutor::OwnedTransaction variant and is_transaction() method

* Add one more test case

Author: tyt2y3

sea-orm-migration/Cargo.toml

@@ -39,7 +39,7 @@ tracing-subscriber = { version = "0.3.17", default-features = false, features =
 ] }
 
 [dev-dependencies]
-async-std = { version = "1", features = ["attributes", "tokio1"] }
+tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
 
 [features]
 cli = ["clap", "dotenvy", "sea-orm-cli/cli"]

sea-orm-migration/src/lib.rs

@@ -31,4 +31,13 @@ pub trait MigrationTrait: MigrationName + Send + Sync {
     async fn down(&self, _manager: &SchemaManager) -> Result<(), DbErr> {
         Err(DbErr::Migration("We Don't Do That Here".to_owned()))
     }
+
+    /// Control whether this migration runs inside a transaction.
+    ///
+    /// - `None` (default): follow backend convention (Postgres = transaction, MySQL/SQLite = no transaction)
+    /// - `Some(true)`: force wrapping in a transaction on any backend
+    /// - `Some(false)`: disable automatic transaction wrapping (use `manager.begin()` for manual control)
+    fn use_transaction(&self) -> Option<bool> {
+        None
+    }
 }

sea-orm-migration/src/manager.rs

@@ -5,7 +5,7 @@ use sea_orm::sea_query::{
     TableRenameStatement, TableTruncateStatement,
     extension::postgres::{TypeAlterStatement, TypeCreateStatement, TypeDropStatement},
 };
-use sea_orm::{ConnectionTrait, DbBackend, DbErr, StatementBuilder};
+use sea_orm::{ConnectionTrait, DbBackend, DbErr, StatementBuilder, TransactionTrait};
 #[allow(unused_imports)]
 use sea_schema::probe::SchemaProbe;
 
@@ -48,6 +48,30 @@ impl<'c> SchemaManager<'c> {
     }
 }
 
+/// Transaction Control
+impl SchemaManager<'_> {
+    /// Begin a new transaction, returning an owned `SchemaManager` backed by it.
+    ///
+    /// Useful in migrations with `use_transaction() -> Some(false)` for manual
+    /// transaction management (e.g., separating DDL and DML into distinct transactions).
+    pub async fn begin(&self) -> Result<SchemaManager<'static>, DbErr> {
+        let txn = self.conn.begin().await?;
+        Ok(SchemaManager {
+            conn: SchemaManagerConnection::OwnedTransaction(txn),
+        })
+    }
+
+    /// Commit the owned transaction. Only valid on a `SchemaManager` created by [`begin()`](Self::begin).
+    pub async fn commit(self) -> Result<(), DbErr> {
+        match self.conn {
+            SchemaManagerConnection::OwnedTransaction(txn) => txn.commit().await,
+            _ => Err(DbErr::Custom(
+                "Cannot commit: SchemaManager does not own a transaction".into(),
+            )),
+        }
+    }
+}
+
 /// Schema Creation
 impl SchemaManager<'_> {
     pub async fn create_table(&self, stmt: TableCreateStatement) -> Result<(), DbErr> {

sea-orm-migration/src/migrator.rs

@@ -194,21 +194,19 @@ pub trait MigratorTrait: Send {
     where
         C: IntoSchemaManagerConnection<'c>,
     {
-        exec_with_connection!(db, async |manager| {
-            exec_up::<Self>(manager, steps).await
-        })
-        .await
+        let db = db.into_database_executor();
+        let manager = SchemaManager::new(db);
+        exec_up::<Self>(&manager, steps).await
     }
 
     /// Rollback applied migrations
     async fn down<'c, C>(db: C, steps: Option<u32>) -> Result<(), DbErr>
     where
         C: IntoSchemaManagerConnection<'c>,
     {
-        exec_with_connection!(db, async |manager| {
-            exec_down::<Self>(manager, steps).await
-        })
-        .await
+        let db = db.into_database_executor();
+        let manager = SchemaManager::new(db);
+        exec_down::<Self>(&manager, steps).await
     }
 }
 

sea-orm-migration/src/migrator/exec.rs

@@ -10,7 +10,7 @@ use sea_orm::sea_query::{
 };
 use sea_orm::{
     ActiveValue, ConnectionTrait, DbBackend, DbErr, DynIden, EntityTrait, FromQueryResult,
-    Iterable, QueryFilter, Schema, Statement,
+    Iterable, QueryFilter, Schema, Statement, TransactionTrait,
 };
 
 pub async fn get_migration_models<C>(
@@ -198,6 +198,48 @@ pub async fn drop_everything<C: ConnectionTrait>(db: &C) -> Result<(), DbErr> {
     Ok(())
 }
 
+fn should_use_transaction(migration: &dyn crate::MigrationTrait, backend: DbBackend) -> bool {
+    match migration.use_transaction() {
+        Some(v) => v,
+        None => backend == DbBackend::Postgres,
+    }
+}
+
+async fn insert_migration_record<C: ConnectionTrait>(
+    db: &C,
+    name: &str,
+    migration_table_name: DynIden,
+) -> Result<(), DbErr> {
+    #[cfg(not(feature = "with-time"))]
+    let applied_at = SystemTime::now()
+        .duration_since(SystemTime::UNIX_EPOCH)
+        .expect("SystemTime before UNIX EPOCH!")
+        .as_secs() as i64;
+    #[cfg(feature = "with-time")]
+    let applied_at = sea_orm::prelude::TimeDateTimeWithTimeZone::now_utc().unix_timestamp();
+    seaql_migrations::Entity::insert(seaql_migrations::ActiveModel {
+        version: ActiveValue::Set(name.to_owned()),
+        applied_at: ActiveValue::Set(applied_at),
+    })
+    .table_name(migration_table_name)
+    .exec(db)
+    .await?;
+    Ok(())
+}
+
+async fn delete_migration_record<C: ConnectionTrait>(
+    db: &C,
+    name: &str,
+    migration_table_name: DynIden,
+) -> Result<(), DbErr> {
+    seaql_migrations::Entity::delete_many()
+        .filter(Expr::col(seaql_migrations::Column::Version).eq(name))
+        .table_name(migration_table_name)
+        .exec(db)
+        .await?;
+    Ok(())
+}
+
 pub async fn exec_up_with(
     manager: &SchemaManager<'_>,
     mut steps: Option<u32>,
@@ -222,23 +264,23 @@ pub async fn exec_up_with(
             }
             *steps -= 1;
         }
+
+        let use_txn = should_use_transaction(migration.as_ref(), db.get_database_backend());
         info!("Applying migration '{}'", migration.name());
-        migration.up(manager).await?;
-        info!("Migration '{}' has been applied", migration.name());
-        #[cfg(not(feature = "with-time"))]
-        let applied_at = SystemTime::now()
-            .duration_since(SystemTime::UNIX_EPOCH)
-            .expect("SystemTime before UNIX EPOCH!")
-            .as_secs() as i64;
-        #[cfg(feature = "with-time")]
-        let applied_at = sea_orm::prelude::TimeDateTimeWithTimeZone::now_utc().unix_timestamp();
-        seaql_migrations::Entity::insert(seaql_migrations::ActiveModel {
-            version: ActiveValue::Set(migration.name().to_owned()),
-            applied_at: ActiveValue::Set(applied_at),
-        })
-        .table_name(migration_table_name.clone())
-        .exec(db)
-        .await?;
+
+        if use_txn {
+            let transaction = db.begin().await?;
+            let txn_manager = SchemaManager::new(&transaction);
+            migration.up(&txn_manager).await?;
+            info!("Migration '{}' has been applied", migration.name());
+            insert_migration_record(&transaction, migration.name(), migration_table_name.clone())
+                .await?;
+            transaction.commit().await?;
+        } else {
+            migration.up(manager).await?;
+            info!("Migration '{}' has been applied", migration.name());
+            insert_migration_record(db, migration.name(), migration_table_name.clone()).await?;
+        }
     }
 
     Ok(())
@@ -268,14 +310,23 @@ pub async fn exec_down_with(
             }
             *steps -= 1;
         }
+
+        let use_txn = should_use_transaction(migration.as_ref(), db.get_database_backend());
         info!("Rolling back migration '{}'", migration.name());
-        migration.down(manager).await?;
-        info!("Migration '{}' has been rollbacked", migration.name());
-        seaql_migrations::Entity::delete_many()
-            .filter(Expr::col(seaql_migrations::Column::Version).eq(migration.name()))
-            .table_name(migration_table_name.clone())
-            .exec(db)
-            .await?;
+
+        if use_txn {
+            let transaction = db.begin().await?;
+            let txn_manager = SchemaManager::new(&transaction);
+            migration.down(&txn_manager).await?;
+            info!("Migration '{}' has been rolled back", migration.name());
+            delete_migration_record(&transaction, migration.name(), migration_table_name.clone())
+                .await?;
+            transaction.commit().await?;
+        } else {
+            migration.down(manager).await?;
+            info!("Migration '{}' has been rolled back", migration.name());
+            delete_migration_record(db, migration.name(), migration_table_name.clone()).await?;
+        }
     }
 
     Ok(())

sea-orm-migration/src/migrator/with_self.rs

@@ -152,18 +152,19 @@ pub trait MigratorTraitSelf: Sized + Send + Sync {
     where
         C: IntoSchemaManagerConnection<'c>,
     {
-        exec_with_connection!(db, async |manager| { exec_up(self, manager, steps).await }).await
+        let db = db.into_database_executor();
+        let manager = SchemaManager::new(db);
+        exec_up(self, &manager, steps).await
     }
 
     /// Rollback applied migrations
     async fn down<'c, C>(&self, db: C, steps: Option<u32>) -> Result<(), DbErr>
     where
         C: IntoSchemaManagerConnection<'c>,
     {
-        exec_with_connection!(db, async |manager| {
-            exec_down(self, manager, steps).await
-        })
-        .await
+        let db = db.into_database_executor();
+        let manager = SchemaManager::new(db);
+        exec_down(self, &manager, steps).await
     }
 }
 

sea-orm-migration/tests/common/migration/m20250101_000001_create_test_table.rs

@@ -0,0 +1,63 @@
+use sea_orm_migration::prelude::*;
+use sea_orm_migration::schema::*;
+use sea_orm_migration::sea_orm::DbBackend;
+
+pub struct Migration {
+    pub use_transaction: Option<bool>,
+    pub should_fail: bool,
+}
+
+impl MigrationName for Migration {
+    fn name(&self) -> &str {
+        "m20250101_000001_create_test_table"
+    }
+}
+
+#[async_trait::async_trait]
+impl MigrationTrait for Migration {
+    fn use_transaction(&self) -> Option<bool> {
+        self.use_transaction
+    }
+
+    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        let expect_txn = self
+            .use_transaction
+            .unwrap_or(manager.get_database_backend() == DbBackend::Postgres);
+        assert_eq!(
+            manager.get_connection().is_transaction(),
+            expect_txn,
+            "up: expected is_transaction() = {expect_txn}"
+        );
+
+        manager
+            .create_table(
+                Table::create()
+                    .table("test_table")
+                    .col(pk_auto("id"))
+                    .col(string("name"))
+                    .to_owned(),
+            )
+            .await?;
+
+        if self.should_fail {
+            return Err(DbErr::Migration("intentional failure".into()));
+        }
+
+        Ok(())
+    }
+
+    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        let expect_txn = self
+            .use_transaction
+            .unwrap_or(manager.get_database_backend() == DbBackend::Postgres);
+        assert_eq!(
+            manager.get_connection().is_transaction(),
+            expect_txn,
+            "down: expected is_transaction() = {expect_txn}"
+        );
+
+        manager
+            .drop_table(Table::drop().table("test_table").to_owned())
+            .await
+    }
+}

sea-orm-migration/tests/common/migration/m20250101_000002_manual_transaction.rs

@@ -0,0 +1,50 @@
+use sea_orm_migration::prelude::*;
+use sea_orm_migration::schema::*;
+
+pub struct Migration;
+
+impl MigrationName for Migration {
+    fn name(&self) -> &str {
+        "m20250101_000002_manual_transaction"
+    }
+}
+
+#[async_trait::async_trait]
+impl MigrationTrait for Migration {
+    fn use_transaction(&self) -> Option<bool> {
+        Some(false)
+    }
+
+    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        assert!(
+            !manager.get_connection().is_transaction(),
+            "outer manager should not be in a transaction"
+        );
+
+        let m = manager.begin().await?;
+        assert!(
+            m.get_connection().is_transaction(),
+            "inner manager should be in a transaction"
+        );
+        m.create_table(
+            Table::create()
+                .table("manual_txn_table")
+                .col(pk_auto("id"))
+                .col(string("name"))
+                .to_owned(),
+        )
+        .await?;
+        m.commit().await?;
+
+        Ok(())
+    }
+
+    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        let m = manager.begin().await?;
+        m.drop_table(Table::drop().table("manual_txn_table").to_owned())
+            .await?;
+        m.commit().await?;
+
+        Ok(())
+    }
+}

sea-orm-migration/tests/common/migration/mod.rs

@@ -4,3 +4,5 @@ pub mod m20220118_000003_seed_cake_table;
 pub mod m20220118_000004_create_tea_enum;
 pub mod m20220923_000001_seed_cake_table;
 pub mod m20230109_000001_seed_cake_table;
+pub mod m20250101_000001_create_test_table;
+pub mod m20250101_000002_manual_transaction;

sea-orm-migration/tests/common/migrator/mod.rs

@@ -1,3 +1,4 @@
 pub mod default;
 pub mod override_migration_table_name;
+pub mod transaction_test;
 pub mod with_self;