Skip to main content

sea_orm/database/
mod.rs

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