Test and fix simple_query

sfackler · sfackler · commit 9938ffff1abc · 2019-07-28T16:34:07.000-07:00
diff --git a/tokio-postgres/Cargo.toml b/tokio-postgres/Cargo.toml
@@ -37,6 +37,7 @@ futures-preview = { version = "0.3.0-alpha.17", features = ["nightly", "async-aw
 log = "0.4"
 parking_lot = "0.9"
 percent-encoding = "1.0"
+pin-utils = "0.1.0-alpha.4"
 phf = "0.7.23"
 postgres-protocol = { version = "0.4.1", path = "../postgres-protocol" }
 tokio = { git = "https://github.com/tokio-rs/tokio", default-features = false, features = ["io", "codec"] }
diff --git a/tokio-postgres/src/connect.rs b/tokio-postgres/src/connect.rs
@@ -1,4 +1,5 @@
 use crate::config::{Host, TargetSessionAttrs};
+use pin_utils::pin_mut;
 use crate::connect_raw::connect_raw;
 use crate::connect_socket::connect_socket;
 use crate::tls::{MakeTlsConnect, TlsConnect};
@@ -55,7 +56,8 @@ where
     let (mut client, connection) = connect_raw(socket, tls, config, Some(idx)).await?;
 
     if let TargetSessionAttrs::ReadWrite = config.target_session_attrs {
-        let mut rows = client.simple_query("SHOW transaction_read_only");
+        let rows = client.simple_query("SHOW transaction_read_only");
+        pin_mut!(rows);
 
         loop {
             match rows.try_next().await? {
diff --git a/tokio-postgres/src/simple_query.rs b/tokio-postgres/src/simple_query.rs
@@ -3,11 +3,10 @@ use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::{Error, SimpleQueryMessage, SimpleQueryRow};
 use fallible_iterator::FallibleIterator;
-use futures::{ready, Stream};
+use futures::{ready, Stream, TryFutureExt};
 use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
 use std::future::Future;
-use std::mem;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
@@ -17,7 +16,18 @@ pub fn simple_query(
     query: &str,
 ) -> impl Stream<Item = Result<SimpleQueryMessage, Error>> {
     let buf = encode(query);
-    SimpleQuery::Start { client, buf }
+
+    let start = async move {
+        let buf = buf?;
+        let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+        Ok(SimpleQuery {
+            responses,
+            columns: None,
+        })
+    };
+
+    start.try_flatten_stream()
 }
 
 pub fn batch_execute(
@@ -49,84 +59,49 @@ fn encode(query: &str) -> Result<Vec<u8>, Error> {
     Ok(buf)
 }
 
-enum SimpleQuery {
-    Start {
-        client: Arc<InnerClient>,
-        buf: Result<Vec<u8>, Error>,
-    },
-    Reading {
-        responses: Responses,
-        columns: Option<Arc<[String]>>,
-    },
-    Done,
+struct SimpleQuery {
+    responses: Responses,
+    columns: Option<Arc<[String]>>
 }
 
 impl Stream for SimpleQuery {
     type Item = Result<SimpleQueryMessage, Error>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         loop {
-            match mem::replace(&mut *self, SimpleQuery::Done) {
-                SimpleQuery::Start { client, buf } => {
-                    let buf = buf?;
-                    let responses =
-                        client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
-
-                    *self = SimpleQuery::Reading {
-                        responses,
-                        columns: None,
+            match ready!(self.responses.poll_next(cx)?) {
+                Message::CommandComplete(body) => {
+                    let rows = body
+                        .tag()
+                        .map_err(Error::parse)?
+                        .rsplit(' ')
+                        .next()
+                        .unwrap()
+                        .parse()
+                        .unwrap_or(0);
+                    return Poll::Ready(Some(Ok(SimpleQueryMessage::CommandComplete(rows))));
+                }
+                Message::EmptyQueryResponse => {
+                    return Poll::Ready(Some(Ok(SimpleQueryMessage::CommandComplete(0))));
+                }
+                Message::RowDescription(body) => {
+                    let columns = body
+                        .fields()
+                        .map(|f| Ok(f.name().to_string()))
+                        .collect::<Vec<_>>()
+                        .map_err(Error::parse)?
+                        .into();
+                    self.columns = Some(columns);
+                }
+                Message::DataRow(body) => {
+                    let row = match &self.columns {
+                        Some(columns) => SimpleQueryRow::new(columns.clone(), body)?,
+                        None => return Poll::Ready(Some(Err(Error::unexpected_message()))),
                     };
+                    return Poll::Ready(Some(Ok(SimpleQueryMessage::Row(row))));
                 }
-                SimpleQuery::Reading {
-                    mut responses,
-                    columns,
-                } => match ready!(responses.poll_next(cx)?) {
-                    Message::CommandComplete(body) => {
-                        let rows = body
-                            .tag()
-                            .map_err(Error::parse)?
-                            .rsplit(' ')
-                            .next()
-                            .unwrap()
-                            .parse()
-                            .unwrap_or(0);
-                        *self = SimpleQuery::Reading {
-                            responses,
-                            columns: None,
-                        };
-                        return Poll::Ready(Some(Ok(SimpleQueryMessage::CommandComplete(rows))));
-                    }
-                    Message::EmptyQueryResponse => {
-                        *self = SimpleQuery::Reading {
-                            responses,
-                            columns: None,
-                        };
-                        return Poll::Ready(Some(Ok(SimpleQueryMessage::CommandComplete(0))));
-                    }
-                    Message::RowDescription(body) => {
-                        let columns = body
-                            .fields()
-                            .map(|f| Ok(f.name().to_string()))
-                            .collect::<Vec<_>>()
-                            .map_err(Error::parse)?
-                            .into();
-                        *self = SimpleQuery::Reading {
-                            responses,
-                            columns: Some(columns),
-                        };
-                    }
-                    Message::DataRow(body) => {
-                        let row = match &columns {
-                            Some(columns) => SimpleQueryRow::new(columns.clone(), body)?,
-                            None => return Poll::Ready(Some(Err(Error::unexpected_message()))),
-                        };
-                        *self = SimpleQuery::Reading { responses, columns };
-                        return Poll::Ready(Some(Ok(SimpleQueryMessage::Row(row))));
-                    }
-                    Message::ReadyForQuery(_) => return Poll::Ready(None),
-                    _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
-                },
-                SimpleQuery::Done => return Poll::Ready(None),
+                Message::ReadyForQuery(_) => return Poll::Ready(None),
+                _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
             }
         }
     }
diff --git a/tokio-postgres/tests/test/main.rs b/tokio-postgres/tests/test/main.rs
@@ -6,7 +6,7 @@ use tokio::net::TcpStream;
 use tokio_postgres::error::SqlState;
 use tokio_postgres::tls::{NoTls, NoTlsStream};
 use tokio_postgres::types::{Kind, Type};
-use tokio_postgres::{Client, Config, Connection, Error};
+use tokio_postgres::{Client, Config, Connection, Error, SimpleQueryMessage};
 
 mod parse;
 #[cfg(feature = "runtime")]
@@ -114,12 +114,10 @@ async fn pipelined_prepare() {
 async fn insert_select() {
     let mut client = connect("user=postgres").await;
 
-    let setup = client
-        .prepare("CREATE TEMPORARY TABLE foo (id SERIAL, name TEXT)")
+    client
+        .batch_execute("CREATE TEMPORARY TABLE foo (id SERIAL, name TEXT)")
         .await
         .unwrap();
-    client.execute(&setup, &[]).await.unwrap();
-    drop(setup);
 
     let insert = client.prepare("INSERT INTO foo (name) VALUES ($1), ($2)");
     let select = client.prepare("SELECT id, name FROM foo ORDER BY id");
@@ -142,8 +140,8 @@ async fn insert_select() {
 async fn custom_enum() {
     let mut client = connect("user=postgres").await;
 
-    let create = client
-        .prepare(
+    client
+        .batch_execute(
             "CREATE TYPE pg_temp.mood AS ENUM (
                 'sad',
                 'ok',
@@ -152,7 +150,6 @@ async fn custom_enum() {
         )
         .await
         .unwrap();
-    client.execute(&create, &[]).await.unwrap();
 
     let select = client.prepare("SELECT $1::mood").await.unwrap();
 
@@ -172,11 +169,10 @@ async fn custom_enum() {
 async fn custom_domain() {
     let mut client = connect("user=postgres").await;
 
-    let create = client
-        .prepare("CREATE DOMAIN pg_temp.session_id AS bytea CHECK(octet_length(VALUE) = 16)")
+    client
+        .batch_execute("CREATE DOMAIN pg_temp.session_id AS bytea CHECK(octet_length(VALUE) = 16)")
         .await
         .unwrap();
-    client.execute(&create, &[]).await.unwrap();
 
     let select = client.prepare("SELECT $1::session_id").await.unwrap();
 
@@ -206,17 +202,16 @@ async fn custom_array() {
 async fn custom_composite() {
     let mut client = connect("user=postgres").await;
 
-    let create = client
-        .prepare(
+    client
+        .batch_execute(
             "CREATE TYPE pg_temp.inventory_item AS (
-            name TEXT,
-            supplier INTEGER,
-            price NUMERIC
-        )",
+                name TEXT,
+                supplier INTEGER,
+                price NUMERIC
+            )",
         )
         .await
         .unwrap();
-    client.execute(&create, &[]).await.unwrap();
 
     let select = client.prepare("SELECT $1::inventory_item").await.unwrap();
 
@@ -239,16 +234,15 @@ async fn custom_composite() {
 async fn custom_range() {
     let mut client = connect("user=postgres").await;
 
-    let create = client
-        .prepare(
+    client
+        .batch_execute(
             "CREATE TYPE pg_temp.floatrange AS RANGE (
-            subtype = float8,
-            subtype_diff = float8mi
-        )",
+                subtype = float8,
+                subtype_diff = float8mi
+            )",
         )
         .await
         .unwrap();
-    client.execute(&create, &[]).await.unwrap();
 
     let select = client.prepare("SELECT $1::floatrange").await.unwrap();
 
@@ -257,6 +251,52 @@ async fn custom_range() {
     assert_eq!(&Kind::Range(Type::FLOAT8), ty.kind());
 }
 
+#[tokio::test]
+async fn simple_query() {
+    let mut client = connect("user=postgres").await;
+
+    let messages = client
+        .simple_query(
+            "CREATE TEMPORARY TABLE foo (
+                id SERIAL,
+                name TEXT
+            );
+            INSERT INTO foo (name) VALUES ('steven'), ('joe');
+            SELECT * FROM foo ORDER BY id;",
+        )
+        .try_collect::<Vec<_>>()
+        .await
+        .unwrap();
+
+    match messages[0] {
+        SimpleQueryMessage::CommandComplete(0) => {}
+        _ => panic!("unexpected message"),
+    }
+    match messages[1] {
+        SimpleQueryMessage::CommandComplete(2) => {}
+        _ => panic!("unexpected message"),
+    }
+    match &messages[2] {
+        SimpleQueryMessage::Row(row) => {
+            assert_eq!(row.get(0), Some("1"));
+            assert_eq!(row.get(1), Some("steven"));
+        }
+        _ => panic!("unexpected message"),
+    }
+    match &messages[3] {
+        SimpleQueryMessage::Row(row) => {
+            assert_eq!(row.get(0), Some("2"));
+            assert_eq!(row.get(1), Some("joe"));
+        }
+        _ => panic!("unexpected message"),
+    }
+    match messages[4] {
+        SimpleQueryMessage::CommandComplete(2) => {}
+        _ => panic!("unexpected message"),
+    }
+    assert_eq!(messages.len(), 5);
+}
+
 /*
 #[test]
 fn query_portal() {
@@ -675,56 +715,6 @@ fn transaction_builder_around_moved_client() {
     runtime.run().unwrap();
 }
 
-#[test]
-fn simple_query() {
-    let _ = env_logger::try_init();
-    let mut runtime = Runtime::new().unwrap();
-
-    let (mut client, connection) = runtime.block_on(connect("user=postgres")).unwrap();
-    let connection = connection.map_err(|e| panic!("{}", e));
-    runtime.handle().spawn(connection).unwrap();
-
-    let f = client
-        .simple_query(
-            "CREATE TEMPORARY TABLE foo (
-                id SERIAL,
-                name TEXT
-            );
-            INSERT INTO foo (name) VALUES ('steven'), ('joe');
-            SELECT * FROM foo ORDER BY id;",
-        )
-        .collect();
-    let messages = runtime.block_on(f).unwrap();
-
-    match messages[0] {
-        SimpleQueryMessage::CommandComplete(0) => {}
-        _ => panic!("unexpected message"),
-    }
-    match messages[1] {
-        SimpleQueryMessage::CommandComplete(2) => {}
-        _ => panic!("unexpected message"),
-    }
-    match &messages[2] {
-        SimpleQueryMessage::Row(row) => {
-            assert_eq!(row.get(0), Some("1"));
-            assert_eq!(row.get(1), Some("steven"));
-        }
-        _ => panic!("unexpected message"),
-    }
-    match &messages[3] {
-        SimpleQueryMessage::Row(row) => {
-            assert_eq!(row.get(0), Some("2"));
-            assert_eq!(row.get(1), Some("joe"));
-        }
-        _ => panic!("unexpected message"),
-    }
-    match messages[4] {
-        SimpleQueryMessage::CommandComplete(2) => {}
-        _ => panic!("unexpected message"),
-    }
-    assert_eq!(messages.len(), 5);
-}
-
 #[test]
 fn poll_idle_running() {
     struct DelayStream(Delay);
