a06673b
pub use sea_orm;
use crate::config::{AppDatabaseConnConfig, AppDatabaseSettings};
pub use sea_orm::DbErr as SqlxError;
use sea_orm::{
ConnectOptions, ConnectionTrait, Database, DatabaseBackend, DatabaseConnection, DbErr,
Statement,
};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use tracing::error;
#[derive(Debug, thiserror::Error)]
pub enum DbError {
#[error("invalid database configuration: {0}")]
InvalidConfig(String),
#[error("failed to build primary pool: {0}")]
PrimaryPool(#[source] DbErr),
#[error("failed to build replica pool ({index}): {source}")]
ReplicaPool {
index: usize,
#[source]
source: DbErr,
},
#[error("failed to reach {pool}: {source}")]
ConnectivityCheck {
pool: String,
#[source]
source: DbErr,
},
}
#[derive(Clone)]
pub struct Db {
inner: Arc<DbInner>,
}
struct DbInner {
primary: DatabaseConnection,
replicas: Vec<DatabaseConnection>,
rr_counter: AtomicUsize,
}
impl Db {
pub async fn from_config(cfg: &AppDatabaseSettings) -> Result<Self, DbError> {
validate_config(cfg)?;
let primary = build_connection(
&cfg.primary,
cfg.sqlx_logging,
cfg.query_timeout,
cfg.slow_query_threshold_ms,
)
.await
.map_err(DbError::PrimaryPool)?;
check_connectivity(&primary, "primary").await?;
let mut replicas = Vec::with_capacity(cfg.replicas.len());
for (i, rc) in cfg.replicas.iter().enumerate() {
match build_connection(
rc,
cfg.sqlx_logging,
cfg.query_timeout,
cfg.slow_query_threshold_ms,
)
.await
{
Ok(conn) => {
check_connectivity(&conn, &format!("replica[{i}]")).await?;
replicas.push(conn);
}
Err(e) => {
error!(index = i, error = %e, "failed to build replica pool");
return Err(DbError::ReplicaPool {
index: i,
source: e,
});
}
}
}
Ok(Self {
inner: Arc::new(DbInner {
primary,
replicas,
rr_counter: AtomicUsize::new(0),
}),
})
}
pub fn from_connection(connection: DatabaseConnection) -> Self {
Self {
inner: Arc::new(DbInner {
primary: connection,
replicas: Vec::new(),
rr_counter: AtomicUsize::new(0),
}),
}
}
pub fn writer(&self) -> DatabaseConnection {
self.inner.primary.clone()
}
pub fn reader(&self) -> DatabaseConnection {
if self.inner.replicas.is_empty() {
return self.inner.primary.clone();
}
let idx = next_reader_index(self.inner.replicas.len(), &self.inner.rr_counter);
self.inner.replicas[idx].clone()
}
pub fn primary_connection(&self) -> &DatabaseConnection {
&self.inner.primary
}
pub fn replica_connections(&self) -> &[DatabaseConnection] {
&self.inner.replicas
}
pub async fn close(&self) {
let _ = self.inner.primary.clone().close().await;
for replica in &self.inner.replicas {
let _ = replica.clone().close().await;
}
}
}
fn validate_config(cfg: &AppDatabaseSettings) -> Result<(), DbError> {
let connections = std::iter::once(&cfg.primary).chain(&cfg.replicas);
let mut total = 0u32;
for (index, connection) in connections.enumerate() {
if connection.url.is_empty() {
return Err(DbError::InvalidConfig(format!(
"pool {index} url must be configured"
)));
}
if connection.max_connections == 0 {
return Err(DbError::InvalidConfig(format!(
"pool {index} max_connections must be greater than zero"
)));
}
if connection.min_connections > connection.max_connections {
return Err(DbError::InvalidConfig(format!(
"pool {index} min_connections exceeds max_connections"
)));
}
total = total
.checked_add(connection.max_connections)
.ok_or_else(|| {
DbError::InvalidConfig("total connection count overflows u32".to_string())
})?;
}
if total > cfg.max_connections_total {
return Err(DbError::InvalidConfig(format!(
"configured pools require {total} connections, exceeding max_connections_total {}",
cfg.max_connections_total
)));
}
Ok(())
}
#[allow(clippy::cast_possible_truncation)]
async fn build_connection(
cfg: &AppDatabaseConnConfig,
sqlx_logging: bool,
query_timeout: u64,
slow_query_threshold_ms: u64,
) -> Result<DatabaseConnection, DbErr> {
let app_name = cfg.application_name.clone();
let schema = cfg.schema.clone();
let mut opts = ConnectOptions::new(&cfg.url);
opts.connect_lazy(true);
opts.max_connections(cfg.max_connections)
.min_connections(cfg.min_connections);
let acquire_timeout = cfg.acquire_timeout.unwrap_or(5);
opts.acquire_timeout(Duration::from_secs(acquire_timeout));
if let Some(lt) = cfg.max_lifetime {
opts.max_lifetime(Duration::from_secs(lt));
}
if let Some(it) = cfg.idle_timeout {
opts.idle_timeout(Duration::from_secs(it));
}
opts.sqlx_logging(sqlx_logging);
opts.test_before_acquire(true);
if sqlx_logging && slow_query_threshold_ms > 0 {
opts.sqlx_slow_statements_logging_settings(
log::LevelFilter::Warn,
Duration::from_millis(slow_query_threshold_ms),
);
}
opts.map_sqlx_postgres_opts(move |pg_opts: sea_orm::sqlx::postgres::PgConnectOptions| {
let mut pg_opts = pg_opts;
if let Some(ref name) = app_name {
pg_opts = pg_opts.application_name(name);
}
if let Some(ref sch) = schema {
pg_opts = pg_opts.options([("search_path", sch.as_str())]);
}
if query_timeout > 0 {
let stmt_timeout = format!("{query_timeout}s");
pg_opts = pg_opts.options([("statement_timeout", stmt_timeout.as_str())]);
}
pg_opts
});
Database::connect(opts).await
}
async fn check_connectivity(conn: &DatabaseConnection, pool: &str) -> Result<(), DbError> {
let stmt = Statement::from_string(DatabaseBackend::Postgres, "SELECT 1");
conn.execute_raw(stmt)
.await
.map(|_| ())
.map_err(|source| DbError::ConnectivityCheck {
pool: pool.to_string(),
source,
})
}
fn next_reader_index(len: usize, counter: &AtomicUsize) -> usize {
debug_assert!(len > 0);
counter.fetch_add(1, Ordering::Relaxed) % len
}
pub async fn rollback_or_log(txn: sea_orm::DatabaseTransaction) {
if let Err(error) = txn.rollback().await {
error!(error = %error, "transaction_rollback_failed");
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::{AppDatabaseConnConfig, AppDatabaseSettings};
fn conn_config(url: &str, max: u32, min: u32) -> AppDatabaseConnConfig {
AppDatabaseConnConfig {
url: url.to_string(),
max_connections: max,
min_connections: min,
..AppDatabaseConnConfig::default()
}
}
fn settings(
primary: AppDatabaseConnConfig,
replicas: Vec<AppDatabaseConnConfig>,
total: u32,
) -> AppDatabaseSettings {
AppDatabaseSettings {
primary,
replicas,
max_connections_total: total,
..AppDatabaseSettings::default()
}
}
#[test]
fn validate_config_rejects_empty_url() {
let cfg = settings(conn_config("", 20, 2), vec![], 100);
let err = validate_config(&cfg).unwrap_err().to_string();
assert!(err.contains("url must be configured"), "{err}");
}
#[test]
fn validate_config_rejects_zero_max_connections() {
let cfg = settings(conn_config("postgres://db", 0, 0), vec![], 100);
let err = validate_config(&cfg).unwrap_err().to_string();
assert!(
err.contains("max_connections must be greater than zero"),
"{err}"
);
}
#[test]
fn validate_config_rejects_min_over_max() {
let cfg = settings(conn_config("postgres://db", 5, 10), vec![], 100);
let err = validate_config(&cfg).unwrap_err().to_string();
assert!(err.contains("min_connections exceeds"), "{err}");
}
#[test]
fn validate_config_rejects_total_overflow() {
let primary = conn_config("postgres://db", u32::MAX, 2);
let replica = conn_config("postgres://db2", u32::MAX, 2);
let cfg = settings(primary, vec![replica], u32::MAX);
let err = validate_config(&cfg).unwrap_err().to_string();
assert!(err.contains("overflows u32"), "{err}");
}
#[test]
fn validate_config_rejects_total_exceeding_budget() {
let primary = conn_config("postgres://db", 80, 2);
let replica = conn_config("postgres://db2", 30, 2);
let cfg = settings(primary, vec![replica], 100);
let err = validate_config(&cfg).unwrap_err().to_string();
assert!(err.contains("exceeding max_connections_total"), "{err}");
}
#[test]
fn validate_config_accepts_valid_configuration() {
let primary = conn_config("postgres://db", 20, 2);
let replica = conn_config("postgres://db2", 10, 2);
let cfg = settings(primary, vec![replica], 100);
assert!(validate_config(&cfg).is_ok());
}
#[test]
fn next_reader_index_round_robins() {
let counter = AtomicUsize::new(0);
let idxs: Vec<usize> = (0..6).map(|_| next_reader_index(3, &counter)).collect();
assert_eq!(idxs, vec![0, 1, 2, 0, 1, 2]);
}
#[test]
fn next_reader_index_single_pool_is_stable() {
let counter = AtomicUsize::new(0);
let idxs: Vec<usize> = (0..5).map(|_| next_reader_index(1, &counter)).collect();
assert_eq!(idxs, vec![0, 0, 0, 0, 0]);
}
}