From 2115d02cb036817fc2f5bc3c70b9e52e9f2693ef Mon Sep 17 00:00:00 2001
From: Julius de Bruijn <julius+github@nauk.io>
Date: Fri, 26 Jun 2020 10:54:52 +0200
Subject: [PATCH] Move pg-specific socket options to its options.

Makes tcp connections explicit.
---
 sqlx-core/src/mysql/connection/stream.rs    |  2 +-
 sqlx-core/src/net/socket.rs                 | 14 +-------------
 sqlx-core/src/postgres/connection/stream.rs |  9 +++------
 sqlx-core/src/postgres/options.rs           | 16 ++++++++++++++++
 4 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/sqlx-core/src/mysql/connection/stream.rs b/sqlx-core/src/mysql/connection/stream.rs
index e9a9317f2..ed7106a59 100644
--- a/sqlx-core/src/mysql/connection/stream.rs
+++ b/sqlx-core/src/mysql/connection/stream.rs
@@ -32,7 +32,7 @@ impl MySqlStream {
     pub(super) async fn connect(options: &MySqlConnectOptions) -> Result<Self, Error> {
         let socket = match options.socket {
             Some(ref path) => Socket::connect_uds(path).await?,
-            None => Socket::connect(&options.host, options.port).await?,
+            None => Socket::connect_tcp(&options.host, options.port).await?,
         };
 
         let mut capabilities = Capabilities::PROTOCOL_41
diff --git a/sqlx-core/src/net/socket.rs b/sqlx-core/src/net/socket.rs
index 7850738fa..e8caab09a 100644
--- a/sqlx-core/src/net/socket.rs
+++ b/sqlx-core/src/net/socket.rs
@@ -17,22 +17,10 @@ pub enum Socket {
 }
 
 impl Socket {
-    #[cfg(not(unix))]
-    pub async fn connect(host: &str, port: u16) -> io::Result<Self> {
+    pub async fn connect_tcp(host: &str, port: u16) -> io::Result<Self> {
         TcpStream::connect((host, port)).await.map(Socket::Tcp)
     }
 
-    #[cfg(unix)]
-    pub async fn connect(host: &str, port: u16) -> io::Result<Self> {
-        if host.starts_with('/') {
-            // if the host starts with a forward slash, assume that this is a request
-            // to connect to a local socket
-            Self::connect_uds(&format!("{}/.s.PGSQL.{}", host, port)).await
-        } else {
-            TcpStream::connect((host, port)).await.map(Socket::Tcp)
-        }
-    }
-
     #[cfg(unix)]
     pub async fn connect_uds(path: impl AsRef<Path>) -> io::Result<Self> {
         sqlx_rt::UnixStream::connect(path.as_ref())
diff --git a/sqlx-core/src/postgres/connection/stream.rs b/sqlx-core/src/postgres/connection/stream.rs
index ea35d1f68..024d95582 100644
--- a/sqlx-core/src/postgres/connection/stream.rs
+++ b/sqlx-core/src/postgres/connection/stream.rs
@@ -31,12 +31,9 @@ pub struct PgStream {
 
 impl PgStream {
     pub(super) async fn connect(options: &PgConnectOptions) -> Result<Self, Error> {
-        let socket = match options.socket {
-            Some(ref path) => {
-                Socket::connect_uds(&format!("{}/.s.PGSQL.{}", path.display(), options.port))
-                    .await?
-            }
-            None => Socket::connect(&options.host, options.port).await?,
+        let socket = match options.fetch_socket() {
+            Some(ref path) => Socket::connect_uds(path).await?,
+            None => Socket::connect_tcp(&options.host, options.port).await?,
         };
 
         let inner = BufStream::new(MaybeTlsStream::Raw(socket));
diff --git a/sqlx-core/src/postgres/options.rs b/sqlx-core/src/postgres/options.rs
index ca355c7f9..1dcfa09b4 100644
--- a/sqlx-core/src/postgres/options.rs
+++ b/sqlx-core/src/postgres/options.rs
@@ -320,6 +320,22 @@ impl PgConnectOptions {
         self.statement_cache_capacity = capacity;
         self
     }
+
+    /// We try using a socket if hostname starts with `/` or if socket parameter
+    /// is specified.
+    pub(crate) fn fetch_socket(&self) -> Option<String> {
+        match self.socket {
+            Some(ref socket) => {
+                let full_path = format!("{}/.s.PGSQL.{}", socket.display(), self.port);
+                Some(full_path)
+            }
+            None if self.host.starts_with('/') => {
+                let full_path = format!("{}/.s.PGSQL.{}", self.host, self.port);
+                Some(full_path)
+            }
+            _ => None,
+        }
+    }
 }
 
 fn default_host(port: u16) -> String {