sea_orm/database/
db_connection.rs

1use crate::{
2    AccessMode, ConnectionTrait, DatabaseTransaction, ExecResult, IsolationLevel, QueryResult,
3    Statement, StatementBuilder, StreamTrait, TransactionError, TransactionTrait, error::*,
4};
5use std::{fmt::Debug, future::Future, pin::Pin};
6use tracing::instrument;
7use url::Url;
8
9#[cfg(feature = "sqlx-dep")]
10use sqlx::pool::PoolConnection;
11
12#[cfg(any(feature = "mock", feature = "proxy"))]
13use std::sync::Arc;
14
15/// Handle a database connection depending on the backend enabled by the feature
16/// flags. This creates a connection pool internally (for SQLx connections),
17/// and so is cheap to clone.
18#[derive(Debug, Clone)]
19#[non_exhaustive]
20pub struct DatabaseConnection {
21    /// `DatabaseConnection` used to be a enum. Now it's moved into inner,
22    /// because we have to attach other contexts.
23    pub inner: DatabaseConnectionType,
24    #[cfg(feature = "rbac")]
25    pub(crate) rbac: crate::RbacEngineMount,
26}
27
28/// The underlying database connection type.
29#[derive(Clone)]
30pub enum DatabaseConnectionType {
31    /// MySql database connection pool
32    #[cfg(feature = "sqlx-mysql")]
33    SqlxMySqlPoolConnection(crate::SqlxMySqlPoolConnection),
34
35    /// PostgreSQL database connection pool
36    #[cfg(feature = "sqlx-postgres")]
37    SqlxPostgresPoolConnection(crate::SqlxPostgresPoolConnection),
38
39    /// SQLite database connection pool
40    #[cfg(feature = "sqlx-sqlite")]
41    SqlxSqlitePoolConnection(crate::SqlxSqlitePoolConnection),
42
43    /// Mock database connection useful for testing
44    #[cfg(feature = "mock")]
45    MockDatabaseConnection(Arc<crate::MockDatabaseConnection>),
46
47    /// Proxy database connection
48    #[cfg(feature = "proxy")]
49    ProxyDatabaseConnection(Arc<crate::ProxyDatabaseConnection>),
50
51    /// The connection has never been established
52    Disconnected,
53}
54
55/// The same as a [DatabaseConnection]
56pub type DbConn = DatabaseConnection;
57
58impl Default for DatabaseConnection {
59    fn default() -> Self {
60        DatabaseConnectionType::Disconnected.into()
61    }
62}
63
64impl From<DatabaseConnectionType> for DatabaseConnection {
65    fn from(inner: DatabaseConnectionType) -> Self {
66        Self {
67            inner,
68            #[cfg(feature = "rbac")]
69            rbac: Default::default(),
70        }
71    }
72}
73
74/// The type of database backend for real world databases.
75/// This is enabled by feature flags as specified in the crate documentation
76#[derive(Debug, Copy, Clone, PartialEq, Eq)]
77#[non_exhaustive]
78pub enum DatabaseBackend {
79    /// A MySQL backend
80    MySql,
81    /// A PostgreSQL backend
82    Postgres,
83    /// A SQLite backend
84    Sqlite,
85}
86
87/// A shorthand for [DatabaseBackend].
88pub type DbBackend = DatabaseBackend;
89
90#[derive(Debug)]
91pub(crate) enum InnerConnection {
92    #[cfg(feature = "sqlx-mysql")]
93    MySql(PoolConnection<sqlx::MySql>),
94    #[cfg(feature = "sqlx-postgres")]
95    Postgres(PoolConnection<sqlx::Postgres>),
96    #[cfg(feature = "sqlx-sqlite")]
97    Sqlite(PoolConnection<sqlx::Sqlite>),
98    #[cfg(feature = "mock")]
99    Mock(Arc<crate::MockDatabaseConnection>),
100    #[cfg(feature = "proxy")]
101    Proxy(Arc<crate::ProxyDatabaseConnection>),
102}
103
104impl Debug for DatabaseConnectionType {
105    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
106        write!(
107            f,
108            "{}",
109            match self {
110                #[cfg(feature = "sqlx-mysql")]
111                Self::SqlxMySqlPoolConnection(_) => "SqlxMySqlPoolConnection",
112                #[cfg(feature = "sqlx-postgres")]
113                Self::SqlxPostgresPoolConnection(_) => "SqlxPostgresPoolConnection",
114                #[cfg(feature = "sqlx-sqlite")]
115                Self::SqlxSqlitePoolConnection(_) => "SqlxSqlitePoolConnection",
116                #[cfg(feature = "mock")]
117                Self::MockDatabaseConnection(_) => "MockDatabaseConnection",
118                #[cfg(feature = "proxy")]
119                Self::ProxyDatabaseConnection(_) => "ProxyDatabaseConnection",
120                Self::Disconnected => "Disconnected",
121            }
122        )
123    }
124}
125
126#[async_trait::async_trait]
127impl ConnectionTrait for DatabaseConnection {
128    fn get_database_backend(&self) -> DbBackend {
129        self.get_database_backend()
130    }
131
132    #[instrument(level = "trace")]
133    #[allow(unused_variables)]
134    async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
135        match &self.inner {
136            #[cfg(feature = "sqlx-mysql")]
137            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.execute(stmt).await,
138            #[cfg(feature = "sqlx-postgres")]
139            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => conn.execute(stmt).await,
140            #[cfg(feature = "sqlx-sqlite")]
141            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.execute(stmt).await,
142            #[cfg(feature = "mock")]
143            DatabaseConnectionType::MockDatabaseConnection(conn) => conn.execute(stmt),
144            #[cfg(feature = "proxy")]
145            DatabaseConnectionType::ProxyDatabaseConnection(conn) => conn.execute(stmt).await,
146            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
147        }
148    }
149
150    #[instrument(level = "trace")]
151    #[allow(unused_variables)]
152    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
153        match &self.inner {
154            #[cfg(feature = "sqlx-mysql")]
155            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => {
156                conn.execute_unprepared(sql).await
157            }
158            #[cfg(feature = "sqlx-postgres")]
159            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => {
160                conn.execute_unprepared(sql).await
161            }
162            #[cfg(feature = "sqlx-sqlite")]
163            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => {
164                conn.execute_unprepared(sql).await
165            }
166            #[cfg(feature = "mock")]
167            DatabaseConnectionType::MockDatabaseConnection(conn) => {
168                let db_backend = conn.get_database_backend();
169                let stmt = Statement::from_string(db_backend, sql);
170                conn.execute(stmt)
171            }
172            #[cfg(feature = "proxy")]
173            DatabaseConnectionType::ProxyDatabaseConnection(conn) => {
174                let db_backend = conn.get_database_backend();
175                let stmt = Statement::from_string(db_backend, sql);
176                conn.execute(stmt).await
177            }
178            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
179        }
180    }
181
182    #[instrument(level = "trace")]
183    #[allow(unused_variables)]
184    async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
185        match &self.inner {
186            #[cfg(feature = "sqlx-mysql")]
187            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.query_one(stmt).await,
188            #[cfg(feature = "sqlx-postgres")]
189            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => conn.query_one(stmt).await,
190            #[cfg(feature = "sqlx-sqlite")]
191            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.query_one(stmt).await,
192            #[cfg(feature = "mock")]
193            DatabaseConnectionType::MockDatabaseConnection(conn) => conn.query_one(stmt),
194            #[cfg(feature = "proxy")]
195            DatabaseConnectionType::ProxyDatabaseConnection(conn) => conn.query_one(stmt).await,
196            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
197        }
198    }
199
200    #[instrument(level = "trace")]
201    #[allow(unused_variables)]
202    async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
203        match &self.inner {
204            #[cfg(feature = "sqlx-mysql")]
205            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.query_all(stmt).await,
206            #[cfg(feature = "sqlx-postgres")]
207            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => conn.query_all(stmt).await,
208            #[cfg(feature = "sqlx-sqlite")]
209            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.query_all(stmt).await,
210            #[cfg(feature = "mock")]
211            DatabaseConnectionType::MockDatabaseConnection(conn) => conn.query_all(stmt),
212            #[cfg(feature = "proxy")]
213            DatabaseConnectionType::ProxyDatabaseConnection(conn) => conn.query_all(stmt).await,
214            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
215        }
216    }
217
218    #[cfg(feature = "mock")]
219    fn is_mock_connection(&self) -> bool {
220        matches!(
221            self,
222            DatabaseConnection {
223                inner: DatabaseConnectionType::MockDatabaseConnection(_),
224                ..
225            }
226        )
227    }
228}
229
230#[async_trait::async_trait]
231impl StreamTrait for DatabaseConnection {
232    type Stream<'a> = crate::QueryStream;
233
234    fn get_database_backend(&self) -> DbBackend {
235        self.get_database_backend()
236    }
237
238    #[instrument(level = "trace")]
239    #[allow(unused_variables)]
240    fn stream_raw<'a>(
241        &'a self,
242        stmt: Statement,
243    ) -> Pin<Box<dyn Future<Output = Result<Self::Stream<'a>, DbErr>> + 'a + Send>> {
244        Box::pin(async move {
245            match &self.inner {
246                #[cfg(feature = "sqlx-mysql")]
247                DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.stream(stmt).await,
248                #[cfg(feature = "sqlx-postgres")]
249                DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => conn.stream(stmt).await,
250                #[cfg(feature = "sqlx-sqlite")]
251                DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.stream(stmt).await,
252                #[cfg(feature = "mock")]
253                DatabaseConnectionType::MockDatabaseConnection(conn) => {
254                    Ok(crate::QueryStream::from((Arc::clone(conn), stmt, None)))
255                }
256                #[cfg(feature = "proxy")]
257                DatabaseConnectionType::ProxyDatabaseConnection(conn) => {
258                    Ok(crate::QueryStream::from((Arc::clone(conn), stmt, None)))
259                }
260                DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
261            }
262        })
263    }
264}
265
266#[async_trait::async_trait]
267impl TransactionTrait for DatabaseConnection {
268    type Transaction = DatabaseTransaction;
269
270    #[instrument(level = "trace")]
271    async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
272        match &self.inner {
273            #[cfg(feature = "sqlx-mysql")]
274            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.begin(None, None).await,
275            #[cfg(feature = "sqlx-postgres")]
276            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => {
277                conn.begin(None, None).await
278            }
279            #[cfg(feature = "sqlx-sqlite")]
280            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.begin(None, None).await,
281            #[cfg(feature = "mock")]
282            DatabaseConnectionType::MockDatabaseConnection(conn) => {
283                DatabaseTransaction::new_mock(Arc::clone(conn), None).await
284            }
285            #[cfg(feature = "proxy")]
286            DatabaseConnectionType::ProxyDatabaseConnection(conn) => {
287                DatabaseTransaction::new_proxy(conn.clone(), None).await
288            }
289            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
290        }
291    }
292
293    #[instrument(level = "trace")]
294    async fn begin_with_config(
295        &self,
296        _isolation_level: Option<IsolationLevel>,
297        _access_mode: Option<AccessMode>,
298    ) -> Result<DatabaseTransaction, DbErr> {
299        match &self.inner {
300            #[cfg(feature = "sqlx-mysql")]
301            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => {
302                conn.begin(_isolation_level, _access_mode).await
303            }
304            #[cfg(feature = "sqlx-postgres")]
305            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => {
306                conn.begin(_isolation_level, _access_mode).await
307            }
308            #[cfg(feature = "sqlx-sqlite")]
309            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => {
310                conn.begin(_isolation_level, _access_mode).await
311            }
312            #[cfg(feature = "mock")]
313            DatabaseConnectionType::MockDatabaseConnection(conn) => {
314                DatabaseTransaction::new_mock(Arc::clone(conn), None).await
315            }
316            #[cfg(feature = "proxy")]
317            DatabaseConnectionType::ProxyDatabaseConnection(conn) => {
318                DatabaseTransaction::new_proxy(conn.clone(), None).await
319            }
320            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
321        }
322    }
323
324    /// Execute the function inside a transaction.
325    /// If the function returns an error, the transaction will be rolled back. If it does not return an error, the transaction will be committed.
326    #[instrument(level = "trace", skip(_callback))]
327    async fn transaction<F, T, E>(&self, _callback: F) -> Result<T, TransactionError<E>>
328    where
329        F: for<'c> FnOnce(
330                &'c DatabaseTransaction,
331            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'c>>
332            + Send,
333        T: Send,
334        E: std::fmt::Display + std::fmt::Debug + Send,
335    {
336        match &self.inner {
337            #[cfg(feature = "sqlx-mysql")]
338            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => {
339                conn.transaction(_callback, None, None).await
340            }
341            #[cfg(feature = "sqlx-postgres")]
342            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => {
343                conn.transaction(_callback, None, None).await
344            }
345            #[cfg(feature = "sqlx-sqlite")]
346            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => {
347                conn.transaction(_callback, None, None).await
348            }
349            #[cfg(feature = "mock")]
350            DatabaseConnectionType::MockDatabaseConnection(conn) => {
351                let transaction = DatabaseTransaction::new_mock(Arc::clone(conn), None)
352                    .await
353                    .map_err(TransactionError::Connection)?;
354                transaction.run(_callback).await
355            }
356            #[cfg(feature = "proxy")]
357            DatabaseConnectionType::ProxyDatabaseConnection(conn) => {
358                let transaction = DatabaseTransaction::new_proxy(conn.clone(), None)
359                    .await
360                    .map_err(TransactionError::Connection)?;
361                transaction.run(_callback).await
362            }
363            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected").into()),
364        }
365    }
366
367    /// Execute the function inside a transaction.
368    /// If the function returns an error, the transaction will be rolled back. If it does not return an error, the transaction will be committed.
369    #[instrument(level = "trace", skip(_callback))]
370    async fn transaction_with_config<F, T, E>(
371        &self,
372        _callback: F,
373        _isolation_level: Option<IsolationLevel>,
374        _access_mode: Option<AccessMode>,
375    ) -> Result<T, TransactionError<E>>
376    where
377        F: for<'c> FnOnce(
378                &'c DatabaseTransaction,
379            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'c>>
380            + Send,
381        T: Send,
382        E: std::fmt::Display + std::fmt::Debug + Send,
383    {
384        match &self.inner {
385            #[cfg(feature = "sqlx-mysql")]
386            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => {
387                conn.transaction(_callback, _isolation_level, _access_mode)
388                    .await
389            }
390            #[cfg(feature = "sqlx-postgres")]
391            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => {
392                conn.transaction(_callback, _isolation_level, _access_mode)
393                    .await
394            }
395            #[cfg(feature = "sqlx-sqlite")]
396            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => {
397                conn.transaction(_callback, _isolation_level, _access_mode)
398                    .await
399            }
400            #[cfg(feature = "mock")]
401            DatabaseConnectionType::MockDatabaseConnection(conn) => {
402                let transaction = DatabaseTransaction::new_mock(Arc::clone(conn), None)
403                    .await
404                    .map_err(TransactionError::Connection)?;
405                transaction.run(_callback).await
406            }
407            #[cfg(feature = "proxy")]
408            DatabaseConnectionType::ProxyDatabaseConnection(conn) => {
409                let transaction = DatabaseTransaction::new_proxy(conn.clone(), None)
410                    .await
411                    .map_err(TransactionError::Connection)?;
412                transaction.run(_callback).await
413            }
414            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected").into()),
415        }
416    }
417}
418
419#[cfg(feature = "mock")]
420impl DatabaseConnection {
421    /// Generate a database connection for testing the Mock database
422    ///
423    /// # Panics
424    ///
425    /// Panics if [DbConn] is not a mock connection.
426    pub fn as_mock_connection(&self) -> &crate::MockDatabaseConnection {
427        match &self.inner {
428            DatabaseConnectionType::MockDatabaseConnection(mock_conn) => mock_conn,
429            _ => panic!("Not mock connection"),
430        }
431    }
432
433    /// Get the transaction log as a collection Vec<[crate::Transaction]>
434    ///
435    /// # Panics
436    ///
437    /// Panics if the mocker mutex is being held by another thread.
438    pub fn into_transaction_log(self) -> Vec<crate::Transaction> {
439        let mut mocker = self
440            .as_mock_connection()
441            .get_mocker_mutex()
442            .lock()
443            .expect("Fail to acquire mocker");
444        mocker.drain_transaction_log()
445    }
446}
447
448#[cfg(feature = "proxy")]
449impl DatabaseConnection {
450    /// Generate a database connection for testing the Proxy database
451    ///
452    /// # Panics
453    ///
454    /// Panics if [DbConn] is not a proxy connection.
455    pub fn as_proxy_connection(&self) -> &crate::ProxyDatabaseConnection {
456        match &self.inner {
457            DatabaseConnectionType::ProxyDatabaseConnection(proxy_conn) => proxy_conn,
458            _ => panic!("Not proxy connection"),
459        }
460    }
461}
462
463#[cfg(feature = "rbac")]
464impl DatabaseConnection {
465    /// Load RBAC data from the same database as this connection and setup RBAC engine.
466    /// If the RBAC engine already exists, it will be replaced.
467    pub async fn load_rbac(&self) -> Result<(), DbErr> {
468        self.load_rbac_from(self).await
469    }
470
471    /// Load RBAC data from the given database connection and setup RBAC engine.
472    /// This could be from another database.
473    pub async fn load_rbac_from(&self, db: &DbConn) -> Result<(), DbErr> {
474        let engine = crate::rbac::RbacEngine::load_from(db).await?;
475        self.rbac.replace(engine);
476        Ok(())
477    }
478
479    /// Replace the internal RBAC engine.
480    pub fn replace_rbac(&self, engine: crate::rbac::RbacEngine) {
481        self.rbac.replace(engine);
482    }
483
484    /// Create a restricted connection with access control specific for the user.
485    pub fn restricted_for(
486        &self,
487        user_id: crate::rbac::RbacUserId,
488    ) -> Result<crate::RestrictedConnection, DbErr> {
489        if self.rbac.is_some() {
490            Ok(crate::RestrictedConnection {
491                user_id,
492                conn: self.clone(),
493            })
494        } else {
495            Err(DbErr::RbacError("engine not set up".into()))
496        }
497    }
498}
499
500impl DatabaseConnection {
501    /// Get the database backend for this connection
502    ///
503    /// # Panics
504    ///
505    /// Panics if [DatabaseConnection] is `Disconnected`.
506    pub fn get_database_backend(&self) -> DbBackend {
507        match &self.inner {
508            #[cfg(feature = "sqlx-mysql")]
509            DatabaseConnectionType::SqlxMySqlPoolConnection(_) => DbBackend::MySql,
510            #[cfg(feature = "sqlx-postgres")]
511            DatabaseConnectionType::SqlxPostgresPoolConnection(_) => DbBackend::Postgres,
512            #[cfg(feature = "sqlx-sqlite")]
513            DatabaseConnectionType::SqlxSqlitePoolConnection(_) => DbBackend::Sqlite,
514            #[cfg(feature = "mock")]
515            DatabaseConnectionType::MockDatabaseConnection(conn) => conn.get_database_backend(),
516            #[cfg(feature = "proxy")]
517            DatabaseConnectionType::ProxyDatabaseConnection(conn) => conn.get_database_backend(),
518            DatabaseConnectionType::Disconnected => panic!("Disconnected"),
519        }
520    }
521
522    /// Sets a callback to metric this connection
523    pub fn set_metric_callback<F>(&mut self, _callback: F)
524    where
525        F: Fn(&crate::metric::Info<'_>) + Send + Sync + 'static,
526    {
527        match &mut self.inner {
528            #[cfg(feature = "sqlx-mysql")]
529            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => {
530                conn.set_metric_callback(_callback)
531            }
532            #[cfg(feature = "sqlx-postgres")]
533            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => {
534                conn.set_metric_callback(_callback)
535            }
536            #[cfg(feature = "sqlx-sqlite")]
537            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => {
538                conn.set_metric_callback(_callback)
539            }
540            _ => {}
541        }
542    }
543
544    /// Checks if a connection to the database is still valid.
545    pub async fn ping(&self) -> Result<(), DbErr> {
546        match &self.inner {
547            #[cfg(feature = "sqlx-mysql")]
548            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.ping().await,
549            #[cfg(feature = "sqlx-postgres")]
550            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => conn.ping().await,
551            #[cfg(feature = "sqlx-sqlite")]
552            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.ping().await,
553            #[cfg(feature = "mock")]
554            DatabaseConnectionType::MockDatabaseConnection(conn) => conn.ping(),
555            #[cfg(feature = "proxy")]
556            DatabaseConnectionType::ProxyDatabaseConnection(conn) => conn.ping().await,
557            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
558        }
559    }
560
561    /// Explicitly close the database connection.
562    /// See [`Self::close_by_ref`] for usage with references.
563    pub async fn close(self) -> Result<(), DbErr> {
564        self.close_by_ref().await
565    }
566
567    /// Explicitly close the database connection
568    pub async fn close_by_ref(&self) -> Result<(), DbErr> {
569        match &self.inner {
570            #[cfg(feature = "sqlx-mysql")]
571            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.close_by_ref().await,
572            #[cfg(feature = "sqlx-postgres")]
573            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => conn.close_by_ref().await,
574            #[cfg(feature = "sqlx-sqlite")]
575            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.close_by_ref().await,
576            #[cfg(feature = "mock")]
577            DatabaseConnectionType::MockDatabaseConnection(_) => {
578                // Nothing to cleanup, we just consume the `DatabaseConnection`
579                Ok(())
580            }
581            #[cfg(feature = "proxy")]
582            DatabaseConnectionType::ProxyDatabaseConnection(_) => {
583                // Nothing to cleanup, we just consume the `DatabaseConnection`
584                Ok(())
585            }
586            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
587        }
588    }
589}
590
591impl DatabaseConnection {
592    /// Get [sqlx::MySqlPool]
593    ///
594    /// # Panics
595    ///
596    /// Panics if [DbConn] is not a MySQL connection.
597    #[cfg(feature = "sqlx-mysql")]
598    pub fn get_mysql_connection_pool(&self) -> &sqlx::MySqlPool {
599        match &self.inner {
600            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => &conn.pool,
601            _ => panic!("Not MySQL Connection"),
602        }
603    }
604
605    /// Get [sqlx::PgPool]
606    ///
607    /// # Panics
608    ///
609    /// Panics if [DbConn] is not a Postgres connection.
610    #[cfg(feature = "sqlx-postgres")]
611    pub fn get_postgres_connection_pool(&self) -> &sqlx::PgPool {
612        match &self.inner {
613            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => &conn.pool,
614            _ => panic!("Not Postgres Connection"),
615        }
616    }
617
618    /// Get [sqlx::SqlitePool]
619    ///
620    /// # Panics
621    ///
622    /// Panics if [DbConn] is not a SQLite connection.
623    #[cfg(feature = "sqlx-sqlite")]
624    pub fn get_sqlite_connection_pool(&self) -> &sqlx::SqlitePool {
625        match &self.inner {
626            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => &conn.pool,
627            _ => panic!("Not SQLite Connection"),
628        }
629    }
630}
631
632impl DbBackend {
633    /// Check if the URI is the same as the specified database backend.
634    /// Returns true if they match.
635    ///
636    /// # Panics
637    ///
638    /// Panics if `base_url` cannot be parsed as `Url`.
639    pub fn is_prefix_of(self, base_url: &str) -> bool {
640        let base_url_parsed = Url::parse(base_url).expect("Fail to parse database URL");
641        match self {
642            Self::Postgres => {
643                base_url_parsed.scheme() == "postgres" || base_url_parsed.scheme() == "postgresql"
644            }
645            Self::MySql => base_url_parsed.scheme() == "mysql",
646            Self::Sqlite => base_url_parsed.scheme() == "sqlite",
647        }
648    }
649
650    /// Build an SQL [Statement]
651    pub fn build<S>(&self, statement: &S) -> Statement
652    where
653        S: StatementBuilder,
654    {
655        statement.build(self)
656    }
657
658    /// Check if the database supports `RETURNING` syntax on insert and update
659    pub fn support_returning(&self) -> bool {
660        match self {
661            Self::Postgres => true,
662            Self::Sqlite if cfg!(feature = "sqlite-use-returning-for-3_35") => true,
663            _ => false,
664        }
665    }
666
667    /// A getter for database dependent boolean value
668    pub fn boolean_value(&self, boolean: bool) -> sea_query::Value {
669        match self {
670            Self::MySql | Self::Postgres | Self::Sqlite => boolean.into(),
671        }
672    }
673}
674
675impl DatabaseBackend {
676    /// Get the display string for this enum
677    pub fn as_str(&self) -> &'static str {
678        match self {
679            DatabaseBackend::MySql => "MySql",
680            DatabaseBackend::Postgres => "Postgres",
681            DatabaseBackend::Sqlite => "Sqlite",
682        }
683    }
684}
685
686#[cfg(test)]
687mod tests {
688    use crate::DatabaseConnection;
689
690    #[test]
691    fn assert_database_connection_traits() {
692        fn assert_send_sync<T: Send + Sync>() {}
693
694        assert_send_sync::<DatabaseConnection>();
695    }
696}
sea_orm/database/db_connection.rs

sea_orm/database/
db_connection.rs
