sea_orm/database/
mod.rs

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