sea_orm/database/
db_connection.rs

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

sea_orm/database/
db_connection.rs
