use crate::{DataLayerError, PostgresPoolConfig}; use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode}; use sqlx::PgPool; use std::str::FromStr; use std::time::Duration; const STATEMENT_TIMEOUT_ENV: &str = "AETHER_GATEWAY_DATA_POSTGRES_STATEMENT_TIMEOUT_MS"; const LOCK_TIMEOUT_ENV: &str = "AETHER_GATEWAY_DATA_POSTGRES_LOCK_TIMEOUT_MS"; #[derive(Debug, Clone, Copy)] struct PostgresSessionTimeouts { statement_ms: u32, lock_ms: u32, } impl PostgresSessionTimeouts { fn from_env() -> Result { Ok(Self { statement_ms: read_timeout_env(STATEMENT_TIMEOUT_ENV, 30_000)?, lock_ms: read_timeout_env(LOCK_TIMEOUT_ENV, 3_000)?, }) } fn apply(self, options: PgConnectOptions) -> PgConnectOptions { options.options([ ("statement_timeout", self.statement_ms), ("lock_timeout", self.lock_ms), ]) } } fn parse_timeout_ms(name: &str, value: &str) -> Result { value .trim() .parse::() .ok() .filter(|value| *value <= i32::MAX as u32) .ok_or_else(|| { DataLayerError::InvalidConfiguration(format!( "{name} must be milliseconds in 0..=2147483647 (0 disables the timeout)" )) }) } fn read_timeout_env(name: &str, default: u32) -> Result { match std::env::var(name) { Ok(value) => parse_timeout_ms(name, &value), Err(std::env::VarError::NotPresent) => Ok(default), Err(std::env::VarError::NotUnicode(_)) => Err(DataLayerError::InvalidConfiguration( format!("{name} must contain a valid integer"), )), } } /// Migration and historical backfill connections are discarded on every exit path, /// including cancellation, so their relaxed deadlines cannot escape into request work. pub async fn acquire_postgres_migration_connection( pool: &PgPool, ) -> Result, sqlx::Error> { let mut conn = pool.acquire().await?; conn.close_on_drop(); sqlx::query("SET statement_timeout = 0") .execute(&mut *conn) .await?; sqlx::query("SET lock_timeout = 0") .execute(&mut *conn) .await?; Ok(conn) } fn connect_options(config: &PostgresPoolConfig) -> Result { config.validate()?; let options = PgConnectOptions::from_str(config.database_url.trim()).map_err(|err| { DataLayerError::InvalidConfiguration(format!("invalid postgres database_url: {err}")) })?; // Preserve an explicit verification mode from the URL. `require_ssl` is // a minimum transport guarantee, so it may upgrade Disable/Allow/Prefer // to Require but must never silently weaken VerifyCa/VerifyFull. let ssl_mode = if config.require_ssl && !matches!( options.get_ssl_mode(), PgSslMode::VerifyCa | PgSslMode::VerifyFull ) { PgSslMode::Require } else { options.get_ssl_mode() }; Ok(options .ssl_mode(ssl_mode) .statement_cache_capacity(config.statement_cache_capacity)) } pub type PostgresPool = PgPool; #[derive(Debug, Clone)] pub struct PostgresPoolFactory { config: PostgresPoolConfig, timeouts: PostgresSessionTimeouts, } impl PostgresPoolFactory { pub fn new(config: PostgresPoolConfig) -> Result { config.validate()?; Ok(Self { config, timeouts: PostgresSessionTimeouts::from_env()?, }) } pub fn config(&self) -> &PostgresPoolConfig { &self.config } pub fn connect_lazy(&self) -> Result { let options = self.timeouts.apply(connect_options(&self.config)?); Ok(PgPoolOptions::new() .min_connections(self.config.min_connections) .max_connections(self.config.max_connections) .acquire_timeout(Duration::from_millis(self.config.acquire_timeout_ms)) .idle_timeout(Duration::from_millis(self.config.idle_timeout_ms)) .max_lifetime(Duration::from_millis(self.config.max_lifetime_ms)) .connect_lazy_with(options)) } } #[cfg(test)] mod tests { use super::{connect_options, parse_timeout_ms, PostgresPoolFactory, PostgresSessionTimeouts}; use crate::PostgresPoolConfig; use sqlx::postgres::PgSslMode; #[tokio::test] async fn migration_connection_future_is_send() { fn assert_send(_: impl Send) {} let pool = sqlx::postgres::PgPoolOptions::new() .connect_lazy("postgres://localhost/aether") .unwrap(); assert_send(super::acquire_postgres_migration_connection(&pool)); } #[test] fn validates_session_timeout_milliseconds() { assert_eq!(parse_timeout_ms("timeout", "0").unwrap(), 0); assert_eq!(parse_timeout_ms("timeout", " 3000 ").unwrap(), 3_000); assert_eq!( parse_timeout_ms("timeout", "2147483647").unwrap(), i32::MAX as u32 ); for invalid in ["", "-1", "3s", "2147483648", "4294967296"] { assert!(parse_timeout_ms("timeout", invalid).is_err()); } } #[test] fn session_deadlines_preserve_unrelated_connection_options() { let options = PostgresSessionTimeouts { statement_ms: 30_000, lock_ms: 3_000, } .apply(sqlx::postgres::PgConnectOptions::new().options([("search_path", "audit")])); assert_eq!( options.get_options(), Some("-c search_path=audit -c statement_timeout=30000 -c lock_timeout=3000") ); } #[tokio::test] #[ignore = "requires an isolated AETHER_TEST_DATABASE_URL"] async fn live_session_deadlines_rollback_transactions_and_isolate_migration_overrides() { use crate::error::SqlxResultExt; use crate::{PostgresTransactionOptions, PostgresTransactionRunner}; let factory = PostgresPoolFactory { config: PostgresPoolConfig { database_url: std::env::var("AETHER_TEST_DATABASE_URL").expect("test database URL"), min_connections: 0, max_connections: 2, ..PostgresPoolConfig::default() }, timeouts: PostgresSessionTimeouts { statement_ms: 100, lock_ms: 40, }, }; let pool = factory.connect_lazy().unwrap(); let table = format!("deadline_test_{}", uuid::Uuid::new_v4().simple()); sqlx::query(&format!( "CREATE TABLE {table} (id INTEGER PRIMARY KEY, value INTEGER NOT NULL)" )) .execute(&pool) .await .unwrap(); sqlx::query(&format!("INSERT INTO {table} VALUES (1, 0)")) .execute(&pool) .await .unwrap(); let mut blocker = pool.begin().await.unwrap(); sqlx::query(&format!("UPDATE {table} SET value = 7 WHERE id = 1")) .execute(&mut *blocker) .await .unwrap(); let runner = PostgresTransactionRunner::new(pool.clone()); let insert = format!("INSERT INTO {table} VALUES (2, 2)"); let update = format!("UPDATE {table} SET value = 9 WHERE id = 1"); let started = std::time::Instant::now(); let error = runner .run_read_write(|tx| { Box::pin(async move { sqlx::query(&insert) .execute(&mut **tx) .await .map_postgres_err()?; sqlx::query(&update) .execute(&mut **tx) .await .map_postgres_err()?; Ok(()) }) }) .await .unwrap_err(); assert!(error.to_string().contains("SQLSTATE 55P03"), "{error}"); assert!(started.elapsed() < std::time::Duration::from_secs(2)); blocker.rollback().await.unwrap(); assert_eq!( sqlx::query_scalar::<_, i64>(&format!("SELECT COUNT(*) FROM {table}")) .fetch_one(&pool) .await .unwrap(), 1 ); assert_eq!( sqlx::query_scalar::<_, i32>(&format!("SELECT value FROM {table} WHERE id = 1")) .fetch_one(&pool) .await .unwrap(), 0 ); let error = sqlx::query("SELECT pg_sleep(0.3)") .execute(&pool) .await .map_postgres_err() .unwrap_err(); assert!(error.to_string().contains("SQLSTATE 57014"), "{error}"); runner .run( PostgresTransactionOptions { statement_timeout_ms: Some(1_000), ..PostgresTransactionOptions::read_write() }, |tx| { Box::pin(async move { sqlx::query("SELECT pg_sleep(0.15)") .execute(&mut **tx) .await .map_postgres_err()?; Ok(()) }) }, ) .await .unwrap(); let mut migration = super::acquire_postgres_migration_connection(&pool) .await .unwrap(); sqlx::query("SELECT pg_sleep(0.15)") .execute(&mut *migration) .await .unwrap(); drop(migration); for _ in 0..2 { let configured: i64 = sqlx::query_scalar( "SELECT setting::BIGINT FROM pg_settings WHERE name = 'statement_timeout'", ) .fetch_one(&pool) .await .unwrap(); assert_eq!( configured, 100, "relaxed/local overrides must not leak into pooled requests" ); } sqlx::query(&format!("DROP TABLE {table}")) .execute(&pool) .await .unwrap(); pool.close().await; } fn ssl_mode(url: &str, require_ssl: bool) -> PgSslMode { connect_options(&PostgresPoolConfig { database_url: url.to_string(), require_ssl, ..PostgresPoolConfig::default() }) .expect("postgres options should parse") .get_ssl_mode() } #[test] fn preserves_explicit_postgres_verification_modes() { assert!(matches!( ssl_mode("postgres://localhost/aether?sslmode=verify-full", false), PgSslMode::VerifyFull )); assert!(matches!( ssl_mode("postgres://localhost/aether?sslmode=verify-ca", true), PgSslMode::VerifyCa )); } #[test] fn require_ssl_only_upgrades_weak_postgres_modes() { for mode in ["disable", "allow", "prefer"] { let url = format!("postgres://localhost/aether?sslmode={mode}"); assert!(matches!(ssl_mode(&url, true), PgSslMode::Require)); } assert!(matches!( ssl_mode("postgres://localhost/aether", false), PgSslMode::Prefer )); } #[tokio::test] async fn factory_builds_lazy_pool_from_valid_config() { let config = PostgresPoolConfig { database_url: "postgres://localhost/aether".to_string(), min_connections: 1, max_connections: 4, acquire_timeout_ms: 1_000, idle_timeout_ms: 5_000, max_lifetime_ms: 30_000, statement_cache_capacity: 64, require_ssl: false, }; let factory = PostgresPoolFactory::new(config).expect("factory should build"); let _pool = factory.connect_lazy().expect("lazy pool should build"); } }