Atomic migration (#1379)

* Running atomic migration

* refactor

* clippy

* visibility

* Execute migration in transaction when it's a Postgres connection

* Testing rolling back on migrate up and down

* lifetimes elision

* typo

* abort migration before committing nested transaction

sea-orm-migration/Cargo.toml

@@ -28,6 +28,7 @@ sea-orm-cli = { version = "0.10.3", path = "../sea-orm-cli", default-features =
 sea-schema = { version = "0.11" }
 tracing = { version = "0.1", default-features = false, features = ["log"] }
 tracing-subscriber = { version = "0.3", default-features = false, features = ["env-filter", "fmt"] }
+futures = { version = "0.3", default-features = false, features = ["std"] }
 
 [dev-dependencies]
 async-std = { version = "1", features = ["attributes", "tokio1"] }

sea-orm-migration/src/lib.rs

@@ -1,12 +1,14 @@
 #[cfg(feature = "cli")]
 pub mod cli;
 pub mod manager;
+pub mod manager_connection;
 pub mod migrator;
 pub mod prelude;
 pub mod seaql_migrations;
 pub mod util;
 
 pub use manager::*;
+pub use manager_connection::*;
 pub use migrator::*;
 
 pub use async_trait;

sea-orm-migration/src/manager.rs

@@ -1,20 +1,26 @@
+use super::{IntoSchemaManagerConnection, SchemaManagerConnection};
 use sea_orm::sea_query::{
     extension::postgres::{TypeAlterStatement, TypeCreateStatement, TypeDropStatement},
     ForeignKeyCreateStatement, ForeignKeyDropStatement, IndexCreateStatement, IndexDropStatement,
     TableAlterStatement, TableCreateStatement, TableDropStatement, TableRenameStatement,
     TableTruncateStatement,
 };
-use sea_orm::{ConnectionTrait, DbBackend, DbConn, DbErr, StatementBuilder};
+use sea_orm::{ConnectionTrait, DbBackend, DbErr, StatementBuilder};
 use sea_schema::{mysql::MySql, postgres::Postgres, probe::SchemaProbe, sqlite::Sqlite};
 
 /// Helper struct for writing migration scripts in migration file
 pub struct SchemaManager<'c> {
-    conn: &'c DbConn,
+    conn: SchemaManagerConnection<'c>,
 }
 
 impl<'c> SchemaManager<'c> {
-    pub fn new(conn: &'c DbConn) -> Self {
+    pub fn new<T>(conn: T) -> Self
-        Self { conn }
+    where
+        T: IntoSchemaManagerConnection<'c>,
+    {
+        Self {
+            conn: conn.into_schema_manager_connection(),
+        }
     }
 
     pub async fn exec_stmt<S>(&self, stmt: S) -> Result<(), DbErr>
@@ -29,8 +35,8 @@ impl<'c> SchemaManager<'c> {
         self.conn.get_database_backend()
     }
 
-    pub fn get_connection(&self) -> &'c DbConn {
+    pub fn get_connection(&self) -> &SchemaManagerConnection<'c> {
-        self.conn
+        &self.conn
     }
 }
 

sea-orm-migration/src/manager_connection.rs

@@ -0,0 +1,145 @@
+use futures::Future;
+use sea_orm::{
+    AccessMode, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr,
+    ExecResult, IsolationLevel, QueryResult, Statement, TransactionError, TransactionTrait,
+};
+use std::pin::Pin;
+
+pub enum SchemaManagerConnection<'c> {
+    Connection(&'c DatabaseConnection),
+    Transaction(&'c DatabaseTransaction),
+}
+
+#[async_trait::async_trait]
+impl<'c> ConnectionTrait for SchemaManagerConnection<'c> {
+    fn get_database_backend(&self) -> DbBackend {
+        match self {
+            SchemaManagerConnection::Connection(conn) => conn.get_database_backend(),
+            SchemaManagerConnection::Transaction(trans) => trans.get_database_backend(),
+        }
+    }
+
+    async fn execute(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
+        match self {
+            SchemaManagerConnection::Connection(conn) => conn.execute(stmt).await,
+            SchemaManagerConnection::Transaction(trans) => trans.execute(stmt).await,
+        }
+    }
+
+    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
+        match self {
+            SchemaManagerConnection::Connection(conn) => conn.execute_unprepared(sql).await,
+            SchemaManagerConnection::Transaction(trans) => trans.execute_unprepared(sql).await,
+        }
+    }
+
+    async fn query_one(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
+        match self {
+            SchemaManagerConnection::Connection(conn) => conn.query_one(stmt).await,
+            SchemaManagerConnection::Transaction(trans) => trans.query_one(stmt).await,
+        }
+    }
+
+    async fn query_all(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
+        match self {
+            SchemaManagerConnection::Connection(conn) => conn.query_all(stmt).await,
+            SchemaManagerConnection::Transaction(trans) => trans.query_all(stmt).await,
+        }
+    }
+
+    fn is_mock_connection(&self) -> bool {
+        match self {
+            SchemaManagerConnection::Connection(conn) => conn.is_mock_connection(),
+            SchemaManagerConnection::Transaction(trans) => trans.is_mock_connection(),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<'c> TransactionTrait for SchemaManagerConnection<'c> {
+    async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
+        match self {
+            SchemaManagerConnection::Connection(conn) => conn.begin().await,
+            SchemaManagerConnection::Transaction(trans) => trans.begin().await,
+        }
+    }
+
+    async fn begin_with_config(
+        &self,
+        isolation_level: Option<IsolationLevel>,
+        access_mode: Option<AccessMode>,
+    ) -> Result<DatabaseTransaction, DbErr> {
+        match self {
+            SchemaManagerConnection::Connection(conn) => {
+                conn.begin_with_config(isolation_level, access_mode).await
+            }
+            SchemaManagerConnection::Transaction(trans) => {
+                trans.begin_with_config(isolation_level, access_mode).await
+            }
+        }
+    }
+
+    async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
+    where
+        F: for<'a> FnOnce(
+                &'a DatabaseTransaction,
+            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>
+            + Send,
+        T: Send,
+        E: std::error::Error + Send,
+    {
+        match self {
+            SchemaManagerConnection::Connection(conn) => conn.transaction(callback).await,
+            SchemaManagerConnection::Transaction(trans) => trans.transaction(callback).await,
+        }
+    }
+
+    async fn transaction_with_config<F, T, E>(
+        &self,
+        callback: F,
+        isolation_level: Option<IsolationLevel>,
+        access_mode: Option<AccessMode>,
+    ) -> Result<T, TransactionError<E>>
+    where
+        F: for<'a> FnOnce(
+                &'a DatabaseTransaction,
+            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>
+            + Send,
+        T: Send,
+        E: std::error::Error + Send,
+    {
+        match self {
+            SchemaManagerConnection::Connection(conn) => {
+                conn.transaction_with_config(callback, isolation_level, access_mode)
+                    .await
+            }
+            SchemaManagerConnection::Transaction(trans) => {
+                trans
+                    .transaction_with_config(callback, isolation_level, access_mode)
+                    .await
+            }
+        }
+    }
+}
+
+pub trait IntoSchemaManagerConnection<'c>: Send {
+    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c>;
+}
+
+impl<'c> IntoSchemaManagerConnection<'c> for SchemaManagerConnection<'c> {
+    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
+        self
+    }
+}
+
+impl<'c> IntoSchemaManagerConnection<'c> for &'c DatabaseConnection {
+    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
+        SchemaManagerConnection::Connection(self)
+    }
+}
+
+impl<'c> IntoSchemaManagerConnection<'c> for &'c DatabaseTransaction {
+    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
+        SchemaManagerConnection::Transaction(self)
+    }
+}

sea-orm-migration/src/migrator.rs

@@ -1,5 +1,7 @@
+use futures::Future;
 use std::collections::HashSet;
 use std::fmt::Display;
+use std::pin::Pin;
 use std::time::SystemTime;
 use tracing::info;
 
@@ -8,12 +10,12 @@ use sea_orm::sea_query::{
     SelectStatement, SimpleExpr, Table,
 };
 use sea_orm::{
-    ActiveModelTrait, ActiveValue, ColumnTrait, Condition, ConnectionTrait, DbBackend, DbConn,
+    ActiveModelTrait, ActiveValue, ColumnTrait, Condition, ConnectionTrait, DbBackend, DbErr,
-    DbErr, EntityTrait, QueryFilter, QueryOrder, Schema, Statement,
+    EntityTrait, QueryFilter, QueryOrder, Schema, Statement, TransactionTrait,
 };
 use sea_schema::{mysql::MySql, postgres::Postgres, probe::SchemaProbe, sqlite::Sqlite};
 
-use super::{seaql_migrations, MigrationTrait, SchemaManager};
+use super::{seaql_migrations, IntoSchemaManagerConnection, MigrationTrait, SchemaManager};
 
 #[derive(Debug, PartialEq, Eq)]
 /// Status of migration
@@ -57,7 +59,10 @@ pub trait MigratorTrait: Send {
     }
 
     /// Get list of applied migrations from database
-    async fn get_migration_models(db: &DbConn) -> Result<Vec<seaql_migrations::Model>, DbErr> {
+    async fn get_migration_models<C>(db: &C) -> Result<Vec<seaql_migrations::Model>, DbErr>
+    where
+        C: ConnectionTrait,
+    {
         Self::install(db).await?;
         seaql_migrations::Entity::find()
             .order_by_asc(seaql_migrations::Column::Version)
@@ -66,7 +71,10 @@ pub trait MigratorTrait: Send {
     }
 
     /// Get list of migrations with status
-    async fn get_migration_with_status(db: &DbConn) -> Result<Vec<Migration>, DbErr> {
+    async fn get_migration_with_status<C>(db: &C) -> Result<Vec<Migration>, DbErr>
+    where
+        C: ConnectionTrait,
+    {
         Self::install(db).await?;
         let mut migration_files = Self::get_migration_files();
         let migration_models = Self::get_migration_models(db).await?;
@@ -102,7 +110,10 @@ pub trait MigratorTrait: Send {
     }
 
     /// Get list of pending migrations
-    async fn get_pending_migrations(db: &DbConn) -> Result<Vec<Migration>, DbErr> {
+    async fn get_pending_migrations<C>(db: &C) -> Result<Vec<Migration>, DbErr>
+    where
+        C: ConnectionTrait,
+    {
         Self::install(db).await?;
         Ok(Self::get_migration_with_status(db)
             .await?
@@ -112,7 +123,10 @@ pub trait MigratorTrait: Send {
     }
 
     /// Get list of applied migrations
-    async fn get_applied_migrations(db: &DbConn) -> Result<Vec<Migration>, DbErr> {
+    async fn get_applied_migrations<C>(db: &C) -> Result<Vec<Migration>, DbErr>
+    where
+        C: ConnectionTrait,
+    {
         Self::install(db).await?;
         Ok(Self::get_migration_with_status(db)
             .await?
@@ -122,7 +136,10 @@ pub trait MigratorTrait: Send {
     }
 
     /// Create migration table `seaql_migrations` in the database
-    async fn install(db: &DbConn) -> Result<(), DbErr> {
+    async fn install<C>(db: &C) -> Result<(), DbErr>
+    where
+        C: ConnectionTrait,
+    {
         let builder = db.get_database_backend();
         let schema = Schema::new(builder);
         let mut stmt = schema.create_table_from_entity(seaql_migrations::Entity);
@@ -130,9 +147,112 @@ pub trait MigratorTrait: Send {
         db.execute(builder.build(&stmt)).await.map(|_| ())
     }
 
-    /// Drop all tables from the database, then reapply all migrations
+    /// Check the status of all migrations
-    async fn fresh(db: &DbConn) -> Result<(), DbErr> {
+    async fn status<C>(db: &C) -> Result<(), DbErr>
+    where
+        C: ConnectionTrait,
+    {
         Self::install(db).await?;
+
+        info!("Checking migration status");
+
+        for Migration { migration, status } in Self::get_migration_with_status(db).await? {
+            info!("Migration '{}'... {}", migration.name(), status);
+        }
+
+        Ok(())
+    }
+
+    /// Drop all tables from the database, then reapply all migrations
+    async fn fresh<'c, C>(db: C) -> Result<(), DbErr>
+    where
+        C: IntoSchemaManagerConnection<'c>,
+    {
+        exec_with_connection::<'_, _, _, Self>(db, move |manager| {
+            Box::pin(async move { exec_fresh::<Self>(manager).await })
+        })
+        .await
+    }
+
+    /// Rollback all applied migrations, then reapply all migrations
+    async fn refresh<'c, C>(db: C) -> Result<(), DbErr>
+    where
+        C: IntoSchemaManagerConnection<'c>,
+    {
+        exec_with_connection::<'_, _, _, Self>(db, move |manager| {
+            Box::pin(async move {
+                exec_down::<Self>(manager, None).await?;
+                exec_up::<Self>(manager, None).await
+            })
+        })
+        .await
+    }
+
+    /// Rollback all applied migrations
+    async fn reset<'c, C>(db: C) -> Result<(), DbErr>
+    where
+        C: IntoSchemaManagerConnection<'c>,
+    {
+        exec_with_connection::<'_, _, _, Self>(db, move |manager| {
+            Box::pin(async move { exec_down::<Self>(manager, None).await })
+        })
+        .await
+    }
+
+    /// Apply pending migrations
+    async fn up<'c, C>(db: C, steps: Option<u32>) -> Result<(), DbErr>
+    where
+        C: IntoSchemaManagerConnection<'c>,
+    {
+        exec_with_connection::<'_, _, _, Self>(db, move |manager| {
+            Box::pin(async move { exec_up::<Self>(manager, steps).await })
+        })
+        .await
+    }
+
+    /// Rollback applied migrations
+    async fn down<'c, C>(db: C, steps: Option<u32>) -> Result<(), DbErr>
+    where
+        C: IntoSchemaManagerConnection<'c>,
+    {
+        exec_with_connection::<'_, _, _, Self>(db, move |manager| {
+            Box::pin(async move { exec_down::<Self>(manager, steps).await })
+        })
+        .await
+    }
+}
+
+async fn exec_with_connection<'c, C, F, M>(db: C, f: F) -> Result<(), DbErr>
+where
+    C: IntoSchemaManagerConnection<'c>,
+    F: for<'b> Fn(
+        &'b SchemaManager<'_>,
+    ) -> Pin<Box<dyn Future<Output = Result<(), DbErr>> + Send + 'b>>,
+    M: MigratorTrait + ?Sized,
+{
+    let db = db.into_schema_manager_connection();
+
+    match db.get_database_backend() {
+        DbBackend::Postgres => {
+            let transaction = db.begin().await?;
+            let manager = SchemaManager::new(&transaction);
+            f(&manager).await?;
+            transaction.commit().await
+        }
+        DbBackend::MySql | DbBackend::Sqlite => {
+            let manager = SchemaManager::new(db);
+            f(&manager).await
+        }
+    }
+}
+
+async fn exec_fresh<M>(manager: &SchemaManager<'_>) -> Result<(), DbErr>
+where
+    M: MigratorTrait + ?Sized,
+{
+    let db = manager.get_connection();
+
+    M::install(db).await?;
     let db_backend = db.get_database_backend();
 
     // Temporarily disable the foreign key check
@@ -208,37 +328,16 @@ pub trait MigratorTrait: Send {
     }
 
     // Reapply all migrations
-        Self::up(db, None).await
+    exec_up::<M>(manager, None).await
 }
 
-    /// Rollback all applied migrations, then reapply all migrations
+async fn exec_up<M>(manager: &SchemaManager<'_>, mut steps: Option<u32>) -> Result<(), DbErr>
-    async fn refresh(db: &DbConn) -> Result<(), DbErr> {
+where
-        Self::down(db, None).await?;
+    M: MigratorTrait + ?Sized,
-        Self::up(db, None).await
+{
-    }
+    let db = manager.get_connection();
 
-    /// Rollback all applied migrations
+    M::install(db).await?;
-    async fn reset(db: &DbConn) -> Result<(), DbErr> {
-        Self::down(db, None).await
-    }
-
-    /// Check the status of all migrations
-    async fn status(db: &DbConn) -> Result<(), DbErr> {
-        Self::install(db).await?;
-
-        info!("Checking migration status");
-
-        for Migration { migration, status } in Self::get_migration_with_status(db).await? {
-            info!("Migration '{}'... {}", migration.name(), status);
-        }
-
-        Ok(())
-    }
-
-    /// Apply pending migrations
-    async fn up(db: &DbConn, mut steps: Option<u32>) -> Result<(), DbErr> {
-        Self::install(db).await?;
-        let manager = SchemaManager::new(db);
 
     if let Some(steps) = steps {
         info!("Applying {} pending migrations", steps);
@@ -246,7 +345,7 @@ pub trait MigratorTrait: Send {
         info!("Applying all pending migrations");
     }
 
-        let migrations = Self::get_pending_migrations(db).await?.into_iter();
+    let migrations = M::get_pending_migrations(db).await?.into_iter();
     if migrations.len() == 0 {
         info!("No pending migrations");
     }
@@ -258,7 +357,7 @@ pub trait MigratorTrait: Send {
             *steps -= 1;
         }
         info!("Applying migration '{}'", migration.name());
-            migration.up(&manager).await?;
+        migration.up(manager).await?;
         info!("Migration '{}' has been applied", migration.name());
         let now = SystemTime::now()
             .duration_since(SystemTime::UNIX_EPOCH)
@@ -274,10 +373,13 @@ pub trait MigratorTrait: Send {
     Ok(())
 }
 
-    /// Rollback applied migrations
+async fn exec_down<M>(manager: &SchemaManager<'_>, mut steps: Option<u32>) -> Result<(), DbErr>
-    async fn down(db: &DbConn, mut steps: Option<u32>) -> Result<(), DbErr> {
+where
-        Self::install(db).await?;
+    M: MigratorTrait + ?Sized,
-        let manager = SchemaManager::new(db);
+{
+    let db = manager.get_connection();
+
+    M::install(db).await?;
 
     if let Some(steps) = steps {
         info!("Rolling back {} applied migrations", steps);
@@ -285,7 +387,7 @@ pub trait MigratorTrait: Send {
         info!("Rolling back all applied migrations");
     }
 
-        let migrations = Self::get_applied_migrations(db).await?.into_iter().rev();
+    let migrations = M::get_applied_migrations(db).await?.into_iter().rev();
     if migrations.len() == 0 {
         info!("No applied migrations");
     }
@@ -297,7 +399,7 @@ pub trait MigratorTrait: Send {
             *steps -= 1;
         }
         info!("Rolling back migration '{}'", migration.name());
-            migration.down(&manager).await?;
+        migration.down(manager).await?;
         info!("Migration '{}' has been rollbacked", migration.name());
         seaql_migrations::Entity::delete_many()
             .filter(seaql_migrations::Column::Version.eq(migration.name()))
@@ -307,9 +409,11 @@ pub trait MigratorTrait: Send {
 
     Ok(())
 }
-}
 
-pub(crate) fn query_tables(db: &DbConn) -> SelectStatement {
+fn query_tables<C>(db: &C) -> SelectStatement
+where
+    C: ConnectionTrait,
+{
     match db.get_database_backend() {
         DbBackend::MySql => MySql::query_tables(),
         DbBackend::Postgres => Postgres::query_tables(),
@@ -317,7 +421,10 @@ pub(crate) fn query_tables(db: &DbConn) -> SelectStatement {
     }
 }
 
-pub(crate) fn get_current_schema(db: &DbConn) -> SimpleExpr {
+fn get_current_schema<C>(db: &C) -> SimpleExpr
+where
+    C: ConnectionTrait,
+{
     match db.get_database_backend() {
         DbBackend::MySql => MySql::get_current_schema(),
         DbBackend::Postgres => Postgres::get_current_schema(),
@@ -338,7 +445,10 @@ enum InformationSchema {
     ConstraintType,
 }
 
-fn query_mysql_foreign_keys(db: &DbConn) -> SelectStatement {
+fn query_mysql_foreign_keys<C>(db: &C) -> SelectStatement
+where
+    C: ConnectionTrait,
+{
     let mut stmt = Query::select();
     stmt.columns([
         InformationSchema::TableName,
@@ -380,7 +490,10 @@ enum PgNamespace {
     Nspname,
 }
 
-fn query_pg_types(db: &DbConn) -> SelectStatement {
+fn query_pg_types<C>(db: &C) -> SelectStatement
+where
+    C: ConnectionTrait,
+{
     let mut stmt = Query::select();
     stmt.column(PgType::Typname)
         .from(PgType::Table)

sea-orm-migration/src/prelude.rs

@@ -1,12 +1,13 @@
 #[cfg(feature = "cli")]
-pub use super::cli;
+pub use crate::cli;
 
-pub use super::manager::SchemaManager;
+pub use crate::{
-pub use super::migrator::MigratorTrait;
+    IntoSchemaManagerConnection, MigrationName, MigrationTrait, MigratorTrait, SchemaManager,
-pub use super::{MigrationName, MigrationTrait};
+    SchemaManagerConnection,
+};
 pub use async_trait;
-pub use sea_orm;
+pub use sea_orm::{
-pub use sea_orm::sea_query;
+    self,
-pub use sea_orm::sea_query::*;
+    sea_query::{self, *},
-pub use sea_orm::DbErr;
+    DbErr, DeriveMigrationName,
-pub use sea_orm::DeriveMigrationName;
+};

sea-orm-migration/tests/main.rs

@@ -107,6 +107,33 @@ async fn run_migration(url: &str, db_name: &str, schema: &str) -> Result<(), DbE
     assert!(!manager.has_table("cake").await?);
     assert!(!manager.has_table("fruit").await?);
 
+    // Tests rolling back changes of "migrate up" when running migration on Postgres
+    if matches!(db.get_database_backend(), DbBackend::Postgres) {
+        println!("\nRoll back changes when encounter errors");
+
+        // Set a flag to throw error inside `m20230109_000001_seed_cake_table.rs`
+        std::env::set_var("ABORT_MIGRATION", "YES");
+
+        // Should throw an error
+        println!("\nMigrator::up");
+        assert_eq!(
+            Migrator::up(db, None).await,
+            Err(DbErr::Migration(
+                "Aboard migration and rollback changes".into()
+            ))
+        );
+
+        println!("\nMigrator::status");
+        Migrator::status(db).await?;
+
+        // Check migrations have been rolled back
+        assert!(!manager.has_table("cake").await?);
+        assert!(!manager.has_table("fruit").await?);
+
+        // Unset the flag
+        std::env::remove_var("ABORT_MIGRATION");
+    }
+
     println!("\nMigrator::up");
     Migrator::up(db, None).await?;
 
@@ -119,6 +146,33 @@ async fn run_migration(url: &str, db_name: &str, schema: &str) -> Result<(), DbE
     assert!(manager.has_column("cake", "name").await?);
     assert!(manager.has_column("fruit", "cake_id").await?);
 
+    // Tests rolling back changes of "migrate down" when running migration on Postgres
+    if matches!(db.get_database_backend(), DbBackend::Postgres) {
+        println!("\nRoll back changes when encounter errors");
+
+        // Set a flag to throw error inside `m20230109_000001_seed_cake_table.rs`
+        std::env::set_var("ABORT_MIGRATION", "YES");
+
+        // Should throw an error
+        println!("\nMigrator::down");
+        assert_eq!(
+            Migrator::down(db, None).await,
+            Err(DbErr::Migration(
+                "Aboard migration and rollback changes".into()
+            ))
+        );
+
+        println!("\nMigrator::status");
+        Migrator::status(db).await?;
+
+        // Check migrations have been rolled back
+        assert!(manager.has_table("cake").await?);
+        assert!(manager.has_table("fruit").await?);
+
+        // Unset the flag
+        std::env::remove_var("ABORT_MIGRATION");
+    }
+
     println!("\nMigrator::down");
     Migrator::down(db, None).await?;
 

sea-orm-migration/tests/migrator/m20220118_000001_create_cake_table.rs

@@ -26,7 +26,15 @@ impl MigrationTrait for Migration {
     async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
         manager
             .drop_table(Table::drop().table(Cake::Table).to_owned())
-            .await
+            .await?;
+
+        if std::env::var_os("ABORT_MIGRATION").eq(&Some("YES".into())) {
+            return Err(DbErr::Migration(
+                "Aboard migration and rollback changes".into(),
+            ));
+        }
+
+        Ok(())
     }
 }
 

sea-orm-migration/tests/migrator/m20220923_000001_seed_cake_table.rs

@@ -1,5 +1,4 @@
 use sea_orm_migration::prelude::*;
-use sea_orm_migration::sea_orm::{entity::*, query::*};
 
 #[derive(DeriveMigrationName)]
 pub struct Migration;
@@ -34,6 +33,5 @@ impl MigrationTrait for Migration {
 #[derive(Iden)]
 pub enum Cake {
     Table,
-    Id,
     Name,
 }

sea-orm-migration/tests/migrator/m20230109_000001_seed_cake_table.rs

@@ -0,0 +1,63 @@
+use sea_orm_migration::prelude::*;
+use sea_orm_migration::sea_orm::{entity::*, query::*};
+
+#[derive(DeriveMigrationName)]
+pub struct Migration;
+
+#[async_trait::async_trait]
+impl MigrationTrait for Migration {
+    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        let db = manager.get_connection();
+
+        let transaction = db.begin().await?;
+
+        cake::ActiveModel {
+            name: Set("Cheesecake".to_owned()),
+            ..Default::default()
+        }
+        .insert(&transaction)
+        .await?;
+
+        if std::env::var_os("ABORT_MIGRATION").eq(&Some("YES".into())) {
+            return Err(DbErr::Migration(
+                "Aboard migration and rollback changes".into(),
+            ));
+        }
+
+        transaction.commit().await?;
+
+        Ok(())
+    }
+
+    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        let db = manager.get_connection();
+
+        let transaction = db.begin().await?;
+
+        cake::Entity::delete_many()
+            .filter(cake::Column::Name.eq("Cheesecake"))
+            .exec(&transaction)
+            .await?;
+
+        transaction.commit().await?;
+
+        Ok(())
+    }
+}
+
+mod cake {
+    use sea_orm_migration::sea_orm::entity::prelude::*;
+
+    #[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+    #[sea_orm(table_name = "cake")]
+    pub struct Model {
+        #[sea_orm(primary_key)]
+        pub id: i32,
+        pub name: String,
+    }
+
+    #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+    pub enum Relation {}
+
+    impl ActiveModelBehavior for ActiveModel {}
+}

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

@@ -5,6 +5,7 @@ mod m20220118_000002_create_fruit_table;
 mod m20220118_000003_seed_cake_table;
 mod m20220118_000004_create_tea_enum;
 mod m20220923_000001_seed_cake_table;
+mod m20230109_000001_seed_cake_table;
 
 pub struct Migrator;
 
@@ -17,6 +18,7 @@ impl MigratorTrait for Migrator {
             Box::new(m20220118_000003_seed_cake_table::Migration),
             Box::new(m20220118_000004_create_tea_enum::Migration),
             Box::new(m20220923_000001_seed_cake_table::Migration),
+            Box::new(m20230109_000001_seed_cake_table::Migration),
         ]
     }
 }