use crate::config::SslMode;
use crate::dialect::{DatabaseDialect, DialectKind};
use crate::error::{Result, WaypointError};
use std::path::PathBuf;
#[cfg(feature = "postgres")]
use fastrand;
#[cfg(feature = "postgres")]
use tokio_postgres::Client;
#[derive(Debug, Clone)]
pub struct TransportConfig {
pub ssl_mode: SslMode,
pub ssl_root_cert: Option<PathBuf>,
pub retries: u32,
pub connect_timeout_secs: u32,
pub statement_timeout_secs: u32,
pub keepalive_secs: u32,
}
impl Default for TransportConfig {
fn default() -> Self {
Self {
ssl_mode: SslMode::Prefer,
ssl_root_cert: None,
retries: 0,
connect_timeout_secs: 30,
statement_timeout_secs: 0,
keepalive_secs: 120,
}
}
}
impl TransportConfig {
pub fn from_database_config(db: &crate::config::DatabaseConfig) -> Self {
Self {
ssl_mode: db.ssl_mode,
ssl_root_cert: db.ssl_root_cert.clone(),
retries: db.connect_retries,
connect_timeout_secs: db.connect_timeout_secs,
statement_timeout_secs: db.statement_timeout_secs,
keepalive_secs: db.keepalive_secs,
}
}
}
pub fn sandbox_name(prefix: &str) -> String {
let millis = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
format!(
"{}_{}_{:x}_{:08x}",
prefix,
millis,
std::process::id(),
fastrand::u32(..)
)
}
pub fn quote_literal(value: &str) -> String {
format!("'{}'", value.replace('\'', "''"))
}
pub fn quote_ident(name: &str) -> String {
format!("\"{}\"", name.replace('"', "\"\""))
}
pub fn quote_ident_mysql(name: &str) -> String {
format!("`{}`", name.replace('`', "``"))
}
pub fn validate_identifier(name: &str) -> Result<()> {
if name.is_empty() {
return Err(WaypointError::ConfigError(
"Identifier cannot be empty".to_string(),
));
}
if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
return Err(WaypointError::ConfigError(format!(
"Identifier '{}' contains invalid characters. Only [a-zA-Z0-9_] are allowed.",
name
)));
}
Ok(())
}
pub enum DbClient {
#[cfg(feature = "postgres")]
Postgres(Client),
#[cfg(feature = "mysql")]
Mysql(mysql_async::Pool),
}
impl DbClient {
#[cfg(feature = "postgres")]
pub fn with_postgres(client: Client) -> Self {
DbClient::Postgres(client)
}
#[cfg(feature = "mysql")]
pub fn with_mysql(pool: mysql_async::Pool) -> Self {
DbClient::Mysql(pool)
}
pub fn dialect_kind(&self) -> DialectKind {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(_) => DialectKind::Postgres,
#[cfg(feature = "mysql")]
DbClient::Mysql(_) => DialectKind::Mysql,
}
}
pub fn dialect(&self) -> &'static dyn DatabaseDialect {
#[cfg(feature = "postgres")]
static PG: crate::dialect::postgres::PostgresDialect =
crate::dialect::postgres::PostgresDialect;
#[cfg(feature = "mysql")]
static MY: crate::dialect::mysql::MysqlDialect = crate::dialect::mysql::MysqlDialect;
match self.dialect_kind() {
#[cfg(feature = "postgres")]
DialectKind::Postgres => &PG,
#[cfg(not(feature = "postgres"))]
DialectKind::Postgres => {
panic!("PostgreSQL connection without `postgres` feature compiled in")
}
#[cfg(feature = "mysql")]
DialectKind::Mysql => &MY,
#[cfg(not(feature = "mysql"))]
DialectKind::Mysql => {
panic!("MySQL connection without `mysql` feature compiled in")
}
}
}
#[cfg(feature = "postgres")]
pub fn as_postgres(&self) -> Result<&Client> {
match self {
DbClient::Postgres(c) => Ok(c),
#[cfg(feature = "mysql")]
DbClient::Mysql(_) => Err(WaypointError::ConfigError(
"This operation is not yet implemented for MySQL".into(),
)),
}
}
#[cfg(feature = "mysql")]
pub fn as_mysql(&self) -> Result<&mysql_async::Pool> {
match self {
DbClient::Mysql(p) => Ok(p),
#[cfg(feature = "postgres")]
DbClient::Postgres(_) => Err(WaypointError::ConfigError(
"This operation requires a MySQL connection".into(),
)),
}
}
pub async fn check_connection(&self) -> Result<()> {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(c) => check_connection(c).await,
#[cfg(feature = "mysql")]
DbClient::Mysql(pool) => {
use mysql_async::prelude::*;
let mut conn =
pool.get_conn()
.await
.map_err(|e| WaypointError::ConnectionLost {
operation: "health check".into(),
detail: e.to_string(),
})?;
conn.query_drop("DO 0")
.await
.map_err(|e| WaypointError::ConnectionLost {
operation: "health check".into(),
detail: e.to_string(),
})?;
Ok(())
}
}
}
pub async fn acquire_lock(&self, schema: &str, table_name: &str) -> Result<()> {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(c) => acquire_advisory_lock(c, schema, table_name).await,
#[cfg(feature = "mysql")]
DbClient::Mysql(pool) => {
use mysql_async::prelude::*;
let key = mysql_lock_key(schema, table_name);
let mut conn = pool.get_conn().await?;
let acquired: Option<i64> = conn
.exec_first("SELECT GET_LOCK(?, -1)", (key.clone(),))
.await?;
match acquired {
Some(1) => {
park_lock_conn(pool, &key, conn);
Ok(())
}
_ => Err(WaypointError::LockError(format!(
"Failed to acquire MySQL named lock {}",
key
))),
}
}
}
}
pub async fn acquire_lock_with_timeout(
&self,
schema: &str,
table_name: &str,
timeout_secs: u32,
) -> Result<()> {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(c) => {
acquire_advisory_lock_with_timeout(c, schema, table_name, timeout_secs).await
}
#[cfg(feature = "mysql")]
DbClient::Mysql(pool) => {
use mysql_async::prelude::*;
let key = mysql_lock_key(schema, table_name);
let mut conn = pool.get_conn().await?;
let acquired: Option<i64> = conn
.exec_first("SELECT GET_LOCK(?, ?)", (key.clone(), timeout_secs as i64))
.await?;
match acquired {
Some(1) => {
park_lock_conn(pool, &key, conn);
Ok(())
}
Some(0) => Err(WaypointError::LockError(format!(
"Timed out waiting for MySQL named lock {} after {}s",
key, timeout_secs
))),
_ => Err(WaypointError::LockError(format!(
"Failed to acquire MySQL named lock {} (NULL result)",
key
))),
}
}
}
}
pub async fn release_lock(&self, schema: &str, table_name: &str) -> Result<()> {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(c) => release_advisory_lock(c, schema, table_name).await,
#[cfg(feature = "mysql")]
DbClient::Mysql(pool) => {
use mysql_async::prelude::*;
let key = mysql_lock_key(schema, table_name);
let mut conn = match unpark_lock_conn(pool, &key) {
Some(conn) => conn,
None => {
return Err(WaypointError::LockError(format!(
"No pinned connection holds MySQL named lock {} — \
release_lock called without a matching acquire_lock",
key
)));
}
};
let released = conn
.exec_first::<Option<i64>, _, _>("SELECT RELEASE_LOCK(?)", (key.clone(),))
.await;
drop(conn);
match released {
Ok(Some(Some(1))) => Ok(()),
Ok(_) => {
log::warn!(
"RELEASE_LOCK({}) did not report success; the lock is released \
regardless because the holding session was returned to the pool",
key
);
Ok(())
}
Err(e) => Err(WaypointError::MysqlError(e)),
}
}
}
}
pub async fn current_user(&self) -> Result<String> {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(c) => get_current_user(c).await,
#[cfg(feature = "mysql")]
DbClient::Mysql(pool) => {
use mysql_async::prelude::*;
let mut conn = pool.get_conn().await?;
let user: Option<String> = conn.query_first("SELECT CURRENT_USER()").await?;
user.ok_or_else(|| {
WaypointError::ConfigError("CURRENT_USER() returned no rows".into())
})
}
}
}
pub async fn current_database(&self) -> Result<String> {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(c) => get_current_database(c).await,
#[cfg(feature = "mysql")]
DbClient::Mysql(pool) => {
use mysql_async::prelude::*;
let mut conn = pool.get_conn().await?;
let db: Option<Option<String>> = conn.query_first("SELECT DATABASE()").await?;
match db.flatten() {
Some(name) => Ok(name),
None => Err(WaypointError::ConfigError(
"MySQL connection has no current database (none selected in URL)".into(),
)),
}
}
}
}
pub async fn resolve_schema(&self, configured: &str) -> Result<String> {
match self.dialect_kind() {
DialectKind::Postgres => Ok(configured.to_string()),
DialectKind::Mysql => {
if configured == "public" {
self.current_database().await
} else {
Ok(configured.to_string())
}
}
}
}
pub async fn execute_raw(&self, sql: &str) -> Result<i32> {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(c) => execute_raw(c, sql).await,
#[cfg(feature = "mysql")]
DbClient::Mysql(pool) => {
use mysql_async::prelude::*;
let start = std::time::Instant::now();
let mut conn = pool.get_conn().await?;
for stmt in crate::sql_parser::split_mysql_statements(sql) {
conn.query_drop(&stmt).await?;
}
Ok(start.elapsed().as_millis() as i32)
}
}
}
pub async fn execute_in_transaction(&self, sql: &str) -> Result<i32> {
match self {
#[cfg(feature = "postgres")]
DbClient::Postgres(c) => execute_in_transaction(c, sql).await,
#[cfg(feature = "mysql")]
DbClient::Mysql(_) => self.execute_raw(sql).await,
}
}
}
pub async fn connect_for_url(
conn_string: &str,
#[cfg_attr(
not(any(feature = "postgres", feature = "mysql")),
allow(unused_variables)
)]
config: &crate::config::WaypointConfig,
) -> Result<DbClient> {
let kind = DialectKind::from_url(conn_string).unwrap_or(config.database.engine);
match kind {
#[cfg(feature = "postgres")]
DialectKind::Postgres => {
let transport = TransportConfig::from_database_config(&config.database);
let client = connect_with_transport(conn_string, &transport).await?;
Ok(DbClient::with_postgres(client))
}
#[cfg(not(feature = "postgres"))]
DialectKind::Postgres => Err(WaypointError::ConfigError(
"PostgreSQL support is not compiled in (enable the `postgres` feature)".into(),
)),
#[cfg(feature = "mysql")]
DialectKind::Mysql => {
let pool = connect_mysql_pool(
conn_string,
config.database.ssl_mode,
config.database.ssl_root_cert.as_deref(),
config.database.statement_timeout_secs,
config.database.keepalive_secs,
)
.await?;
Ok(DbClient::with_mysql(pool))
}
#[cfg(not(feature = "mysql"))]
DialectKind::Mysql => Err(WaypointError::ConfigError(
"MySQL support is not compiled in (enable the `mysql` feature)".into(),
)),
}
}
#[cfg(feature = "mysql")]
async fn connect_mysql_pool(
conn_string: &str,
ssl_mode: SslMode,
ssl_root_cert: Option<&std::path::Path>,
statement_timeout_secs: u32,
keepalive_secs: u32,
) -> Result<mysql_async::Pool> {
let base = mysql_async::Opts::from_url(conn_string)
.map_err(|e| WaypointError::ConfigError(format!("Invalid MySQL connection URL: {}", e)))?;
let mut builder = mysql_async::OptsBuilder::from_opts(base);
if statement_timeout_secs > 0 {
let millis = u64::from(statement_timeout_secs).saturating_mul(1000);
log::debug!(
"Setting MySQL MAX_EXECUTION_TIME={}ms (bounds SELECTs only; DDL is not interruptible \
by it)",
millis
);
builder = builder.setup(vec![format!("SET SESSION MAX_EXECUTION_TIME = {}", millis)]);
}
if keepalive_secs > 0 {
builder = builder.tcp_keepalive(Some(std::time::Duration::from_secs(u64::from(
keepalive_secs,
))));
}
let base = mysql_async::Opts::from(builder);
if ssl_mode.requires_tls() && base.socket().is_some() {
return Err(WaypointError::ConfigError(format!(
"ssl_mode = '{}' requires TLS, but this MySQL connection uses a Unix \
socket, which the driver cannot secure. Use a TCP host:port, or set \
ssl_mode = 'disable'.",
ssl_mode
)));
}
if ssl_mode == SslMode::Prefer && base.ssl_opts().is_some() {
log::debug!(
"Using the TLS options from the MySQL connection URL (ssl_mode is at its default)."
);
return Ok(mysql_async::Pool::new(base));
}
let Some(ssl_opts) = crate::tls::make_mysql_ssl_opts(ssl_mode, ssl_root_cert) else {
return Ok(mysql_async::Pool::new(base));
};
let secure = mysql_async::Pool::new(
mysql_async::OptsBuilder::from_opts(base.clone()).ssl_opts(Some(ssl_opts)),
);
if ssl_mode != SslMode::Prefer {
return Ok(secure);
}
match secure.get_conn().await {
Ok(conn) => {
drop(conn);
Ok(secure)
}
Err(e) if mysql_tls_unavailable(&e) => {
log::warn!(
"MySQL server does not support TLS ({}); continuing with an UNENCRYPTED \
connection because ssl_mode is 'prefer'. Set ssl_mode to 'require' or \
higher to refuse this.",
e
);
let _ = secure.disconnect().await;
Ok(mysql_async::Pool::new(base))
}
Err(e) => Err(WaypointError::MysqlError(e)),
}
}
#[cfg(feature = "mysql")]
fn mysql_tls_unavailable(e: &mysql_async::Error) -> bool {
matches!(
e,
mysql_async::Error::Driver(mysql_async::DriverError::NoClientSslFlagFromServer)
) || matches!(e, mysql_async::Error::Io(mysql_async::IoError::Tls(_)))
}
#[cfg(feature = "mysql")]
fn mysql_lock_key(schema: &str, table_name: &str) -> String {
let full = format!("waypoint_{}_{}", schema, table_name);
if full.len() <= 64 {
full
} else {
format!("waypoint_{:08x}", crc32fast::hash(full.as_bytes()))
}
}
#[cfg(feature = "mysql")]
type MysqlLockRegistry = std::collections::HashMap<(usize, String), mysql_async::Conn>;
#[cfg(feature = "mysql")]
static MYSQL_LOCK_CONNS: std::sync::LazyLock<std::sync::Mutex<MysqlLockRegistry>> =
std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
#[cfg(feature = "mysql")]
fn mysql_pool_ident(pool: &mysql_async::Pool) -> usize {
pool as *const mysql_async::Pool as usize
}
#[cfg(feature = "mysql")]
fn park_lock_conn(pool: &mysql_async::Pool, key: &str, conn: mysql_async::Conn) {
let registry_key = (mysql_pool_ident(pool), key.to_string());
match MYSQL_LOCK_CONNS.lock() {
Ok(mut guard) => {
guard.insert(registry_key, conn);
}
Err(poisoned) => {
poisoned.into_inner().insert(registry_key, conn);
}
}
}
#[cfg(feature = "mysql")]
fn unpark_lock_conn(pool: &mysql_async::Pool, key: &str) -> Option<mysql_async::Conn> {
let registry_key = (mysql_pool_ident(pool), key.to_string());
match MYSQL_LOCK_CONNS.lock() {
Ok(mut guard) => guard.remove(®istry_key),
Err(poisoned) => poisoned.into_inner().remove(®istry_key),
}
}
#[cfg(feature = "postgres")]
fn to_pg_ssl_mode(mode: SslMode) -> tokio_postgres::config::SslMode {
match mode {
SslMode::Disable => tokio_postgres::config::SslMode::Disable,
SslMode::Prefer => tokio_postgres::config::SslMode::Prefer,
SslMode::Require | SslMode::VerifyCa | SslMode::VerifyFull => {
tokio_postgres::config::SslMode::Require
}
}
}
#[cfg(feature = "postgres")]
fn is_permanent_error(e: &tokio_postgres::Error) -> bool {
if let Some(db_err) = e.as_db_error() {
let code = db_err.code().code();
return code == "28P01" || code == "28000";
}
false
}
pub fn inject_keepalive(conn_string: &str, keepalive_secs: u32) -> String {
if keepalive_secs == 0 {
return conn_string.to_string();
}
let lower = conn_string.to_lowercase();
if lower.contains("keepalives") {
return conn_string.to_string();
}
let params = format!("keepalives=1&keepalives_idle={}", keepalive_secs);
if conn_string.starts_with("postgres://") || conn_string.starts_with("postgresql://") {
if conn_string.contains('?') {
format!("{}&{}", conn_string, params)
} else {
format!("{}?{}", conn_string, params)
}
} else {
format!(
"{} keepalives=1 keepalives_idle={}",
conn_string, keepalive_secs
)
}
}
#[cfg(feature = "postgres")]
fn spawn_connection_task<F>(connection: F)
where
F: std::future::Future<Output = std::result::Result<(), tokio_postgres::Error>>
+ Send
+ 'static,
{
tokio::spawn(async move {
if let Err(e) = connection.await {
log::error!("Database connection error: {}", e);
}
});
}
#[cfg(feature = "postgres")]
async fn connect_once(
pg_config: &tokio_postgres::Config,
tls_config: Option<&rustls::ClientConfig>,
connect_timeout_secs: u32,
) -> std::result::Result<Client, tokio_postgres::Error> {
let connect_fut = async {
match tls_config {
None => {
let (client, connection) = pg_config.connect(tokio_postgres::NoTls).await?;
spawn_connection_task(connection);
Ok(client)
}
Some(tls_config) => {
let tls = tokio_postgres_rustls::MakeRustlsConnect::new(tls_config.clone());
let (client, connection) = pg_config.connect(tls).await?;
spawn_connection_task(connection);
Ok(client)
}
}
};
if connect_timeout_secs > 0 {
match tokio::time::timeout(
std::time::Duration::from_secs(connect_timeout_secs as u64),
connect_fut,
)
.await
{
Ok(result) => result,
Err(_) => Err(tokio_postgres::Error::__private_api_timeout()),
}
} else {
connect_fut.await
}
}
#[cfg(feature = "postgres")]
#[deprecated(
since = "0.7.0",
note = "Use connect_with_transport, which supports the full sslmode ladder and a custom CA. Will be removed in 1.0."
)]
pub async fn connect(conn_string: &str) -> Result<Client> {
connect_with_transport(conn_string, &TransportConfig::default()).await
}
#[cfg(feature = "postgres")]
#[deprecated(
since = "0.7.0",
note = "Use connect_with_transport, which supports the full sslmode ladder and a custom CA. Will be removed in 1.0."
)]
pub async fn connect_with_config(
conn_string: &str,
ssl_mode: &SslMode,
retries: u32,
connect_timeout_secs: u32,
statement_timeout_secs: u32,
) -> Result<Client> {
connect_with_transport(
conn_string,
&TransportConfig {
ssl_mode: *ssl_mode,
retries,
connect_timeout_secs,
statement_timeout_secs,
..TransportConfig::default()
},
)
.await
}
#[cfg(feature = "postgres")]
#[deprecated(
since = "0.7.0",
note = "Use connect_with_transport — this signature cannot express ssl_root_cert. Will be removed in 1.0."
)]
pub async fn connect_with_full_config(
conn_string: &str,
ssl_mode: &SslMode,
retries: u32,
connect_timeout_secs: u32,
statement_timeout_secs: u32,
keepalive_secs: u32,
) -> Result<Client> {
connect_with_transport(
conn_string,
&TransportConfig {
ssl_mode: *ssl_mode,
ssl_root_cert: None,
retries,
connect_timeout_secs,
statement_timeout_secs,
keepalive_secs,
},
)
.await
}
#[cfg(feature = "postgres")]
pub async fn connect_with_transport(
conn_string: &str,
transport: &TransportConfig,
) -> Result<Client> {
let conn_string = inject_keepalive(conn_string, transport.keepalive_secs);
let (conn_string, embedded) = crate::tls::parse_url_sslmode(&conn_string);
let ssl_mode = crate::tls::reconcile_ssl_mode(transport.ssl_mode, embedded.mode);
let ssl_root_cert =
crate::tls::reconcile_root_cert(transport.ssl_root_cert.as_deref(), embedded.root_cert);
let mut pg_config: tokio_postgres::Config = conn_string.parse().map_err(|e| {
WaypointError::ConfigError(format!("Invalid PostgreSQL connection string: {}", e))
})?;
pg_config.ssl_mode(to_pg_ssl_mode(ssl_mode));
let tls_config = match ssl_mode {
SslMode::Disable => None,
_ => Some(crate::tls::make_rustls_config(
ssl_mode,
ssl_root_cert.as_deref(),
)?),
};
let retries = transport.retries;
let mut last_err = None;
for attempt in 0..=retries {
if attempt > 0 {
let base_delay = std::cmp::min(1u64 << attempt, 30);
let jitter_ms = fastrand::u64(0..1000);
let delay = std::time::Duration::from_secs(base_delay)
+ std::time::Duration::from_millis(jitter_ms);
log::info!(
"Connection attempt failed, retrying; attempt={}, max_attempts={}, delay_ms={}",
attempt + 1,
retries + 1,
delay.as_millis() as u64
);
tokio::time::sleep(delay).await;
}
match connect_once(
&pg_config,
tls_config.as_ref(),
transport.connect_timeout_secs,
)
.await
{
Ok(client) => {
if attempt > 0 {
log::info!(
"Connected successfully after retry; attempt={}, max_attempts={}",
attempt + 1,
retries + 1
);
}
if transport.statement_timeout_secs > 0 {
let timeout_sql = format!(
"SET statement_timeout = '{}s'",
transport.statement_timeout_secs
);
client.batch_execute(&timeout_sql).await?;
}
return Ok(client);
}
Err(e) => {
if is_permanent_error(&e) {
log::error!("Permanent connection error, not retrying: {}", e);
return Err(WaypointError::DatabaseError(e));
}
last_err = Some(e);
}
}
}
Err(WaypointError::DatabaseError(last_err.unwrap()))
}
#[cfg(feature = "postgres")]
pub async fn acquire_advisory_lock(client: &Client, schema: &str, table_name: &str) -> Result<()> {
let lock_id = advisory_lock_id(schema, table_name);
log::info!(
"Acquiring advisory lock; lock_id={}, table={}",
lock_id,
table_name
);
client
.execute("SELECT pg_advisory_lock($1)", &[&lock_id])
.await
.map_err(|e| WaypointError::LockError(format!("Failed to acquire advisory lock: {}", e)))?;
Ok(())
}
#[cfg(feature = "postgres")]
pub async fn acquire_advisory_lock_with_timeout(
client: &Client,
schema: &str,
table_name: &str,
timeout_secs: u32,
) -> Result<()> {
let lock_id = advisory_lock_id(schema, table_name);
log::info!(
"Trying to acquire advisory lock with timeout; lock_id={}, table={}, timeout_secs={}",
lock_id,
table_name,
timeout_secs
);
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(timeout_secs as u64);
loop {
let row = client
.query_one("SELECT pg_try_advisory_lock($1)", &[&lock_id])
.await
.map_err(|e| WaypointError::LockError(format!("Failed to try advisory lock: {}", e)))?;
let acquired: bool = row.get(0);
if acquired {
return Ok(());
}
if std::time::Instant::now() >= deadline {
return Err(WaypointError::LockError(format!(
"Timed out waiting for advisory lock after {}s (table: {}). Another migration may be running.",
timeout_secs, table_name
)));
}
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
}
}
#[cfg(feature = "postgres")]
pub async fn release_advisory_lock(client: &Client, schema: &str, table_name: &str) -> Result<()> {
let lock_id = advisory_lock_id(schema, table_name);
log::info!(
"Releasing advisory lock; lock_id={}, table={}",
lock_id,
table_name
);
client
.execute("SELECT pg_advisory_unlock($1)", &[&lock_id])
.await
.map_err(|e| WaypointError::LockError(format!("Failed to release advisory lock: {}", e)))?;
Ok(())
}
pub fn advisory_lock_id(schema: &str, table_name: &str) -> i64 {
let key = format!("{}\0{}", schema, table_name);
crc32fast::hash(key.as_bytes()) as i64
}
#[cfg(feature = "postgres")]
pub async fn get_current_user(client: &Client) -> Result<String> {
let row = client.query_one("SELECT current_user", &[]).await?;
Ok(row.get::<_, String>(0))
}
#[cfg(feature = "postgres")]
pub async fn get_current_database(client: &Client) -> Result<String> {
let row = client.query_one("SELECT current_database()", &[]).await?;
Ok(row.get::<_, String>(0))
}
#[cfg(feature = "postgres")]
pub async fn execute_in_transaction(client: &Client, sql: &str) -> Result<i32> {
let start = std::time::Instant::now();
client.batch_execute("BEGIN").await?;
match client.batch_execute(sql).await {
Ok(()) => {
client.batch_execute("COMMIT").await?;
}
Err(e) => {
if let Err(rollback_err) = client.batch_execute("ROLLBACK").await {
log::warn!("Failed to rollback transaction: {}", rollback_err);
}
return Err(WaypointError::DatabaseError(e));
}
}
let elapsed = start.elapsed().as_millis() as i32;
Ok(elapsed)
}
#[cfg(feature = "postgres")]
pub async fn execute_raw(client: &Client, sql: &str) -> Result<i32> {
let start = std::time::Instant::now();
client.batch_execute(sql).await?;
let elapsed = start.elapsed().as_millis() as i32;
Ok(elapsed)
}
pub fn is_transient_error(e: &WaypointError) -> bool {
match e {
#[cfg(feature = "postgres")]
WaypointError::DatabaseError(pg_err) => {
if pg_err.is_closed() {
return true;
}
if let Some(db_err) = pg_err.as_db_error() {
let code = db_err.code().code();
return matches!(
code,
"57P01" | "57P02" | "57P03" | "08000" | "08003" | "08006"
);
}
let msg = pg_err.to_string().to_lowercase();
msg.contains("connection reset")
|| msg.contains("broken pipe")
|| msg.contains("connection closed")
|| msg.contains("unexpected eof")
}
#[cfg(feature = "mysql")]
WaypointError::MysqlError(my_err) => {
let msg = my_err.to_string().to_lowercase();
msg.contains("connection reset")
|| msg.contains("broken pipe")
|| msg.contains("connection closed")
|| msg.contains("server has gone away")
|| msg.contains("lost connection")
|| msg.contains("io error")
}
WaypointError::ConnectionLost { .. } => true,
_ => false,
}
}
#[cfg(feature = "postgres")]
pub async fn check_connection(client: &Client) -> Result<()> {
client
.simple_query("")
.await
.map_err(|e| WaypointError::ConnectionLost {
operation: "health check".to_string(),
detail: e.to_string(),
})?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_inject_keepalive_url_style() {
let result = inject_keepalive("postgres://user:pass@localhost/db", 120);
assert_eq!(
result,
"postgres://user:pass@localhost/db?keepalives=1&keepalives_idle=120"
);
}
#[test]
fn test_inject_keepalive_url_with_existing_params() {
let result = inject_keepalive("postgres://user:pass@localhost/db?sslmode=require", 60);
assert_eq!(
result,
"postgres://user:pass@localhost/db?sslmode=require&keepalives=1&keepalives_idle=60"
);
}
#[test]
fn test_inject_keepalive_kv_style() {
let result = inject_keepalive("host=localhost port=5432 user=admin dbname=mydb", 90);
assert_eq!(
result,
"host=localhost port=5432 user=admin dbname=mydb keepalives=1 keepalives_idle=90"
);
}
#[test]
fn test_inject_keepalive_zero_disables() {
let result = inject_keepalive("postgres://user:pass@localhost/db", 0);
assert_eq!(result, "postgres://user:pass@localhost/db");
}
#[test]
fn test_inject_keepalive_already_present() {
let result = inject_keepalive("postgres://user:pass@localhost/db?keepalives=1", 120);
assert_eq!(result, "postgres://user:pass@localhost/db?keepalives=1");
}
#[test]
fn test_transient_error_connection_lost() {
let err = WaypointError::ConnectionLost {
operation: "test".to_string(),
detail: "gone".to_string(),
};
assert!(is_transient_error(&err));
}
#[test]
fn test_transient_error_config_is_not_transient() {
let err = WaypointError::ConfigError("bad config".to_string());
assert!(!is_transient_error(&err));
}
#[test]
fn test_transient_error_migration_failed_is_not_transient() {
let err = WaypointError::MigrationFailed {
script: "V1__test.sql".to_string(),
reason: "syntax error".to_string(),
};
assert!(!is_transient_error(&err));
}
#[test]
fn test_advisory_lock_id_stability() {
let id1 = advisory_lock_id("public", "waypoint_schema_history");
let id2 = advisory_lock_id("public", "waypoint_schema_history");
assert_eq!(id1, id2);
let id3 = advisory_lock_id("public", "other_table");
assert_ne!(id1, id3);
}
#[test]
fn test_advisory_lock_id_is_scoped_per_schema() {
let a = advisory_lock_id("tenant_a", "waypoint_schema_history");
let b = advisory_lock_id("tenant_b", "waypoint_schema_history");
assert_ne!(a, b, "schemas in one database must not share a lock");
}
#[test]
fn test_advisory_lock_id_separator_cannot_be_forged() {
assert_ne!(
advisory_lock_id("a", "b_c"),
advisory_lock_id("a_b", "c"),
"schema/table boundary must be unambiguous"
);
}
#[test]
fn test_transient_error_lock_error_is_not_transient() {
let err = WaypointError::LockError("lock failed".to_string());
assert!(!is_transient_error(&err));
}
#[test]
fn test_transient_error_io_error_is_not_transient() {
let err = WaypointError::IoError(std::io::Error::new(
std::io::ErrorKind::NotFound,
"file not found",
));
assert!(!is_transient_error(&err));
}
#[test]
fn test_validate_identifier_valid() {
assert!(validate_identifier("users").is_ok());
assert!(validate_identifier("my_table").is_ok());
assert!(validate_identifier("Table123").is_ok());
assert!(validate_identifier("a").is_ok());
}
#[test]
fn test_validate_identifier_invalid() {
assert!(validate_identifier("").is_err());
assert!(validate_identifier("my-table").is_err());
assert!(validate_identifier("my table").is_err());
assert!(validate_identifier("table.name").is_err());
assert!(validate_identifier("table;drop").is_err());
}
#[test]
fn test_quote_ident_simple() {
assert_eq!(quote_ident("users"), "\"users\"");
}
#[test]
fn test_quote_ident_embedded_quotes() {
assert_eq!(quote_ident("my\"table"), "\"my\"\"table\"");
}
#[test]
fn test_quote_ident_empty() {
assert_eq!(quote_ident(""), "\"\"");
}
#[test]
fn test_inject_keepalive_postgresql_prefix() {
let result = inject_keepalive("postgresql://user:pass@localhost/db", 120);
assert_eq!(
result,
"postgresql://user:pass@localhost/db?keepalives=1&keepalives_idle=120"
);
}
#[cfg(feature = "mysql")]
#[test]
fn mysql_lock_key_is_scoped_per_database() {
let a = mysql_lock_key("app_prod", "waypoint_schema_history");
let b = mysql_lock_key("app_staging", "waypoint_schema_history");
assert_ne!(a, b);
assert_eq!(a, "waypoint_app_prod_waypoint_schema_history");
}
#[cfg(feature = "mysql")]
#[test]
fn mysql_lock_key_respects_the_64_char_limit() {
let long_db = "d".repeat(60);
let long_tbl = "t".repeat(60);
let k = mysql_lock_key(&long_db, &long_tbl);
assert!(
k.len() <= 64,
"GET_LOCK names are capped at 64: {}",
k.len()
);
}
#[cfg(feature = "mysql")]
#[test]
fn mysql_lock_key_does_not_collide_after_shortening() {
let prefix = "x".repeat(60);
let a = mysql_lock_key(&prefix, "alpha");
let b = mysql_lock_key(&prefix, "beta");
assert!(a.len() <= 64 && b.len() <= 64);
assert_ne!(a, b, "distinct tables collapsed onto one lock key");
}
#[cfg(feature = "mysql")]
#[test]
fn mysql_lock_key_is_stable() {
assert_eq!(mysql_lock_key("db", "tbl"), mysql_lock_key("db", "tbl"));
}
#[test]
fn test_sandbox_name_is_unique_across_rapid_calls() {
let names: std::collections::HashSet<String> =
(0..2000).map(|_| sandbox_name("waypoint_sim")).collect();
assert_eq!(
names.len(),
2000,
"sandbox names collided within a single tight loop"
);
}
#[test]
fn test_sandbox_name_fits_identifier_limits() {
for prefix in ["waypoint_sim", "waypoint_drift_check"] {
let name = sandbox_name(prefix);
assert!(
name.len() <= 63,
"{} is {} bytes, over PostgreSQL's 63-byte limit",
name,
name.len()
);
assert!(name.starts_with(prefix));
}
}
#[test]
fn test_quote_literal_escapes_embedded_single_quotes() {
assert_eq!(quote_literal("fine"), "'fine'");
assert_eq!(quote_literal("it's bad"), "'it''s bad'");
assert_eq!(quote_literal("''"), r"''''''");
assert_eq!(quote_literal(""), r"''");
}
#[test]
fn test_quote_literal_leaves_other_characters_alone() {
assert_eq!(quote_literal(r"back\slash"), r"'back\slash'");
assert_eq!(quote_literal("multi\nline"), "'multi\nline'");
}
}