use axum::extract::FromRequestParts;
use diesel;
use diesel_async::AsyncPgConnection;
use diesel_async::pooled_connection::AsyncDieselConnectionManager;
use diesel_async::pooled_connection::deadpool::Pool;
use futures::FutureExt as _;
use std::any::Any;
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tracing::Instrument as _;
use crate::config::DatabaseConfig;
use crate::error::AutumnError;
#[cfg(not(feature = "sqlite"))]
pub type RuntimeConnection = diesel_async::AsyncPgConnection;
#[cfg(feature = "sqlite")]
pub type RuntimeConnection =
diesel_async::sync_connection_wrapper::SyncConnectionWrapper<diesel::SqliteConnection>;
#[cfg(not(feature = "sqlite"))]
pub type RuntimeBackend = diesel::pg::Pg;
#[cfg(feature = "sqlite")]
pub type RuntimeBackend = diesel::sqlite::Sqlite;
pub type CommitCallback = Box<
dyn FnOnce() -> Pin<Box<dyn Future<Output = crate::AutumnResult<()>> + Send + 'static>>
+ Send
+ 'static,
>;
tokio::task_local! {
pub static AFTER_COMMIT_REGISTRY: Arc<Mutex<Vec<CommitCallback>>>;
}
#[derive(Default, Debug)]
pub(crate) struct RequestDbTimings {
pub(crate) total_us: AtomicU64,
pub(crate) query_count: std::sync::atomic::AtomicUsize,
}
tokio::task_local! {
pub(crate) static REQUEST_DB_TIMINGS: Arc<RequestDbTimings>;
}
tokio::task_local! {
#[cfg(feature = "db")]
pub(crate) static REQUEST_QUERY_CAPTURE: Arc<Mutex<Vec<crate::inspector::QueryRecord>>>;
}
#[cfg(feature = "db")]
fn strip_bind_annotation(sql: &str) -> &str {
sql.rfind("-- binds:")
.map_or(sql, |idx| sql[..idx].trim_end())
}
pub(crate) fn record_request_db_query(elapsed: Duration, sql: Option<&str>) {
let _ = REQUEST_DB_TIMINGS.try_with(|t| {
let micros = u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX);
t.total_us.fetch_add(micros, Ordering::Relaxed);
t.query_count
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
});
#[cfg(feature = "db")]
if let Some(sql) = sql {
let _ = REQUEST_QUERY_CAPTURE.try_with(|sink| {
if let Ok(mut list) = sink.lock() {
list.push(crate::inspector::QueryRecord {
sql: strip_bind_annotation(sql).to_owned(),
params: Vec::new(),
elapsed_ms: u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX),
location: String::new(),
});
}
});
}
}
#[cfg(feature = "db")]
pub(crate) fn request_db_timing_active() -> bool {
REQUEST_DB_TIMINGS.try_with(|_| ()).is_ok()
}
#[cfg(feature = "db")]
pub(crate) fn request_query_capture_active() -> bool {
REQUEST_QUERY_CAPTURE.try_with(|_| ()).is_ok()
}
#[cfg(feature = "db")]
#[derive(Default, Debug)]
pub(crate) struct RequestQueryTimer {
pending: Option<PendingQuery>,
}
#[cfg(feature = "db")]
#[derive(Debug)]
struct PendingQuery {
started_at: std::time::Instant,
sql: String,
}
#[cfg(feature = "db")]
impl RequestQueryTimer {
fn is_uncounted_statement(sql: &str) -> bool {
let mut tokens = sql.split_whitespace();
let Some(first) = tokens.next() else {
return false;
};
match first.to_ascii_uppercase().as_str() {
"BEGIN" | "COMMIT" | "ROLLBACK" | "SAVEPOINT" | "RELEASE" | "SET" => true,
"START" => tokens
.next()
.is_some_and(|second| second.eq_ignore_ascii_case("TRANSACTION")),
_ => false,
}
}
fn on_start(&mut self, now: std::time::Instant, sql: impl FnOnce() -> String) {
if !request_db_timing_active() && !request_query_capture_active() {
self.pending = None;
return;
}
let sql = sql();
self.pending = if Self::is_uncounted_statement(&sql) {
None
} else {
Some(PendingQuery {
started_at: now,
sql,
})
};
}
fn on_finish(&mut self, now: std::time::Instant) {
if let Some(p) = self.pending.take() {
record_request_db_query(now.saturating_duration_since(p.started_at), Some(&p.sql));
}
}
}
#[cfg(feature = "db")]
impl diesel::connection::Instrumentation for RequestQueryTimer {
fn on_connection_event(&mut self, event: diesel::connection::InstrumentationEvent<'_>) {
use diesel::connection::InstrumentationEvent;
match event {
InstrumentationEvent::StartQuery { query, .. } => {
self.on_start(std::time::Instant::now(), || query.to_string());
}
InstrumentationEvent::FinishQuery { .. } => self.on_finish(std::time::Instant::now()),
_ => {}
}
}
}
pub static AFTER_COMMIT_FAILURES_TOTAL: AtomicU64 = AtomicU64::new(0);
pub(crate) fn record_after_commit_failure() -> u64 {
AFTER_COMMIT_FAILURES_TOTAL.fetch_add(1, Ordering::Relaxed) + 1
}
pub static TX_RETRIES_TOTAL: AtomicU64 = AtomicU64::new(0);
pub static TX_RETRY_EXHAUSTED_TOTAL: AtomicU64 = AtomicU64::new(0);
#[cfg_attr(feature = "sqlite", allow(dead_code))]
pub(crate) fn record_tx_retry() -> u64 {
TX_RETRIES_TOTAL.fetch_add(1, Ordering::Relaxed) + 1
}
#[cfg_attr(feature = "sqlite", allow(dead_code))]
pub(crate) fn record_tx_retry_exhausted() -> u64 {
TX_RETRY_EXHAUSTED_TOTAL.fetch_add(1, Ordering::Relaxed) + 1
}
const NESTED_TX_MESSAGE: &str = "Nested Db::tx calls are not supported; use \
autumn_web::db::savepoint(conn, ..) inside the closure for a same-connection savepoint";
pub(crate) fn reject_ambient_after_commit_registry_for_tx() -> Result<(), AutumnError> {
if AFTER_COMMIT_REGISTRY.try_with(|_| ()).is_ok() {
return Err(AutumnError::bad_request_msg(NESTED_TX_MESSAGE));
}
Ok(())
}
pub(crate) fn spawn_committed_after_commit_callbacks(
callbacks: Vec<CommitCallback>,
) -> Option<tokio::task::JoinHandle<()>> {
if callbacks.is_empty() {
return None;
}
Some(tokio::task::spawn(async move {
for cb in callbacks {
let result = match std::panic::catch_unwind(AssertUnwindSafe(cb)) {
Ok(callback) => AssertUnwindSafe(callback).catch_unwind().await,
Err(panic) => Err(panic),
};
match result {
Ok(Ok(())) => {}
Ok(Err(e)) => {
let failures_total = record_after_commit_failure();
tracing::error!(
autumn.after_commit.failures_total = failures_total,
"after_commit callback failed (tx already committed): {e}"
);
}
Err(panic) => {
let failures_total = record_after_commit_failure();
let panic = after_commit_panic_message(&*panic);
tracing::error!(
autumn.after_commit.failures_total = failures_total,
"after_commit callback panicked (tx already committed): {panic}"
);
}
}
}
}))
}
fn after_commit_panic_message(payload: &(dyn Any + Send)) -> String {
match (
payload.downcast_ref::<&'static str>(),
payload.downcast_ref::<String>(),
) {
(Some(message), _) => (*message).to_owned(),
(_, Some(message)) => message.clone(),
(None, None) => "non-string panic payload".to_owned(),
}
}
pub async fn register_after_commit<F, Fut>(f: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = crate::AutumnResult<()>> + Send + 'static,
{
let mut f_opt = Some(f);
AFTER_COMMIT_REGISTRY
.try_with(|registry| {
let f = f_opt.take().expect("closure only entered once");
let boxed: CommitCallback = Box::new(move || Box::pin(f()));
registry.lock().expect("registry lock").push(boxed);
})
.ok();
if let Some(f) = f_opt {
tracing::debug!("register_after_commit: no active transaction; running callback eagerly");
if let Err(e) = f().await {
let failures_total = record_after_commit_failure();
tracing::error!(
autumn.after_commit.failures_total = failures_total,
"register_after_commit eager callback failed: {e}"
);
}
}
}
pub trait DbState {
fn pool(&self) -> Option<&Pool<RuntimeConnection>>;
fn metrics(&self) -> Option<&crate::middleware::MetricsCollector> {
None
}
fn replica_pool(&self) -> Option<&Pool<RuntimeConnection>> {
None
}
fn read_pool(&self) -> Option<&Pool<RuntimeConnection>> {
self.replica_pool().or_else(|| self.pool())
}
fn shards(&self) -> Option<&crate::sharding::ShardSet> {
None
}
fn db_interceptors(
&self,
) -> Vec<std::sync::Arc<dyn crate::interceptor::DbConnectionInterceptor>> {
Vec::new()
}
fn statement_timeout(&self) -> Option<std::time::Duration> {
None
}
fn slow_query_threshold(&self) -> std::time::Duration {
std::time::Duration::from_millis(500)
}
}
fn consume_estring_body(chars: &mut std::iter::Peekable<std::str::Chars<'_>>) {
loop {
match chars.next() {
None => break,
Some('\'') => {
if chars.peek() == Some(&'\'') {
chars.next(); } else {
break;
}
}
Some('\\') => {
chars.next(); }
Some(_) => {}
}
}
}
fn consume_dollar_quoted_body(chars: &mut std::iter::Peekable<std::str::Chars<'_>>, tag: &str) {
let closing: Vec<char> = format!("${tag}$").chars().collect();
let clen = closing.len();
let mut match_count = 0usize;
for sc in chars.by_ref() {
if sc == closing[match_count] {
match_count += 1;
if match_count == clen {
break; }
} else {
match_count = 0;
if sc == closing[0] {
match_count = 1;
}
}
}
}
#[inline]
const fn is_separator(c: char) -> bool {
matches!(
c,
' ' | '\t' | '\n' | '=' | '<' | '>' | '!' | '+' | '-' | '*' | '/' | '%' | '(' | ',' )
}
#[must_use]
pub fn scrub_sql(sql: &str) -> String {
let mut out = String::with_capacity(sql.len());
let mut prev_is_sep = true;
let mut chars = sql.chars().peekable();
while let Some(c) = chars.next() {
if (c == 'E' || c == 'e') && chars.peek() == Some(&'\'') {
chars.next(); out.push_str("'?'");
prev_is_sep = false;
consume_estring_body(&mut chars);
continue;
}
if c == '\'' {
out.push_str("'?'");
prev_is_sep = false;
loop {
match chars.next() {
None => break,
Some('\'') => {
if chars.peek() == Some(&'\'') {
chars.next();
} else {
break;
}
}
Some(_) => {}
}
}
continue;
}
if c == '$' {
let next_ch = chars.peek().copied();
if next_ch.is_some_and(|nc| nc.is_ascii_digit()) {
out.push('$');
prev_is_sep = false;
while chars.peek().is_some_and(char::is_ascii_digit) {
if let Some(d) = chars.next() {
out.push(d);
}
}
continue;
}
let mut tag = String::new();
let mut found_closing_dollar = false;
if next_ch == Some('$') {
chars.next();
found_closing_dollar = true;
} else if next_ch.is_some_and(|nc| nc.is_alphabetic() || nc == '_') {
while let Some(&tc) = chars.peek() {
if tc == '$' {
chars.next(); found_closing_dollar = true;
break;
} else if tc.is_alphanumeric() || tc == '_' {
tag.push(tc);
chars.next();
} else {
break;
}
}
}
if found_closing_dollar {
out.push_str("'?'");
prev_is_sep = false;
consume_dollar_quoted_body(&mut chars, &tag);
} else {
out.push('$');
out.push_str(&tag);
prev_is_sep = false;
}
continue;
}
let is_leading_dot =
c == '.' && prev_is_sep && chars.peek().is_some_and(char::is_ascii_digit);
if (c.is_ascii_digit() && prev_is_sep) || is_leading_dot {
out.push('?');
if is_leading_dot {
chars.next(); }
while chars
.peek()
.is_some_and(|d| d.is_ascii_digit() || *d == '.' || *d == '_')
{
chars.next();
}
if chars.peek().is_some_and(|e| *e == 'e' || *e == 'E') {
chars.next(); if chars.peek().is_some_and(|s| *s == '+' || *s == '-') {
chars.next(); }
while chars.peek().is_some_and(char::is_ascii_digit) {
chars.next();
}
}
prev_is_sep = false;
continue;
}
out.push(c);
prev_is_sep = is_separator(c);
}
out
}
pub async fn run_instrumented<F, Fut, T>(
sql: &str,
route_key: &str,
slow_threshold: std::time::Duration,
metrics: &crate::middleware::metrics::MetricsCollector,
query: F,
) -> Result<T, AutumnError>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<T, diesel::result::Error>>,
{
let start = std::time::Instant::now();
let result = query().await;
let elapsed = start.elapsed();
let elapsed_ms = u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX);
let verb = sql.split_whitespace().next().unwrap_or("?");
let metric_key = format!("{route_key} {verb}");
metrics.record_db_query(&metric_key, elapsed_ms);
if elapsed >= slow_threshold {
let fingerprint = scrub_sql(sql);
tracing::warn!(
route = %route_key,
sql = %fingerprint,
duration_ms = elapsed_ms,
"slow database query"
);
}
result.map_err(|db_err| {
if is_query_canceled(&db_err) {
tracing::warn!(
route = %route_key,
duration_ms = elapsed_ms,
"database query cancelled: statement_timeout exceeded"
);
AutumnError::query_timeout(format!(
"Database query timed out after {elapsed_ms}ms (statement_timeout exceeded)"
))
} else {
AutumnError::from(db_err)
}
})
}
fn source_chain_has_sqlstate(
err: &(dyn std::error::Error + 'static),
predicate: impl Fn(&tokio_postgres::error::SqlState) -> bool,
) -> bool {
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(err);
while let Some(e) = source {
if e.downcast_ref::<tokio_postgres::Error>()
.and_then(tokio_postgres::Error::code)
.is_some_and(&predicate)
{
return true;
}
if e.downcast_ref::<tokio_postgres::error::DbError>()
.is_some_and(|db_err| predicate(db_err.code()))
{
return true;
}
source = e.source();
}
false
}
fn is_query_canceled(err: &diesel::result::Error) -> bool {
let err_str = err.to_string().to_lowercase();
if err_str.contains("57014")
|| err_str.contains("query_canceled")
|| err_str.contains("canceling statement due to statement timeout")
|| err_str.contains("statement timeout")
|| err_str.contains("query canceled")
{
return true;
}
source_chain_has_sqlstate(err, |state| {
*state == tokio_postgres::error::SqlState::QUERY_CANCELED
})
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum PoolError {
#[error(transparent)]
Build(#[from] diesel_async::pooled_connection::deadpool::BuildError),
#[error("{0}")]
UnsupportedBackend(String),
}
#[derive(Clone)]
pub struct DatabaseTopology {
primary: Pool<RuntimeConnection>,
replica: Option<Pool<RuntimeConnection>>,
migration_url: Option<String>,
}
impl DatabaseTopology {
#[must_use]
pub const fn from_pools(
primary: Pool<RuntimeConnection>,
replica: Option<Pool<RuntimeConnection>>,
) -> Self {
Self {
primary,
replica,
migration_url: None,
}
}
#[must_use]
pub const fn primary_only(primary: Pool<RuntimeConnection>) -> Self {
Self {
primary,
replica: None,
migration_url: None,
}
}
#[must_use]
pub fn with_migration_url(mut self, url: Option<String>) -> Self {
self.migration_url = url;
self
}
#[must_use]
pub fn migration_url(&self) -> Option<&str> {
self.migration_url.as_deref()
}
#[must_use]
pub const fn primary(&self) -> &Pool<RuntimeConnection> {
&self.primary
}
#[must_use]
pub const fn replica(&self) -> Option<&Pool<RuntimeConnection>> {
self.replica.as_ref()
}
#[must_use]
pub fn read(&self) -> &Pool<RuntimeConnection> {
self.replica.as_ref().unwrap_or(&self.primary)
}
}
fn build_pool(
url: &str,
pool_size: usize,
connect_timeout_secs: u64,
) -> Result<Pool<RuntimeConnection>, PoolError> {
#[cfg(feature = "sqlite")]
{
build_sqlite_pool(url, pool_size, connect_timeout_secs)
}
#[cfg(not(feature = "sqlite"))]
if crate::config::DatabaseBackend::detect(url) == Some(crate::config::DatabaseBackend::Sqlite) {
return Err(PoolError::UnsupportedBackend(format!(
"SQLite is a recognized database backend but its runtime pool is only available in \
a build of autumn-web compiled with `--features sqlite`; this is a default \
(Postgres) build (target: {url:?})"
)));
}
#[cfg(not(feature = "sqlite"))]
{
let timeout = Duration::from_secs(connect_timeout_secs);
let manager = match tls::TlsPosture::from_database_url(url) {
tls::TlsPosture::Off => AsyncDieselConnectionManager::<AsyncPgConnection>::new(url),
posture => {
let mut config =
diesel_async::pooled_connection::ManagerConfig::<AsyncPgConnection>::default();
config.custom_setup = tls::setup_callback(posture);
AsyncDieselConnectionManager::<AsyncPgConnection>::new_with_config(url, config)
}
};
Ok(Pool::builder(manager)
.max_size(pool_size.max(1))
.wait_timeout(Some(timeout))
.create_timeout(Some(timeout))
.runtime(deadpool::Runtime::Tokio1)
.build()?)
}
}
#[cfg(feature = "sqlite")]
fn normalize_sqlite_target(url: &str) -> String {
if url.starts_with("file:") {
return url.to_owned();
}
let rest = url
.strip_prefix("sqlite://")
.or_else(|| url.strip_prefix("sqlite:"))
.unwrap_or(url);
if rest.is_empty() || rest == ":memory:" {
return String::from(":memory:");
}
rest.to_owned()
}
#[cfg(feature = "sqlite")]
fn sqlite_target_is_memory(target: &str) -> bool {
if target.contains("cache=shared") {
return false;
}
target == ":memory:"
|| target == "file::memory:"
|| target.starts_with("file::memory:?")
|| target.contains("mode=memory")
}
#[cfg(feature = "sqlite")]
pub(crate) fn sqlite_target_is_any_in_memory(url: &str) -> bool {
let target = normalize_sqlite_target(url);
target == ":memory:"
|| target == "file::memory:"
|| target.starts_with("file::memory:?")
|| target.contains("mode=memory")
}
#[cfg(feature = "sqlite")]
fn sqlite_target_is_read_only(target: &str) -> bool {
let Some((_, query)) = target.split_once('?') else {
return false;
};
for pair in query.split('&') {
let (key, value) = pair.split_once('=').unwrap_or((pair, ""));
let (key, value) = (key.trim(), value.trim());
if key.eq_ignore_ascii_case("mode") && value.eq_ignore_ascii_case("ro") {
return true;
}
if key.eq_ignore_ascii_case("immutable")
&& (value == "1" || value.eq_ignore_ascii_case("true"))
{
return true;
}
}
false
}
#[cfg(feature = "sqlite")]
fn build_sqlite_pool(
url: &str,
pool_size: usize,
connect_timeout_secs: u64,
) -> Result<Pool<RuntimeConnection>, PoolError> {
if crate::config::DatabaseBackend::detect(url) == Some(crate::config::DatabaseBackend::Postgres)
{
return Err(PoolError::UnsupportedBackend(format!(
"this build of autumn-web targets SQLite (compiled with `--features sqlite`) but the \
configured database URL is a Postgres target; configure a `sqlite:` URL instead \
(target: {url:?})"
)));
}
let timeout = Duration::from_secs(connect_timeout_secs);
let target = normalize_sqlite_target(url);
let max_size = if sqlite_target_is_memory(&target) {
1
} else {
pool_size.max(1)
};
let mut config = diesel_async::pooled_connection::ManagerConfig::<RuntimeConnection>::default();
config.custom_setup = Box::new(|url: &str| {
use diesel_async::{AsyncConnection as _, SimpleAsyncConnection as _};
let url = url.to_owned();
async move {
let mut conn = RuntimeConnection::establish(&url).await?;
let pragmas = if sqlite_target_is_read_only(&url) {
"PRAGMA busy_timeout = 5000; \
PRAGMA foreign_keys = ON;"
} else {
"PRAGMA busy_timeout = 5000; \
PRAGMA journal_mode = WAL; \
PRAGMA synchronous = NORMAL; \
PRAGMA foreign_keys = ON;"
};
conn.batch_execute(pragmas)
.await
.map_err(diesel::ConnectionError::CouldntSetupConfiguration)?;
if let Err(e) = conn
.batch_execute(
"CREATE VIRTUAL TABLE temp.__autumn_fts5_probe USING fts5(x); \
DROP TABLE temp.__autumn_fts5_probe;",
)
.await
{
return Err(diesel::ConnectionError::CouldntSetupConfiguration(
diesel::result::Error::QueryBuilderError(
format!(
"SQLite FTS5 is not available in the linked SQLite library, but \
autumn-web full-text search (searchable repositories / the \
`--search` scaffold, issue #1910) requires it. Build with the \
bundled, FTS5-enabled SQLite by enabling autumn-web's `sqlite` \
feature (it turns on `libsqlite3-sys/bundled`, whose amalgamation \
defines SQLITE_ENABLE_FTS5). Underlying probe error: {e}"
)
.into(),
),
));
}
Ok(conn)
}
.boxed()
});
let manager =
AsyncDieselConnectionManager::<RuntimeConnection>::new_with_config(target, config);
Ok(Pool::builder(manager)
.max_size(max_size)
.wait_timeout(Some(timeout))
.create_timeout(Some(timeout))
.runtime(deadpool::Runtime::Tokio1)
.build()?)
}
pub fn create_pool(config: &DatabaseConfig) -> Result<Option<Pool<RuntimeConnection>>, PoolError> {
let Some(url) = config.effective_primary_url() else {
return Ok(None);
};
#[cfg(feature = "sqlite")]
reject_sqlite_statement_timeout(config.statement_timeout)?;
let pool = build_pool(
url,
config.effective_primary_pool_size(),
config.connect_timeout_secs,
)?;
Ok(Some(pool))
}
#[cfg(feature = "sqlite")]
pub(crate) fn reject_sqlite_statement_timeout(
statement_timeout: Option<Duration>,
) -> Result<(), PoolError> {
let Some(timeout) = statement_timeout.filter(|t| !t.is_zero()) else {
return Ok(());
};
Err(PoolError::UnsupportedBackend(format!(
"SQLite backend cannot enforce database.statement_timeout ({}ms): diesel's \
SqliteConnection exposes no interrupt/progress-handler hook through the async \
connection wrapper, so a runaway query cannot be aborted. Unset \
database.statement_timeout for the SQLite backend (busy_timeout already bounds \
lock waits), or run on Postgres. Tracking issue: #1996/#1910.",
timeout.as_millis()
)))
}
#[cfg(feature = "sqlite")]
fn reject_unusable_sqlite_replica(primary_url: &str, replica_url: &str) -> Result<(), PoolError> {
let replica_target = normalize_sqlite_target(replica_url);
let primary_target = normalize_sqlite_target(primary_url);
if sqlite_target_is_memory(&replica_target) || replica_target != primary_target {
return Err(PoolError::UnsupportedBackend(format!(
"SQLite does not support a separate read replica: replica_url {replica_url:?} \
is in-memory or differs from primary_url {primary_url:?}. Configure only a \
primary, or point the replica at the same database file as the primary."
)));
}
Ok(())
}
pub fn create_topology(config: &DatabaseConfig) -> Result<Option<DatabaseTopology>, PoolError> {
let Some(primary_url) = config.effective_primary_url() else {
return Ok(None);
};
#[cfg(feature = "sqlite")]
reject_sqlite_statement_timeout(config.statement_timeout)?;
let primary = build_pool(
primary_url,
config.effective_primary_pool_size(),
config.connect_timeout_secs,
)?;
#[cfg(feature = "sqlite")]
if let Some(replica_url) = config.replica_url.as_deref() {
reject_unusable_sqlite_replica(primary_url, replica_url)?;
}
let replica = config
.replica_url
.as_deref()
.map(|url| {
build_pool(
url,
config.effective_replica_pool_size(),
config.connect_timeout_secs,
)
})
.transpose()?;
Ok(Some(DatabaseTopology::from_pools(primary, replica)))
}
pub fn create_shard_topology(
shard: &crate::config::ShardConfig,
defaults: &DatabaseConfig,
) -> Result<DatabaseTopology, PoolError> {
#[cfg(feature = "sqlite")]
reject_sqlite_statement_timeout(defaults.statement_timeout)?;
let primary = build_pool(
&shard.primary_url,
shard.effective_primary_pool_size(defaults),
defaults.connect_timeout_secs,
)?;
#[cfg(feature = "sqlite")]
if let Some(replica_url) = shard.replica_url.as_deref() {
reject_unusable_sqlite_replica(&shard.primary_url, replica_url)?;
}
let replica = shard
.replica_url
.as_deref()
.map(|url| {
build_pool(
url,
shard.effective_replica_pool_size(defaults),
defaults.connect_timeout_secs,
)
})
.transpose()?;
Ok(DatabaseTopology::from_pools(primary, replica))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum IsolationLevel {
#[default]
ReadCommitted,
RepeatableRead,
Serializable,
}
impl IsolationLevel {
const fn as_str(self) -> &'static str {
match self {
Self::ReadCommitted => "read_committed",
Self::RepeatableRead => "repeatable_read",
Self::Serializable => "serializable",
}
}
}
const DEFAULT_TX_MAX_ATTEMPTS: u32 = 5;
#[derive(Debug, Clone, Copy)]
pub struct TxOptions {
pub isolation: IsolationLevel,
pub read_only: bool,
pub deferrable: bool,
pub max_attempts: u32,
pub initial_backoff: Duration,
pub max_backoff: Duration,
}
impl Default for TxOptions {
fn default() -> Self {
Self {
isolation: IsolationLevel::ReadCommitted,
read_only: false,
deferrable: false,
max_attempts: 1,
initial_backoff: Duration::from_millis(5),
max_backoff: Duration::from_millis(500),
}
}
}
impl TxOptions {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn read_committed() -> Self {
Self {
isolation: IsolationLevel::ReadCommitted,
..Self::default()
}
}
#[must_use]
pub fn repeatable_read() -> Self {
Self {
isolation: IsolationLevel::RepeatableRead,
max_attempts: DEFAULT_TX_MAX_ATTEMPTS,
..Self::default()
}
}
#[must_use]
pub fn serializable() -> Self {
Self {
isolation: IsolationLevel::Serializable,
max_attempts: DEFAULT_TX_MAX_ATTEMPTS,
..Self::default()
}
}
#[must_use]
pub const fn isolation(mut self, level: IsolationLevel) -> Self {
self.isolation = level;
self
}
#[must_use]
pub const fn read_only(mut self) -> Self {
self.read_only = true;
self
}
#[must_use]
pub const fn deferrable(mut self) -> Self {
self.deferrable = true;
self
}
#[must_use]
pub const fn max_attempts(mut self, attempts: u32) -> Self {
self.max_attempts = if attempts < 1 { 1 } else { attempts };
self
}
#[cfg_attr(feature = "sqlite", allow(dead_code))]
const fn effective_max_attempts(&self) -> u32 {
if self.max_attempts < 1 {
1
} else {
self.max_attempts
}
}
#[must_use]
pub const fn initial_backoff(mut self, delay: Duration) -> Self {
self.initial_backoff = delay;
self
}
#[must_use]
pub const fn max_backoff(mut self, delay: Duration) -> Self {
self.max_backoff = delay;
self
}
}
#[cfg_attr(feature = "sqlite", allow(dead_code))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RetryDecision {
Retry,
Stop,
}
#[cfg_attr(feature = "sqlite", allow(dead_code))]
const fn retry_decision(attempt: u32, max_attempts: u32, retryable: bool) -> RetryDecision {
if retryable && attempt < max_attempts {
RetryDecision::Retry
} else {
RetryDecision::Stop
}
}
#[cfg_attr(feature = "sqlite", allow(dead_code))]
fn retry_backoff_base(initial: Duration, max: Duration, attempt: u32) -> Duration {
let shift = attempt.saturating_sub(1).min(31);
let multiplier = 1u32 << shift;
initial.checked_mul(multiplier).unwrap_or(max).min(max)
}
#[cfg_attr(feature = "sqlite", allow(dead_code))]
fn retry_backoff_delay(initial: Duration, max: Duration, attempt: u32) -> Duration {
let base = retry_backoff_base(initial, max, attempt);
crate::cache::jittered_ttl(base, 0.2).min(max)
}
#[cfg_attr(feature = "sqlite", allow(dead_code))]
fn is_retryable_txn_error(err: &AutumnError) -> bool {
use tokio_postgres::error::SqlState;
fn is_retryable_sqlstate(state: &SqlState) -> bool {
*state == SqlState::T_R_SERIALIZATION_FAILURE || *state == SqlState::T_R_DEADLOCK_DETECTED
}
if let Some(diesel_err) = err.downcast_chain_ref::<diesel::result::Error>() {
return match diesel_err {
diesel::result::Error::DatabaseError(
diesel::result::DatabaseErrorKind::SerializationFailure,
_,
) => true,
diesel::result::Error::DatabaseError(
diesel::result::DatabaseErrorKind::Unknown,
info,
) => {
source_chain_has_sqlstate(diesel_err, is_retryable_sqlstate)
|| info
.message()
.trim()
.eq_ignore_ascii_case("deadlock detected")
}
_ => source_chain_has_sqlstate(diesel_err, is_retryable_sqlstate),
};
}
err.downcast_chain_ref::<tokio_postgres::Error>()
.and_then(tokio_postgres::Error::code)
.is_some_and(is_retryable_sqlstate)
|| err
.downcast_chain_ref::<tokio_postgres::error::DbError>()
.is_some_and(|db| is_retryable_sqlstate(db.code()))
}
#[doc(hidden)]
pub async fn scoped_transaction<'a, T, E, C, F>(conn: &'a mut C, f: F) -> Result<T, E>
where
C: diesel_async::AsyncConnection + Send,
T: Send + 'a,
E: From<diesel::result::Error> + Send + 'a,
F: for<'r> FnOnce(&'r mut C) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
+ Send
+ 'a,
{
use diesel_async::TransactionManager as _;
C::TransactionManager::begin_transaction(conn).await?;
match f(&mut *conn).await {
Ok(value) => {
C::TransactionManager::commit_transaction(conn).await?;
Ok(value)
}
Err(user_error) => match C::TransactionManager::rollback_transaction(conn).await {
Ok(()) | Err(diesel::result::Error::BrokenTransactionManager) => Err(user_error),
Err(rollback_error) => Err(rollback_error.into()),
},
}
}
#[doc(hidden)]
pub async fn scoped_immediate_transaction<'a, T, E, F>(
conn: &'a mut RuntimeConnection,
f: F,
) -> Result<T, E>
where
T: Send + 'a,
E: From<diesel::result::Error> + Send + 'a,
F: for<'r> FnOnce(
&'r mut RuntimeConnection,
) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
+ Send
+ 'a,
{
crate::backend_select! {
pg => { scoped_transaction(conn, f).await },
sqlite => {{
use diesel::connection::{AnsiTransactionManager, TransactionManager};
conn.spawn_blocking(|inner| {
AnsiTransactionManager::begin_transaction_sql(inner, "BEGIN IMMEDIATE")
})
.await
.map_err(E::from)?;
let outcome = AssertUnwindSafe(f(&mut *conn)).catch_unwind().await;
match outcome {
Ok(Ok(value)) => {
conn.spawn_blocking(|inner| {
<AnsiTransactionManager as TransactionManager<
diesel::SqliteConnection,
>>::commit_transaction(inner)
})
.await
.map_err(E::from)?;
Ok(value)
}
Ok(Err(user_error)) => {
if let Err(e) = conn
.spawn_blocking(|inner| {
<AnsiTransactionManager as TransactionManager<
diesel::SqliteConnection,
>>::rollback_transaction(inner)
})
.await
{
tracing::warn!(
"failed to roll back immediate transaction after error: {e}"
);
}
Err(user_error)
}
Err(panic) => {
if let Err(e) = conn
.spawn_blocking(|inner| {
<AnsiTransactionManager as TransactionManager<
diesel::SqliteConnection,
>>::rollback_transaction(inner)
})
.await
{
tracing::error!(
"failed to roll back immediate transaction during panic: {e}"
);
}
std::panic::resume_unwind(panic);
}
}
}},
}
}
pub async fn savepoint<'a, C, T, E, F>(conn: &'a mut C, f: F) -> Result<T, E>
where
C: diesel_async::AsyncConnection + Send,
T: Send + 'a,
E: From<diesel::result::Error> + Send + 'a,
F: for<'r> FnOnce(&'r mut C) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
+ Send
+ 'a,
{
scoped_transaction(conn, f).await
}
pub type PooledConnection = diesel_async::pooled_connection::deadpool::Object<RuntimeConnection>;
struct TxDepthGuard<'a> {
depth: &'a mut usize,
poisoned: &'a mut bool,
disarmed: bool,
}
impl Drop for TxDepthGuard<'_> {
fn drop(&mut self) {
*self.depth -= 1;
if !self.disarmed {
*self.poisoned = true;
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StatementTimeout(pub std::time::Duration);
pub struct Db {
conn: PooledConnection,
span: tracing::Span,
tx_depth: usize,
tx_poisoned: bool,
route_key: Option<String>,
metrics: Option<crate::middleware::MetricsCollector>,
slow_query_threshold: std::time::Duration,
start_time: std::time::Instant,
is_test_tx: bool,
}
impl Db {
#[must_use]
pub const fn span(&self) -> &tracing::Span {
&self.span
}
pub async fn tx<'a, T, E, F>(&'a mut self, f: F) -> Result<T, crate::error::AutumnError>
where
T: Send + 'a,
E: From<diesel::result::Error> + Send + Sync + 'a,
crate::error::AutumnError: From<E>,
F: for<'r> FnOnce(
&'r mut PooledConnection,
) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
+ Send
+ 'a,
{
if self.tx_poisoned {
return Err(crate::error::AutumnError::service_unavailable_msg(
"Database connection is in an invalid transaction state",
));
}
if self.tx_depth > 0 {
return Err(crate::error::AutumnError::bad_request_msg(
NESTED_TX_MESSAGE,
));
}
reject_ambient_after_commit_registry_for_tx()?;
self.tx_depth += 1;
let mut guard = TxDepthGuard {
depth: &mut self.tx_depth,
poisoned: &mut self.tx_poisoned,
disarmed: false,
};
let registry: Arc<Mutex<Vec<CommitCallback>>> = Arc::new(Mutex::new(Vec::new()));
let result = AFTER_COMMIT_REGISTRY
.scope(
registry.clone(),
scoped_transaction::<T, E, _, _>(&mut self.conn, f),
)
.await
.map_err(Into::into);
guard.disarmed = true;
if result.is_ok() {
let callbacks: Vec<CommitCallback> = {
let mut reg = registry.lock().expect("registry lock");
std::mem::take(&mut *reg)
};
if !callbacks.is_empty() && !self.is_test_tx {
let _ = spawn_committed_after_commit_callbacks(callbacks);
}
}
result
}
#[allow(clippy::too_many_lines)]
#[cfg_attr(feature = "sqlite", allow(unused_mut))]
pub async fn tx_with<'a, T, E, F>(
&'a mut self,
opts: TxOptions,
mut f: F,
) -> Result<T, crate::error::AutumnError>
where
T: Send + 'a,
E: From<diesel::result::Error> + Send + Sync + 'a,
crate::error::AutumnError: From<E>,
F: for<'r> FnMut(
&'r mut RuntimeConnection,
) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
+ Send
+ 'a,
{
if self.tx_poisoned {
return Err(crate::error::AutumnError::service_unavailable_msg(
"Database connection is in an invalid transaction state",
));
}
if self.tx_depth > 0 {
return Err(crate::error::AutumnError::bad_request_msg(
NESTED_TX_MESSAGE,
));
}
reject_ambient_after_commit_registry_for_tx()?;
#[cfg(feature = "sqlite")]
if opts.read_only {
return Err(crate::error::AutumnError::bad_request_msg(
"SQLite runtime does not support read-only transactions \
(TxOptions::read_only); this build cannot enforce read-only \
semantics on SQLite. Remove the read_only option or run on \
Postgres.",
));
}
self.tx_depth += 1;
let mut guard = TxDepthGuard {
depth: &mut self.tx_depth,
poisoned: &mut self.tx_poisoned,
disarmed: false,
};
let span = tracing::info_span!(
"db.transaction",
db.system = "postgresql",
db.isolation = opts.isolation.as_str(),
db.tx.attempts = tracing::field::Empty,
);
if self.is_test_tx {
let registry: Arc<Mutex<Vec<CommitCallback>>> = Arc::new(Mutex::new(Vec::new()));
let conn: &mut RuntimeConnection = &mut self.conn;
let result = AFTER_COMMIT_REGISTRY
.scope(registry, scoped_transaction::<T, E, _, _>(conn, f))
.instrument(span.clone())
.await
.map_err(Into::into);
span.record("db.tx.attempts", 1u32);
guard.disarmed = true;
return result;
}
#[cfg(feature = "sqlite")]
{
let registry: Arc<Mutex<Vec<CommitCallback>>> = Arc::new(Mutex::new(Vec::new()));
let conn: &mut RuntimeConnection = &mut self.conn;
let result: Result<T, E> = AFTER_COMMIT_REGISTRY
.scope(registry.clone(), scoped_transaction::<T, E, _, _>(conn, f))
.instrument(span.clone())
.await;
span.record("db.tx.attempts", 1u32);
guard.disarmed = true;
match result {
Ok(value) => {
let callbacks: Vec<CommitCallback> = {
let mut reg = registry.lock().expect("registry lock");
std::mem::take(&mut *reg)
};
if !callbacks.is_empty() {
let _ = spawn_committed_after_commit_callbacks(callbacks);
}
Ok(value)
}
Err(user_error) => Err(crate::error::AutumnError::from(user_error)),
}
}
#[cfg(not(feature = "sqlite"))]
{
let max_attempts = opts.effective_max_attempts();
let mut attempt: u32 = 0;
let outcome: Result<T, crate::error::AutumnError> = loop {
attempt += 1;
let registry: Arc<Mutex<Vec<CommitCallback>>> = Arc::new(Mutex::new(Vec::new()));
let attempt_result: Result<T, E> = {
let f_ref = &mut f;
let mut builder = self.conn.build_transaction();
builder = match opts.isolation {
IsolationLevel::ReadCommitted => builder.read_committed(),
IsolationLevel::RepeatableRead => builder.repeatable_read(),
IsolationLevel::Serializable => builder.serializable(),
};
if opts.read_only {
builder = builder.read_only();
}
if opts.deferrable {
builder = builder.deferrable();
}
AFTER_COMMIT_REGISTRY
.scope(
registry.clone(),
builder.run::<T, E, _>(async move |conn| f_ref(conn).await),
)
.instrument(span.clone())
.await
};
match attempt_result {
Ok(value) => {
let callbacks: Vec<CommitCallback> = {
let mut reg = registry.lock().expect("registry lock");
std::mem::take(&mut *reg)
};
if !callbacks.is_empty() {
let _ = spawn_committed_after_commit_callbacks(callbacks);
}
break Ok(value);
}
Err(e) => {
let ae = crate::error::AutumnError::from(e);
let retryable = is_retryable_txn_error(&ae);
match retry_decision(attempt, max_attempts, retryable) {
RetryDecision::Retry => {
let retries_total = record_tx_retry();
let delay = retry_backoff_delay(
opts.initial_backoff,
opts.max_backoff,
attempt,
);
tracing::debug!(
parent: &span,
attempt,
max_attempts,
delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
autumn.tx.retries_total = retries_total,
"retrying transaction after serialization/deadlock failure"
);
tokio::time::sleep(delay).await;
}
RetryDecision::Stop => {
if retryable {
let exhausted_total = record_tx_retry_exhausted();
tracing::warn!(
parent: &span,
attempt,
max_attempts,
autumn.tx.retry_exhausted_total = exhausted_total,
"transaction retry budget exhausted; returning final error"
);
}
break Err(ae);
}
}
}
}
};
span.record("db.tx.attempts", attempt);
guard.disarmed = true;
outcome
}
}
}
impl std::ops::Deref for Db {
type Target = RuntimeConnection;
fn deref(&self) -> &Self::Target {
assert!(
!self.tx_poisoned,
"Db connection is poisoned due to a cancelled/dropped transaction"
);
&self.conn
}
}
impl std::ops::DerefMut for Db {
fn deref_mut(&mut self) -> &mut Self::Target {
assert!(
!self.tx_poisoned,
"Db connection is poisoned due to a cancelled/dropped transaction"
);
&mut self.conn
}
}
pub(crate) struct DbCheckoutParams<'a> {
pub pool: &'a Pool<RuntimeConnection>,
pub pool_name: &'a str,
pub shard: Option<&'a str>,
pub statement_timeout: Option<std::time::Duration>,
pub route_key: Option<String>,
pub metrics: Option<crate::middleware::MetricsCollector>,
pub slow_query_threshold: std::time::Duration,
pub interceptors: Vec<std::sync::Arc<dyn crate::interceptor::DbConnectionInterceptor>>,
}
impl Db {
pub(crate) async fn checkout(params: DbCheckoutParams<'_>) -> Result<Self, AutumnError> {
#[cfg(not(feature = "sqlite"))]
const PG_TIMEOUT_MAX_MS: u64 = i32::MAX as u64;
#[cfg(not(feature = "sqlite"))]
use diesel_async::RunQueryDsl as _;
let span = tracing::info_span!(
"db.connection",
otel.kind = "client",
db.system = "postgresql",
db.shard = tracing::field::Empty,
);
if let Some(shard) = params.shard {
span.record("db.shard", shard);
}
let pool = params.pool;
let mut checkout_future: std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<PooledConnection, AutumnError>> + Send + '_,
>,
> = Box::pin(async move {
pool.get().await.map_err(|e| {
tracing::error!("Failed to acquire database connection: {e}");
AutumnError::service_unavailable_msg(e.to_string())
})
});
for interceptor in ¶ms.interceptors {
let ctx = crate::interceptor::DbCheckoutContext {
pool_name: params.pool_name.to_string(),
};
checkout_future = interceptor.intercept_checkout(ctx, checkout_future);
}
let mut conn = checkout_future.instrument(span.clone()).await?;
#[cfg(feature = "sqlite")]
let _ = params.statement_timeout;
#[cfg(not(feature = "sqlite"))]
let timeout_ms = params.statement_timeout.map_or(0u64, |d| {
u64::try_from(d.as_millis())
.unwrap_or(PG_TIMEOUT_MAX_MS)
.min(PG_TIMEOUT_MAX_MS)
});
#[cfg(feature = "db")]
{
use diesel_async::AsyncConnection as _;
if request_db_timing_active() || request_query_capture_active() {
conn.set_instrumentation(RequestQueryTimer::default());
}
}
#[cfg(not(feature = "sqlite"))]
diesel::sql_query(format!("SET statement_timeout = {timeout_ms}"))
.execute(&mut conn)
.await
.map_err(|e| {
tracing::error!("Failed to set database statement_timeout to {timeout_ms}ms: {e}");
AutumnError::service_unavailable_msg(format!("Database initialization error: {e}"))
})?;
let start_time = std::time::Instant::now();
let is_test_tx = params
.interceptors
.iter()
.any(|i| i.is_transactional_test());
Ok(Self {
conn,
span,
tx_depth: 0,
tx_poisoned: false,
route_key: params.route_key,
metrics: params.metrics,
slow_query_threshold: params.slow_query_threshold,
start_time,
is_test_tx,
})
}
#[cfg(feature = "test-support")]
pub async fn connect_for_test(pool: &Pool<RuntimeConnection>) -> Result<Self, AutumnError> {
Self::checkout(DbCheckoutParams {
pool,
pool_name: "test",
shard: None,
statement_timeout: None,
route_key: None,
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(500),
interceptors: Vec::new(),
})
.await
}
}
#[derive(Clone)]
pub(crate) struct RequestDbContext {
pub statement_timeout: Option<std::time::Duration>,
pub route_key: Option<String>,
pub metrics: Option<crate::middleware::MetricsCollector>,
pub slow_query_threshold: std::time::Duration,
pub interceptors: Vec<std::sync::Arc<dyn crate::interceptor::DbConnectionInterceptor>>,
}
impl RequestDbContext {
pub(crate) fn from_parts<S: DbState>(parts: &axum::http::request::Parts, state: &S) -> Self {
let timeout_override = parts.extensions.get::<StatementTimeout>().copied();
let matched_path = parts
.extensions
.get::<axum::extract::MatchedPath>()
.map_or_else(|| parts.uri.path(), axum::extract::MatchedPath::as_str);
Self {
statement_timeout: timeout_override
.map(|t| t.0)
.or_else(|| state.statement_timeout()),
route_key: Some(format!("{} {}", parts.method, matched_path)),
metrics: state.metrics().cloned(),
slow_query_threshold: state.slow_query_threshold(),
interceptors: state.db_interceptors(),
}
}
}
impl<S> FromRequestParts<S> for Db
where
S: DbState + Send + Sync,
{
type Rejection = AutumnError;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let pool = state
.pool()
.ok_or_else(|| AutumnError::service_unavailable_msg("Database not configured"))?;
let ctx = RequestDbContext::from_parts(parts, state);
let result = Self::checkout(DbCheckoutParams {
pool,
pool_name: "primary",
shard: None,
statement_timeout: ctx.statement_timeout,
route_key: ctx.route_key,
metrics: ctx.metrics,
slow_query_threshold: ctx.slow_query_threshold,
interceptors: ctx.interceptors,
})
.await;
if result.is_ok() {
crate::read_your_writes::mark_write();
}
result
}
}
impl Drop for Db {
fn drop(&mut self) {
if let (Some(route_key), Some(metrics)) = (&self.route_key, &self.metrics) {
let elapsed = self.start_time.elapsed();
let elapsed_ms = u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX);
let metric_key = format!("{route_key} SELECT");
metrics.record_db_query(&metric_key, elapsed_ms);
if elapsed >= self.slow_query_threshold {
tracing::warn!(
route = %route_key,
sql = "SELECT ?",
duration_ms = elapsed_ms,
"slow database query"
);
}
}
}
}
pub trait DatabasePoolProvider: Send + Sync + 'static {
fn create_pool(
&self,
config: &DatabaseConfig,
) -> impl std::future::Future<Output = Result<Option<Pool<RuntimeConnection>>, PoolError>> + Send;
fn create_topology(
&self,
config: &DatabaseConfig,
) -> impl std::future::Future<Output = Result<Option<DatabaseTopology>, PoolError>> + Send {
async move {
let Some(primary) = self.create_pool(config).await? else {
return Ok(None);
};
#[cfg(feature = "sqlite")]
if let (Some(primary_url), Some(replica_url)) = (
config.effective_primary_url(),
config.replica_url.as_deref(),
) {
reject_unusable_sqlite_replica(primary_url, replica_url)?;
}
let replica = config
.replica_url
.as_deref()
.map(|url| {
build_pool(
url,
config.effective_replica_pool_size(),
config.connect_timeout_secs,
)
})
.transpose()?;
Ok(Some(DatabaseTopology::from_pools(primary, replica)))
}
}
fn create_shard_topology(
&self,
shard: &crate::config::ShardConfig,
defaults: &DatabaseConfig,
) -> impl std::future::Future<Output = Result<DatabaseTopology, PoolError>> + Send {
async move { create_shard_topology(shard, defaults) }
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct DieselDeadpoolPoolProvider;
impl DieselDeadpoolPoolProvider {
#[must_use]
pub const fn new() -> Self {
Self
}
}
impl DatabasePoolProvider for DieselDeadpoolPoolProvider {
async fn create_pool(
&self,
config: &DatabaseConfig,
) -> Result<Option<Pool<RuntimeConnection>>, PoolError> {
create_pool(config)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::DatabaseConfig;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;
#[cfg(feature = "db")]
#[tokio::test]
async fn request_query_timer_accumulates_finished_queries() {
let timings = Arc::new(RequestDbTimings::default());
REQUEST_DB_TIMINGS
.scope(Arc::clone(&timings), async {
let mut timer = RequestQueryTimer::default();
let t0 = std::time::Instant::now();
timer.on_start(t0, || "SELECT 1".to_string());
timer.on_finish(t0 + Duration::from_micros(1_200));
let t1 = std::time::Instant::now();
timer.on_start(t1, || "UPDATE users SET name = $1".to_string());
timer.on_finish(t1 + Duration::from_micros(800));
timer.on_finish(std::time::Instant::now());
})
.await;
assert_eq!(
timings.query_count.load(Ordering::Relaxed),
2,
"only completed StartQuery/FinishQuery pairs are counted"
);
assert_eq!(
timings.total_us.load(Ordering::Relaxed),
2_000,
"elapsed times accumulate (1200µs + 800µs)"
);
}
#[cfg(feature = "db")]
#[tokio::test]
async fn request_query_timer_is_noop_off_request() {
let mut timer = RequestQueryTimer::default();
let t0 = std::time::Instant::now();
timer.on_start(t0, || "SELECT 1".to_string());
timer.on_finish(t0 + Duration::from_micros(500));
}
#[cfg(feature = "db")]
#[tokio::test]
async fn request_query_timer_on_start_is_cheap_noop_off_request() {
let mut timer = RequestQueryTimer::default();
let invoked = std::cell::Cell::new(false);
let t0 = std::time::Instant::now();
timer.on_start(t0, || {
invoked.set(true);
"SELECT 1".to_string()
});
assert!(
!invoked.get(),
"off-request, on_start must not format the SQL (no allocation)"
);
timer.on_finish(t0 + Duration::from_micros(500));
}
#[cfg(feature = "db")]
#[tokio::test]
async fn request_query_timer_excludes_checkout_statement_timeout_set() {
let timings = Arc::new(RequestDbTimings::default());
REQUEST_DB_TIMINGS
.scope(Arc::clone(&timings), async {
let mut timer = RequestQueryTimer::default();
let t0 = std::time::Instant::now();
timer.on_start(t0, || "SET statement_timeout = 5000".to_string());
timer.on_finish(t0 + Duration::from_millis(3));
assert_eq!(
timings.query_count.load(Ordering::Relaxed),
0,
"the checkout SET statement_timeout must not be counted"
);
assert_eq!(
timings.total_us.load(Ordering::Relaxed),
0,
"the checkout SET contributes no latency to db;dur"
);
let t1 = std::time::Instant::now();
timer.on_start(t1, || "SELECT * FROM users".to_string());
timer.on_finish(t1 + Duration::from_micros(600));
assert_eq!(
timings.query_count.load(Ordering::Relaxed),
1,
"the real SELECT is the first counted query"
);
assert_eq!(
timings.total_us.load(Ordering::Relaxed),
600,
"only the SELECT's 600µs accumulates; the SET is excluded"
);
})
.await;
}
#[cfg(feature = "db")]
#[tokio::test]
async fn request_db_timing_active_reflects_scope() {
assert!(
!request_db_timing_active(),
"no REQUEST_DB_TIMINGS scope active outside the ServerTimingLayer: \
the timer records nothing"
);
let timings = Arc::new(RequestDbTimings::default());
REQUEST_DB_TIMINGS
.scope(Arc::clone(&timings), async {
assert!(
request_db_timing_active(),
"inside a REQUEST_DB_TIMINGS scope the probe is active: \
the timer records queries"
);
})
.await;
assert!(
!request_db_timing_active(),
"the gate is false again once the scope is dropped"
);
}
#[cfg(feature = "db")]
#[tokio::test]
async fn request_query_timer_excludes_transaction_control_statements() {
let timings = Arc::new(RequestDbTimings::default());
REQUEST_DB_TIMINGS
.scope(Arc::clone(&timings), async {
let mut timer = RequestQueryTimer::default();
let t0 = std::time::Instant::now();
timer.on_start(t0, || "BEGIN".to_string());
timer.on_finish(t0 + Duration::from_millis(10));
let t1 = std::time::Instant::now();
timer.on_start(t1, || "SELECT * FROM users".to_string());
timer.on_finish(t1 + Duration::from_micros(700));
let t2 = std::time::Instant::now();
timer.on_start(t2, || "COMMIT".to_string());
timer.on_finish(t2 + Duration::from_millis(10));
})
.await;
assert_eq!(
timings.query_count.load(Ordering::Relaxed),
1,
"only the real SELECT is counted; BEGIN/COMMIT are excluded"
);
assert_eq!(
timings.total_us.load(Ordering::Relaxed),
700,
"only the SELECT's 700µs accumulates; begin/commit latency is excluded"
);
}
#[cfg(feature = "db")]
#[tokio::test]
async fn checkout_installs_timer_only_when_scope_active() {
fn run_checkout_gate(slot: &mut &'static str) {
if request_db_timing_active() {
*slot = "timer";
}
}
assert!(!request_db_timing_active());
let mut slot = "app";
run_checkout_gate(&mut slot);
assert_eq!(
slot, "app",
"off-scope checkout must NOT clobber the app's instrumentation"
);
let timings = Arc::new(RequestDbTimings::default());
REQUEST_DB_TIMINGS
.scope(Arc::clone(&timings), async {
let mut slot = "app";
run_checkout_gate(&mut slot);
assert_eq!(
slot, "timer",
"inside a REQUEST_DB_TIMINGS scope autumn installs its timer"
);
})
.await;
}
#[cfg(feature = "db")]
#[test]
fn is_uncounted_statement_classifies_statements() {
for sql in [
"BEGIN",
"begin",
" COMMIT",
"ROLLBACK",
"ROLLBACK TO SAVEPOINT diesel_savepoint_0",
"SAVEPOINT diesel_savepoint_1",
"RELEASE SAVEPOINT diesel_savepoint_0",
"start transaction",
"SET TRANSACTION ISOLATION LEVEL SERIALIZABLE",
"SET statement_timeout = 5000",
"set statement_timeout = 0",
] {
assert!(
RequestQueryTimer::is_uncounted_statement(sql),
"{sql:?} should be treated as an uncounted housekeeping/tx statement"
);
}
for sql in [
"SELECT 1",
"UPDATE users SET name = $1",
"INSERT INTO t VALUES (1)",
"DELETE FROM t WHERE id = 1",
"",
] {
assert!(
!RequestQueryTimer::is_uncounted_statement(sql),
"{sql:?} should be treated as a countable query"
);
}
}
#[cfg(feature = "db")]
#[test]
fn strip_bind_annotation_lets_detect_n_plus_one_see_per_row_repetition() {
use crate::inspector::{QueryRecord, detect_n_plus_one};
fn record(diesel_display: &str) -> QueryRecord {
QueryRecord {
sql: strip_bind_annotation(diesel_display).to_owned(),
params: Vec::new(),
elapsed_ms: 1,
location: String::new(),
}
}
assert_eq!(
strip_bind_annotation("SELECT * FROM books WHERE author_id = $1 -- binds: [1]"),
"SELECT * FROM books WHERE author_id = $1"
);
assert_eq!(
strip_bind_annotation("SELECT * FROM authors -- binds: []"),
"SELECT * FROM authors"
);
assert_eq!(strip_bind_annotation("SELECT 1"), "SELECT 1");
assert_eq!(strip_bind_annotation("BEGIN"), "BEGIN");
let n_plus_one = vec![
record("SELECT * FROM authors -- binds: []"),
record("SELECT * FROM books WHERE author_id = $1 -- binds: [1]"),
record("SELECT * FROM books WHERE author_id = $1 -- binds: [2]"),
record("SELECT * FROM books WHERE author_id = $1 -- binds: [3]"),
];
let warning = detect_n_plus_one(&n_plus_one, 3)
.expect("three identical per-row templates should trip the N+1 detector");
assert_eq!(
warning.count, 3,
"all three per-row executions collapse to a single template"
);
assert_eq!(
warning.sql_template, "select * from books where author_id = $1",
"the reported template is the parameterised statement, bind-free"
);
let distinct = vec![
record("SELECT * FROM authors -- binds: []"),
record("SELECT * FROM books WHERE id = $1 -- binds: [1]"),
record("UPDATE users SET name = $1 WHERE id = $2 -- binds: [\"x\", 2]"),
];
assert!(
detect_n_plus_one(&distinct, 3).is_none(),
"three distinct query templates must not be reported as N+1"
);
}
#[tokio::test]
async fn register_after_commit_outside_tx_runs_eagerly() {
let counter = Arc::new(AtomicUsize::new(0));
let c = counter.clone();
register_after_commit(move || async move {
c.fetch_add(1, Ordering::SeqCst);
Ok(())
})
.await;
assert_eq!(counter.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn register_after_commit_eager_failure_increments_failure_counter() {
let before = AFTER_COMMIT_FAILURES_TOTAL.load(Ordering::Relaxed);
register_after_commit(|| async {
Err(crate::AutumnError::internal_server_error_msg(
"deliberate eager after-commit failure",
))
})
.await;
let after = AFTER_COMMIT_FAILURES_TOTAL.load(Ordering::Relaxed);
assert!(
after > before,
"eager after_commit failures should be counted for recovery signals"
);
}
#[tokio::test]
async fn register_after_commit_inside_scope_defers_until_drained() {
let counter = Arc::new(AtomicUsize::new(0));
let c = counter.clone();
let registry = Arc::new(std::sync::Mutex::new(Vec::<CommitCallback>::new()));
AFTER_COMMIT_REGISTRY
.scope(registry.clone(), async {
register_after_commit(move || async move {
c.fetch_add(1, Ordering::SeqCst);
Ok(())
})
.await;
})
.await;
assert_eq!(counter.load(Ordering::SeqCst), 0);
let callbacks: Vec<CommitCallback> = {
let mut reg = registry.lock().unwrap();
std::mem::take(&mut *reg)
};
for cb in callbacks {
cb().await.unwrap();
}
assert_eq!(counter.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn register_after_commit_on_rollback_callbacks_dropped() {
let counter = Arc::new(AtomicUsize::new(0));
let c = counter.clone();
let registry = Arc::new(std::sync::Mutex::new(Vec::<CommitCallback>::new()));
AFTER_COMMIT_REGISTRY
.scope(registry.clone(), async {
register_after_commit(move || async move {
c.fetch_add(1, Ordering::SeqCst);
Ok(())
})
.await;
})
.await;
drop(registry);
assert_eq!(counter.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn register_after_commit_callbacks_run_in_registration_order() {
let order = Arc::new(std::sync::Mutex::new(Vec::<u32>::new()));
let registry = Arc::new(std::sync::Mutex::new(Vec::<CommitCallback>::new()));
let o1 = order.clone();
let o2 = order.clone();
let o3 = order.clone();
AFTER_COMMIT_REGISTRY
.scope(registry.clone(), async {
register_after_commit(move || async move {
o1.lock().unwrap().push(1);
Ok(())
})
.await;
register_after_commit(move || async move {
o2.lock().unwrap().push(2);
Ok(())
})
.await;
register_after_commit(move || async move {
o3.lock().unwrap().push(3);
Ok(())
})
.await;
})
.await;
let callbacks: Vec<CommitCallback> = {
let mut reg = registry.lock().unwrap();
std::mem::take(&mut *reg)
};
for cb in callbacks {
cb().await.unwrap();
}
assert_eq!(*order.lock().unwrap(), vec![1, 2, 3]);
}
#[tokio::test]
async fn production_after_commit_drain_preserves_registration_order() {
let order = Arc::new(std::sync::Mutex::new(Vec::<u32>::new()));
let (release_first, wait_first) = tokio::sync::oneshot::channel::<()>();
let first_order = order.clone();
let second_order = order.clone();
let callbacks: Vec<CommitCallback> = vec![
Box::new(move || {
Box::pin(async move {
wait_first
.await
.expect("test should release first callback");
first_order.lock().unwrap().push(1);
Ok(())
})
}),
Box::new(move || {
Box::pin(async move {
second_order.lock().unwrap().push(2);
Ok(())
})
}),
];
let drain = spawn_committed_after_commit_callbacks(callbacks)
.expect("non-empty callback list should spawn a drain task");
tokio::task::yield_now().await;
assert_eq!(
*order.lock().unwrap(),
Vec::<u32>::new(),
"later callbacks must wait for earlier callbacks to finish"
);
release_first
.send(())
.expect("first callback receiver alive");
drain.await.expect("drain task should not panic");
assert_eq!(*order.lock().unwrap(), vec![1, 2]);
}
#[tokio::test]
async fn production_after_commit_drain_isolates_panicking_callbacks() {
let before = AFTER_COMMIT_FAILURES_TOTAL.load(Ordering::Relaxed);
let ran_later = Arc::new(AtomicU64::new(0));
let later = ran_later.clone();
let callbacks: Vec<CommitCallback> = vec![
Box::new(|| Box::pin(async { panic!("deliberate after_commit panic") })),
Box::new(move || {
Box::pin(async move {
later.fetch_add(1, Ordering::SeqCst);
Ok(())
})
}),
];
let drain = spawn_committed_after_commit_callbacks(callbacks)
.expect("non-empty callback list should spawn a drain task");
drain.await.expect("panicking callback should be isolated");
assert_eq!(
ran_later.load(Ordering::SeqCst),
1,
"later callbacks must still run after an earlier callback panics"
);
let after = AFTER_COMMIT_FAILURES_TOTAL.load(Ordering::Relaxed);
assert!(
after > before,
"panicking after_commit callbacks must increment the failure counter"
);
}
#[tokio::test]
async fn db_tx_rejects_ambient_after_commit_registry() {
let registry = Arc::new(std::sync::Mutex::new(Vec::<CommitCallback>::new()));
let err = AFTER_COMMIT_REGISTRY
.scope(registry, async {
reject_ambient_after_commit_registry_for_tx().expect_err(
"starting Db::tx inside an ambient transaction registry should fail",
)
})
.await;
assert!(
err.to_string().contains("Nested Db::tx calls"),
"unexpected nested transaction error: {err}"
);
}
#[tokio::test]
async fn register_after_commit_callback_error_is_swallowed() {
let registry = Arc::new(std::sync::Mutex::new(Vec::<CommitCallback>::new()));
AFTER_COMMIT_REGISTRY
.scope(registry.clone(), async {
register_after_commit(|| async {
Err(crate::AutumnError::internal_server_error_msg(
"deliberate error",
))
})
.await;
})
.await;
let callbacks: Vec<CommitCallback> = {
let mut reg = registry.lock().unwrap();
std::mem::take(&mut *reg)
};
for cb in callbacks {
let _ = cb().await;
}
}
struct NoOpPoolProvider;
impl DatabasePoolProvider for NoOpPoolProvider {
async fn create_pool(
&self,
_config: &DatabaseConfig,
) -> Result<Option<Pool<crate::db::RuntimeConnection>>, PoolError> {
Ok(None)
}
}
#[tokio::test]
async fn pool_provider_trait_returns_supplied_pool() {
let config = DatabaseConfig {
url: Some("postgres://localhost/ignored".to_owned()),
..Default::default()
};
let provider = NoOpPoolProvider;
let pool = provider
.create_pool(&config)
.await
.expect("no-op provider should succeed");
assert!(
pool.is_none(),
"no-op provider must override default behaviour"
);
}
#[tokio::test]
async fn default_pool_provider_matches_free_function() {
let config = DatabaseConfig::default();
let via_provider = DieselDeadpoolPoolProvider::new()
.create_pool(&config)
.await
.expect("default provider should succeed");
let via_function = create_pool(&config).expect("free fn should succeed");
assert_eq!(via_provider.is_none(), via_function.is_none());
}
#[tokio::test]
async fn default_pool_provider_respects_url_config() {
let config = DatabaseConfig {
url: Some("postgres://localhost/test".into()),
..Default::default()
};
let provider = DieselDeadpoolPoolProvider::new();
let pool = provider
.create_pool(&config)
.await
.expect("default provider should succeed");
assert!(
pool.is_some(),
"default provider should return Some when url is provided"
);
}
#[test]
fn create_pool_with_no_url_returns_none() {
let config = DatabaseConfig::default();
let pool = create_pool(&config).expect("should not fail with no URL");
assert!(pool.is_none());
}
#[test]
fn create_pool_with_url_returns_some() {
let config = DatabaseConfig {
url: Some("postgres://localhost/test".into()),
..Default::default()
};
let pool = create_pool(&config).expect("should build pool from valid config");
assert!(pool.is_some());
}
#[cfg(not(feature = "sqlite"))]
#[test]
fn create_pool_with_sqlite_url_fails_fast() {
let config = DatabaseConfig {
url: Some("sqlite:///var/lib/app.db".into()),
..Default::default()
};
let Err(err) = create_pool(&config) else {
panic!("sqlite target must refuse at pool build");
};
assert!(
matches!(err, PoolError::UnsupportedBackend(_)),
"expected UnsupportedBackend, got: {err:?}"
);
let msg = err.to_string();
assert!(
msg.contains("SQLite") && msg.contains("--features sqlite"),
"message must be actionable and name the sqlite build, got: {msg}"
);
}
#[cfg(not(feature = "sqlite"))]
#[test]
fn create_topology_with_sqlite_url_fails_fast() {
let config = DatabaseConfig {
primary_url: Some("sqlite::memory:".into()),
..Default::default()
};
let Err(err) = create_topology(&config) else {
panic!("sqlite target must refuse at topology build");
};
assert!(matches!(err, PoolError::UnsupportedBackend(_)), "{err:?}");
}
#[cfg(feature = "sqlite")]
#[test]
fn create_pool_with_sqlite_url_builds_under_feature() {
let config = DatabaseConfig {
primary_url: Some("sqlite::memory:".into()),
..Default::default()
};
let pool = create_pool(&config).expect("sqlite pool builds under the feature");
assert!(pool.is_some(), "a configured sqlite url yields a pool");
}
#[cfg(feature = "sqlite")]
#[test]
fn create_pool_with_postgres_url_refuses_under_sqlite_feature() {
let config = DatabaseConfig {
url: Some("postgres://localhost/test".into()),
..Default::default()
};
let Err(err) = create_pool(&config) else {
panic!("a Postgres url must refuse under the sqlite feature");
};
assert!(matches!(err, PoolError::UnsupportedBackend(_)), "{err:?}");
}
#[cfg(feature = "sqlite")]
#[test]
fn create_topology_rejects_in_memory_sqlite_replica() {
let config = DatabaseConfig {
primary_url: Some("sqlite:///var/lib/app.db".into()),
replica_url: Some("sqlite::memory:".into()),
..Default::default()
};
let Err(err) = create_topology(&config) else {
panic!("an in-memory sqlite replica must be rejected");
};
assert!(matches!(err, PoolError::UnsupportedBackend(_)), "{err:?}");
let msg = err.to_string();
assert!(
msg.contains("separate read replica") && msg.contains("primary"),
"message must be actionable: {msg}"
);
}
#[cfg(feature = "sqlite")]
#[test]
fn create_topology_rejects_distinct_file_sqlite_replica() {
let config = DatabaseConfig {
primary_url: Some("sqlite:///var/lib/primary.db".into()),
replica_url: Some("sqlite:///var/lib/replica.db".into()),
..Default::default()
};
let Err(err) = create_topology(&config) else {
panic!("a distinct-file sqlite replica must be rejected");
};
assert!(matches!(err, PoolError::UnsupportedBackend(_)), "{err:?}");
}
#[cfg(feature = "sqlite")]
#[test]
fn create_topology_allows_same_file_sqlite_replica() {
let config = DatabaseConfig {
primary_url: Some("sqlite:///var/lib/app.db".into()),
replica_url: Some("sqlite:///var/lib/app.db".into()),
..Default::default()
};
let topology = create_topology(&config)
.expect("same-file sqlite replica must be allowed")
.expect("a configured primary yields a topology");
assert!(
topology.replica().is_some(),
"same-file replica pool should be built"
);
}
#[cfg(feature = "sqlite")]
#[test]
fn create_topology_sqlite_primary_only_builds() {
let config = DatabaseConfig {
primary_url: Some("sqlite::memory:".into()),
..Default::default()
};
let topology = create_topology(&config)
.expect("sqlite primary-only topology builds")
.expect("a configured primary yields a topology");
assert!(topology.replica().is_none(), "no replica configured");
}
#[cfg(feature = "sqlite")]
struct PrimaryOnlyProvider;
#[cfg(feature = "sqlite")]
impl DatabasePoolProvider for PrimaryOnlyProvider {
async fn create_pool(
&self,
config: &DatabaseConfig,
) -> Result<Option<Pool<RuntimeConnection>>, PoolError> {
create_pool(config)
}
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn default_provider_topology_rejects_distinct_file_sqlite_replica() {
let config = DatabaseConfig {
primary_url: Some("sqlite:///var/lib/primary.db".into()),
replica_url: Some("sqlite:///var/lib/replica.db".into()),
..Default::default()
};
let Err(err) = PrimaryOnlyProvider.create_topology(&config).await else {
panic!("the default-topology path must reject a distinct-file sqlite replica");
};
assert!(matches!(err, PoolError::UnsupportedBackend(_)), "{err:?}");
let msg = err.to_string();
assert!(
msg.contains("separate read replica") && msg.contains("primary"),
"message must be actionable: {msg}"
);
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn default_provider_topology_rejects_in_memory_sqlite_replica() {
let config = DatabaseConfig {
primary_url: Some("sqlite:///var/lib/app.db".into()),
replica_url: Some("sqlite::memory:".into()),
..Default::default()
};
let Err(err) = PrimaryOnlyProvider.create_topology(&config).await else {
panic!("an in-memory sqlite replica must be rejected by the default topology");
};
assert!(matches!(err, PoolError::UnsupportedBackend(_)), "{err:?}");
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn default_provider_topology_allows_same_file_and_primary_only() {
let same_file = DatabaseConfig {
primary_url: Some("sqlite:///var/lib/app.db".into()),
replica_url: Some("sqlite:///var/lib/app.db".into()),
..Default::default()
};
let topology = PrimaryOnlyProvider
.create_topology(&same_file)
.await
.expect("same-file replica must be allowed by the default topology")
.expect("a configured primary yields a topology");
assert!(
topology.replica().is_some(),
"same-file replica pool should be built"
);
let primary_only = DatabaseConfig {
primary_url: Some("sqlite::memory:".into()),
..Default::default()
};
let topology = PrimaryOnlyProvider
.create_topology(&primary_only)
.await
.expect("primary-only default topology builds")
.expect("a configured primary yields a topology");
assert!(topology.replica().is_none(), "no replica configured");
}
#[cfg(feature = "sqlite")]
#[test]
fn create_shard_topology_rejects_in_memory_sqlite_replica() {
let defaults = DatabaseConfig::default();
let shard = crate::config::ShardConfig {
name: "shard0".into(),
primary_url: "sqlite:///var/lib/shard0.db".into(),
replica_url: Some("sqlite::memory:".into()),
..Default::default()
};
let Err(err) = create_shard_topology(&shard, &defaults) else {
panic!("an in-memory sqlite shard replica must be rejected");
};
assert!(matches!(err, PoolError::UnsupportedBackend(_)), "{err:?}");
let msg = err.to_string();
assert!(
msg.contains("separate read replica") && msg.contains("primary"),
"message must be actionable: {msg}"
);
}
#[cfg(feature = "sqlite")]
#[test]
fn create_shard_topology_rejects_distinct_file_sqlite_replica() {
let defaults = DatabaseConfig::default();
let shard = crate::config::ShardConfig {
name: "shard0".into(),
primary_url: "sqlite:///var/lib/shard0-primary.db".into(),
replica_url: Some("sqlite:///var/lib/shard0-replica.db".into()),
..Default::default()
};
let Err(err) = create_shard_topology(&shard, &defaults) else {
panic!("a distinct-file sqlite shard replica must be rejected");
};
assert!(matches!(err, PoolError::UnsupportedBackend(_)), "{err:?}");
}
#[cfg(feature = "sqlite")]
#[test]
fn create_shard_topology_allows_same_file_sqlite_replica() {
let defaults = DatabaseConfig::default();
let shard = crate::config::ShardConfig {
name: "shard0".into(),
primary_url: "sqlite:///var/lib/shard0.db".into(),
replica_url: Some("sqlite:///var/lib/shard0.db".into()),
..Default::default()
};
let topology = create_shard_topology(&shard, &defaults)
.expect("same-file sqlite shard replica must be allowed");
assert!(
topology.replica().is_some(),
"same-file shard replica pool should be built"
);
}
#[cfg(feature = "sqlite")]
#[test]
fn create_shard_topology_sqlite_primary_only_builds() {
let defaults = DatabaseConfig::default();
let shard = crate::config::ShardConfig {
name: "shard0".into(),
primary_url: "sqlite::memory:".into(),
..Default::default()
};
let topology = create_shard_topology(&shard, &defaults)
.expect("sqlite primary-only shard topology builds");
assert!(topology.replica().is_none(), "no shard replica configured");
}
#[cfg(feature = "sqlite")]
#[test]
fn sqlite_target_is_memory_covers_file_memory() {
assert!(sqlite_target_is_memory(":memory:"));
assert!(sqlite_target_is_memory("file::memory:"));
assert!(sqlite_target_is_memory("file::memory:?foo=bar"));
assert!(sqlite_target_is_memory("file:app?mode=memory"));
assert!(!sqlite_target_is_memory("file::memory:?cache=shared"));
assert!(!sqlite_target_is_memory(
"file:app?mode=memory&cache=shared"
));
assert!(!sqlite_target_is_memory("/var/lib/app.db"));
}
#[cfg(feature = "sqlite")]
#[test]
fn sqlite_target_is_any_in_memory_covers_shared_cache() {
assert!(sqlite_target_is_any_in_memory(":memory:"));
assert!(sqlite_target_is_any_in_memory("sqlite::memory:"));
assert!(sqlite_target_is_any_in_memory("sqlite://:memory:"));
assert!(sqlite_target_is_any_in_memory("sqlite://"));
assert!(sqlite_target_is_any_in_memory("file::memory:"));
assert!(sqlite_target_is_any_in_memory("file::memory:?foo=bar"));
assert!(sqlite_target_is_any_in_memory("file::memory:?cache=shared"));
assert!(sqlite_target_is_any_in_memory(
"file:app?mode=memory&cache=shared"
));
assert!(!sqlite_target_is_memory("file::memory:?cache=shared"));
assert!(!sqlite_target_is_memory(
"file:app?mode=memory&cache=shared"
));
assert!(!sqlite_target_is_any_in_memory("sqlite:///var/lib/app.db"));
assert!(!sqlite_target_is_any_in_memory("/var/lib/app.db"));
assert!(!sqlite_target_is_any_in_memory("file:/var/lib/app.db"));
}
#[cfg(feature = "sqlite")]
#[test]
fn sqlite_migration_connection_sets_busy_timeout() {
use diesel::RunQueryDsl as _;
#[derive(diesel::QueryableByName)]
struct BusyTimeout {
#[diesel(sql_type = diesel::sql_types::Integer)]
timeout: i32,
}
let tmp = tempfile::TempDir::new().expect("temp dir");
let db_path = tmp.path().join("migration.db");
let url = format!("sqlite://{}", db_path.display());
let mut conn = super::establish_sqlite_migration_connection(&url)
.expect("establish sqlite migration connection");
let rows: Vec<BusyTimeout> = diesel::sql_query("PRAGMA busy_timeout")
.load(&mut conn)
.expect("read busy_timeout pragma");
let timeout = rows
.into_iter()
.next()
.expect("busy_timeout pragma returns a row")
.timeout;
assert_eq!(
timeout, 5000,
"the migration connection must carry the configured busy_timeout (ms) \
so concurrent/locked migrators wait rather than hitting SQLITE_BUSY"
);
}
#[cfg(feature = "sqlite")]
#[test]
fn sqlite_target_is_read_only_detects_ro_and_immutable() {
assert!(sqlite_target_is_read_only("file:/srv/ref.db?mode=ro"));
assert!(sqlite_target_is_read_only("file:/srv/ref.db?MODE=RO"));
assert!(sqlite_target_is_read_only("file:/srv/ref.db?immutable=1"));
assert!(sqlite_target_is_read_only(
"file:/srv/ref.db?immutable=true"
));
assert!(sqlite_target_is_read_only(
"file:/srv/ref.db?immutable=TRUE"
));
assert!(sqlite_target_is_read_only(
"file:/srv/ref.db?cache=shared&mode=ro"
));
assert!(!sqlite_target_is_read_only("/var/lib/app.db"));
assert!(!sqlite_target_is_read_only(":memory:"));
assert!(!sqlite_target_is_read_only("file::memory:?cache=shared"));
assert!(!sqlite_target_is_read_only("file:app?mode=memory"));
assert!(!sqlite_target_is_read_only("file:/srv/ref.db?mode=rwc"));
assert!(!sqlite_target_is_read_only("file:/srv/ref.db?immutable=0"));
assert!(!sqlite_target_is_read_only("file:/srv/ro-data.db"));
}
#[cfg(feature = "sqlite")]
#[test]
fn file_memory_sqlite_pool_is_single_slot() {
let config = DatabaseConfig {
primary_url: Some("file::memory:".into()),
pool_size: 5,
..Default::default()
};
let pool = create_pool(&config)
.expect("file::memory: sqlite pool builds")
.expect("a configured url yields a pool");
assert_eq!(
pool.status().max_size,
1,
"a private in-memory target must be forced single-slot"
);
}
#[cfg(feature = "sqlite")]
#[test]
fn shared_cache_in_memory_sqlite_pool_is_not_forced_single_slot() {
let config = DatabaseConfig {
primary_url: Some("file::memory:?cache=shared".into()),
pool_size: 5,
..Default::default()
};
let pool = create_pool(&config)
.expect("shared-cache in-memory sqlite pool builds")
.expect("a configured url yields a pool");
assert_eq!(
pool.status().max_size,
5,
"a shared-cache in-memory target must respect the configured size"
);
}
#[cfg(feature = "sqlite")]
#[test]
fn normalize_sqlite_target_strips_schemes() {
assert_eq!(normalize_sqlite_target("sqlite::memory:"), ":memory:");
assert_eq!(normalize_sqlite_target("sqlite://:memory:"), ":memory:");
assert_eq!(normalize_sqlite_target("sqlite://"), ":memory:");
assert_eq!(
normalize_sqlite_target("sqlite:///var/lib/app.db"),
"/var/lib/app.db"
);
assert_eq!(normalize_sqlite_target("sqlite:app.db"), "app.db");
assert_eq!(
normalize_sqlite_target("file:/var/lib/app.db"),
"file:/var/lib/app.db"
);
}
#[test]
fn pool_respects_max_size() {
let config = DatabaseConfig {
url: Some("postgres://localhost/test".into()),
pool_size: 5,
..Default::default()
};
let pool = create_pool(&config)
.expect("should build pool")
.expect("should be Some");
assert_eq!(pool.status().max_size, 5);
}
#[test]
fn pool_clamps_size_to_one_if_zero() {
let config = DatabaseConfig {
url: Some("postgres://localhost/test".into()),
pool_size: 0,
..Default::default()
};
let pool = create_pool(&config)
.expect("should build pool")
.expect("should be Some");
assert_eq!(
pool.status().max_size,
1,
"Pool size should be clamped to 1"
);
}
#[test]
fn database_topology_builds_primary_and_replica_pools() {
let config = DatabaseConfig {
primary_url: Some("postgres://localhost/primary".into()),
replica_url: Some("postgres://localhost/replica".into()),
primary_pool_size: Some(6),
replica_pool_size: Some(2),
..Default::default()
};
let topology = create_topology(&config)
.expect("topology should build")
.expect("topology should be configured");
assert_eq!(topology.primary().status().max_size, 6);
assert_eq!(
topology.replica().expect("replica pool").status().max_size,
2
);
assert_eq!(topology.read().status().max_size, 2);
}
#[test]
fn database_topology_single_url_builds_only_primary_pool() {
let config = DatabaseConfig {
url: Some("postgres://localhost/single".into()),
pool_size: 5,
..Default::default()
};
let topology = create_topology(&config)
.expect("topology should build")
.expect("topology should be configured");
assert_eq!(topology.primary().status().max_size, 5);
assert!(topology.replica().is_none());
assert_eq!(topology.read().status().max_size, 5);
}
#[test]
fn config_runtime_drift_pool_applies_connect_timeout_to_wait_and_create() {
let config = DatabaseConfig {
url: Some("postgres://localhost/test".into()),
connect_timeout_secs: 7,
..Default::default()
};
let pool = create_pool(&config)
.expect("should build pool")
.expect("should be Some");
let timeouts = pool.timeouts();
assert_eq!(timeouts.wait, Some(Duration::from_secs(7)));
assert_eq!(timeouts.create, Some(Duration::from_secs(7)));
}
#[derive(Clone)]
struct TestDbState;
impl DbState for TestDbState {
fn pool(&self) -> Option<&Pool<crate::db::RuntimeConnection>> {
None
}
}
#[derive(Clone)]
struct TestReadState {
primary: Pool<crate::db::RuntimeConnection>,
}
impl DbState for TestReadState {
fn pool(&self) -> Option<&Pool<crate::db::RuntimeConnection>> {
Some(&self.primary)
}
}
#[test]
fn database_topology_read_pool_falls_back_to_primary() {
let config = DatabaseConfig {
url: Some("postgres://localhost/read-fallback".into()),
pool_size: 3,
..Default::default()
};
let primary = create_pool(&config).unwrap().unwrap();
let state = TestReadState { primary };
assert_eq!(state.read_pool().expect("read pool").status().max_size, 3);
}
#[tokio::test]
async fn db_extractor_rejects_when_no_pool() {
use axum::Router;
use axum::body::Body;
use axum::http::{Request, StatusCode};
use axum::routing::get;
use tower::ServiceExt;
async fn handler(_db: Db) -> &'static str {
"ok"
}
let app = Router::new()
.route("/", get(handler))
.with_state(TestDbState);
let response = app
.oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}
#[tokio::test]
async fn database_topology_primary_only_has_no_replica() {
let config = DatabaseConfig {
primary_url: Some("postgres://user:pass@localhost/db".to_string()),
..DatabaseConfig::default()
};
let topology = create_topology(&config).unwrap().unwrap();
let primary = topology.primary().clone();
let new_topology = DatabaseTopology::primary_only(primary);
assert!(
new_topology.replica().is_none(),
"primary_only must set replica to None"
);
}
#[tokio::test]
async fn database_topology_from_pools_retains_replica() {
let config = DatabaseConfig {
primary_url: Some("postgres://user:pass@localhost/db".to_string()),
replica_url: Some("postgres://user:pass@localhost/db_replica".to_string()),
..DatabaseConfig::default()
};
let topology = create_topology(&config).unwrap().unwrap();
let primary = topology.primary().clone();
let replica = topology.replica().cloned();
let new_topology = DatabaseTopology::from_pools(primary, replica);
assert!(
new_topology.replica().is_some(),
"from_pools must preserve the replica pool"
);
}
#[test]
fn scrub_sql_strips_string_literals() {
assert_eq!(
super::scrub_sql("SELECT * FROM users WHERE name = 'Alice'"),
"SELECT * FROM users WHERE name = '?'"
);
}
#[test]
fn scrub_sql_strips_numeric_literals() {
assert_eq!(
super::scrub_sql("SELECT * FROM orders WHERE id = 42"),
"SELECT * FROM orders WHERE id = ?"
);
}
#[test]
fn scrub_sql_preserves_pg_positional_params() {
assert_eq!(
super::scrub_sql("SELECT * FROM t WHERE x = $1 AND y = $2"),
"SELECT * FROM t WHERE x = $1 AND y = $2"
);
}
#[test]
fn scrub_sql_does_not_stomp_identifiers() {
assert_eq!(
super::scrub_sql("SELECT * FROM table1 WHERE active = true"),
"SELECT * FROM table1 WHERE active = true"
);
}
#[test]
fn scrub_sql_multiple_literals_in_one_query() {
assert_eq!(
super::scrub_sql("INSERT INTO users (name, age) VALUES ('Bob', 30)"),
"INSERT INTO users (name, age) VALUES ('?', ?)"
);
}
#[test]
fn scrub_sql_handles_escaped_single_quotes() {
assert_eq!(
super::scrub_sql("SELECT * FROM t WHERE s = 'it''s a test'"),
"SELECT * FROM t WHERE s = '?'"
);
}
#[test]
fn scrub_sql_empty_string() {
assert_eq!(super::scrub_sql(""), "");
}
#[test]
fn scrub_sql_scientific_notation_integer_exponent() {
assert_eq!(
super::scrub_sql("SELECT * FROM t WHERE n = 1e6"),
"SELECT * FROM t WHERE n = ?"
);
}
#[test]
fn scrub_sql_scientific_notation_float_exponent() {
assert_eq!(
super::scrub_sql("SELECT * FROM t WHERE n = 2.5E-4"),
"SELECT * FROM t WHERE n = ?"
);
}
#[test]
fn scrub_sql_scientific_notation_uppercase_positive_exponent() {
assert_eq!(
super::scrub_sql("SELECT * FROM t WHERE n = 3E+10"),
"SELECT * FROM t WHERE n = ?"
);
}
#[test]
fn scrub_sql_dollar_quoted_anonymous() {
assert_eq!(super::scrub_sql("SELECT $$secret value$$"), "SELECT '?'");
}
#[test]
fn scrub_sql_dollar_quoted_with_tag() {
assert_eq!(
super::scrub_sql("SELECT $body$hello world$body$"),
"SELECT '?'"
);
}
#[test]
fn scrub_sql_dollar_quoted_does_not_affect_positional_params() {
assert_eq!(
super::scrub_sql("SELECT $1, $2 FROM $$secret$$ WHERE id = $3"),
"SELECT $1, $2 FROM '?' WHERE id = $3"
);
}
#[test]
fn scrub_sql_estring_backslash_escaped_quote() {
assert_eq!(
super::scrub_sql(r"SELECT E'it\'s secret' FROM t"),
"SELECT '?' FROM t"
);
}
#[test]
fn scrub_sql_estring_uppercase() {
assert_eq!(
super::scrub_sql("SELECT E'hello world' FROM t"),
"SELECT '?' FROM t"
);
}
#[test]
fn scrub_sql_estring_multiple_backslash_escapes() {
assert_eq!(
super::scrub_sql(r"SELECT E'line1\nline2' FROM t"),
"SELECT '?' FROM t"
);
}
#[test]
fn scrub_sql_leading_dot_numeric_literals() {
assert_eq!(super::scrub_sql("SELECT .5"), "SELECT ?");
assert_eq!(super::scrub_sql("SELECT .25 + .75"), "SELECT ? + ?");
assert_eq!(super::scrub_sql("SELECT t.col"), "SELECT t.col");
assert_eq!(
super::scrub_sql("SELECT schema.table.col"),
"SELECT schema.table.col"
);
}
#[test]
fn scrub_sql_estring_doubled_quote_escape() {
assert_eq!(
super::scrub_sql("SELECT E'it''s secret' FROM t"),
"SELECT '?' FROM t"
);
}
#[test]
fn scrub_sql_numeric_literal_underscore_grouping() {
assert_eq!(super::scrub_sql("SELECT 5_432_000"), "SELECT ?");
assert_eq!(super::scrub_sql("SELECT 1_000.5_0"), "SELECT ?");
assert_eq!(super::scrub_sql("SELECT col_5_val"), "SELECT col_5_val");
assert_eq!(super::scrub_sql("SELECT col_5"), "SELECT col_5");
}
#[derive(Debug)]
struct FakeDbErrorInfo {
message: &'static str,
}
impl diesel::result::DatabaseErrorInformation for FakeDbErrorInfo {
fn message(&self) -> &str {
self.message
}
fn details(&self) -> Option<&str> {
None
}
fn hint(&self) -> Option<&str> {
None
}
fn table_name(&self) -> Option<&str> {
None
}
fn column_name(&self) -> Option<&str> {
None
}
fn constraint_name(&self) -> Option<&str> {
None
}
fn statement_position(&self) -> Option<i32> {
None
}
}
fn diesel_db_error(
kind: diesel::result::DatabaseErrorKind,
message: &'static str,
) -> diesel::result::Error {
diesel::result::Error::DatabaseError(kind, Box::new(FakeDbErrorInfo { message }))
}
#[test]
fn tx_options_default_is_read_committed_no_retry() {
let opts = TxOptions::default();
assert_eq!(opts.isolation, IsolationLevel::ReadCommitted);
assert!(!opts.read_only);
assert!(!opts.deferrable);
assert_eq!(
opts.max_attempts, 1,
"default must behave exactly like today's tx(): one attempt, no retry"
);
}
#[test]
fn tx_options_serializable_enables_retry() {
let opts = TxOptions::serializable();
assert_eq!(opts.isolation, IsolationLevel::Serializable);
assert!(
opts.max_attempts > 1,
"serializable() should default to retrying since retry is the point"
);
}
#[test]
fn tx_options_repeatable_read_enables_retry() {
let opts = TxOptions::repeatable_read();
assert_eq!(opts.isolation, IsolationLevel::RepeatableRead);
assert!(opts.max_attempts > 1);
}
#[test]
fn tx_options_builder_chains() {
let opts = TxOptions::serializable()
.read_only()
.deferrable()
.max_attempts(9)
.initial_backoff(Duration::from_millis(3))
.max_backoff(Duration::from_millis(30));
assert!(opts.read_only);
assert!(opts.deferrable);
assert_eq!(opts.max_attempts, 9);
assert_eq!(opts.initial_backoff, Duration::from_millis(3));
assert_eq!(opts.max_backoff, Duration::from_millis(30));
}
#[test]
fn tx_options_max_attempts_clamps_to_at_least_one() {
assert_eq!(TxOptions::default().max_attempts(0).max_attempts, 1);
}
#[test]
fn tx_options_effective_max_attempts_clamps_struct_literal_bypass() {
let opts = TxOptions {
max_attempts: 0,
..TxOptions::default()
};
assert_eq!(opts.max_attempts, 0, "the raw field is not itself clamped");
assert_eq!(
opts.effective_max_attempts(),
1,
"the retry loop's accessor must clamp a struct-literal 0 to 1"
);
}
#[test]
fn retry_backoff_base_doubles_per_attempt() {
let initial = Duration::from_millis(5);
let max = Duration::from_secs(60);
assert_eq!(
retry_backoff_base(initial, max, 1),
Duration::from_millis(5)
);
assert_eq!(
retry_backoff_base(initial, max, 2),
Duration::from_millis(10)
);
assert_eq!(
retry_backoff_base(initial, max, 3),
Duration::from_millis(20)
);
assert_eq!(
retry_backoff_base(initial, max, 4),
Duration::from_millis(40)
);
}
#[test]
fn retry_backoff_base_caps_at_max() {
let initial = Duration::from_millis(5);
let max = Duration::from_millis(50);
assert_eq!(
retry_backoff_base(initial, max, 5),
Duration::from_millis(50)
);
assert_eq!(retry_backoff_base(initial, max, 40), max);
assert_eq!(retry_backoff_base(initial, max, u32::MAX), max);
}
#[test]
fn retry_backoff_delay_stays_within_jitter_bounds_and_cap() {
let initial = Duration::from_millis(100);
let max = Duration::from_secs(10);
for attempt in 1..=6u32 {
let base = retry_backoff_base(initial, max, attempt);
for _ in 0..64 {
let d = retry_backoff_delay(initial, max, attempt);
assert!(
d >= base.mul_f64(0.8) && d <= base.mul_f64(1.2),
"delay {d:?} out of jitter bounds for base {base:?}"
);
assert!(d <= max, "delay {d:?} exceeded cap {max:?}");
}
}
}
#[test]
fn serialization_failure_is_retryable() {
let err: AutumnError = AutumnError::internal_server_error(diesel_db_error(
diesel::result::DatabaseErrorKind::SerializationFailure,
"could not serialize access due to read/write dependencies among transactions",
));
assert!(is_retryable_txn_error(&err));
}
#[test]
fn deadlock_is_retryable() {
let err: AutumnError = AutumnError::internal_server_error(diesel_db_error(
diesel::result::DatabaseErrorKind::Unknown,
"deadlock detected",
));
assert!(is_retryable_txn_error(&err));
}
#[test]
fn deadlock_message_is_matched_case_and_whitespace_insensitively() {
let err: AutumnError = AutumnError::internal_server_error(diesel_db_error(
diesel::result::DatabaseErrorKind::Unknown,
" Deadlock Detected ",
));
assert!(is_retryable_txn_error(&err));
}
#[test]
fn unknown_kind_message_merely_mentioning_deadlock_is_not_retryable() {
let err: AutumnError = AutumnError::internal_server_error(diesel_db_error(
diesel::result::DatabaseErrorKind::Unknown,
"order 40001 could not be processed: deadlock detected in workflow",
));
assert!(!is_retryable_txn_error(&err));
}
#[test]
fn unknown_kind_40001_text_alone_is_not_retryable() {
let err: AutumnError = AutumnError::internal_server_error(diesel_db_error(
diesel::result::DatabaseErrorKind::Unknown,
"ERROR: 40001: could not serialize access",
));
assert!(!is_retryable_txn_error(&err));
}
#[test]
fn unique_violation_is_not_retryable() {
let err: AutumnError = AutumnError::internal_server_error(diesel_db_error(
diesel::result::DatabaseErrorKind::UniqueViolation,
"duplicate key value violates unique constraint",
));
assert!(!is_retryable_txn_error(&err));
}
#[test]
fn unique_violation_with_magic_substring_in_message_is_not_retryable() {
let err: AutumnError = AutumnError::internal_server_error(diesel_db_error(
diesel::result::DatabaseErrorKind::UniqueViolation,
"duplicate key value violates unique constraint \"orders_40001_key\"",
));
assert!(!is_retryable_txn_error(&err));
}
#[test]
fn plain_internal_error_is_not_retryable() {
let err = AutumnError::internal_server_error_msg("something unrelated broke");
assert!(!is_retryable_txn_error(&err));
}
#[test]
fn domain_error_with_magic_substring_and_no_db_error_is_not_retryable() {
let err = AutumnError::bad_request_msg(
"order 40001 could not be processed: deadlock detected in workflow",
);
assert!(!is_retryable_txn_error(&err));
}
#[derive(Debug, thiserror::Error)]
#[error("app error")]
struct WrappedDbError(#[source] diesel::result::Error);
#[test]
fn serialization_failure_wrapped_in_custom_error_is_retryable() {
let err: AutumnError = AutumnError::internal_server_error(WrappedDbError(diesel_db_error(
diesel::result::DatabaseErrorKind::SerializationFailure,
"could not serialize access due to read/write dependencies among transactions",
)));
assert!(is_retryable_txn_error(&err));
}
#[test]
fn unique_violation_wrapped_in_custom_error_is_not_retryable() {
let err: AutumnError = AutumnError::internal_server_error(WrappedDbError(diesel_db_error(
diesel::result::DatabaseErrorKind::UniqueViolation,
"duplicate key value violates unique constraint",
)));
assert!(!is_retryable_txn_error(&err));
}
#[test]
fn retry_decision_retries_until_attempts_exhausted() {
let max = 3;
assert_eq!(retry_decision(1, max, true), RetryDecision::Retry);
assert_eq!(retry_decision(2, max, true), RetryDecision::Retry);
assert_eq!(
retry_decision(3, max, true),
RetryDecision::Stop,
"the final attempt must stop even when the error is retryable (exhausted)"
);
}
#[test]
fn retry_decision_stops_immediately_on_non_retryable() {
assert_eq!(retry_decision(1, 5, false), RetryDecision::Stop);
}
#[test]
fn retry_decision_single_attempt_never_retries() {
assert_eq!(retry_decision(1, 1, true), RetryDecision::Stop);
}
}
mod tls {
use std::sync::Arc;
use diesel::{ConnectionError, ConnectionResult};
use diesel_async::pooled_connection::SetupCallback;
use diesel_async::{AsyncConnection as _, AsyncPgConnection};
use futures::FutureExt as _;
use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::crypto::CryptoProvider;
use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
use rustls::{DigitallySignedStruct, SignatureScheme};
use tokio_postgres_rustls::MakeRustlsConnect;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum TlsPosture {
Off,
Require,
VerifyFull {
root_cert: Option<String>,
},
Unsupported { reason: String },
}
impl TlsPosture {
pub(super) fn from_database_url(database_url: &str) -> Self {
let params = ssl_params(database_url);
let get = |key: &str| {
params
.iter()
.rev()
.find(|(k, _)| k == key)
.map(|(_, v)| v.as_str())
};
let root_cert = get("sslrootcert").map(str::to_owned);
match get("sslmode") {
Some("require") => {
if root_cert.is_some() {
Self::Unsupported {
reason: "sslmode=require with sslrootcert is not supported: \
PostgreSQL treats that combination as requiring \
certificate verification against the CA file, which \
this pool implements only for sslmode=verify-full. \
Use sslmode=verify-full to verify the certificate, or \
sslmode=require without sslrootcert to encrypt without \
verifying it."
.to_owned(),
}
} else {
Self::Require
}
}
Some("verify-full") => Self::VerifyFull { root_cert },
Some("verify-ca") => Self::Unsupported {
reason: "sslmode=verify-ca is not supported: chain-only verification \
(without hostname checking) is not implemented. Use \
sslmode=verify-full (stricter) or sslmode=require (encrypts \
without verifying the certificate)."
.to_owned(),
},
_ => Self::Off,
}
}
}
#[cfg_attr(feature = "sqlite", allow(dead_code))]
pub(super) fn setup_callback(posture: TlsPosture) -> SetupCallback<AsyncPgConnection> {
Box::new(move |url: &str| {
let posture = posture.clone();
let url = url.to_owned();
async move { establish(&url, posture).await }.boxed()
})
}
pub(super) async fn establish(
url: &str,
posture: TlsPosture,
) -> ConnectionResult<AsyncPgConnection> {
match posture {
TlsPosture::Off => AsyncPgConnection::establish(url).await,
TlsPosture::Require => connect_with(url, relaxed_connector()?).await,
TlsPosture::VerifyFull { root_cert } => {
let sanitized = sanitize_for_verify_full(url);
connect_with(&sanitized, verifying_connector(root_cert.as_deref())?).await
}
TlsPosture::Unsupported { reason } => {
Err(ConnectionError::InvalidConnectionUrl(reason))
}
}
}
async fn connect_with(
url: &str,
tls: MakeRustlsConnect,
) -> ConnectionResult<AsyncPgConnection> {
let (client, connection) = tokio_postgres::connect(url, tls).await.map_err(|e| {
let msg = std::error::Error::source(&e)
.map_or_else(|| e.to_string(), |source| format!("{e}: {source}"));
ConnectionError::BadConnection(msg)
})?;
AsyncPgConnection::try_from_client_and_connection(client, connection).await
}
fn relaxed_connector() -> Result<MakeRustlsConnect, ConnectionError> {
let provider = Arc::new(rustls::crypto::ring::default_provider());
let config = rustls::ClientConfig::builder_with_provider(provider.clone())
.with_safe_default_protocol_versions()
.map_err(|e| {
ConnectionError::BadConnection(format!("failed to build TLS config: {e}"))
})?
.dangerous()
.with_custom_certificate_verifier(Arc::new(NoServerCertVerification(
(*provider).clone(),
)))
.with_no_client_auth();
Ok(MakeRustlsConnect::new(config))
}
fn verifying_connector(root_cert: Option<&str>) -> Result<MakeRustlsConnect, ConnectionError> {
let mut roots = rustls::RootCertStore::empty();
roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
if let Some(path) = root_cert {
use rustls_pki_types::pem::PemObject as _;
let certs = CertificateDer::pem_file_iter(path).map_err(|e| {
ConnectionError::BadConnection(format!("failed to read sslrootcert {path}: {e}"))
})?;
for cert in certs {
let cert = cert.map_err(|e| {
ConnectionError::BadConnection(format!(
"failed to parse sslrootcert {path}: {e}"
))
})?;
roots.add(cert).map_err(|e| {
ConnectionError::BadConnection(format!(
"failed to trust sslrootcert {path}: {e}"
))
})?;
}
}
let provider = Arc::new(rustls::crypto::ring::default_provider());
let config = rustls::ClientConfig::builder_with_provider(provider)
.with_safe_default_protocol_versions()
.map_err(|e| {
ConnectionError::BadConnection(format!("failed to build TLS config: {e}"))
})?
.with_root_certificates(roots)
.with_no_client_auth();
Ok(MakeRustlsConnect::new(config))
}
use crate::pg_conn_str::is_url as is_url_connection_string;
use crate::pg_conn_str::keyword_value_pairs;
fn ssl_params(database_url: &str) -> Vec<(String, String)> {
if is_url_connection_string(database_url) {
url::Url::parse(database_url)
.map(|u| {
u.query_pairs()
.map(|(k, v)| (k.into_owned(), v.into_owned()))
.collect()
})
.unwrap_or_default()
} else {
keyword_value_pairs(database_url).unwrap_or_default()
}
}
fn quote_keyword_value(value: &str) -> String {
if !value.is_empty()
&& !value.contains(|c: char| c.is_whitespace() || c == '\'' || c == '\\')
{
return value.to_owned();
}
let mut quoted = String::with_capacity(value.len() + 2);
quoted.push('\'');
for c in value.chars() {
if c == '\'' || c == '\\' {
quoted.push('\\');
}
quoted.push(c);
}
quoted.push('\'');
quoted
}
fn sanitize_for_verify_full(database_url: &str) -> String {
if is_url_connection_string(database_url) {
let Ok(mut parsed) = url::Url::parse(database_url) else {
return database_url.to_owned();
};
let raw = parsed
.query()
.unwrap_or("")
.split('&')
.filter(|component| !component.is_empty())
.filter_map(|component| {
let key = component.split('=').next().unwrap_or(component);
match key {
"sslmode" => Some("sslmode=require"),
"sslrootcert" => None,
_ => Some(component),
}
})
.collect::<Vec<_>>()
.join("&");
if raw.is_empty() {
parsed.set_query(None);
} else {
parsed.set_query(Some(&raw));
}
parsed.to_string()
} else {
let Some(pairs) = keyword_value_pairs(database_url) else {
return database_url.to_owned();
};
pairs
.into_iter()
.filter_map(|(k, v)| match k.as_str() {
"sslmode" => Some("sslmode=require".to_owned()),
"sslrootcert" => None,
_ => Some(format!("{k}={}", quote_keyword_value(&v))),
})
.collect::<Vec<_>>()
.join(" ")
}
}
#[derive(Debug)]
struct NoServerCertVerification(CryptoProvider);
impl ServerCertVerifier for NoServerCertVerification {
fn verify_server_cert(
&self,
_end_entity: &CertificateDer<'_>,
_intermediates: &[CertificateDer<'_>],
_server_name: &ServerName<'_>,
_ocsp_response: &[u8],
_now: UnixTime,
) -> Result<ServerCertVerified, rustls::Error> {
Ok(ServerCertVerified::assertion())
}
fn verify_tls12_signature(
&self,
message: &[u8],
cert: &CertificateDer<'_>,
dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
rustls::crypto::verify_tls12_signature(
message,
cert,
dss,
&self.0.signature_verification_algorithms,
)
}
fn verify_tls13_signature(
&self,
message: &[u8],
cert: &CertificateDer<'_>,
dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
rustls::crypto::verify_tls13_signature(
message,
cert,
dss,
&self.0.signature_verification_algorithms,
)
}
fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
self.0.signature_verification_algorithms.supported_schemes()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn absent_disable_and_prefer_keep_the_default_notls_path() {
for url in [
"postgres://user:pass@db.example.com:5432/app",
"postgres://user:pass@db.example.com:5432/app?sslmode=disable",
"postgres://user:pass@db.example.com:5432/app?sslmode=prefer",
"host=db.example.com user=user",
"host=db.example.com user=user sslmode=disable",
] {
assert_eq!(
TlsPosture::from_database_url(url),
TlsPosture::Off,
"{url} must keep the default NoTls path"
);
}
}
#[test]
fn unrecognized_sslmode_values_keep_the_default_path_and_its_errors() {
assert_eq!(
TlsPosture::from_database_url("postgres://u@h/db?sslmode=allow"),
TlsPosture::Off
);
assert_eq!(
TlsPosture::from_database_url("postgres://u@h/db?sslmode=bogus"),
TlsPosture::Off
);
}
#[test]
fn require_selects_the_relaxed_tls_connector() {
assert_eq!(
TlsPosture::from_database_url(
"postgres://user:pass@db.example.com:5432/app?sslmode=require"
),
TlsPosture::Require
);
assert_eq!(
TlsPosture::from_database_url("host=db.example.com sslmode=require"),
TlsPosture::Require
);
}
#[test]
fn verify_full_selects_the_verifying_connector_with_optional_root() {
assert_eq!(
TlsPosture::from_database_url("postgres://u@h/db?sslmode=verify-full"),
TlsPosture::VerifyFull { root_cert: None }
);
assert_eq!(
TlsPosture::from_database_url(
"postgres://u@h/db?sslmode=verify-full&sslrootcert=/etc/ca.pem"
),
TlsPosture::VerifyFull {
root_cert: Some("/etc/ca.pem".to_owned())
}
);
}
#[test]
fn keyword_strings_with_spaces_around_equals_are_parsed() {
assert_eq!(
TlsPosture::from_database_url("host=db user=u sslmode = require"),
TlsPosture::Require
);
assert_eq!(
TlsPosture::from_database_url("sslmode = require"),
TlsPosture::Require
);
assert_eq!(
TlsPosture::from_database_url("host = db sslmode\t=\nverify-full user=u"),
TlsPosture::VerifyFull { root_cert: None }
);
}
#[test]
fn keyword_strings_with_quoted_values_are_parsed() {
assert_eq!(
TlsPosture::from_database_url("sslmode='verify-full'"),
TlsPosture::VerifyFull { root_cert: None }
);
assert_eq!(
TlsPosture::from_database_url(
"host=db sslmode='verify-full' sslrootcert='/path with space/ca.pem'"
),
TlsPosture::VerifyFull {
root_cert: Some("/path with space/ca.pem".to_owned())
}
);
assert_eq!(
TlsPosture::from_database_url(r"sslmode=verify-full sslrootcert='/pki/it\'s.pem'"),
TlsPosture::VerifyFull {
root_cert: Some("/pki/it's.pem".to_owned())
}
);
assert_eq!(
TlsPosture::from_database_url(r"sslmode=verify-full sslrootcert=/pki/my\ ca.pem"),
TlsPosture::VerifyFull {
root_cert: Some("/pki/my ca.pem".to_owned())
}
);
}
#[test]
fn malformed_keyword_strings_keep_the_default_path_and_its_errors() {
for url in [
"host=db sslmode", "host=db sslmode=", "host=db sslmode='require", r"host=db sslmode='require\", ] {
assert_eq!(
TlsPosture::from_database_url(url),
TlsPosture::Off,
"{url:?} must keep the default NoTls path"
);
}
}
#[test]
fn keyword_strings_with_url_like_values_stay_on_the_keyword_path() {
assert_eq!(
TlsPosture::from_database_url("host=db password=https://secret sslmode=require"),
TlsPosture::Require
);
assert_eq!(
TlsPosture::from_database_url(
"host=db options='-c foo=bar://baz' sslmode='verify-full'"
),
TlsPosture::VerifyFull { root_cert: None }
);
assert_eq!(
TlsPosture::from_database_url("postgresql://u@h/db?sslmode=require"),
TlsPosture::Require
);
}
#[test]
fn verify_ca_and_require_with_rootcert_fail_loudly() {
assert!(matches!(
TlsPosture::from_database_url("postgres://u@h/db?sslmode=verify-ca"),
TlsPosture::Unsupported { .. }
));
assert!(matches!(
TlsPosture::from_database_url(
"postgres://u@h/db?sslmode=require&sslrootcert=/etc/ca.pem"
),
TlsPosture::Unsupported { .. }
));
}
#[test]
fn last_sslmode_occurrence_wins() {
assert_eq!(
TlsPosture::from_database_url("postgres://u@h/db?sslmode=disable&sslmode=require"),
TlsPosture::Require
);
}
#[test]
fn sanitize_rewrites_verify_full_and_drops_sslrootcert() {
let sanitized = sanitize_for_verify_full(
"postgres://u@h:5432/db?application_name=app&sslmode=verify-full&sslrootcert=/etc/ca.pem",
);
assert!(sanitized.contains("sslmode=require"), "{sanitized}");
assert!(!sanitized.contains("verify-full"), "{sanitized}");
assert!(!sanitized.contains("sslrootcert"), "{sanitized}");
assert!(sanitized.contains("application_name=app"), "{sanitized}");
let sanitized = sanitize_for_verify_full(
"postgres://u@h/db?options=-c%20search_path%3Dtenant&sslmode=verify-full&sslrootcert=/etc/ca.pem",
);
assert_eq!(
sanitized, "postgres://u@h/db?options=-c%20search_path%3Dtenant&sslmode=require",
"raw components must be preserved verbatim"
);
assert_eq!(
sanitize_for_verify_full("postgres://u@h/db?sslrootcert=/etc/ca.pem"),
"postgres://u@h/db"
);
let kv = sanitize_for_verify_full(
"host=h user=u sslmode=verify-full sslrootcert=/etc/ca.pem dbname=db",
);
assert_eq!(kv, "host=h user=u sslmode=require dbname=db");
let kv = sanitize_for_verify_full(
r"host=h sslmode = 'verify-full' sslrootcert='/path with space/ca.pem' password='it\'s'",
);
assert_eq!(kv, r"host=h sslmode=require password='it\'s'");
let kv = sanitize_for_verify_full(
"host=h password=https://secret sslmode=verify-full sslrootcert=/etc/ca.pem",
);
assert_eq!(kv, "host=h password=https://secret sslmode=require");
}
#[test]
fn verifying_connector_builds_with_and_without_root_cert_file() {
assert!(verifying_connector(None).is_ok());
assert!(
verifying_connector(Some("/nonexistent/ca.pem")).is_err(),
"an unreadable sslrootcert must fail loudly"
);
}
#[test]
fn relaxed_connector_builds() {
assert!(relaxed_connector().is_ok());
}
}
}
#[allow(clippy::large_enum_variant)] pub(crate) enum MigrationConnection {
Native(diesel::PgConnection),
Rustls {
runtime: Option<tokio::runtime::Runtime>,
conn: diesel_async::async_connection_wrapper::AsyncConnectionWrapper<AsyncPgConnection>,
},
}
fn migration_connection_needs_rustls(database_url: &str) -> bool {
!matches!(
tls::TlsPosture::from_database_url(database_url),
tls::TlsPosture::Off
)
}
pub(crate) fn establish_migration_connection(
database_url: &str,
) -> Result<MigrationConnection, diesel::ConnectionError> {
use diesel::Connection as _;
if !migration_connection_needs_rustls(database_url) {
return diesel::PgConnection::establish(database_url).map(MigrationConnection::Native);
}
let posture = tls::TlsPosture::from_database_url(database_url);
let (runtime, handle) = if let Ok(handle) = tokio::runtime::Handle::try_current() {
(None, handle)
} else {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.map_err(|e| {
diesel::ConnectionError::BadConnection(format!(
"failed to build a tokio runtime for the TLS migration connection: {e}"
))
})?;
let handle = runtime.handle().clone();
(Some(runtime), handle)
};
let inner = handle.block_on(tls::establish(database_url, posture))?;
Ok(MigrationConnection::Rustls {
runtime,
conn: diesel_async::async_connection_wrapper::AsyncConnectionWrapper::from(inner),
})
}
#[cfg(feature = "sqlite")]
pub(crate) fn establish_sqlite_migration_connection(
database_url: &str,
) -> Result<diesel::SqliteConnection, diesel::ConnectionError> {
use diesel::Connection as _;
use diesel::connection::SimpleConnection as _;
let mut conn = diesel::SqliteConnection::establish(&normalize_sqlite_target(database_url))?;
conn.batch_execute("PRAGMA busy_timeout = 5000;")
.map_err(diesel::ConnectionError::CouldntSetupConfiguration)?;
Ok(conn)
}
#[cfg(test)]
mod migration_connection_tests {
use super::migration_connection_needs_rustls;
#[test]
fn migration_path_selection_follows_the_tls_posture() {
for url in [
"postgres://u@h/db",
"postgres://u@h/db?sslmode=disable",
"postgres://u@h/db?sslmode=prefer",
"host=db user=u",
] {
assert!(
!migration_connection_needs_rustls(url),
"must stay on the native path: {url}"
);
}
for url in [
"postgres://u@h/db?sslmode=require",
"postgres://u@h/db?sslmode=verify-full",
"postgres://u@h/db?sslmode=verify-full&sslrootcert=/etc/ca.pem",
"postgres://u@h/db?sslmode=verify-ca",
"host=db user=u sslmode=require",
"host=db sslmode = 'verify-full'",
] {
assert!(
migration_connection_needs_rustls(url),
"must use the rustls wrapper: {url}"
);
}
}
}