Add text protocol based query method (sfackler#14)

Add query_raw_txt client method

It takes all the extended protocol params as text and passes them to
postgres to sort out types. With that we can avoid situations when
postgres derived different type compared to what was passed in arguments.
There is also propare_typed method, but since we receive data in text
format anyway it makes more sense to avoid dealing with types in params.

This way we also can save on roundtrip and send Parse+Bind+Describe+Execute
right away without waiting for params description before Bind.

Use text protocol for responses -- that allows to grab
postgres-provided serializations for types.

Catch command tag.

Expose row buffer size and add `max_backend_message_size` option to
prevent handling and storing in memory large messages from the backend.

Co-authored-by: Arthur Petukhovsky <[email protected]>

Author: 2 people

.github/workflows/ci.yml

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

postgres-types/src/lib.rs

@@ -449,6 +449,22 @@ impl WrongType {
     }
 }
 
+/// An error indicating that a as_text conversion was attempted on a binary
+/// result.
+#[derive(Debug)]
+pub struct WrongFormat {}
+
+impl Error for WrongFormat {}
+
+impl fmt::Display for WrongFormat {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            fmt,
+            "cannot read column as text while it is in binary format"
+        )
+    }
+}
+
 /// A trait for types that can be created from a Postgres value.
 ///
 /// # Types
@@ -900,7 +916,7 @@ pub trait ToSql: fmt::Debug {
 /// Supported Postgres message format types
 ///
 /// Using Text format in a message assumes a Postgres `SERVER_ENCODING` of `UTF8`
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum Format {
     /// Text format (UTF-8)
     Text,

tokio-postgres/src/client.rs

@@ -4,8 +4,10 @@ use crate::connection::{Request, RequestMessages};
 use crate::copy_out::CopyOutStream;
 #[cfg(feature = "runtime")]
 use crate::keepalive::KeepaliveConfig;
+use crate::prepare::get_type;
 use crate::query::RowStream;
 use crate::simple_query::SimpleQueryStream;
+use crate::statement::Column;
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
 use crate::tls::TlsConnect;
@@ -16,7 +18,7 @@ use crate::{
     copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, CopyInSink, Error,
     Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
 };
-use bytes::{Buf, BytesMut};
+use bytes::{Buf, BufMut, BytesMut};
 use fallible_iterator::FallibleIterator;
 use futures_channel::mpsc;
 use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
@@ -368,6 +370,87 @@ impl Client {
         query::query(&self.inner, statement, params).await
     }
 
+    /// 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>
+    where
+        S: AsRef<str>,
+        I: IntoIterator<Item = S>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let params = params.into_iter();
+        let params_len = params.len();
+
+        let buf = self.inner.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
+                std::iter::empty(), // all parameters use the default format (text)
+                params,
+                |param, buf| {
+                    buf.put_slice(param.as_ref().as_bytes());
+                    Ok(postgres_protocol::IsNull::No)
+                },
+                Some(0), // all text
+                buf,
+            ) {
+                Ok(()) => Ok(()),
+                Err(frontend::BindError::Conversion(e)) => Err(Error::to_sql(e, 0)),
+                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
+            frontend::sync(buf);
+
+            Ok(buf.split().freeze())
+        })?;
+
+        let mut responses = self
+            .inner
+            .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)? {
+                let type_ = get_type(&self.inner, field.type_oid()).await?;
+                let column = Column::new(field.name().to_string(), type_);
+                columns.push(column);
+            }
+        }
+
+        let statement = Statement::new_text(&self.inner, "".to_owned(), parameters, columns);
+
+        Ok(RowStream::new(statement, responses))
+    }
+
     /// Executes a statement, returning the number of rows modified.
     ///
     /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list

tokio-postgres/src/codec.rs

@@ -35,7 +35,9 @@ impl FallibleIterator for BackendMessages {
     }
 }
 
-pub struct PostgresCodec;
+pub struct PostgresCodec {
+    pub max_message_size: Option<usize>,
+}
 
 impl Encoder<FrontendMessage> for PostgresCodec {
     type Error = io::Error;
@@ -64,6 +66,15 @@ impl Decoder for PostgresCodec {
                 break;
             }
 
+            if let Some(max) = self.max_message_size {
+                if len > max {
+                    return Err(io::Error::new(
+                        io::ErrorKind::InvalidInput,
+                        "message too large",
+                    ));
+                }
+            }
+
             match header.tag() {
                 backend::NOTICE_RESPONSE_TAG
                 | backend::NOTIFICATION_RESPONSE_TAG

tokio-postgres/src/config.rs

@@ -209,6 +209,8 @@ pub struct Config {
     pub(crate) target_session_attrs: TargetSessionAttrs,
     pub(crate) channel_binding: ChannelBinding,
     pub(crate) load_balance_hosts: LoadBalanceHosts,
+    pub(crate) replication_mode: Option<ReplicationMode>,
+    pub(crate) max_backend_message_size: Option<usize>,
 }
 
 impl Default for Config {
@@ -242,6 +244,8 @@ impl Config {
             target_session_attrs: TargetSessionAttrs::Any,
             channel_binding: ChannelBinding::Prefer,
             load_balance_hosts: LoadBalanceHosts::Disable,
+            replication_mode: None,
+            max_backend_message_size: None,
         }
     }
 
@@ -522,6 +526,17 @@ impl Config {
         self.load_balance_hosts
     }
 
+    /// Set limit for backend messages size.
+    pub fn max_backend_message_size(&mut self, max_backend_message_size: usize) -> &mut Config {
+        self.max_backend_message_size = Some(max_backend_message_size);
+        self
+    }
+
+    /// Get limit for backend messages size.
+    pub fn get_max_backend_message_size(&self) -> Option<usize> {
+        self.max_backend_message_size
+    }
+
     fn param(&mut self, key: &str, value: &str) -> Result<(), Error> {
         match key {
             "user" => {
@@ -658,6 +673,14 @@ impl Config {
                 };
                 self.load_balance_hosts(load_balance_hosts);
             }
+            "max_backend_message_size" => {
+                let limit = value.parse::<usize>().map_err(|_| {
+                    Error::config_parse(Box::new(InvalidValue("max_backend_message_size")))
+                })?;
+                if limit > 0 {
+                    self.max_backend_message_size(limit);
+                }
+            }
             key => {
                 return Err(Error::config_parse(Box::new(UnknownOption(
                     key.to_string(),

tokio-postgres/src/connect_raw.rs

@@ -92,7 +92,12 @@ where
     let stream = connect_tls(stream, config.ssl_mode, tls, has_hostname).await?;
 
     let mut stream = StartupStream {
-        inner: Framed::new(stream, PostgresCodec),
+        inner: Framed::new(
+            stream,
+            PostgresCodec {
+                max_message_size: config.max_backend_message_size,
+            },
+        ),
         buf: BackendMessages::empty(),
         delayed: VecDeque::new(),
     };

tokio-postgres/src/prepare.rs

@@ -131,7 +131,7 @@ fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Resu
     })
 }
 
-async fn get_type(client: &Arc<InnerClient>, oid: Oid) -> Result<Type, Error> {
+pub async fn get_type(client: &Arc<InnerClient>, oid: Oid) -> Result<Type, Error> {
     if let Some(type_) = Type::from_oid(oid) {
         return Ok(type_);
     }

tokio-postgres/src/query.rs

@@ -53,6 +53,7 @@ where
         statement,
         responses,
         rows_affected: None,
+        command_tag: None,
         _p: PhantomPinned,
     })
 }
@@ -74,6 +75,7 @@ pub async fn query_portal(
         statement: portal.statement().clone(),
         responses,
         rows_affected: None,
+        command_tag: None,
         _p: PhantomPinned,
     })
 }
@@ -208,11 +210,24 @@ pin_project! {
         statement: Statement,
         responses: Responses,
         rows_affected: Option<u64>,
+        command_tag: Option<String>,
         #[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>;
 
@@ -225,6 +240,10 @@ impl Stream for RowStream {
                 }
                 Message::CommandComplete(body) => {
                     *this.rows_affected = Some(extract_row_affected(&body)?);
+
+                    if let Ok(tag) = body.tag() {
+                        *this.command_tag = Some(tag.to_string());
+                    }
                 }
                 Message::EmptyQueryResponse | Message::PortalSuspended => {}
                 Message::ReadyForQuery(_) => return Poll::Ready(None),
@@ -241,4 +260,11 @@ impl RowStream {
     pub fn rows_affected(&self) -> Option<u64> {
         self.rows_affected
     }
+
+    /// Returns the command tag of this query.
+    ///
+    /// This is only available after the stream has been exhausted.
+    pub fn command_tag(&self) -> Option<String> {
+        self.command_tag.clone()
+    }
 }

tokio-postgres/src/row.rs

@@ -7,6 +7,7 @@ use crate::types::{FromSql, Type, WrongType};
 use crate::{Error, Statement};
 use fallible_iterator::FallibleIterator;
 use postgres_protocol::message::backend::DataRowBody;
+use postgres_types::{Format, WrongFormat};
 use std::fmt;
 use std::ops::Range;
 use std::str;
@@ -188,6 +189,27 @@ impl Row {
         let range = self.ranges[idx].to_owned()?;
         Some(&self.body.buffer()[range])
     }
+
+    /// Interpret the column at the given index as text
+    ///
+    /// 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 {
+            match self.col_buffer(idx) {
+                Some(raw) => {
+                    FromSql::from_sql(&Type::TEXT, raw).map_err(|e| Error::from_sql(e, idx))
+                }
+                None => Ok(None),
+            }
+        } else {
+            Err(Error::from_sql(Box::new(WrongFormat {}), idx))
+        }
+    }
+
+    /// Row byte size
+    pub fn body_len(&self) -> usize {
+        self.body.buffer().len()
+    }
 }
 
 impl AsName for SimpleColumn {