Make internal simple query future a stream

Author: sfackler

tokio-postgres/src/lib.rs

@@ -9,7 +9,7 @@ use tokio_io::{AsyncRead, AsyncWrite};
 pub use crate::builder::*;
 pub use crate::error::*;
 use crate::proto::CancelFuture;
-pub use crate::row::{Row, RowIndex};
+pub use crate::row::*;
 #[cfg(feature = "runtime")]
 pub use crate::socket::Socket;
 pub use crate::stmt::Column;
@@ -358,14 +358,16 @@ where
 }
 
 #[must_use = "futures do nothing unless polled"]
-pub struct BatchExecute(proto::SimpleQueryFuture);
+pub struct BatchExecute(proto::SimpleQueryStream);
 
 impl Future for BatchExecute {
     type Item = ();
     type Error = Error;
 
     fn poll(&mut self) -> Poll<(), Error> {
-        self.0.poll()
+        while let Some(_) = try_ready!(self.0.poll()) {}
+
+        Ok(Async::Ready(()))
     }
 }
 

tokio-postgres/src/proto/client.rs

@@ -18,7 +18,7 @@ use crate::proto::idle::{IdleGuard, IdleState};
 use crate::proto::portal::Portal;
 use crate::proto::prepare::PrepareFuture;
 use crate::proto::query::QueryStream;
-use crate::proto::simple_query::SimpleQueryFuture;
+use crate::proto::simple_query::SimpleQueryStream;
 use crate::proto::statement::Statement;
 use crate::types::{IsNull, Oid, ToSql, Type};
 use crate::Error;
@@ -121,13 +121,13 @@ impl Client {
             .map_err(|_| Error::closed())
     }
 
-    pub fn batch_execute(&self, query: &str) -> SimpleQueryFuture {
+    pub fn batch_execute(&self, query: &str) -> SimpleQueryStream {
         let pending = self.pending(|buf| {
             frontend::query(query, buf).map_err(Error::parse)?;
             Ok(())
         });
 
-        SimpleQueryFuture::new(self.clone(), pending)
+        SimpleQueryStream::new(self.clone(), pending)
     }
 
     pub fn prepare(&self, name: String, query: &str, param_types: &[Type]) -> PrepareFuture {

tokio-postgres/src/proto/mod.rs

@@ -59,7 +59,7 @@ pub use crate::proto::handshake::HandshakeFuture;
 pub use crate::proto::portal::Portal;
 pub use crate::proto::prepare::PrepareFuture;
 pub use crate::proto::query::QueryStream;
-pub use crate::proto::simple_query::SimpleQueryFuture;
+pub use crate::proto::simple_query::SimpleQueryStream;
 pub use crate::proto::statement::Statement;
 pub use crate::proto::tls::TlsFuture;
 pub use crate::proto::transaction::TransactionFuture;

tokio-postgres/src/proto/simple_query.rs

@@ -1,56 +1,70 @@
 use futures::sync::mpsc;
-use futures::{Poll, Stream};
+use futures::{Async, Poll, Stream};
 use postgres_protocol::message::backend::Message;
-use state_machine_future::{transition, RentToOwn, StateMachineFuture};
+use std::mem;
 
 use crate::proto::client::{Client, PendingRequest};
-use crate::Error;
+use crate::{Error, StringRow};
 
-#[derive(StateMachineFuture)]
-pub enum SimpleQuery {
-    #[state_machine_future(start, transitions(ReadResponse))]
+pub enum State {
     Start {
         client: Client,
         request: PendingRequest,
     },
-    #[state_machine_future(transitions(Finished))]
-    ReadResponse { receiver: mpsc::Receiver<Message> },
-    #[state_machine_future(ready)]
-    Finished(()),
-    #[state_machine_future(error)]
-    Failed(Error),
+    ReadResponse {
+        receiver: mpsc::Receiver<Message>,
+    },
+    Done,
 }
 
-impl PollSimpleQuery for SimpleQuery {
-    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
-        let state = state.take();
-        let receiver = state.client.send(state.request)?;
+pub struct SimpleQueryStream(State);
 
-        transition!(ReadResponse { receiver })
-    }
+impl Stream for SimpleQueryStream {
+    type Item = StringRow;
+    type Error = Error;
 
-    fn poll_read_response<'a>(
-        state: &'a mut RentToOwn<'a, ReadResponse>,
-    ) -> Poll<AfterReadResponse, Error> {
+    fn poll(&mut self) -> Poll<Option<StringRow>, Error> {
         loop {
-            let message = try_ready_receive!(state.receiver.poll());
+            match mem::replace(&mut self.0, State::Done) {
+                State::Start { client, request } => {
+                    let receiver = client.send(request)?;
+                    self.0 = State::ReadResponse { receiver };
+                }
+                State::ReadResponse { mut receiver } => {
+                    let message = match receiver.poll() {
+                        Ok(Async::Ready(message)) => message,
+                        Ok(Async::NotReady) => {
+                            self.0 = State::ReadResponse { receiver };
+                            return Ok(Async::NotReady);
+                        }
+                        Err(()) => unreachable!("mpsc receiver can't panic"),
+                    };
 
-            match message {
-                Some(Message::CommandComplete(_))
-                | Some(Message::RowDescription(_))
-                | Some(Message::DataRow(_))
-                | Some(Message::EmptyQueryResponse) => {}
-                Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
-                Some(Message::ReadyForQuery(_)) => transition!(Finished(())),
-                Some(_) => return Err(Error::unexpected_message()),
-                None => return Err(Error::closed()),
+                    match message {
+                        Some(Message::CommandComplete(_))
+                        | Some(Message::RowDescription(_))
+                        | Some(Message::EmptyQueryResponse) => {
+                            self.0 = State::ReadResponse { receiver };
+                        }
+                        Some(Message::DataRow(body)) => {
+                            self.0 = State::ReadResponse { receiver };
+                            let row = StringRow::new(body)?;
+                            return Ok(Async::Ready(Some(row)));
+                        }
+                        Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
+                        Some(Message::ReadyForQuery(_)) => return Ok(Async::Ready(None)),
+                        Some(_) => return Err(Error::unexpected_message()),
+                        None => return Err(Error::closed()),
+                    }
+                }
+                State::Done => return Ok(Async::Ready(None)),
             }
         }
     }
 }
 
-impl SimpleQueryFuture {
-    pub fn new(client: Client, request: PendingRequest) -> SimpleQueryFuture {
-        SimpleQuery::start(client, request)
+impl SimpleQueryStream {
+    pub fn new(client: Client, request: PendingRequest) -> SimpleQueryStream {
+        SimpleQueryStream(State::Start { client, request })
     }
 }

tokio-postgres/src/proto/transaction.rs

@@ -1,6 +1,6 @@
 use crate::proto::client::Client;
-use crate::proto::simple_query::SimpleQueryFuture;
-use futures::{try_ready, Async, Future, Poll};
+use crate::proto::simple_query::SimpleQueryStream;
+use futures::{try_ready, Async, Future, Poll, Stream};
 use state_machine_future::{transition, RentToOwn, StateMachineFuture};
 
 use crate::Error;
@@ -16,14 +16,14 @@ where
     #[state_machine_future(transitions(Running))]
     Beginning {
         client: Client,
-        begin: SimpleQueryFuture,
+        begin: SimpleQueryStream,
         future: F,
     },
     #[state_machine_future(transitions(Finishing))]
     Running { client: Client, future: F },
     #[state_machine_future(transitions(Finished))]
     Finishing {
-        future: SimpleQueryFuture,
+        future: SimpleQueryStream,
         result: Result<T, E>,
     },
     #[state_machine_future(ready)]
@@ -51,7 +51,8 @@ where
     fn poll_beginning<'a>(
         state: &'a mut RentToOwn<'a, Beginning<F, T, E>>,
     ) -> Poll<AfterBeginning<F, T, E>, E> {
-        try_ready!(state.begin.poll());
+        while let Some(_) = try_ready!(state.begin.poll()) {}
+
         let state = state.take();
         transition!(Running {
             client: state.client,
@@ -78,17 +79,20 @@ where
     fn poll_finishing<'a>(
         state: &'a mut RentToOwn<'a, Finishing<T, E>>,
     ) -> Poll<AfterFinishing<T>, E> {
-        match state.future.poll() {
-            Ok(Async::NotReady) => Ok(Async::NotReady),
-            Ok(Async::Ready(())) => {
-                let t = state.take().result?;
-                transition!(Finished(t))
+        loop {
+            match state.future.poll() {
+                Ok(Async::NotReady) => return Ok(Async::NotReady),
+                Ok(Async::Ready(Some(_))) => {}
+                Ok(Async::Ready(None)) => {
+                    let t = state.take().result?;
+                    transition!(Finished(t))
+                }
+                Err(e) => match state.take().result {
+                    Ok(_) => return Err(e.into()),
+                    // prioritize the future's error over the rollback error
+                    Err(e) => return Err(e),
+                },
             }
-            Err(e) => match state.take().result {
-                Ok(_) => Err(e.into()),
-                // prioritize the future's error over the rollback error
-                Err(e) => Err(e),
-            },
         }
     }
 }

tokio-postgres/src/row.rs

@@ -2,6 +2,7 @@ use fallible_iterator::FallibleIterator;
 use postgres_protocol::message::backend::DataRowBody;
 use std::fmt;
 use std::ops::Range;
+use std::str;
 
 use crate::proto;
 use crate::row::sealed::Sealed;
@@ -133,3 +134,50 @@ impl Row {
         value.map(Some).map_err(Error::from_sql)
     }
 }
+
+pub struct StringRow {
+    body: DataRowBody,
+    ranges: Vec<Option<Range<usize>>>,
+}
+
+impl StringRow {
+    #[allow(clippy::new_ret_no_self)]
+    pub(crate) fn new(body: DataRowBody) -> Result<StringRow, Error> {
+        let ranges = body.ranges().collect().map_err(Error::parse)?;
+        Ok(StringRow { body, ranges })
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    pub fn len(&self) -> usize {
+        self.ranges.len()
+    }
+
+    pub fn get(&self, idx: usize) -> Option<&str> {
+        match self.try_get(idx) {
+            Ok(Some(ok)) => ok,
+            Err(err) => panic!("error retrieving column {}: {}", idx, err),
+            Ok(None) => panic!("no such column {}", idx),
+        }
+    }
+
+    #[allow(clippy::option_option)] // FIXME
+    pub fn try_get(&self, idx: usize) -> Result<Option<Option<&str>>, Error> {
+        let buf = match self.ranges.get(idx) {
+            Some(range) => range.clone().map(|r| &self.body.buffer()[r]),
+            None => return Ok(None),
+        };
+
+        let v = match buf {
+            Some(buf) => {
+                let s = str::from_utf8(buf).map_err(|e| Error::from_sql(Box::new(e)))?;
+                Some(s)
+            }
+            None => None,
+        };
+
+        Ok(Some(v))
+    }
+}

tokio-postgres/src/stmt.rs

@@ -8,8 +8,7 @@ pub struct Column {
 }
 
 impl Column {
-    #[doc(hidden)]
-    pub fn new(name: String, type_: Type) -> Column {
+    pub(crate) fn new(name: String, type_: Type) -> Column {
         Column { name, type_ }
     }