sea_orm_migration/
connection.rs

1use sea_orm::{
2    AccessMode, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr,
3    ExecResult, IsolationLevel, QueryResult, Schema, SchemaBuilder, Statement, TransactionError,
4    TransactionTrait,
5};
6use std::future::Future;
7use std::pin::Pin;
8
9pub enum SchemaManagerConnection<'c> {
10    Connection(&'c DatabaseConnection),
11    Transaction(&'c DatabaseTransaction),
12}
13
14#[async_trait::async_trait]
15impl ConnectionTrait for SchemaManagerConnection<'_> {
16    fn get_database_backend(&self) -> DbBackend {
17        match self {
18            SchemaManagerConnection::Connection(conn) => conn.get_database_backend(),
19            SchemaManagerConnection::Transaction(trans) => trans.get_database_backend(),
20        }
21    }
22
23    async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
24        match self {
25            SchemaManagerConnection::Connection(conn) => conn.execute_raw(stmt).await,
26            SchemaManagerConnection::Transaction(trans) => trans.execute_raw(stmt).await,
27        }
28    }
29
30    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
31        match self {
32            SchemaManagerConnection::Connection(conn) => conn.execute_unprepared(sql).await,
33            SchemaManagerConnection::Transaction(trans) => trans.execute_unprepared(sql).await,
34        }
35    }
36
37    async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
38        match self {
39            SchemaManagerConnection::Connection(conn) => conn.query_one_raw(stmt).await,
40            SchemaManagerConnection::Transaction(trans) => trans.query_one_raw(stmt).await,
41        }
42    }
43
44    async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
45        match self {
46            SchemaManagerConnection::Connection(conn) => conn.query_all_raw(stmt).await,
47            SchemaManagerConnection::Transaction(trans) => trans.query_all_raw(stmt).await,
48        }
49    }
50}
51
52#[async_trait::async_trait]
53impl TransactionTrait for SchemaManagerConnection<'_> {
54    type Transaction = DatabaseTransaction;
55
56    async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
57        match self {
58            SchemaManagerConnection::Connection(conn) => conn.begin().await,
59            SchemaManagerConnection::Transaction(trans) => trans.begin().await,
60        }
61    }
62
63    async fn begin_with_config(
64        &self,
65        isolation_level: Option<IsolationLevel>,
66        access_mode: Option<AccessMode>,
67    ) -> Result<DatabaseTransaction, DbErr> {
68        match self {
69            SchemaManagerConnection::Connection(conn) => {
70                conn.begin_with_config(isolation_level, access_mode).await
71            }
72            SchemaManagerConnection::Transaction(trans) => {
73                trans.begin_with_config(isolation_level, access_mode).await
74            }
75        }
76    }
77
78    async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
79    where
80        F: for<'a> FnOnce(
81                &'a DatabaseTransaction,
82            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>
83            + Send,
84        T: Send,
85        E: std::fmt::Display + std::fmt::Debug + Send,
86    {
87        match self {
88            SchemaManagerConnection::Connection(conn) => conn.transaction(callback).await,
89            SchemaManagerConnection::Transaction(trans) => trans.transaction(callback).await,
90        }
91    }
92
93    async fn transaction_with_config<F, T, E>(
94        &self,
95        callback: F,
96        isolation_level: Option<IsolationLevel>,
97        access_mode: Option<AccessMode>,
98    ) -> Result<T, TransactionError<E>>
99    where
100        F: for<'a> FnOnce(
101                &'a DatabaseTransaction,
102            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>
103            + Send,
104        T: Send,
105        E: std::fmt::Display + std::fmt::Debug + Send,
106    {
107        match self {
108            SchemaManagerConnection::Connection(conn) => {
109                conn.transaction_with_config(callback, isolation_level, access_mode)
110                    .await
111            }
112            SchemaManagerConnection::Transaction(trans) => {
113                trans
114                    .transaction_with_config(callback, isolation_level, access_mode)
115                    .await
116            }
117        }
118    }
119}
120
121#[cfg(feature = "sqlx-dep")]
122mod sea_schema_shim {
123    use super::{DatabaseConnection, DatabaseTransaction, SchemaManagerConnection};
124    use sea_orm::sea_query::SelectStatement;
125    use sea_schema::sqlx_types::{SqlxError, SqlxRow};
126
127    #[async_trait::async_trait]
128    impl sea_schema::Connection for SchemaManagerConnection<'_> {
129        async fn query_all(&self, select: SelectStatement) -> Result<Vec<SqlxRow>, SqlxError> {
130            match self {
131                Self::Connection(conn) => {
132                    <DatabaseConnection as sea_schema::Connection>::query_all(conn, select).await
133                }
134                Self::Transaction(txn) => {
135                    <DatabaseTransaction as sea_schema::Connection>::query_all(txn, select).await
136                }
137            }
138        }
139
140        async fn query_all_raw(&self, sql: String) -> Result<Vec<SqlxRow>, SqlxError> {
141            match self {
142                Self::Connection(conn) => {
143                    <DatabaseConnection as sea_schema::Connection>::query_all_raw(conn, sql).await
144                }
145                Self::Transaction(txn) => {
146                    <DatabaseTransaction as sea_schema::Connection>::query_all_raw(txn, sql).await
147                }
148            }
149        }
150    }
151}
152
153impl SchemaManagerConnection<'_> {
154    /// Creates a [`SchemaBuilder`] for this backend
155    pub fn get_schema_builder(&self) -> SchemaBuilder {
156        Schema::new(self.get_database_backend()).builder()
157    }
158
159    #[cfg(feature = "entity-registry")]
160    #[cfg_attr(docsrs, doc(cfg(feature = "entity-registry")))]
161    /// Builds a schema for all the entites in the given module
162    pub fn get_schema_registry(&self, prefix: &str) -> SchemaBuilder {
163        let schema = Schema::new(self.get_database_backend());
164        sea_orm::EntityRegistry::build_schema(schema, prefix)
165    }
166}
167
168pub trait IntoSchemaManagerConnection<'c>: Send
169where
170    Self: 'c,
171{
172    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c>;
173}
174
175impl<'c> IntoSchemaManagerConnection<'c> for SchemaManagerConnection<'c> {
176    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
177        self
178    }
179}
180
181impl<'c> IntoSchemaManagerConnection<'c> for &'c DatabaseConnection {
182    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
183        SchemaManagerConnection::Connection(self)
184    }
185}
186
187impl<'c> IntoSchemaManagerConnection<'c> for &'c DatabaseTransaction {
188    fn into_schema_manager_connection(self) -> SchemaManagerConnection<'c> {
189        SchemaManagerConnection::Transaction(self)
190    }
191}
sea_orm_migration/connection.rs

sea_orm_migration/
connection.rs
