Add unix domain socket support for Postgres

sqlx-core/Cargo.toml

@@ -22,7 +22,7 @@ all-type = ["bigdecimal", "json", "time", "chrono", "ipnetwork", "uuid"]
 # we need a feature which activates `num-bigint` as well because
 # `bigdecimal` uses types from it but does not reexport (tsk tsk)
 bigdecimal = ["bigdecimal_", "num-bigint"]
-postgres = [ "md-5", "sha2", "base64", "sha-1", "rand", "hmac", "futures-channel/sink", "futures-util/sink" ]
+postgres = [ "md-5", "sha2", "base64", "sha-1", "rand", "hmac", "futures-channel/sink", "futures-util/sink", "tokio/uds" ]
 json = ["serde", "serde_json"]
 mysql = [ "sha-1", "sha2", "generic-array", "num-bigint", "base64", "digest", "rand" ]
 sqlite = [ "libsqlite3-sys" ]

sqlx-core/src/io/tls.rs

@@ -4,7 +4,6 @@ use std::pin::Pin;
 use std::task::{Context, Poll};
 
 use crate::runtime::{AsyncRead, AsyncWrite, TcpStream};
-use crate::url::Url;
 
 use self::Inner::*;
 
@@ -14,15 +13,24 @@ pub struct MaybeTlsStream {
 
 enum Inner {
     NotTls(TcpStream),
+    #[cfg(all(feature = "postgres", unix))]
+    UnixStream(crate::runtime::UnixStream),
     #[cfg(feature = "tls")]
     Tls(async_native_tls::TlsStream<TcpStream>),
     #[cfg(feature = "tls")]
     Upgrading,
 }
 
 impl MaybeTlsStream {
-    pub async fn connect(url: &Url, default_port: u16) -> crate::Result<Self> {
-        let conn = TcpStream::connect((url.host(), url.port(default_port))).await?;
+    #[cfg(all(feature = "postgres", unix))]
+    pub async fn connect_uds<S: AsRef<std::ffi::OsStr>>(p: S) -> crate::Result<Self> {
+        let conn = crate::runtime::UnixStream::connect(p.as_ref()).await?;
+        Ok(Self {
+            inner: Inner::UnixStream(conn),
+        })
+    }
+    pub async fn connect(host: &str, port: u16) -> crate::Result<Self> {
+        let conn = TcpStream::connect((host, port)).await?;
         Ok(Self {
             inner: Inner::NotTls(conn),
         })
@@ -32,6 +40,8 @@ impl MaybeTlsStream {
     pub fn is_tls(&self) -> bool {
         match self.inner {
             Inner::NotTls(_) => false,
+            #[cfg(all(feature = "postgres", unix))]
+            Inner::UnixStream(_) => false,
             #[cfg(feature = "tls")]
             Inner::Tls(_) => true,
             #[cfg(feature = "tls")]
@@ -43,23 +53,29 @@ impl MaybeTlsStream {
     #[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
     pub async fn upgrade(
         &mut self,
-        url: &Url,
+        host: &str,
         connector: async_native_tls::TlsConnector,
     ) -> crate::Result<()> {
         let conn = match std::mem::replace(&mut self.inner, Upgrading) {
             NotTls(conn) => conn,
+            #[cfg(all(feature = "postgres", unix))]
+            UnixStream(_) => {
+                return Err(tls_err!("TLS is not supported with unix domain sockets").into())
+            }
             Tls(_) => return Err(tls_err!("connection already upgraded").into()),
             Upgrading => return Err(tls_err!("connection already failed to upgrade").into()),
         };
 
-        self.inner = Tls(connector.connect(url.host(), conn).await?);
+        self.inner = Tls(connector.connect(host, conn).await?);
 
         Ok(())
     }
 
     pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
         match self.inner {
             NotTls(ref conn) => conn.shutdown(how),
+            #[cfg(all(feature = "postgres", unix))]
+            UnixStream(ref conn) => conn.shutdown(how),
             #[cfg(feature = "tls")]
             Tls(ref conn) => conn.get_ref().shutdown(how),
             #[cfg(feature = "tls")]
@@ -73,6 +89,8 @@ macro_rules! forward_pin (
     ($self:ident.$method:ident($($arg:ident),*)) => (
         match &mut $self.inner {
             NotTls(ref mut conn) => Pin::new(conn).$method($($arg),*),
+            #[cfg(all(feature = "postgres", unix))]
+            UnixStream(ref mut conn) => Pin::new(conn).$method($($arg),*),
             #[cfg(feature = "tls")]
             Tls(ref mut conn) => Pin::new(conn).$method($($arg),*),
             #[cfg(feature = "tls")]

sqlx-core/src/mysql/stream.rs

@@ -34,7 +34,9 @@ pub(crate) struct MySqlStream {
 
 impl MySqlStream {
     pub(super) async fn new(url: &Url) -> crate::Result<Self> {
-        let stream = MaybeTlsStream::connect(&url, 3306).await?;
+        let host = url.host().unwrap_or("localhost");
+        let port = url.port(3306);
+        let stream = MaybeTlsStream::connect(host, port).await?;
 
         let mut capabilities = Capabilities::PROTOCOL_41
             | Capabilities::IGNORE_SPACE

sqlx-core/src/mysql/tls.rs

@@ -116,5 +116,8 @@ async fn try_upgrade(
         )
         .await?;
 
-    stream.stream.upgrade(url, connector).await
+    stream
+        .stream
+        .upgrade(url.host().unwrap_or("localhost"), connector)
+        .await
 }

sqlx-core/src/postgres/connection.rs

@@ -110,9 +110,13 @@ pub struct PgConnection {
 
 // https://www.postgresql.org/docs/12/protocol-flow.html#id-1.10.5.7.3
 async fn startup(stream: &mut PgStream, url: &Url) -> crate::Result<BackendKeyData> {
-    // Defaults to postgres@.../postgres
-    let username = url.username().unwrap_or(Cow::Borrowed("postgres"));
-    let database = url.database().unwrap_or("postgres");
+    // Defaults to $USER@.../$USER
+    // and falls back to postgres@.../postgres
+    let username = url
+        .username()
+        .or_else(|| std::env::var("USER").map(Cow::Owned).ok())
+        .unwrap_or(Cow::Borrowed("postgres"));
+    let database = url.database().unwrap_or(&username);
 
     // See this doc for more runtime parameters
     // https://www.postgresql.org/docs/12/runtime-config-client.html

sqlx-core/src/postgres/stream.rs

@@ -23,7 +23,27 @@ pub struct PgStream {
 
 impl PgStream {
     pub(super) async fn new(url: &Url) -> crate::Result<Self> {
-        let stream = MaybeTlsStream::connect(&url, 5432).await?;
+        let host = url.host();
+        let port = url.port(5432);
+        #[cfg(unix)]
+        let stream = {
+            let host = host
+                .map(|host| {
+                    percent_encoding::percent_decode_str(host)
+                        .decode_utf8()
+                        .expect("percent-encoded hostname contained non-UTF-8 bytes")
+                })
+                .or_else(|| url.param("host"))
+                .unwrap_or("/var/run/postgresql".into());
+            if host.starts_with("/") {
+                let path = format!("{}/.s.PGSQL.{}", host, port);
+                MaybeTlsStream::connect_uds(&path).await?
+            } else {
+                MaybeTlsStream::connect(&host, port).await?
+            }
+        };
+        #[cfg(not(unix))]
+        let stream = MaybeTlsStream::connect(host.unwrap_or("localhost"), port).await?;
 
         Ok(Self {
             notifications: None,

sqlx-core/src/postgres/tls.rs

@@ -69,6 +69,7 @@ async fn try_upgrade(
     accept_invalid_host_names: bool,
 ) -> crate::Result<bool> {
     use async_native_tls::TlsConnector;
+    use std::borrow::Cow;
 
     stream.write(crate::postgres::protocol::SslRequest);
     stream.flush().await?;
@@ -105,7 +106,12 @@ async fn try_upgrade(
         }
     }
 
-    stream.stream.upgrade(url, connector).await?;
+    let host = url
+        .host()
+        .map(Cow::Borrowed)
+        .or_else(|| url.param("host"))
+        .unwrap_or("localhost".into());
+    stream.stream.upgrade(&host, connector).await?;
 
     Ok(true)
 }

sqlx-core/src/runtime.rs

@@ -17,6 +17,9 @@ pub(crate) use async_std::{
     task::spawn,
 };
 
+#[cfg(all(feature = "runtime-async-std", feature = "postgres", unix))]
+pub(crate) use async_std::os::unix::net::UnixStream;
+
 #[cfg(feature = "runtime-tokio")]
 pub(crate) use tokio::{
     fs,
@@ -26,3 +29,6 @@ pub(crate) use tokio::{
     time::delay_for as sleep,
     time::timeout,
 };
+
+#[cfg(all(feature = "runtime-tokio", feature = "postgres", unix))]
+pub(crate) use tokio::net::UnixStream;

sqlx-core/src/url.rs

@@ -34,13 +34,10 @@ impl Url {
         self.0.as_str()
     }
 
-    pub fn host(&self) -> &str {
-        let host = self.0.host_str();
-
-        match host {
-            Some(host) if !host.is_empty() => host,
-
-            _ => "localhost",
+    pub fn host(&self) -> Option<&str> {
+        match self.0.host_str()? {
+            "" => None,
+            host => Some(host),
         }
     }