A less stringy builder

This allows us to support things like non-utf8 passwords and unix socket
directories.

Author: sfackler

postgres/src/builder.rs

@@ -1,7 +1,9 @@
 use futures::sync::oneshot;
 use futures::Future;
 use log::error;
+use std::path::Path;
 use std::str::FromStr;
+use std::time::Duration;
 use tokio_postgres::{Error, MakeTlsMode, Socket, TlsMode};
 
 use crate::{Client, RUNTIME};
@@ -19,11 +21,43 @@ impl Builder {
         Builder(tokio_postgres::Builder::new())
     }
 
+    pub fn host(&mut self, host: &str) -> &mut Builder {
+        self.0.host(host);
+        self
+    }
+
+    #[cfg(unix)]
+    pub fn host_path<T>(&mut self, host: T) -> &mut Builder
+    where
+        T: AsRef<Path>,
+    {
+        self.0.host_path(host);
+        self
+    }
+
+    pub fn port(&mut self, port: u16) -> &mut Builder {
+        self.0.port(port);
+        self
+    }
+
     pub fn param(&mut self, key: &str, value: &str) -> &mut Builder {
         self.0.param(key, value);
         self
     }
 
+    pub fn connect_timeout(&mut self, connect_timeout: Duration) -> &mut Builder {
+        self.0.connect_timeout(connect_timeout);
+        self
+    }
+
+    pub fn password<T>(&mut self, password: T) -> &mut Builder
+    where
+        T: AsRef<[u8]>,
+    {
+        self.0.password(password);
+        self
+    }
+
     pub fn connect<T>(&self, tls_mode: T) -> Result<Client, Error>
     where
         T: MakeTlsMode<Socket> + 'static + Send,

tokio-postgres-native-tls/src/test.rs

@@ -6,13 +6,15 @@ use tokio_postgres::{self, PreferTls, RequireTls, TlsMode};
 
 use crate::TlsConnector;
 
-fn smoke_test<T>(builder: &tokio_postgres::Builder, tls: T)
+fn smoke_test<T>(s: &str, tls: T)
 where
     T: TlsMode<TcpStream>,
     T::Stream: 'static,
 {
     let mut runtime = Runtime::new().unwrap();
 
+    let builder = s.parse::<tokio_postgres::Builder>().unwrap();
+
     let handshake = TcpStream::connect(&"127.0.0.1:5433".parse().unwrap())
         .map_err(|e| panic!("{}", e))
         .and_then(|s| builder.handshake(s, tls));
@@ -42,9 +44,7 @@ fn require() {
         .build()
         .unwrap();
     smoke_test(
-        tokio_postgres::Builder::new()
-            .user("ssl_user")
-            .dbname("postgres"),
+        "user=ssl_user dbname=postgres",
         RequireTls(TlsConnector::with_connector(connector, "localhost")),
     );
 }
@@ -58,9 +58,7 @@ fn prefer() {
         .build()
         .unwrap();
     smoke_test(
-        tokio_postgres::Builder::new()
-            .user("ssl_user")
-            .dbname("postgres"),
+        "user=ssl_user dbname=postgres",
         PreferTls(TlsConnector::with_connector(connector, "localhost")),
     );
 }
@@ -74,10 +72,7 @@ fn scram_user() {
         .build()
         .unwrap();
     smoke_test(
-        tokio_postgres::Builder::new()
-            .user("scram_user")
-            .password("password")
-            .dbname("postgres"),
+        "user=scram_user password=password dbname=postgres",
         RequireTls(TlsConnector::with_connector(connector, "localhost")),
     );
 }

tokio-postgres-openssl/src/test.rs

@@ -6,13 +6,15 @@ use tokio_postgres::{self, PreferTls, RequireTls, TlsMode};
 
 use super::*;
 
-fn smoke_test<T>(builder: &tokio_postgres::Builder, tls: T)
+fn smoke_test<T>(s: &str, tls: T)
 where
     T: TlsMode<TcpStream>,
     T::Stream: 'static,
 {
     let mut runtime = Runtime::new().unwrap();
 
+    let builder = s.parse::<tokio_postgres::Builder>().unwrap();
+
     let handshake = TcpStream::connect(&"127.0.0.1:5433".parse().unwrap())
         .map_err(|e| panic!("{}", e))
         .and_then(|s| builder.handshake(s, tls));
@@ -39,9 +41,7 @@ fn require() {
     builder.set_ca_file("../test/server.crt").unwrap();
     let ctx = builder.build();
     smoke_test(
-        tokio_postgres::Builder::new()
-            .user("ssl_user")
-            .dbname("postgres"),
+        "user=ssl_user dbname=postgres",
         RequireTls(TlsConnector::new(ctx.configure().unwrap(), "localhost")),
     );
 }
@@ -52,9 +52,7 @@ fn prefer() {
     builder.set_ca_file("../test/server.crt").unwrap();
     let ctx = builder.build();
     smoke_test(
-        tokio_postgres::Builder::new()
-            .user("ssl_user")
-            .dbname("postgres"),
+        "user=ssl_user dbname=postgres",
         PreferTls(TlsConnector::new(ctx.configure().unwrap(), "localhost")),
     );
 }
@@ -65,10 +63,7 @@ fn scram_user() {
     builder.set_ca_file("../test/server.crt").unwrap();
     let ctx = builder.build();
     smoke_test(
-        tokio_postgres::Builder::new()
-            .user("scram_user")
-            .password("password")
-            .dbname("postgres"),
+        "user=scram_user password=password dbname=postgres",
         RequireTls(TlsConnector::new(ctx.configure().unwrap(), "localhost")),
     );
 }

tokio-postgres/src/builder.rs

@@ -1,6 +1,10 @@
 use std::collections::hash_map::{self, HashMap};
 use std::iter;
+#[cfg(all(feature = "runtime", unix))]
+use std::path::{Path, PathBuf};
 use std::str::{self, FromStr};
+#[cfg(feature = "runtime")]
+use std::time::Duration;
 use tokio_io::{AsyncRead, AsyncWrite};
 
 #[cfg(feature = "runtime")]
@@ -10,9 +14,24 @@ use crate::proto::HandshakeFuture;
 use crate::{Connect, MakeTlsMode, Socket};
 use crate::{Error, Handshake, TlsMode};
 
-#[derive(Clone)]
+#[cfg(feature = "runtime")]
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) enum Host {
+    Tcp(String),
+    #[cfg(unix)]
+    Unix(PathBuf),
+}
+
+#[derive(Debug, Clone, PartialEq)]
 pub struct Builder {
-    params: HashMap<String, String>,
+    pub(crate) params: HashMap<String, String>,
+    pub(crate) password: Option<Vec<u8>>,
+    #[cfg(feature = "runtime")]
+    pub(crate) host: Vec<Host>,
+    #[cfg(feature = "runtime")]
+    pub(crate) port: Vec<u16>,
+    #[cfg(feature = "runtime")]
+    pub(crate) connect_timeout: Option<Duration>,
 }
 
 impl Default for Builder {
@@ -27,45 +46,80 @@ impl Builder {
         params.insert("client_encoding".to_string(), "UTF8".to_string());
         params.insert("timezone".to_string(), "GMT".to_string());
 
-        Builder { params }
+        Builder {
+            params,
+            password: None,
+            #[cfg(feature = "runtime")]
+            host: vec![],
+            #[cfg(feature = "runtime")]
+            port: vec![],
+            #[cfg(feature = "runtime")]
+            connect_timeout: None,
+        }
     }
 
-    pub fn user(&mut self, user: &str) -> &mut Builder {
-        self.param("user", user)
+    #[cfg(feature = "runtime")]
+    pub fn host(&mut self, host: &str) -> &mut Builder {
+        #[cfg(unix)]
+        {
+            if host.starts_with('/') {
+                self.host.push(Host::Unix(PathBuf::from(host)));
+                return self;
+            }
+        }
+
+        self.host.push(Host::Tcp(host.to_string()));
+        self
     }
 
-    pub fn dbname(&mut self, database: &str) -> &mut Builder {
-        self.param("dbname", database)
+    #[cfg(all(feature = "runtime", unix))]
+    pub fn host_path<T>(&mut self, host: T) -> &mut Builder
+    where
+        T: AsRef<Path>,
+    {
+        self.host.push(Host::Unix(host.as_ref().to_path_buf()));
+        self
     }
 
-    pub fn password(&mut self, password: &str) -> &mut Builder {
-        self.param("password", password)
+    #[cfg(feature = "runtime")]
+    pub fn port(&mut self, port: u16) -> &mut Builder {
+        self.port.push(port);
+        self
     }
 
-    pub fn param(&mut self, key: &str, value: &str) -> &mut Builder {
-        self.params.insert(key.to_string(), value.to_string());
+    #[cfg(feature = "runtime")]
+    pub fn connect_timeout(&mut self, connect_timeout: Duration) -> &mut Builder {
+        self.connect_timeout = Some(connect_timeout);
         self
     }
 
-    /// FIXME do we want this?
-    pub fn iter(&self) -> Iter<'_> {
-        Iter(self.params.iter())
+    pub fn password<T>(&mut self, password: T) -> &mut Builder
+    where
+        T: AsRef<[u8]>,
+    {
+        self.password = Some(password.as_ref().to_vec());
+        self
+    }
+
+    pub fn param(&mut self, key: &str, value: &str) -> &mut Builder {
+        self.params.insert(key.to_string(), value.to_string());
+        self
     }
 
     pub fn handshake<S, T>(&self, stream: S, tls_mode: T) -> Handshake<S, T>
     where
         S: AsyncRead + AsyncWrite,
         T: TlsMode<S>,
     {
-        Handshake(HandshakeFuture::new(stream, tls_mode, self.params.clone()))
+        Handshake(HandshakeFuture::new(stream, tls_mode, self.clone()))
     }
 
     #[cfg(feature = "runtime")]
     pub fn connect<T>(&self, make_tls_mode: T) -> Connect<T>
     where
         T: MakeTlsMode<Socket>,
     {
-        Connect(ConnectFuture::new(make_tls_mode, self.params.clone()))
+        Connect(ConnectFuture::new(make_tls_mode, self.clone()))
     }
 }
 
@@ -77,7 +131,40 @@ impl FromStr for Builder {
         let mut builder = Builder::new();
 
         while let Some((key, value)) = parser.parameter()? {
-            builder.params.insert(key.to_string(), value);
+            match key {
+                "password" => {
+                    builder.password(value);
+                }
+                #[cfg(feature = "runtime")]
+                "host" => {
+                    for host in value.split(',') {
+                        builder.host(host);
+                    }
+                }
+                #[cfg(feature = "runtime")]
+                "port" => {
+                    for port in value.split(',') {
+                        let port = if port.is_empty() {
+                            5432
+                        } else {
+                            port.parse().map_err(Error::invalid_port)?
+                        };
+                        builder.port(port);
+                    }
+                }
+                #[cfg(feature = "runtime")]
+                "connect_timeout" => {
+                    let timeout = value
+                        .parse::<i64>()
+                        .map_err(Error::invalid_connect_timeout)?;
+                    if timeout > 0 {
+                        builder.connect_timeout(Duration::from_secs(timeout as u64));
+                    }
+                }
+                key => {
+                    builder.param(key, &value);
+                }
+            }
         }
 
         Ok(builder)