Switch query_raw_txt() to text format

That allows to grab postgres-provided serializations for types.

Author: kelvich

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

@@ -10,6 +10,7 @@ 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;
@@ -21,12 +22,11 @@ use crate::{
     CopyInSink, Error, Row, SimpleQueryMessage, Statement, ToStatement, Transaction,
     TransactionBuilder,
 };
-use bytes::{Buf, BytesMut, BufMut};
+use bytes::{Buf, BufMut, BytesMut};
 use fallible_iterator::FallibleIterator;
 use futures_channel::mpsc;
 use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
-use crate::statement::Column;
 use postgres_protocol::message::{backend::Message, frontend};
 use postgres_types::BorrowToSql;
 use std::collections::HashMap;
@@ -392,15 +392,15 @@ impl Client {
             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
+                "",                 // 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(1), // all binary
+                Some(0), // all text
                 buf,
             ) {
                 Ok(()) => Ok(()),
@@ -418,7 +418,9 @@ impl Client {
             Ok(buf.split().freeze())
         })?;
 
-        let mut responses = self.inner.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+        let mut responses = self
+            .inner
+            .send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
 
         // now read the responses
 
@@ -450,12 +452,7 @@ impl Client {
             }
         }
 
-        let statement = Statement::new(
-            &self.inner,
-            "".to_owned(),
-            parameters,
-            columns
-        );
+        let statement = Statement::new_text(&self.inner, "".to_owned(), parameters, columns);
 
         Ok(RowStream::new(statement, responses))
     }

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,22 @@ 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))
+        }
+    }
 }
 
 impl AsName for SimpleColumn {

tokio-postgres/src/statement.rs

@@ -3,6 +3,7 @@ use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::types::Type;
 use postgres_protocol::message::frontend;
+use postgres_types::Format;
 use std::{
     fmt,
     sync::{Arc, Weak},
@@ -13,6 +14,7 @@ struct StatementInner {
     name: String,
     params: Vec<Type>,
     columns: Vec<Column>,
+    output_format: Format,
 }
 
 impl Drop for StatementInner {
@@ -46,6 +48,22 @@ 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,
         }))
     }
 
@@ -62,6 +80,11 @@ 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/tests/test/main.rs

@@ -257,18 +257,23 @@ async fn query_raw_txt() {
 
     let rows: Vec<tokio_postgres::Row> = client
         .query_raw_txt("SELECT 55 * $1", ["42"])
-        .await.unwrap()
+        .await
+        .unwrap()
         .try_collect()
-        .await.unwrap();
+        .await
+        .unwrap();
 
     assert_eq!(rows.len(), 1);
-    assert_eq!(rows[0].get::<_, i32>(0), 55 * 42);
+    let res: i32 = rows[0].as_text(0).unwrap().parse::<i32>().unwrap();
+    assert_eq!(res, 55 * 42);
 
     let rows: Vec<tokio_postgres::Row> = client
         .query_raw_txt("SELECT $1", ["42"])
-        .await.unwrap()
+        .await
+        .unwrap()
         .try_collect()
-        .await.unwrap();
+        .await
+        .unwrap();
 
     assert_eq!(rows.len(), 1);
     assert_eq!(rows[0].get::<_, &str>(0), "42");