Skip to main content

sea_orm/database/
executor.rs

1use crate::{
2    AccessMode, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr,
3    ExecResult, IsolationLevel, QueryResult, Statement, TransactionError, TransactionOptions,
4    TransactionTrait,
5};
6use crate::{Schema, SchemaBuilder};
7use std::future::Future;
8use std::pin::Pin;
9
10/// A wrapper that holds either a reference to a [`DatabaseConnection`] or [`DatabaseTransaction`],
11/// or an owned [`DatabaseTransaction`].
12#[derive(Debug)]
13pub enum DatabaseExecutor<'c> {
14    /// A reference to a database connection
15    Connection(&'c DatabaseConnection),
16    /// A reference to a database transaction
17    Transaction(&'c DatabaseTransaction),
18    /// An owned database transaction (used by migration's `SchemaManager::begin()`)
19    OwnedTransaction(DatabaseTransaction),
20}
21
22impl<'c> From<&'c DatabaseConnection> for DatabaseExecutor<'c> {
23    fn from(conn: &'c DatabaseConnection) -> Self {
24        Self::Connection(conn)
25    }
26}
27
28impl<'c> From<&'c DatabaseTransaction> for DatabaseExecutor<'c> {
29    fn from(trans: &'c DatabaseTransaction) -> Self {
30        Self::Transaction(trans)
31    }
32}
33
34impl ConnectionTrait for DatabaseExecutor<'_> {
35    fn get_database_backend(&self) -> DbBackend {
36        match self {
37            DatabaseExecutor::Connection(conn) => conn.get_database_backend(),
38            DatabaseExecutor::Transaction(trans) => trans.get_database_backend(),
39            DatabaseExecutor::OwnedTransaction(trans) => trans.get_database_backend(),
40        }
41    }
42
43    fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
44        match self {
45            DatabaseExecutor::Connection(conn) => conn.execute_raw(stmt),
46            DatabaseExecutor::Transaction(trans) => trans.execute_raw(stmt),
47            DatabaseExecutor::OwnedTransaction(trans) => trans.execute_raw(stmt),
48        }
49    }
50
51    fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
52        match self {
53            DatabaseExecutor::Connection(conn) => conn.execute_unprepared(sql),
54            DatabaseExecutor::Transaction(trans) => trans.execute_unprepared(sql),
55            DatabaseExecutor::OwnedTransaction(trans) => trans.execute_unprepared(sql),
56        }
57    }
58
59    fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
60        match self {
61            DatabaseExecutor::Connection(conn) => conn.query_one_raw(stmt),
62            DatabaseExecutor::Transaction(trans) => trans.query_one_raw(stmt),
63            DatabaseExecutor::OwnedTransaction(trans) => trans.query_one_raw(stmt),
64        }
65    }
66
67    fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
68        match self {
69            DatabaseExecutor::Connection(conn) => conn.query_all_raw(stmt),
70            DatabaseExecutor::Transaction(trans) => trans.query_all_raw(stmt),
71            DatabaseExecutor::OwnedTransaction(trans) => trans.query_all_raw(stmt),
72        }
73    }
74}
75
76impl TransactionTrait for DatabaseExecutor<'_> {
77    type Transaction = DatabaseTransaction;
78
79    fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
80        match self {
81            DatabaseExecutor::Connection(conn) => conn.begin(),
82            DatabaseExecutor::Transaction(trans) => trans.begin(),
83            DatabaseExecutor::OwnedTransaction(trans) => trans.begin(),
84        }
85    }
86
87    fn begin_with_config(
88        &self,
89        isolation_level: Option<IsolationLevel>,
90        access_mode: Option<AccessMode>,
91    ) -> Result<DatabaseTransaction, DbErr> {
92        match self {
93            DatabaseExecutor::Connection(conn) => {
94                conn.begin_with_config(isolation_level, access_mode)
95            }
96            DatabaseExecutor::Transaction(trans) => {
97                trans.begin_with_config(isolation_level, access_mode)
98            }
99            DatabaseExecutor::OwnedTransaction(trans) => {
100                trans.begin_with_config(isolation_level, access_mode)
101            }
102        }
103    }
104
105    fn begin_with_options(
106        &self,
107        options: TransactionOptions,
108    ) -> Result<DatabaseTransaction, DbErr> {
109        match self {
110            DatabaseExecutor::Connection(conn) => conn.begin_with_options(options),
111            DatabaseExecutor::Transaction(trans) => trans.begin_with_options(options),
112            DatabaseExecutor::OwnedTransaction(trans) => trans.begin_with_options(options),
113        }
114    }
115
116    fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
117    where
118        F: for<'c> FnOnce(&'c DatabaseTransaction) -> Result<T, E>,
119        E: std::fmt::Display + std::fmt::Debug,
120    {
121        match self {
122            DatabaseExecutor::Connection(conn) => conn.transaction(callback),
123            DatabaseExecutor::Transaction(trans) => trans.transaction(callback),
124            DatabaseExecutor::OwnedTransaction(trans) => trans.transaction(callback),
125        }
126    }
127
128    fn transaction_with_config<F, T, E>(
129        &self,
130        callback: F,
131        isolation_level: Option<IsolationLevel>,
132        access_mode: Option<AccessMode>,
133    ) -> Result<T, TransactionError<E>>
134    where
135        F: for<'c> FnOnce(&'c DatabaseTransaction) -> Result<T, E>,
136        E: std::fmt::Display + std::fmt::Debug,
137    {
138        match self {
139            DatabaseExecutor::Connection(conn) => {
140                conn.transaction_with_config(callback, isolation_level, access_mode)
141            }
142            DatabaseExecutor::Transaction(trans) => {
143                trans.transaction_with_config(callback, isolation_level, access_mode)
144            }
145            DatabaseExecutor::OwnedTransaction(trans) => {
146                trans.transaction_with_config(callback, isolation_level, access_mode)
147            }
148        }
149    }
150}
151
152/// A trait for converting into [`DatabaseExecutor`]
153pub trait IntoDatabaseExecutor<'c>
154where
155    Self: 'c,
156{
157    /// Convert into a [`DatabaseExecutor`]
158    fn into_database_executor(self) -> DatabaseExecutor<'c>;
159}
160
161impl<'c> IntoDatabaseExecutor<'c> for DatabaseExecutor<'c> {
162    fn into_database_executor(self) -> DatabaseExecutor<'c> {
163        self
164    }
165}
166
167impl<'c> IntoDatabaseExecutor<'c> for &'c DatabaseConnection {
168    fn into_database_executor(self) -> DatabaseExecutor<'c> {
169        DatabaseExecutor::Connection(self)
170    }
171}
172
173impl<'c> IntoDatabaseExecutor<'c> for &'c DatabaseTransaction {
174    fn into_database_executor(self) -> DatabaseExecutor<'c> {
175        DatabaseExecutor::Transaction(self)
176    }
177}
178
179impl IntoDatabaseExecutor<'static> for DatabaseTransaction {
180    fn into_database_executor(self) -> DatabaseExecutor<'static> {
181        DatabaseExecutor::OwnedTransaction(self)
182    }
183}
184
185impl DatabaseExecutor<'_> {
186    /// Returns `true` if this executor is backed by a transaction (borrowed or owned).
187    pub fn is_transaction(&self) -> bool {
188        matches!(
189            self,
190            DatabaseExecutor::Transaction(_) | DatabaseExecutor::OwnedTransaction(_)
191        )
192    }
193
194    /// Creates a [`SchemaBuilder`] for this backend
195    pub fn get_schema_builder(&self) -> SchemaBuilder {
196        Schema::new(self.get_database_backend()).builder()
197    }
198
199    #[cfg(feature = "entity-registry")]
200    #[cfg_attr(docsrs, doc(cfg(feature = "entity-registry")))]
201    /// Builds a schema for all the entities in the given module
202    pub fn get_schema_registry(&self, prefix: &str) -> SchemaBuilder {
203        let schema = Schema::new(self.get_database_backend());
204        crate::EntityRegistry::build_schema(schema, prefix)
205    }
206}