From 919012d0c9292369eed828af4811e7cc6274b27f Mon Sep 17 00:00:00 2001 From: Steven Fackler Date: Mon, 17 Dec 2018 21:25:21 -0800 Subject: [PATCH] Finish convenience API --- tokio-postgres/Cargo.toml | 5 +- tokio-postgres/src/builder.rs | 12 ++ tokio-postgres/src/error/mod.rs | 27 ++++ tokio-postgres/src/lib.rs | 21 ++++ tokio-postgres/src/proto/connect.rs | 177 +++++++++++++++++++++++++++ tokio-postgres/src/proto/mod.rs | 4 + tokio-postgres/tests/test/main.rs | 2 + tokio-postgres/tests/test/runtime.rs | 34 +++++ 8 files changed, 280 insertions(+), 2 deletions(-) create mode 100644 tokio-postgres/src/proto/connect.rs create mode 100644 tokio-postgres/tests/test/runtime.rs diff --git a/tokio-postgres/Cargo.toml b/tokio-postgres/Cargo.toml index c317c5d6..349d4cdc 100644 --- a/tokio-postgres/Cargo.toml +++ b/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 } diff --git a/tokio-postgres/src/builder.rs b/tokio-postgres/src/builder.rs index 85981e9a..3090f09a 100644 --- a/tokio-postgres/src/builder.rs +++ b/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(&self, tls_mode: T) -> Connect + where + T: TlsMode, + { + Connect(ConnectFuture::new(tls_mode, self.params.clone())) + } } impl FromStr for Builder { diff --git a/tokio-postgres/src/error/mod.rs b/tokio-postgres/src/error/mod.rs index 13a8149b..d83c50dd 100644 --- a/tokio-postgres/src/error/mod.rs +++ b/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) -> 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))) + } } diff --git a/tokio-postgres/src/lib.rs b/tokio-postgres/src/lib.rs index c9d77d1c..d2bfec1a 100644 --- a/tokio-postgres/src/lib.rs +++ b/tokio-postgres/src/lib.rs @@ -180,6 +180,27 @@ where } } +#[cfg(feature = "runtime")] +#[must_use = "futures do nothing unless polled"] +pub struct Connect(proto::ConnectFuture) +where + T: TlsMode; + +#[cfg(feature = "runtime")] +impl Future for Connect +where + T: TlsMode, +{ + type Item = (Client, Connection); + type Error = Error; + + fn poll(&mut self) -> Poll<(Client, Connection), 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); diff --git a/tokio-postgres/src/proto/connect.rs b/tokio-postgres/src/proto/connect.rs new file mode 100644 index 00000000..ad21e9df --- /dev/null +++ b/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 +where + T: TlsMode, +{ + #[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, + }, + #[cfg(unix)] + #[state_machine_future(transitions(Handshaking))] + ConnectingUnix { + future: tokio_uds::ConnectFuture, + tls_mode: T, + params: HashMap, + }, + #[state_machine_future(transitions(ConnectingTcp))] + ResolvingDns { + future: CpuFuture, io::Error>, + tls_mode: T, + params: HashMap, + }, + #[state_machine_future(transitions(Handshaking))] + ConnectingTcp { + future: tokio_tcp::ConnectFuture, + addrs: vec::IntoIter, + tls_mode: T, + params: HashMap, + }, + #[state_machine_future(transitions(Finished))] + Handshaking { future: HandshakeFuture }, + #[state_machine_future(ready)] + Finished((Client, Connection)), + #[state_machine_future(error)] + Failed(Error), +} + +impl PollConnect for Connect +where + T: TlsMode, +{ + fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll, 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::().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>, + ) -> Poll, 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>, + ) -> Poll, 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>, + ) -> Poll, 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>, + ) -> Poll, Error> { + let r = try_ready!(state.future.poll()); + + transition!(Finished(r)) + } +} + +impl ConnectFuture +where + T: TlsMode, +{ + pub fn new(tls_mode: T, params: HashMap) -> ConnectFuture { + Connect::start(tls_mode, params) + } +} diff --git a/tokio-postgres/src/proto/mod.rs b/tokio-postgres/src/proto/mod.rs index 9d19fa0e..3a13e6bc 100644 --- a/tokio-postgres/src/proto/mod.rs +++ b/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; diff --git a/tokio-postgres/tests/test/main.rs b/tokio-postgres/tests/test/main.rs index ad1736f0..41918e65 100644 --- a/tokio-postgres/tests/test/main.rs +++ b/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( diff --git a/tokio-postgres/tests/test/runtime.rs b/tokio-postgres/tests/test/runtime.rs new file mode 100644 index 00000000..f723be7b --- /dev/null +++ b/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), Error = Error> { + s.parse::().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(); +}