sqlx_postgres/io/
mod.rs

1mod buf_mut;
2
3pub use buf_mut::PgBufMutExt;
4use std::fmt;
5use std::fmt::{Display, Formatter};
6use std::num::{NonZeroU32, Saturating};
7
8pub(crate) use sqlx_core::io::*;
9
10#[derive(Debug, Copy, Clone, PartialEq, Eq)]
11pub(crate) struct StatementId(IdInner);
12
13#[derive(Debug, Copy, Clone, PartialEq, Eq)]
14pub(crate) struct PortalId(IdInner);
15
16#[derive(Debug, Copy, Clone, PartialEq, Eq)]
17struct IdInner(Option<NonZeroU32>);
18
19pub(crate) struct DisplayId {
20    prefix: &'static str,
21    id: NonZeroU32,
22}
23
24impl StatementId {
25    #[allow(dead_code)]
26    pub const UNNAMED: Self = Self(IdInner::UNNAMED);
27
28    pub const NAMED_START: Self = Self(IdInner::NAMED_START);
29
30    #[cfg(test)]
31    pub const TEST_VAL: Self = Self(IdInner::TEST_VAL);
32
33    const NAME_PREFIX: &'static str = "sqlx_s_";
34
35    pub fn next(&self) -> Self {
36        Self(self.0.next())
37    }
38
39    pub fn name_len(&self) -> Saturating<usize> {
40        self.0.name_len(Self::NAME_PREFIX)
41    }
42
43    /// Get a type to format this statement ID with [`Display`].
44    ///
45    /// Returns `None` if this is the unnamed statement.
46    #[inline(always)]
47    pub fn display(&self) -> Option<DisplayId> {
48        self.0.display(Self::NAME_PREFIX)
49    }
50
51    pub fn put_name_with_nul(&self, buf: &mut Vec<u8>) {
52        self.0.put_name_with_nul(Self::NAME_PREFIX, buf)
53    }
54}
55
56impl Display for DisplayId {
57    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
58        write!(f, "{}{}", self.prefix, self.id)
59    }
60}
61
62#[allow(dead_code)]
63impl PortalId {
64    // None selects the unnamed portal
65    pub const UNNAMED: Self = PortalId(IdInner::UNNAMED);
66
67    pub const NAMED_START: Self = PortalId(IdInner::NAMED_START);
68
69    #[cfg(test)]
70    pub const TEST_VAL: Self = Self(IdInner::TEST_VAL);
71
72    const NAME_PREFIX: &'static str = "sqlx_p_";
73
74    /// If ID represents a named portal, return the next ID, wrapping on overflow.
75    ///
76    /// If this ID represents the unnamed portal, return the same.
77    pub fn next(&self) -> Self {
78        Self(self.0.next())
79    }
80
81    /// Calculate the number of bytes that will be written by [`Self::put_name_with_nul()`].
82    pub fn name_len(&self) -> Saturating<usize> {
83        self.0.name_len(Self::NAME_PREFIX)
84    }
85
86    pub fn put_name_with_nul(&self, buf: &mut Vec<u8>) {
87        self.0.put_name_with_nul(Self::NAME_PREFIX, buf)
88    }
89}
90
91impl IdInner {
92    const UNNAMED: Self = Self(None);
93
94    const NAMED_START: Self = Self(Some(NonZeroU32::MIN));
95
96    #[cfg(test)]
97    pub const TEST_VAL: Self = Self(NonZeroU32::new(1234567890));
98
99    #[inline(always)]
100    fn next(&self) -> Self {
101        Self(
102            self.0
103                .map(|id| id.checked_add(1).unwrap_or(NonZeroU32::MIN)),
104        )
105    }
106
107    #[inline(always)]
108    fn display(&self, prefix: &'static str) -> Option<DisplayId> {
109        self.0.map(|id| DisplayId { prefix, id })
110    }
111
112    #[inline(always)]
113    fn name_len(&self, name_prefix: &str) -> Saturating<usize> {
114        let mut len = Saturating(0);
115
116        if let Some(id) = self.0 {
117            len += name_prefix.len();
118            // estimate the length of the ID in decimal
119            // `.ilog10()` can't panic since the value is never zero
120            len += id.get().ilog10() as usize;
121            // add one to compensate for `ilog10()` rounding down.
122            len += 1;
123        }
124
125        // count the NUL terminator
126        len += 1;
127
128        len
129    }
130
131    #[inline(always)]
132    fn put_name_with_nul(&self, name_prefix: &str, buf: &mut Vec<u8>) {
133        if let Some(id) = self.0 {
134            buf.extend_from_slice(name_prefix.as_bytes());
135            buf.extend_from_slice(itoa::Buffer::new().format(id.get()).as_bytes());
136        }
137
138        buf.push(0);
139    }
140}
141
142#[test]
143fn statement_id_display_matches_encoding() {
144    const EXPECTED_STR: &str = "sqlx_s_1234567890";
145    const EXPECTED_BYTES: &[u8] = b"sqlx_s_1234567890\0";
146
147    let mut bytes = Vec::new();
148
149    StatementId::TEST_VAL.put_name_with_nul(&mut bytes);
150
151    assert_eq!(bytes, EXPECTED_BYTES);
152
153    let str = StatementId::TEST_VAL.display().unwrap().to_string();
154
155    assert_eq!(str, EXPECTED_STR);
156}
sqlx_postgres/io/mod.rs

sqlx_postgres/io/
mod.rs
