Add text protocol based query method

.github/workflows/ci.yml

@@ -57,7 +57,7 @@ jobs:
       - run: docker compose up -d
       - uses: sfackler/actions/rustup@master
         with:
-          version: 1.62.0
+          version: 1.65.0
       - run: echo "::set-output name=version::$(rustc --version)"
         id: rust-version
       - uses: actions/cache@v1

postgres-derive-test/src/lib.rs

@@ -14,7 +14,7 @@ where
     T: PartialEq + FromSqlOwned + ToSql + Sync,
     S: fmt::Display,
 {
-    for &(ref val, ref repr) in checks.iter() {
+    for (val, repr) in checks.iter() {
         let stmt = conn
             .prepare(&format!("SELECT {}::{}", *repr, sql_type))
             .unwrap();
@@ -38,7 +38,7 @@ pub fn test_type_asymmetric<T, F, S, C>(
     S: fmt::Display,
     C: Fn(&T, &F) -> bool,
 {
-    for &(ref val, ref repr) in checks.iter() {
+    for (val, repr) in checks.iter() {
         let stmt = conn
             .prepare(&format!("SELECT {}::{}", *repr, sql_type))
             .unwrap();

postgres-protocol/src/authentication/sasl.rs

@@ -389,7 +389,7 @@ impl<'a> Parser<'a> {
     }
 
     fn posit_number(&mut self) -> io::Result<u32> {
-        let n = self.take_while(|c| matches!(c, '0'..='9'))?;
+        let n = self.take_while(|c| c.is_ascii_digit())?;
         n.parse()
             .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))
     }

postgres-protocol/src/message/backend.rs

@@ -707,7 +707,7 @@ impl<'a> FallibleIterator for DataRowRanges<'a> {
                 ));
             }
             let base = self.len - self.buf.len();
-            self.buf = &self.buf[len as usize..];
+            self.buf = &self.buf[len..];
             Ok(Some(Some(base..base + len)))
         }
     }

postgres-types/src/lib.rs

@@ -395,6 +395,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
@@ -846,7 +862,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

@@ -7,8 +7,10 @@ use crate::copy_both::CopyBothDuplex;
 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;
@@ -20,7 +22,7 @@ use crate::{
     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};
@@ -374,6 +376,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

@@ -185,6 +185,7 @@ pub struct Config {
     pub(crate) target_session_attrs: TargetSessionAttrs,
     pub(crate) channel_binding: ChannelBinding,
     pub(crate) replication_mode: Option<ReplicationMode>,
+    pub(crate) max_backend_message_size: Option<usize>,
 }
 
 impl Default for Config {
@@ -217,6 +218,7 @@ impl Config {
             target_session_attrs: TargetSessionAttrs::Any,
             channel_binding: ChannelBinding::Prefer,
             replication_mode: None,
+            max_backend_message_size: None,
         }
     }
 
@@ -472,6 +474,17 @@ impl Config {
         self.replication_mode
     }
 
+    /// 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" => {
@@ -586,6 +599,14 @@ impl Config {
                     self.replication_mode(mode);
                 }
             }
+            "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

@@ -90,7 +90,12 @@ where
     let stream = connect_tls(stream, config.ssl_mode, tls).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

@@ -126,7 +126,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

@@ -52,6 +52,7 @@ where
     Ok(RowStream {
         statement,
         responses,
+        command_tag: None,
         _p: PhantomPinned,
     })
 }
@@ -72,6 +73,7 @@ pub async fn query_portal(
     Ok(RowStream {
         statement: portal.statement().clone(),
         responses,
+        command_tag: None,
         _p: PhantomPinned,
     })
 }
@@ -202,11 +204,24 @@ pin_project! {
     pub struct RowStream {
         statement: Statement,
         responses: Responses,
+        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>;
 
@@ -217,12 +232,24 @@ impl Stream for RowStream {
                 Message::DataRow(body) => {
                     return Poll::Ready(Some(Ok(Row::new(this.statement.clone(), body)?)))
                 }
-                Message::EmptyQueryResponse
-                | Message::CommandComplete(_)
-                | Message::PortalSuspended => {}
+                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),
                 _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
             }
         }
     }
 }
+
+impl RowStream {
+    /// 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;
@@ -187,6 +188,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 {