Skip to main content

pg/
encode.rs

1use crate::{
2    error::{PgError, Result},
3    types::PgType,
4};
5
6pub trait ToSql: Send + Sync {
7    fn to_sql(&self) -> Result<Vec<u8>>;
8    fn pg_type(&self) -> &'static PgType;
9}
10
11impl ToSql for i16 {
12    fn to_sql(&self) -> Result<Vec<u8>> {
13        Ok(self.to_be_bytes().to_vec())
14    }
15    fn pg_type(&self) -> &'static PgType {
16        &crate::types::INT2
17    }
18}
19
20impl ToSql for i32 {
21    fn to_sql(&self) -> Result<Vec<u8>> {
22        Ok(self.to_be_bytes().to_vec())
23    }
24    fn pg_type(&self) -> &'static PgType {
25        &crate::types::INT4
26    }
27}
28
29impl ToSql for i64 {
30    fn to_sql(&self) -> Result<Vec<u8>> {
31        Ok(self.to_be_bytes().to_vec())
32    }
33    fn pg_type(&self) -> &'static PgType {
34        &crate::types::INT8
35    }
36}
37
38impl ToSql for f32 {
39    fn to_sql(&self) -> Result<Vec<u8>> {
40        Ok(self.to_be_bytes().to_vec())
41    }
42    fn pg_type(&self) -> &'static PgType {
43        &crate::types::FLOAT4
44    }
45}
46
47impl ToSql for f64 {
48    fn to_sql(&self) -> Result<Vec<u8>> {
49        Ok(self.to_be_bytes().to_vec())
50    }
51    fn pg_type(&self) -> &'static PgType {
52        &crate::types::FLOAT8
53    }
54}
55
56impl ToSql for bool {
57    fn to_sql(&self) -> Result<Vec<u8>> {
58        Ok(vec![if *self { 1 } else { 0 }])
59    }
60    fn pg_type(&self) -> &'static PgType {
61        &crate::types::BOOL
62    }
63}
64
65impl ToSql for String {
66    fn to_sql(&self) -> Result<Vec<u8>> {
67        Ok(self.as_bytes().to_vec())
68    }
69    fn pg_type(&self) -> &'static PgType {
70        &crate::types::TEXT
71    }
72}
73
74impl ToSql for &str {
75    fn to_sql(&self) -> Result<Vec<u8>> {
76        Ok(self.as_bytes().to_vec())
77    }
78    fn pg_type(&self) -> &'static PgType {
79        &crate::types::TEXT
80    }
81}
82
83impl ToSql for uuid::Uuid {
84    fn to_sql(&self) -> Result<Vec<u8>> {
85        Ok(self.as_bytes().to_vec())
86    }
87    fn pg_type(&self) -> &'static PgType {
88        &crate::types::UUID
89    }
90}
91
92impl ToSql for chrono::DateTime<chrono::Utc> {
93    fn to_sql(&self) -> Result<Vec<u8>> {
94        let pg_epoch = chrono::NaiveDate::from_ymd_opt(2000, 1, 1)
95            .and_then(|d| d.and_hms_opt(0, 0, 0))
96            .map(|d| d.and_utc())
97            .unwrap();
98        let diff = *self - pg_epoch;
99        let micros = diff.num_microseconds().ok_or_else(|| {
100            PgError::Encode("timestamptz value out of range for PostgreSQL microsecond encoding".into())
101        })?;
102        Ok(micros.to_be_bytes().to_vec())
103    }
104    fn pg_type(&self) -> &'static PgType {
105        &crate::types::TIMESTAMPTZ
106    }
107}
108
109impl<T: ToSql> ToSql for Option<T> {
110    fn to_sql(&self) -> Result<Vec<u8>> {
111        match self {
112            Some(val) => val.to_sql(),
113            None => Ok(Vec::new()),
114        }
115    }
116    fn pg_type(&self) -> &'static PgType {
117        self.as_ref()
118            .and_then(|v| Some(v.pg_type()))
119            .unwrap_or(&crate::types::TEXT)
120    }
121}
122
123impl<T: ToSql> ToSql for Vec<T> {
124    fn to_sql(&self) -> Result<Vec<u8>> {
125        let elem_type = self.first().map(|e| e.pg_type()).unwrap_or(&crate::types::INT4);
126        let elem_oid = elem_type.oid;
127        let mut buf = Vec::new();
128        buf.extend_from_slice(&1i32.to_be_bytes());
129        buf.extend_from_slice(&0i32.to_be_bytes());
130        buf.extend_from_slice(&elem_oid.to_be_bytes());
131        buf.extend_from_slice(&(self.len() as i32).to_be_bytes());
132        buf.extend_from_slice(&1i32.to_be_bytes());
133        for elem in self {
134            let data = elem.to_sql()?;
135            buf.extend_from_slice(&(data.len() as i32).to_be_bytes());
136            buf.extend_from_slice(&data);
137        }
138        Ok(buf)
139    }
140    fn pg_type(&self) -> &'static PgType {
141        self.first()
142            .map(|e| crate::types::element_to_array(e.pg_type()))
143            .unwrap_or(&crate::types::INT4_ARRAY)
144    }
145}
146
147impl ToSql for Vec<u8> {
148    fn to_sql(&self) -> Result<Vec<u8>> {
149        Ok(self.clone())
150    }
151    fn pg_type(&self) -> &'static PgType {
152        &crate::types::BYTEA
153    }
154}
155
156impl ToSql for &[u8] {
157    fn to_sql(&self) -> Result<Vec<u8>> {
158        Ok(self.to_vec())
159    }
160    fn pg_type(&self) -> &'static PgType {
161        &crate::types::BYTEA
162    }
163}
164
165impl<T: ToSql> ToSql for &[T] {
166    fn to_sql(&self) -> Result<Vec<u8>> {
167        let elem_type = self.first().map(|e| e.pg_type()).unwrap_or(&crate::types::INT4);
168        let elem_oid = elem_type.oid;
169        let mut buf = Vec::new();
170        buf.extend_from_slice(&1i32.to_be_bytes());
171        buf.extend_from_slice(&0i32.to_be_bytes());
172        buf.extend_from_slice(&elem_oid.to_be_bytes());
173        buf.extend_from_slice(&(self.len() as i32).to_be_bytes());
174        buf.extend_from_slice(&1i32.to_be_bytes());
175        for elem in *self {
176            let data = elem.to_sql()?;
177            buf.extend_from_slice(&(data.len() as i32).to_be_bytes());
178            buf.extend_from_slice(&data);
179        }
180        Ok(buf)
181    }
182    fn pg_type(&self) -> &'static PgType {
183        self.first()
184            .map(|e| crate::types::element_to_array(e.pg_type()))
185            .unwrap_or(&crate::types::INT4_ARRAY)
186    }
187}
188
189/// Wraps an iterator to produce a PG array without allocating an intermediate `Vec`.
190/// Uses `size_hint` from the iterator for pre-allocation when available.
191///
192/// # Example
193/// ```ignore
194/// use pg::{BindIter, ToSql, types::UUID};
195/// let ids: Vec<uuid::Uuid> = vec![/* ... */];
196/// let param = BindIter::new(ids.iter().copied(), &UUID);
197/// conn.execute_raw("SELECT * FROM unnest($1::uuid[])", &[&param]).await?;
198/// ```
199pub struct BindIter<I> {
200    inner: std::sync::Mutex<Option<I>>,
201    elem_type: &'static PgType,
202}
203
204impl<I> BindIter<I> {
205    pub fn new(iter: I, elem_type: &'static PgType) -> Self {
206        BindIter {
207            inner: std::sync::Mutex::new(Some(iter)),
208            elem_type,
209        }
210    }
211}
212
213impl<I, T> ToSql for BindIter<I>
214where
215    I: Iterator<Item = T> + Send,
216    T: ToSql,
217{
218    fn to_sql(&self) -> Result<Vec<u8>> {
219        let mut buf = Vec::new();
220        buf.extend_from_slice(&1i32.to_be_bytes());
221        buf.extend_from_slice(&0i32.to_be_bytes());
222        buf.extend_from_slice(&self.elem_type.oid.to_be_bytes());
223        buf.extend_from_slice(&0i32.to_be_bytes());
224        buf.extend_from_slice(&1i32.to_be_bytes());
225
226        let per_elem_guess = match self.elem_type.oid {
227            crate::types::INT2OID => 4 + 2,
228            crate::types::INT4OID | crate::types::FLOAT4OID => 4 + 4,
229            crate::types::INT8OID | crate::types::FLOAT8OID | crate::types::TIMESTAMPTZOID => 4 + 8,
230            crate::types::UUIDOID => 4 + 16,
231            crate::types::BOOLOID => 4 + 1,
232            _ => 4 + 64,
233        };
234        if let Some(ref iter) = *self.inner.lock().unwrap() {
235            let (lower, _) = iter.size_hint();
236            if lower > 0 {
237                buf.reserve(20 + lower * per_elem_guess);
238            }
239        }
240
241        let mut count = 0i32;
242        let mut iter_guard = self.inner.lock().unwrap();
243        if let Some(ref mut iter) = *iter_guard {
244            while let Some(item) = iter.next() {
245                let data = item.to_sql()?;
246                buf.extend_from_slice(&(data.len() as i32).to_be_bytes());
247                buf.extend_from_slice(&data);
248                count += 1;
249            }
250        }
251
252        buf[12..16].copy_from_slice(&count.to_be_bytes());
253        Ok(buf)
254    }
255
256    fn pg_type(&self) -> &'static PgType {
257        crate::types::element_to_array(self.elem_type)
258    }
259}