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
This commit is contained in:
Billy Chan 2023-01-30 22:55:46 +08:00 committed by GitHub
parent 1b24d42bec
commit 6d7bcb35e5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 583 additions and 190 deletions

View File

@ -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"] }

View File

@ -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;

View File

@ -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<T>(conn: T) -> Self
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 {
self.conn
pub fn get_connection(&self) -> &SchemaManagerConnection<'c> {
&self.conn
}
}

View File

@ -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)
}
}

View File

@ -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<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,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<C>(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::<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(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);
} 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<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(db: &DbConn, mut steps: Option<u32>) -> 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<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
}
}
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<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
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::<M>(manager, None).await
}
async fn exec_up<M>(manager: &SchemaManager<'_>, mut steps: Option<u32>) -> 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<M>(manager: &SchemaManager<'_>, mut steps: Option<u32>) -> 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<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)

View File

@ -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,
};

View File

@ -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?;

View File

@ -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(())
}
}

View File

@ -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,
}

View File

@ -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 {}
}

View File

@ -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),
]
}
}