tokio-postgres: prepare and execute unnamed statements in one roundtrip (#6)

* tokio-postgres: prepare and execute unnamed statements in one roundtrip

* cleanup

---------

Co-authored-by: Julius de Bruijn <julius@nauk.io>

Author: 2 people

postgres/src/client.rs

@@ -521,15 +521,6 @@ impl Client {
         CancelToken::new(self.client.cancel_token())
     }
 
-    /// Clears the client's type information cache.
-    ///
-    /// When user-defined types are used in a query, the client loads their definitions from the database and caches
-    /// them for the lifetime of the client. If those definitions are changed in the database, this method can be used
-    /// to flush the local cache and allow the new, updated definitions to be loaded.
-    pub fn clear_type_cache(&self) {
-        self.client.clear_type_cache();
-    }
-
     /// Determines if the client's connection has already closed.
     ///
     /// If this returns `true`, the client is no longer usable.

tokio-postgres/src/client.rs

@@ -9,7 +9,7 @@ use crate::simple_query::SimpleQueryStream;
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
 use crate::tls::TlsConnect;
-use crate::types::{Oid, ToSql, Type};
+use crate::types::{ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
 use crate::{
@@ -23,7 +23,6 @@ use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 use postgres_protocol::message::{backend::Message, frontend};
 use postgres_types::BorrowToSql;
-use std::collections::HashMap;
 use std::fmt;
 #[cfg(feature = "runtime")]
 use std::net::IpAddr;
@@ -63,8 +62,6 @@ impl Responses {
 
 pub struct InnerClient {
     sender: mpsc::UnboundedSender<Request>,
-    cached_typeinfo: Mutex<HashMap<Oid, Type>>,
-
     /// A buffer to use when writing out postgres commands.
     buffer: Mutex<BytesMut>,
 }
@@ -83,14 +80,6 @@ impl InnerClient {
         })
     }
 
-    pub fn type_(&self, oid: Oid) -> Option<Type> {
-        self.cached_typeinfo.lock().get(&oid).cloned()
-    }
-
-    pub fn clear_type_cache(&self) {
-        self.cached_typeinfo.lock().clear();
-    }
-
     /// Call the given function with a buffer to be used when writing out
     /// postgres commands.
     pub fn with_buf<F, R>(&self, f: F) -> R
@@ -146,7 +135,6 @@ impl Client {
         Client {
             inner: Arc::new(InnerClient {
                 sender,
-                cached_typeinfo: Default::default(),
                 buffer: Default::default(),
             }),
             #[cfg(feature = "runtime")]
@@ -322,19 +310,13 @@ 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, T, S, I>(
-        &self,
-        statement: &T,
-        params: I,
-    ) -> Result<RowStream, Error>
+    pub async fn query_raw_txt<'a, S, I>(&self, query: &str, params: I) -> Result<RowStream, Error>
     where
-        T: ?Sized + ToStatement,
         S: AsRef<str>,
         I: IntoIterator<Item = Option<S>>,
         I::IntoIter: ExactSizeIterator,
     {
-        let statement = statement.__convert().into_statement(self).await?;
-        query::query_txt(&self.inner, statement, params).await
+        query::query_txt(&self.inner, query, params).await
     }
 
     /// Executes a statement, returning the number of rows modified.
@@ -527,15 +509,6 @@ impl Client {
         self.cancel_token().cancel_query_raw(stream, tls).await
     }
 
-    /// Clears the client's type information cache.
-    ///
-    /// When user-defined types are used in a query, the client loads their definitions from the database and caches
-    /// them for the lifetime of the client. If those definitions are changed in the database, this method can be used
-    /// to flush the local cache and allow the new, updated definitions to be loaded.
-    pub fn clear_type_cache(&self) {
-        self.inner().clear_type_cache();
-    }
-
     /// Determines if the connection to the server has already closed.
     ///
     /// In that case, all future queries will fail.

tokio-postgres/src/generic_client.rs

@@ -57,13 +57,12 @@ pub trait GenericClient: private::Sealed {
         I::IntoIter: ExactSizeIterator;
 
     /// Like `Client::query_raw_txt`.
-    async fn query_raw_txt<'a, T, S, I>(
+    async fn query_raw_txt<'a,  S, I>(
         &self,
-        statement: &T,
+        statement: &str,
         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;
@@ -148,9 +147,8 @@ impl GenericClient for Client {
         self.query_raw(statement, params).await
     }
 
-    async fn query_raw_txt<'a, T, S, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
+    async fn query_raw_txt<'a,  S, I>(&self, statement: &str, 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,
@@ -244,9 +242,8 @@ impl GenericClient for Transaction<'_> {
         self.query_raw(statement, params).await
     }
 
-    async fn query_raw_txt<'a, T, S, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
+    async fn query_raw_txt<'a,  S, I>(&self, statement: &str, 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,

tokio-postgres/src/prepare.rs

@@ -47,7 +47,7 @@ pub async fn prepare(
     let mut parameters = vec![];
     let mut it = parameter_description.parameters();
     while let Some(oid) = it.next().map_err(Error::parse)? {
-        let type_ = get_type(client, oid).await?;
+        let type_ = get_type(oid);
         parameters.push(type_);
     }
 
@@ -67,13 +67,18 @@ pub async fn prepare(
     }
 
     if unnamed {
-        Ok(Statement::unnamed(query.to_owned(), parameters, columns))
+        Ok(Statement::unnamed(parameters, columns))
     } else {
         Ok(Statement::named(client, name, parameters, columns))
     }
 }
 
-fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Result<Bytes, Error> {
+pub(crate) fn encode(
+    client: &InnerClient,
+    name: &str,
+    query: &str,
+    types: &[Type],
+) -> Result<Bytes, Error> {
     if types.is_empty() {
         debug!("preparing query {}: {}", name, query);
     } else {
@@ -88,14 +93,10 @@ fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Resu
     })
 }
 
-pub async fn get_type(client: &Arc<InnerClient>, oid: Oid) -> Result<Type, Error> {
+pub fn get_type(oid: Oid) -> Type {
     if let Some(type_) = Type::from_oid(oid) {
-        return Ok(type_);
-    }
-
-    if let Some(type_) = client.type_(oid) {
-        return Ok(type_);
+        return type_;
     }
 
-    Ok(Type::TEXT)
+    Type::TEXT
 }

tokio-postgres/src/query.rs

@@ -2,12 +2,15 @@ use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::types::{BorrowToSql, IsNull};
-use crate::{Error, Portal, Row, Statement};
+use crate::{Column, 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::backend::{
+    CommandCompleteBody, Message, ParameterDescriptionBody, RowDescriptionBody,
+};
 use postgres_protocol::message::frontend;
 use postgres_types::Format;
 use std::fmt;
@@ -50,21 +53,22 @@ where
     } else {
         encode(client, &statement, params)?
     };
-    let responses = start(client, buf).await?;
+    let (statement, responses) = start(client, buf).await?;
     Ok(RowStream {
         statement,
         responses,
         rows_affected: None,
         command_tag: None,
         status: None,
         output_format: Format::Binary,
+        parameter_description: None,
         _p: PhantomPinned,
     })
 }
 
 pub async fn query_txt<S, I>(
     client: &Arc<InnerClient>,
-    statement: Statement,
+    query: &str,
     params: I,
 ) -> Result<RowStream, Error>
 where
@@ -75,10 +79,15 @@ where
     let params = params.into_iter();
 
     let buf = client.with_buf(|buf| {
+        // prepare
+        frontend::parse("", query, std::iter::empty(), buf).map_err(Error::encode)?;
+        frontend::describe(b'S', "", buf).map_err(Error::encode)?;
+        frontend::flush(buf);
+
         // Bind, pass params as text, retrieve as binary
         match frontend::bind(
             "",                 // empty string selects the unnamed portal
-            statement.name(),   // named prepared statement
+            "",                 // unnamed prepared statement
             std::iter::empty(), // all parameters use the default format (text)
             params,
             |param, buf| match param {
@@ -105,9 +114,10 @@ where
     })?;
 
     // now read the responses
-    let responses = start(client, buf).await?;
+    let (statement, responses) = start(client, buf).await?;
 
     Ok(RowStream {
+        parameter_description: None,
         statement,
         responses,
         command_tag: None,
@@ -132,7 +142,8 @@ pub async fn query_portal(
     let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
 
     Ok(RowStream {
-        statement: portal.statement().clone(),
+        parameter_description: None,
+        statement: Some(portal.statement().clone()),
         responses,
         rows_affected: None,
         command_tag: None,
@@ -176,7 +187,7 @@ where
     } else {
         encode(client, &statement, params)?
     };
-    let mut responses = start(client, buf).await?;
+    let (_statement, mut responses) = start(client, buf).await?;
 
     let mut rows = 0;
     loop {
@@ -192,19 +203,57 @@ where
     }
 }
 
-async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
+async fn start(client: &InnerClient, buf: Bytes) -> Result<(Option<Statement>, Responses), Error> {
+    let mut parameter_description: Option<ParameterDescriptionBody> = None;
+    let mut statement = None;
     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
 
-    match responses.next().await? {
-        Message::ParseComplete => match responses.next().await? {
-            Message::BindComplete => {}
+    loop {
+        match responses.next().await? {
+            Message::ParseComplete => {}
+            Message::BindComplete => return Ok((statement, responses)),
+            Message::ParameterDescription(body) => {
+                parameter_description = Some(body); // tooo-o-ooo-o loooove
+            }
+            Message::NoData => {
+                statement = Some(make_statement(parameter_description.take().unwrap(), None)?);
+            }
+            Message::RowDescription(body) => {
+                statement = Some(make_statement(
+                    parameter_description.take().unwrap(),
+                    Some(body),
+                )?);
+            }
             m => return Err(Error::unexpected_message(m)),
-        },
-        Message::BindComplete => {}
-        m => return Err(Error::unexpected_message(m)),
+        }
+    }
+}
+
+fn make_statement(
+    parameter_description: ParameterDescriptionBody,
+    row_description: Option<RowDescriptionBody>,
+) -> Result<Statement, Error> {
+    let mut parameters = vec![];
+    let mut it = parameter_description.parameters();
+
+    while let Some(oid) = it.next().map_err(Error::parse).unwrap() {
+        let type_ = crate::prepare::get_type(oid);
+        parameters.push(type_);
     }
 
-    Ok(responses)
+    let mut columns = Vec::new();
+
+    if let Some(row_description) = row_description {
+        let mut it = row_description.fields();
+
+        while let Some(field) = it.next().map_err(Error::parse)? {
+            let type_ = crate::prepare::get_type(field.type_oid());
+            let column = Column::new(field.name().to_string(), type_, field);
+            columns.push(column);
+        }
+    }
+
+    Ok(Statement::unnamed(parameters, columns))
 }
 
 pub fn encode<P, I>(client: &InnerClient, statement: &Statement, params: I) -> Result<Bytes, Error>
@@ -214,12 +263,10 @@ where
     I::IntoIter: ExactSizeIterator,
 {
     client.with_buf(|buf| {
-        if let Some(query) = statement.query() {
-            frontend::parse("", query, [], buf).unwrap();
-        }
         encode_bind(statement, params, "", buf)?;
         frontend::execute("", 0, buf).map_err(Error::encode)?;
         frontend::sync(buf);
+
         Ok(buf.split().freeze())
     })
 }
@@ -276,12 +323,14 @@ where
 pin_project! {
     /// A stream of table rows.
     pub struct RowStream {
-        statement: Statement,
+        statement: Option<Statement>,
         responses: Responses,
         rows_affected: Option<u64>,
         command_tag: Option<String>,
         output_format: Format,
         status: Option<u8>,
+        parameter_description: Option<ParameterDescriptionBody>,
+
         #[pin]
         _p: PhantomPinned,
     }
@@ -292,11 +341,12 @@ impl Stream for RowStream {
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.project();
+
         loop {
             match ready!(this.responses.poll_next(cx)?) {
                 Message::DataRow(body) => {
                     return Poll::Ready(Some(Ok(Row::new(
-                        this.statement.clone(),
+                        this.statement.as_ref().unwrap().clone(),
                         body,
                         *this.output_format,
                     )?)))