sqlx_postgres/message/
data_row.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
use byteorder::{BigEndian, ByteOrder};
use sqlx_core::bytes::Bytes;
use std::ops::Range;

use crate::error::Error;
use crate::message::{BackendMessage, BackendMessageFormat};

/// A row of data from the database.
#[derive(Debug)]
pub struct DataRow {
    pub(crate) storage: Bytes,

    /// Ranges into the stored row data.
    /// This uses `u32` instead of usize to reduce the size of this type. Values cannot be larger
    /// than `i32` in postgres.
    pub(crate) values: Vec<Option<Range<u32>>>,
}

impl DataRow {
    #[inline]
    pub(crate) fn get(&self, index: usize) -> Option<&'_ [u8]> {
        self.values[index]
            .as_ref()
            .map(|col| &self.storage[(col.start as usize)..(col.end as usize)])
    }
}

impl BackendMessage for DataRow {
    const FORMAT: BackendMessageFormat = BackendMessageFormat::DataRow;

    fn decode_body(buf: Bytes) -> Result<Self, Error> {
        if buf.len() < 2 {
            return Err(err_protocol!(
                "expected at least 2 bytes, got {}",
                buf.len()
            ));
        }

        let cnt = BigEndian::read_u16(&buf) as usize;

        let mut values = Vec::with_capacity(cnt);
        let mut offset: u32 = 2;

        for _ in 0..cnt {
            let value_start = offset
                .checked_add(4)
                .ok_or_else(|| err_protocol!("next value start out of range (offset: {offset})"))?;

            // widen both to a larger type for a safe comparison
            if (buf.len() as u64) < (value_start as u64) {
                return Err(err_protocol!(
                    "expected 4 bytes at offset {offset}, got {}",
                    (value_start as u64) - (buf.len() as u64)
                ));
            }

            // Length of the column value, in bytes (this count does not include itself).
            // Can be zero. As a special case, -1 indicates a NULL column value.
            // No value bytes follow in the NULL case.
            //
            // we know `offset` is within range of `buf.len()` from the above check
            #[allow(clippy::cast_possible_truncation)]
            let length = BigEndian::read_i32(&buf[(offset as usize)..]);

            if let Ok(length) = u32::try_from(length) {
                let value_end = value_start.checked_add(length).ok_or_else(|| {
                    err_protocol!("value_start + length out of range ({offset} + {length})")
                })?;

                values.push(Some(value_start..value_end));
                offset = value_end;
            } else {
                // Negative values signify NULL
                values.push(None);
                // `value_start` is actually the next value now.
                offset = value_start;
            }
        }

        Ok(Self {
            storage: buf,
            values,
        })
    }
}

#[test]
fn test_decode_data_row() {
    const DATA: &[u8] = b"\
        \x00\x08\
        \xff\xff\xff\xff\
        \x00\x00\x00\x04\
        \x00\x00\x00\n\
        \xff\xff\xff\xff\
        \x00\x00\x00\x04\
        \x00\x00\x00\x14\
        \xff\xff\xff\xff\
        \x00\x00\x00\x04\
        \x00\x00\x00(\
        \xff\xff\xff\xff\
        \x00\x00\x00\x04\
        \x00\x00\x00P";

    let row = DataRow::decode_body(DATA.into()).unwrap();

    assert_eq!(row.values.len(), 8);

    assert!(row.get(0).is_none());
    assert_eq!(row.get(1).unwrap(), &[0_u8, 0, 0, 10][..]);
    assert!(row.get(2).is_none());
    assert_eq!(row.get(3).unwrap(), &[0_u8, 0, 0, 20][..]);
    assert!(row.get(4).is_none());
    assert_eq!(row.get(5).unwrap(), &[0_u8, 0, 0, 40][..]);
    assert!(row.get(6).is_none());
    assert_eq!(row.get(7).unwrap(), &[0_u8, 0, 0, 80][..]);
}

#[cfg(all(test, not(debug_assertions)))]
#[bench]
fn bench_data_row_get(b: &mut test::Bencher) {
    const DATA: &[u8] = b"\x00\x08\xff\xff\xff\xff\x00\x00\x00\x04\x00\x00\x00\n\xff\xff\xff\xff\x00\x00\x00\x04\x00\x00\x00\x14\xff\xff\xff\xff\x00\x00\x00\x04\x00\x00\x00(\xff\xff\xff\xff\x00\x00\x00\x04\x00\x00\x00P";

    let row = DataRow::decode_body(test::black_box(Bytes::from_static(DATA))).unwrap();

    b.iter(|| {
        let _value = test::black_box(&row).get(3);
    });
}

#[cfg(all(test, not(debug_assertions)))]
#[bench]
fn bench_decode_data_row(b: &mut test::Bencher) {
    const DATA: &[u8] = b"\x00\x08\xff\xff\xff\xff\x00\x00\x00\x04\x00\x00\x00\n\xff\xff\xff\xff\x00\x00\x00\x04\x00\x00\x00\x14\xff\xff\xff\xff\x00\x00\x00\x04\x00\x00\x00(\xff\xff\xff\xff\x00\x00\x00\x04\x00\x00\x00P";

    b.iter(|| {
        let _ = DataRow::decode_body(test::black_box(Bytes::from_static(DATA)));
    });
}