Skip to content

Commit

Permalink
feat: Support for lazy database connections (#2268)
Browse files Browse the repository at this point in the history
Add support for creating DB connection pools without establishing
connections up front. This is already supported by SQLx via the
[Pool::connect_lazy](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html#method.connect_lazy)
method.

This PR adds a new `connect_lazy` option to `database::ConnectOptions`.
If set to `true`, the SQLx `Pool` will be created using the
`Pool::connect_lazy_with` method; otherwise, the `Pool::connect_with`
method will be used (e.g., the existing behavior). This "lazy" behavior
is implemented for each DB variant (Postgres/MySQL/SQLite).

This was discussed previously
[here](#1645), but it
appears support was never added to SEA ORM directly.
  • Loading branch information
spencewenski authored Aug 21, 2024
1 parent 0b956dc commit 4553160
Show file tree
Hide file tree
Showing 4 changed files with 72 additions and 35 deletions.
20 changes: 19 additions & 1 deletion src/database/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,10 +63,15 @@ pub struct ConnectOptions {
/// Schema search path (PostgreSQL only)
pub(crate) schema_search_path: Option<String>,
pub(crate) test_before_acquire: bool,
/// Only establish connections to the DB as needed. If set to `true`, the db connection will
/// be created using SQLx's [connect_lazy](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html#method.connect_lazy)
/// method.
pub(crate) connect_lazy: bool,
}

impl Database {
/// Method to create a [DatabaseConnection] on a database
/// Method to create a [DatabaseConnection] on a database. This method will return an error
/// if the database is not available.
#[instrument(level = "trace", skip(opt))]
pub async fn connect<C>(opt: C) -> Result<DatabaseConnection, DbErr>
where
Expand Down Expand Up @@ -157,6 +162,7 @@ impl ConnectOptions {
sqlcipher_key: None,
schema_search_path: None,
test_before_acquire: true,
connect_lazy: false,
}
}

Expand Down Expand Up @@ -297,4 +303,16 @@ impl ConnectOptions {
self.test_before_acquire = value;
self
}

/// If set to `true`, the db connection pool will be created using SQLx's
/// [connect_lazy](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html#method.connect_lazy) method.
pub fn connect_lazy(&mut self, value: bool) -> &mut Self {
self.connect_lazy = value;
self
}

/// Get whether DB connections will be established when the pool is created or only as needed.
pub fn get_connect_lazy(&self) -> bool {
self.connect_lazy
}
}
24 changes: 15 additions & 9 deletions src/driver/sqlx_mysql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,15 +61,21 @@ impl SqlxMySqlConnector {
);
}
}
match options.sqlx_pool_options().connect_with(opt).await {
Ok(pool) => Ok(DatabaseConnection::SqlxMySqlPoolConnection(
SqlxMySqlPoolConnection {
pool,
metric_callback: None,
},
)),
Err(e) => Err(sqlx_error_to_conn_err(e)),
}
let pool = if options.connect_lazy {
options.sqlx_pool_options().connect_lazy_with(opt)
} else {
options
.sqlx_pool_options()
.connect_with(opt)
.await
.map_err(sqlx_error_to_conn_err)?
};
Ok(DatabaseConnection::SqlxMySqlPoolConnection(
SqlxMySqlPoolConnection {
pool,
metric_callback: None,
},
))
}
}

Expand Down
24 changes: 15 additions & 9 deletions src/driver/sqlx_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ impl SqlxPostgresConnector {
.schema_search_path
.as_ref()
.map(|schema| format!("SET search_path = {schema}"));
let lazy = options.connect_lazy;
let mut pool_options = options.sqlx_pool_options();
if let Some(sql) = set_search_path_sql {
pool_options = pool_options.after_connect(move |conn, _| {
Expand All @@ -76,15 +77,20 @@ impl SqlxPostgresConnector {
})
});
}
match pool_options.connect_with(opt).await {
Ok(pool) => Ok(DatabaseConnection::SqlxPostgresPoolConnection(
SqlxPostgresPoolConnection {
pool,
metric_callback: None,
},
)),
Err(e) => Err(sqlx_error_to_conn_err(e)),
}
let pool = if lazy {
pool_options.connect_lazy_with(opt)
} else {
pool_options
.connect_with(opt)
.await
.map_err(sqlx_error_to_conn_err)?
};
Ok(DatabaseConnection::SqlxPostgresPoolConnection(
SqlxPostgresPoolConnection {
pool,
metric_callback: None,
},
))
}
}

Expand Down
39 changes: 23 additions & 16 deletions src/driver/sqlx_sqlite.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,26 +66,33 @@ impl SqlxSqliteConnector {
);
}
}

if options.get_max_connections().is_none() {
options.max_connections(1);
}
match options.sqlx_pool_options().connect_with(opt).await {
Ok(pool) => {
let pool = SqlxSqlitePoolConnection {
pool,
metric_callback: None,
};

#[cfg(feature = "sqlite-use-returning-for-3_35")]
{
let version = get_version(&pool).await?;
ensure_returning_version(&version)?;
}

Ok(DatabaseConnection::SqlxSqlitePoolConnection(pool))
}
Err(e) => Err(sqlx_error_to_conn_err(e)),

let pool = if options.connect_lazy {
options.sqlx_pool_options().connect_lazy_with(opt)
} else {
options
.sqlx_pool_options()
.connect_with(opt)
.await
.map_err(sqlx_error_to_conn_err)?
};

let pool = SqlxSqlitePoolConnection {
pool,
metric_callback: None,
};

#[cfg(feature = "sqlite-use-returning-for-3_35")]
{
let version = get_version(&pool).await?;
ensure_returning_version(&version)?;
}

Ok(DatabaseConnection::SqlxSqlitePoolConnection(pool))
}
}

Expand Down

0 comments on commit 4553160

Please sign in to comment.