From abb4560636bedcae6572742c1a2810b0bef8caf2 Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Sat, 8 Aug 2026 04:13:25 +0100 Subject: [PATCH 01/14] feat: add unix socket to stream --- pgdog/src/net/stream.rs | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/pgdog/src/net/stream.rs b/pgdog/src/net/stream.rs index 10ee0d3e9..1d6ee6d64 100644 --- a/pgdog/src/net/stream.rs +++ b/pgdog/src/net/stream.rs @@ -3,18 +3,38 @@ use bytes::{BufMut, BytesMut}; use futures::FutureExt; use pin_project::pin_project; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream, ReadBuf}; -use tokio::net::TcpStream; +use tokio::io::{self, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream, ReadBuf}; +use tokio::net::{TcpStream, UnixStream}; use tracing::trace; use std::io::{Error, ErrorKind}; use std::net::SocketAddr; use std::ops::Deref; +use std::os::fd::AsRawFd; use std::pin::Pin; use std::task::Context; use super::messages::{ErrorResponse, Message, Protocol, ReadyForQuery}; +fn unix_peek(stream: &UnixStream, buffer: &mut [u8]) -> Option> { + let n = unsafe { + libc::recv( + stream.as_raw_fd(), + buffer.as_mut_ptr().cast(), + buffer.len(), + libc::MSG_PEEK, + ) + }; + + if n >= 0 { + return Some(Ok(n as usize)); + } else if io::Error::last_os_error().kind() == io::ErrorKind::WouldBlock { + return None; + } else { + return Some(Err(io::Error::last_os_error())); + } +} + /// Inner stream types. #[pin_project(project = StreamInnerProjection)] #[derive(Debug)] @@ -23,6 +43,7 @@ enum StreamInner { Plain(#[pin] BufStream), Tls(#[pin] BufStream>), DevNull, + UnixSockets(#[pin] BufStream), } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -54,6 +75,7 @@ impl AsyncRead for Stream { match project.inner.project() { StreamInnerProjection::Plain(stream) => stream.poll_read(cx, buf), StreamInnerProjection::Tls(stream) => stream.poll_read(cx, buf), + StreamInnerProjection::UnixSockets(stream) => stream.poll_read(cx, buf), StreamInnerProjection::DevNull => std::task::Poll::Ready(Ok(())), } } @@ -70,6 +92,7 @@ impl AsyncWrite for Stream { StreamInnerProjection::Plain(stream) => stream.poll_write(cx, buf), StreamInnerProjection::Tls(stream) => stream.poll_write(cx, buf), StreamInnerProjection::DevNull => std::task::Poll::Ready(Ok(buf.len())), + StreamInnerProjection::UnixSockets(stream) => stream.poll_write(cx, buf), } } @@ -82,6 +105,7 @@ impl AsyncWrite for Stream { StreamInnerProjection::Plain(stream) => stream.poll_flush(cx), StreamInnerProjection::Tls(stream) => stream.poll_flush(cx), StreamInnerProjection::DevNull => std::task::Poll::Ready(Ok(())), + StreamInnerProjection::UnixSockets(stream) => stream.poll_flush(cx), } } @@ -93,6 +117,7 @@ impl AsyncWrite for Stream { match project.inner.project() { StreamInnerProjection::Plain(stream) => stream.poll_shutdown(cx), StreamInnerProjection::Tls(stream) => stream.poll_shutdown(cx), + StreamInnerProjection::UnixSockets(stream) => stream.poll_shutdown(cx), StreamInnerProjection::DevNull => std::task::Poll::Ready(Ok(())), } } @@ -165,6 +190,7 @@ impl Stream { StreamInner::Plain(stream) => stream.get_ref().peer_addr().ok().into(), StreamInner::Tls(stream) => stream.get_ref().get_ref().0.peer_addr().ok().into(), StreamInner::DevNull => PeerAddr { addr: None }, + StreamInner::UnixSockets(stream) => PeerAddr { addr: None }, } } @@ -175,6 +201,10 @@ impl Stream { StreamInner::Plain(plain) => eof(plain.get_mut().peek(&mut buf).await)?, StreamInner::Tls(tls) => eof(tls.get_mut().get_mut().0.peek(&mut buf).await)?, StreamInner::DevNull => 0, + StreamInner::UnixSockets(stream) => match unix_peek(stream.get_ref(), &mut buf) { + None => 0, + Some(res) => eof(res)?, + }, }; Ok(()) @@ -186,6 +216,7 @@ impl Stream { StreamInner::Plain(plain) => plain.get_mut().peek(&mut buf).now_or_never(), StreamInner::Tls(tls) => tls.get_mut().get_mut().0.peek(&mut buf).now_or_never(), StreamInner::DevNull => return Liveness::Clean, + StreamInner::UnixSockets(stream) => unix_peek(stream.get_ref(), &mut buf), }; match peeked { @@ -216,6 +247,7 @@ impl Stream { StreamInner::Plain(stream) => eof(stream.write_all(&bytes).await)?, StreamInner::Tls(stream) => eof(stream.write_all(&bytes).await)?, StreamInner::DevNull => (), + StreamInner::UnixSockets(stream) => eof(stream.write_all(&bytes).await)?, } #[cfg(debug_assertions)] From 82fbe24881704cad8aa5805f969a3dff1e7685ef Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 01:14:57 +0100 Subject: [PATCH 02/14] feat: create transport enum for address --- pgdog/src/admin/show_bans.rs | 2 +- pgdog/src/admin/show_pools.rs | 2 +- pgdog/src/admin/show_replication.rs | 2 +- pgdog/src/admin/show_server_memory.rs | 2 +- pgdog/src/admin/show_servers.rs | 2 +- pgdog/src/admin/show_stats.rs | 2 +- pgdog/src/backend/auth/rds_iam.rs | 6 ++- pgdog/src/backend/pool/address.rs | 67 +++++++++++++++++++++--- pgdog/src/backend/pool/token_cache.rs | 2 +- pgdog/src/backend/schema/sync/pg_dump.rs | 2 +- pgdog/src/backend/server.rs | 7 ++- pgdog/src/net/stream.rs | 61 +++++++++++++-------- pgdog/src/stats/pools.rs | 2 +- 13 files changed, 116 insertions(+), 43 deletions(-) diff --git a/pgdog/src/admin/show_bans.rs b/pgdog/src/admin/show_bans.rs index 91eb75272..963a8f72a 100644 --- a/pgdog/src/admin/show_bans.rs +++ b/pgdog/src/admin/show_bans.rs @@ -55,7 +55,7 @@ impl Command for ShowBans { row.add(pool.id() as i64) .add(user.database.as_str()) .add(user.user.as_str()) - .add(pool.addr().host.as_str()) + .add(pool.addr().host.to_string()) .add(pool.addr().port as i64) .add(shard_num as i64) .add(role.to_string()) diff --git a/pgdog/src/admin/show_pools.rs b/pgdog/src/admin/show_pools.rs index 5366dc5c4..9a68993c5 100644 --- a/pgdog/src/admin/show_pools.rs +++ b/pgdog/src/admin/show_pools.rs @@ -58,7 +58,7 @@ impl Command for ShowPools { row.add(pool.id() as i64) .add(user.database.as_str()) .add(user.user.as_str()) - .add(pool.addr().host.as_str()) + .add(pool.addr().host.to_string()) .add(pool.addr().port as i64) .add(shard_num as i64) .add(role.to_string()) diff --git a/pgdog/src/admin/show_replication.rs b/pgdog/src/admin/show_replication.rs index 3a8ca15f2..68190e6cc 100644 --- a/pgdog/src/admin/show_replication.rs +++ b/pgdog/src/admin/show_replication.rs @@ -53,7 +53,7 @@ impl Command for ShowReplication { row.add(pool.id() as i64) .add(user.database.as_str()) .add(user.user.as_str()) - .add(pool.addr().host.as_str()) + .add(pool.addr().host.to_string()) .add(pool.addr().port as i64) .add(shard_num as i64) .add(role.to_string()) diff --git a/pgdog/src/admin/show_server_memory.rs b/pgdog/src/admin/show_server_memory.rs index 88f81ae50..bcb3915fb 100644 --- a/pgdog/src/admin/show_server_memory.rs +++ b/pgdog/src/admin/show_server_memory.rs @@ -43,7 +43,7 @@ impl Command for ShowServerMemory { row.add(server.stats.pool_id as i64) .add(server.addr.database_name.as_str()) .add(server.addr.user.as_str()) - .add(server.addr.host.as_str()) + .add(server.addr.host.to_string()) .add(server.addr.port as i64) .add(server.stats.id) .add(memory.buffer.reallocs as i64) diff --git a/pgdog/src/admin/show_servers.rs b/pgdog/src/admin/show_servers.rs index 287db00c6..d6dab7fe2 100644 --- a/pgdog/src/admin/show_servers.rs +++ b/pgdog/src/admin/show_servers.rs @@ -90,7 +90,7 @@ impl Command for ShowServers { .add("pool_id", server.stats.pool_id) .add("database", server.addr.database_name) .add("user", server.addr.user) - .add("addr", server.addr.host.as_str()) + .add("addr", server.addr.host.to_string()) .add("port", server.addr.port.to_string()) .add("state", server.stats.state.to_string()) .add( diff --git a/pgdog/src/admin/show_stats.rs b/pgdog/src/admin/show_stats.rs index 740095edc..8011fe95c 100644 --- a/pgdog/src/admin/show_stats.rs +++ b/pgdog/src/admin/show_stats.rs @@ -78,7 +78,7 @@ impl Command for ShowStats { dr.add(user.database.as_str()) .add(user.user.as_str()) - .add(&pool.addr().host) + .add(&pool.addr().host.to_string()) .add(pool.addr().port as i64) .add(shard_num) .add(role.to_string()); diff --git a/pgdog/src/backend/auth/rds_iam.rs b/pgdog/src/backend/auth/rds_iam.rs index 8e3d3aa76..2ce28b657 100644 --- a/pgdog/src/backend/auth/rds_iam.rs +++ b/pgdog/src/backend/auth/rds_iam.rs @@ -37,7 +37,8 @@ fn resolve_region(addr: &Address) -> Result { return Ok(region.clone()); } - infer_region_from_rds_host(&addr.host).ok_or_else(|| { + let host = addr.host.tcp().expect("Host must be a TCP address"); + infer_region_from_rds_host(&host).ok_or_else(|| { Error::RdsIamToken(format!( "unable to infer AWS region from host \"{}\"; set \"server_iam_region\"", addr.host @@ -53,9 +54,10 @@ fn resolve_region(addr: &Address) -> Result { pub(crate) async fn token(addr: Address) -> Result<(String, SystemTime), Error> { let region = resolve_region(&addr)?; let sdk_config = aws_config::load_defaults(BehaviorVersion::latest()).await; + let host = addr.host.tcp().expect("host must be a TCP address"); let config = AuthTokenConfig::builder() - .hostname(addr.host.as_str()) + .hostname(host) .port(addr.port.into()) .username(addr.user.as_str()) .region(Region::new(region.clone())) diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 9443a9069..80501d93f 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -1,6 +1,10 @@ //! Server address. +use std::default::Default; +use std::fmt::Display; use std::net::{SocketAddr, ToSocketAddrs}; use std::ops::Deref; +use std::path::Path; +use std::path::PathBuf; use pgdog_config::Role; use pgdog_config::users::PasswordKind; @@ -14,11 +18,57 @@ use crate::backend::pool::dns_cache::DnsCache; use crate::backend::pool::token_cache::TokenCache; use crate::config::{Database, ServerAuth, User, config}; +/// Transport enum +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] +pub enum Transport { + TCP(String), + Unix(PathBuf), +} + +impl Transport { + pub fn new(value: &str) -> Self { + if value.starts_with('/') { + Transport::Unix(value.into()) + } else { + Transport::TCP(value.to_string()) + } + } + + pub fn tcp(&self) -> Option<&str> { + match self { + Transport::TCP(host) => Some(host), + Transport::Unix(_) => None, + } + } + + pub fn unix(&self) -> Option<&Path> { + match self { + Transport::TCP(_) => None, + Transport::Unix(path_buf) => Some(path_buf), + } + } +} + +impl Display for Transport { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Transport::TCP(addr) => write!(f, "{}", addr), + Transport::Unix(path_buf) => write!(f, "{}", path_buf.display()), + } + } +} + +impl Default for Transport { + fn default() -> Self { + Transport::TCP("127.0.0.1".to_string()) + } +} + /// Server address. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default, Eq, Hash)] pub struct Address { /// Server host. - pub host: String, + pub host: Transport, /// Server port. pub port: u16, /// PostgreSQL database name. @@ -48,7 +98,7 @@ pub struct Address { impl From
for pgdog_stats::Address { fn from(value: Address) -> Self { pgdog_stats::Address { - host: value.host, + host: value.host.to_string(), port: value.port, database_name: value.database_name, user: value.user, @@ -66,7 +116,7 @@ impl Address { let server_auth = user.server_auth; Address { - host: database.host.clone(), + host: Transport::new(&database.host), port: database.port, database_name: if let Some(database_name) = database.database_name.clone() { database_name @@ -168,9 +218,10 @@ impl Address { /// pub(crate) async fn addr(&self) -> Result { let dns_cache_override_enabled = config().config.general.dns_ttl().is_some(); + let host = self.host.tcp().expect("addr must be a TCP address"); if dns_cache_override_enabled { - let ip = DnsCache::global().resolve(&self.host).await?; + let ip = DnsCache::global().resolve(&host).await?; return Ok(SocketAddr::new(ip, self.port)); } @@ -179,7 +230,7 @@ impl Address { socket_addrs .next() - .ok_or(Error::DnsResolutionFailed(self.host.clone())) + .ok_or(Error::DnsResolutionFailed(host.to_string())) } /// A replacement for [`PartialEq`] which accounts for @@ -241,7 +292,7 @@ impl TryFrom for Address { // via the derived `Default`. The PROBE command (the only caller) // never reads `configured_role` anyway. Ok(Self { - host, + host: Transport::new(&host), port, passwords: vec![password.into()], user, @@ -280,7 +331,7 @@ mod test { let address = Address::new(&database, &user, 0); - assert_eq!(address.host, "127.0.0.1"); + assert_eq!(address.host.into(), "127.0.0.1"); assert_eq!(address.port, 6432); assert_eq!(address.database_name, "pgdog"); assert_eq!(address.user, "pgdog"); @@ -361,7 +412,7 @@ mod test { let addr = Address::try_from(Url::parse("postgres://user:password@127.0.0.1:6432/pgdb").unwrap()) .unwrap(); - assert_eq!(addr.host, "127.0.0.1"); + assert_eq!(addr.host.into(), "127.0.0.1"); assert_eq!(addr.port, 6432); assert_eq!(addr.database_name, "pgdb"); assert_eq!(addr.user, "user"); diff --git a/pgdog/src/backend/pool/token_cache.rs b/pgdog/src/backend/pool/token_cache.rs index f751c12f3..61905554d 100644 --- a/pgdog/src/backend/pool/token_cache.rs +++ b/pgdog/src/backend/pool/token_cache.rs @@ -76,7 +76,7 @@ impl From<&Address> for CacheKey { fn from(addr: &Address) -> Self { Self { user: addr.user.clone(), - host: addr.host.clone(), + host: addr.host.to_string(), port: addr.port, } } diff --git a/pgdog/src/backend/schema/sync/pg_dump.rs b/pgdog/src/backend/schema/sync/pg_dump.rs index f1a44a97e..5c2981080 100644 --- a/pgdog/src/backend/schema/sync/pg_dump.rs +++ b/pgdog/src/backend/schema/sync/pg_dump.rs @@ -186,7 +186,7 @@ fn build_pg_dump_command( command .arg("--schema-only") .arg("-h") - .arg(&addr.host) + .arg(&addr.host.to_string()) .arg("-p") .arg(addr.port.to_string()) .arg("-U") diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 96d9d0e7f..2b28e4cdb 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -278,7 +278,12 @@ impl Server { )?; let plain = stream.take()?; - let server_name = ServerName::try_from(addr.host.clone())?; + let host = addr + .host + .tcp() + .expect("host address must be a TCP address") + .to_owned(); + let server_name = ServerName::try_from(host)?; debug!("connecting with TLS to server name: {:?}", server_name); match connector.connect(server_name.clone(), plain).await { diff --git a/pgdog/src/net/stream.rs b/pgdog/src/net/stream.rs index 1d6ee6d64..3cc34e030 100644 --- a/pgdog/src/net/stream.rs +++ b/pgdog/src/net/stream.rs @@ -4,13 +4,14 @@ use bytes::{BufMut, BytesMut}; use futures::FutureExt; use pin_project::pin_project; use tokio::io::{self, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream, ReadBuf}; -use tokio::net::{TcpStream, UnixStream}; +use tokio::net::{TcpStream, UnixStream, unix}; use tracing::trace; use std::io::{Error, ErrorKind}; use std::net::SocketAddr; use std::ops::Deref; use std::os::fd::AsRawFd; +use std::path::PathBuf; use std::pin::Pin; use std::task::Context; @@ -27,11 +28,11 @@ fn unix_peek(stream: &UnixStream, buffer: &mut [u8]) -> Option= 0 { - return Some(Ok(n as usize)); + Some(Ok(n as usize)) } else if io::Error::last_os_error().kind() == io::ErrorKind::WouldBlock { - return None; + None } else { - return Some(Err(io::Error::last_os_error())); + Some(Err(io::Error::last_os_error())) } } @@ -183,14 +184,14 @@ impl Stream { matches!(self.inner, StreamInner::Tls(_)) } - /// Get peer address if any. We're not using UNIX sockets (yet) + /// Get peer address/unix socket if any. /// so the peer address should always be available. pub fn peer_addr(&self) -> PeerAddr { match &self.inner { - StreamInner::Plain(stream) => stream.get_ref().peer_addr().ok().into(), - StreamInner::Tls(stream) => stream.get_ref().get_ref().0.peer_addr().ok().into(), - StreamInner::DevNull => PeerAddr { addr: None }, - StreamInner::UnixSockets(stream) => PeerAddr { addr: None }, + StreamInner::Plain(stream) => stream.get_ref().peer_addr().into(), + StreamInner::Tls(stream) => stream.get_ref().get_ref().0.peer_addr().into(), + StreamInner::DevNull => PeerAddr::Empty, + StreamInner::UnixSockets(stream) => stream.get_ref().peer_addr().into(), } } @@ -397,30 +398,44 @@ pub fn eof(result: std::io::Result) -> Result { /// Wrapper around SocketAddr /// to make it easier to debug. -pub struct PeerAddr { - addr: Option, +pub enum PeerAddr { + /// TCP peer: Ip and Port + TCP(SocketAddr), + /// Unix Socket file path + Unix(PathBuf), + /// No Identifiable Peer + Empty, } -impl Deref for PeerAddr { - type Target = Option; - - fn deref(&self) -> &Self::Target { - &self.addr +impl From> for PeerAddr { + fn from(addr: io::Result) -> Self { + match addr { + Ok(addr) => PeerAddr::TCP(addr), + Err(_) => PeerAddr::Empty, + } } } -impl From> for PeerAddr { - fn from(value: Option) -> Self { - Self { addr: value } +impl From> for PeerAddr { + fn from(addr: io::Result) -> Self { + match addr { + Ok(addr) => { + let Some(path) = addr.as_pathname() else { + return PeerAddr::Empty; + }; + PeerAddr::Unix(path.into()) + } + Err(_) => PeerAddr::Empty, + } } } impl std::fmt::Debug for PeerAddr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - if let Some(addr) = &self.addr { - write!(f, "[{}]", addr) - } else { - write!(f, "") + match self { + PeerAddr::TCP(addr) => write!(f, "[{}]", addr), + PeerAddr::Unix(file_path) => write!(f, "{}", file_path.display()), + PeerAddr::Empty => write!(f, "No address"), } } } diff --git a/pgdog/src/stats/pools.rs b/pgdog/src/stats/pools.rs index 8f8dbee44..3082347b1 100644 --- a/pgdog/src/stats/pools.rs +++ b/pgdog/src/stats/pools.rs @@ -103,7 +103,7 @@ impl Pools { let labels = vec![ ("user".into(), user.user.clone()), ("database".into(), user.database.clone()), - ("host".into(), pool.addr().host.clone()), + ("host".into(), pool.addr().host.to_string()), ("port".into(), pool.addr().port.to_string()), ("shard".into(), shard_num.to_string()), ("role".into(), role.to_string()), From 51b082c6f1572708039b805a8aef031011ff8549 Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 02:32:50 +0100 Subject: [PATCH 03/14] fix: connection to unix socket --- pgdog/src/backend/pool/address.rs | 7 ++++ pgdog/src/backend/server.rs | 55 +++++++++++++++++++++++-------- pgdog/src/net/stream.rs | 11 +++++++ 3 files changed, 60 insertions(+), 13 deletions(-) diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 80501d93f..2aabfc987 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -34,6 +34,13 @@ impl Transport { } } + pub fn unix_socket_path(&self, port: &u16) -> Option { + match self { + Transport::Unix(dir) => Some(dir.join(format!(".s.PGSQL.{}", port))), + Transport::TCP(_) => None, + } + } + pub fn tcp(&self) -> Option<&str> { match self { Transport::TCP(host) => Some(host), diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 2b28e4cdb..211764d01 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -6,7 +6,7 @@ use bytes::{BufMut, BytesMut}; use rustls_pki_types::ServerName; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, - net::TcpStream, + net::{TcpStream, UnixStream}, spawn, time::Instant, }; @@ -18,7 +18,7 @@ use super::{ }; use crate::{ auth::{md5, scram::Client}, - backend::pool::stats::MemoryStats, + backend::pool::{address::Transport, stats::MemoryStats}, config::AuthType, frontend::ClientRequest, net::{ @@ -238,17 +238,32 @@ impl Server { oids: Arc, ) -> Result { debug!("=> {}", addr); - let stream = TcpStream::connect(addr.addr().await?).await?; let config = config(); - - if let Err(err) = tweak(&stream, &config.config.tcp) { - warn!( - "keepalive settings ({}) are not supported on this system, ignoring, error: {} [{}]", - config.config.tcp, err, addr, - ); - } - - let mut stream = Stream::plain(stream, config.config.memory.net_buffer); + let mut stream = match &addr.host { + Transport::TCP(_) => { + let socket = addr.addr().await?; + let tcp = TcpStream::connect(socket).await?; + + if let Err(err) = tweak(&tcp, &config.config.tcp) { + warn!( + "keepalive settings ({}) are not supported on this system, ignoring, error: {} [{}]", + config.config.tcp, err, addr, + ); + } + Stream::plain(tcp, config.config.memory.net_buffer) + } + Transport::Unix(_) => { + let path = addr + .host + .unix_socket_path(&addr.port) + .expect("unix transport"); + debug!("connecting to Unix socket {}", path.display()); + Stream::unix( + UnixStream::connect(&path).await?, + config.config.memory.net_buffer, + ) + } + }; let tls_mode = config.config.general.tls_verify; @@ -455,7 +470,21 @@ impl Server { /// Request query cancellation for the given backend server identifier. pub async fn cancel(addr: &Address, id: BackendKeyData) -> Result<(), Error> { - let mut stream = TcpStream::connect(addr.addr().await?).await?; + let mut stream = match &addr.host { + Transport::TCP(_) => { + let tcp = TcpStream::connect(addr.addr().await?).await?; + Stream::plain(tcp, config().config.memory.net_buffer) + } + Transport::Unix(_) => { + let path = addr + .host + .unix_socket_path(&addr.port) + .expect("Unix transport"); + let unix = UnixStream::connect(&path).await?; + Stream::unix(unix, config().config.memory.net_buffer) + } + }; + stream.write_all(&Startup::Cancel { id }.to_bytes()).await?; stream.flush().await?; diff --git a/pgdog/src/net/stream.rs b/pgdog/src/net/stream.rs index 3cc34e030..1791e6438 100644 --- a/pgdog/src/net/stream.rs +++ b/pgdog/src/net/stream.rs @@ -141,6 +141,17 @@ impl Stream { } } + /// Wrap a unix socket stream. + pub fn unix(stream: UnixStream, capacity: usize) -> Self { + Self { + inner: StreamInner::UnixSockets(BufStream::with_capacity(capacity, capacity, stream)), + io_in_progress: false, + capacity, + tls_identity: None, + tls_client_certificate: false, + } + } + /// Wrap an encrypted TCP stream. pub fn tls( stream: tokio_rustls::TlsStream, From cc8520620903d2d1a0c2e233c5a33a7293a87d80 Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 03:10:47 +0100 Subject: [PATCH 04/14] fix: tests with new address --- .../backend/auth/azure_workload_identity.rs | 3 +- pgdog/src/backend/auth/rds_iam.rs | 5 +- pgdog/src/backend/auth/vault.rs | 3 +- pgdog/src/backend/pool/address.rs | 95 +++++++++---------- pgdog/src/backend/pool/cluster.rs | 12 +-- pgdog/src/backend/pool/guard.rs | 2 +- pgdog/src/backend/pool/lb/test.rs | 18 ++-- pgdog/src/backend/pool/lsn_monitor.rs | 2 +- pgdog/src/backend/pool/monitor.rs | 6 +- pgdog/src/backend/pool/pool_impl.rs | 2 +- pgdog/src/backend/pool/shard/mod.rs | 10 +- pgdog/src/backend/pool/shard/monitor.rs | 8 +- pgdog/src/backend/pool/shard/role_detector.rs | 4 +- pgdog/src/backend/pool/test/mod.rs | 43 +++++---- pgdog/src/backend/pool/token_cache.rs | 3 +- pgdog/src/backend/pool/waiting.rs | 5 +- pgdog/src/backend/schema/sync/pg_dump.rs | 6 +- pgdog/src/backend/server.rs | 10 +- pgdog/src/net/stream.rs | 1 - 19 files changed, 121 insertions(+), 117 deletions(-) diff --git a/pgdog/src/backend/auth/azure_workload_identity.rs b/pgdog/src/backend/auth/azure_workload_identity.rs index 6346b2123..00dee62dd 100644 --- a/pgdog/src/backend/auth/azure_workload_identity.rs +++ b/pgdog/src/backend/auth/azure_workload_identity.rs @@ -40,6 +40,7 @@ mod tests { use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use super::*; + use crate::backend::pool::address::Transport; use crate::config::ServerAuth; use crate::test_utils::set_env_var; use pgdog_config::Role; @@ -52,7 +53,7 @@ mod tests { let _azure_token_file_path = set_env_var("AZURE_FEDERATED_TOKEN_FILE", "/tmp/example"); let addr = Address { - host: "my-awesome-db.postgres.database.azure.com".into(), + host: Transport::new("my-awesome-db.postgres.database.azure.com"), port: 5432, database_name: "postgres".into(), user: "db_user".into(), diff --git a/pgdog/src/backend/auth/rds_iam.rs b/pgdog/src/backend/auth/rds_iam.rs index 2ce28b657..440e0528f 100644 --- a/pgdog/src/backend/auth/rds_iam.rs +++ b/pgdog/src/backend/auth/rds_iam.rs @@ -90,12 +90,13 @@ mod tests { use pgdog_config::Role; use super::*; + use crate::backend::pool::address::Transport; use crate::config::ServerAuth; use crate::test_utils::set_env_var; fn make_addr() -> Address { Address { - host: "db.cluster-abc123.us-east-1.rds.amazonaws.com".into(), + host: Transport::new("db.cluster-abc123.us-east-1.rds.amazonaws.com"), port: 5432, database_name: "postgres".into(), user: "db_user".into(), @@ -170,7 +171,7 @@ mod tests { #[test] fn resolve_region_errors_when_neither_override_nor_inference() { let addr = Address { - host: "postgres.internal.example.com".into(), + host: Transport::new("postgres.internal.example.com"), port: 5432, user: "u".into(), server_iam_region: None, diff --git a/pgdog/src/backend/auth/vault.rs b/pgdog/src/backend/auth/vault.rs index 297a74af0..d70da2987 100644 --- a/pgdog/src/backend/auth/vault.rs +++ b/pgdog/src/backend/auth/vault.rs @@ -151,6 +151,7 @@ mod tests { use super::*; use crate::auth::vault::{VAULT_TOKEN, VaultToken}; + use crate::backend::pool::address::Transport; use crate::config::ConfigAndUsers; fn setup() { @@ -179,7 +180,7 @@ mod tests { fn make_addr(vault_path: Option<&str>) -> Address { Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "testdb".into(), user: "testuser".into(), diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 2aabfc987..3e1ee49d2 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -1,5 +1,4 @@ //! Server address. -use std::default::Default; use std::fmt::Display; use std::net::{SocketAddr, ToSocketAddrs}; use std::ops::Deref; @@ -65,14 +64,8 @@ impl Display for Transport { } } -impl Default for Transport { - fn default() -> Self { - Transport::TCP("127.0.0.1".to_string()) - } -} - /// Server address. -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default, Eq, Hash)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] pub struct Address { /// Server host. pub host: Transport, @@ -102,6 +95,27 @@ pub struct Address { pub configured_role: Role, } +impl Default for Address { + /// Local development defaults: `pgdog` on `127.0.0.1:5432`. `Transport` + /// deliberately has no `Default` — the transport decision must always be + /// explicit — so the TCP variant is chosen here, at this one visible site. + fn default() -> Self { + Address { + host: Transport::TCP("127.0.0.1".to_string()), + port: 5432, + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + database_name: "pgdog".into(), + server_auth: ServerAuth::Password, + server_iam_region: None, + vault_path: None, + vault_refresh_percent: None, + database_number: 0, + configured_role: Role::Primary, + } + } +} + impl From
for pgdog_stats::Address { fn from(value: Address) -> Self { pgdog_stats::Address { @@ -254,25 +268,6 @@ impl Address { true } } - - /// Test convention: `new_test()` represents a primary. Tests that need - /// a replica do `Address { configured_role: Role::Replica, ..new_test() }`. - #[cfg(test)] - pub fn new_test() -> Self { - Self { - host: "127.0.0.1".into(), - port: 5432, - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - database_name: "pgdog".into(), - server_auth: ServerAuth::Password, - server_iam_region: None, - vault_path: None, - vault_refresh_percent: None, - database_number: 0, - configured_role: Role::Primary, - } - } } impl std::fmt::Display for Address { @@ -295,9 +290,9 @@ impl TryFrom for Address { let password = value.password().ok_or(())?.to_string(); let database_name = value.path().replace("/", "").to_string(); - // A URL says nothing about role; fall through to `Role::Auto` - // via the derived `Default`. The PROBE command (the only caller) - // never reads `configured_role` anyway. + // A URL says nothing about role; fall through to the explicit + // `Role::Auto` below. The PROBE command (the only caller) never + // reads `configured_role` anyway. Ok(Self { host: Transport::new(&host), port, @@ -305,7 +300,11 @@ impl TryFrom for Address { user, database_name, server_auth: ServerAuth::Password, - ..Default::default() + server_iam_region: None, + vault_path: None, + vault_refresh_percent: None, + database_number: 0, + configured_role: Role::Auto, }) } } @@ -338,7 +337,7 @@ mod test { let address = Address::new(&database, &user, 0); - assert_eq!(address.host.into(), "127.0.0.1"); + assert_eq!(address.host.to_string(), "127.0.0.1"); assert_eq!(address.port, 6432); assert_eq!(address.database_name, "pgdog"); assert_eq!(address.user, "pgdog"); @@ -419,7 +418,7 @@ mod test { let addr = Address::try_from(Url::parse("postgres://user:password@127.0.0.1:6432/pgdb").unwrap()) .unwrap(); - assert_eq!(addr.host.into(), "127.0.0.1"); + assert_eq!(addr.host.to_string(), "127.0.0.1"); assert_eq!(addr.port, 6432); assert_eq!(addr.database_name, "pgdb"); assert_eq!(addr.user, "user"); @@ -430,7 +429,7 @@ mod test { #[test] fn test_compatible_ignores_password_changes() { - let address = Address::new_test(); + let address = Address::default(); let mut rotated = address.clone(); rotated.passwords = vec!["rotated".into()]; @@ -439,7 +438,7 @@ mod test { #[test] fn test_compatible_rejects_other_field_changes() { - let address = Address::new_test(); + let address = Address::default(); let mut moved = address.clone(); moved.port += 1; @@ -450,13 +449,13 @@ mod test { #[tokio::test] async fn test_auth_secret_password_mode() { - let addr = Address::new_test(); + let addr = Address::default(); assert_eq!(addr.auth_secrets().await.unwrap().first().unwrap(), "pgdog"); } #[tokio::test] async fn test_auth_secrets_returns_valid_password_first() { - let mut addr = Address::new_test(); + let mut addr = Address::default(); let invalid1: Password = "invalid1".into(); let invalid2: Password = "invalid2".into(); let valid: Password = "valid".into(); @@ -470,7 +469,7 @@ mod test { assert!(secrets.first().unwrap().is_valid()); // Even if the valid password is last, it should still come first. - let mut addr = Address::new_test(); + let mut addr = Address::default(); let invalid1: Password = "invalid1".into(); let invalid2: Password = "invalid2".into(); let valid: Password = "valid".into(); @@ -482,7 +481,7 @@ mod test { assert_eq!(secrets.first().unwrap(), "valid"); // With multiple valid passwords, a valid one is still first. - let mut addr = Address::new_test(); + let mut addr = Address::default(); let invalid: Password = "invalid".into(); invalid.valid(false); addr.passwords = vec![invalid, "valid_a".into(), "valid_b".into()]; @@ -493,7 +492,7 @@ mod test { assert!(head == "valid_a" || head == "valid_b"); // Flipping validity at runtime changes which password comes first. - let mut addr = Address::new_test(); + let mut addr = Address::default(); let first: Password = "first".into(); let second: Password = "second".into(); addr.passwords = vec![first.clone(), second.clone()]; @@ -536,7 +535,7 @@ mod test { #[tokio::test] async fn test_auth_secret_rds_iam_serves_token_from_cache() { let addr = Address { - host: "auth-secrets-rds.internal".into(), + host: Transport::new("auth-secrets-rds.internal"), port: 15432, user: "rds_user".into(), server_auth: ServerAuth::RdsIam, @@ -563,7 +562,7 @@ mod test { #[tokio::test] async fn test_auth_secret_azure_workload_identity_serves_token_from_cache() { let addr = Address { - host: "auth-secrets-azure.internal".into(), + host: Transport::new("auth-secrets-azure.internal"), port: 15433, user: "azure_user".into(), server_auth: ServerAuth::AzureWorkloadIdentity, @@ -591,7 +590,7 @@ mod test { use crate::backend::pool::token_cache::{Credentials, FetchedCredentials}; let addr = Address { - host: "auth-secrets-vault.internal".into(), + host: Transport::new("auth-secrets-vault.internal"), port: 15435, user: "configured_user".into(), server_auth: ServerAuth::VaultDynamic, @@ -623,7 +622,7 @@ mod test { #[tokio::test] async fn test_auth_credentials_password_mode_uses_configured_user() { - let addr = Address::new_test(); + let addr = Address::default(); let (user, secrets) = addr.auth_credentials().await.unwrap(); assert_eq!(user, "pgdog"); assert_eq!(secrets.first().unwrap(), "pgdog"); @@ -661,7 +660,7 @@ mod test { #[tokio::test] async fn test_auth_credentials_vault_static_serves_password_from_cache() { let addr = Address { - host: "auth-secrets-vault-static.internal".into(), + host: Transport::new("auth-secrets-vault-static.internal"), port: 15436, user: "pgdog_static".into(), server_auth: ServerAuth::VaultStatic, @@ -718,7 +717,7 @@ mod test { // The monitor is responsible for refreshing it; auth_secrets never // blocks on a refresh. let addr = Address { - host: "auth-secrets-stale.internal".into(), + host: Transport::new("auth-secrets-stale.internal"), port: 15434, user: "stale_user".into(), server_auth: ServerAuth::RdsIam, @@ -753,9 +752,9 @@ mod test { cache.clear_cache_for_testing(); let addr = Address { - host: hostname.into(), + host: Transport::new(&hostname), port: 15432, - ..Address::new_test() + ..Default::default() }; let socket_addr = addr.addr().await.expect("resolve address"); diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index d10dfb85d..fc03323cc 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -814,13 +814,13 @@ mod test { database: "pgdog".into(), }); let primary = Some(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config: Config::default(), }); let replicas = &[PoolConfig { address: Address { configured_role: Role::Replica, - ..Address::new_test() + ..Default::default() }, config: Config::default(), }]; @@ -959,7 +959,7 @@ mod test { primary: Some(&PoolConfig { address: Address { database_name: "pgdog1".into(), - ..Address::new_test() + ..Default::default() }, config: Config::default(), }), @@ -967,7 +967,7 @@ mod test { address: Address { database_name: "pgdog1".into(), configured_role: Role::Replica, - ..Address::new_test() + ..Default::default() }, config: Config::default(), }], @@ -987,7 +987,7 @@ mod test { Cluster { shards: vec![Shard::new(ShardConfig { primary: Some(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config: Config::default(), }), identifier: identifier.clone(), @@ -1018,7 +1018,7 @@ mod test { replicas: &[PoolConfig { address: Address { configured_role: Role::Replica, - ..Address::new_test() + ..Default::default() }, config: Config::default(), }], diff --git a/pgdog/src/backend/pool/guard.rs b/pgdog/src/backend/pool/guard.rs index 03d27cf76..5833187ea 100644 --- a/pgdog/src/backend/pool/guard.rs +++ b/pgdog/src/backend/pool/guard.rs @@ -349,7 +349,7 @@ mod test { }; let pool = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); pool.launch(); diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index 50fd6c1b4..534e1a523 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -2,7 +2,7 @@ use std::collections::HashSet; use std::time::Duration; use tokio::time::sleep; -use crate::backend::pool::{Address, Config, Error, PoolConfig, Request}; +use crate::backend::pool::{Address, Config, Error, PoolConfig, Request, address::Transport}; use crate::backend::replication::publisher::Lsn; use crate::config::{LoadBalancingStrategy, Role}; use pgdog_stats::{LsnStats as StatsLsnStats, ReplicaLag}; @@ -13,7 +13,7 @@ use monitor::Monitor; fn create_test_pool_config(host: &str, port: u16) -> PoolConfig { PoolConfig { address: Address { - host: host.into(), + host: Transport::new(&host), port, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -879,7 +879,7 @@ async fn test_monitor_unbans_all_when_all_unhealthy() { async fn test_monitor_does_not_ban_with_zero_ban_timeout() { let pool_config1 = PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -898,7 +898,7 @@ async fn test_monitor_does_not_ban_with_zero_ban_timeout() { let pool_config2 = PoolConfig { address: Address { - host: "localhost".into(), + host: Transport::new("localhost"), port: 5432, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -1250,7 +1250,7 @@ async fn test_move_conns_to_with_added_replica_matches_by_address() { let new_target_for_existing = lb_new .targets .iter() - .find(|t| t.pool.addr().host == "127.0.0.1") + .find(|t| t.pool.addr().host == Transport::new("127.0.0.1")) .expect("should have target for 127.0.0.1"); assert_eq!(new_target_for_existing.role(), Role::Primary); @@ -1258,7 +1258,7 @@ async fn test_move_conns_to_with_added_replica_matches_by_address() { let new_target_for_added = lb_new .targets .iter() - .find(|t| t.pool.addr().host == "localhost") + .find(|t| t.pool.addr().host == Transport::new("localhost")) .expect("should have target for localhost"); assert_eq!(new_target_for_added.role(), Role::Replica); @@ -1507,7 +1507,7 @@ async fn test_monitor_unbans_all_when_second_target_becomes_unhealthy_after_firs fn create_test_pool_config_weighted(host: &str, port: u16, lb_weight: u8) -> PoolConfig { PoolConfig { address: Address { - host: host.into(), + host: Transport::new(&host), port, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -1928,7 +1928,7 @@ fn test_ban_check_does_not_ban_single_target() { fn test_ban_check_does_not_ban_with_zero_ban_timeout() { let pool_config1 = PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -1947,7 +1947,7 @@ fn test_ban_check_does_not_ban_with_zero_ban_timeout() { let pool_config2 = PoolConfig { address: Address { - host: "localhost".into(), + host: Transport::new("localhost"), port: 5432, user: "pgdog".into(), passwords: vec!["pgdog".into()], diff --git a/pgdog/src/backend/pool/lsn_monitor.rs b/pgdog/src/backend/pool/lsn_monitor.rs index 891d14607..8c60f92e1 100644 --- a/pgdog/src/backend/pool/lsn_monitor.rs +++ b/pgdog/src/backend/pool/lsn_monitor.rs @@ -450,7 +450,7 @@ mod test { }; let pool = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); pool.launch(); diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index 7cbe0831d..42a96d9f9 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -509,7 +509,7 @@ impl Monitor { #[cfg(test)] mod test { use crate::backend::pool::test::pool; - use crate::backend::pool::{Address, Config, PoolConfig}; + use crate::backend::pool::{Address, Config, PoolConfig, address::Transport}; use super::*; @@ -566,7 +566,7 @@ mod test { let pool = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 1, database_name: "pgdog".into(), user: "pgdog".into(), @@ -600,7 +600,7 @@ mod test { let pool = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 376b3cfd3..6bb4a268d 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -89,7 +89,7 @@ impl Pool { #[cfg(test)] pub fn new_test() -> Self { let config = PoolConfig { - address: Address::new_test(), + address: Address::default(), config: Config::default(), }; diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index 1dab2e3af..915cb89ab 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -390,14 +390,14 @@ mod test { crate::logger(); let primary = Some(&PoolConfig { - address: Address::new_test(), + address: Address::default(), ..Default::default() }); let replicas = &[PoolConfig { address: Address { configured_role: Role::Replica, - ..Address::new_test() + ..Default::default() }, ..Default::default() }]; @@ -430,12 +430,12 @@ mod test { crate::logger(); let primary = Some(&PoolConfig { - address: Address::new_test(), + address: Address::default(), ..Default::default() }); let replicas = &[PoolConfig { - address: Address::new_test(), + address: Address::default(), ..Default::default() }]; @@ -468,7 +468,7 @@ mod test { let replicas = &[PoolConfig { address: Address { configured_role: Role::Auto, - ..Address::new_test() + ..Default::default() }, config: super::super::Config { inner: pgdog_stats::Config { diff --git a/pgdog/src/backend/pool/shard/monitor.rs b/pgdog/src/backend/pool/shard/monitor.rs index af3df05e0..826d3bd5e 100644 --- a/pgdog/src/backend/pool/shard/monitor.rs +++ b/pgdog/src/backend/pool/shard/monitor.rs @@ -250,13 +250,13 @@ mod test { #[test] fn test_update_replica_lag_assigns_primary_minus_replica_to_replica_pool() { let primary = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config: Config::default(), }); let replica = Pool::new(&PoolConfig { address: Address { configured_role: Role::Replica, - ..Address::new_test() + ..Default::default() }, config: Config::default(), }); @@ -282,10 +282,10 @@ mod test { async fn test_monitor_updates_roles_on_failover() { crate::logger(); - let primary = Some(&pool_config(Address::new_test())); + let primary = Some(&pool_config(Address::default())); let replicas = [pool_config(Address { configured_role: Role::Auto, - ..Address::new_test() + ..Default::default() })]; let shard = Shard::new(ShardConfig { diff --git a/pgdog/src/backend/pool/shard/role_detector.rs b/pgdog/src/backend/pool/shard/role_detector.rs index 6fcc8eadc..e9d1ca718 100644 --- a/pgdog/src/backend/pool/shard/role_detector.rs +++ b/pgdog/src/backend/pool/shard/role_detector.rs @@ -44,7 +44,7 @@ mod test { use crate::backend::databases::User; use crate::backend::pool::lsn_monitor::LsnStats; - use crate::backend::pool::{Address, Config, PoolConfig}; + use crate::backend::pool::{Address, Config, PoolConfig, address::Transport}; use crate::backend::replication::publisher::Lsn; use crate::config::{ReadWriteSplit, Role}; use pgdog_stats::LsnStats as StatsLsnStats; @@ -55,7 +55,7 @@ mod test { fn create_test_pool_config(host: &str, port: u16, role_detection: bool) -> PoolConfig { PoolConfig { address: Address { - host: host.into(), + host: Transport::new(&host), port, user: "pgdog".into(), passwords: vec!["pgdog".into()], diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index 0d213ded9..6cf0c9316 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -12,6 +12,7 @@ use tokio::time::{Instant, sleep, timeout}; use tokio_util::task::TaskTracker; use crate::backend::ConnectReason; +use crate::backend::pool::address::Transport; use crate::backend::pool::token_cache::TokenCache; use crate::net::ProtocolMessage; use crate::net::{Parse, Protocol, Query, Sync}; @@ -30,7 +31,7 @@ pub fn pool() -> Pool { let pool = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -55,7 +56,7 @@ pub fn pool_with_prepared_capacity(capacity: usize) -> Pool { let pool = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -370,7 +371,7 @@ async fn test_server_force_close_discards_connection() { let pool = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -541,7 +542,7 @@ async fn test_idle_healthcheck_loop() { let pool = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -589,7 +590,7 @@ async fn test_idle_healthcheck_loop_disabled_with_zero_interval() { let pool = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 1, database_name: "pgdog".into(), user: "pgdog".into(), @@ -624,7 +625,7 @@ async fn test_checkout_timeout() { }; let pool = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); pool.launch(); @@ -653,7 +654,7 @@ async fn test_move_conns_to() { let source = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -666,7 +667,7 @@ async fn test_move_conns_to() { let destination = Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -721,13 +722,13 @@ async fn test_move_conns_all_idle() { }; let source = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); source.launch(); let destination = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); @@ -768,13 +769,13 @@ async fn test_move_conns_all_checked_out() { }; let source = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); source.launch(); let destination = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); @@ -823,13 +824,13 @@ async fn test_move_conns_destination_serves_after_launch() { }; let source = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); source.launch(); let destination = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); @@ -866,7 +867,7 @@ fn auth_pool(passwords: Vec) -> Pool { Pool::new(&PoolConfig { address: Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -1004,7 +1005,7 @@ async fn test_lsn_monitor() { }; let pool = Pool::new(&PoolConfig { - address: Address::new_test(), + address: Address::default(), config, }); @@ -1050,7 +1051,7 @@ async fn test_token_refresh_loop_primes_cache_on_cold_start() { }; let addr = Address { - host: "token-refresh-test.internal".into(), + host: Transport::new("token-refresh-test.internal"), port: 15500, user: "refresh_user".into(), server_auth: ServerAuth::RdsIam, @@ -1088,7 +1089,7 @@ async fn test_token_refresh_loop_refreshes_before_expiry() { }; let addr = Address { - host: "token-refresh-expiry.internal".into(), + host: Transport::new("token-refresh-expiry.internal"), port: 15501, user: "refresh_user".into(), server_auth: ServerAuth::RdsIam, @@ -1133,7 +1134,7 @@ async fn test_token_refresh_loop_evicts_on_failed_refresh() { }; let addr = Address { - host: "token-refresh-fail.internal".into(), + host: Transport::new("token-refresh-fail.internal"), port: 15502, user: "refresh_user".into(), server_auth: ServerAuth::RdsIam, @@ -1178,7 +1179,7 @@ async fn test_token_refresh_loop_not_spawned_for_password_auth() { }; let addr = Address { - host: "token-refresh-password.internal".into(), + host: Transport::new("token-refresh-password.internal"), port: 15503, user: "refresh_user".into(), server_auth: ServerAuth::Password, @@ -1218,7 +1219,7 @@ async fn test_token_refresh_loop_stops_on_shutdown() { }; let addr = Address { - host: "token-refresh-shutdown.internal".into(), + host: Transport::new("token-refresh-shutdown.internal"), port: 15504, user: "refresh_user".into(), server_auth: ServerAuth::AzureWorkloadIdentity, diff --git a/pgdog/src/backend/pool/token_cache.rs b/pgdog/src/backend/pool/token_cache.rs index 61905554d..5d72f3ff4 100644 --- a/pgdog/src/backend/pool/token_cache.rs +++ b/pgdog/src/backend/pool/token_cache.rs @@ -280,6 +280,7 @@ impl TokenCache { #[cfg(test)] mod tests { use super::*; + use crate::backend::pool::address::Transport; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -287,7 +288,7 @@ mod tests { /// in the global cache. fn addr(port: u16) -> Address { Address { - host: "token-cache-test.internal".into(), + host: Transport::new("token-cache-test.internal"), port, user: "test_user".into(), ..Default::default() diff --git a/pgdog/src/backend/pool/waiting.rs b/pgdog/src/backend/pool/waiting.rs index 8b02aad2e..5eabbab9c 100644 --- a/pgdog/src/backend/pool/waiting.rs +++ b/pgdog/src/backend/pool/waiting.rs @@ -91,6 +91,7 @@ pub(super) struct Waiter { mod tests { use super::*; use crate::backend::pool::Pool; + use crate::backend::pool::address::Transport; use crate::net::messages::FrontendPid; use tokio::time::{Duration, sleep, timeout}; @@ -152,12 +153,12 @@ mod tests { let pool = Pool::new(&crate::backend::pool::PoolConfig { address: crate::backend::pool::Address { - host: "127.0.0.1".into(), + host: Transport::new("127.0.0.1"), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), passwords: vec!["pgdog".into()], - ..Default::default() + ..crate::backend::pool::Address::default() }, config, }); diff --git a/pgdog/src/backend/schema/sync/pg_dump.rs b/pgdog/src/backend/schema/sync/pg_dump.rs index 5c2981080..43c28b905 100644 --- a/pgdog/src/backend/schema/sync/pg_dump.rs +++ b/pgdog/src/backend/schema/sync/pg_dump.rs @@ -1886,7 +1886,7 @@ mod test { #[test] fn test_build_pg_dump_command_sets_password_env() { - let addr = backend::pool::Address::new_test(); + let addr = backend::pool::Address::default(); let command = build_pg_dump_command("pg_dump", &addr, "secret"); let env = command @@ -1908,7 +1908,7 @@ mod test { #[test] fn test_build_pg_dump_command_sets_tls_for_rds_iam() { - let mut addr = backend::pool::Address::new_test(); + let mut addr = backend::pool::Address::default(); addr.server_auth = ServerAuth::RdsIam; let command = build_pg_dump_command("pg_dump", &addr, "token"); @@ -1923,7 +1923,7 @@ mod test { #[test] fn test_build_pg_dump_command_sets_tls_for_azure_workload_identity() { - let mut addr = backend::pool::Address::new_test(); + let mut addr = backend::pool::Address::default(); addr.server_auth = ServerAuth::AzureWorkloadIdentity; let command = build_pg_dump_command("pg_dump", &addr, "token"); diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 211764d01..7695f166b 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -1425,7 +1425,7 @@ pub mod test { pub(crate) async fn test_server() -> Server { Server::connect( - &Address::new_test(), + &Address::default(), ServerOptions::default(), ConnectReason::Other, Default::default(), @@ -1441,7 +1441,7 @@ pub mod test { Server::connect( &Address { database_name: "pgdog1".into(), - ..Address::new_test() + ..Default::default() }, ServerOptions::default(), ConnectReason::Other, @@ -1453,7 +1453,7 @@ pub mod test { pub async fn test_replication_server() -> Server { Server::connect( - &Address::new_test(), + &Address::default(), ServerOptions::new_replication(), ConnectReason::Replication, Default::default(), @@ -1546,7 +1546,7 @@ pub mod test { } }); - let mut addr = Address::new_test(); + let mut addr = Address::default(); addr.port = port; addr.server_auth = crate::config::ServerAuth::RdsIam; addr.server_iam_region = Some("us-east-1".into()); @@ -1613,7 +1613,7 @@ pub mod test { } }); - let mut addr = Address::new_test(); + let mut addr = Address::default(); addr.port = port; addr.server_auth = crate::config::ServerAuth::AzureWorkloadIdentity; addr.passwords = vec!["wrong-password".into()]; diff --git a/pgdog/src/net/stream.rs b/pgdog/src/net/stream.rs index 1791e6438..2febc1ac3 100644 --- a/pgdog/src/net/stream.rs +++ b/pgdog/src/net/stream.rs @@ -9,7 +9,6 @@ use tracing::trace; use std::io::{Error, ErrorKind}; use std::net::SocketAddr; -use std::ops::Deref; use std::os::fd::AsRawFd; use std::path::PathBuf; use std::pin::Pin; From 4ca476e213d1e13dbb0aab38bff914a6015bc67f Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 03:42:37 +0100 Subject: [PATCH 05/14] feat: add more tests --- integration/unix_sockets/run.sh | 97 +++++++++++++++++++++++++++++++ pgdog/src/backend/pool/address.rs | 25 ++++++++ pgdog/src/backend/server.rs | 65 ++++++++++++++++++++- 3 files changed, 184 insertions(+), 3 deletions(-) create mode 100755 integration/unix_sockets/run.sh diff --git a/integration/unix_sockets/run.sh b/integration/unix_sockets/run.sh new file mode 100755 index 000000000..e080ea397 --- /dev/null +++ b/integration/unix_sockets/run.sh @@ -0,0 +1,97 @@ +#!/bin/bash +# End-to-end test: pgdog connecting to Postgres over a Unix domain socket. +# +# Prerequisites: +# - Postgres listening on a Unix socket (default: /tmp), trust or peer auth. +# - pgdog built in the workspace (the script builds it if needed). +# +# Env overrides: +# UNIX_SOCKET_DIR - Postgres unix_socket_directories entry (default /tmp) +# PG_PORT - Postgres port (default 5432) +# PROXY_PORT - pgdog proxy listen port (default 6432) +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +SOCKET_DIR="${UNIX_SOCKET_DIR:-/tmp}" +PG_PORT="${PG_PORT:-5432}" +PROXY_PORT="${PROXY_PORT:-6432}" + +CONFIG="$(mktemp)" +PGOUT="$(mktemp)" +PGOOD_PID="" + +cleanup() { + if [[ -n "$PGOOD_PID" ]]; then + kill "$PGOOD_PID" 2>/dev/null || true + wait "$PGOOD_PID" 2>/dev/null || true + fi + rm -f "$CONFIG" "$PGOUT" +} +trap cleanup EXIT + +pass() { echo "PASS: $*"; } +fail() { echo "FAIL: $*" >&2; exit 1; } + +# --- 0. Preconditions --------------------------------------------------------- +if ! command -v pg_isready >/dev/null 2>&1; then + fail "pg_isready not found (is Postgres client installed?)" +fi +if ! pg_isready -h "$SOCKET_DIR" -p "$PG_PORT" >/dev/null 2>&1; then + fail "Postgres is not listening on unix socket $SOCKET_DIR:$PG_PORT" +fi +pass "Postgres is listening on unix socket $SOCKET_DIR (port $PG_PORT)" + +# --- 1. Build pgdog ----------------------------------------------------------- +echo "==> building pgdog" +cargo build --manifest-path "$ROOT/Cargo.toml" --bin pgdog + +# --- 2. Config: backend pointed at the socket dir ----------------------------- +cat > "$CONFIG" < starting pgdog" +"$ROOT/target/debug/pgdog" --config "$CONFIG" --users "$ROOT/integration/users.toml" \ + >"$PGOUT" 2>&1 & +PGOOD_PID=$! + +for _ in $(seq 1 30); do + pg_isready -h 127.0.0.1 -p "$PROXY_PORT" -U pgdog -d pgdog >/dev/null 2>&1 && break + sleep 1 +done +if ! pg_isready -h 127.0.0.1 -p "$PROXY_PORT" -U pgdog -d pgdog >/dev/null 2>&1; then + fail "pgdog did not become ready on 127.0.0.1:$PROXY_PORT" +fi +pass "pgdog is accepting connections on 127.0.0.1:$PROXY_PORT" + +# --- 4. Query through pgdog --------------------------------------------------- +echo "==> querying through pgdog" +psql -h 127.0.0.1 -p "$PROXY_PORT" -U pgdog -d pgdog -v ON_ERROR_STOP=1 \ + -c "select version()" >/dev/null \ + || fail "query through pgdog failed" +pass "query through pgdog succeeded" + +# --- 5. Server-side proof: unix socket connections have client_addr NULL ------ +echo "==> backend connections as seen by Postgres" +CONNS="$(psql -h 127.0.0.1 -p "$PG_PORT" -U pgdog -d postgres -t -A \ + -c "select count(*) from pg_stat_activity where usename = 'pgdog' and backend_type = 'client backend' and client_addr is null")" +[[ "$CONNS" != "0" && -n "$CONNS" ]] || fail "no backend connections over unix socket (client_addr NULL)" +pass "$CONNS backend connection(s) over unix socket (client_addr IS NULL)" +psql -h 127.0.0.1 -p "$PG_PORT" -U pgdog -d postgres \ + -c "select pid, client_addr, client_hostname from pg_stat_activity where usename = 'pgdog' and backend_type = 'client backend'" + +# --- 6. pgdog's own view of the backend address (best effort) ----------------- +psql -h 127.0.0.1 -p "$PROXY_PORT" -U pgdog -d pgdog -c "SHOW SERVERS" \ + || echo "(note: SHOW SERVERS not available, skipping)" + +echo +echo "ALL PASSED: pgdog -> Postgres over Unix domain socket" diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 3e1ee49d2..9ad395d2d 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -765,4 +765,29 @@ mod test { Some(socket_addr.ip()) ); } + + #[test] + fn test_transport_parsing() { + // Absolute paths are Unix socket directories. + assert!(matches!(Transport::new("/tmp"), Transport::Unix(_))); + assert!(matches!( + Transport::new("/var/run/postgresql"), + Transport::Unix(_) + )); + // Everything else is a TCP hostname. + assert!(matches!(Transport::new("127.0.0.1"), Transport::TCP(_))); + assert!(matches!(Transport::new("localhost"), Transport::TCP(_))); + } + + #[test] + fn test_unix_socket_path() { + let unix = Transport::new("/tmp"); + assert_eq!( + unix.unix_socket_path(&5432), + Some(PathBuf::from("/tmp/.s.PGSQL.5432")) + ); + + let tcp = Transport::new("127.0.0.1"); + assert_eq!(tcp.unix_socket_path(&5432), None); + } } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 7695f166b..c6ec93235 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -267,8 +267,8 @@ impl Server { let tls_mode = config.config.general.tls_verify; - // Only attempt TLS if not in Disabled mode - if tls_mode != TlsVerifyMode::Disabled { + // Only attempt TLS if not in Disabled mode and its not connecting to a unix socket + if tls_mode != TlsVerifyMode::Disabled && addr.host.tcp().is_some() { debug!( "requesting TLS connection with verify mode: {:?} [{}]", tls_mode, addr, @@ -1350,7 +1350,7 @@ pub mod test { use bytes::{BufMut, BytesMut}; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, + net::{TcpListener, UnixListener}, }; use crate::{ @@ -1637,6 +1637,65 @@ pub mod test { server_task.await.unwrap(); } + #[tokio::test] + async fn test_connect_over_unix_socket() { + let dir = std::env::temp_dir().join(format!("pgdog-unix-test-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let port = 6543; + let socket_path = dir.join(format!(".s.PGSQL.{}", port)); + + // Bind the listener at the path pgdog derives from the socket + // directory and port. If the derivation is wrong, connect fails. + let listener = UnixListener::bind(&socket_path).unwrap(); + + let server_task = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + + // TLS is skipped for Unix sockets: the first packet must be the + // Startup message, never an SSLRequest. + let startup = Startup::from_stream(&mut socket).await.unwrap(); + assert!( + matches!(startup, Startup::Startup { .. }), + "expected Startup, got {:?}", + startup + ); + + // Peer/trust auth replies: no password is ever exchanged. + socket + .write_all(&Authentication::Ok.to_bytes()) + .await + .unwrap(); + socket + .write_all(&BackendKeyData::random_legacy().to_bytes()) + .await + .unwrap(); + socket + .write_all(&ReadyForQuery::idle().to_bytes()) + .await + .unwrap(); + }); + + let addr = Address { + host: Transport::new(dir.to_str().unwrap()), + port, + ..Address::default() + }; + + let server = Server::connect( + &addr, + ServerOptions::default(), + ConnectReason::Other, + Default::default(), + ) + .await + .expect("connect over unix socket"); + + drop(server); + server_task.await.unwrap(); + + std::fs::remove_dir_all(&dir).unwrap(); + } + #[tokio::test] async fn test_simple_query() { let mut server = test_server().await; From f8a1a1eded778538e8c3c117da62b240f36c186f Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 03:45:32 +0100 Subject: [PATCH 06/14] fix: clippy warnings --- pgdog/src/admin/show_stats.rs | 2 +- pgdog/src/backend/auth/rds_iam.rs | 2 +- pgdog/src/backend/pool/address.rs | 2 +- pgdog/src/backend/schema/sync/pg_dump.rs | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pgdog/src/admin/show_stats.rs b/pgdog/src/admin/show_stats.rs index 8011fe95c..7f253ede3 100644 --- a/pgdog/src/admin/show_stats.rs +++ b/pgdog/src/admin/show_stats.rs @@ -78,7 +78,7 @@ impl Command for ShowStats { dr.add(user.database.as_str()) .add(user.user.as_str()) - .add(&pool.addr().host.to_string()) + .add(pool.addr().host.to_string()) .add(pool.addr().port as i64) .add(shard_num) .add(role.to_string()); diff --git a/pgdog/src/backend/auth/rds_iam.rs b/pgdog/src/backend/auth/rds_iam.rs index 440e0528f..07c9e6491 100644 --- a/pgdog/src/backend/auth/rds_iam.rs +++ b/pgdog/src/backend/auth/rds_iam.rs @@ -38,7 +38,7 @@ fn resolve_region(addr: &Address) -> Result { } let host = addr.host.tcp().expect("Host must be a TCP address"); - infer_region_from_rds_host(&host).ok_or_else(|| { + infer_region_from_rds_host(host).ok_or_else(|| { Error::RdsIamToken(format!( "unable to infer AWS region from host \"{}\"; set \"server_iam_region\"", addr.host diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 9ad395d2d..e118ee7d0 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -242,7 +242,7 @@ impl Address { let host = self.host.tcp().expect("addr must be a TCP address"); if dns_cache_override_enabled { - let ip = DnsCache::global().resolve(&host).await?; + let ip = DnsCache::global().resolve(host).await?; return Ok(SocketAddr::new(ip, self.port)); } diff --git a/pgdog/src/backend/schema/sync/pg_dump.rs b/pgdog/src/backend/schema/sync/pg_dump.rs index 43c28b905..a4808b4f7 100644 --- a/pgdog/src/backend/schema/sync/pg_dump.rs +++ b/pgdog/src/backend/schema/sync/pg_dump.rs @@ -186,7 +186,7 @@ fn build_pg_dump_command( command .arg("--schema-only") .arg("-h") - .arg(&addr.host.to_string()) + .arg(addr.host.to_string()) .arg("-p") .arg(addr.port.to_string()) .arg("-U") From d9f3926bf9475a5576111cc7cd19e82b3ff46d5d Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:17:32 +0100 Subject: [PATCH 07/14] fix: remove panics --- pgdog/src/backend/auth/rds_iam.rs | 4 +- pgdog/src/backend/error.rs | 5 ++- pgdog/src/backend/pool/address.rs | 57 ++------------------------ pgdog/src/backend/pool/mod.rs | 1 + pgdog/src/backend/pool/transport.rs | 62 +++++++++++++++++++++++++++++ pgdog/src/backend/server.rs | 16 +++----- 6 files changed, 77 insertions(+), 68 deletions(-) create mode 100644 pgdog/src/backend/pool/transport.rs diff --git a/pgdog/src/backend/auth/rds_iam.rs b/pgdog/src/backend/auth/rds_iam.rs index 07c9e6491..d16227e9e 100644 --- a/pgdog/src/backend/auth/rds_iam.rs +++ b/pgdog/src/backend/auth/rds_iam.rs @@ -37,7 +37,7 @@ fn resolve_region(addr: &Address) -> Result { return Ok(region.clone()); } - let host = addr.host.tcp().expect("Host must be a TCP address"); + let host = addr.host.tcp()?; infer_region_from_rds_host(host).ok_or_else(|| { Error::RdsIamToken(format!( "unable to infer AWS region from host \"{}\"; set \"server_iam_region\"", @@ -54,7 +54,7 @@ fn resolve_region(addr: &Address) -> Result { pub(crate) async fn token(addr: Address) -> Result<(String, SystemTime), Error> { let region = resolve_region(&addr)?; let sdk_config = aws_config::load_defaults(BehaviorVersion::latest()).await; - let host = addr.host.tcp().expect("host must be a TCP address"); + let host = addr.host.tcp()?; let config = AuthTokenConfig::builder() .hostname(host) diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 6b344c820..01f10bb5d 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -1,6 +1,6 @@ use thiserror::Error; -use crate::net::messages::ErrorResponse; +use crate::{backend::pool::transport::TransportError, net::messages::ErrorResponse}; use super::databases::User; @@ -146,6 +146,9 @@ pub enum Error { #[error("missing canonical oid for type {0}")] MissingCanonicalOid(String), + + #[error(transparent)] + Transport(#[from] TransportError), } impl From for Error { diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index e118ee7d0..6f7233944 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -1,9 +1,6 @@ //! Server address. -use std::fmt::Display; use std::net::{SocketAddr, ToSocketAddrs}; use std::ops::Deref; -use std::path::Path; -use std::path::PathBuf; use pgdog_config::Role; use pgdog_config::users::PasswordKind; @@ -15,55 +12,9 @@ use crate::backend::Error; use crate::backend::auth::{azure_workload_identity, rds_iam, vault}; use crate::backend::pool::dns_cache::DnsCache; use crate::backend::pool::token_cache::TokenCache; +use crate::backend::pool::transport::Transport; use crate::config::{Database, ServerAuth, User, config}; -/// Transport enum -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] -pub enum Transport { - TCP(String), - Unix(PathBuf), -} - -impl Transport { - pub fn new(value: &str) -> Self { - if value.starts_with('/') { - Transport::Unix(value.into()) - } else { - Transport::TCP(value.to_string()) - } - } - - pub fn unix_socket_path(&self, port: &u16) -> Option { - match self { - Transport::Unix(dir) => Some(dir.join(format!(".s.PGSQL.{}", port))), - Transport::TCP(_) => None, - } - } - - pub fn tcp(&self) -> Option<&str> { - match self { - Transport::TCP(host) => Some(host), - Transport::Unix(_) => None, - } - } - - pub fn unix(&self) -> Option<&Path> { - match self { - Transport::TCP(_) => None, - Transport::Unix(path_buf) => Some(path_buf), - } - } -} - -impl Display for Transport { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Transport::TCP(addr) => write!(f, "{}", addr), - Transport::Unix(path_buf) => write!(f, "{}", path_buf.display()), - } - } -} - /// Server address. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] pub struct Address { @@ -96,9 +47,7 @@ pub struct Address { } impl Default for Address { - /// Local development defaults: `pgdog` on `127.0.0.1:5432`. `Transport` - /// deliberately has no `Default` — the transport decision must always be - /// explicit — so the TCP variant is chosen here, at this one visible site. + /// Local development defaults: `pgdog` on `127.0.0.1:5432`. ` fn default() -> Self { Address { host: Transport::TCP("127.0.0.1".to_string()), @@ -239,7 +188,7 @@ impl Address { /// pub(crate) async fn addr(&self) -> Result { let dns_cache_override_enabled = config().config.general.dns_ttl().is_some(); - let host = self.host.tcp().expect("addr must be a TCP address"); + let host = self.host.tcp()?; if dns_cache_override_enabled { let ip = DnsCache::global().resolve(host).await?; diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index 99430d448..eec4f229a 100644 --- a/pgdog/src/backend/pool/mod.rs +++ b/pgdog/src/backend/pool/mod.rs @@ -25,6 +25,7 @@ pub mod state; pub mod stats; pub mod taken; pub mod token_cache; +pub mod transport; pub mod waiting; pub use address::Address; diff --git a/pgdog/src/backend/pool/transport.rs b/pgdog/src/backend/pool/transport.rs new file mode 100644 index 000000000..015adc1cf --- /dev/null +++ b/pgdog/src/backend/pool/transport.rs @@ -0,0 +1,62 @@ +use serde::Deserialize; +use serde::Serialize; +use std::fmt::Display; +use std::path::Path; +use std::path::PathBuf; +use thiserror::Error; + +/// Transport enum +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] +pub enum Transport { + TCP(String), + Unix(PathBuf), +} + +#[derive(Debug, Error, Clone, PartialEq, Eq)] +pub enum TransportError { + #[error("Expected a TCP host but address is a unix socket directory {0}")] + ExpectedTCP(PathBuf), + + #[error("Expected Unix socket directory but address is a TCP host {0}")] + ExpectedUnix(String), +} + +impl Transport { + pub fn new(value: &str) -> Self { + if value.starts_with('/') { + Transport::Unix(value.into()) + } else { + Transport::TCP(value.to_string()) + } + } + + pub fn unix_socket_path(&self, port: &u16) -> Result { + match self { + Transport::TCP(host) => Err(TransportError::ExpectedUnix(host.clone())), + Transport::Unix(dir) => Ok(dir.join(format!(".s.PGSQL.{}", port))), + } + } + + pub fn tcp(&self) -> Result<&str, TransportError> { + match self { + Transport::TCP(host) => Ok(host), + Transport::Unix(path) => Err(TransportError::ExpectedTCP(path.clone())), + } + } + + pub fn unix(&self) -> Result<&Path, TransportError> { + match self { + Transport::TCP(host) => Err(TransportError::ExpectedUnix(host.clone())), + Transport::Unix(path_buf) => Ok(path_buf), + } + } +} + +impl Display for Transport { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Transport::TCP(addr) => write!(f, "{}", addr), + Transport::Unix(path_buf) => write!(f, "{}", path_buf.display()), + } + } +} diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index c6ec93235..812d91080 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -18,7 +18,7 @@ use super::{ }; use crate::{ auth::{md5, scram::Client}, - backend::pool::{address::Transport, stats::MemoryStats}, + backend::pool::{stats::MemoryStats, transport::Transport}, config::AuthType, frontend::ClientRequest, net::{ @@ -253,10 +253,8 @@ impl Server { Stream::plain(tcp, config.config.memory.net_buffer) } Transport::Unix(_) => { - let path = addr - .host - .unix_socket_path(&addr.port) - .expect("unix transport"); + let path = addr.host.unix_socket_path(&addr.port)?; + debug!("connecting to Unix socket {}", path.display()); Stream::unix( UnixStream::connect(&path).await?, @@ -268,7 +266,7 @@ impl Server { let tls_mode = config.config.general.tls_verify; // Only attempt TLS if not in Disabled mode and its not connecting to a unix socket - if tls_mode != TlsVerifyMode::Disabled && addr.host.tcp().is_some() { + if tls_mode != TlsVerifyMode::Disabled && addr.host.tcp().is_ok() { debug!( "requesting TLS connection with verify mode: {:?} [{}]", tls_mode, addr, @@ -293,11 +291,7 @@ impl Server { )?; let plain = stream.take()?; - let host = addr - .host - .tcp() - .expect("host address must be a TCP address") - .to_owned(); + let host = addr.host.tcp()?.to_owned(); let server_name = ServerName::try_from(host)?; debug!("connecting with TLS to server name: {:?}", server_name); From 4f81140ec0c9783a656c35101ec8e6f2367c23f8 Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:38:41 +0100 Subject: [PATCH 08/14] fix: integration test to use setup scripts --- integration/unix_sockets/pgdog.toml | 15 +++++ integration/unix_sockets/run.sh | 99 ++++------------------------- integration/unix_sockets/users.toml | 4 ++ 3 files changed, 33 insertions(+), 85 deletions(-) create mode 100644 integration/unix_sockets/pgdog.toml create mode 100644 integration/unix_sockets/users.toml diff --git a/integration/unix_sockets/pgdog.toml b/integration/unix_sockets/pgdog.toml new file mode 100644 index 000000000..601b937f1 --- /dev/null +++ b/integration/unix_sockets/pgdog.toml @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------------------ +# ----- General ---------------------------------------------------------------- + +[general] +auth_type = "trust" + +# ------------------------------------------------------------------------------ +# ----- Database :: pgdog ------------------------------------------------------ + +[[databases]] +name = "pgdog" +host = "/tmp" +port = 5432 +database_name = "pgdog" +user = "pgdog" diff --git a/integration/unix_sockets/run.sh b/integration/unix_sockets/run.sh index e080ea397..f433118ad 100755 --- a/integration/unix_sockets/run.sh +++ b/integration/unix_sockets/run.sh @@ -3,95 +3,24 @@ # # Prerequisites: # - Postgres listening on a Unix socket (default: /tmp), trust or peer auth. -# - pgdog built in the workspace (the script builds it if needed). # -# Env overrides: -# UNIX_SOCKET_DIR - Postgres unix_socket_directories entry (default /tmp) -# PG_PORT - Postgres port (default 5432) -# PROXY_PORT - pgdog proxy listen port (default 6432) +# Verifies: +# - pgdog's backend connections reach Postgres over the socket: +# pg_stat_activity shows client_addr IS NULL for them. set -euo pipefail +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +source ${SCRIPT_DIR}/../common.sh -ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" -SOCKET_DIR="${UNIX_SOCKET_DIR:-/tmp}" -PG_PORT="${PG_PORT:-5432}" -PROXY_PORT="${PROXY_PORT:-6432}" +run_pgdog $SCRIPT_DIR +wait_for_pgdog -CONFIG="$(mktemp)" -PGOUT="$(mktemp)" -PGOOD_PID="" -cleanup() { - if [[ -n "$PGOOD_PID" ]]; then - kill "$PGOOD_PID" 2>/dev/null || true - wait "$PGOOD_PID" 2>/dev/null || true - fi - rm -f "$CONFIG" "$PGOUT" -} -trap cleanup EXIT - -pass() { echo "PASS: $*"; } -fail() { echo "FAIL: $*" >&2; exit 1; } - -# --- 0. Preconditions --------------------------------------------------------- -if ! command -v pg_isready >/dev/null 2>&1; then - fail "pg_isready not found (is Postgres client installed?)" -fi -if ! pg_isready -h "$SOCKET_DIR" -p "$PG_PORT" >/dev/null 2>&1; then - fail "Postgres is not listening on unix socket $SOCKET_DIR:$PG_PORT" -fi -pass "Postgres is listening on unix socket $SOCKET_DIR (port $PG_PORT)" - -# --- 1. Build pgdog ----------------------------------------------------------- -echo "==> building pgdog" -cargo build --manifest-path "$ROOT/Cargo.toml" --bin pgdog - -# --- 2. Config: backend pointed at the socket dir ----------------------------- -cat > "$CONFIG" < starting pgdog" -"$ROOT/target/debug/pgdog" --config "$CONFIG" --users "$ROOT/integration/users.toml" \ - >"$PGOUT" 2>&1 & -PGOOD_PID=$! - -for _ in $(seq 1 30); do - pg_isready -h 127.0.0.1 -p "$PROXY_PORT" -U pgdog -d pgdog >/dev/null 2>&1 && break - sleep 1 -done -if ! pg_isready -h 127.0.0.1 -p "$PROXY_PORT" -U pgdog -d pgdog >/dev/null 2>&1; then - fail "pgdog did not become ready on 127.0.0.1:$PROXY_PORT" +CONNS=$(psql -h 127.0.0.1 -p 5432 -U pgdog -d postgres -t -A -c \ + "select count(*) from pg_stat_activity where usename = 'pgdog' and backend_type = 'client backend' and client_addr is null") +if [ -z "${CONNS}" ] || [ "${CONNS}" = "0" ]; then + echo "FAIL: no backend connections over Unix socket (client_addr IS NULL)" >&2 + exit 1 fi -pass "pgdog is accepting connections on 127.0.0.1:$PROXY_PORT" - -# --- 4. Query through pgdog --------------------------------------------------- -echo "==> querying through pgdog" -psql -h 127.0.0.1 -p "$PROXY_PORT" -U pgdog -d pgdog -v ON_ERROR_STOP=1 \ - -c "select version()" >/dev/null \ - || fail "query through pgdog failed" -pass "query through pgdog succeeded" - -# --- 5. Server-side proof: unix socket connections have client_addr NULL ------ -echo "==> backend connections as seen by Postgres" -CONNS="$(psql -h 127.0.0.1 -p "$PG_PORT" -U pgdog -d postgres -t -A \ - -c "select count(*) from pg_stat_activity where usename = 'pgdog' and backend_type = 'client backend' and client_addr is null")" -[[ "$CONNS" != "0" && -n "$CONNS" ]] || fail "no backend connections over unix socket (client_addr NULL)" -pass "$CONNS backend connection(s) over unix socket (client_addr IS NULL)" -psql -h 127.0.0.1 -p "$PG_PORT" -U pgdog -d postgres \ - -c "select pid, client_addr, client_hostname from pg_stat_activity where usename = 'pgdog' and backend_type = 'client backend'" - -# --- 6. pgdog's own view of the backend address (best effort) ----------------- -psql -h 127.0.0.1 -p "$PROXY_PORT" -U pgdog -d pgdog -c "SHOW SERVERS" \ - || echo "(note: SHOW SERVERS not available, skipping)" +echo "PASS: ${CONNS} backend connection(s) over Unix socket" -echo -echo "ALL PASSED: pgdog -> Postgres over Unix domain socket" +stop_pgdog diff --git a/integration/unix_sockets/users.toml b/integration/unix_sockets/users.toml new file mode 100644 index 000000000..581cdb75b --- /dev/null +++ b/integration/unix_sockets/users.toml @@ -0,0 +1,4 @@ +[[users]] +name = "pgdog" +database = "pgdog" +password = "pgdog" From 9ecc15c1ffde4ddec9cc4f445bb3d41b5afac19d Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:49:50 +0100 Subject: [PATCH 09/14] fix: failing tests --- integration/unix_sockets/run.sh | 44 +++++++++++++++---- .../backend/auth/azure_workload_identity.rs | 2 +- pgdog/src/backend/auth/rds_iam.rs | 2 +- pgdog/src/backend/auth/vault.rs | 2 +- pgdog/src/backend/pool/address.rs | 4 +- pgdog/src/backend/pool/lb/test.rs | 2 +- pgdog/src/backend/pool/monitor.rs | 2 +- pgdog/src/backend/pool/shard/role_detector.rs | 2 +- pgdog/src/backend/pool/test/mod.rs | 2 +- pgdog/src/backend/pool/token_cache.rs | 2 +- pgdog/src/backend/pool/waiting.rs | 2 +- 11 files changed, 46 insertions(+), 20 deletions(-) diff --git a/integration/unix_sockets/run.sh b/integration/unix_sockets/run.sh index f433118ad..342550451 100755 --- a/integration/unix_sockets/run.sh +++ b/integration/unix_sockets/run.sh @@ -1,21 +1,46 @@ #!/bin/bash # End-to-end test: pgdog connecting to Postgres over a Unix domain socket. -# -# Prerequisites: -# - Postgres listening on a Unix socket (default: /tmp), trust or peer auth. -# -# Verifies: -# - pgdog's backend connections reach Postgres over the socket: -# pg_stat_activity shows client_addr IS NULL for them. set -euo pipefail SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) source ${SCRIPT_DIR}/../common.sh -run_pgdog $SCRIPT_DIR +export PGPASSWORD=pgdog +CTL_PSQL=(psql -h 127.0.0.1 -p 5432 -U pgdog -d postgres -t -A) + +# Detect socket dir + port from the running Postgres. +SOCKET_DIR=$("${CTL_PSQL[@]}" -c "show unix_socket_directories" | cut -d, -f1) +PG_BACKEND_PORT=$("${CTL_PSQL[@]}" -c "show port") +if [ -z "${SOCKET_DIR}" ]; then + echo "FAIL: could not detect unix_socket_directories from Postgres" >&2 + exit 1 +fi +echo "Postgres unix socket directory: ${SOCKET_DIR} (port ${PG_BACKEND_PORT})" + +# Pre-flight: pgdog must be able to reach Postgres over the socket as pgdog. +if ! psql -h "${SOCKET_DIR}" -p "${PG_BACKEND_PORT}" -U pgdog -d postgres \ + -c "select 1" >/dev/null 2>&1; then + echo "FAIL: cannot connect to Postgres over unix socket ${SOCKET_DIR} as pgdog." >&2 + echo " Check pg_hba.conf 'local' lines: trust, or peer with OS user pgdog." >&2 + exit 1 +fi + +# Patch the static config with the detected socket dir/port into a temp dir. +TMP_CFG_DIR=$(mktemp -d /tmp/pgdog-unix-cfg.XXXXXX) +sed -e "s|^host = .*|host = \"${SOCKET_DIR}\"|" \ + -e "s|^port = .*|port = ${PG_BACKEND_PORT}|" \ + "${SCRIPT_DIR}/pgdog.toml" > "${TMP_CFG_DIR}/pgdog.toml" +cp "${SCRIPT_DIR}/users.toml" "${TMP_CFG_DIR}/" + +run_pgdog "${TMP_CFG_DIR}" wait_for_pgdog +# 1. Query through the proxy. +psql -h 127.0.0.1 -p 6432 -U pgdog -d pgdog -v ON_ERROR_STOP=1 \ + -c "select version()" >/dev/null +echo "PASS: query through pgdog" -CONNS=$(psql -h 127.0.0.1 -p 5432 -U pgdog -d postgres -t -A -c \ +# 2. Backend connections over the Unix socket have client_addr IS NULL. +CONNS=$("${CTL_PSQL[@]}" -c \ "select count(*) from pg_stat_activity where usename = 'pgdog' and backend_type = 'client backend' and client_addr is null") if [ -z "${CONNS}" ] || [ "${CONNS}" = "0" ]; then echo "FAIL: no backend connections over Unix socket (client_addr IS NULL)" >&2 @@ -24,3 +49,4 @@ fi echo "PASS: ${CONNS} backend connection(s) over Unix socket" stop_pgdog +rm -rf "${TMP_CFG_DIR}" diff --git a/pgdog/src/backend/auth/azure_workload_identity.rs b/pgdog/src/backend/auth/azure_workload_identity.rs index 00dee62dd..0f2086543 100644 --- a/pgdog/src/backend/auth/azure_workload_identity.rs +++ b/pgdog/src/backend/auth/azure_workload_identity.rs @@ -40,7 +40,7 @@ mod tests { use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use super::*; - use crate::backend::pool::address::Transport; + use crate::backend::pool::transport::Transport; use crate::config::ServerAuth; use crate::test_utils::set_env_var; use pgdog_config::Role; diff --git a/pgdog/src/backend/auth/rds_iam.rs b/pgdog/src/backend/auth/rds_iam.rs index d16227e9e..e0563ac41 100644 --- a/pgdog/src/backend/auth/rds_iam.rs +++ b/pgdog/src/backend/auth/rds_iam.rs @@ -90,7 +90,7 @@ mod tests { use pgdog_config::Role; use super::*; - use crate::backend::pool::address::Transport; + use crate::backend::pool::transport::Transport; use crate::config::ServerAuth; use crate::test_utils::set_env_var; diff --git a/pgdog/src/backend/auth/vault.rs b/pgdog/src/backend/auth/vault.rs index d70da2987..ad304821f 100644 --- a/pgdog/src/backend/auth/vault.rs +++ b/pgdog/src/backend/auth/vault.rs @@ -151,7 +151,7 @@ mod tests { use super::*; use crate::auth::vault::{VAULT_TOKEN, VaultToken}; - use crate::backend::pool::address::Transport; + use crate::backend::pool::transport::Transport; use crate::config::ConfigAndUsers; fn setup() { diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 6f7233944..143aea924 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -733,10 +733,10 @@ mod test { let unix = Transport::new("/tmp"); assert_eq!( unix.unix_socket_path(&5432), - Some(PathBuf::from("/tmp/.s.PGSQL.5432")) + Ok(std::path::PathBuf::from("/tmp/.s.PGSQL.5432")) ); let tcp = Transport::new("127.0.0.1"); - assert_eq!(tcp.unix_socket_path(&5432), None); + assert!(tcp.unix_socket_path(&5432).is_err()); } } diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index 534e1a523..b25ccd10a 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -2,7 +2,7 @@ use std::collections::HashSet; use std::time::Duration; use tokio::time::sleep; -use crate::backend::pool::{Address, Config, Error, PoolConfig, Request, address::Transport}; +use crate::backend::pool::{Address, Config, Error, PoolConfig, Request, transport::Transport}; use crate::backend::replication::publisher::Lsn; use crate::config::{LoadBalancingStrategy, Role}; use pgdog_stats::{LsnStats as StatsLsnStats, ReplicaLag}; diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index 42a96d9f9..6bfc05548 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -509,7 +509,7 @@ impl Monitor { #[cfg(test)] mod test { use crate::backend::pool::test::pool; - use crate::backend::pool::{Address, Config, PoolConfig, address::Transport}; + use crate::backend::pool::{Address, Config, PoolConfig, transport::Transport}; use super::*; diff --git a/pgdog/src/backend/pool/shard/role_detector.rs b/pgdog/src/backend/pool/shard/role_detector.rs index e9d1ca718..cf566047b 100644 --- a/pgdog/src/backend/pool/shard/role_detector.rs +++ b/pgdog/src/backend/pool/shard/role_detector.rs @@ -44,7 +44,7 @@ mod test { use crate::backend::databases::User; use crate::backend::pool::lsn_monitor::LsnStats; - use crate::backend::pool::{Address, Config, PoolConfig, address::Transport}; + use crate::backend::pool::{Address, Config, PoolConfig, transport::Transport}; use crate::backend::replication::publisher::Lsn; use crate::config::{ReadWriteSplit, Role}; use pgdog_stats::LsnStats as StatsLsnStats; diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index 6cf0c9316..a36ed8394 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -12,7 +12,7 @@ use tokio::time::{Instant, sleep, timeout}; use tokio_util::task::TaskTracker; use crate::backend::ConnectReason; -use crate::backend::pool::address::Transport; +use crate::backend::pool::transport::Transport; use crate::backend::pool::token_cache::TokenCache; use crate::net::ProtocolMessage; use crate::net::{Parse, Protocol, Query, Sync}; diff --git a/pgdog/src/backend/pool/token_cache.rs b/pgdog/src/backend/pool/token_cache.rs index 5d72f3ff4..68dc654d6 100644 --- a/pgdog/src/backend/pool/token_cache.rs +++ b/pgdog/src/backend/pool/token_cache.rs @@ -280,7 +280,7 @@ impl TokenCache { #[cfg(test)] mod tests { use super::*; - use crate::backend::pool::address::Transport; + use crate::backend::pool::transport::Transport; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; diff --git a/pgdog/src/backend/pool/waiting.rs b/pgdog/src/backend/pool/waiting.rs index 5eabbab9c..819845ed0 100644 --- a/pgdog/src/backend/pool/waiting.rs +++ b/pgdog/src/backend/pool/waiting.rs @@ -91,7 +91,7 @@ pub(super) struct Waiter { mod tests { use super::*; use crate::backend::pool::Pool; - use crate::backend::pool::address::Transport; + use crate::backend::pool::transport::Transport; use crate::net::messages::FrontendPid; use tokio::time::{Duration, sleep, timeout}; From ad1075c1b139ccce6ba7e77d42737345c2d56e26 Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:51:11 +0100 Subject: [PATCH 10/14] feat: add further tests for cov --- pgdog/src/backend/pool/test/mod.rs | 2 +- pgdog/src/net/stream.rs | 214 ++++++++++++++++++++++++++++- 2 files changed, 214 insertions(+), 2 deletions(-) diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index a36ed8394..a52a46424 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -12,8 +12,8 @@ use tokio::time::{Instant, sleep, timeout}; use tokio_util::task::TaskTracker; use crate::backend::ConnectReason; -use crate::backend::pool::transport::Transport; use crate::backend::pool::token_cache::TokenCache; +use crate::backend::pool::transport::Transport; use crate::net::ProtocolMessage; use crate::net::{Parse, Protocol, Query, Sync}; use crate::state::State; diff --git a/pgdog/src/net/stream.rs b/pgdog/src/net/stream.rs index 2febc1ac3..23c41f69d 100644 --- a/pgdog/src/net/stream.rs +++ b/pgdog/src/net/stream.rs @@ -455,7 +455,10 @@ mod tests { use std::time::Duration; use super::*; - use tokio::net::TcpListener; + use tokio::{ + io::AsyncWriteExt, + net::{TcpListener, UnixListener, UnixStream}, + }; #[tokio::test] async fn test_io_in_progress_initially_false() { @@ -573,4 +576,213 @@ mod tests { assert_eq!(stream.liveness(), Liveness::Clean); } + + // ── Unix domain socket streams ────────────────────────────────────────── + + /// Bind a Unix socket pair and wrap the *connecting* end in a [`Stream`] + /// (pgdog is the connecting client; the listener is the Postgres side). + async fn unix_pair(name: &str) -> (Stream, UnixStream, std::path::PathBuf) { + let dir = + std::env::temp_dir().join(format!("pgdog-stream-{}-{}", name, std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("test.sock"); + let listener = UnixListener::bind(&path).unwrap(); + let server = tokio::spawn(async move { + let (accepted, _) = listener.accept().await.unwrap(); + accepted + }); + let client = UnixStream::connect(&path).await.unwrap(); + let stream = Stream::unix(client, 4096); + (stream, server.await.unwrap(), dir) + } + + #[tokio::test] + async fn test_unix_stream_peer_addr() { + let (stream, _peer, dir) = unix_pair("peer").await; + + let peer_addr = stream.peer_addr(); + let expected = dir.join("test.sock"); + match &peer_addr { + PeerAddr::Unix(path) => assert_eq!(path, &expected), + other => panic!("expected PeerAddr::Unix({:?}), got {:?}", expected, other), + } + + drop(stream); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[tokio::test] + async fn test_unix_stream_liveness() { + let (mut stream, mut peer, dir) = unix_pair("liveness").await; + + // Idle: the socket is open but has no data. + assert_eq!(stream.liveness(), Liveness::Clean); + + // Unsolicited data is visible to liveness without consuming it. + peer.write_all(b"x").await.unwrap(); + peer.flush().await.unwrap(); + tokio::task::yield_now().await; + assert_eq!(stream.liveness(), Liveness::DataPending); + + // The peek must not consume the byte. + let mut buf = [0u8; 1]; + stream.read_exact(&mut buf).await.unwrap(); + assert_eq!(&buf, b"x"); + + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[tokio::test] + async fn test_unix_stream_shutdown() { + let (mut stream, _peer, dir) = unix_pair("shutdown").await; + + // Graceful shutdown polls the UnixSockets shutdown arm. + stream.shutdown().await.unwrap(); + + drop(stream); + std::fs::remove_dir_all(&dir).unwrap(); + } + + // ── unix_peek edge cases ──────────────────────────────────────────────── + + /// Raw Unix socket pair (connecting end, accepted end) for direct + /// [`unix_peek`] tests. + async fn unix_stream_pair(name: &str) -> (UnixStream, UnixStream, std::path::PathBuf) { + let dir = std::env::temp_dir().join(format!("pgdog-peek-{}-{}", name, std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("test.sock"); + let listener = UnixListener::bind(&path).unwrap(); + let server = tokio::spawn(async move { + let (accepted, _) = listener.accept().await.unwrap(); + accepted + }); + let client = UnixStream::connect(&path).await.unwrap(); + (client, server.await.unwrap(), dir) + } + + #[tokio::test] + async fn test_unix_peek_no_data_is_none() { + let (client, _server, dir) = unix_stream_pair("no-data").await; + + // Idle non-blocking socket: recv returns WouldBlock → None. + let mut buf = [0u8; 1]; + assert!(unix_peek(&client, &mut buf).is_none()); + + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[tokio::test] + async fn test_unix_peek_reports_data_without_consuming() { + let (mut client, mut server, dir) = unix_stream_pair("peek").await; + + server.write_all(b"hello").await.unwrap(); + server.flush().await.unwrap(); + tokio::task::yield_now().await; + + let mut buf = [0u8; 1]; + assert!(matches!(unix_peek(&client, &mut buf), Some(Ok(1)))); + + // MSG_PEEK must not consume: a full read still gets everything. + let mut out = [0u8; 5]; + client.read_exact(&mut out).await.unwrap(); + assert_eq!(&out, b"hello"); + + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[tokio::test] + async fn test_unix_peek_limited_by_buffer_size() { + let (mut client, mut server, dir) = unix_stream_pair("limit").await; + + server.write_all(&[0u8; 100]).await.unwrap(); + server.flush().await.unwrap(); + tokio::task::yield_now().await; + + // Peek returns min(available, buffer length), not the full payload. + let mut buf = [0u8; 10]; + assert!(matches!(unix_peek(&client, &mut buf), Some(Ok(10)))); + + // And the remaining 90 bytes are still there. + let mut out = vec![0u8; 100]; + client.read_exact(&mut out).await.unwrap(); + + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[tokio::test] + async fn test_unix_peek_eof_after_peer_close() { + let (client, server, dir) = unix_stream_pair("eof").await; + drop(server); + tokio::task::yield_now().await; + + // Peer closed: recv returns 0 → Some(Ok(0)) → liveness() = Closed. + let mut buf = [0u8; 1]; + assert!(matches!(unix_peek(&client, &mut buf), Some(Ok(0)))); + + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[tokio::test] + async fn test_unix_peek_error_on_bad_fd() { + let (client, _server, dir) = unix_stream_pair("bad-fd").await; + + // Close the fd behind the stream's back: recv then fails with EBADF, + // exercising the Err branch of unix_peek. + unsafe { libc::close(client.as_raw_fd()) }; + + let mut buf = [0u8; 1]; + let result = unix_peek(&client, &mut buf); + assert!( + matches!(result, Some(Err(ref e)) if e.raw_os_error() == Some(libc::EBADF)), + "expected Some(Err(EBADF)), got {:?}", + result + ); + + // The fd is already closed; forget the stream so its Drop doesn't + // close a possibly-reused fd number. + std::mem::forget(client); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[tokio::test] + async fn test_peer_addr_conversions() { + let addr: SocketAddr = "127.0.0.1:5432".parse().unwrap(); + assert!(matches!(PeerAddr::from(Ok(addr)), PeerAddr::TCP(_))); + + let tcp_err: io::Result = Err(io::Error::new(io::ErrorKind::Other, "no peer")); + assert!(matches!(PeerAddr::from(tcp_err), PeerAddr::Empty)); + + let unix_err: io::Result = + Err(io::Error::new(io::ErrorKind::Other, "no peer")); + assert!(matches!(PeerAddr::from(unix_err), PeerAddr::Empty)); + + // A Unix socket whose peer has no pathname (e.g. the accepted end of + // an unbound connecting client) maps to Empty, not Unix. + let dir = std::env::temp_dir().join(format!("pgdog-stream-conv-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("test.sock"); + let listener = UnixListener::bind(&path).unwrap(); + let server = tokio::spawn(async move { + let (accepted, _) = listener.accept().await.unwrap(); + accepted.peer_addr().unwrap() + }); + let _client = UnixStream::connect(&path).await.unwrap(); + let unnamed = server.await.unwrap(); + assert!(unnamed.as_pathname().is_none()); + assert!(matches!(PeerAddr::from(Ok(unnamed)), PeerAddr::Empty)); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_peer_addr_debug() { + let addr: SocketAddr = "127.0.0.1:5432".parse().unwrap(); + assert_eq!(format!("{:?}", PeerAddr::TCP(addr)), "[127.0.0.1:5432]"); + assert_eq!(format!("{:?}", PeerAddr::Empty), "No address"); + + let path = PathBuf::from("/var/run/postgresql/.s.PGSQL.5432"); + assert_eq!( + format!("{:?}", PeerAddr::Unix(path)), + "/var/run/postgresql/.s.PGSQL.5432" + ); + } } From 21c45fc235536d0c26105d3378df58dc6d4fa34d Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:13:33 +0100 Subject: [PATCH 11/14] chore: cleanup --- pgdog/src/backend/pool/address.rs | 33 +++++++++++++++++++++-------- pgdog/src/backend/pool/transport.rs | 11 ++++------ pgdog/src/backend/server.rs | 17 +++++++-------- 3 files changed, 36 insertions(+), 25 deletions(-) diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 143aea924..4ccba58d2 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -12,7 +12,7 @@ use crate::backend::Error; use crate::backend::auth::{azure_workload_identity, rds_iam, vault}; use crate::backend::pool::dns_cache::DnsCache; use crate::backend::pool::token_cache::TokenCache; -use crate::backend::pool::transport::Transport; +use crate::backend::pool::transport::{Transport, unix_socket_path}; use crate::config::{Database, ServerAuth, User, config}; /// Server address. @@ -221,11 +221,19 @@ impl Address { impl std::fmt::Display for Address { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "{}@{}:{}/{}", - self.user, self.host, self.port, self.database_name - ) + match &self.host { + Transport::TCP(host) => { + write!( + f, + "{}@{}:{}/{}", + self.user, host, self.port, self.database_name + ) + } + Transport::Unix(dir) => { + let file = unix_socket_path(dir, &self.port); + write!(f, "{}", file.display()) + } + } } } @@ -260,9 +268,14 @@ impl TryFrom for Address { #[cfg(test)] mod test { - use std::time::{Duration, Instant, SystemTime}; + use std::{ + path::PathBuf, + time::{Duration, Instant, SystemTime}, + }; + + use pg_query::protobuf::Token::Path; - use crate::config; + use crate::{backend::pool::transport::unix_socket_path, config}; use super::*; @@ -730,9 +743,11 @@ mod test { #[test] fn test_unix_socket_path() { + let mut dir = PathBuf::new(); + dir.push("/tmp"); let unix = Transport::new("/tmp"); assert_eq!( - unix.unix_socket_path(&5432), + unix_socket_path(&dir, &5432), Ok(std::path::PathBuf::from("/tmp/.s.PGSQL.5432")) ); diff --git a/pgdog/src/backend/pool/transport.rs b/pgdog/src/backend/pool/transport.rs index 015adc1cf..9c78fe074 100644 --- a/pgdog/src/backend/pool/transport.rs +++ b/pgdog/src/backend/pool/transport.rs @@ -21,6 +21,10 @@ pub enum TransportError { ExpectedUnix(String), } +pub fn unix_socket_path(dir: &Path, port: &u16) -> PathBuf { + dir.join(format!(".s.PGSQL.{}", port)) +} + impl Transport { pub fn new(value: &str) -> Self { if value.starts_with('/') { @@ -30,13 +34,6 @@ impl Transport { } } - pub fn unix_socket_path(&self, port: &u16) -> Result { - match self { - Transport::TCP(host) => Err(TransportError::ExpectedUnix(host.clone())), - Transport::Unix(dir) => Ok(dir.join(format!(".s.PGSQL.{}", port))), - } - } - pub fn tcp(&self) -> Result<&str, TransportError> { match self { Transport::TCP(host) => Ok(host), diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 812d91080..3aaf89d9f 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -18,7 +18,10 @@ use super::{ }; use crate::{ auth::{md5, scram::Client}, - backend::pool::{stats::MemoryStats, transport::Transport}, + backend::pool::{ + stats::MemoryStats, + transport::{Transport, unix_socket_path}, + }, config::AuthType, frontend::ClientRequest, net::{ @@ -252,9 +255,8 @@ impl Server { } Stream::plain(tcp, config.config.memory.net_buffer) } - Transport::Unix(_) => { - let path = addr.host.unix_socket_path(&addr.port)?; - + Transport::Unix(dir) => { + let path = unix_socket_path(dir, &addr.port); debug!("connecting to Unix socket {}", path.display()); Stream::unix( UnixStream::connect(&path).await?, @@ -469,11 +471,8 @@ impl Server { let tcp = TcpStream::connect(addr.addr().await?).await?; Stream::plain(tcp, config().config.memory.net_buffer) } - Transport::Unix(_) => { - let path = addr - .host - .unix_socket_path(&addr.port) - .expect("Unix transport"); + Transport::Unix(dir) => { + let path = unix_socket_path(dir, &addr.port); let unix = UnixStream::connect(&path).await?; Stream::unix(unix, config().config.memory.net_buffer) } From 44a8ec3085930f0986d8a0876373cdfa5269f045 Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:20:41 +0100 Subject: [PATCH 12/14] chore: fix tests --- pgdog/src/backend/pool/address.rs | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 4ccba58d2..98d9735a2 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -268,12 +268,7 @@ impl TryFrom for Address { #[cfg(test)] mod test { - use std::{ - path::PathBuf, - time::{Duration, Instant, SystemTime}, - }; - - use pg_query::protobuf::Token::Path; + use std::time::{Duration, Instant, SystemTime}; use crate::{backend::pool::transport::unix_socket_path, config}; @@ -743,15 +738,16 @@ mod test { #[test] fn test_unix_socket_path() { - let mut dir = PathBuf::new(); - dir.push("/tmp"); - let unix = Transport::new("/tmp"); + let dir = std::path::PathBuf::from("/tmp"); assert_eq!( unix_socket_path(&dir, &5432), - Ok(std::path::PathBuf::from("/tmp/.s.PGSQL.5432")) + std::path::PathBuf::from("/tmp/.s.PGSQL.5432") ); - let tcp = Transport::new("127.0.0.1"); - assert!(tcp.unix_socket_path(&5432).is_err()); + // Any port formats into the socket file name. + assert_eq!( + unix_socket_path(&dir, &54321), + std::path::PathBuf::from("/tmp/.s.PGSQL.54321") + ); } } From 5920e3913a15ee0b96a2887e4b71cf209c07bbb5 Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Tue, 11 Aug 2026 03:16:37 +0100 Subject: [PATCH 13/14] feat: add into helper --- .../backend/auth/azure_workload_identity.rs | 2 +- pgdog/src/backend/auth/rds_iam.rs | 4 +-- pgdog/src/backend/auth/vault.rs | 2 +- pgdog/src/backend/pool/address.rs | 17 ++++++------ pgdog/src/backend/pool/lb/test.rs | 12 ++++----- pgdog/src/backend/pool/monitor.rs | 4 +-- pgdog/src/backend/pool/shard/role_detector.rs | 2 +- pgdog/src/backend/pool/test/mod.rs | 26 +++++++++---------- pgdog/src/backend/pool/token_cache.rs | 2 +- pgdog/src/backend/pool/transport.rs | 12 +++++++++ pgdog/src/backend/pool/waiting.rs | 2 +- 11 files changed, 49 insertions(+), 36 deletions(-) diff --git a/pgdog/src/backend/auth/azure_workload_identity.rs b/pgdog/src/backend/auth/azure_workload_identity.rs index 0f2086543..9f25dbac9 100644 --- a/pgdog/src/backend/auth/azure_workload_identity.rs +++ b/pgdog/src/backend/auth/azure_workload_identity.rs @@ -53,7 +53,7 @@ mod tests { let _azure_token_file_path = set_env_var("AZURE_FEDERATED_TOKEN_FILE", "/tmp/example"); let addr = Address { - host: Transport::new("my-awesome-db.postgres.database.azure.com"), + host: "my-awesome-db.postgres.database.azure.com".into(), port: 5432, database_name: "postgres".into(), user: "db_user".into(), diff --git a/pgdog/src/backend/auth/rds_iam.rs b/pgdog/src/backend/auth/rds_iam.rs index e0563ac41..cbef7ef2e 100644 --- a/pgdog/src/backend/auth/rds_iam.rs +++ b/pgdog/src/backend/auth/rds_iam.rs @@ -96,7 +96,7 @@ mod tests { fn make_addr() -> Address { Address { - host: Transport::new("db.cluster-abc123.us-east-1.rds.amazonaws.com"), + host: "db.cluster-abc123.us-east-1.rds.amazonaws.com".into(), port: 5432, database_name: "postgres".into(), user: "db_user".into(), @@ -171,7 +171,7 @@ mod tests { #[test] fn resolve_region_errors_when_neither_override_nor_inference() { let addr = Address { - host: Transport::new("postgres.internal.example.com"), + host: "postgres.internal.example.com".into(), port: 5432, user: "u".into(), server_iam_region: None, diff --git a/pgdog/src/backend/auth/vault.rs b/pgdog/src/backend/auth/vault.rs index ad304821f..102070281 100644 --- a/pgdog/src/backend/auth/vault.rs +++ b/pgdog/src/backend/auth/vault.rs @@ -180,7 +180,7 @@ mod tests { fn make_addr(vault_path: Option<&str>) -> Address { Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "testdb".into(), user: "testuser".into(), diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 98d9735a2..cb65bdf94 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/pgdog/src/backend/pool/address.rs @@ -84,9 +84,10 @@ impl Address { /// Create new address from config values. pub(crate) fn new(database: &Database, user: &User, database_number: usize) -> Self { let server_auth = user.server_auth; + let host = database.host.clone(); Address { - host: Transport::new(&database.host), + host: host.into(), port: database.port, database_name: if let Some(database_name) = database.database_name.clone() { database_name @@ -251,7 +252,7 @@ impl TryFrom for Address { // `Role::Auto` below. The PROBE command (the only caller) never // reads `configured_role` anyway. Ok(Self { - host: Transport::new(&host), + host: host.into(), port, passwords: vec![password.into()], user, @@ -492,7 +493,7 @@ mod test { #[tokio::test] async fn test_auth_secret_rds_iam_serves_token_from_cache() { let addr = Address { - host: Transport::new("auth-secrets-rds.internal"), + host: "auth-secrets-rds.internal".into(), port: 15432, user: "rds_user".into(), server_auth: ServerAuth::RdsIam, @@ -519,7 +520,7 @@ mod test { #[tokio::test] async fn test_auth_secret_azure_workload_identity_serves_token_from_cache() { let addr = Address { - host: Transport::new("auth-secrets-azure.internal"), + host: "auth-secrets-azure.internal".into(), port: 15433, user: "azure_user".into(), server_auth: ServerAuth::AzureWorkloadIdentity, @@ -547,7 +548,7 @@ mod test { use crate::backend::pool::token_cache::{Credentials, FetchedCredentials}; let addr = Address { - host: Transport::new("auth-secrets-vault.internal"), + host: "auth-secrets-vault.internal".into(), port: 15435, user: "configured_user".into(), server_auth: ServerAuth::VaultDynamic, @@ -617,7 +618,7 @@ mod test { #[tokio::test] async fn test_auth_credentials_vault_static_serves_password_from_cache() { let addr = Address { - host: Transport::new("auth-secrets-vault-static.internal"), + host: "auth-secrets-vault-static.internal".into(), port: 15436, user: "pgdog_static".into(), server_auth: ServerAuth::VaultStatic, @@ -674,7 +675,7 @@ mod test { // The monitor is responsible for refreshing it; auth_secrets never // blocks on a refresh. let addr = Address { - host: Transport::new("auth-secrets-stale.internal"), + host: "auth-secrets-stale.internal".into(), port: 15434, user: "stale_user".into(), server_auth: ServerAuth::RdsIam, @@ -709,7 +710,7 @@ mod test { cache.clear_cache_for_testing(); let addr = Address { - host: Transport::new(&hostname), + host: hostname.into(), port: 15432, ..Default::default() }; diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index b25ccd10a..41a292cf3 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -13,7 +13,7 @@ use monitor::Monitor; fn create_test_pool_config(host: &str, port: u16) -> PoolConfig { PoolConfig { address: Address { - host: Transport::new(&host), + host: host.into(), port, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -879,7 +879,7 @@ async fn test_monitor_unbans_all_when_all_unhealthy() { async fn test_monitor_does_not_ban_with_zero_ban_timeout() { let pool_config1 = PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -898,7 +898,7 @@ async fn test_monitor_does_not_ban_with_zero_ban_timeout() { let pool_config2 = PoolConfig { address: Address { - host: Transport::new("localhost"), + host: "localhost".into(), port: 5432, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -1507,7 +1507,7 @@ async fn test_monitor_unbans_all_when_second_target_becomes_unhealthy_after_firs fn create_test_pool_config_weighted(host: &str, port: u16, lb_weight: u8) -> PoolConfig { PoolConfig { address: Address { - host: Transport::new(&host), + host: host.into(), port, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -1928,7 +1928,7 @@ fn test_ban_check_does_not_ban_single_target() { fn test_ban_check_does_not_ban_with_zero_ban_timeout() { let pool_config1 = PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, user: "pgdog".into(), passwords: vec!["pgdog".into()], @@ -1947,7 +1947,7 @@ fn test_ban_check_does_not_ban_with_zero_ban_timeout() { let pool_config2 = PoolConfig { address: Address { - host: Transport::new("localhost"), + host: "localhost".into(), port: 5432, user: "pgdog".into(), passwords: vec!["pgdog".into()], diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index 6bfc05548..813786916 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -566,7 +566,7 @@ mod test { let pool = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 1, database_name: "pgdog".into(), user: "pgdog".into(), @@ -600,7 +600,7 @@ mod test { let pool = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), diff --git a/pgdog/src/backend/pool/shard/role_detector.rs b/pgdog/src/backend/pool/shard/role_detector.rs index cf566047b..9b3a8d7a2 100644 --- a/pgdog/src/backend/pool/shard/role_detector.rs +++ b/pgdog/src/backend/pool/shard/role_detector.rs @@ -55,7 +55,7 @@ mod test { fn create_test_pool_config(host: &str, port: u16, role_detection: bool) -> PoolConfig { PoolConfig { address: Address { - host: Transport::new(&host), + host: host.into(), port, user: "pgdog".into(), passwords: vec!["pgdog".into()], diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index a52a46424..d69b1f9e0 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -31,7 +31,7 @@ pub fn pool() -> Pool { let pool = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -56,7 +56,7 @@ pub fn pool_with_prepared_capacity(capacity: usize) -> Pool { let pool = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -371,7 +371,7 @@ async fn test_server_force_close_discards_connection() { let pool = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -542,7 +542,7 @@ async fn test_idle_healthcheck_loop() { let pool = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -590,7 +590,7 @@ async fn test_idle_healthcheck_loop_disabled_with_zero_interval() { let pool = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 1, database_name: "pgdog".into(), user: "pgdog".into(), @@ -654,7 +654,7 @@ async fn test_move_conns_to() { let source = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -667,7 +667,7 @@ async fn test_move_conns_to() { let destination = Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -867,7 +867,7 @@ fn auth_pool(passwords: Vec) -> Pool { Pool::new(&PoolConfig { address: Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), @@ -1051,7 +1051,7 @@ async fn test_token_refresh_loop_primes_cache_on_cold_start() { }; let addr = Address { - host: Transport::new("token-refresh-test.internal"), + host: "token-refresh-test.internal".into(), port: 15500, user: "refresh_user".into(), server_auth: ServerAuth::RdsIam, @@ -1089,7 +1089,7 @@ async fn test_token_refresh_loop_refreshes_before_expiry() { }; let addr = Address { - host: Transport::new("token-refresh-expiry.internal"), + host: "token-refresh-expiry.internal".into(), port: 15501, user: "refresh_user".into(), server_auth: ServerAuth::RdsIam, @@ -1134,7 +1134,7 @@ async fn test_token_refresh_loop_evicts_on_failed_refresh() { }; let addr = Address { - host: Transport::new("token-refresh-fail.internal"), + host: "token-refresh-fail.internal".into(), port: 15502, user: "refresh_user".into(), server_auth: ServerAuth::RdsIam, @@ -1179,7 +1179,7 @@ async fn test_token_refresh_loop_not_spawned_for_password_auth() { }; let addr = Address { - host: Transport::new("token-refresh-password.internal"), + host: "token-refresh-password.internal".into(), port: 15503, user: "refresh_user".into(), server_auth: ServerAuth::Password, @@ -1219,7 +1219,7 @@ async fn test_token_refresh_loop_stops_on_shutdown() { }; let addr = Address { - host: Transport::new("token-refresh-shutdown.internal"), + host: "token-refresh-shutdown.internal".into(), port: 15504, user: "refresh_user".into(), server_auth: ServerAuth::AzureWorkloadIdentity, diff --git a/pgdog/src/backend/pool/token_cache.rs b/pgdog/src/backend/pool/token_cache.rs index 68dc654d6..309c2f745 100644 --- a/pgdog/src/backend/pool/token_cache.rs +++ b/pgdog/src/backend/pool/token_cache.rs @@ -288,7 +288,7 @@ mod tests { /// in the global cache. fn addr(port: u16) -> Address { Address { - host: Transport::new("token-cache-test.internal"), + host: "token-cache-test.internal".into(), port, user: "test_user".into(), ..Default::default() diff --git a/pgdog/src/backend/pool/transport.rs b/pgdog/src/backend/pool/transport.rs index 9c78fe074..31c683e3b 100644 --- a/pgdog/src/backend/pool/transport.rs +++ b/pgdog/src/backend/pool/transport.rs @@ -57,3 +57,15 @@ impl Display for Transport { } } } + +impl From<&str> for Transport { + fn from(value: &str) -> Self { + Transport::new(value) + } +} + +impl From for Transport { + fn from(value: String) -> Self { + Transport::new(&value) + } +} diff --git a/pgdog/src/backend/pool/waiting.rs b/pgdog/src/backend/pool/waiting.rs index 819845ed0..6480e1bc4 100644 --- a/pgdog/src/backend/pool/waiting.rs +++ b/pgdog/src/backend/pool/waiting.rs @@ -153,7 +153,7 @@ mod tests { let pool = Pool::new(&crate::backend::pool::PoolConfig { address: crate::backend::pool::Address { - host: Transport::new("127.0.0.1"), + host: "127.0.0.1".into(), port: 5432, database_name: "pgdog".into(), user: "pgdog".into(), From cbb7b244d9c8a2ae8caba8f394c1292b05211e64 Mon Sep 17 00:00:00 2001 From: Hamza Hamud <53880692+hhamud@users.noreply.github.com> Date: Tue, 11 Aug 2026 03:32:54 +0100 Subject: [PATCH 14/14] fix: reduce single allocation --- pgdog/src/backend/server.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 3aaf89d9f..e4e866fda 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -268,7 +268,7 @@ impl Server { let tls_mode = config.config.general.tls_verify; // Only attempt TLS if not in Disabled mode and its not connecting to a unix socket - if tls_mode != TlsVerifyMode::Disabled && addr.host.tcp().is_ok() { + if tls_mode != TlsVerifyMode::Disabled && matches!(addr.host, Transport::TCP(_)) { debug!( "requesting TLS connection with verify mode: {:?} [{}]", tls_mode, addr,