pgwire/
types.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::iter::TrustedLen;
16use std::ops::Index;
17use std::slice::Iter;
18
19use bytes::Bytes;
20
21use crate::error::{PsqlError, PsqlResult};
22
23/// A row of data returned from the database by a query.
24#[derive(Debug, Clone)]
25// NOTE: Since we only support simple query protocol, the values are represented as strings.
26pub struct Row(Vec<Option<Bytes>>);
27
28impl Row {
29    /// Create a row from values.
30    pub fn new(row: Vec<Option<Bytes>>) -> Self {
31        Self(row)
32    }
33
34    /// Returns the number of values in the row.
35    pub fn len(&self) -> usize {
36        self.0.len()
37    }
38
39    /// Returns `true` if the row contains no values. Required by clippy.
40    pub fn is_empty(&self) -> bool {
41        self.0.is_empty()
42    }
43
44    /// Returns the values.
45    pub fn values(&self) -> &[Option<Bytes>] {
46        &self.0
47    }
48
49    pub fn take(self) -> Vec<Option<Bytes>> {
50        self.0
51    }
52}
53
54impl Index<usize> for Row {
55    type Output = Option<Bytes>;
56
57    fn index(&self, index: usize) -> &Self::Output {
58        &self.0[index]
59    }
60}
61
62/// <https://www.postgresql.org/docs/current/protocol-overview.html#PROTOCOL-FORMAT-CODES>
63#[derive(Debug, Clone, Copy, PartialEq, Eq)]
64pub enum Format {
65    Binary,
66    Text,
67}
68
69impl Format {
70    pub fn from_i16(format_code: i16) -> PsqlResult<Self> {
71        match format_code {
72            0 => Ok(Format::Text),
73            1 => Ok(Format::Binary),
74            _ => Err(PsqlError::Uncategorized(
75                format!("Unknown format code: {}", format_code).into(),
76            )),
77        }
78    }
79}
80
81/// FormatIterator used to generate formats of actual length given the provided format.
82/// According Postgres Document: <https://www.postgresql.org/docs/current/protocol-message-formats.html#:~:text=The%20number%20of,number%20of%20parameters>
83/// - If the length of provided format is 0, all format will be default format(TEXT).
84/// - If the length of provided format is 1, all format will be the same as this only format.
85/// - If the length of provided format > 1, provided format should be the actual format.
86#[derive(Debug, Clone)]
87pub struct FormatIterator<'a, 'b>
88where
89    'a: 'b,
90{
91    _formats: &'a [Format],
92    format_iter: Iter<'b, Format>,
93    actual_len: usize,
94    default_format: Format,
95}
96
97impl<'a> FormatIterator<'a, '_> {
98    pub fn new(provided_formats: &'a [Format], actual_len: usize) -> Result<Self, String> {
99        if !provided_formats.is_empty()
100            && provided_formats.len() != 1
101            && provided_formats.len() != actual_len
102        {
103            return Err(format!(
104                "format codes length {} is not 0, 1 or equal to actual length {}",
105                provided_formats.len(),
106                actual_len
107            ));
108        }
109
110        let default_format = provided_formats.first().copied().unwrap_or(Format::Text);
111
112        Ok(Self {
113            _formats: provided_formats,
114            default_format,
115            format_iter: provided_formats.iter(),
116            actual_len,
117        })
118    }
119}
120
121impl Iterator for FormatIterator<'_, '_> {
122    type Item = Format;
123
124    fn next(&mut self) -> Option<Self::Item> {
125        if self.actual_len == 0 {
126            return None;
127        }
128
129        self.actual_len -= 1;
130
131        Some(
132            self.format_iter
133                .next()
134                .copied()
135                .unwrap_or(self.default_format),
136        )
137    }
138
139    fn size_hint(&self) -> (usize, Option<usize>) {
140        (self.actual_len, Some(self.actual_len))
141    }
142}
143
144impl ExactSizeIterator for FormatIterator<'_, '_> {}
145unsafe impl TrustedLen for FormatIterator<'_, '_> {}
pgwire/types.rs

pgwire/
types.rs
