Skip to main content

sea_orm/database/
mod.rs

1use std::{sync::Arc, time::Duration};
2
3#[cfg(not(feature = "sync"))]
4#[cfg(feature = "sqlx-mysql")]
5use sqlx::mysql::MySqlConnectOptions;
6#[cfg(feature = "sqlx-postgres")]
7use sqlx::postgres::PgConnectOptions;
8#[cfg(feature = "sqlx-sqlite")]
9use sqlx::sqlite::SqliteConnectOptions;
10
11mod connection;
12mod db_connection;
13mod executor;
14#[cfg(feature = "mock")]
15#[cfg_attr(docsrs, doc(cfg(feature = "mock")))]
16mod mock;
17#[cfg(feature = "proxy")]
18#[cfg_attr(docsrs, doc(cfg(feature = "proxy")))]
19mod proxy;
20#[cfg(feature = "rbac")]
21mod restricted_connection;
22#[cfg(all(feature = "schema-sync", feature = "rusqlite"))]
23mod sea_schema_rusqlite;
24#[cfg(all(feature = "schema-sync", feature = "sqlx-dep"))]
25mod sea_schema_shim;
26mod statement;
27mod stream;
28mod tracing_spans;
29mod transaction;
30
31pub use connection::*;
32pub use db_connection::*;
33pub use executor::*;
34#[cfg(feature = "mock")]
35#[cfg_attr(docsrs, doc(cfg(feature = "mock")))]
36pub use mock::*;
37#[cfg(feature = "proxy")]
38#[cfg_attr(docsrs, doc(cfg(feature = "proxy")))]
39pub use proxy::*;
40#[cfg(feature = "rbac")]
41pub use restricted_connection::*;
42pub use statement::*;
43use std::borrow::Cow;
44pub use stream::*;
45use tracing::instrument;
46pub use transaction::*;
47
48use crate::error::*;
49
50/// Defines a database
51#[derive(Debug, Default)]
52pub struct Database;
53
54#[cfg(feature = "sync")]
55type BoxFuture<'a, T> = T;
56
57#[cfg(feature = "sqlx-mysql")]
58type MapMySqlPoolOptsFn =
59    Arc<dyn Fn(sqlx::pool::PoolOptions<sqlx::MySql>) -> sqlx::pool::PoolOptions<sqlx::MySql>>;
60
61#[cfg(feature = "sqlx-postgres")]
62type MapPgPoolOptsFn =
63    Arc<dyn Fn(sqlx::pool::PoolOptions<sqlx::Postgres>) -> sqlx::pool::PoolOptions<sqlx::Postgres>>;
64
65#[cfg(feature = "sqlx-sqlite")]
66type MapSqlitePoolOptsFn = Option<
67    Arc<dyn Fn(sqlx::pool::PoolOptions<sqlx::Sqlite>) -> sqlx::pool::PoolOptions<sqlx::Sqlite>>,
68>;
69
70type AfterConnectCallback =
71    Option<Arc<dyn Fn(DatabaseConnection) -> BoxFuture<'static, Result<(), DbErr>> + 'static>>;
72
73/// Defines the configuration options of a database
74#[derive(derive_more::Debug, Clone)]
75pub struct ConnectOptions {
76    /// The URI of the database
77    pub(crate) url: String,
78    /// Maximum number of connections for a pool
79    pub(crate) max_connections: Option<u32>,
80    /// Minimum number of connections for a pool
81    pub(crate) min_connections: Option<u32>,
82    /// The connection timeout for a packet connection
83    pub(crate) connect_timeout: Option<Duration>,
84    /// Maximum idle time for a particular connection to prevent
85    /// network resource exhaustion
86    pub(crate) idle_timeout: Option<Option<Duration>>,
87    /// Set the maximum amount of time to spend waiting for acquiring a connection
88    pub(crate) acquire_timeout: Option<Duration>,
89    /// Set the maximum lifetime of individual connections
90    pub(crate) max_lifetime: Option<Option<Duration>>,
91    /// Enable SQLx statement logging
92    pub(crate) sqlx_logging: bool,
93    /// Record SQL statements in tracing spans
94    pub(crate) record_stmt_in_spans: bool,
95    /// SQLx statement logging level (ignored if `sqlx_logging` is false)
96    pub(crate) sqlx_logging_level: log::LevelFilter,
97    /// SQLx slow statements logging level (ignored if `sqlx_logging` is false)
98    pub(crate) sqlx_slow_statements_logging_level: log::LevelFilter,
99    /// SQLx slow statements duration threshold (ignored if `sqlx_logging` is false)
100    pub(crate) sqlx_slow_statements_logging_threshold: Duration,
101    /// set sqlcipher key
102    pub(crate) sqlcipher_key: Option<Cow<'static, str>>,
103    /// Schema search path (PostgreSQL only)
104    pub(crate) schema_search_path: Option<String>,
105    /// Application name (PostgreSQL only)
106    pub(crate) application_name: Option<String>,
107    /// Statement timeout (PostgreSQL only)
108    pub(crate) statement_timeout: Option<Duration>,
109    pub(crate) test_before_acquire: bool,
110    /// Only establish connections to the DB as needed. If set to `true`, the db connection will
111    /// be created using SQLx's [connect_lazy](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html#method.connect_lazy)
112    /// method.
113    pub(crate) connect_lazy: bool,
114
115    #[debug(skip)]
116    pub(crate) after_connect: AfterConnectCallback,
117
118    #[cfg(feature = "sqlx-mysql")]
119    #[debug(skip)]
120    pub(crate) mysql_pool_opts_fn: Option<MapMySqlPoolOptsFn>,
121    #[cfg(feature = "sqlx-postgres")]
122    #[debug(skip)]
123    pub(crate) pg_pool_opts_fn: Option<MapPgPoolOptsFn>,
124    #[cfg(feature = "sqlx-sqlite")]
125    #[debug(skip)]
126    pub(crate) sqlite_pool_opts_fn: MapSqlitePoolOptsFn,
127    #[cfg(feature = "sqlx-mysql")]
128    #[debug(skip)]
129    pub(crate) mysql_opts_fn: Option<Arc<dyn Fn(MySqlConnectOptions) -> MySqlConnectOptions>>,
130    #[cfg(feature = "sqlx-postgres")]
131    #[debug(skip)]
132    pub(crate) pg_opts_fn: Option<Arc<dyn Fn(PgConnectOptions) -> PgConnectOptions>>,
133    #[cfg(feature = "sqlx-sqlite")]
134    #[debug(skip)]
135    pub(crate) sqlite_opts_fn: Option<Arc<dyn Fn(SqliteConnectOptions) -> SqliteConnectOptions>>,
136}
137
138impl Database {
139    /// Method to create a [DatabaseConnection] on a database. This method will return an error
140    /// if the database is not available.
141    #[instrument(level = "trace", skip(opt))]
142    pub fn connect<C>(opt: C) -> Result<DatabaseConnection, DbErr>
143    where
144        C: Into<ConnectOptions>,
145    {
146        let opt: ConnectOptions = opt.into();
147
148        if url::Url::parse(&opt.url).is_err() {
149            return Err(conn_err(format!(
150                "The connection string '{}' cannot be parsed.",
151                opt.url
152            )));
153        }
154
155        #[cfg(feature = "sqlx-mysql")]
156        if DbBackend::MySql.is_prefix_of(&opt.url) {
157            return crate::SqlxMySqlConnector::connect(opt);
158        }
159        #[cfg(feature = "sqlx-postgres")]
160        if DbBackend::Postgres.is_prefix_of(&opt.url) {
161            return crate::SqlxPostgresConnector::connect(opt);
162        }
163        #[cfg(feature = "sqlx-sqlite")]
164        if DbBackend::Sqlite.is_prefix_of(&opt.url) {
165            return crate::SqlxSqliteConnector::connect(opt);
166        }
167        #[cfg(feature = "rusqlite")]
168        if DbBackend::Sqlite.is_prefix_of(&opt.url) {
169            return crate::driver::rusqlite::RusqliteConnector::connect(opt);
170        }
171        #[cfg(feature = "mock")]
172        if crate::MockDatabaseConnector::accepts(&opt.url) {
173            return crate::MockDatabaseConnector::connect(&opt.url);
174        }
175
176        Err(conn_err(format!(
177            "The connection string '{}' has no supporting driver.",
178            opt.url
179        )))
180    }
181
182    /// Method to create a [DatabaseConnection] on a proxy database
183    #[cfg(feature = "proxy")]
184    #[instrument(level = "trace", skip(proxy_func_arc))]
185    pub fn connect_proxy(
186        db_type: DbBackend,
187        proxy_func_arc: std::sync::Arc<Box<dyn ProxyDatabaseTrait>>,
188    ) -> Result<DatabaseConnection, DbErr> {
189        match db_type {
190            DbBackend::MySql => {
191                return crate::ProxyDatabaseConnector::connect(
192                    DbBackend::MySql,
193                    proxy_func_arc.to_owned(),
194                );
195            }
196            DbBackend::Postgres => {
197                return crate::ProxyDatabaseConnector::connect(
198                    DbBackend::Postgres,
199                    proxy_func_arc.to_owned(),
200                );
201            }
202            DbBackend::Sqlite => {
203                return crate::ProxyDatabaseConnector::connect(
204                    DbBackend::Sqlite,
205                    proxy_func_arc.to_owned(),
206                );
207            }
208        }
209    }
210}
211
212impl<T> From<T> for ConnectOptions
213where
214    T: Into<String>,
215{
216    fn from(s: T) -> ConnectOptions {
217        ConnectOptions::new(s.into())
218    }
219}
220
221impl ConnectOptions {
222    /// Create new [ConnectOptions] for a [Database] by passing in a URI string
223    pub fn new<T>(url: T) -> Self
224    where
225        T: Into<String>,
226    {
227        Self {
228            url: url.into(),
229            max_connections: None,
230            min_connections: None,
231            connect_timeout: None,
232            idle_timeout: None,
233            acquire_timeout: None,
234            max_lifetime: None,
235            sqlx_logging: true,
236            record_stmt_in_spans: true,
237            sqlx_logging_level: log::LevelFilter::Info,
238            sqlx_slow_statements_logging_level: log::LevelFilter::Off,
239            sqlx_slow_statements_logging_threshold: Duration::from_secs(1),
240            sqlcipher_key: None,
241            schema_search_path: None,
242            application_name: None,
243            statement_timeout: None,
244            test_before_acquire: true,
245            connect_lazy: false,
246            after_connect: None,
247            #[cfg(feature = "sqlx-mysql")]
248            mysql_pool_opts_fn: None,
249            #[cfg(feature = "sqlx-postgres")]
250            pg_pool_opts_fn: None,
251            #[cfg(feature = "sqlx-sqlite")]
252            sqlite_pool_opts_fn: None,
253            #[cfg(feature = "sqlx-mysql")]
254            mysql_opts_fn: None,
255            #[cfg(feature = "sqlx-postgres")]
256            pg_opts_fn: None,
257            #[cfg(feature = "sqlx-sqlite")]
258            sqlite_opts_fn: None,
259        }
260    }
261
262    /// Get the database URL of the pool
263    pub fn get_url(&self) -> &str {
264        &self.url
265    }
266
267    /// Set the maximum number of connections of the pool
268    pub fn max_connections(&mut self, value: u32) -> &mut Self {
269        self.max_connections = Some(value);
270        self
271    }
272
273    /// Get the maximum number of connections of the pool, if set
274    pub fn get_max_connections(&self) -> Option<u32> {
275        self.max_connections
276    }
277
278    /// Set the minimum number of connections of the pool
279    pub fn min_connections(&mut self, value: u32) -> &mut Self {
280        self.min_connections = Some(value);
281        self
282    }
283
284    /// Get the minimum number of connections of the pool, if set
285    pub fn get_min_connections(&self) -> Option<u32> {
286        self.min_connections
287    }
288
289    /// Set the timeout duration when acquiring a connection
290    pub fn connect_timeout(&mut self, value: Duration) -> &mut Self {
291        self.connect_timeout = Some(value);
292        self
293    }
294
295    /// Get the timeout duration when acquiring a connection, if set
296    pub fn get_connect_timeout(&self) -> Option<Duration> {
297        self.connect_timeout
298    }
299
300    /// Set the idle duration before closing a connection.
301    pub fn idle_timeout<T>(&mut self, value: T) -> &mut Self
302    where
303        T: Into<Option<Duration>>,
304    {
305        self.idle_timeout = Some(value.into());
306        self
307    }
308
309    /// Get the idle duration before closing a connection, if set
310    pub fn get_idle_timeout(&self) -> Option<Option<Duration>> {
311        self.idle_timeout
312    }
313
314    /// Set the maximum amount of time to spend waiting for acquiring a connection
315    pub fn acquire_timeout(&mut self, value: Duration) -> &mut Self {
316        self.acquire_timeout = Some(value);
317        self
318    }
319
320    /// Get the maximum amount of time to spend waiting for acquiring a connection
321    pub fn get_acquire_timeout(&self) -> Option<Duration> {
322        self.acquire_timeout
323    }
324
325    /// Set the maximum lifetime of individual connections.
326    pub fn max_lifetime<T>(&mut self, lifetime: T) -> &mut Self
327    where
328        T: Into<Option<Duration>>,
329    {
330        self.max_lifetime = Some(lifetime.into());
331        self
332    }
333
334    /// Get the maximum lifetime of individual connections, if set
335    pub fn get_max_lifetime(&self) -> Option<Option<Duration>> {
336        self.max_lifetime
337    }
338
339    /// Enable SQLx statement logging (default true)
340    pub fn sqlx_logging(&mut self, value: bool) -> &mut Self {
341        self.sqlx_logging = value;
342        self
343    }
344
345    /// Get whether SQLx statement logging is enabled
346    pub fn get_sqlx_logging(&self) -> bool {
347        self.sqlx_logging
348    }
349
350    /// Enable recording `db.statement` in tracing spans (default true).
351    pub fn record_stmt_in_spans(&mut self, value: bool) -> &mut Self {
352        self.record_stmt_in_spans = value;
353        self
354    }
355
356    /// Get whether `db.statement` recording in tracing spans is enabled.
357    pub fn get_record_stmt_in_spans(&self) -> bool {
358        self.record_stmt_in_spans
359    }
360
361    /// Set SQLx statement logging level (default INFO).
362    /// (ignored if `sqlx_logging` is `false`)
363    pub fn sqlx_logging_level(&mut self, level: log::LevelFilter) -> &mut Self {
364        self.sqlx_logging_level = level;
365        self
366    }
367
368    /// Set SQLx slow statements logging level and duration threshold (default `LevelFilter::Off`).
369    /// (ignored if `sqlx_logging` is `false`)
370    pub fn sqlx_slow_statements_logging_settings(
371        &mut self,
372        level: log::LevelFilter,
373        duration: Duration,
374    ) -> &mut Self {
375        self.sqlx_slow_statements_logging_level = level;
376        self.sqlx_slow_statements_logging_threshold = duration;
377        self
378    }
379
380    /// Get the level of SQLx statement logging
381    pub fn get_sqlx_logging_level(&self) -> log::LevelFilter {
382        self.sqlx_logging_level
383    }
384
385    /// Get the SQLx slow statements logging settings
386    pub fn get_sqlx_slow_statements_logging_settings(&self) -> (log::LevelFilter, Duration) {
387        (
388            self.sqlx_slow_statements_logging_level,
389            self.sqlx_slow_statements_logging_threshold,
390        )
391    }
392
393    /// set key for sqlcipher
394    pub fn sqlcipher_key<T>(&mut self, value: T) -> &mut Self
395    where
396        T: Into<Cow<'static, str>>,
397    {
398        self.sqlcipher_key = Some(value.into());
399        self
400    }
401
402    /// Set schema search path (PostgreSQL only)
403    pub fn set_schema_search_path<T>(&mut self, schema_search_path: T) -> &mut Self
404    where
405        T: Into<String>,
406    {
407        self.schema_search_path = Some(schema_search_path.into());
408        self
409    }
410
411    /// Set application name (PostgreSQL only)
412    pub fn set_application_name<T>(&mut self, application_name: T) -> &mut Self
413    where
414        T: Into<String>,
415    {
416        self.application_name = Some(application_name.into());
417        self
418    }
419
420    /// Set the statement timeout (PostgreSQL only).
421    ///
422    /// This sets the PostgreSQL `statement_timeout` parameter via the connection options,
423    /// causing the server to abort any statement that exceeds the specified duration.
424    /// The timeout is applied at connection time and does not require an extra roundtrip.
425    ///
426    /// Has no effect on MySQL or SQLite connections.
427    pub fn statement_timeout(&mut self, value: Duration) -> &mut Self {
428        self.statement_timeout = Some(value);
429        self
430    }
431
432    /// Get the statement timeout, if set
433    pub fn get_statement_timeout(&self) -> Option<Duration> {
434        self.statement_timeout
435    }
436
437    /// If true, the connection will be pinged upon acquiring from the pool (default true).
438    pub fn test_before_acquire(&mut self, value: bool) -> &mut Self {
439        self.test_before_acquire = value;
440        self
441    }
442
443    /// If set to `true`, the db connection pool will be created using SQLx's
444    /// [connect_lazy](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html#method.connect_lazy) method.
445    pub fn connect_lazy(&mut self, value: bool) -> &mut Self {
446        self.connect_lazy = value;
447        self
448    }
449
450    /// Get whether DB connections will be established when the pool is created or only as needed.
451    pub fn get_connect_lazy(&self) -> bool {
452        self.connect_lazy
453    }
454
455    /// Set a callback function that will be called after a new connection is established.
456    pub fn after_connect<F>(&mut self, f: F) -> &mut Self
457    where
458        F: Fn(DatabaseConnection) -> BoxFuture<'static, Result<(), DbErr>> + 'static,
459    {
460        self.after_connect = Some(Arc::new(f));
461
462        self
463    }
464
465    #[cfg(feature = "sqlx-mysql")]
466    #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-mysql")))]
467    /// Apply a function to modify the underlying [`MySqlConnectOptions`] before
468    /// creating the connection pool.
469    pub fn map_sqlx_mysql_opts<F>(&mut self, f: F) -> &mut Self
470    where
471        F: Fn(MySqlConnectOptions) -> MySqlConnectOptions + 'static,
472    {
473        self.mysql_opts_fn = Some(Arc::new(f));
474        self
475    }
476
477    #[cfg(feature = "sqlx-mysql")]
478    #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-mysql")))]
479    /// Apply a function to modify the underlying [`sqlx::pool::PoolOptions<sqlx::MySql>`]
480    /// before creating the connection pool.
481    pub fn map_sqlx_mysql_pool_opts<F>(&mut self, f: F) -> &mut Self
482    where
483        F: Fn(sqlx::pool::PoolOptions<sqlx::MySql>) -> sqlx::pool::PoolOptions<sqlx::MySql>
484            + 'static,
485    {
486        self.mysql_pool_opts_fn = Some(Arc::new(f));
487        self
488    }
489
490    #[cfg(feature = "sqlx-postgres")]
491    #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-postgres")))]
492    /// Apply a function to modify the underlying [`PgConnectOptions`] before
493    /// creating the connection pool.
494    pub fn map_sqlx_postgres_opts<F>(&mut self, f: F) -> &mut Self
495    where
496        F: Fn(PgConnectOptions) -> PgConnectOptions + 'static,
497    {
498        self.pg_opts_fn = Some(Arc::new(f));
499        self
500    }
501
502    #[cfg(feature = "sqlx-postgres")]
503    #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-postgres")))]
504    /// Apply a function to modify the underlying [`sqlx::pool::PoolOptions<sqlx::Postgres>`]
505    /// before creating the connection pool.
506    pub fn map_sqlx_postgres_pool_opts<F>(&mut self, f: F) -> &mut Self
507    where
508        F: Fn(sqlx::pool::PoolOptions<sqlx::Postgres>) -> sqlx::pool::PoolOptions<sqlx::Postgres>
509            + 'static,
510    {
511        self.pg_pool_opts_fn = Some(Arc::new(f));
512        self
513    }
514
515    #[cfg(feature = "sqlx-sqlite")]
516    #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-sqlite")))]
517    /// Apply a function to modify the underlying [`SqliteConnectOptions`] before
518    /// creating the connection pool.
519    pub fn map_sqlx_sqlite_opts<F>(&mut self, f: F) -> &mut Self
520    where
521        F: Fn(SqliteConnectOptions) -> SqliteConnectOptions + 'static,
522    {
523        self.sqlite_opts_fn = Some(Arc::new(f));
524        self
525    }
526
527    #[cfg(feature = "sqlx-sqlite")]
528    #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-sqlite")))]
529    /// Apply a function to modify the underlying [`sqlx::pool::PoolOptions<sqlx::Sqlite>`]
530    /// before creating the connection pool.
531    pub fn map_sqlx_sqlite_pool_opts<F>(&mut self, f: F) -> &mut Self
532    where
533        F: Fn(sqlx::pool::PoolOptions<sqlx::Sqlite>) -> sqlx::pool::PoolOptions<sqlx::Sqlite>
534            + 'static,
535    {
536        self.sqlite_pool_opts_fn = Some(Arc::new(f));
537        self
538    }
539}