Introduce RowBuilder to support writing basic unit tests

Added RowDescription trait, and let rows to share the same description
rather than having a copy in each row (think when there are thousand of
them in the result).

Added RowBuilder to support adding stubs of row data in unit tests.

Currently, the library users have no chooice but have to use integration
tests for testing Postgres data access code. With the changes in this
commit, the `tokio-postgres` lib users can use RowBuilder to create
sutbs to verify the deserialization from database result (Rows) to
custom stucts in unit tests.

It can also serves as a base for future implementation of certain kind
of mocks of the db connection.

Related-to #910 #950

Author: hunts

postgres-protocol/src/message/backend.rs

@@ -164,7 +164,7 @@ impl Message {
             DATA_ROW_TAG => {
                 let len = buf.read_u16::<BigEndian>()?;
                 let storage = buf.read_all();
-                Message::DataRow(DataRowBody { storage, len })
+                Message::DataRow(DataRowBody::new(storage, len))
             }
             ERROR_RESPONSE_TAG => {
                 let storage = buf.read_all();
@@ -531,6 +531,11 @@ pub struct DataRowBody {
 }
 
 impl DataRowBody {
+    /// Constructs a new data row body.
+    pub fn new(storage: Bytes, len: u16) -> Self {
+        Self { storage, len }
+    }
+
     #[inline]
     pub fn ranges(&self) -> DataRowRanges<'_> {
         DataRowRanges {

postgres-types/src/lib.rs

@@ -735,6 +735,7 @@ impl<'a> FromSql<'a> for IpAddr {
 }
 
 /// An enum representing the nullability of a Postgres value.
+#[derive(Debug, Eq, PartialEq)]
 pub enum IsNull {
     /// The value is NULL.
     Yes,

postgres/src/test.rs

@@ -5,7 +5,7 @@ use std::thread;
 use std::time::Duration;
 use tokio_postgres::error::SqlState;
 use tokio_postgres::types::Type;
-use tokio_postgres::NoTls;
+use tokio_postgres::{NoTls, RowDescription};
 
 use super::*;
 use crate::binary_copy::{BinaryCopyInWriter, BinaryCopyOutIter};

tokio-postgres/src/lib.rs

@@ -130,11 +130,11 @@ pub use crate::error::Error;
 pub use crate::generic_client::GenericClient;
 pub use crate::portal::Portal;
 pub use crate::query::RowStream;
-pub use crate::row::{Row, SimpleQueryRow};
+pub use crate::row::{Row, RowBuilder, SimpleQueryRow};
 pub use crate::simple_query::SimpleQueryStream;
 #[cfg(feature = "runtime")]
 pub use crate::socket::Socket;
-pub use crate::statement::{Column, Statement};
+pub use crate::statement::{Column, RowDescription, Statement};
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
 pub use crate::tls::NoTls;

tokio-postgres/src/query.rs

@@ -12,6 +12,7 @@ use postgres_protocol::message::frontend;
 use std::fmt;
 use std::marker::PhantomPinned;
 use std::pin::Pin;
+use std::sync::Arc;
 use std::task::{Context, Poll};
 
 struct BorrowToSqlParamsDebug<'a, T>(&'a [T]);
@@ -50,7 +51,7 @@ where
     };
     let responses = start(client, buf).await?;
     Ok(RowStream {
-        statement,
+        statement: Arc::new(statement),
         responses,
         _p: PhantomPinned,
     })
@@ -70,7 +71,7 @@ pub async fn query_portal(
     let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
 
     Ok(RowStream {
-        statement: portal.statement().clone(),
+        statement: Arc::new(portal.statement().clone()),
         responses,
         _p: PhantomPinned,
     })
@@ -200,7 +201,7 @@ where
 pin_project! {
     /// A stream of table rows.
     pub struct RowStream {
-        statement: Statement,
+        statement: Arc<Statement>,
         responses: Responses,
         #[pin]
         _p: PhantomPinned,

tokio-postgres/src/row.rs

@@ -2,11 +2,13 @@
 
 use crate::row::sealed::{AsName, Sealed};
 use crate::simple_query::SimpleColumn;
-use crate::statement::Column;
+use crate::statement::{Column, RowDescription};
 use crate::types::{FromSql, Type, WrongType};
-use crate::{Error, Statement};
+use crate::Error;
+use bytes::{BufMut, BytesMut};
 use fallible_iterator::FallibleIterator;
 use postgres_protocol::message::backend::DataRowBody;
+use postgres_types::{IsNull, ToSql};
 use std::fmt;
 use std::ops::Range;
 use std::str;
@@ -96,7 +98,7 @@ where
 
 /// A row of data returned from the database by a query.
 pub struct Row {
-    statement: Statement,
+    description: Arc<dyn RowDescription>,
     body: DataRowBody,
     ranges: Vec<Option<Range<usize>>>,
 }
@@ -110,18 +112,26 @@ impl fmt::Debug for Row {
 }
 
 impl Row {
-    pub(crate) fn new(statement: Statement, body: DataRowBody) -> Result<Row, Error> {
+    pub(crate) fn new(
+        description: Arc<dyn RowDescription>,
+        body: DataRowBody,
+    ) -> Result<Row, Error> {
         let ranges = body.ranges().collect().map_err(Error::parse)?;
         Ok(Row {
-            statement,
+            description,
             body,
             ranges,
         })
     }
 
+    /// Returns description about the data in the row.
+    pub fn description(&self) -> Arc<dyn RowDescription> {
+        self.description.clone()
+    }
+
     /// Returns information about the columns of data in the row.
     pub fn columns(&self) -> &[Column] {
-        self.statement.columns()
+        self.description.columns()
     }
 
     /// Determines if the row contains no values.
@@ -270,3 +280,111 @@ impl SimpleQueryRow {
         FromSql::from_sql_nullable(&Type::TEXT, buf).map_err(|e| Error::from_sql(e, idx))
     }
 }
+/// Builder for building a [`Row`].
+pub struct RowBuilder {
+    desc: Arc<dyn RowDescription>,
+    buf: BytesMut,
+    n: usize,
+}
+
+impl RowBuilder {
+    /// Creates a new builder using the provided row description.
+    pub fn new(desc: Arc<dyn RowDescription>) -> Self {
+        Self {
+            desc,
+            buf: BytesMut::new(),
+            n: 0,
+        }
+    }
+
+    /// Appends a column's value and returns a value indicates if this value should be represented
+    /// as NULL.
+    pub fn push(&mut self, value: Option<impl ToSql>) -> Result<IsNull, Error> {
+        let columns = self.desc.columns();
+
+        if columns.len() == self.n {
+            return Err(Error::column(
+                "exceeded expected number of columns".to_string(),
+            ));
+        }
+
+        let db_type = columns[self.n].type_();
+        let start = self.buf.len();
+
+        // Reserve 4 bytes for the length of the binary data to be written
+        self.buf.put_i32(-1i32);
+
+        let is_null = value
+            .to_sql(db_type, &mut self.buf)
+            .map_err(|e| Error::to_sql(e, self.n))?;
+
+        // Calculate the length of data just written.
+        if is_null == IsNull::No {
+            let len = (self.buf.len() - start - 4) as i32;
+            // Update the length of data
+            self.buf[start..start + 4].copy_from_slice(&len.to_be_bytes());
+        };
+
+        self.n += 1;
+        Ok(is_null)
+    }
+
+    /// Builds the row.
+    pub fn build(self) -> Result<Row, Error> {
+        Row::new(
+            self.desc.clone(),
+            DataRowBody::new(self.buf.freeze(), self.n as u16),
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use postgres_types::IsNull;
+
+    use super::*;
+    use std::net::IpAddr;
+
+    struct TestRowDescription {
+        columns: Vec<Column>,
+    }
+
+    impl RowDescription for TestRowDescription {
+        fn columns(&self) -> &[Column] {
+            &self.columns
+        }
+    }
+
+    #[test]
+    fn test_row_builder() {
+        let mut builder = RowBuilder::new(Arc::new(TestRowDescription {
+            columns: vec![
+                Column::new("id".to_string(), Type::INT8),
+                Column::new("name".to_string(), Type::VARCHAR),
+                Column::new("ip".to_string(), Type::INET),
+            ],
+        }));
+
+        let expected_id = 1234i64;
+        let is_null = builder.push(Some(expected_id)).unwrap();
+        assert_eq!(IsNull::No, is_null);
+
+        let expected_name = "row builder";
+        let is_null = builder.push(Some(expected_name)).unwrap();
+        assert_eq!(IsNull::No, is_null);
+
+        let is_null = builder.push(None::<IpAddr>).unwrap();
+        assert_eq!(IsNull::Yes, is_null);
+
+        let row = builder.build().unwrap();
+
+        let actual_id: i64 = row.try_get("id").unwrap();
+        assert_eq!(expected_id, actual_id);
+
+        let actual_name: String = row.try_get("name").unwrap();
+        assert_eq!(expected_name, actual_name);
+
+        let actual_dt: Option<IpAddr> = row.try_get("ip").unwrap();
+        assert_eq!(None, actual_dt);
+    }
+}

tokio-postgres/src/statement.rs

@@ -8,6 +8,12 @@ use std::{
     sync::{Arc, Weak},
 };
 
+/// Describes the data (columns) in a row.
+pub trait RowDescription: Sync + Send {
+    /// Returns information about the columns returned when the statement is queried.
+    fn columns(&self) -> &[Column];
+}
+
 struct StatementInner {
     client: Weak<InnerClient>,
     name: String,
@@ -57,9 +63,16 @@ impl Statement {
     pub fn params(&self) -> &[Type] {
         &self.0.params
     }
+}
 
-    /// Returns information about the columns returned when the statement is queried.
-    pub fn columns(&self) -> &[Column] {
+impl RowDescription for Statement {
+    fn columns(&self) -> &[Column] {
+        &self.0.columns
+    }
+}
+
+impl RowDescription for Arc<Statement> {
+    fn columns(&self) -> &[Column] {
         &self.0.columns
     }
 }
@@ -71,7 +84,8 @@ pub struct Column {
 }
 
 impl Column {
-    pub(crate) fn new(name: String, type_: Type) -> Column {
+    /// Constructs a new column.
+    pub fn new(name: String, type_: Type) -> Column {
         Column { name, type_ }
     }
 

tokio-postgres/tests/test/main.rs

@@ -16,7 +16,8 @@ use tokio_postgres::error::SqlState;
 use tokio_postgres::tls::{NoTls, NoTlsStream};
 use tokio_postgres::types::{Kind, Type};
 use tokio_postgres::{
-    AsyncMessage, Client, Config, Connection, Error, IsolationLevel, SimpleQueryMessage,
+    AsyncMessage, Client, Config, Connection, Error, IsolationLevel, RowDescription,
+    SimpleQueryMessage,
 };
 
 mod binary_copy;