Change batch_execute into simple_query

Closes rust-postgres#413

Author: sfackler

postgres/src/client.rs

@@ -1,3 +1,4 @@
+use fallible_iterator::FallibleIterator;
 use futures::{Async, Future, Poll, Stream};
 use std::io::{self, Read};
 use tokio_postgres::types::{ToSql, Type};
@@ -7,7 +8,7 @@ use tokio_postgres::{MakeTlsConnect, Socket, TlsConnect};
 
 #[cfg(feature = "runtime")]
 use crate::Config;
-use crate::{CopyOutReader, Query, Statement, ToStatement, Transaction};
+use crate::{CopyOutReader, Query, SimpleQuery, Statement, ToStatement, Transaction};
 
 pub struct Client(tokio_postgres::Client);
 
@@ -81,12 +82,12 @@ impl Client {
         CopyOutReader::new(stream)
     }
 
-    pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
-        self.0.batch_execute(query).wait()
+    pub fn simple_query(&mut self, query: &str) -> Result<SimpleQuery<'_>, Error> {
+        Ok(SimpleQuery::new(self.0.simple_query(query)))
     }
 
     pub fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
-        self.batch_execute("BEGIN")?;
+        self.simple_query("BEGIN")?.count()?;
         Ok(Transaction::new(self))
     }
 

postgres/src/lib.rs

@@ -10,6 +10,7 @@ mod copy_out_reader;
 mod portal;
 mod query;
 mod query_portal;
+mod simple_query;
 mod statement;
 mod to_statement;
 mod transaction;
@@ -25,6 +26,7 @@ pub use crate::copy_out_reader::*;
 pub use crate::portal::*;
 pub use crate::query::*;
 pub use crate::query_portal::*;
+pub use crate::simple_query::*;
 pub use crate::statement::*;
 pub use crate::to_statement::*;
 pub use crate::transaction::*;

postgres/src/simple_query.rs

@@ -0,0 +1,42 @@
+use fallible_iterator::FallibleIterator;
+use futures::stream::{self, Stream};
+use std::marker::PhantomData;
+use tokio_postgres::impls;
+use tokio_postgres::{Error, SimpleQueryMessage};
+
+pub struct SimpleQuery<'a> {
+    it: stream::Wait<impls::SimpleQuery>,
+    _p: PhantomData<&'a mut ()>,
+}
+
+// no-op impl to extend borrow until drop
+impl<'a> Drop for SimpleQuery<'a> {
+    fn drop(&mut self) {}
+}
+
+impl<'a> SimpleQuery<'a> {
+    pub(crate) fn new(stream: impls::SimpleQuery) -> SimpleQuery<'a> {
+        SimpleQuery {
+            it: stream.wait(),
+            _p: PhantomData,
+        }
+    }
+
+    /// A convenience API which collects the resulting messages into a `Vec` and returns them.
+    pub fn into_vec(self) -> Result<Vec<SimpleQueryMessage>, Error> {
+        self.collect()
+    }
+}
+
+impl<'a> FallibleIterator for SimpleQuery<'a> {
+    type Item = SimpleQueryMessage;
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<SimpleQueryMessage>, Error> {
+        match self.it.next() {
+            Some(Ok(row)) => Ok(Some(row)),
+            Some(Err(e)) => Err(e),
+            None => Ok(None),
+        }
+    }
+}

postgres/src/test.rs

@@ -1,3 +1,4 @@
+use fallible_iterator::FallibleIterator;
 use std::io::Read;
 use tokio_postgres::types::Type;
 use tokio_postgres::NoTls;
@@ -47,7 +48,9 @@ fn transaction_commit() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
 
     client
-        .batch_execute("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
+        .simple_query("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
+        .unwrap()
+        .count()
         .unwrap();
 
     let mut transaction = client.transaction().unwrap();
@@ -72,7 +75,9 @@ fn transaction_rollback() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
 
     client
-        .batch_execute("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
+        .simple_query("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
+        .unwrap()
+        .count()
         .unwrap();
 
     let mut transaction = client.transaction().unwrap();
@@ -96,7 +101,9 @@ fn transaction_drop() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
 
     client
-        .batch_execute("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
+        .simple_query("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
+        .unwrap()
+        .count()
         .unwrap();
 
     let mut transaction = client.transaction().unwrap();
@@ -120,7 +127,9 @@ fn nested_transactions() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
 
     client
-        .batch_execute("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY)")
+        .simple_query("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY)")
+        .unwrap()
+        .count()
         .unwrap();
 
     let mut transaction = client.transaction().unwrap();
@@ -177,7 +186,9 @@ fn copy_in() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
 
     client
-        .batch_execute("CREATE TEMPORARY TABLE foo (id INT, name TEXT)")
+        .simple_query("CREATE TEMPORARY TABLE foo (id INT, name TEXT)")
+        .unwrap()
+        .count()
         .unwrap();
 
     client
@@ -206,13 +217,12 @@ fn copy_out() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
 
     client
-        .batch_execute(
-            "
-            CREATE TEMPORARY TABLE foo (id INT, name TEXT);
-
-            INSERT INTO foo (id, name) VALUES (1, 'steven'), (2, 'timothy');
-            ",
+        .simple_query(
+            "CREATE TEMPORARY TABLE foo (id INT, name TEXT);
+             INSERT INTO foo (id, name) VALUES (1, 'steven'), (2, 'timothy');",
         )
+        .unwrap()
+        .count()
         .unwrap();
 
     let mut reader = client
@@ -224,21 +234,20 @@ fn copy_out() {
 
     assert_eq!(s, "1\tsteven\n2\ttimothy\n");
 
-    client.batch_execute("SELECT 1").unwrap();
+    client.simple_query("SELECT 1").unwrap().count().unwrap();
 }
 
 #[test]
 fn portal() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
 
     client
-        .batch_execute(
-            "
-            CREATE TEMPORARY TABLE foo (id INT);
-
-            INSERT INTO foo (id) VALUES (1), (2), (3);
-            ",
+        .simple_query(
+            "CREATE TEMPORARY TABLE foo (id INT);
+             INSERT INTO foo (id) VALUES (1), (2), (3);",
         )
+        .unwrap()
+        .count()
         .unwrap();
 
     let mut transaction = client.transaction().unwrap();

postgres/src/transaction.rs

@@ -1,9 +1,12 @@
+use fallible_iterator::FallibleIterator;
 use futures::Future;
 use std::io::Read;
 use tokio_postgres::types::{ToSql, Type};
 use tokio_postgres::Error;
 
-use crate::{Client, CopyOutReader, Portal, Query, QueryPortal, Statement, ToStatement};
+use crate::{
+    Client, CopyOutReader, Portal, Query, QueryPortal, SimpleQuery, Statement, ToStatement,
+};
 
 pub struct Transaction<'a> {
     client: &'a mut Client,
@@ -30,12 +33,14 @@ impl<'a> Transaction<'a> {
 
     pub fn commit(mut self) -> Result<(), Error> {
         self.done = true;
-        if self.depth == 0 {
-            self.client.batch_execute("COMMIT")
+        let it = if self.depth == 0 {
+            self.client.simple_query("COMMIT")?
         } else {
             self.client
-                .batch_execute(&format!("RELEASE sp{}", self.depth))
-        }
+                .simple_query(&format!("RELEASE sp{}", self.depth))?
+        };
+        it.count()?;
+        Ok(())
     }
 
     pub fn rollback(mut self) -> Result<(), Error> {
@@ -44,12 +49,14 @@ impl<'a> Transaction<'a> {
     }
 
     fn rollback_inner(&mut self) -> Result<(), Error> {
-        if self.depth == 0 {
-            self.client.batch_execute("ROLLBACK")
+        let it = if self.depth == 0 {
+            self.client.simple_query("ROLLBACK")?
         } else {
             self.client
-                .batch_execute(&format!("ROLLBACK TO sp{}", self.depth))
-        }
+                .simple_query(&format!("ROLLBACK TO sp{}", self.depth))?
+        };
+        it.count()?;
+        Ok(())
     }
 
     pub fn prepare(&mut self, query: &str) -> Result<Statement, Error> {
@@ -120,14 +127,15 @@ impl<'a> Transaction<'a> {
         self.client.copy_out(query, params)
     }
 
-    pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
-        self.client.batch_execute(query)
+    pub fn simple_query(&mut self, query: &str) -> Result<SimpleQuery<'_>, Error> {
+        self.client.simple_query(query)
     }
 
     pub fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
         let depth = self.depth + 1;
         self.client
-            .batch_execute(&format!("SAVEPOINT sp{}", depth))?;
+            .simple_query(&format!("SAVEPOINT sp{}", depth))?
+            .count()?;
         Ok(Transaction {
             client: self.client,
             depth,

tokio-postgres-native-tls/src/test.rs

@@ -100,6 +100,6 @@ fn runtime() {
     let connection = connection.map_err(|e| panic!("{}", e));
     runtime.spawn(connection);
 
-    let execute = client.batch_execute("SELECT 1");
+    let execute = client.simple_query("SELECT 1").for_each(|_| Ok(()));
     runtime.block_on(execute).unwrap();
 }

tokio-postgres-openssl/src/test.rs

@@ -85,6 +85,6 @@ fn runtime() {
     let connection = connection.map_err(|e| panic!("{}", e));
     runtime.spawn(connection);
 
-    let execute = client.batch_execute("SELECT 1");
+    let execute = client.simple_query("SELECT 1").for_each(|_| Ok(()));
     runtime.block_on(execute).unwrap();
 }

tokio-postgres/src/error/mod.rs

@@ -336,6 +336,7 @@ enum Kind {
     Tls,
     ToSql,
     FromSql,
+    Column,
     CopyInStream,
     Closed,
     Db,
@@ -373,6 +374,7 @@ impl fmt::Display for Error {
             Kind::Tls => "error performing TLS handshake",
             Kind::ToSql => "error serializing a value",
             Kind::FromSql => "error deserializing a value",
+            Kind::Column => "invalid column",
             Kind::CopyInStream => "error from a copy_in stream",
             Kind::Closed => "connection closed",
             Kind::Db => "db error",
@@ -451,6 +453,10 @@ impl Error {
         Error::new(Kind::FromSql, Some(e))
     }
 
+    pub(crate) fn column() -> Error {
+        Error::new(Kind::Column, None)
+    }
+
     pub(crate) fn copy_in_stream<E>(e: E) -> Error
     where
         E: Into<Box<dyn error::Error + Sync + Send>>,

tokio-postgres/src/impls.rs

@@ -5,7 +5,7 @@ use std::error;
 use tokio_io::{AsyncRead, AsyncWrite};
 
 use crate::proto;
-use crate::{Client, Connection, Error, Portal, Row, Statement, TlsConnect};
+use crate::{Client, Connection, Error, Portal, Row, SimpleQueryMessage, Statement, TlsConnect};
 #[cfg(feature = "runtime")]
 use crate::{MakeTlsConnect, Socket};
 
@@ -187,3 +187,16 @@ impl Stream for CopyOut {
         self.0.poll()
     }
 }
+
+/// The stream returned by `Client::simple_query`.
+#[must_use = "streams do nothing unless polled"]
+pub struct SimpleQuery(pub(crate) proto::SimpleQueryStream);
+
+impl Stream for SimpleQuery {
+    type Item = SimpleQueryMessage;
+    type Error = Error;
+
+    fn poll(&mut self) -> Poll<Option<SimpleQueryMessage>, Error> {
+        self.0.poll()
+    }
+}

tokio-postgres/src/lib.rs

@@ -102,7 +102,7 @@
 #![warn(rust_2018_idioms, clippy::all)]
 
 use bytes::IntoBuf;
-use futures::{try_ready, Async, Future, Poll, Stream};
+use futures::{Future, Poll, Stream};
 use std::error::Error as StdError;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use tokio_io::{AsyncRead, AsyncWrite};
@@ -240,19 +240,21 @@ impl Client {
         impls::CopyOut(self.0.copy_out(&statement.0, params))
     }
 
-    /// Executes a sequence of SQL statements.
+    /// Executes a sequence of SQL statements using the simple query protocol.
     ///
     /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that
-    /// point. This is intended for the execution of batches of non-dynamic statements, for example, the creation of
-    /// a schema for a fresh database.
+    /// point. The simple query protocol returns the values in rows as strings rather than in their binary encodings,
+    /// so the associated row type doesn't work with the `FromSql` trait. Rather than simply returning a stream over the
+    /// rows, this method returns a stream over an enum which indicates either the completion of one of the commands,
+    /// or a row of data. This preserves the framing between the separate statements in the request.
     ///
     /// # Warning
     ///
     /// Prepared statements should be use for any query which contains user-specified data, as they provided the
     /// functionality to safely imbed that data in the request. Do not form statements via string concatenation and pass
     /// them to this method!
-    pub fn batch_execute(&mut self, query: &str) -> BatchExecute {
-        BatchExecute(self.0.batch_execute(query))
+    pub fn simple_query(&mut self, query: &str) -> impls::SimpleQuery {
+        impls::SimpleQuery(self.0.simple_query(query))
     }
 
     /// A utility method to wrap a future in a database transaction.
@@ -445,18 +447,11 @@ where
     }
 }
 
-#[must_use = "futures do nothing unless polled"]
-pub struct BatchExecute(proto::SimpleQueryStream);
-
-impl Future for BatchExecute {
-    type Item = ();
-    type Error = Error;
-
-    fn poll(&mut self) -> Poll<(), Error> {
-        while let Some(_) = try_ready!(self.0.poll()) {}
-
-        Ok(Async::Ready(()))
-    }
+pub enum SimpleQueryMessage {
+    Row(SimpleQueryRow),
+    CommandComplete(u64),
+    #[doc(hidden)]
+    __NonExhaustive,
 }
 
 /// An asynchronous notification.