Skip to main content

sea_orm/database/
mod.rs

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