sqlx_postgres/message/
data_row.rs

1use byteorder::{BigEndian, ByteOrder};
2use sqlx_core::bytes::Bytes;
3use std::ops::Range;
4
5use crate::error::Error;
6use crate::message::{BackendMessage, BackendMessageFormat};
7
8/// A row of data from the database.
9#[derive(Debug)]
10pub struct DataRow {
11    pub(crate) storage: Bytes,
12
13    /// Ranges into the stored row data.
14    /// This uses `u32` instead of usize to reduce the size of this type. Values cannot be larger
15    /// than `i32` in postgres.
16    pub(crate) values: Vec<Option<Range<u32>>>,
17}
18
19impl DataRow {
20    #[inline]
21    pub(crate) fn get(&self, index: usize) -> Option<&'_ [u8]> {
22        self.values[index]
23            .as_ref()
24            .map(|col| &self.storage[(col.start as usize)..(col.end as usize)])
25    }
26}
27
28impl BackendMessage for DataRow {
29    const FORMAT: BackendMessageFormat = BackendMessageFormat::DataRow;
30
31    fn decode_body(buf: Bytes) -> Result<Self, Error> {
32        if buf.len() < 2 {
33            return Err(err_protocol!(
34                "expected at least 2 bytes, got {}",
35                buf.len()
36            ));
37        }
38
39        let cnt = BigEndian::read_u16(&buf) as usize;
40
41        let mut values = Vec::with_capacity(cnt);
42        let mut offset: u32 = 2;
43
44        for _ in 0..cnt {
45            let value_start = offset
46                .checked_add(4)
47                .ok_or_else(|| err_protocol!("next value start out of range (offset: {offset})"))?;
48
49            // widen both to a larger type for a safe comparison
50            if (buf.len() as u64) < (value_start as u64) {
51                return Err(err_protocol!(
52                    "expected 4 bytes at offset {offset}, got {}",
53                    (value_start as u64) - (buf.len() as u64)
54                ));
55            }
56
57            // Length of the column value, in bytes (this count does not include itself).
58            // Can be zero. As a special case, -1 indicates a NULL column value.
59            // No value bytes follow in the NULL case.
60            //
61            // we know `offset` is within range of `buf.len()` from the above check
62            #[allow(clippy::cast_possible_truncation)]
63            let length = BigEndian::read_i32(&buf[(offset as usize)..]);
64
65            if let Ok(length) = u32::try_from(length) {
66                let value_end = value_start.checked_add(length).ok_or_else(|| {
67                    err_protocol!("value_start + length out of range ({offset} + {length})")
68                })?;
69
70                values.push(Some(value_start..value_end));
71                offset = value_end;
72            } else {
73                // Negative values signify NULL
74                values.push(None);
75                // `value_start` is actually the next value now.
76                offset = value_start;
77            }
78        }
79
80        Ok(Self {
81            storage: buf,
82            values,
83        })
84    }
85}
86
87#[test]
88fn test_decode_data_row() {
89    const DATA: &[u8] = b"\
90        \x00\x08\
91        \xff\xff\xff\xff\
92        \x00\x00\x00\x04\
93        \x00\x00\x00\n\
94        \xff\xff\xff\xff\
95        \x00\x00\x00\x04\
96        \x00\x00\x00\x14\
97        \xff\xff\xff\xff\
98        \x00\x00\x00\x04\
99        \x00\x00\x00(\
100        \xff\xff\xff\xff\
101        \x00\x00\x00\x04\
102        \x00\x00\x00P";
103
104    let row = DataRow::decode_body(DATA.into()).unwrap();
105
106    assert_eq!(row.values.len(), 8);
107
108    assert!(row.get(0).is_none());
109    assert_eq!(row.get(1).unwrap(), &[0_u8, 0, 0, 10][..]);
110    assert!(row.get(2).is_none());
111    assert_eq!(row.get(3).unwrap(), &[0_u8, 0, 0, 20][..]);
112    assert!(row.get(4).is_none());
113    assert_eq!(row.get(5).unwrap(), &[0_u8, 0, 0, 40][..]);
114    assert!(row.get(6).is_none());
115    assert_eq!(row.get(7).unwrap(), &[0_u8, 0, 0, 80][..]);
116}
117
118#[cfg(all(test, not(debug_assertions)))]
119#[bench]
120fn bench_data_row_get(b: &mut test::Bencher) {
121    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";
122
123    let row = DataRow::decode_body(test::black_box(Bytes::from_static(DATA))).unwrap();
124
125    b.iter(|| {
126        let _value = test::black_box(&row).get(3);
127    });
128}
129
130#[cfg(all(test, not(debug_assertions)))]
131#[bench]
132fn bench_decode_data_row(b: &mut test::Bencher) {
133    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";
134
135    b.iter(|| {
136        let _ = DataRow::decode_body(test::black_box(Bytes::from_static(DATA)));
137    });
138}