Implement changes to support passing through raw query result rows

tokio-postgres/src/bind.rs

@@ -2,7 +2,7 @@ use crate::client::InnerClient;
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::types::BorrowToSql;
-use crate::{query, Error, Portal, Statement};
+use crate::{query, Error, Portal, Statement, DEFAULT_RESULT_FORMATS};
 use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -22,7 +22,7 @@ where
 {
     let name = format!("p{}", NEXT_ID.fetch_add(1, Ordering::SeqCst));
     let buf = client.with_buf(|buf| {
-        query::encode_bind(&statement, params, &name, buf)?;
+        query::encode_bind(&statement, params, &name, buf, DEFAULT_RESULT_FORMATS)?;
         frontend::sync(buf);
         Ok(buf.split().freeze())
     })?;

tokio-postgres/src/client.rs

@@ -18,6 +18,7 @@ use crate::Socket;
 use crate::{
     copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, CopyInSink, Error,
     Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
+    DEFAULT_RESULT_FORMATS,
 };
 use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
@@ -259,7 +260,7 @@ impl Client {
         T: ?Sized + ToStatement,
     {
         // try_collect will wait for the ReadyForQuery message which is not strictly necessary.
-        self.generic_query_raw(statement, slice_iter(params))
+        self.generic_query_raw(statement, slice_iter(params), DEFAULT_RESULT_FORMATS)
             .await?
             .try_collect()
             .await
@@ -389,20 +390,32 @@ impl Client {
         query::query(&self.inner, statement, params).await
     }
 
-    /// Executes a generic query
-    pub async fn generic_query_raw<T, P, I>(
+    /// Executes a generic query.
+    ///
+    /// Note that the result_formats parameter allows us to specify result formats other than the
+    /// default Some(1) (i.e. the default of using the binary format for everything) *but* this
+    /// crate is mostly written under the assumption that we only use binary format. As such,
+    /// trying to use functions like `Row::get` on the result will fail if the FromSql impl on the
+    /// result type tries to decode the result; FromSql doesn't pass in a result format to its
+    /// `from_sql` callback because it's assumed by the trait that we always use binary format. As
+    /// such, passing in non-binary formats here is only useful if we're *not* planning on decoding
+    /// the result (such as in a case where we are merely proxying the raw row to a different
+    /// client) or if we've written custom FromSql impls that expect text instead of binary format.
+    pub async fn generic_query_raw<T, P, I, J>(
         &self,
         statement: &T,
         params: I,
+        result_formats: J,
     ) -> Result<ResultStream, Error>
     where
         T: ?Sized + ToStatement,
         P: BorrowToSql,
         I: IntoIterator<Item = P>,
         I::IntoIter: ExactSizeIterator,
+        J: IntoIterator<Item = i16>,
     {
         let statement = statement.__convert().into_statement(self).await?;
-        query::generic_query(&self.inner, statement, params).await
+        query::generic_query(&self.inner, statement, params, result_formats).await
     }
 
     /// Executes a statement, returning the number of rows modified.

tokio-postgres/src/lib.rs

@@ -185,6 +185,15 @@ mod transaction;
 mod transaction_builder;
 pub mod types;
 
+/// This crate was originally written to use binary result format for pretty much everything, and
+/// it's still the case in most parts of the code that we hardcode to use binary. There are some
+/// small use cases we've hacked in where you can specify different result formats, though.
+/// Since Postgres lets you specify different result formats for different columns, we generally
+/// accept an IntoIterator<Item = i16> since each result format is specified as an i16 value of 0
+/// for text or 1 for binary. However, if you only specify a single format, Postgres interprets
+/// that as "use this format for everything", hence Some(1) being the default here:
+pub(crate) const DEFAULT_RESULT_FORMATS: Option<i16> = Some(1);
+
 /// A convenience function which parses a connection string and connects to the database.
 ///
 /// See the documentation for [`Config`] for details on the connection string format.

tokio-postgres/src/query.rs

@@ -2,7 +2,7 @@ use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::types::{BorrowToSql, IsNull};
-use crate::{Error, GenericResult, Portal, Row, Statement};
+use crate::{Error, GenericResult, Portal, Row, Statement, DEFAULT_RESULT_FORMATS};
 use bytes::{Bytes, BytesMut};
 use futures_util::{ready, Stream};
 use log::{debug, log_enabled, Level};
@@ -56,15 +56,17 @@ where
     })
 }
 
-pub async fn generic_query<P, I>(
+pub async fn generic_query<P, I, J>(
     client: &InnerClient,
     statement: Statement,
     params: I,
+    result_formats: J,
 ) -> Result<ResultStream, Error>
 where
     P: BorrowToSql,
     I: IntoIterator<Item = P>,
     I::IntoIter: ExactSizeIterator,
+    J: IntoIterator<Item = i16>,
 {
     let buf = if log_enabled!(Level::Debug) {
         let params = params.into_iter().collect::<Vec<_>>();
@@ -73,9 +75,9 @@ where
             statement.name(),
             BorrowToSqlParamsDebug(params.as_slice()),
         );
-        encode(client, &statement, params)?
+        encode_with_result_formats(client, &statement, params, result_formats)?
     } else {
-        encode(client, &statement, params)?
+        encode_with_result_formats(client, &statement, params, result_formats)?
     };
     let responses = start(client, buf).await?;
     Ok(ResultStream {
@@ -165,25 +167,42 @@ where
     P: BorrowToSql,
     I: IntoIterator<Item = P>,
     I::IntoIter: ExactSizeIterator,
+{
+    encode_with_result_formats(client, statement, params, DEFAULT_RESULT_FORMATS)
+}
+
+pub fn encode_with_result_formats<P, I, J>(
+    client: &InnerClient,
+    statement: &Statement,
+    params: I,
+    result_formats: J,
+) -> Result<Bytes, Error>
+where
+    P: BorrowToSql,
+    I: IntoIterator<Item = P>,
+    I::IntoIter: ExactSizeIterator,
+    J: IntoIterator<Item = i16>,
 {
     client.with_buf(|buf| {
-        encode_bind(statement, params, "", buf)?;
+        encode_bind(statement, params, "", buf, result_formats)?;
         frontend::execute("", 0, buf).map_err(Error::encode)?;
         frontend::sync(buf);
         Ok(buf.split().freeze())
     })
 }
 
-pub fn encode_bind<P, I>(
+pub fn encode_bind<P, I, J>(
     statement: &Statement,
     params: I,
     portal: &str,
     buf: &mut BytesMut,
+    result_formats: J,
 ) -> Result<(), Error>
 where
     P: BorrowToSql,
     I: IntoIterator<Item = P>,
     I::IntoIter: ExactSizeIterator,
+    J: IntoIterator<Item = i16>,
 {
     let param_types = statement.params();
     let params = params.into_iter();
@@ -216,7 +235,7 @@ where
                 Err(e)
             }
         },
-        Some(1),
+        result_formats,
         buf,
     );
     match r {

tokio-postgres/src/row.rs

@@ -140,6 +140,11 @@ impl Row {
         self.columns().len()
     }
 
+    /// Returns the DataRowBody which stores the row as Bytes.
+    pub fn body(&self) -> &DataRowBody {
+        &self.body
+    }
+
     /// Deserializes a value from the row.
     ///
     /// The value can be specified either by its numeric index in the row, or by its column name.