sea_orm/database/
executor.rs

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