Working select/execute

Author: sfackler

tokio-postgres/src/lib.rs

@@ -20,7 +20,9 @@ extern crate state_machine_future;
 #[cfg(unix)]
 extern crate tokio_uds;
 
-use futures::{Async, Future, Poll};
+use futures::{Async, Future, Poll, Stream};
+use postgres_shared::rows::RowIndex;
+use std::fmt;
 use std::io;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
@@ -33,7 +35,7 @@ pub use postgres_shared::{CancelData, Notification};
 
 use error::Error;
 use params::ConnectParams;
-use types::{ToSql, Type};
+use types::{FromSql, ToSql, Type};
 
 mod proto;
 
@@ -72,6 +74,10 @@ impl Client {
     pub fn execute(&mut self, statement: &Statement, params: &[&ToSql]) -> Execute {
         Execute(self.0.execute(&statement.0, params))
     }
+
+    pub fn query(&mut self, statement: &Statement, params: &[&ToSql]) -> Query {
+        Query(self.0.query(&statement.0, params))
+    }
 }
 
 #[must_use = "futures do nothing unless polled"]
@@ -147,3 +153,48 @@ impl Future for Execute {
         self.0.poll()
     }
 }
+
+#[must_use = "streams do nothing unless polled"]
+pub struct Query(proto::QueryStream);
+
+impl Stream for Query {
+    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 Row(proto::Row);
+
+impl Row {
+    pub fn columns(&self) -> &[Column] {
+        self.0.columns()
+    }
+
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    pub fn get<'a, I, T>(&'a self, idx: I) -> T
+    where
+        I: RowIndex + fmt::Debug,
+        T: FromSql<'a>,
+    {
+        self.0.get(idx)
+    }
+
+    pub fn try_get<'a, I, T>(&'a self, idx: I) -> Result<Option<T>, Error>
+    where
+        I: RowIndex,
+        T: FromSql<'a>,
+    {
+        self.0.try_get(idx)
+    }
+}

tokio-postgres/src/proto/client.rs

@@ -8,6 +8,7 @@ use error::{self, Error};
 use proto::connection::Request;
 use proto::execute::ExecuteFuture;
 use proto::prepare::PrepareFuture;
+use proto::query::QueryStream;
 use proto::statement::Statement;
 use types::{IsNull, ToSql, Type};
 
@@ -48,7 +49,17 @@ impl Client {
     }
 
     pub fn execute(&mut self, statement: &Statement, params: &[&ToSql]) -> ExecuteFuture {
-        let pending = self.pending(|buf| {
+        let pending = self.pending_execute(statement, params);
+        ExecuteFuture::new(pending, statement.clone())
+    }
+
+    pub fn query(&mut self, statement: &Statement, params: &[&ToSql]) -> QueryStream {
+        let pending = self.pending_execute(statement, params);
+        QueryStream::new(pending, statement.clone())
+    }
+
+    fn pending_execute(&self, statement: &Statement, params: &[&ToSql]) -> PendingRequest {
+        self.pending(|buf| {
             let r = frontend::bind(
                 "",
                 statement.name(),
@@ -70,9 +81,7 @@ impl Client {
             frontend::execute("", 0, buf)?;
             frontend::sync(buf);
             Ok(())
-        });
-
-        ExecuteFuture::new(pending, statement.clone())
+        })
     }
 
     fn pending<F>(&self, messages: F) -> PendingRequest

tokio-postgres/src/proto/mod.rs

@@ -14,6 +14,8 @@ mod connection;
 mod execute;
 mod handshake;
 mod prepare;
+mod query;
+mod row;
 mod socket;
 mod statement;
 
@@ -23,5 +25,7 @@ pub use proto::connection::Connection;
 pub use proto::execute::ExecuteFuture;
 pub use proto::handshake::HandshakeFuture;
 pub use proto::prepare::PrepareFuture;
+pub use proto::query::QueryStream;
+pub use proto::row::Row;
 pub use proto::socket::Socket;
 pub use proto::statement::Statement;

tokio-postgres/src/proto/prepare.rs

@@ -45,7 +45,7 @@ pub enum Prepare {
         receiver: mpsc::Receiver<Message>,
         name: String,
         parameters: ParameterDescriptionBody,
-        columns: RowDescriptionBody,
+        columns: Option<RowDescriptionBody>,
     },
     #[state_machine_future(ready)]
     Finished(Statement),
@@ -96,7 +96,6 @@ impl PollPrepare for Prepare {
                 name: state.name,
                 parameters: body,
             }),
-            Some(Message::ErrorResponse(body)) => Err(error::__db(body)),
             Some(_) => Err(bad_response()),
             None => Err(disconnected()),
         }
@@ -108,18 +107,20 @@ impl PollPrepare for Prepare {
         let message = try_receive!(state.receiver.poll());
         let state = state.take();
 
-        match message {
-            Some(Message::RowDescription(body)) => transition!(ReadReadyForQuery {
-                sender: state.sender,
-                receiver: state.receiver,
-                name: state.name,
-                parameters: state.parameters,
-                columns: body,
-            }),
-            Some(Message::ErrorResponse(body)) => Err(error::__db(body)),
-            Some(_) => Err(bad_response()),
-            None => Err(disconnected()),
-        }
+        let body = match message {
+            Some(Message::RowDescription(body)) => Some(body),
+            Some(Message::NoData) => None,
+            Some(_) => return Err(bad_response()),
+            None => return Err(disconnected()),
+        };
+
+        transition!(ReadReadyForQuery {
+            sender: state.sender,
+            receiver: state.receiver,
+            name: state.name,
+            parameters: state.parameters,
+            columns: body,
+        })
     }
 
     fn poll_read_ready_for_query<'a>(
@@ -136,13 +137,15 @@ impl PollPrepare for Prepare {
                     .parameters()
                     .map(|oid| Type::from_oid(oid).unwrap())
                     .collect()?;
-                let columns = state
-                    .columns
-                    .fields()
-                    .map(|f| {
-                        Column::new(f.name().to_string(), Type::from_oid(f.type_oid()).unwrap())
-                    })
-                    .collect()?;
+                let columns = match state.columns {
+                    Some(body) => body
+                        .fields()
+                        .map(|f| {
+                            Column::new(f.name().to_string(), Type::from_oid(f.type_oid()).unwrap())
+                        })
+                        .collect()?,
+                    None => vec![],
+                };
 
                 transition!(Finished(Statement::new(
                     state.sender,
@@ -151,7 +154,6 @@ impl PollPrepare for Prepare {
                     columns
                 )))
             }
-            Some(Message::ErrorResponse(body)) => Err(error::__db(body)),
             Some(_) => Err(bad_response()),
             None => Err(disconnected()),
         }

tokio-postgres/src/proto/query.rs

@@ -0,0 +1,108 @@
+use futures::sync::mpsc;
+use futures::{Async, Poll, Stream};
+use postgres_protocol::message::backend::Message;
+use std::mem;
+
+use error::{self, Error};
+use proto::client::PendingRequest;
+use proto::row::Row;
+use proto::statement::Statement;
+use {bad_response, disconnected};
+
+enum State {
+    Start {
+        request: PendingRequest,
+        statement: Statement,
+    },
+    ReadingResponse {
+        receiver: mpsc::Receiver<Message>,
+        statement: Statement,
+    },
+    ReadingReadyForQuery {
+        receiver: mpsc::Receiver<Message>,
+    },
+    Done,
+}
+
+pub struct QueryStream(State);
+
+impl Stream for QueryStream {
+    type Item = Row;
+    type Error = Error;
+
+    fn poll(&mut self) -> Poll<Option<Row>, Error> {
+        loop {
+            match mem::replace(&mut self.0, State::Done) {
+                State::Start { request, statement } => {
+                    let receiver = request.send()?;
+                    self.0 = State::ReadingResponse {
+                        receiver,
+                        statement,
+                    };
+                }
+                State::ReadingResponse {
+                    mut receiver,
+                    statement,
+                } => {
+                    let message = match receiver.poll() {
+                        Ok(Async::Ready(message)) => message,
+                        Ok(Async::NotReady) => {
+                            self.0 = State::ReadingResponse {
+                                receiver,
+                                statement,
+                            };
+                            break Ok(Async::NotReady);
+                        }
+                        Err(()) => unreachable!("mpsc::Receiver doesn't return errors"),
+                    };
+
+                    match message {
+                        Some(Message::BindComplete) => {
+                            self.0 = State::ReadingResponse {
+                                receiver,
+                                statement,
+                            };
+                        }
+                        Some(Message::ErrorResponse(body)) => break Err(error::__db(body)),
+                        Some(Message::DataRow(body)) => {
+                            let row = Row::new(statement.clone(), body)?;
+                            self.0 = State::ReadingResponse {
+                                receiver,
+                                statement,
+                            };
+                            break Ok(Async::Ready(Some(row)));
+                        }
+                        Some(Message::EmptyQueryResponse) | Some(Message::CommandComplete(_)) => {
+                            self.0 = State::ReadingReadyForQuery { receiver };
+                        }
+                        Some(_) => break Err(bad_response()),
+                        None => break Err(disconnected()),
+                    }
+                }
+                State::ReadingReadyForQuery { mut receiver } => {
+                    let message = match receiver.poll() {
+                        Ok(Async::Ready(message)) => message,
+                        Ok(Async::NotReady) => {
+                            self.0 = State::ReadingReadyForQuery { receiver };
+                            break Ok(Async::NotReady);
+                        }
+                        Err(()) => unreachable!("mpsc::Receiver doesn't return errors"),
+                    };
+
+                    match message {
+                        Some(Message::ReadyForQuery(_)) => break Ok(Async::Ready(None)),
+                        Some(_) => break Err(bad_response()),
+                        None => break Err(disconnected()),
+                    }
+                }
+                State::Done => break Ok(Async::Ready(None)),
+            }
+        }
+    }
+}
+
+impl QueryStream {
+    pub fn new(request: PendingRequest, statement: Statement) -> QueryStream {
+        QueryStream(State::Start { request, statement })
+    }
+}

tokio-postgres/src/proto/row.rs

@@ -0,0 +1,66 @@
+use postgres_protocol::message::backend::DataRowBody;
+use postgres_shared::rows::{RowData, RowIndex};
+use std::fmt;
+
+use error::{self, Error};
+use proto::statement::Statement;
+use types::{FromSql, WrongType};
+use Column;
+
+pub struct Row {
+    statement: Statement,
+    data: RowData,
+}
+
+impl Row {
+    pub fn new(statement: Statement, data: DataRowBody) -> Result<Row, Error> {
+        let data = RowData::new(data)?;
+        Ok(Row { statement, data })
+    }
+
+    pub fn columns(&self) -> &[Column] {
+        self.statement.columns()
+    }
+
+    pub fn len(&self) -> usize {
+        self.columns().len()
+    }
+
+    pub fn get<'b, I, T>(&'b self, idx: I) -> T
+    where
+        I: RowIndex + fmt::Debug,
+        T: FromSql<'b>,
+    {
+        match self.get_inner(&idx) {
+            Ok(Some(ok)) => ok,
+            Err(err) => panic!("error retrieving column {:?}: {:?}", idx, err),
+            Ok(None) => panic!("no such column {:?}", idx),
+        }
+    }
+
+    pub fn try_get<'b, I, T>(&'b self, idx: I) -> Result<Option<T>, Error>
+    where
+        I: RowIndex,
+        T: FromSql<'b>,
+    {
+        self.get_inner(&idx)
+    }
+
+    fn get_inner<'b, I, T>(&'b self, idx: &I) -> Result<Option<T>, Error>
+    where
+        I: RowIndex,
+        T: FromSql<'b>,
+    {
+        let idx = match idx.__idx(&self.columns()) {
+            Some(idx) => idx,
+            None => return Ok(None),
+        };
+
+        let ty = self.statement.columns()[idx].type_();
+        if !<T as FromSql>::accepts(ty) {
+            return Err(error::conversion(Box::new(WrongType::new(ty.clone()))));
+        }
+        let value = FromSql::from_sql_nullable(ty, self.data.get(idx));
+        value.map(Some).map_err(error::conversion)
+    }
+}