Finish convenience API

tokio-postgres/Cargo.toml

@@ -28,7 +28,7 @@ circle-ci = { repository = "sfackler/rust-postgres" }
 
 [features]
 default = ["runtime"]
-runtime = ["tokio-tcp", "tokio-uds"]
+runtime = ["tokio-tcp", "tokio-uds", "futures-cpupool", "lazy_static"]
 
 "with-bit-vec-0.5" = ["bit-vec-05"]
 "with-chrono-0.4" = ["chrono-04"]
@@ -42,7 +42,6 @@ antidote = "1.0"
 bytes = "0.4"
 fallible-iterator = "0.1.6"
 futures = "0.1.7"
-futures-cpupool = "0.1"
 log = "0.4"
 phf = "0.7.23"
 postgres-protocol = { version = "0.3.0", path = "../postgres-protocol" }
@@ -52,6 +51,8 @@ tokio-io = "0.1"
 void = "1.0"
 
 tokio-tcp = { version = "0.1", optional = true }
+futures-cpupool = { version = "0.1", optional = true }
+lazy_static = { version = "1.0", optional = true }
 
 bit-vec-05 = { version = "0.5", package = "bit-vec", optional = true }
 chrono-04 = { version = "0.4", package = "chrono", optional = true }

tokio-postgres/src/builder.rs

@@ -3,7 +3,11 @@ use std::iter;
 use std::str::{self, FromStr};
 use tokio_io::{AsyncRead, AsyncWrite};
 
+#[cfg(feature = "runtime")]
+use crate::proto::ConnectFuture;
 use crate::proto::HandshakeFuture;
+#[cfg(feature = "runtime")]
+use crate::{Connect, Socket};
 use crate::{Error, Handshake, TlsMode};
 
 #[derive(Clone)]
@@ -55,6 +59,14 @@ impl Builder {
     {
         Handshake(HandshakeFuture::new(stream, tls_mode, self.params.clone()))
     }
+
+    #[cfg(feature = "runtime")]
+    pub fn connect<T>(&self, tls_mode: T) -> Connect<T>
+    where
+        T: TlsMode<Socket>,
+    {
+        Connect(ConnectFuture::new(tls_mode, self.params.clone()))
+    }
 }
 
 impl FromStr for Builder {

tokio-postgres/src/error/mod.rs

@@ -5,6 +5,8 @@ use postgres_protocol::message::backend::{ErrorFields, ErrorResponseBody};
 use std::error::{self, Error as _Error};
 use std::fmt;
 use std::io;
+#[cfg(feature = "runtime")]
+use std::num::ParseIntError;
 
 pub use self::sqlstate::*;
 
@@ -346,6 +348,11 @@ enum Kind {
     UnsupportedAuthentication,
     Authentication,
     ConnectionSyntax,
+    Connect,
+    #[cfg(feature = "runtime")]
+    MissingHost,
+    #[cfg(feature = "runtime")]
+    InvalidPort,
 }
 
 struct ErrorInner {
@@ -383,6 +390,11 @@ impl fmt::Display for Error {
             Kind::UnsupportedAuthentication => "unsupported authentication method requested",
             Kind::Authentication => "authentication error",
             Kind::ConnectionSyntax => "invalid connection string",
+            Kind::Connect => "error connecting to server",
+            #[cfg(feature = "runtime")]
+            Kind::MissingHost => "host not provided",
+            #[cfg(feature = "runtime")]
+            Kind::InvalidPort => "invalid port",
         };
         fmt.write_str(s)?;
         if let Some(ref cause) = self.0.cause {
@@ -485,4 +497,19 @@ impl Error {
     pub(crate) fn connection_syntax(e: Box<dyn error::Error + Sync + Send>) -> Error {
         Error::new(Kind::ConnectionSyntax, Some(e))
     }
+
+    #[cfg(feature = "runtime")]
+    pub(crate) fn connect(e: io::Error) -> Error {
+        Error::new(Kind::Connect, Some(Box::new(e)))
+    }
+
+    #[cfg(feature = "runtime")]
+    pub(crate) fn missing_host() -> Error {
+        Error::new(Kind::MissingHost, None)
+    }
+
+    #[cfg(feature = "runtime")]
+    pub(crate) fn invalid_port(e: ParseIntError) -> Error {
+        Error::new(Kind::InvalidPort, Some(Box::new(e)))
+    }
 }

tokio-postgres/src/lib.rs

@@ -180,6 +180,27 @@ where
     }
 }
 
+#[cfg(feature = "runtime")]
+#[must_use = "futures do nothing unless polled"]
+pub struct Connect<T>(proto::ConnectFuture<T>)
+where
+    T: TlsMode<Socket>;
+
+#[cfg(feature = "runtime")]
+impl<T> Future for Connect<T>
+where
+    T: TlsMode<Socket>,
+{
+    type Item = (Client, Connection<T::Stream>);
+    type Error = Error;
+
+    fn poll(&mut self) -> Poll<(Client, Connection<T::Stream>), Error> {
+        let (client, connection) = try_ready!(self.0.poll());
+
+        Ok(Async::Ready((Client(client), Connection(connection))))
+    }
+}
+
 #[must_use = "futures do nothing unless polled"]
 pub struct Prepare(proto::PrepareFuture);
 

tokio-postgres/src/proto/connect.rs

@@ -0,0 +1,177 @@
+use futures::{try_ready, Async, Future, Poll};
+use futures_cpupool::{CpuFuture, CpuPool};
+use lazy_static::lazy_static;
+use state_machine_future::{transition, RentToOwn, StateMachineFuture};
+use std::collections::HashMap;
+use std::io;
+use std::net::{SocketAddr, ToSocketAddrs};
+#[cfg(unix)]
+use std::path::Path;
+use std::vec;
+use tokio_tcp::TcpStream;
+#[cfg(unix)]
+use tokio_uds::UnixStream;
+
+use crate::proto::{Client, Connection, HandshakeFuture};
+use crate::{Error, Socket, TlsMode};
+
+lazy_static! {
+    static ref DNS_POOL: CpuPool = futures_cpupool::Builder::new()
+        .name_prefix("postgres-dns-")
+        .pool_size(2)
+        .create();
+}
+
+#[derive(StateMachineFuture)]
+pub enum Connect<T>
+where
+    T: TlsMode<Socket>,
+{
+    #[state_machine_future(start)]
+    #[cfg_attr(unix, state_machine_future(transitions(ConnectingUnix, ResolvingDns)))]
+    #[cfg_attr(not(unix), state_machine_future(transitions(ConnectingTcp)))]
+    Start {
+        tls_mode: T,
+        params: HashMap<String, String>,
+    },
+    #[cfg(unix)]
+    #[state_machine_future(transitions(Handshaking))]
+    ConnectingUnix {
+        future: tokio_uds::ConnectFuture,
+        tls_mode: T,
+        params: HashMap<String, String>,
+    },
+    #[state_machine_future(transitions(ConnectingTcp))]
+    ResolvingDns {
+        future: CpuFuture<vec::IntoIter<SocketAddr>, io::Error>,
+        tls_mode: T,
+        params: HashMap<String, String>,
+    },
+    #[state_machine_future(transitions(Handshaking))]
+    ConnectingTcp {
+        future: tokio_tcp::ConnectFuture,
+        addrs: vec::IntoIter<SocketAddr>,
+        tls_mode: T,
+        params: HashMap<String, String>,
+    },
+    #[state_machine_future(transitions(Finished))]
+    Handshaking { future: HandshakeFuture<Socket, T> },
+    #[state_machine_future(ready)]
+    Finished((Client, Connection<T::Stream>)),
+    #[state_machine_future(error)]
+    Failed(Error),
+}
+
+impl<T> PollConnect<T> for Connect<T>
+where
+    T: TlsMode<Socket>,
+{
+    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<T>>) -> Poll<AfterStart<T>, Error> {
+        let mut state = state.take();
+
+        let host = match state.params.remove("host") {
+            Some(host) => host,
+            None => return Err(Error::missing_host()),
+        };
+
+        let port = match state.params.remove("port") {
+            Some(port) => port.parse::<u16>().map_err(Error::invalid_port)?,
+            None => 5432,
+        };
+
+        #[cfg(unix)]
+        {
+            if host.starts_with('/') {
+                let path = Path::new(&host).join(format!(".s.PGSQL.{}", port));
+                transition!(ConnectingUnix {
+                    future: UnixStream::connect(path),
+                    tls_mode: state.tls_mode,
+                    params: state.params,
+                })
+            }
+        }
+
+        transition!(ResolvingDns {
+            future: DNS_POOL.spawn_fn(move || (&*host, port).to_socket_addrs()),
+            tls_mode: state.tls_mode,
+            params: state.params,
+        })
+    }
+
+    #[cfg(unix)]
+    fn poll_connecting_unix<'a>(
+        state: &'a mut RentToOwn<'a, ConnectingUnix<T>>,
+    ) -> Poll<AfterConnectingUnix<T>, Error> {
+        let stream = try_ready!(state.future.poll().map_err(Error::connect));
+        let stream = Socket::new_unix(stream);
+        let state = state.take();
+
+        transition!(Handshaking {
+            future: HandshakeFuture::new(stream, state.tls_mode, state.params)
+        })
+    }
+
+    fn poll_resolving_dns<'a>(
+        state: &'a mut RentToOwn<'a, ResolvingDns<T>>,
+    ) -> Poll<AfterResolvingDns<T>, Error> {
+        let mut addrs = try_ready!(state.future.poll().map_err(Error::connect));
+        let state = state.take();
+
+        let addr = match addrs.next() {
+            Some(addr) => addr,
+            None => {
+                return Err(Error::connect(io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    "resolved 0 addresses",
+                )))
+            }
+        };
+
+        transition!(ConnectingTcp {
+            future: TcpStream::connect(&addr),
+            addrs,
+            tls_mode: state.tls_mode,
+            params: state.params,
+        })
+    }
+
+    fn poll_connecting_tcp<'a>(
+        state: &'a mut RentToOwn<'a, ConnectingTcp<T>>,
+    ) -> Poll<AfterConnectingTcp<T>, Error> {
+        let stream = loop {
+            match state.future.poll() {
+                Ok(Async::Ready(stream)) => break Socket::new_tcp(stream),
+                Ok(Async::NotReady) => return Ok(Async::NotReady),
+                Err(e) => {
+                    let addr = match state.addrs.next() {
+                        Some(addr) => addr,
+                        None => return Err(Error::connect(e)),
+                    };
+                    state.future = TcpStream::connect(&addr);
+                }
+            }
+        };
+        let state = state.take();
+
+        transition!(Handshaking {
+            future: HandshakeFuture::new(stream, state.tls_mode, state.params),
+        })
+    }
+
+    fn poll_handshaking<'a>(
+        state: &'a mut RentToOwn<'a, Handshaking<T>>,
+    ) -> Poll<AfterHandshaking<T>, Error> {
+        let r = try_ready!(state.future.poll());
+
+        transition!(Finished(r))
+    }
+}
+
+impl<T> ConnectFuture<T>
+where
+    T: TlsMode<Socket>,
+{
+    pub fn new(tls_mode: T, params: HashMap<String, String>) -> ConnectFuture<T> {
+        Connect::start(tls_mode, params)
+    }
+}

tokio-postgres/src/proto/mod.rs

@@ -22,6 +22,8 @@ mod bind;
 mod cancel;
 mod client;
 mod codec;
+#[cfg(feature = "runtime")]
+mod connect;
 mod connection;
 mod copy_in;
 mod copy_out;
@@ -42,6 +44,8 @@ pub use crate::proto::bind::BindFuture;
 pub use crate::proto::cancel::CancelFuture;
 pub use crate::proto::client::Client;
 pub use crate::proto::codec::PostgresCodec;
+#[cfg(feature = "runtime")]
+pub use crate::proto::connect::ConnectFuture;
 pub use crate::proto::connection::Connection;
 pub use crate::proto::copy_in::CopyInFuture;
 pub use crate::proto::copy_out::CopyOutStream;

tokio-postgres/tests/test/main.rs

@@ -14,6 +14,8 @@ use tokio_postgres::types::{Kind, Type};
 use tokio_postgres::{AsyncMessage, Client, Connection, NoTls};
 
 mod parse;
+#[cfg(feature = "runtime")]
+mod runtime;
 mod types;
 
 fn connect(

tokio-postgres/tests/test/runtime.rs

@@ -0,0 +1,34 @@
+use futures::Future;
+use tokio::runtime::current_thread::Runtime;
+use tokio_postgres::{Client, Connection, Error, NoTls, Socket};
+
+fn connect(s: &str) -> impl Future<Item = (Client, Connection<Socket>), Error = Error> {
+    s.parse::<tokio_postgres::Builder>().unwrap().connect(NoTls)
+}
+
+#[test]
+#[ignore] // FIXME doesn't work with our docker-based tests :(
+fn unix_socket() {
+    let mut runtime = Runtime::new().unwrap();
+
+    let connect = connect("host=/var/run/postgresql port=5433 user=postgres");
+    let (mut client, connection) = runtime.block_on(connect).unwrap();
+    let connection = connection.map_err(|e| panic!("{}", e));
+    runtime.spawn(connection);
+
+    let execute = client.batch_execute("SELECT 1");
+    runtime.block_on(execute).unwrap();
+}
+
+#[test]
+fn tcp() {
+    let mut runtime = Runtime::new().unwrap();
+
+    let connect = connect("host=localhost port=5433 user=postgres");
+    let (mut client, connection) = runtime.block_on(connect).unwrap();
+    let connection = connection.map_err(|e| panic!("{}", e));
+    runtime.spawn(connection);
+
+    let execute = client.batch_execute("SELECT 1");
+    runtime.block_on(execute).unwrap();
+}