From 6d7bcb35e5d9e1fe2cb8a62578f36332582725bc Mon Sep 17 00:00:00 2001 From: Billy Chan Date: Mon, 30 Jan 2023 22:55:46 +0800 Subject: [PATCH] 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 | 1 + sea-orm-migration/src/lib.rs | 2 + sea-orm-migration/src/manager.rs | 18 +- sea-orm-migration/src/manager_connection.rs | 145 ++++++ sea-orm-migration/src/migrator.rs | 457 +++++++++++------- sea-orm-migration/src/prelude.rs | 19 +- sea-orm-migration/tests/main.rs | 54 +++ .../m20220118_000001_create_cake_table.rs | 10 +- .../m20220923_000001_seed_cake_table.rs | 2 - .../m20230109_000001_seed_cake_table.rs | 63 +++ sea-orm-migration/tests/migrator/mod.rs | 2 + 11 files changed, 583 insertions(+), 190 deletions(-) create mode 100644 sea-orm-migration/src/manager_connection.rs create mode 100644 sea-orm-migration/tests/migrator/m20230109_000001_seed_cake_table.rs diff --git a/sea-orm-migration/Cargo.toml b/sea-orm-migration/Cargo.toml index bb8e5e88..555849f6 100644 --- a/sea-orm-migration/Cargo.toml +++ b/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"] } diff --git a/sea-orm-migration/src/lib.rs b/sea-orm-migration/src/lib.rs index 6a7f8547..38b0dd1b 100644 --- a/sea-orm-migration/src/lib.rs +++ b/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; diff --git a/sea-orm-migration/src/manager.rs b/sea-orm-migration/src/manager.rs index 71e91b37..f70c9212 100644 --- a/sea-orm-migration/src/manager.rs +++ b/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 { - Self { conn } + pub fn new(conn: T) -> Self + where + T: IntoSchemaManagerConnection<'c>, + { + Self { + conn: conn.into_schema_manager_connection(), + } } pub async fn exec_stmt(&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 { - self.conn + pub fn get_connection(&self) -> &SchemaManagerConnection<'c> { + &self.conn } } diff --git a/sea-orm-migration/src/manager_connection.rs b/sea-orm-migration/src/manager_connection.rs new file mode 100644 index 00000000..eda6f628 --- /dev/null +++ b/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 { + match self { + SchemaManagerConnection::Connection(conn) => conn.execute(stmt).await, + SchemaManagerConnection::Transaction(trans) => trans.execute(stmt).await, + } + } + + async fn execute_unprepared(&self, sql: &str) -> Result { + 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, 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, 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 { + match self { + SchemaManagerConnection::Connection(conn) => conn.begin().await, + SchemaManagerConnection::Transaction(trans) => trans.begin().await, + } + } + + async fn begin_with_config( + &self, + isolation_level: Option, + access_mode: Option, + ) -> Result { + 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(&self, callback: F) -> Result> + where + F: for<'a> FnOnce( + &'a DatabaseTransaction, + ) -> Pin> + 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( + &self, + callback: F, + isolation_level: Option, + access_mode: Option, + ) -> Result> + where + F: for<'a> FnOnce( + &'a DatabaseTransaction, + ) -> Pin> + 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) + } +} diff --git a/sea-orm-migration/src/migrator.rs b/sea-orm-migration/src/migrator.rs index 349da99d..c62041b8 100644 --- a/sea-orm-migration/src/migrator.rs +++ b/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, - DbErr, EntityTrait, QueryFilter, QueryOrder, Schema, Statement, + ActiveModelTrait, ActiveValue, ColumnTrait, Condition, ConnectionTrait, DbBackend, DbErr, + 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, DbErr> { + async fn get_migration_models(db: &C) -> Result, 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, DbErr> { + async fn get_migration_with_status(db: &C) -> Result, 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, DbErr> { + async fn get_pending_migrations(db: &C) -> Result, 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, DbErr> { + async fn get_applied_migrations(db: &C) -> Result, 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(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,100 +147,11 @@ pub trait MigratorTrait: Send { db.execute(builder.build(&stmt)).await.map(|_| ()) } - /// Drop all tables from the database, then reapply all migrations - async fn fresh(db: &DbConn) -> Result<(), DbErr> { - Self::install(db).await?; - let db_backend = db.get_database_backend(); - - // Temporarily disable the foreign key check - if db_backend == DbBackend::Sqlite { - info!("Disabling foreign key check"); - db.execute(Statement::from_string( - db_backend, - "PRAGMA foreign_keys = OFF".to_owned(), - )) - .await?; - info!("Foreign key check disabled"); - } - - // Drop all foreign keys - if db_backend == DbBackend::MySql { - info!("Dropping all foreign keys"); - let stmt = query_mysql_foreign_keys(db); - let rows = db.query_all(db_backend.build(&stmt)).await?; - for row in rows.into_iter() { - let constraint_name: String = row.try_get("", "CONSTRAINT_NAME")?; - let table_name: String = row.try_get("", "TABLE_NAME")?; - info!( - "Dropping foreign key '{}' from table '{}'", - constraint_name, table_name - ); - let mut stmt = ForeignKey::drop(); - stmt.table(Alias::new(table_name.as_str())) - .name(constraint_name.as_str()); - db.execute(db_backend.build(&stmt)).await?; - info!("Foreign key '{}' has been dropped", constraint_name); - } - info!("All foreign keys dropped"); - } - - // Drop all tables - let stmt = query_tables(db); - let rows = db.query_all(db_backend.build(&stmt)).await?; - for row in rows.into_iter() { - let table_name: String = row.try_get("", "table_name")?; - info!("Dropping table '{}'", table_name); - let mut stmt = Table::drop(); - stmt.table(Alias::new(table_name.as_str())) - .if_exists() - .cascade(); - db.execute(db_backend.build(&stmt)).await?; - info!("Table '{}' has been dropped", table_name); - } - - // Drop all types - if db_backend == DbBackend::Postgres { - info!("Dropping all types"); - let stmt = query_pg_types(db); - let rows = db.query_all(db_backend.build(&stmt)).await?; - for row in rows { - let type_name: String = row.try_get("", "typname")?; - info!("Dropping type '{}'", type_name); - let mut stmt = Type::drop(); - stmt.name(Alias::new(&type_name as &str)); - db.execute(db_backend.build(&stmt)).await?; - info!("Type '{}' has been dropped", type_name); - } - } - - // Restore the foreign key check - if db_backend == DbBackend::Sqlite { - info!("Restoring foreign key check"); - db.execute(Statement::from_string( - db_backend, - "PRAGMA foreign_keys = ON".to_owned(), - )) - .await?; - info!("Foreign key check restored"); - } - - // Reapply all migrations - Self::up(db, None).await - } - - /// Rollback all applied migrations, then reapply all migrations - async fn refresh(db: &DbConn) -> Result<(), DbErr> { - Self::down(db, None).await?; - Self::up(db, None).await - } - - /// Rollback all applied migrations - 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> { + async fn status(db: &C) -> Result<(), DbErr> + where + C: ConnectionTrait, + { Self::install(db).await?; info!("Checking migration status"); @@ -235,81 +163,257 @@ pub trait MigratorTrait: Send { 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::(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::(manager, None).await?; + exec_up::(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::(manager, None).await }) + }) + .await + } + /// Apply pending migrations - async fn up(db: &DbConn, mut steps: Option) -> Result<(), DbErr> { - Self::install(db).await?; - let manager = SchemaManager::new(db); - - if let Some(steps) = steps { - info!("Applying {} pending migrations", steps); - } else { - info!("Applying all pending migrations"); - } - - let migrations = Self::get_pending_migrations(db).await?.into_iter(); - if migrations.len() == 0 { - info!("No pending migrations"); - } - for Migration { migration, .. } in migrations { - if let Some(steps) = steps.as_mut() { - if steps == &0 { - break; - } - *steps -= 1; - } - info!("Applying migration '{}'", migration.name()); - migration.up(&manager).await?; - info!("Migration '{}' has been applied", migration.name()); - let now = SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .expect("SystemTime before UNIX EPOCH!"); - seaql_migrations::ActiveModel { - version: ActiveValue::Set(migration.name().to_owned()), - applied_at: ActiveValue::Set(now.as_secs() as i64), - } - .insert(db) - .await?; - } - - Ok(()) + async fn up<'c, C>(db: C, steps: Option) -> Result<(), DbErr> + where + C: IntoSchemaManagerConnection<'c>, + { + exec_with_connection::<'_, _, _, Self>(db, move |manager| { + Box::pin(async move { exec_up::(manager, steps).await }) + }) + .await } /// Rollback applied migrations - async fn down(db: &DbConn, mut steps: Option) -> Result<(), DbErr> { - Self::install(db).await?; - let manager = SchemaManager::new(db); - - if let Some(steps) = steps { - info!("Rolling back {} applied migrations", steps); - } else { - info!("Rolling back all applied migrations"); - } - - let migrations = Self::get_applied_migrations(db).await?.into_iter().rev(); - if migrations.len() == 0 { - info!("No applied migrations"); - } - for Migration { migration, .. } in migrations { - if let Some(steps) = steps.as_mut() { - if steps == &0 { - break; - } - *steps -= 1; - } - info!("Rolling back migration '{}'", migration.name()); - migration.down(&manager).await?; - info!("Migration '{}' has been rollbacked", migration.name()); - seaql_migrations::Entity::delete_many() - .filter(seaql_migrations::Column::Version.eq(migration.name())) - .exec(db) - .await?; - } - - Ok(()) + async fn down<'c, C>(db: C, steps: Option) -> Result<(), DbErr> + where + C: IntoSchemaManagerConnection<'c>, + { + exec_with_connection::<'_, _, _, Self>(db, move |manager| { + Box::pin(async move { exec_down::(manager, steps).await }) + }) + .await } } -pub(crate) fn query_tables(db: &DbConn) -> SelectStatement { +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> + 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(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 + if db_backend == DbBackend::Sqlite { + info!("Disabling foreign key check"); + db.execute(Statement::from_string( + db_backend, + "PRAGMA foreign_keys = OFF".to_owned(), + )) + .await?; + info!("Foreign key check disabled"); + } + + // Drop all foreign keys + if db_backend == DbBackend::MySql { + info!("Dropping all foreign keys"); + let stmt = query_mysql_foreign_keys(db); + let rows = db.query_all(db_backend.build(&stmt)).await?; + for row in rows.into_iter() { + let constraint_name: String = row.try_get("", "CONSTRAINT_NAME")?; + let table_name: String = row.try_get("", "TABLE_NAME")?; + info!( + "Dropping foreign key '{}' from table '{}'", + constraint_name, table_name + ); + let mut stmt = ForeignKey::drop(); + stmt.table(Alias::new(table_name.as_str())) + .name(constraint_name.as_str()); + db.execute(db_backend.build(&stmt)).await?; + info!("Foreign key '{}' has been dropped", constraint_name); + } + info!("All foreign keys dropped"); + } + + // Drop all tables + let stmt = query_tables(db); + let rows = db.query_all(db_backend.build(&stmt)).await?; + for row in rows.into_iter() { + let table_name: String = row.try_get("", "table_name")?; + info!("Dropping table '{}'", table_name); + let mut stmt = Table::drop(); + stmt.table(Alias::new(table_name.as_str())) + .if_exists() + .cascade(); + db.execute(db_backend.build(&stmt)).await?; + info!("Table '{}' has been dropped", table_name); + } + + // Drop all types + if db_backend == DbBackend::Postgres { + info!("Dropping all types"); + let stmt = query_pg_types(db); + let rows = db.query_all(db_backend.build(&stmt)).await?; + for row in rows { + let type_name: String = row.try_get("", "typname")?; + info!("Dropping type '{}'", type_name); + let mut stmt = Type::drop(); + stmt.name(Alias::new(&type_name as &str)); + db.execute(db_backend.build(&stmt)).await?; + info!("Type '{}' has been dropped", type_name); + } + } + + // Restore the foreign key check + if db_backend == DbBackend::Sqlite { + info!("Restoring foreign key check"); + db.execute(Statement::from_string( + db_backend, + "PRAGMA foreign_keys = ON".to_owned(), + )) + .await?; + info!("Foreign key check restored"); + } + + // Reapply all migrations + exec_up::(manager, None).await +} + +async fn exec_up(manager: &SchemaManager<'_>, mut steps: Option) -> Result<(), DbErr> +where + M: MigratorTrait + ?Sized, +{ + let db = manager.get_connection(); + + M::install(db).await?; + + if let Some(steps) = steps { + info!("Applying {} pending migrations", steps); + } else { + info!("Applying all pending migrations"); + } + + let migrations = M::get_pending_migrations(db).await?.into_iter(); + if migrations.len() == 0 { + info!("No pending migrations"); + } + for Migration { migration, .. } in migrations { + if let Some(steps) = steps.as_mut() { + if steps == &0 { + break; + } + *steps -= 1; + } + info!("Applying migration '{}'", migration.name()); + migration.up(manager).await?; + info!("Migration '{}' has been applied", migration.name()); + let now = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .expect("SystemTime before UNIX EPOCH!"); + seaql_migrations::ActiveModel { + version: ActiveValue::Set(migration.name().to_owned()), + applied_at: ActiveValue::Set(now.as_secs() as i64), + } + .insert(db) + .await?; + } + + Ok(()) +} + +async fn exec_down(manager: &SchemaManager<'_>, mut steps: Option) -> Result<(), DbErr> +where + M: MigratorTrait + ?Sized, +{ + let db = manager.get_connection(); + + M::install(db).await?; + + if let Some(steps) = steps { + info!("Rolling back {} applied migrations", steps); + } else { + info!("Rolling back all applied migrations"); + } + + let migrations = M::get_applied_migrations(db).await?.into_iter().rev(); + if migrations.len() == 0 { + info!("No applied migrations"); + } + for Migration { migration, .. } in migrations { + if let Some(steps) = steps.as_mut() { + if steps == &0 { + break; + } + *steps -= 1; + } + info!("Rolling back migration '{}'", migration.name()); + migration.down(manager).await?; + info!("Migration '{}' has been rollbacked", migration.name()); + seaql_migrations::Entity::delete_many() + .filter(seaql_migrations::Column::Version.eq(migration.name())) + .exec(db) + .await?; + } + + Ok(()) +} + +fn query_tables(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(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(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(db: &C) -> SelectStatement +where + C: ConnectionTrait, +{ let mut stmt = Query::select(); stmt.column(PgType::Typname) .from(PgType::Table) diff --git a/sea-orm-migration/src/prelude.rs b/sea-orm-migration/src/prelude.rs index fd78117e..25bf9d5a 100644 --- a/sea-orm-migration/src/prelude.rs +++ b/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 super::migrator::MigratorTrait; -pub use super::{MigrationName, MigrationTrait}; +pub use crate::{ + IntoSchemaManagerConnection, MigrationName, MigrationTrait, MigratorTrait, SchemaManager, + SchemaManagerConnection, +}; pub use async_trait; -pub use sea_orm; -pub use sea_orm::sea_query; -pub use sea_orm::sea_query::*; -pub use sea_orm::DbErr; -pub use sea_orm::DeriveMigrationName; +pub use sea_orm::{ + self, + sea_query::{self, *}, + DbErr, DeriveMigrationName, +}; diff --git a/sea-orm-migration/tests/main.rs b/sea-orm-migration/tests/main.rs index 82af5717..2279f7ec 100644 --- a/sea-orm-migration/tests/main.rs +++ b/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?; diff --git a/sea-orm-migration/tests/migrator/m20220118_000001_create_cake_table.rs b/sea-orm-migration/tests/migrator/m20220118_000001_create_cake_table.rs index 3ad76eab..926ba146 100644 --- a/sea-orm-migration/tests/migrator/m20220118_000001_create_cake_table.rs +++ b/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(()) } } diff --git a/sea-orm-migration/tests/migrator/m20220923_000001_seed_cake_table.rs b/sea-orm-migration/tests/migrator/m20220923_000001_seed_cake_table.rs index ece73e1a..16c019bf 100644 --- a/sea-orm-migration/tests/migrator/m20220923_000001_seed_cake_table.rs +++ b/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, } diff --git a/sea-orm-migration/tests/migrator/m20230109_000001_seed_cake_table.rs b/sea-orm-migration/tests/migrator/m20230109_000001_seed_cake_table.rs new file mode 100644 index 00000000..91314f5e --- /dev/null +++ b/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 {} +} diff --git a/sea-orm-migration/tests/migrator/mod.rs b/sea-orm-migration/tests/migrator/mod.rs index 042e5323..7e7f6802 100644 --- a/sea-orm-migration/tests/migrator/mod.rs +++ b/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), ] } }