Add bind/portal

Author: sfackler

tokio-postgres/src/lib.rs

@@ -45,10 +45,14 @@ pub mod error;
 mod proto;
 pub mod tls;
 
-static NEXT_STATEMENT_ID: AtomicUsize = AtomicUsize::new(0);
-
 fn next_statement() -> String {
-    format!("s{}", NEXT_STATEMENT_ID.fetch_add(1, Ordering::SeqCst))
+    static ID: AtomicUsize = AtomicUsize::new(0);
+    format!("s{}", ID.fetch_add(1, Ordering::SeqCst))
+}
+
+fn next_portal() -> String {
+    static ID: AtomicUsize = AtomicUsize::new(0);
+    format!("p{}", ID.fetch_add(1, Ordering::SeqCst))
 }
 
 pub enum TlsMode {
@@ -84,6 +88,10 @@ impl Client {
         Query(self.0.query(&statement.0, params))
     }
 
+    pub fn bind(&mut self, statement: &Statement, params: &[&ToSql]) -> Bind {
+        Bind(self.0.bind(&statement.0, next_portal(), params))
+    }
+
     pub fn copy_in<S>(&mut self, statement: &Statement, params: &[&ToSql], stream: S) -> CopyIn<S>
     where
         S: Stream,
@@ -226,6 +234,24 @@ impl Stream for Query {
     }
 }
 
+#[must_use = "futures do nothing unless polled"]
+pub struct Bind(proto::BindFuture);
+
+impl Future for Bind {
+    type Item = Portal;
+    type Error = Error;
+
+    fn poll(&mut self) -> Poll<Portal, Error> {
+        match self.0.poll() {
+            Ok(Async::Ready(portal)) => Ok(Async::Ready(Portal(portal))),
+            Ok(Async::NotReady) => Ok(Async::NotReady),
+            Err(e) => Err(e),
+        }
+    }
+}
+
+pub struct Portal(proto::Portal);
+
 #[must_use = "futures do nothing unless polled"]
 pub struct CopyIn<S>(proto::CopyInFuture<S>)
 where

tokio-postgres/src/proto/bind.rs

@@ -0,0 +1,72 @@
+use futures::sync::mpsc;
+use futures::{Poll, Stream};
+use postgres_protocol::message::backend::Message;
+use proto::client::{Client, PendingRequest};
+use proto::portal::Portal;
+use proto::statement::Statement;
+use state_machine_future::RentToOwn;
+use Error;
+
+#[derive(StateMachineFuture)]
+pub enum Bind {
+    #[state_machine_future(start, transitions(ReadBindComplete))]
+    Start {
+        client: Client,
+        request: PendingRequest,
+        name: String,
+        statement: Statement,
+    },
+    #[state_machine_future(transitions(Finished))]
+    ReadBindComplete {
+        receiver: mpsc::Receiver<Message>,
+        client: Client,
+        name: String,
+        statement: Statement,
+    },
+    #[state_machine_future(ready)]
+    Finished(Portal),
+    #[state_machine_future(error)]
+    Failed(Error),
+}
+
+impl PollBind for Bind {
+    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
+        let state = state.take();
+        let receiver = state.client.send(state.request)?;
+
+        transition!(ReadBindComplete {
+            receiver,
+            client: state.client,
+            name: state.name,
+            statement: state.statement,
+        })
+    }
+
+    fn poll_read_bind_complete<'a>(
+        state: &'a mut RentToOwn<'a, ReadBindComplete>,
+    ) -> Poll<AfterReadBindComplete, Error> {
+        let message = try_ready_receive!(state.receiver.poll());
+        let state = state.take();
+
+        match message {
+            Some(Message::BindComplete) => transition!(Finished(Portal::new(
+                state.client.downgrade(),
+                state.name,
+                state.statement,
+            ))),
+            Some(_) => Err(Error::unexpected_message()),
+            None => Err(Error::closed()),
+        }
+    }
+}
+
+impl BindFuture {
+    pub fn new(
+        client: Client,
+        request: PendingRequest,
+        name: String,
+        statement: Statement,
+    ) -> BindFuture {
+        Bind::start(client, request, name, statement)
+    }
+}

tokio-postgres/src/proto/client.rs

@@ -8,6 +8,7 @@ use std::collections::HashMap;
 use std::error::Error as StdError;
 use std::sync::{Arc, Weak};
 
+use proto::bind::BindFuture;
 use proto::connection::{Request, RequestMessages};
 use proto::copy_in::{CopyInFuture, CopyInReceiver, CopyMessage};
 use proto::copy_out::CopyOutStream;
@@ -140,6 +141,15 @@ impl Client {
         QueryStream::new(self.clone(), pending, statement.clone())
     }
 
+    pub fn bind(&self, statement: &Statement, name: String, params: &[&ToSql]) -> BindFuture {
+        let mut buf = self.bind_message(statement, &name, params);
+        if let Ok(ref mut buf) = buf {
+            frontend::sync(buf);
+        }
+        let pending = PendingRequest(buf.map(RequestMessages::Single));
+        BindFuture::new(self.clone(), pending, name, statement.clone())
+    }
+
     pub fn copy_in<S>(&self, statement: &Statement, params: &[&ToSql], stream: S) -> CopyInFuture<S>
     where
         S: Stream,
@@ -169,8 +179,16 @@ impl Client {
     }
 
     pub fn close_statement(&self, name: &str) {
+        self.close(b'S', name)
+    }
+
+    pub fn close_portal(&self, name: &str) {
+        self.close(b'P', name)
+    }
+
+    fn close(&self, ty: u8, name: &str) {
         let mut buf = vec![];
-        frontend::close(b'S', name, &mut buf).expect("statement name not valid");
+        frontend::close(ty, name, &mut buf).expect("statement name not valid");
         frontend::sync(&mut buf);
         let (sender, _) = mpsc::channel(0);
         let _ = self.0.sender.unbounded_send(Request {
@@ -179,10 +197,15 @@ impl Client {
         });
     }
 
-    fn excecute_message(&self, statement: &Statement, params: &[&ToSql]) -> Result<Vec<u8>, Error> {
+    fn bind_message(
+        &self,
+        statement: &Statement,
+        name: &str,
+        params: &[&ToSql],
+    ) -> Result<Vec<u8>, Error> {
         let mut buf = vec![];
         let r = frontend::bind(
-            "",
+            name,
             statement.name(),
             Some(1),
             params.iter().zip(statement.params()),
@@ -195,10 +218,14 @@ impl Client {
             &mut buf,
         );
         match r {
-            Ok(()) => {}
+            Ok(()) => Ok(buf),
             Err(frontend::BindError::Conversion(e)) => return Err(Error::to_sql(e)),
             Err(frontend::BindError::Serialization(e)) => return Err(Error::encode(e)),
         }
+    }
+
+    fn excecute_message(&self, statement: &Statement, params: &[&ToSql]) -> Result<Vec<u8>, Error> {
+        let mut buf = self.bind_message(statement, "", params)?;
         frontend::execute("", 0, &mut buf).map_err(Error::parse)?;
         frontend::sync(&mut buf);
         Ok(buf)

tokio-postgres/src/proto/mod.rs

@@ -18,6 +18,7 @@ macro_rules! try_ready_closed {
     };
 }
 
+mod bind;
 mod cancel;
 mod client;
 mod codec;
@@ -27,6 +28,7 @@ mod copy_in;
 mod copy_out;
 mod execute;
 mod handshake;
+mod portal;
 mod prepare;
 mod query;
 mod row;
@@ -38,6 +40,7 @@ mod typeinfo;
 mod typeinfo_composite;
 mod typeinfo_enum;
 
+pub use proto::bind::BindFuture;
 pub use proto::cancel::CancelFuture;
 pub use proto::client::Client;
 pub use proto::codec::PostgresCodec;
@@ -46,6 +49,7 @@ pub use proto::copy_in::CopyInFuture;
 pub use proto::copy_out::CopyOutStream;
 pub use proto::execute::ExecuteFuture;
 pub use proto::handshake::HandshakeFuture;
+pub use proto::portal::Portal;
 pub use proto::prepare::PrepareFuture;
 pub use proto::query::QueryStream;
 pub use proto::row::Row;

tokio-postgres/src/proto/portal.rs

@@ -0,0 +1,38 @@
+use std::sync::Arc;
+
+use proto::client::WeakClient;
+use proto::statement::Statement;
+
+struct Inner {
+    client: WeakClient,
+    name: String,
+    statement: Statement,
+}
+
+pub struct Portal(Arc<Inner>);
+
+impl Drop for Portal {
+    fn drop(&mut self) {
+        if let Some(client) = self.0.client.upgrade() {
+            client.close_portal(&self.0.name);
+        }
+    }
+}
+
+impl Portal {
+    pub fn new(client: WeakClient, name: String, statement: Statement) -> Portal {
+        Portal(Arc::new(Inner {
+            client,
+            name,
+            statement,
+        }))
+    }
+
+    pub fn name(&self) -> &str {
+        &self.0.name
+    }
+
+    pub fn statement(&self) -> &Statement {
+        &self.0.statement
+    }
+}