Add execute

Author: sfackler

tokio-postgres/src/lib.rs

@@ -33,7 +33,7 @@ pub use postgres_shared::{CancelData, Notification};
 
 use error::Error;
 use params::ConnectParams;
-use types::Type;
+use types::{ToSql, Type};
 
 mod proto;
 
@@ -68,6 +68,10 @@ impl Client {
         let name = format!("s{}", NEXT_STATEMENT_ID.fetch_add(1, Ordering::SeqCst));
         Prepare(self.0.prepare(name, query, param_types))
     }
+
+    pub fn execute(&mut self, statement: &Statement, params: &[&ToSql]) -> Execute {
+        Execute(self.0.execute(&statement.0, params))
+    }
 }
 
 #[must_use = "futures do nothing unless polled"]
@@ -131,3 +135,15 @@ impl Statement {
         self.0.columns()
     }
 }
+
+#[must_use = "futures do nothing unless polled"]
+pub struct Execute(proto::ExecuteFuture);
+
+impl Future for Execute {
+    type Item = u64;
+    type Error = Error;
+
+    fn poll(&mut self) -> Poll<u64, Error> {
+        self.0.poll()
+    }
+}

tokio-postgres/src/proto/client.rs

@@ -1,26 +1,27 @@
 use futures::sync::mpsc;
+use postgres_protocol;
 use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
 
 use disconnected;
-use error::Error;
+use error::{self, Error};
 use proto::connection::Request;
+use proto::execute::ExecuteFuture;
 use proto::prepare::PrepareFuture;
-use types::Type;
+use proto::statement::Statement;
+use types::{IsNull, ToSql, Type};
 
 pub struct PendingRequest {
     sender: mpsc::UnboundedSender<Request>,
-    messages: Vec<u8>,
+    messages: Result<Vec<u8>, Error>,
 }
 
 impl PendingRequest {
     pub fn send(self) -> Result<mpsc::Receiver<Message>, Error> {
+        let messages = self.messages?;
         let (sender, receiver) = mpsc::channel(0);
         self.sender
-            .unbounded_send(Request {
-                messages: self.messages,
-                sender,
-            })
+            .unbounded_send(Request { messages, sender })
             .map(|_| receiver)
             .map_err(|_| disconnected())
     }
@@ -36,16 +37,52 @@ impl Client {
     }
 
     pub fn prepare(&mut self, name: String, query: &str, param_types: &[Type]) -> PrepareFuture {
+        let pending = self.pending(|buf| {
+            frontend::parse(&name, query, param_types.iter().map(|t| t.oid()), buf)?;
+            frontend::describe(b'S', &name, buf)?;
+            frontend::sync(buf);
+            Ok(())
+        });
+
+        PrepareFuture::new(pending, self.sender.clone(), name)
+    }
+
+    pub fn execute(&mut self, statement: &Statement, params: &[&ToSql]) -> ExecuteFuture {
+        let pending = self.pending(|buf| {
+            let r = frontend::bind(
+                "",
+                statement.name(),
+                Some(1),
+                params.iter().zip(statement.params()),
+                |(param, ty), buf| match param.to_sql_checked(ty, buf) {
+                    Ok(IsNull::No) => Ok(postgres_protocol::IsNull::No),
+                    Ok(IsNull::Yes) => Ok(postgres_protocol::IsNull::Yes),
+                    Err(e) => Err(e),
+                },
+                Some(1),
+                buf,
+            );
+            match r {
+                Ok(()) => {}
+                Err(frontend::BindError::Conversion(e)) => return Err(error::conversion(e)),
+                Err(frontend::BindError::Serialization(e)) => return Err(Error::from(e)),
+            }
+            frontend::execute("", 0, buf)?;
+            frontend::sync(buf);
+            Ok(())
+        });
+
+        ExecuteFuture::new(pending, statement.clone())
+    }
+
+    fn pending<F>(&self, messages: F) -> PendingRequest
+    where
+        F: FnOnce(&mut Vec<u8>) -> Result<(), Error>,
+    {
         let mut buf = vec![];
-        let request = frontend::parse(&name, query, param_types.iter().map(|t| t.oid()), &mut buf)
-            .and_then(|()| frontend::describe(b'S', &name, &mut buf))
-            .and_then(|()| Ok(frontend::sync(&mut buf)))
-            .map(|()| PendingRequest {
-                sender: self.sender.clone(),
-                messages: buf,
-            })
-            .map_err(Into::into);
-
-        PrepareFuture::new(request, self.sender.clone(), name)
+        PendingRequest {
+            sender: self.sender.clone(),
+            messages: messages(&mut buf).map(|()| buf),
+        }
     }
 }

tokio-postgres/src/proto/execute.rs

@@ -0,0 +1,88 @@
+use futures::sync::mpsc;
+use futures::{Poll, Stream};
+use postgres_protocol::message::backend::Message;
+use state_machine_future::RentToOwn;
+
+use error::{self, Error};
+use proto::client::PendingRequest;
+use proto::statement::Statement;
+use {bad_response, disconnected};
+
+#[derive(StateMachineFuture)]
+pub enum Execute {
+    #[state_machine_future(start, transitions(ReadResponse))]
+    Start {
+        request: PendingRequest,
+        statement: Statement,
+    },
+    #[state_machine_future(transitions(ReadReadyForQuery))]
+    ReadResponse { receiver: mpsc::Receiver<Message> },
+    #[state_machine_future(transitions(Finished))]
+    ReadReadyForQuery {
+        receiver: mpsc::Receiver<Message>,
+        rows: u64,
+    },
+    #[state_machine_future(ready)]
+    Finished(u64),
+    #[state_machine_future(error)]
+    Failed(Error),
+}
+
+impl PollExecute for Execute {
+    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
+        let state = state.take();
+        let receiver = state.request.send()?;
+
+        // the statement can drop after this point, since its close will queue up after the execution
+        transition!(ReadResponse { receiver })
+    }
+
+    fn poll_read_response<'a>(
+        state: &'a mut RentToOwn<'a, ReadResponse>,
+    ) -> Poll<AfterReadResponse, Error> {
+        loop {
+            let message = try_receive!(state.receiver.poll());
+
+            match message {
+                Some(Message::BindComplete) => {}
+                Some(Message::DataRow(_)) => {}
+                Some(Message::ErrorResponse(body)) => return Err(error::__db(body)),
+                Some(Message::CommandComplete(body)) => {
+                    let rows = body.tag()?.rsplit(' ').next().unwrap().parse().unwrap_or(0);
+                    let state = state.take();
+                    transition!(ReadReadyForQuery {
+                        receiver: state.receiver,
+                        rows,
+                    });
+                }
+                Some(Message::EmptyQueryResponse) => {
+                    let state = state.take();
+                    transition!(ReadReadyForQuery {
+                        receiver: state.receiver,
+                        rows: 0,
+                    });
+                }
+                Some(_) => return Err(bad_response()),
+                None => return Err(disconnected()),
+            }
+        }
+    }
+
+    fn poll_read_ready_for_query<'a>(
+        state: &'a mut RentToOwn<'a, ReadReadyForQuery>,
+    ) -> Poll<AfterReadReadyForQuery, Error> {
+        let message = try_receive!(state.receiver.poll());
+
+        match message {
+            Some(Message::ReadyForQuery(_)) => transition!(Finished(state.rows)),
+            Some(_) => Err(bad_response()),
+            None => Err(disconnected()),
+        }
+    }
+}
+
+impl ExecuteFuture {
+    pub fn new(request: PendingRequest, statement: Statement) -> ExecuteFuture {
+        Execute::start(request, statement)
+    }
+}

tokio-postgres/src/proto/mod.rs

@@ -11,6 +11,7 @@ macro_rules! try_receive {
 mod client;
 mod codec;
 mod connection;
+mod execute;
 mod handshake;
 mod prepare;
 mod socket;
@@ -19,6 +20,7 @@ mod statement;
 pub use proto::client::Client;
 pub use proto::codec::PostgresCodec;
 pub use proto::connection::Connection;
+pub use proto::execute::ExecuteFuture;
 pub use proto::handshake::HandshakeFuture;
 pub use proto::prepare::PrepareFuture;
 pub use proto::socket::Socket;

tokio-postgres/src/proto/prepare.rs

@@ -16,7 +16,7 @@ use {bad_response, disconnected};
 pub enum Prepare {
     #[state_machine_future(start, transitions(ReadParseComplete))]
     Start {
-        request: Result<PendingRequest, Error>,
+        request: PendingRequest,
         sender: mpsc::UnboundedSender<Request>,
         name: String,
     },
@@ -56,7 +56,7 @@ pub enum Prepare {
 impl PollPrepare for Prepare {
     fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
         let state = state.take();
-        let receiver = state.request?.send()?;
+        let receiver = state.request.send()?;
 
         transition!(ReadParseComplete {
             sender: state.sender,
@@ -160,7 +160,7 @@ impl PollPrepare for Prepare {
 
 impl PrepareFuture {
     pub fn new(
-        request: Result<PendingRequest, Error>,
+        request: PendingRequest,
         sender: mpsc::UnboundedSender<Request>,
         name: String,
     ) -> PrepareFuture {

tokio-postgres/src/proto/statement.rs

@@ -1,18 +1,19 @@
 use futures::sync::mpsc;
 use postgres_protocol::message::frontend;
 use postgres_shared::stmt::Column;
+use std::sync::Arc;
 
 use proto::connection::Request;
 use types::Type;
 
-pub struct Statement {
+pub struct StatementInner {
     sender: mpsc::UnboundedSender<Request>,
     name: String,
     params: Vec<Type>,
     columns: Vec<Column>,
 }
 
-impl Drop for Statement {
+impl Drop for StatementInner {
     fn drop(&mut self) {
         let mut buf = vec![];
         frontend::close(b'S', &self.name, &mut buf).expect("statement name not valid");
@@ -25,26 +26,33 @@ impl Drop for Statement {
     }
 }
 
+#[derive(Clone)]
+pub struct Statement(Arc<StatementInner>);
+
 impl Statement {
     pub fn new(
         sender: mpsc::UnboundedSender<Request>,
         name: String,
         params: Vec<Type>,
         columns: Vec<Column>,
     ) -> Statement {
-        Statement {
+        Statement(Arc::new(StatementInner {
             sender,
             name,
             params,
             columns,
-        }
+        }))
+    }
+
+    pub fn name(&self) -> &str {
+        &self.0.name
     }
 
     pub fn params(&self) -> &[Type] {
-        &self.params
+        &self.0.params
     }
 
     pub fn columns(&self) -> &[Column] {
-        &self.columns
+        &self.0.columns
     }
 }