sea_orm_migration/
connection.rs

1use sea_orm::{
2    AccessMode, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr,
3    ExecResult, IsolationLevel, QueryResult, Statement, TransactionError, TransactionTrait,
4};
5use std::future::Future;
6use std::pin::Pin;
7
8pub enum SchemaManagerConnection<'c> {
9    Connection(&'c DatabaseConnection),
10    Transaction(&'c DatabaseTransaction),
11}
12
13#[async_trait::async_trait]
14impl ConnectionTrait for SchemaManagerConnection<'_> {
15    fn get_database_backend(&self) -> DbBackend {
16        match self {
17            SchemaManagerConnection::Connection(conn) => conn.get_database_backend(),
18            SchemaManagerConnection::Transaction(trans) => trans.get_database_backend(),
19        }
20    }
21
22    async fn execute(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
23        match self {
24            SchemaManagerConnection::Connection(conn) => conn.execute(stmt).await,
25            SchemaManagerConnection::Transaction(trans) => trans.execute(stmt).await,
26        }
27    }
28
29    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
30        match self {
31            SchemaManagerConnection::Connection(conn) => conn.execute_unprepared(sql).await,
32            SchemaManagerConnection::Transaction(trans) => trans.execute_unprepared(sql).await,
33        }
34    }
35
36    async fn query_one(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
37        match self {
38            SchemaManagerConnection::Connection(conn) => conn.query_one(stmt).await,
39            SchemaManagerConnection::Transaction(trans) => trans.query_one(stmt).await,
40        }
41    }
42
43    async fn query_all(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
44        match self {
45            SchemaManagerConnection::Connection(conn) => conn.query_all(stmt).await,
46            SchemaManagerConnection::Transaction(trans) => trans.query_all(stmt).await,
47        }
48    }
49
50    fn is_mock_connection(&self) -> bool {
51        match self {
52            SchemaManagerConnection::Connection(conn) => conn.is_mock_connection(),
53            SchemaManagerConnection::Transaction(trans) => trans.is_mock_connection(),
54        }
55    }
56}
57
58#[async_trait::async_trait]
59impl TransactionTrait for SchemaManagerConnection<'_> {
60    async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
61        match self {
62            SchemaManagerConnection::Connection(conn) => conn.begin().await,
63            SchemaManagerConnection::Transaction(trans) => trans.begin().await,
64        }
65    }
66
67    async fn begin_with_config(
68        &self,
69        isolation_level: Option<IsolationLevel>,
70        access_mode: Option<AccessMode>,
71    ) -> Result<DatabaseTransaction, DbErr> {
72        match self {
73            SchemaManagerConnection::Connection(conn) => {
74                conn.begin_with_config(isolation_level, access_mode).await
75            }
76            SchemaManagerConnection::Transaction(trans) => {
77                trans.begin_with_config(isolation_level, access_mode).await
78            }
79        }
80    }
81
82    async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
83    where
84        F: for<'a> FnOnce(
85                &'a DatabaseTransaction,
86            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>
87            + Send,
88        T: Send,
89        E: std::fmt::Display + std::fmt::Debug + Send,
90    {
91        match self {
92            SchemaManagerConnection::Connection(conn) => conn.transaction(callback).await,
93            SchemaManagerConnection::Transaction(trans) => trans.transaction(callback).await,
94        }
95    }
96
97    async fn transaction_with_config<F, T, E>(
98        &self,
99        callback: F,
100        isolation_level: Option<IsolationLevel>,
101        access_mode: Option<AccessMode>,
102    ) -> Result<T, TransactionError<E>>
103    where
104        F: for<'a> FnOnce(
105                &'a DatabaseTransaction,
106            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>
107            + Send,
108        T: Send,
109        E: std::fmt::Display + std::fmt::Debug + Send,
110    {
111        match self {
112            SchemaManagerConnection::Connection(conn) => {
113                conn.transaction_with_config(callback, isolation_level, access_mode)
114                    .await
115            }
116            SchemaManagerConnection::Transaction(trans) => {
117                trans
118                    .transaction_with_config(callback, isolation_level, access_mode)
119                    .await
120            }
121        }
122    }
123}
124
125pub trait IntoSchemaManagerConnection<'c>: Send
126where
127    Self: 'c,
128{
129    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c>;
130}
131
132impl<'c> IntoSchemaManagerConnection<'c> for SchemaManagerConnection<'c> {
133    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
134        self
135    }
136}
137
138impl<'c> IntoSchemaManagerConnection<'c> for &'c DatabaseConnection {
139    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
140        SchemaManagerConnection::Connection(self)
141    }
142}
143
144impl<'c> IntoSchemaManagerConnection<'c> for &'c DatabaseTransaction {
145    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
146        SchemaManagerConnection::Transaction(self)
147    }
148}
sea_orm_migration/connection.rs

sea_orm_migration/
connection.rs
