Skip to content

Commit baf27a5

Browse files
levkklevkk
committed
support NULLs
1 parent 3d4b920 commit baf27a5

File tree

7 files changed

+85
-38
lines changed

7 files changed

+85
-38
lines changed

‎pgdog/src/net/messages/data_row.rs

+68-25
Original file line numberDiff line numberDiff line change
@@ -2,72 +2,110 @@
22
33
use super::{code, prelude::*, Datum, Format, FromDataType, Numeric, RowDescription};
44
use bytes::BytesMut;
5-
use std::ops::Deref;
5+
use std::ops::{Deref, DerefMut};
6+
7+
#[derive(Debug, Clone)]
8+
pub struct Data {
9+
data: Bytes,
10+
is_null: bool,
11+
}
12+
13+
impl Deref for Data {
14+
type Target = Bytes;
15+
16+
fn deref(&self) -> &Self::Target {
17+
&self.data
18+
}
19+
}
20+
21+
impl DerefMut for Data {
22+
fn deref_mut(&mut self) -> &mut Self::Target {
23+
&mut self.data
24+
}
25+
}
26+
27+
impl From<Bytes> for Data {
28+
fn from(value: Bytes) -> Self {
29+
Self {
30+
data: value,
31+
is_null: false,
32+
}
33+
}
34+
}
35+
36+
impl Data {
37+
pub fn null() -> Self {
38+
Self {
39+
data: Bytes::new(),
40+
is_null: true,
41+
}
42+
}
43+
}
644

745
/// DataRow message.
846
#[derive(Debug, Clone)]
947
pub struct DataRow {
10-
columns: Vec<Bytes>,
48+
columns: Vec<Data>,
1149
}
1250

1351
/// Convert value to data row column
1452
/// using text formatting.
1553
pub trait ToDataRowColumn {
16-
fn to_data_row_column(&self) -> Bytes;
54+
fn to_data_row_column(&self) -> Data;
1755
}
1856

1957
impl ToDataRowColumn for String {
20-
fn to_data_row_column(&self) -> Bytes {
21-
Bytes::copy_from_slice(self.as_bytes())
58+
fn to_data_row_column(&self) -> Data {
59+
Bytes::copy_from_slice(self.as_bytes()).into()
2260
}
2361
}
2462

2563
impl ToDataRowColumn for &String {
26-
fn to_data_row_column(&self) -> Bytes {
27-
Bytes::copy_from_slice(self.as_bytes())
64+
fn to_data_row_column(&self) -> Data {
65+
Bytes::copy_from_slice(self.as_bytes()).into()
2866
}
2967
}
3068

3169
impl ToDataRowColumn for &str {
32-
fn to_data_row_column(&self) -> Bytes {
33-
Bytes::copy_from_slice(self.as_bytes())
70+
fn to_data_row_column(&self) -> Data {
71+
Bytes::copy_from_slice(self.as_bytes()).into()
3472
}
3573
}
3674

3775
impl ToDataRowColumn for i64 {
38-
fn to_data_row_column(&self) -> Bytes {
39-
Bytes::copy_from_slice(self.to_string().as_bytes())
76+
fn to_data_row_column(&self) -> Data {
77+
Bytes::copy_from_slice(self.to_string().as_bytes()).into()
4078
}
4179
}
4280

4381
impl ToDataRowColumn for usize {
44-
fn to_data_row_column(&self) -> Bytes {
45-
Bytes::copy_from_slice(self.to_string().as_bytes())
82+
fn to_data_row_column(&self) -> Data {
83+
Bytes::copy_from_slice(self.to_string().as_bytes()).into()
4684
}
4785
}
4886

4987
impl ToDataRowColumn for u64 {
50-
fn to_data_row_column(&self) -> Bytes {
51-
Bytes::copy_from_slice(self.to_string().as_bytes())
88+
fn to_data_row_column(&self) -> Data {
89+
Bytes::copy_from_slice(self.to_string().as_bytes()).into()
5290
}
5391
}
5492

5593
impl ToDataRowColumn for bool {
56-
fn to_data_row_column(&self) -> Bytes {
57-
Bytes::copy_from_slice(if *self { b"t" } else { b"f" })
94+
fn to_data_row_column(&self) -> Data {
95+
Bytes::copy_from_slice(if *self { b"t" } else { b"f" }).into()
5896
}
5997
}
6098

6199
impl ToDataRowColumn for f64 {
62-
fn to_data_row_column(&self) -> Bytes {
100+
fn to_data_row_column(&self) -> Data {
63101
let number = format!("{:.5}", self);
64-
Bytes::copy_from_slice(number.as_bytes())
102+
Bytes::copy_from_slice(number.as_bytes()).into()
65103
}
66104
}
67105

68106
impl ToDataRowColumn for u128 {
69-
fn to_data_row_column(&self) -> Bytes {
70-
Bytes::copy_from_slice(self.to_string().as_bytes())
107+
fn to_data_row_column(&self) -> Data {
108+
Bytes::copy_from_slice(self.to_string().as_bytes()).into()
71109
}
72110
}
73111

@@ -93,7 +131,7 @@ impl DataRow {
93131
/// columns will be prefilled with NULLs.
94132
pub fn insert(&mut self, index: usize, value: impl ToDataRowColumn) -> &mut Self {
95133
while self.columns.len() <= index {
96-
self.columns.push(Bytes::new());
134+
self.columns.push(Data::null());
97135
}
98136
self.columns[index] = value.to_data_row_column();
99137
self
@@ -111,7 +149,7 @@ impl DataRow {
111149
/// Get data for column at index.
112150
#[inline]
113151
pub fn column(&self, index: usize) -> Option<Bytes> {
114-
self.columns.get(index).cloned()
152+
self.columns.get(index).cloned().map(|d| d.data)
115153
}
116154

117155
/// Get integer at index with text/binary encoding.
@@ -207,6 +245,7 @@ impl FromBytes for DataRow {
207245

208246
column.freeze()
209247
})
248+
.map(Data::from)
210249
.collect();
211250

212251
Ok(Self { columns })
@@ -219,8 +258,12 @@ impl ToBytes for DataRow {
219258
payload.put_i16(self.columns.len() as i16);
220259

221260
for column in &self.columns {
222-
payload.put_i32(column.len() as i32);
223-
payload.put(&column[..]);
261+
if column.is_null {
262+
payload.put_i32(-1);
263+
} else {
264+
payload.put_i32(column.len() as i32);
265+
payload.put(&column[..]);
266+
}
224267
}
225268

226269
Ok(payload.freeze())

‎pgdog/src/net/messages/data_types/interval.rs

+4-2
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
use std::num::ParseIntError;
22

3+
use crate::net::messages::data_row::Data;
4+
35
use super::*;
46
use bytes::Bytes;
57

@@ -32,8 +34,8 @@ impl Add for Interval {
3234
}
3335

3436
impl ToDataRowColumn for Interval {
35-
fn to_data_row_column(&self) -> Bytes {
36-
self.encode(Format::Text).unwrap()
37+
fn to_data_row_column(&self) -> Data {
38+
self.encode(Format::Text).unwrap().into()
3739
}
3840
}
3941

‎pgdog/src/net/messages/data_types/mod.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use std::ops::Add;
22

3-
use super::{bind::Format, Error, ToDataRowColumn};
3+
use super::{bind::Format, data_row::Data, Error, ToDataRowColumn};
44
use ::uuid::Uuid;
55
use bytes::Bytes;
66

@@ -48,7 +48,7 @@ pub enum Datum {
4848
}
4949

5050
impl ToDataRowColumn for Datum {
51-
fn to_data_row_column(&self) -> Bytes {
51+
fn to_data_row_column(&self) -> Data {
5252
use Datum::*;
5353

5454
match self {
@@ -61,7 +61,7 @@ impl ToDataRowColumn for Datum {
6161
TimestampTz(tz) => tz.to_data_row_column(),
6262
Uuid(uuid) => uuid.to_data_row_column(),
6363
Numeric(num) => num.to_data_row_column(),
64-
Null => Bytes::new(),
64+
Null => Data::null(),
6565
}
6666
}
6767
}

‎pgdog/src/net/messages/data_types/numeric.rs

+4-2
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ use std::{
66

77
use bytes::Buf;
88

9+
use crate::net::messages::data_row::Data;
10+
911
use super::*;
1012

1113
/// We don't expect NaN's so we're going to implement Ord for this below.
@@ -89,7 +91,7 @@ impl FromDataType for Numeric {
8991
}
9092

9193
impl ToDataRowColumn for Numeric {
92-
fn to_data_row_column(&self) -> Bytes {
93-
self.encode(Format::Text).unwrap()
94+
fn to_data_row_column(&self) -> Data {
95+
self.encode(Format::Text).unwrap().into()
9496
}
9597
}

‎pgdog/src/net/messages/data_types/timestamp.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@ pub struct Timestamp {
1717
}
1818

1919
impl ToDataRowColumn for Timestamp {
20-
fn to_data_row_column(&self) -> Bytes {
21-
self.encode(Format::Text).unwrap()
20+
fn to_data_row_column(&self) -> Data {
21+
self.encode(Format::Text).unwrap().into()
2222
}
2323
}
2424

‎pgdog/src/net/messages/data_types/timestamptz.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,8 @@ impl DerefMut for TimestampTz {
3737
}
3838

3939
impl ToDataRowColumn for TimestampTz {
40-
fn to_data_row_column(&self) -> Bytes {
41-
self.encode(Format::Text).unwrap()
40+
fn to_data_row_column(&self) -> Data {
41+
self.encode(Format::Text).unwrap().into()
4242
}
4343
}
4444

‎pgdog/src/net/messages/data_types/uuid.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ impl FromDataType for Uuid {
2424
}
2525

2626
impl ToDataRowColumn for Uuid {
27-
fn to_data_row_column(&self) -> Bytes {
28-
self.encode(Format::Text).unwrap()
27+
fn to_data_row_column(&self) -> Data {
28+
self.encode(Format::Text).unwrap().into()
2929
}
3030
}

0 commit comments

Comments
 (0)