Skip to content

Commit f712d07

Browse files
author
Hunts Chen
committed
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
1 parent 0c05614 commit f712d07

File tree

8 files changed

+157
-17
lines changed

8 files changed

+157
-17
lines changed

postgres-protocol/src/message/backend.rs

+6-1
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ impl Message {
164164
DATA_ROW_TAG => {
165165
let len = buf.read_u16::<BigEndian>()?;
166166
let storage = buf.read_all();
167-
Message::DataRow(DataRowBody { storage, len })
167+
Message::DataRow(DataRowBody::new(storage, len))
168168
}
169169
ERROR_RESPONSE_TAG => {
170170
let storage = buf.read_all();
@@ -531,6 +531,11 @@ pub struct DataRowBody {
531531
}
532532

533533
impl DataRowBody {
534+
/// Constructs a new data row body.
535+
pub fn new(storage: Bytes, len: u16) -> Self {
536+
Self { storage, len }
537+
}
538+
534539
#[inline]
535540
pub fn ranges(&self) -> DataRowRanges<'_> {
536541
DataRowRanges {

postgres-types/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -735,6 +735,7 @@ impl<'a> FromSql<'a> for IpAddr {
735735
}
736736

737737
/// An enum representing the nullability of a Postgres value.
738+
#[derive(Debug, Eq, PartialEq)]
738739
pub enum IsNull {
739740
/// The value is NULL.
740741
Yes,

postgres/src/test.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use std::thread;
55
use std::time::Duration;
66
use tokio_postgres::error::SqlState;
77
use tokio_postgres::types::Type;
8-
use tokio_postgres::NoTls;
8+
use tokio_postgres::{NoTls, RowDescription};
99

1010
use super::*;
1111
use crate::binary_copy::{BinaryCopyInWriter, BinaryCopyOutIter};

tokio-postgres/src/lib.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -130,11 +130,11 @@ pub use crate::error::Error;
130130
pub use crate::generic_client::GenericClient;
131131
pub use crate::portal::Portal;
132132
pub use crate::query::RowStream;
133-
pub use crate::row::{Row, SimpleQueryRow};
133+
pub use crate::row::{Row, RowBuilder, SimpleQueryRow};
134134
pub use crate::simple_query::SimpleQueryStream;
135135
#[cfg(feature = "runtime")]
136136
pub use crate::socket::Socket;
137-
pub use crate::statement::{Column, Statement};
137+
pub use crate::statement::{Column, RowDescription, Statement};
138138
#[cfg(feature = "runtime")]
139139
use crate::tls::MakeTlsConnect;
140140
pub use crate::tls::NoTls;

tokio-postgres/src/query.rs

+4-3
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use postgres_protocol::message::frontend;
1212
use std::fmt;
1313
use std::marker::PhantomPinned;
1414
use std::pin::Pin;
15+
use std::sync::Arc;
1516
use std::task::{Context, Poll};
1617

1718
struct BorrowToSqlParamsDebug<'a, T>(&'a [T]);
@@ -50,7 +51,7 @@ where
5051
};
5152
let responses = start(client, buf).await?;
5253
Ok(RowStream {
53-
statement,
54+
statement: Arc::new(statement),
5455
responses,
5556
_p: PhantomPinned,
5657
})
@@ -70,7 +71,7 @@ pub async fn query_portal(
7071
let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
7172

7273
Ok(RowStream {
73-
statement: portal.statement().clone(),
74+
statement: Arc::new(portal.statement().clone()),
7475
responses,
7576
_p: PhantomPinned,
7677
})
@@ -200,7 +201,7 @@ where
200201
pin_project! {
201202
/// A stream of table rows.
202203
pub struct RowStream {
203-
statement: Statement,
204+
statement: Arc<Statement>,
204205
responses: Responses,
205206
#[pin]
206207
_p: PhantomPinned,

tokio-postgres/src/row.rs

+124-6
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,13 @@
22
33
use crate::row::sealed::{AsName, Sealed};
44
use crate::simple_query::SimpleColumn;
5-
use crate::statement::Column;
5+
use crate::statement::{Column, RowDescription};
66
use crate::types::{FromSql, Type, WrongType};
7-
use crate::{Error, Statement};
7+
use crate::Error;
8+
use bytes::{BufMut, BytesMut};
89
use fallible_iterator::FallibleIterator;
910
use postgres_protocol::message::backend::DataRowBody;
11+
use postgres_types::{IsNull, ToSql};
1012
use std::fmt;
1113
use std::ops::Range;
1214
use std::str;
@@ -96,7 +98,7 @@ where
9698

9799
/// A row of data returned from the database by a query.
98100
pub struct Row {
99-
statement: Statement,
101+
description: Arc<dyn RowDescription>,
100102
body: DataRowBody,
101103
ranges: Vec<Option<Range<usize>>>,
102104
}
@@ -110,18 +112,26 @@ impl fmt::Debug for Row {
110112
}
111113

112114
impl Row {
113-
pub(crate) fn new(statement: Statement, body: DataRowBody) -> Result<Row, Error> {
115+
pub(crate) fn new(
116+
description: Arc<dyn RowDescription>,
117+
body: DataRowBody,
118+
) -> Result<Row, Error> {
114119
let ranges = body.ranges().collect().map_err(Error::parse)?;
115120
Ok(Row {
116-
statement,
121+
description,
117122
body,
118123
ranges,
119124
})
120125
}
121126

127+
/// Returns description about the data in the row.
128+
pub fn description(&self) -> Arc<dyn RowDescription> {
129+
self.description.clone()
130+
}
131+
122132
/// Returns information about the columns of data in the row.
123133
pub fn columns(&self) -> &[Column] {
124-
self.statement.columns()
134+
self.description.columns()
125135
}
126136

127137
/// Determines if the row contains no values.
@@ -270,3 +280,111 @@ impl SimpleQueryRow {
270280
FromSql::from_sql_nullable(&Type::TEXT, buf).map_err(|e| Error::from_sql(e, idx))
271281
}
272282
}
283+
/// Builder for building a [`Row`].
284+
pub struct RowBuilder {
285+
desc: Arc<dyn RowDescription>,
286+
buf: BytesMut,
287+
n: usize,
288+
}
289+
290+
impl RowBuilder {
291+
/// Creates a new builder using the provided row description.
292+
pub fn new(desc: Arc<dyn RowDescription>) -> Self {
293+
Self {
294+
desc,
295+
buf: BytesMut::new(),
296+
n: 0,
297+
}
298+
}
299+
300+
/// Appends a column's value and returns a value indicates if this value should be represented
301+
/// as NULL.
302+
pub fn push(&mut self, value: Option<impl ToSql>) -> Result<IsNull, Error> {
303+
let columns = self.desc.columns();
304+
305+
if columns.len() == self.n {
306+
return Err(Error::column(
307+
"exceeded expected number of columns".to_string(),
308+
));
309+
}
310+
311+
let db_type = columns[self.n].type_();
312+
let start = self.buf.len();
313+
314+
// Reserve 4 bytes for the length of the binary data to be written
315+
self.buf.put_i32(-1i32);
316+
317+
let is_null = value
318+
.to_sql(db_type, &mut self.buf)
319+
.map_err(|e| Error::to_sql(e, self.n))?;
320+
321+
// Calculate the length of data just written.
322+
if is_null == IsNull::No {
323+
let len = (self.buf.len() - start - 4) as i32;
324+
// Update the length of data
325+
self.buf[start..start + 4].copy_from_slice(&len.to_be_bytes());
326+
};
327+
328+
self.n += 1;
329+
Ok(is_null)
330+
}
331+
332+
/// Builds the row.
333+
pub fn build(self) -> Result<Row, Error> {
334+
Row::new(
335+
self.desc.clone(),
336+
DataRowBody::new(self.buf.freeze(), self.n as u16),
337+
)
338+
}
339+
}
340+
341+
#[cfg(test)]
342+
mod tests {
343+
use postgres_types::IsNull;
344+
345+
use super::*;
346+
use std::net::IpAddr;
347+
348+
struct TestRowDescription {
349+
columns: Vec<Column>,
350+
}
351+
352+
impl RowDescription for TestRowDescription {
353+
fn columns(&self) -> &[Column] {
354+
&self.columns
355+
}
356+
}
357+
358+
#[test]
359+
fn test_row_builder() {
360+
let mut builder = RowBuilder::new(Arc::new(TestRowDescription {
361+
columns: vec![
362+
Column::new("id".to_string(), Type::INT8),
363+
Column::new("name".to_string(), Type::VARCHAR),
364+
Column::new("ip".to_string(), Type::INET),
365+
],
366+
}));
367+
368+
let expected_id = 1234i64;
369+
let is_null = builder.push(Some(expected_id)).unwrap();
370+
assert_eq!(IsNull::No, is_null);
371+
372+
let expected_name = "row builder";
373+
let is_null = builder.push(Some(expected_name)).unwrap();
374+
assert_eq!(IsNull::No, is_null);
375+
376+
let is_null = builder.push(None::<IpAddr>).unwrap();
377+
assert_eq!(IsNull::Yes, is_null);
378+
379+
let row = builder.build().unwrap();
380+
381+
let actual_id: i64 = row.try_get("id").unwrap();
382+
assert_eq!(expected_id, actual_id);
383+
384+
let actual_name: String = row.try_get("name").unwrap();
385+
assert_eq!(expected_name, actual_name);
386+
387+
let actual_dt: Option<IpAddr> = row.try_get("ip").unwrap();
388+
assert_eq!(None, actual_dt);
389+
}
390+
}

tokio-postgres/src/statement.rs

+17-3
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,12 @@ use std::{
88
sync::{Arc, Weak},
99
};
1010

11+
/// Describes the data (columns) in a row.
12+
pub trait RowDescription: Sync + Send {
13+
/// Returns information about the columns returned when the statement is queried.
14+
fn columns(&self) -> &[Column];
15+
}
16+
1117
struct StatementInner {
1218
client: Weak<InnerClient>,
1319
name: String,
@@ -57,9 +63,16 @@ impl Statement {
5763
pub fn params(&self) -> &[Type] {
5864
&self.0.params
5965
}
66+
}
6067

61-
/// Returns information about the columns returned when the statement is queried.
62-
pub fn columns(&self) -> &[Column] {
68+
impl RowDescription for Statement {
69+
fn columns(&self) -> &[Column] {
70+
&self.0.columns
71+
}
72+
}
73+
74+
impl RowDescription for Arc<Statement> {
75+
fn columns(&self) -> &[Column] {
6376
&self.0.columns
6477
}
6578
}
@@ -71,7 +84,8 @@ pub struct Column {
7184
}
7285

7386
impl Column {
74-
pub(crate) fn new(name: String, type_: Type) -> Column {
87+
/// Constructs a new column.
88+
pub fn new(name: String, type_: Type) -> Column {
7589
Column { name, type_ }
7690
}
7791

tokio-postgres/tests/test/main.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,8 @@ use tokio_postgres::error::SqlState;
1616
use tokio_postgres::tls::{NoTls, NoTlsStream};
1717
use tokio_postgres::types::{Kind, Type};
1818
use tokio_postgres::{
19-
AsyncMessage, Client, Config, Connection, Error, IsolationLevel, SimpleQueryMessage,
19+
AsyncMessage, Client, Config, Connection, Error, IsolationLevel, RowDescription,
20+
SimpleQueryMessage,
2021
};
2122

2223
mod binary_copy;

0 commit comments

Comments
 (0)