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#[derive(Debug)]
12pub enum DatabaseExecutor<'c> {
13    /// A reference to a database connection
14    Connection(&'c DatabaseConnection),
15    /// A reference to a database transaction
16    Transaction(&'c DatabaseTransaction),
17}
18
19impl<'c> From<&'c DatabaseConnection> for DatabaseExecutor<'c> {
20    fn from(conn: &'c DatabaseConnection) -> Self {
21        Self::Connection(conn)
22    }
23}
24
25impl<'c> From<&'c DatabaseTransaction> for DatabaseExecutor<'c> {
26    fn from(trans: &'c DatabaseTransaction) -> Self {
27        Self::Transaction(trans)
28    }
29}
30
31impl ConnectionTrait for DatabaseExecutor<'_> {
32    fn get_database_backend(&self) -> DbBackend {
33        match self {
34            DatabaseExecutor::Connection(conn) => conn.get_database_backend(),
35            DatabaseExecutor::Transaction(trans) => trans.get_database_backend(),
36        }
37    }
38
39    fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
40        match self {
41            DatabaseExecutor::Connection(conn) => conn.execute_raw(stmt),
42            DatabaseExecutor::Transaction(trans) => trans.execute_raw(stmt),
43        }
44    }
45
46    fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
47        match self {
48            DatabaseExecutor::Connection(conn) => conn.execute_unprepared(sql),
49            DatabaseExecutor::Transaction(trans) => trans.execute_unprepared(sql),
50        }
51    }
52
53    fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
54        match self {
55            DatabaseExecutor::Connection(conn) => conn.query_one_raw(stmt),
56            DatabaseExecutor::Transaction(trans) => trans.query_one_raw(stmt),
57        }
58    }
59
60    fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
61        match self {
62            DatabaseExecutor::Connection(conn) => conn.query_all_raw(stmt),
63            DatabaseExecutor::Transaction(trans) => trans.query_all_raw(stmt),
64        }
65    }
66}
67
68impl TransactionTrait for DatabaseExecutor<'_> {
69    type Transaction = DatabaseTransaction;
70
71    fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
72        match self {
73            DatabaseExecutor::Connection(conn) => conn.begin(),
74            DatabaseExecutor::Transaction(trans) => trans.begin(),
75        }
76    }
77
78    fn begin_with_config(
79        &self,
80        isolation_level: Option<IsolationLevel>,
81        access_mode: Option<AccessMode>,
82    ) -> Result<DatabaseTransaction, DbErr> {
83        match self {
84            DatabaseExecutor::Connection(conn) => {
85                conn.begin_with_config(isolation_level, access_mode)
86            }
87            DatabaseExecutor::Transaction(trans) => {
88                trans.begin_with_config(isolation_level, access_mode)
89            }
90        }
91    }
92
93    fn begin_with_options(
94        &self,
95        options: TransactionOptions,
96    ) -> Result<DatabaseTransaction, DbErr> {
97        match self {
98            DatabaseExecutor::Connection(conn) => conn.begin_with_options(options),
99            DatabaseExecutor::Transaction(trans) => trans.begin_with_options(options),
100        }
101    }
102
103    fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
104    where
105        F: for<'c> FnOnce(&'c DatabaseTransaction) -> Result<T, E>,
106        E: std::fmt::Display + std::fmt::Debug,
107    {
108        match self {
109            DatabaseExecutor::Connection(conn) => conn.transaction(callback),
110            DatabaseExecutor::Transaction(trans) => trans.transaction(callback),
111        }
112    }
113
114    fn transaction_with_config<F, T, E>(
115        &self,
116        callback: F,
117        isolation_level: Option<IsolationLevel>,
118        access_mode: Option<AccessMode>,
119    ) -> Result<T, TransactionError<E>>
120    where
121        F: for<'c> FnOnce(&'c DatabaseTransaction) -> Result<T, E>,
122        E: std::fmt::Display + std::fmt::Debug,
123    {
124        match self {
125            DatabaseExecutor::Connection(conn) => {
126                conn.transaction_with_config(callback, isolation_level, access_mode)
127            }
128            DatabaseExecutor::Transaction(trans) => {
129                trans.transaction_with_config(callback, isolation_level, access_mode)
130            }
131        }
132    }
133}
134
135/// A trait for converting into [`DatabaseExecutor`]
136pub trait IntoDatabaseExecutor<'c>
137where
138    Self: 'c,
139{
140    /// Convert into a [`DatabaseExecutor`]
141    fn into_database_executor(self) -> DatabaseExecutor<'c>;
142}
143
144impl<'c> IntoDatabaseExecutor<'c> for DatabaseExecutor<'c> {
145    fn into_database_executor(self) -> DatabaseExecutor<'c> {
146        self
147    }
148}
149
150impl<'c> IntoDatabaseExecutor<'c> for &'c DatabaseConnection {
151    fn into_database_executor(self) -> DatabaseExecutor<'c> {
152        DatabaseExecutor::Connection(self)
153    }
154}
155
156impl<'c> IntoDatabaseExecutor<'c> for &'c DatabaseTransaction {
157    fn into_database_executor(self) -> DatabaseExecutor<'c> {
158        DatabaseExecutor::Transaction(self)
159    }
160}
161
162impl DatabaseExecutor<'_> {
163    /// Creates a [`SchemaBuilder`] for this backend
164    pub fn get_schema_builder(&self) -> SchemaBuilder {
165        Schema::new(self.get_database_backend()).builder()
166    }
167
168    #[cfg(feature = "entity-registry")]
169    #[cfg_attr(docsrs, doc(cfg(feature = "entity-registry")))]
170    /// Builds a schema for all the entities in the given module
171    pub fn get_schema_registry(&self, prefix: &str) -> SchemaBuilder {
172        let schema = Schema::new(self.get_database_backend());
173        crate::EntityRegistry::build_schema(schema, prefix)
174    }
175}