Atomic migration

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

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