Connection changes (sfackler#21)

* refactor query_raw_txt to use a pre-prepared statement

* expose ready_status on RowStream

Author: conradludgate

.github/workflows/ci.yml

@@ -56,7 +56,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: sfackler/actions/rustup@master
         with:
-          version: 1.65.0
+          version: 1.67.0
       - run: echo "::set-output name=version::$(rustc --version)"
         id: rust-version
       - run: rustup target add wasm32-unknown-unknown

tokio-postgres/src/client.rs

@@ -372,13 +372,19 @@ impl Client {
 
     /// Pass text directly to the Postgres backend to allow it to sort out typing itself and
     /// to save a roundtrip
-    pub async fn query_raw_txt<'a, S, I>(&self, query: S, params: I) -> Result<RowStream, Error>
+    pub async fn query_raw_txt<'a, T, S, I>(
+        &self,
+        statement: &T,
+        params: I,
+    ) -> Result<RowStream, Error>
     where
-        S: AsRef<str> + Sync + Send,
+        T: ?Sized + ToStatement,
+        S: AsRef<str>,
         I: IntoIterator<Item = Option<S>>,
-        I::IntoIter: ExactSizeIterator + Sync + Send,
+        I::IntoIter: ExactSizeIterator,
     {
-        query::query_txt(&self.inner, query, params).await
+        let statement = statement.__convert().into_statement(self).await?;
+        query::query_txt(&self.inner, statement, params).await
     }
 
     /// Executes a statement, returning the number of rows modified.

tokio-postgres/src/generic_client.rs

@@ -57,8 +57,13 @@ pub trait GenericClient: private::Sealed {
         I::IntoIter: ExactSizeIterator;
 
     /// Like `Client::query_raw_txt`.
-    async fn query_raw_txt<'a, S, I>(&self, query: S, params: I) -> Result<RowStream, Error>
+    async fn query_raw_txt<'a, T, S, I>(
+        &self,
+        statement: &T,
+        params: I,
+    ) -> Result<RowStream, Error>
     where
+        T: ?Sized + ToStatement + Sync + Send,
         S: AsRef<str> + Sync + Send,
         I: IntoIterator<Item = Option<S>> + Sync + Send,
         I::IntoIter: ExactSizeIterator + Sync + Send;
@@ -143,13 +148,14 @@ impl GenericClient for Client {
         self.query_raw(statement, params).await
     }
 
-    async fn query_raw_txt<'a, S, I>(&self, query: S, params: I) -> Result<RowStream, Error>
+    async fn query_raw_txt<'a, T, S, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
     where
+        T: ?Sized + ToStatement + Sync + Send,
         S: AsRef<str> + Sync + Send,
         I: IntoIterator<Item = Option<S>> + Sync + Send,
         I::IntoIter: ExactSizeIterator + Sync + Send,
     {
-        self.query_raw_txt(query, params).await
+        self.query_raw_txt(statement, params).await
     }
 
     async fn prepare(&self, query: &str) -> Result<Statement, Error> {
@@ -238,13 +244,14 @@ impl GenericClient for Transaction<'_> {
         self.query_raw(statement, params).await
     }
 
-    async fn query_raw_txt<'a, S, I>(&self, query: S, params: I) -> Result<RowStream, Error>
+    async fn query_raw_txt<'a, T, S, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
     where
+        T: ?Sized + ToStatement + Sync + Send,
         S: AsRef<str> + Sync + Send,
         I: IntoIterator<Item = Option<S>> + Sync + Send,
         I::IntoIter: ExactSizeIterator + Sync + Send,
     {
-        self.query_raw_txt(query, params).await
+        self.query_raw_txt(statement, params).await
     }
 
     async fn prepare(&self, query: &str) -> Result<Statement, Error> {

tokio-postgres/src/query.rs

@@ -1,17 +1,15 @@
 use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
-use crate::prepare::get_type;
 use crate::types::{BorrowToSql, IsNull};
-use crate::{Column, Error, Portal, Row, Statement};
+use crate::{Error, Portal, Row, Statement};
 use bytes::{BufMut, Bytes, BytesMut};
-use fallible_iterator::FallibleIterator;
 use futures_util::{ready, Stream};
 use log::{debug, log_enabled, Level};
 use pin_project_lite::pin_project;
 use postgres_protocol::message::backend::{CommandCompleteBody, Message};
 use postgres_protocol::message::frontend;
-use postgres_types::Type;
+use postgres_types::Format;
 use std::fmt;
 use std::marker::PhantomPinned;
 use std::pin::Pin;
@@ -57,30 +55,29 @@ where
         statement,
         responses,
         command_tag: None,
+        status: None,
+        output_format: Format::Binary,
         _p: PhantomPinned,
     })
 }
 
 pub async fn query_txt<S, I>(
     client: &Arc<InnerClient>,
-    query: S,
+    statement: Statement,
     params: I,
 ) -> Result<RowStream, Error>
 where
-    S: AsRef<str> + Sync + Send,
+    S: AsRef<str>,
     I: IntoIterator<Item = Option<S>>,
     I::IntoIter: ExactSizeIterator,
 {
     let params = params.into_iter();
-    let params_len = params.len();
 
     let buf = client.with_buf(|buf| {
-        // Parse, anonymous portal
-        frontend::parse("", query.as_ref(), std::iter::empty(), buf).map_err(Error::encode)?;
         // Bind, pass params as text, retrieve as binary
         match frontend::bind(
             "",                 // empty string selects the unnamed portal
-            "",                 // empty string selects the unnamed prepared statement
+            statement.name(),   // named prepared statement
             std::iter::empty(), // all parameters use the default format (text)
             params,
             |param, buf| match param {
@@ -98,8 +95,6 @@ where
             Err(frontend::BindError::Serialization(e)) => Err(Error::encode(e)),
         }?;
 
-        // Describe portal to typecast results
-        frontend::describe(b'P', "", buf).map_err(Error::encode)?;
         // Execute
         frontend::execute("", 0, buf).map_err(Error::encode)?;
         // Sync
@@ -108,43 +103,16 @@ where
         Ok(buf.split().freeze())
     })?;
 
-    let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
-
     // now read the responses
-
-    match responses.next().await? {
-        Message::ParseComplete => {}
-        _ => return Err(Error::unexpected_message()),
-    }
-    match responses.next().await? {
-        Message::BindComplete => {}
-        _ => return Err(Error::unexpected_message()),
-    }
-    let row_description = match responses.next().await? {
-        Message::RowDescription(body) => Some(body),
-        Message::NoData => None,
-        _ => return Err(Error::unexpected_message()),
-    };
-
-    // construct statement object
-
-    let parameters = vec![Type::UNKNOWN; params_len];
-
-    let mut columns = vec![];
-    if let Some(row_description) = row_description {
-        let mut it = row_description.fields();
-        while let Some(field) = it.next().map_err(Error::parse)? {
-            // NB: for some types that function may send a query to the server. At least in
-            // raw text mode we don't need that info and can skip this.
-            let type_ = get_type(client, field.type_oid()).await?;
-            let column = Column::new(field.name().to_string(), type_, field);
-            columns.push(column);
-        }
-    }
-
-    let statement = Statement::new_text(client, "".to_owned(), parameters, columns);
-
-    Ok(RowStream::new(statement, responses))
+    let responses = start(client, buf).await?;
+    Ok(RowStream {
+        statement,
+        responses,
+        command_tag: None,
+        status: None,
+        output_format: Format::Text,
+        _p: PhantomPinned,
+    })
 }
 
 pub async fn query_portal(
@@ -164,6 +132,8 @@ pub async fn query_portal(
         statement: portal.statement().clone(),
         responses,
         command_tag: None,
+        status: None,
+        output_format: Format::Binary,
         _p: PhantomPinned,
     })
 }
@@ -298,23 +268,13 @@ pin_project! {
         statement: Statement,
         responses: Responses,
         command_tag: Option<String>,
+        output_format: Format,
+        status: Option<u8>,
         #[pin]
         _p: PhantomPinned,
     }
 }
 
-impl RowStream {
-    /// Creates a new `RowStream`.
-    pub fn new(statement: Statement, responses: Responses) -> Self {
-        RowStream {
-            statement,
-            responses,
-            command_tag: None,
-            _p: PhantomPinned,
-        }
-    }
-}
-
 impl Stream for RowStream {
     type Item = Result<Row, Error>;
 
@@ -323,15 +283,22 @@ impl Stream for RowStream {
         loop {
             match ready!(this.responses.poll_next(cx)?) {
                 Message::DataRow(body) => {
-                    return Poll::Ready(Some(Ok(Row::new(this.statement.clone(), body)?)))
+                    return Poll::Ready(Some(Ok(Row::new(
+                        this.statement.clone(),
+                        body,
+                        *this.output_format,
+                    )?)))
                 }
                 Message::EmptyQueryResponse | Message::PortalSuspended => {}
                 Message::CommandComplete(body) => {
                     if let Ok(tag) = body.tag() {
                         *this.command_tag = Some(tag.to_string());
                     }
                 }
-                Message::ReadyForQuery(_) => return Poll::Ready(None),
+                Message::ReadyForQuery(status) => {
+                    *this.status = Some(status.status());
+                    return Poll::Ready(None);
+                }
                 _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
             }
         }
@@ -345,4 +312,11 @@ impl RowStream {
     pub fn command_tag(&self) -> Option<String> {
         self.command_tag.clone()
     }
+
+    /// Returns if the connection is ready for querying, with the status of the connection.
+    ///
+    /// This might be available only after the stream has been exhausted.
+    pub fn ready_status(&self) -> Option<u8> {
+        self.status
+    }
 }

tokio-postgres/src/row.rs

@@ -98,6 +98,7 @@ where
 /// A row of data returned from the database by a query.
 pub struct Row {
     statement: Statement,
+    output_format: Format,
     body: DataRowBody,
     ranges: Vec<Option<Range<usize>>>,
 }
@@ -111,12 +112,17 @@ impl fmt::Debug for Row {
 }
 
 impl Row {
-    pub(crate) fn new(statement: Statement, body: DataRowBody) -> Result<Row, Error> {
+    pub(crate) fn new(
+        statement: Statement,
+        body: DataRowBody,
+        output_format: Format,
+    ) -> Result<Row, Error> {
         let ranges = body.ranges().collect().map_err(Error::parse)?;
         Ok(Row {
             statement,
             body,
             ranges,
+            output_format,
         })
     }
 
@@ -193,7 +199,7 @@ impl Row {
     ///
     /// Useful when using query_raw_txt() which sets text transfer mode
     pub fn as_text(&self, idx: usize) -> Result<Option<&str>, Error> {
-        if self.statement.output_format() == Format::Text {
+        if self.output_format == Format::Text {
             match self.col_buffer(idx) {
                 Some(raw) => {
                     FromSql::from_sql(&Type::TEXT, raw).map_err(|e| Error::from_sql(e, idx))

tokio-postgres/src/statement.rs

@@ -6,7 +6,6 @@ use postgres_protocol::{
     message::{backend::Field, frontend},
     Oid,
 };
-use postgres_types::Format;
 use std::{
     fmt,
     sync::{Arc, Weak},
@@ -17,7 +16,6 @@ struct StatementInner {
     name: String,
     params: Vec<Type>,
     columns: Vec<Column>,
-    output_format: Format,
 }
 
 impl Drop for StatementInner {
@@ -51,22 +49,6 @@ impl Statement {
             name,
             params,
             columns,
-            output_format: Format::Binary,
-        }))
-    }
-
-    pub(crate) fn new_text(
-        inner: &Arc<InnerClient>,
-        name: String,
-        params: Vec<Type>,
-        columns: Vec<Column>,
-    ) -> Statement {
-        Statement(Arc::new(StatementInner {
-            client: Arc::downgrade(inner),
-            name,
-            params,
-            columns,
-            output_format: Format::Text,
         }))
     }
 
@@ -83,11 +65,6 @@ impl Statement {
     pub fn columns(&self) -> &[Column] {
         &self.0.columns
     }
-
-    /// Returns output format for the statement.
-    pub fn output_format(&self) -> Format {
-        self.0.output_format
-    }
 }
 
 /// Information about a column of a query.

tokio-postgres/src/transaction.rs

@@ -150,13 +150,14 @@ impl<'a> Transaction<'a> {
     }
 
     /// Like `Client::query_raw_txt`.
-    pub async fn query_raw_txt<S, I>(&self, query: S, params: I) -> Result<RowStream, Error>
+    pub async fn query_raw_txt<T, S, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
     where
-        S: AsRef<str> + Sync + Send,
+        T: ?Sized + ToStatement,
+        S: AsRef<str>,
         I: IntoIterator<Item = Option<S>>,
-        I::IntoIter: ExactSizeIterator + Sync + Send,
+        I::IntoIter: ExactSizeIterator,
     {
-        self.client.query_raw_txt(query, params).await
+        self.client.query_raw_txt(statement, params).await
     }
 
     /// Like `Client::execute`.