Add query_portal

Author: sfackler

tokio-postgres/src/lib.rs

@@ -92,6 +92,10 @@ impl Client {
         Bind(self.0.bind(&statement.0, next_portal(), params))
     }
 
+    pub fn query_portal(&mut self, portal: &Portal, max_rows: i32) -> QueryPortal {
+        QueryPortal(self.0.query_portal(&portal.0, max_rows))
+    }
+
     pub fn copy_in<S>(&mut self, statement: &Statement, params: &[&ToSql], stream: S) -> CopyIn<S>
     where
         S: Stream,
@@ -218,7 +222,7 @@ impl Future for Execute {
 }
 
 #[must_use = "streams do nothing unless polled"]
-pub struct Query(proto::QueryStream);
+pub struct Query(proto::QueryStream<proto::Statement>);
 
 impl Stream for Query {
     type Item = Row;
@@ -250,6 +254,23 @@ impl Future for Bind {
     }
 }
 
+#[must_use = "streams do nothing unless polled"]
+pub struct QueryPortal(proto::QueryStream<proto::Portal>);
+
+impl Stream for QueryPortal {
+    type Item = Row;
+    type Error = Error;
+
+    fn poll(&mut self) -> Poll<Option<Row>, Error> {
+        match self.0.poll() {
+            Ok(Async::Ready(Some(row))) => Ok(Async::Ready(Some(Row(row)))),
+            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
+            Ok(Async::NotReady) => Ok(Async::NotReady),
+            Err(e) => Err(e),
+        }
+    }
+}
+
 pub struct Portal(proto::Portal);
 
 #[must_use = "futures do nothing unless polled"]

tokio-postgres/src/proto/client.rs

@@ -13,6 +13,7 @@ use proto::connection::{Request, RequestMessages};
 use proto::copy_in::{CopyInFuture, CopyInReceiver, CopyMessage};
 use proto::copy_out::CopyOutStream;
 use proto::execute::ExecuteFuture;
+use proto::portal::Portal;
 use proto::prepare::PrepareFuture;
 use proto::query::QueryStream;
 use proto::simple_query::SimpleQueryFuture;
@@ -133,7 +134,7 @@ impl Client {
         ExecuteFuture::new(self.clone(), pending, statement.clone())
     }
 
-    pub fn query(&self, statement: &Statement, params: &[&ToSql]) -> QueryStream {
+    pub fn query(&self, statement: &Statement, params: &[&ToSql]) -> QueryStream<Statement> {
         let pending = PendingRequest(
             self.excecute_message(statement, params)
                 .map(RequestMessages::Single),
@@ -150,6 +151,15 @@ impl Client {
         BindFuture::new(self.clone(), pending, name, statement.clone())
     }
 
+    pub fn query_portal(&self, portal: &Portal, rows: i32) -> QueryStream<Portal> {
+        let pending = self.pending(|buf| {
+            frontend::execute(portal.name(), rows, buf).map_err(Error::parse)?;
+            frontend::sync(buf);
+            Ok(())
+        });
+        QueryStream::new(self.clone(), pending, portal.clone())
+    }
+
     pub fn copy_in<S>(&self, statement: &Statement, params: &[&ToSql], stream: S) -> CopyInFuture<S>
     where
         S: Stream,

tokio-postgres/src/proto/portal.rs

@@ -9,16 +9,17 @@ struct Inner {
     statement: Statement,
 }
 
-pub struct Portal(Arc<Inner>);
-
-impl Drop for Portal {
+impl Drop for Inner {
     fn drop(&mut self) {
-        if let Some(client) = self.0.client.upgrade() {
-            client.close_portal(&self.0.name);
+        if let Some(client) = self.client.upgrade() {
+            client.close_portal(&self.name);
         }
     }
 }
 
+#[derive(Clone)]
+pub struct Portal(Arc<Inner>);
+
 impl Portal {
     pub fn new(client: WeakClient, name: String, statement: Statement) -> Portal {
         Portal(Arc::new(Inner {

tokio-postgres/src/proto/query.rs

@@ -4,26 +4,46 @@ use postgres_protocol::message::backend::Message;
 use std::mem;
 
 use proto::client::{Client, PendingRequest};
+use proto::portal::Portal;
 use proto::row::Row;
 use proto::statement::Statement;
 use Error;
 
-enum State {
+pub trait StatementHolder {
+    fn statement(&self) -> &Statement;
+}
+
+impl StatementHolder for Statement {
+    fn statement(&self) -> &Statement {
+        self
+    }
+}
+
+impl StatementHolder for Portal {
+    fn statement(&self) -> &Statement {
+        self.statement()
+    }
+}
+
+enum State<T> {
     Start {
         client: Client,
         request: PendingRequest,
-        statement: Statement,
+        statement: T,
     },
     ReadingResponse {
         receiver: mpsc::Receiver<Message>,
-        statement: Statement,
+        statement: T,
     },
     Done,
 }
 
-pub struct QueryStream(State);
+pub struct QueryStream<T>(State<T>);
 
-impl Stream for QueryStream {
+impl<T> Stream for QueryStream<T>
+where
+    T: StatementHolder,
+{
     type Item = Row;
     type Error = Error;
 
@@ -66,14 +86,16 @@ impl Stream for QueryStream {
                         }
                         Some(Message::ErrorResponse(body)) => break Err(Error::db(body)),
                         Some(Message::DataRow(body)) => {
-                            let row = Row::new(statement.clone(), body)?;
+                            let row = Row::new(statement.statement().clone(), body)?;
                             self.0 = State::ReadingResponse {
                                 receiver,
                                 statement,
                             };
                             break Ok(Async::Ready(Some(row)));
                         }
-                        Some(Message::EmptyQueryResponse) | Some(Message::CommandComplete(_)) => {
+                        Some(Message::EmptyQueryResponse)
+                        | Some(Message::PortalSuspended)
+                        | Some(Message::CommandComplete(_)) => {
                             break Ok(Async::Ready(None));
                         }
                         Some(_) => break Err(Error::unexpected_message()),
@@ -86,8 +108,11 @@ impl Stream for QueryStream {
     }
 }
 
-impl QueryStream {
-    pub fn new(client: Client, request: PendingRequest, statement: Statement) -> QueryStream {
+impl<T> QueryStream<T>
+where
+    T: StatementHolder,
+{
+    pub fn new(client: Client, request: PendingRequest, statement: T) -> QueryStream<T> {
         QueryStream(State::Start {
             client,
             request,

tokio-postgres/src/proto/typeinfo.rs

@@ -7,6 +7,7 @@ use next_statement;
 use proto::client::Client;
 use proto::prepare::PrepareFuture;
 use proto::query::QueryStream;
+use proto::statement::Statement;
 use proto::typeinfo_composite::TypeinfoCompositeFuture;
 use proto::typeinfo_enum::TypeinfoEnumFuture;
 use types::{Kind, Oid, Type};
@@ -55,7 +56,7 @@ pub enum Typeinfo {
         QueryingRangeSubtype
     ))]
     QueryingTypeinfo {
-        future: stream::Collect<QueryStream>,
+        future: stream::Collect<QueryStream<Statement>>,
         oid: Oid,
         client: Client,
     },

tokio-postgres/src/proto/typeinfo_composite.rs

@@ -9,6 +9,7 @@ use next_statement;
 use proto::client::Client;
 use proto::prepare::PrepareFuture;
 use proto::query::QueryStream;
+use proto::statement::Statement;
 use proto::typeinfo::TypeinfoFuture;
 use types::{Field, Oid};
 
@@ -36,7 +37,7 @@ pub enum TypeinfoComposite {
     },
     #[state_machine_future(transitions(QueryingCompositeFieldTypes, Finished))]
     QueryingCompositeFields {
-        future: stream::Collect<QueryStream>,
+        future: stream::Collect<QueryStream<Statement>>,
         client: Client,
     },
     #[state_machine_future(transitions(Finished))]

tokio-postgres/src/proto/typeinfo_enum.rs

@@ -7,6 +7,7 @@ use next_statement;
 use proto::client::Client;
 use proto::prepare::PrepareFuture;
 use proto::query::QueryStream;
+use proto::statement::Statement;
 use types::Oid;
 
 const TYPEINFO_ENUM_QUERY: &'static str = "
@@ -45,7 +46,7 @@ pub enum TypeinfoEnum {
     },
     #[state_machine_future(transitions(Finished))]
     QueryingEnumVariants {
-        future: stream::Collect<QueryStream>,
+        future: stream::Collect<QueryStream<Statement>>,
         client: Client,
     },
     #[state_machine_future(ready)]

tokio-postgres/tests/test.rs

@@ -198,6 +198,49 @@ fn insert_select() {
     runtime.block_on(tests).unwrap();
 }
 
+#[test]
+fn query_portal() {
+    let _ = env_logger::try_init();
+    let mut runtime = Runtime::new().unwrap();
+
+    let handshake = tokio_postgres::connect(
+        "postgres://postgres@localhost:5433".parse().unwrap(),
+        TlsMode::None,
+    );
+    let (mut client, connection) = runtime.block_on(handshake).unwrap();
+    let connection = connection.map_err(|e| panic!("{}", e));
+    runtime.handle().spawn(connection).unwrap();
+
+    runtime
+        .block_on(client.batch_execute(
+            "CREATE TEMPORARY TABLE foo (id SERIAL, name TEXT);
+             INSERT INTO foo (name) VALUES ('alice'), ('bob'), ('charlie');
+             BEGIN;",
+        )).unwrap();
+
+    let statement = runtime
+        .block_on(client.prepare("SELECT id, name FROM foo ORDER BY id"))
+        .unwrap();
+    let portal = runtime.block_on(client.bind(&statement, &[])).unwrap();
+
+    let f1 = client.query_portal(&portal, 2).collect();
+    let f2 = client.query_portal(&portal, 2).collect();
+    let f3 = client.query_portal(&portal, 2).collect();
+    let (r1, r2, r3) = runtime.block_on(f1.join3(f2, f3)).unwrap();
+
+    assert_eq!(r1.len(), 2);
+    assert_eq!(r1[0].get::<_, i32>(0), 1);
+    assert_eq!(r1[0].get::<_, &str>(1), "alice");
+    assert_eq!(r1[1].get::<_, i32>(0), 2);
+    assert_eq!(r1[1].get::<_, &str>(1), "bob");
+
+    assert_eq!(r2.len(), 1);
+    assert_eq!(r2[0].get::<_, i32>(0), 3);
+    assert_eq!(r2[0].get::<_, &str>(1), "charlie");
+
+    assert_eq!(r3.len(), 0);
+}
+
 #[test]
 fn cancel_query() {
     let _ = env_logger::try_init();