1use crate::config::SslMode;
9use crate::dialect::{DatabaseDialect, DialectKind};
10use crate::error::{Result, WaypointError};
11use std::path::PathBuf;
12
13#[cfg(feature = "postgres")]
14use fastrand;
15
16#[cfg(feature = "postgres")]
17use tokio_postgres::Client;
18
19#[derive(Debug, Clone)]
26pub struct TransportConfig {
27 pub ssl_mode: SslMode,
29 pub ssl_root_cert: Option<PathBuf>,
31 pub retries: u32,
33 pub connect_timeout_secs: u32,
35 pub statement_timeout_secs: u32,
37 pub keepalive_secs: u32,
39}
40
41impl Default for TransportConfig {
42 fn default() -> Self {
43 Self {
45 ssl_mode: SslMode::Prefer,
46 ssl_root_cert: None,
47 retries: 0,
48 connect_timeout_secs: 30,
49 statement_timeout_secs: 0,
50 keepalive_secs: 120,
51 }
52 }
53}
54
55impl TransportConfig {
56 pub fn from_database_config(db: &crate::config::DatabaseConfig) -> Self {
58 Self {
59 ssl_mode: db.ssl_mode,
60 ssl_root_cert: db.ssl_root_cert.clone(),
61 retries: db.connect_retries,
62 connect_timeout_secs: db.connect_timeout_secs,
63 statement_timeout_secs: db.statement_timeout_secs,
64 keepalive_secs: db.keepalive_secs,
65 }
66 }
67}
68
69pub fn quote_ident(name: &str) -> String {
75 format!("\"{}\"", name.replace('"', "\"\""))
76}
77
78pub fn quote_ident_mysql(name: &str) -> String {
87 format!("`{}`", name.replace('`', "``"))
88}
89
90pub fn validate_identifier(name: &str) -> Result<()> {
95 if name.is_empty() {
96 return Err(WaypointError::ConfigError(
97 "Identifier cannot be empty".to_string(),
98 ));
99 }
100 if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
101 return Err(WaypointError::ConfigError(format!(
102 "Identifier '{}' contains invalid characters. Only [a-zA-Z0-9_] are allowed.",
103 name
104 )));
105 }
106 Ok(())
107}
108
109pub enum DbClient {
119 #[cfg(feature = "postgres")]
121 Postgres(Client),
122 #[cfg(feature = "mysql")]
127 Mysql(mysql_async::Pool),
128}
129
130impl DbClient {
131 #[cfg(feature = "postgres")]
133 pub fn with_postgres(client: Client) -> Self {
134 DbClient::Postgres(client)
135 }
136
137 #[cfg(feature = "mysql")]
139 pub fn with_mysql(pool: mysql_async::Pool) -> Self {
140 DbClient::Mysql(pool)
141 }
142
143 pub fn dialect_kind(&self) -> DialectKind {
145 match self {
146 #[cfg(feature = "postgres")]
147 DbClient::Postgres(_) => DialectKind::Postgres,
148 #[cfg(feature = "mysql")]
149 DbClient::Mysql(_) => DialectKind::Mysql,
150 }
151 }
152
153 pub fn dialect(&self) -> &'static dyn DatabaseDialect {
158 #[cfg(feature = "postgres")]
159 static PG: crate::dialect::postgres::PostgresDialect =
160 crate::dialect::postgres::PostgresDialect;
161 #[cfg(feature = "mysql")]
162 static MY: crate::dialect::mysql::MysqlDialect = crate::dialect::mysql::MysqlDialect;
163 match self.dialect_kind() {
164 #[cfg(feature = "postgres")]
165 DialectKind::Postgres => &PG,
166 #[cfg(not(feature = "postgres"))]
167 DialectKind::Postgres => {
168 panic!("PostgreSQL connection without `postgres` feature compiled in")
169 }
170 #[cfg(feature = "mysql")]
171 DialectKind::Mysql => &MY,
172 #[cfg(not(feature = "mysql"))]
173 DialectKind::Mysql => {
174 panic!("MySQL connection without `mysql` feature compiled in")
175 }
176 }
177 }
178
179 #[cfg(feature = "postgres")]
183 pub fn as_postgres(&self) -> Result<&Client> {
184 match self {
185 DbClient::Postgres(c) => Ok(c),
186 #[cfg(feature = "mysql")]
187 DbClient::Mysql(_) => Err(WaypointError::ConfigError(
188 "This operation is not yet implemented for MySQL".into(),
189 )),
190 }
191 }
192
193 #[cfg(feature = "mysql")]
196 pub fn as_mysql(&self) -> Result<&mysql_async::Pool> {
197 match self {
198 DbClient::Mysql(p) => Ok(p),
199 #[cfg(feature = "postgres")]
200 DbClient::Postgres(_) => Err(WaypointError::ConfigError(
201 "This operation requires a MySQL connection".into(),
202 )),
203 }
204 }
205
206 pub async fn check_connection(&self) -> Result<()> {
208 match self {
209 #[cfg(feature = "postgres")]
210 DbClient::Postgres(c) => check_connection(c).await,
211 #[cfg(feature = "mysql")]
212 DbClient::Mysql(pool) => {
213 use mysql_async::prelude::*;
214 let mut conn =
215 pool.get_conn()
216 .await
217 .map_err(|e| WaypointError::ConnectionLost {
218 operation: "health check".into(),
219 detail: e.to_string(),
220 })?;
221 conn.query_drop("DO 0")
222 .await
223 .map_err(|e| WaypointError::ConnectionLost {
224 operation: "health check".into(),
225 detail: e.to_string(),
226 })?;
227 Ok(())
228 }
229 }
230 }
231
232 pub async fn acquire_lock(&self, table_name: &str) -> Result<()> {
237 match self {
238 #[cfg(feature = "postgres")]
239 DbClient::Postgres(c) => acquire_advisory_lock(c, table_name).await,
240 #[cfg(feature = "mysql")]
241 DbClient::Mysql(pool) => {
242 use mysql_async::prelude::*;
243 let key = mysql_lock_key(&mysql_lock_scope(self).await, table_name);
244 let mut conn = pool.get_conn().await?;
245 let acquired: Option<i64> = conn
246 .exec_first("SELECT GET_LOCK(?, -1)", (key.clone(),))
247 .await?;
248 match acquired {
249 Some(1) => {
250 park_lock_conn(pool, &key, conn);
251 Ok(())
252 }
253 _ => Err(WaypointError::LockError(format!(
254 "Failed to acquire MySQL named lock {}",
255 key
256 ))),
257 }
258 }
259 }
260 }
261
262 pub async fn acquire_lock_with_timeout(
264 &self,
265 table_name: &str,
266 timeout_secs: u32,
267 ) -> Result<()> {
268 match self {
269 #[cfg(feature = "postgres")]
270 DbClient::Postgres(c) => {
271 acquire_advisory_lock_with_timeout(c, table_name, timeout_secs).await
272 }
273 #[cfg(feature = "mysql")]
274 DbClient::Mysql(pool) => {
275 use mysql_async::prelude::*;
276 let key = mysql_lock_key(&mysql_lock_scope(self).await, table_name);
277 let mut conn = pool.get_conn().await?;
278 let acquired: Option<i64> = conn
279 .exec_first("SELECT GET_LOCK(?, ?)", (key.clone(), timeout_secs as i64))
280 .await?;
281 match acquired {
282 Some(1) => {
283 park_lock_conn(pool, &key, conn);
284 Ok(())
285 }
286 Some(0) => Err(WaypointError::LockError(format!(
287 "Timed out waiting for MySQL named lock {} after {}s",
288 key, timeout_secs
289 ))),
290 _ => Err(WaypointError::LockError(format!(
291 "Failed to acquire MySQL named lock {} (NULL result)",
292 key
293 ))),
294 }
295 }
296 }
297 }
298
299 pub async fn release_lock(&self, table_name: &str) -> Result<()> {
301 match self {
302 #[cfg(feature = "postgres")]
303 DbClient::Postgres(c) => release_advisory_lock(c, table_name).await,
304 #[cfg(feature = "mysql")]
305 DbClient::Mysql(pool) => {
306 use mysql_async::prelude::*;
307 let key = mysql_lock_key(&mysql_lock_scope(self).await, table_name);
308 let mut conn = match unpark_lock_conn(pool, &key) {
312 Some(conn) => conn,
313 None => {
314 return Err(WaypointError::LockError(format!(
315 "No pinned connection holds MySQL named lock {} — \
316 release_lock called without a matching acquire_lock",
317 key
318 )));
319 }
320 };
321 let released = conn
322 .exec_first::<Option<i64>, _, _>("SELECT RELEASE_LOCK(?)", (key.clone(),))
323 .await;
324 drop(conn);
328 match released {
329 Ok(Some(Some(1))) => Ok(()),
330 Ok(_) => {
331 log::warn!(
332 "RELEASE_LOCK({}) did not report success; the lock is released \
333 regardless because the holding session was returned to the pool",
334 key
335 );
336 Ok(())
337 }
338 Err(e) => Err(WaypointError::MysqlError(e)),
339 }
340 }
341 }
342 }
343
344 pub async fn current_user(&self) -> Result<String> {
346 match self {
347 #[cfg(feature = "postgres")]
348 DbClient::Postgres(c) => get_current_user(c).await,
349 #[cfg(feature = "mysql")]
350 DbClient::Mysql(pool) => {
351 use mysql_async::prelude::*;
352 let mut conn = pool.get_conn().await?;
353 let user: Option<String> = conn.query_first("SELECT CURRENT_USER()").await?;
354 user.ok_or_else(|| {
355 WaypointError::ConfigError("CURRENT_USER() returned no rows".into())
356 })
357 }
358 }
359 }
360
361 pub async fn current_database(&self) -> Result<String> {
363 match self {
364 #[cfg(feature = "postgres")]
365 DbClient::Postgres(c) => get_current_database(c).await,
366 #[cfg(feature = "mysql")]
367 DbClient::Mysql(pool) => {
368 use mysql_async::prelude::*;
369 let mut conn = pool.get_conn().await?;
370 let db: Option<Option<String>> = conn.query_first("SELECT DATABASE()").await?;
372 match db.flatten() {
373 Some(name) => Ok(name),
374 None => Err(WaypointError::ConfigError(
375 "MySQL connection has no current database (none selected in URL)".into(),
376 )),
377 }
378 }
379 }
380 }
381
382 pub async fn resolve_schema(&self, configured: &str) -> Result<String> {
389 match self.dialect_kind() {
390 DialectKind::Postgres => Ok(configured.to_string()),
391 DialectKind::Mysql => {
392 if configured == "public" {
393 self.current_database().await
394 } else {
395 Ok(configured.to_string())
396 }
397 }
398 }
399 }
400
401 pub async fn execute_raw(&self, sql: &str) -> Result<i32> {
410 match self {
411 #[cfg(feature = "postgres")]
412 DbClient::Postgres(c) => execute_raw(c, sql).await,
413 #[cfg(feature = "mysql")]
414 DbClient::Mysql(pool) => {
415 use mysql_async::prelude::*;
416 let start = std::time::Instant::now();
417 let mut conn = pool.get_conn().await?;
418 for stmt in crate::sql_parser::split_mysql_statements(sql) {
419 conn.query_drop(&stmt).await?;
420 }
421 Ok(start.elapsed().as_millis() as i32)
422 }
423 }
424 }
425
426 pub async fn execute_in_transaction(&self, sql: &str) -> Result<i32> {
435 match self {
436 #[cfg(feature = "postgres")]
437 DbClient::Postgres(c) => execute_in_transaction(c, sql).await,
438 #[cfg(feature = "mysql")]
439 DbClient::Mysql(_) => self.execute_raw(sql).await,
440 }
441 }
442}
443
444pub async fn connect_for_url(
454 conn_string: &str,
455 #[cfg_attr(
456 not(any(feature = "postgres", feature = "mysql")),
457 allow(unused_variables)
458 )]
459 config: &crate::config::WaypointConfig,
460) -> Result<DbClient> {
461 let kind = DialectKind::from_url(conn_string).unwrap_or(config.database.engine);
462 match kind {
463 #[cfg(feature = "postgres")]
464 DialectKind::Postgres => {
465 let transport = TransportConfig::from_database_config(&config.database);
466 let client = connect_with_transport(conn_string, &transport).await?;
467 Ok(DbClient::with_postgres(client))
468 }
469 #[cfg(not(feature = "postgres"))]
470 DialectKind::Postgres => Err(WaypointError::ConfigError(
471 "PostgreSQL support is not compiled in (enable the `postgres` feature)".into(),
472 )),
473 #[cfg(feature = "mysql")]
474 DialectKind::Mysql => {
475 let pool = connect_mysql_pool(
476 conn_string,
477 config.database.ssl_mode,
478 config.database.ssl_root_cert.as_deref(),
479 )
480 .await?;
481 Ok(DbClient::with_mysql(pool))
482 }
483 #[cfg(not(feature = "mysql"))]
484 DialectKind::Mysql => Err(WaypointError::ConfigError(
485 "MySQL support is not compiled in (enable the `mysql` feature)".into(),
486 )),
487 }
488}
489
490#[cfg(feature = "mysql")]
505async fn connect_mysql_pool(
506 conn_string: &str,
507 ssl_mode: SslMode,
508 ssl_root_cert: Option<&std::path::Path>,
509) -> Result<mysql_async::Pool> {
510 let base = mysql_async::Opts::from_url(conn_string)
511 .map_err(|e| WaypointError::ConfigError(format!("Invalid MySQL connection URL: {}", e)))?;
512
513 if ssl_mode.requires_tls() && base.socket().is_some() {
518 return Err(WaypointError::ConfigError(format!(
519 "ssl_mode = '{}' requires TLS, but this MySQL connection uses a Unix \
520 socket, which the driver cannot secure. Use a TCP host:port, or set \
521 ssl_mode = 'disable'.",
522 ssl_mode
523 )));
524 }
525
526 if ssl_mode == SslMode::Prefer && base.ssl_opts().is_some() {
530 log::debug!(
531 "Using the TLS options from the MySQL connection URL (ssl_mode is at its default)."
532 );
533 return Ok(mysql_async::Pool::new(base));
534 }
535
536 let Some(ssl_opts) = crate::tls::make_mysql_ssl_opts(ssl_mode, ssl_root_cert) else {
537 return Ok(mysql_async::Pool::new(base));
539 };
540
541 let secure = mysql_async::Pool::new(
542 mysql_async::OptsBuilder::from_opts(base.clone()).ssl_opts(Some(ssl_opts)),
543 );
544
545 if ssl_mode != SslMode::Prefer {
546 return Ok(secure);
547 }
548
549 match secure.get_conn().await {
550 Ok(conn) => {
551 drop(conn);
552 Ok(secure)
553 }
554 Err(e) if mysql_tls_unavailable(&e) => {
559 log::warn!(
560 "MySQL server does not support TLS ({}); continuing with an UNENCRYPTED \
561 connection because ssl_mode is 'prefer'. Set ssl_mode to 'require' or \
562 higher to refuse this.",
563 e
564 );
565 let _ = secure.disconnect().await;
566 Ok(mysql_async::Pool::new(base))
567 }
568 Err(e) => Err(WaypointError::MysqlError(e)),
569 }
570}
571
572#[cfg(feature = "mysql")]
579fn mysql_tls_unavailable(e: &mysql_async::Error) -> bool {
580 matches!(
581 e,
582 mysql_async::Error::Driver(mysql_async::DriverError::NoClientSslFlagFromServer)
583 ) || matches!(e, mysql_async::Error::Io(mysql_async::IoError::Tls(_)))
584}
585
586#[cfg(feature = "mysql")]
606fn mysql_lock_key(schema: &str, table_name: &str) -> String {
607 let full = format!("waypoint_{}_{}", schema, table_name);
608 if full.len() <= 64 {
609 full
610 } else {
611 format!("waypoint_{:08x}", crc32fast::hash(full.as_bytes()))
612 }
613}
614
615#[cfg(feature = "mysql")]
621async fn mysql_lock_scope(client: &DbClient) -> String {
622 client
623 .current_database()
624 .await
625 .unwrap_or_else(|_| "_nodb".to_string())
626}
627
628#[cfg(feature = "mysql")]
647type MysqlLockRegistry = std::collections::HashMap<(usize, String), mysql_async::Conn>;
648
649#[cfg(feature = "mysql")]
650static MYSQL_LOCK_CONNS: std::sync::LazyLock<std::sync::Mutex<MysqlLockRegistry>> =
651 std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
652
653#[cfg(feature = "mysql")]
667fn mysql_pool_ident(pool: &mysql_async::Pool) -> usize {
668 pool as *const mysql_async::Pool as usize
669}
670
671#[cfg(feature = "mysql")]
673fn park_lock_conn(pool: &mysql_async::Pool, key: &str, conn: mysql_async::Conn) {
674 let registry_key = (mysql_pool_ident(pool), key.to_string());
675 match MYSQL_LOCK_CONNS.lock() {
676 Ok(mut guard) => {
677 guard.insert(registry_key, conn);
678 }
679 Err(poisoned) => {
680 poisoned.into_inner().insert(registry_key, conn);
684 }
685 }
686}
687
688#[cfg(feature = "mysql")]
690fn unpark_lock_conn(pool: &mysql_async::Pool, key: &str) -> Option<mysql_async::Conn> {
691 let registry_key = (mysql_pool_ident(pool), key.to_string());
692 match MYSQL_LOCK_CONNS.lock() {
693 Ok(mut guard) => guard.remove(®istry_key),
694 Err(poisoned) => poisoned.into_inner().remove(®istry_key),
695 }
696}
697
698#[cfg(feature = "postgres")]
712fn to_pg_ssl_mode(mode: SslMode) -> tokio_postgres::config::SslMode {
713 match mode {
714 SslMode::Disable => tokio_postgres::config::SslMode::Disable,
715 SslMode::Prefer => tokio_postgres::config::SslMode::Prefer,
716 SslMode::Require | SslMode::VerifyCa | SslMode::VerifyFull => {
717 tokio_postgres::config::SslMode::Require
718 }
719 }
720}
721
722#[cfg(feature = "postgres")]
724fn is_permanent_error(e: &tokio_postgres::Error) -> bool {
725 if let Some(db_err) = e.as_db_error() {
726 let code = db_err.code().code();
727 return code == "28P01" || code == "28000";
729 }
730 false
731}
732
733pub fn inject_keepalive(conn_string: &str, keepalive_secs: u32) -> String {
739 if keepalive_secs == 0 {
740 return conn_string.to_string();
741 }
742 let lower = conn_string.to_lowercase();
743 if lower.contains("keepalives") {
744 return conn_string.to_string();
745 }
746 let params = format!("keepalives=1&keepalives_idle={}", keepalive_secs);
747 if conn_string.starts_with("postgres://") || conn_string.starts_with("postgresql://") {
748 if conn_string.contains('?') {
749 format!("{}&{}", conn_string, params)
750 } else {
751 format!("{}?{}", conn_string, params)
752 }
753 } else {
754 format!(
756 "{} keepalives=1 keepalives_idle={}",
757 conn_string, keepalive_secs
758 )
759 }
760}
761
762#[cfg(feature = "postgres")]
768fn spawn_connection_task<F>(connection: F)
769where
770 F: std::future::Future<Output = std::result::Result<(), tokio_postgres::Error>>
771 + Send
772 + 'static,
773{
774 tokio::spawn(async move {
775 if let Err(e) = connection.await {
776 log::error!("Database connection error: {}", e);
777 }
778 });
779}
780
781#[cfg(feature = "postgres")]
797async fn connect_once(
798 pg_config: &tokio_postgres::Config,
799 tls_config: Option<&rustls::ClientConfig>,
800 connect_timeout_secs: u32,
801) -> std::result::Result<Client, tokio_postgres::Error> {
802 let connect_fut = async {
803 match tls_config {
804 None => {
805 let (client, connection) = pg_config.connect(tokio_postgres::NoTls).await?;
806 spawn_connection_task(connection);
807 Ok(client)
808 }
809 Some(tls_config) => {
810 let tls = tokio_postgres_rustls::MakeRustlsConnect::new(tls_config.clone());
811 let (client, connection) = pg_config.connect(tls).await?;
812 spawn_connection_task(connection);
813 Ok(client)
814 }
815 }
816 };
817
818 if connect_timeout_secs > 0 {
819 match tokio::time::timeout(
820 std::time::Duration::from_secs(connect_timeout_secs as u64),
821 connect_fut,
822 )
823 .await
824 {
825 Ok(result) => result,
826 Err(_) => Err(tokio_postgres::Error::__private_api_timeout()),
827 }
828 } else {
829 connect_fut.await
830 }
831}
832
833#[cfg(feature = "postgres")]
837#[deprecated(
838 since = "0.7.0",
839 note = "Use connect_with_transport, which supports the full sslmode ladder and a custom CA. Will be removed in 1.0."
840)]
841pub async fn connect(conn_string: &str) -> Result<Client> {
842 connect_with_transport(conn_string, &TransportConfig::default()).await
843}
844
845#[cfg(feature = "postgres")]
850#[deprecated(
851 since = "0.7.0",
852 note = "Use connect_with_transport, which supports the full sslmode ladder and a custom CA. Will be removed in 1.0."
853)]
854pub async fn connect_with_config(
855 conn_string: &str,
856 ssl_mode: &SslMode,
857 retries: u32,
858 connect_timeout_secs: u32,
859 statement_timeout_secs: u32,
860) -> Result<Client> {
861 connect_with_transport(
862 conn_string,
863 &TransportConfig {
864 ssl_mode: *ssl_mode,
865 retries,
866 connect_timeout_secs,
867 statement_timeout_secs,
868 ..TransportConfig::default()
869 },
870 )
871 .await
872}
873
874#[cfg(feature = "postgres")]
876#[deprecated(
877 since = "0.7.0",
878 note = "Use connect_with_transport — this signature cannot express ssl_root_cert. Will be removed in 1.0."
879)]
880pub async fn connect_with_full_config(
881 conn_string: &str,
882 ssl_mode: &SslMode,
883 retries: u32,
884 connect_timeout_secs: u32,
885 statement_timeout_secs: u32,
886 keepalive_secs: u32,
887) -> Result<Client> {
888 connect_with_transport(
889 conn_string,
890 &TransportConfig {
891 ssl_mode: *ssl_mode,
892 ssl_root_cert: None,
893 retries,
894 connect_timeout_secs,
895 statement_timeout_secs,
896 keepalive_secs,
897 },
898 )
899 .await
900}
901
902#[cfg(feature = "postgres")]
909pub async fn connect_with_transport(
910 conn_string: &str,
911 transport: &TransportConfig,
912) -> Result<Client> {
913 let conn_string = inject_keepalive(conn_string, transport.keepalive_secs);
914
915 let (conn_string, embedded) = crate::tls::parse_url_sslmode(&conn_string);
919 let ssl_mode = crate::tls::reconcile_ssl_mode(transport.ssl_mode, embedded.mode);
920 let ssl_root_cert =
921 crate::tls::reconcile_root_cert(transport.ssl_root_cert.as_deref(), embedded.root_cert);
922
923 let mut pg_config: tokio_postgres::Config = conn_string.parse().map_err(|e| {
924 WaypointError::ConfigError(format!("Invalid PostgreSQL connection string: {}", e))
925 })?;
926 pg_config.ssl_mode(to_pg_ssl_mode(ssl_mode));
927
928 let tls_config = match ssl_mode {
931 SslMode::Disable => None,
932 _ => Some(crate::tls::make_rustls_config(
933 ssl_mode,
934 ssl_root_cert.as_deref(),
935 )?),
936 };
937
938 let retries = transport.retries;
939 let mut last_err = None;
940
941 for attempt in 0..=retries {
942 if attempt > 0 {
943 let base_delay = std::cmp::min(1u64 << attempt, 30);
944 let jitter_ms = fastrand::u64(0..1000);
945 let delay = std::time::Duration::from_secs(base_delay)
946 + std::time::Duration::from_millis(jitter_ms);
947 log::info!(
948 "Connection attempt failed, retrying; attempt={}, max_attempts={}, delay_ms={}",
949 attempt + 1,
950 retries + 1,
951 delay.as_millis() as u64
952 );
953 tokio::time::sleep(delay).await;
954 }
955
956 match connect_once(
957 &pg_config,
958 tls_config.as_ref(),
959 transport.connect_timeout_secs,
960 )
961 .await
962 {
963 Ok(client) => {
964 if attempt > 0 {
965 log::info!(
966 "Connected successfully after retry; attempt={}, max_attempts={}",
967 attempt + 1,
968 retries + 1
969 );
970 }
971
972 if transport.statement_timeout_secs > 0 {
974 let timeout_sql = format!(
975 "SET statement_timeout = '{}s'",
976 transport.statement_timeout_secs
977 );
978 client.batch_execute(&timeout_sql).await?;
979 }
980
981 return Ok(client);
982 }
983 Err(e) => {
984 if is_permanent_error(&e) {
986 log::error!("Permanent connection error, not retrying: {}", e);
987 return Err(WaypointError::DatabaseError(e));
988 }
989 last_err = Some(e);
990 }
991 }
992 }
993
994 Err(WaypointError::DatabaseError(last_err.unwrap()))
995}
996
997#[cfg(feature = "postgres")]
1001pub async fn acquire_advisory_lock(client: &Client, table_name: &str) -> Result<()> {
1002 let lock_id = advisory_lock_id(table_name);
1003 log::info!(
1004 "Acquiring advisory lock; lock_id={}, table={}",
1005 lock_id,
1006 table_name
1007 );
1008
1009 client
1010 .execute("SELECT pg_advisory_lock($1)", &[&lock_id])
1011 .await
1012 .map_err(|e| WaypointError::LockError(format!("Failed to acquire advisory lock: {}", e)))?;
1013
1014 Ok(())
1015}
1016
1017#[cfg(feature = "postgres")]
1022pub async fn acquire_advisory_lock_with_timeout(
1023 client: &Client,
1024 table_name: &str,
1025 timeout_secs: u32,
1026) -> Result<()> {
1027 let lock_id = advisory_lock_id(table_name);
1028 log::info!(
1029 "Trying to acquire advisory lock with timeout; lock_id={}, table={}, timeout_secs={}",
1030 lock_id,
1031 table_name,
1032 timeout_secs
1033 );
1034
1035 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(timeout_secs as u64);
1036
1037 loop {
1038 let row = client
1039 .query_one("SELECT pg_try_advisory_lock($1)", &[&lock_id])
1040 .await
1041 .map_err(|e| WaypointError::LockError(format!("Failed to try advisory lock: {}", e)))?;
1042
1043 let acquired: bool = row.get(0);
1044 if acquired {
1045 return Ok(());
1046 }
1047
1048 if std::time::Instant::now() >= deadline {
1049 return Err(WaypointError::LockError(format!(
1050 "Timed out waiting for advisory lock after {}s (table: {}). Another migration may be running.",
1051 timeout_secs, table_name
1052 )));
1053 }
1054
1055 tokio::time::sleep(std::time::Duration::from_millis(500)).await;
1057 }
1058}
1059
1060#[cfg(feature = "postgres")]
1062pub async fn release_advisory_lock(client: &Client, table_name: &str) -> Result<()> {
1063 let lock_id = advisory_lock_id(table_name);
1064 log::info!(
1065 "Releasing advisory lock; lock_id={}, table={}",
1066 lock_id,
1067 table_name
1068 );
1069
1070 client
1071 .execute("SELECT pg_advisory_unlock($1)", &[&lock_id])
1072 .await
1073 .map_err(|e| WaypointError::LockError(format!("Failed to release advisory lock: {}", e)))?;
1074
1075 Ok(())
1076}
1077
1078pub fn advisory_lock_id(table_name: &str) -> i64 {
1084 crc32fast::hash(table_name.as_bytes()) as i64
1085}
1086
1087#[cfg(feature = "postgres")]
1089pub async fn get_current_user(client: &Client) -> Result<String> {
1090 let row = client.query_one("SELECT current_user", &[]).await?;
1091 Ok(row.get::<_, String>(0))
1092}
1093
1094#[cfg(feature = "postgres")]
1096pub async fn get_current_database(client: &Client) -> Result<String> {
1097 let row = client.query_one("SELECT current_database()", &[]).await?;
1098 Ok(row.get::<_, String>(0))
1099}
1100
1101#[cfg(feature = "postgres")]
1104pub async fn execute_in_transaction(client: &Client, sql: &str) -> Result<i32> {
1105 let start = std::time::Instant::now();
1106
1107 client.batch_execute("BEGIN").await?;
1108
1109 match client.batch_execute(sql).await {
1110 Ok(()) => {
1111 client.batch_execute("COMMIT").await?;
1112 }
1113 Err(e) => {
1114 if let Err(rollback_err) = client.batch_execute("ROLLBACK").await {
1115 log::warn!("Failed to rollback transaction: {}", rollback_err);
1116 }
1117 return Err(WaypointError::DatabaseError(e));
1118 }
1119 }
1120
1121 let elapsed = start.elapsed().as_millis() as i32;
1122 Ok(elapsed)
1123}
1124
1125#[cfg(feature = "postgres")]
1127pub async fn execute_raw(client: &Client, sql: &str) -> Result<i32> {
1128 let start = std::time::Instant::now();
1129 client.batch_execute(sql).await?;
1130 let elapsed = start.elapsed().as_millis() as i32;
1131 Ok(elapsed)
1132}
1133
1134pub fn is_transient_error(e: &WaypointError) -> bool {
1139 match e {
1140 #[cfg(feature = "postgres")]
1141 WaypointError::DatabaseError(pg_err) => {
1142 if pg_err.is_closed() {
1144 return true;
1145 }
1146 if let Some(db_err) = pg_err.as_db_error() {
1148 let code = db_err.code().code();
1149 return matches!(
1153 code,
1154 "57P01" | "57P02" | "57P03" | "08000" | "08003" | "08006"
1155 );
1156 }
1157 let msg = pg_err.to_string().to_lowercase();
1159 msg.contains("connection reset")
1160 || msg.contains("broken pipe")
1161 || msg.contains("connection closed")
1162 || msg.contains("unexpected eof")
1163 }
1164 #[cfg(feature = "mysql")]
1165 WaypointError::MysqlError(my_err) => {
1166 let msg = my_err.to_string().to_lowercase();
1170 msg.contains("connection reset")
1171 || msg.contains("broken pipe")
1172 || msg.contains("connection closed")
1173 || msg.contains("server has gone away")
1174 || msg.contains("lost connection")
1175 || msg.contains("io error")
1176 }
1177 WaypointError::ConnectionLost { .. } => true,
1178 _ => false,
1179 }
1180}
1181
1182#[cfg(feature = "postgres")]
1184pub async fn check_connection(client: &Client) -> Result<()> {
1185 client
1186 .simple_query("")
1187 .await
1188 .map_err(|e| WaypointError::ConnectionLost {
1189 operation: "health check".to_string(),
1190 detail: e.to_string(),
1191 })?;
1192 Ok(())
1193}
1194
1195#[cfg(test)]
1196mod tests {
1197 use super::*;
1198
1199 #[test]
1202 fn test_inject_keepalive_url_style() {
1203 let result = inject_keepalive("postgres://user:pass@localhost/db", 120);
1204 assert_eq!(
1205 result,
1206 "postgres://user:pass@localhost/db?keepalives=1&keepalives_idle=120"
1207 );
1208 }
1209
1210 #[test]
1211 fn test_inject_keepalive_url_with_existing_params() {
1212 let result = inject_keepalive("postgres://user:pass@localhost/db?sslmode=require", 60);
1213 assert_eq!(
1214 result,
1215 "postgres://user:pass@localhost/db?sslmode=require&keepalives=1&keepalives_idle=60"
1216 );
1217 }
1218
1219 #[test]
1220 fn test_inject_keepalive_kv_style() {
1221 let result = inject_keepalive("host=localhost port=5432 user=admin dbname=mydb", 90);
1222 assert_eq!(
1223 result,
1224 "host=localhost port=5432 user=admin dbname=mydb keepalives=1 keepalives_idle=90"
1225 );
1226 }
1227
1228 #[test]
1229 fn test_inject_keepalive_zero_disables() {
1230 let result = inject_keepalive("postgres://user:pass@localhost/db", 0);
1231 assert_eq!(result, "postgres://user:pass@localhost/db");
1232 }
1233
1234 #[test]
1235 fn test_inject_keepalive_already_present() {
1236 let result = inject_keepalive("postgres://user:pass@localhost/db?keepalives=1", 120);
1237 assert_eq!(result, "postgres://user:pass@localhost/db?keepalives=1");
1238 }
1239
1240 #[test]
1243 fn test_transient_error_connection_lost() {
1244 let err = WaypointError::ConnectionLost {
1245 operation: "test".to_string(),
1246 detail: "gone".to_string(),
1247 };
1248 assert!(is_transient_error(&err));
1249 }
1250
1251 #[test]
1252 fn test_transient_error_config_is_not_transient() {
1253 let err = WaypointError::ConfigError("bad config".to_string());
1254 assert!(!is_transient_error(&err));
1255 }
1256
1257 #[test]
1258 fn test_transient_error_migration_failed_is_not_transient() {
1259 let err = WaypointError::MigrationFailed {
1260 script: "V1__test.sql".to_string(),
1261 reason: "syntax error".to_string(),
1262 };
1263 assert!(!is_transient_error(&err));
1264 }
1265
1266 #[test]
1267 fn test_advisory_lock_id_stability() {
1268 let id1 = advisory_lock_id("waypoint_schema_history");
1270 let id2 = advisory_lock_id("waypoint_schema_history");
1271 assert_eq!(id1, id2);
1272 let id3 = advisory_lock_id("other_table");
1274 assert_ne!(id1, id3);
1275 }
1276
1277 #[test]
1278 fn test_transient_error_lock_error_is_not_transient() {
1279 let err = WaypointError::LockError("lock failed".to_string());
1280 assert!(!is_transient_error(&err));
1281 }
1282
1283 #[test]
1284 fn test_transient_error_io_error_is_not_transient() {
1285 let err = WaypointError::IoError(std::io::Error::new(
1286 std::io::ErrorKind::NotFound,
1287 "file not found",
1288 ));
1289 assert!(!is_transient_error(&err));
1290 }
1291
1292 #[test]
1293 fn test_validate_identifier_valid() {
1294 assert!(validate_identifier("users").is_ok());
1295 assert!(validate_identifier("my_table").is_ok());
1296 assert!(validate_identifier("Table123").is_ok());
1297 assert!(validate_identifier("a").is_ok());
1298 }
1299
1300 #[test]
1301 fn test_validate_identifier_invalid() {
1302 assert!(validate_identifier("").is_err());
1303 assert!(validate_identifier("my-table").is_err());
1304 assert!(validate_identifier("my table").is_err());
1305 assert!(validate_identifier("table.name").is_err());
1306 assert!(validate_identifier("table;drop").is_err());
1307 }
1308
1309 #[test]
1310 fn test_quote_ident_simple() {
1311 assert_eq!(quote_ident("users"), "\"users\"");
1312 }
1313
1314 #[test]
1315 fn test_quote_ident_embedded_quotes() {
1316 assert_eq!(quote_ident("my\"table"), "\"my\"\"table\"");
1317 }
1318
1319 #[test]
1320 fn test_quote_ident_empty() {
1321 assert_eq!(quote_ident(""), "\"\"");
1322 }
1323
1324 #[test]
1325 fn test_inject_keepalive_postgresql_prefix() {
1326 let result = inject_keepalive("postgresql://user:pass@localhost/db", 120);
1327 assert_eq!(
1328 result,
1329 "postgresql://user:pass@localhost/db?keepalives=1&keepalives_idle=120"
1330 );
1331 }
1332
1333 #[cfg(feature = "mysql")]
1334 #[test]
1335 fn mysql_lock_key_is_scoped_per_database() {
1336 let a = mysql_lock_key("app_prod", "waypoint_schema_history");
1340 let b = mysql_lock_key("app_staging", "waypoint_schema_history");
1341 assert_ne!(a, b);
1342 assert_eq!(a, "waypoint_app_prod_waypoint_schema_history");
1343 }
1344
1345 #[cfg(feature = "mysql")]
1346 #[test]
1347 fn mysql_lock_key_respects_the_64_char_limit() {
1348 let long_db = "d".repeat(60);
1349 let long_tbl = "t".repeat(60);
1350 let k = mysql_lock_key(&long_db, &long_tbl);
1351 assert!(
1352 k.len() <= 64,
1353 "GET_LOCK names are capped at 64: {}",
1354 k.len()
1355 );
1356 }
1357
1358 #[cfg(feature = "mysql")]
1359 #[test]
1360 fn mysql_lock_key_does_not_collide_after_shortening() {
1361 let prefix = "x".repeat(60);
1364 let a = mysql_lock_key(&prefix, "alpha");
1365 let b = mysql_lock_key(&prefix, "beta");
1366 assert!(a.len() <= 64 && b.len() <= 64);
1367 assert_ne!(a, b, "distinct tables collapsed onto one lock key");
1368 }
1369
1370 #[cfg(feature = "mysql")]
1371 #[test]
1372 fn mysql_lock_key_is_stable() {
1373 assert_eq!(mysql_lock_key("db", "tbl"), mysql_lock_key("db", "tbl"));
1376 }
1377}