Query!

Author: sfackler

postgres-shared/src/lib.rs

@@ -113,7 +113,7 @@ impl<'a> RowIndex for str {
     }
 }
 
-impl<'a, T> RowIndex for &'a T
+impl<'a, T: ?Sized> RowIndex for &'a T
     where T: RowIndex
 {
     #[inline]

postgres-tokio/Cargo.toml

@@ -6,6 +6,7 @@ authors = ["Steven Fackler <[email protected]>"]
 [dependencies]
 fallible-iterator = "0.1.3"
 futures = "0.1.7"
+futures-state-stream = { git = "https://github.com/sfackler/futures-state-stream" }
 postgres-shared = { path = "../postgres-shared" }
 postgres-protocol = "0.2"
 tokio-core = "0.1"

postgres-tokio/src/lib.rs

@@ -1,5 +1,6 @@
 extern crate fallible_iterator;
 extern crate futures;
+extern crate futures_state_stream;
 extern crate postgres_shared;
 extern crate postgres_protocol;
 extern crate tokio_core;
@@ -9,6 +10,7 @@ extern crate tokio_uds;
 use fallible_iterator::FallibleIterator;
 use futures::{Future, IntoFuture, BoxFuture, Stream, Sink, Poll, StartSend};
 use futures::future::Either;
+use futures_state_stream::{StreamEvent, StateStream, BoxStateStream, FutureExt};
 use postgres_protocol::authentication;
 use postgres_protocol::message::{backend, frontend};
 use postgres_protocol::message::backend::{ErrorResponseBody, ErrorFields};
@@ -562,6 +564,30 @@ impl Connection {
             .boxed()
     }
 
+    fn read_row(self) -> BoxFuture<(Option<RowData>, Connection), Error> {
+        self.0.read()
+            .map_err(Error::Io)
+            .and_then(|(m, s)| {
+                let c = Connection(s);
+                match m {
+                    backend::Message::DataRow(body) => {
+                        Either::A(body.values()
+                            .collect()
+                            .map(|r| (Some(r), c))
+                            .map_err(Error::Io)
+                            .into_future())
+                    }
+                    backend::Message::EmptyQueryResponse |
+                    backend::Message::CommandComplete(_) => Either::A(Ok((None, c)).into_future()),
+                    backend::Message::ErrorResponse(body) => {
+                        Either::B(c.ready_err(body))
+                    }
+                    _ => Either::A(Err(bad_message()).into_future()),
+                }
+            })
+            .boxed()
+    }
+
     pub fn prepare(mut self, query: &str) -> BoxFuture<(Statement, Connection), Error> {
         let id = self.0.next_stmt_id;
         self.0.next_stmt_id += 1;
@@ -585,6 +611,34 @@ impl Connection {
             .boxed()
     }
 
+    pub fn query(self,
+                 statement: &Statement,
+                 params: &[&ToSql])
+                 -> BoxStateStream<Row, Connection, Error> {
+        let columns = statement.columns.clone();
+        self.raw_execute(&statement.name, "", &statement.params, params)
+            .map(|c| {
+                futures_state_stream::unfold((c, columns), |(c, columns)| {
+                    c.read_row()
+                        .and_then(|(r, c)| {
+                            match r {
+                                Some(data) => {
+                                    let row = Row {
+                                        columns: columns.clone(),
+                                        data: data,
+                                    };
+                                    let event = StreamEvent::Next((row, (c, columns)));
+                                    Either::A(Ok(event).into_future())
+                                },
+                                None => Either::B(c.ready(()).map(|((), c)| StreamEvent::Done(c))),
+                            }
+                        })
+                })
+            })
+            .flatten_state_stream()
+            .boxed()
+    }
+
     pub fn close(self) -> BoxFuture<(), Error> {
         let mut terminate = vec![];
         frontend::terminate(&mut terminate);

postgres-tokio/src/test.rs

@@ -1,4 +1,5 @@
 use futures::Future;
+use futures_state_stream::StateStream;
 use tokio_core::reactor::Core;
 
 use super::*;
@@ -121,3 +122,20 @@ fn prepare_execute() {
         .map(|(n, _)| assert_eq!(n, 2));
     l.run(done).unwrap();
 }
+
+#[test]
+fn query() {
+    let mut l = Core::new().unwrap();
+    let done = Connection::connect("postgres://postgres@localhost", &l.handle())
+        .then(|c| {
+            c.unwrap().batch_execute("CREATE TEMPORARY TABLE foo (id SERIAL, name VARCHAR);
+                                      INSERT INTO foo (name) VALUES ('joe'), ('bob')")
+        })
+        .and_then(|c| c.prepare("SELECT id, name FROM foo ORDER BY id"))
+        .and_then(|(s, c)| c.query(&s, &[]).collect())
+        .map(|(r, _)| {
+            assert_eq!(r[0].get::<String, _>("name"), "joe");
+            assert_eq!(r[1].get::<String, _>("name"), "bob");
+        });
+    l.run(done).unwrap();
+}