Skip to main content

sea_orm/database/
executor.rs

1use super::transaction::run_async_transaction_callback;
2use crate::{
3    AccessMode, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr,
4    ExecResult, IsolationLevel, QueryResult, Statement, TransactionError, TransactionOptions,
5    TransactionTrait,
6};
7use crate::{Schema, SchemaBuilder};
8use std::future::Future;
9use std::pin::Pin;
10
11/// Either a borrowed [`DatabaseConnection`] / [`DatabaseTransaction`], or an
12/// owned [`DatabaseTransaction`].
13///
14/// Implements [`ConnectionTrait`] and [`TransactionTrait`], so APIs that
15/// need to accept "any of those three" can take a `DatabaseExecutor` and
16/// not worry about which variant the caller had. Used in particular by
17/// `sea-orm-migration`'s `SchemaManager`.
18#[derive(Debug)]
19pub enum DatabaseExecutor<'c> {
20    /// Borrowed connection — use against a long-lived pool.
21    Connection(&'c DatabaseConnection),
22    /// Borrowed transaction — caller still owns the transaction handle.
23    Transaction(&'c DatabaseTransaction),
24    /// Owned transaction — used by migration's `SchemaManager::begin()` so
25    /// the transaction can be committed/rolled back at the end of the call.
26    OwnedTransaction(DatabaseTransaction),
27}
28
29impl<'c> From<&'c DatabaseConnection> for DatabaseExecutor<'c> {
30    fn from(conn: &'c DatabaseConnection) -> Self {
31        Self::Connection(conn)
32    }
33}
34
35impl<'c> From<&'c DatabaseTransaction> for DatabaseExecutor<'c> {
36    fn from(trans: &'c DatabaseTransaction) -> Self {
37        Self::Transaction(trans)
38    }
39}
40
41impl ConnectionTrait for DatabaseExecutor<'_> {
42    fn get_database_backend(&self) -> DbBackend {
43        match self {
44            DatabaseExecutor::Connection(conn) => conn.get_database_backend(),
45            DatabaseExecutor::Transaction(trans) => trans.get_database_backend(),
46            DatabaseExecutor::OwnedTransaction(trans) => trans.get_database_backend(),
47        }
48    }
49
50    fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
51        match self {
52            DatabaseExecutor::Connection(conn) => conn.execute_raw(stmt),
53            DatabaseExecutor::Transaction(trans) => trans.execute_raw(stmt),
54            DatabaseExecutor::OwnedTransaction(trans) => trans.execute_raw(stmt),
55        }
56    }
57
58    fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
59        match self {
60            DatabaseExecutor::Connection(conn) => conn.execute_unprepared(sql),
61            DatabaseExecutor::Transaction(trans) => trans.execute_unprepared(sql),
62            DatabaseExecutor::OwnedTransaction(trans) => trans.execute_unprepared(sql),
63        }
64    }
65
66    fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
67        match self {
68            DatabaseExecutor::Connection(conn) => conn.query_one_raw(stmt),
69            DatabaseExecutor::Transaction(trans) => trans.query_one_raw(stmt),
70            DatabaseExecutor::OwnedTransaction(trans) => trans.query_one_raw(stmt),
71        }
72    }
73
74    fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
75        match self {
76            DatabaseExecutor::Connection(conn) => conn.query_all_raw(stmt),
77            DatabaseExecutor::Transaction(trans) => trans.query_all_raw(stmt),
78            DatabaseExecutor::OwnedTransaction(trans) => trans.query_all_raw(stmt),
79        }
80    }
81}
82
83impl TransactionTrait for DatabaseExecutor<'_> {
84    type Transaction = DatabaseTransaction;
85
86    fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
87        match self {
88            DatabaseExecutor::Connection(conn) => conn.begin(),
89            DatabaseExecutor::Transaction(trans) => trans.begin(),
90            DatabaseExecutor::OwnedTransaction(trans) => trans.begin(),
91        }
92    }
93
94    fn begin_with_config(
95        &self,
96        isolation_level: Option<IsolationLevel>,
97        access_mode: Option<AccessMode>,
98    ) -> Result<DatabaseTransaction, DbErr> {
99        match self {
100            DatabaseExecutor::Connection(conn) => {
101                conn.begin_with_config(isolation_level, access_mode)
102            }
103            DatabaseExecutor::Transaction(trans) => {
104                trans.begin_with_config(isolation_level, access_mode)
105            }
106            DatabaseExecutor::OwnedTransaction(trans) => {
107                trans.begin_with_config(isolation_level, access_mode)
108            }
109        }
110    }
111
112    fn begin_with_options(
113        &self,
114        options: TransactionOptions,
115    ) -> Result<DatabaseTransaction, DbErr> {
116        match self {
117            DatabaseExecutor::Connection(conn) => conn.begin_with_options(options),
118            DatabaseExecutor::Transaction(trans) => trans.begin_with_options(options),
119            DatabaseExecutor::OwnedTransaction(trans) => trans.begin_with_options(options),
120        }
121    }
122
123    fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
124    where
125        F: for<'c> FnOnce(&'c DatabaseTransaction) -> Result<T, E>,
126        E: std::fmt::Display + std::fmt::Debug,
127    {
128        match self {
129            DatabaseExecutor::Connection(conn) => conn.transaction(callback),
130            DatabaseExecutor::Transaction(trans) => trans.transaction(callback),
131            DatabaseExecutor::OwnedTransaction(trans) => trans.transaction(callback),
132        }
133    }
134
135    fn transaction_with_config<F, T, E>(
136        &self,
137        callback: F,
138        isolation_level: Option<IsolationLevel>,
139        access_mode: Option<AccessMode>,
140    ) -> Result<T, TransactionError<E>>
141    where
142        F: for<'c> FnOnce(&'c DatabaseTransaction) -> Result<T, E>,
143        E: std::fmt::Display + std::fmt::Debug,
144    {
145        match self {
146            DatabaseExecutor::Connection(conn) => {
147                conn.transaction_with_config(callback, isolation_level, access_mode)
148            }
149            DatabaseExecutor::Transaction(trans) => {
150                trans.transaction_with_config(callback, isolation_level, access_mode)
151            }
152            DatabaseExecutor::OwnedTransaction(trans) => {
153                trans.transaction_with_config(callback, isolation_level, access_mode)
154            }
155        }
156    }
157}
158
159/// Conversion into a [`DatabaseExecutor`]. Implemented for
160/// `&DatabaseConnection`, `&DatabaseTransaction`, and owned
161/// `DatabaseTransaction` — let users hand any of them to functions that
162/// take `impl IntoDatabaseExecutor<'_>`.
163pub trait IntoDatabaseExecutor<'c>
164where
165    Self: 'c,
166{
167    /// Build the [`DatabaseExecutor`].
168    fn into_database_executor(self) -> DatabaseExecutor<'c>;
169}
170
171impl<'c> IntoDatabaseExecutor<'c> for DatabaseExecutor<'c> {
172    fn into_database_executor(self) -> DatabaseExecutor<'c> {
173        self
174    }
175}
176
177impl<'c> IntoDatabaseExecutor<'c> for &'c DatabaseConnection {
178    fn into_database_executor(self) -> DatabaseExecutor<'c> {
179        DatabaseExecutor::Connection(self)
180    }
181}
182
183impl<'c> IntoDatabaseExecutor<'c> for &'c DatabaseTransaction {
184    fn into_database_executor(self) -> DatabaseExecutor<'c> {
185        DatabaseExecutor::Transaction(self)
186    }
187}
188
189impl IntoDatabaseExecutor<'static> for DatabaseTransaction {
190    fn into_database_executor(self) -> DatabaseExecutor<'static> {
191        DatabaseExecutor::OwnedTransaction(self)
192    }
193}
194
195impl DatabaseExecutor<'_> {
196    /// Execute the function inside a transaction.
197    /// If the function returns an error, the transaction will be rolled back.
198    /// Otherwise, the transaction will be committed.
199    pub fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
200    where
201        F: for<'c> FnOnce(&'c DatabaseTransaction) -> Result<T, E>,
202        E: std::fmt::Display + std::fmt::Debug,
203    {
204        let transaction = self.begin().map_err(TransactionError::Connection)?;
205        run_async_transaction_callback(transaction, callback)
206    }
207
208    /// Execute the function inside a transaction with isolation level and/or access mode.
209    /// If the function returns an error, the transaction will be rolled back.
210    /// Otherwise, the transaction will be committed.
211    pub fn transaction_with_config<F, T, E>(
212        &self,
213        callback: F,
214        isolation_level: Option<IsolationLevel>,
215        access_mode: Option<AccessMode>,
216    ) -> Result<T, TransactionError<E>>
217    where
218        F: for<'c> FnOnce(&'c DatabaseTransaction) -> Result<T, E>,
219        E: std::fmt::Display + std::fmt::Debug,
220    {
221        let transaction = self
222            .begin_with_config(isolation_level, access_mode)
223            .map_err(TransactionError::Connection)?;
224        run_async_transaction_callback(transaction, callback)
225    }
226
227    /// Returns `true` if this executor is backed by a transaction (borrowed or owned).
228    pub fn is_transaction(&self) -> bool {
229        matches!(
230            self,
231            DatabaseExecutor::Transaction(_) | DatabaseExecutor::OwnedTransaction(_)
232        )
233    }
234
235    /// Creates a [`SchemaBuilder`] for this backend
236    pub fn get_schema_builder(&self) -> SchemaBuilder {
237        Schema::new(self.get_database_backend()).builder()
238    }
239
240    #[cfg(feature = "entity-registry")]
241    #[cfg_attr(docsrs, doc(cfg(feature = "entity-registry")))]
242    /// Builds a schema for all the entities in the given module
243    pub fn get_schema_registry(&self, prefix: &str) -> SchemaBuilder {
244        let schema = Schema::new(self.get_database_backend());
245        crate::EntityRegistry::build_schema(schema, prefix)
246    }
247}