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    #[cfg(feature = "entity-registry")]
529    #[cfg_attr(docsrs, doc(cfg(feature = "entity-registry")))]
530    /// Builds a schema for all the entites in the given module
531    pub fn get_schema_registry(&self, prefix: &str) -> SchemaBuilder {
532        let schema = Schema::new(self.get_database_backend());
533        crate::EntityRegistry::build_schema(schema, prefix)
534    }
535
536    /// Sets a callback to metric this connection
537    pub fn set_metric_callback<F>(&mut self, _callback: F)
538    where
539        F: Fn(&crate::metric::Info<'_>) + Send + Sync + 'static,
540    {
541        match &mut self.inner {
542            #[cfg(feature = "sqlx-mysql")]
543            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => {
544                conn.set_metric_callback(_callback)
545            }
546            #[cfg(feature = "sqlx-postgres")]
547            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => {
548                conn.set_metric_callback(_callback)
549            }
550            #[cfg(feature = "sqlx-sqlite")]
551            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => {
552                conn.set_metric_callback(_callback)
553            }
554            _ => {}
555        }
556    }
557
558    /// Checks if a connection to the database is still valid.
559    pub async fn ping(&self) -> Result<(), DbErr> {
560        match &self.inner {
561            #[cfg(feature = "sqlx-mysql")]
562            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.ping().await,
563            #[cfg(feature = "sqlx-postgres")]
564            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => conn.ping().await,
565            #[cfg(feature = "sqlx-sqlite")]
566            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.ping().await,
567            #[cfg(feature = "mock")]
568            DatabaseConnectionType::MockDatabaseConnection(conn) => conn.ping(),
569            #[cfg(feature = "proxy")]
570            DatabaseConnectionType::ProxyDatabaseConnection(conn) => conn.ping().await,
571            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
572        }
573    }
574
575    /// Explicitly close the database connection.
576    /// See [`Self::close_by_ref`] for usage with references.
577    pub async fn close(self) -> Result<(), DbErr> {
578        self.close_by_ref().await
579    }
580
581    /// Explicitly close the database connection
582    pub async fn close_by_ref(&self) -> Result<(), DbErr> {
583        match &self.inner {
584            #[cfg(feature = "sqlx-mysql")]
585            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => conn.close_by_ref().await,
586            #[cfg(feature = "sqlx-postgres")]
587            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => conn.close_by_ref().await,
588            #[cfg(feature = "sqlx-sqlite")]
589            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => conn.close_by_ref().await,
590            #[cfg(feature = "mock")]
591            DatabaseConnectionType::MockDatabaseConnection(_) => {
592                // Nothing to cleanup, we just consume the `DatabaseConnection`
593                Ok(())
594            }
595            #[cfg(feature = "proxy")]
596            DatabaseConnectionType::ProxyDatabaseConnection(_) => {
597                // Nothing to cleanup, we just consume the `DatabaseConnection`
598                Ok(())
599            }
600            DatabaseConnectionType::Disconnected => Err(conn_err("Disconnected")),
601        }
602    }
603}
604
605impl DatabaseConnection {
606    /// Get [sqlx::MySqlPool]
607    ///
608    /// # Panics
609    ///
610    /// Panics if [DbConn] is not a MySQL connection.
611    #[cfg(feature = "sqlx-mysql")]
612    pub fn get_mysql_connection_pool(&self) -> &sqlx::MySqlPool {
613        match &self.inner {
614            DatabaseConnectionType::SqlxMySqlPoolConnection(conn) => &conn.pool,
615            _ => panic!("Not MySQL Connection"),
616        }
617    }
618
619    /// Get [sqlx::PgPool]
620    ///
621    /// # Panics
622    ///
623    /// Panics if [DbConn] is not a Postgres connection.
624    #[cfg(feature = "sqlx-postgres")]
625    pub fn get_postgres_connection_pool(&self) -> &sqlx::PgPool {
626        match &self.inner {
627            DatabaseConnectionType::SqlxPostgresPoolConnection(conn) => &conn.pool,
628            _ => panic!("Not Postgres Connection"),
629        }
630    }
631
632    /// Get [sqlx::SqlitePool]
633    ///
634    /// # Panics
635    ///
636    /// Panics if [DbConn] is not a SQLite connection.
637    #[cfg(feature = "sqlx-sqlite")]
638    pub fn get_sqlite_connection_pool(&self) -> &sqlx::SqlitePool {
639        match &self.inner {
640            DatabaseConnectionType::SqlxSqlitePoolConnection(conn) => &conn.pool,
641            _ => panic!("Not SQLite Connection"),
642        }
643    }
644}
645
646impl DbBackend {
647    /// Check if the URI is the same as the specified database backend.
648    /// Returns true if they match.
649    ///
650    /// # Panics
651    ///
652    /// Panics if `base_url` cannot be parsed as `Url`.
653    pub fn is_prefix_of(self, base_url: &str) -> bool {
654        let base_url_parsed = Url::parse(base_url).expect("Fail to parse database URL");
655        match self {
656            Self::Postgres => {
657                base_url_parsed.scheme() == "postgres" || base_url_parsed.scheme() == "postgresql"
658            }
659            Self::MySql => base_url_parsed.scheme() == "mysql",
660            Self::Sqlite => base_url_parsed.scheme() == "sqlite",
661        }
662    }
663
664    /// Build an SQL [Statement]
665    pub fn build<S>(&self, statement: &S) -> Statement
666    where
667        S: StatementBuilder,
668    {
669        statement.build(self)
670    }
671
672    /// Check if the database supports `RETURNING` syntax on insert and update
673    pub fn support_returning(&self) -> bool {
674        match self {
675            Self::Postgres => true,
676            Self::Sqlite if cfg!(feature = "sqlite-use-returning-for-3_35") => true,
677            Self::MySql if cfg!(feature = "mariadb-use-returning") => true,
678            _ => false,
679        }
680    }
681
682    /// A getter for database dependent boolean value
683    pub fn boolean_value(&self, boolean: bool) -> sea_query::Value {
684        match self {
685            Self::MySql | Self::Postgres | Self::Sqlite => boolean.into(),
686        }
687    }
688
689    /// Get the display string for this enum
690    pub fn as_str(&self) -> &'static str {
691        match self {
692            DatabaseBackend::MySql => "MySql",
693            DatabaseBackend::Postgres => "Postgres",
694            DatabaseBackend::Sqlite => "Sqlite",
695        }
696    }
697}
698
699#[cfg(test)]
700mod tests {
701    use crate::DatabaseConnection;
702
703    #[test]
704    fn assert_database_connection_traits() {
705        fn assert_send_sync<T: Send + Sync>() {}
706
707        assert_send_sync::<DatabaseConnection>();
708    }
709}
sea_orm/database/db_connection.rs

sea_orm/database/
db_connection.rs
