Skip to main content

postgresql/
encode.rs

1//! Encoding Rust values into the PostgreSQL binary wire format.
2
3use bytes::{BufMut, BytesMut};
4use postgresql_protocol::oid::{
5    BOOLOID, BYTEAOID, DATEOID, FLOAT4OID, FLOAT8OID, INETOID, INT2OID, INT4OID, INT8OID, INTERVALOID, JSONBOID, Oid,
6    TEXTOID, TIMESTAMPOID, TIMESTAMPTZOID, UUIDOID, array_oid, array_oid_checked,
7};
8
9use crate::{error::Error, interval::Interval, json::Json};
10
11/// Whether an encoded value is SQL `NULL`.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum IsNull {
14    /// The value is SQL `NULL`.
15    Null,
16    /// The value is not null.
17    NotNull,
18}
19
20/// A value that can be bound as a query parameter in the binary format.
21///
22/// This trait is object-safe, so `&[&dyn ToSql]` is the parameter list the
23/// macros build. See [`StaticType`] for the compile-time type OID.
24pub trait ToSql: Send + Sync {
25    /// Appends the binary representation of this value to `buf`, without a
26    /// length prefix.
27    fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error>;
28
29    /// The PostgreSQL type OID of this value.
30    fn oid(&self) -> Oid;
31}
32
33/// A value with a statically known PostgreSQL type OID.
34///
35/// Used by the `query!` / `query_as!` macros to check parameter types at
36/// compile time. This trait is deliberately separate from [`ToSql`] because an
37/// associated constant would make `ToSql` not object-safe.
38pub trait StaticType {
39    /// The PostgreSQL type OID of this Rust type.
40    const OID: Oid;
41    /// Whether this Rust type can represent SQL `NULL` (i.e. it is `Option`).
42    const NULLABLE: bool = false;
43}
44
45macro_rules! impl_pg_type {
46    ($ty:ty, $oid:expr, |$this:ident, $buf:ident| $body:block) => {
47        impl StaticType for $ty {
48            const OID: Oid = $oid;
49        }
50
51        impl ToSql for $ty {
52            fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error> {
53                let $this = self;
54                let $buf = buf;
55                $body
56            }
57
58            fn oid(&self) -> Oid {
59                $oid
60            }
61        }
62    };
63}
64
65/// Appends a complete one-dimensional PostgreSQL array.
66pub(crate) fn encode_array<T: ToSql + StaticType>(
67    iter: impl Iterator<Item = T>,
68    element_oid: Oid,
69    buf: &mut BytesMut,
70) -> Result<IsNull, Error> {
71    buf.put_i32(1); // number of dimensions
72    let nulls_pos = buf.len();
73    buf.put_i32(0); // has-nulls flag, patched below
74    buf.put_u32(element_oid);
75    let dim_pos = buf.len();
76    buf.put_i32(0); // element count, patched below
77    buf.put_i32(1); // lower bound
78
79    let mut count = 0i32;
80    let mut has_nulls = false;
81    for item in iter {
82        let len_pos = buf.len();
83        buf.put_i32(0);
84        let start = buf.len();
85        match item.encode(buf)? {
86            IsNull::Null => {
87                buf.truncate(len_pos);
88                buf.put_i32(-1);
89                has_nulls = true;
90            }
91            IsNull::NotNull => {
92                let n = i32::try_from(buf.len() - start)
93                    .map_err(|_| Error::Encode("array element is larger than 2 GiB".into()))?;
94                buf[len_pos..len_pos + 4].copy_from_slice(&n.to_be_bytes());
95            }
96        }
97        count = count
98            .checked_add(1)
99            .ok_or_else(|| Error::Encode("array has too many elements".into()))?;
100    }
101
102    buf[nulls_pos..nulls_pos + 4].copy_from_slice(&(has_nulls as i32).to_be_bytes());
103    buf[dim_pos..dim_pos + 4].copy_from_slice(&count.to_be_bytes());
104    Ok(IsNull::NotNull)
105}
106
107/// Unix timestamp of the PostgreSQL epoch (2000-01-01T00:00:00Z).
108pub(crate) const PG_EPOCH_UNIX: i64 = 946_684_800;
109
110pub(crate) fn pg_epoch() -> Result<time::DateTime, Error> {
111    Ok(time::DateTime::from_unix(PG_EPOCH_UNIX)?)
112}
113
114impl_pg_type!(i16, INT2OID, |this, buf| {
115    buf.put_slice(&this.to_be_bytes());
116    Ok(IsNull::NotNull)
117});
118
119impl_pg_type!(i32, INT4OID, |this, buf| {
120    buf.put_slice(&this.to_be_bytes());
121    Ok(IsNull::NotNull)
122});
123
124impl_pg_type!(i64, INT8OID, |this, buf| {
125    buf.put_slice(&this.to_be_bytes());
126    Ok(IsNull::NotNull)
127});
128
129impl_pg_type!(f32, FLOAT4OID, |this, buf| {
130    buf.put_slice(&this.to_be_bytes());
131    Ok(IsNull::NotNull)
132});
133
134impl_pg_type!(f64, FLOAT8OID, |this, buf| {
135    buf.put_slice(&this.to_be_bytes());
136    Ok(IsNull::NotNull)
137});
138
139impl_pg_type!(bool, BOOLOID, |this, buf| {
140    buf.put_u8(u8::from(*this));
141    Ok(IsNull::NotNull)
142});
143
144impl_pg_type!(str, TEXTOID, |this, buf| {
145    buf.put_slice(this.as_bytes());
146    Ok(IsNull::NotNull)
147});
148
149impl_pg_type!(String, TEXTOID, |this, buf| {
150    buf.put_slice(this.as_bytes());
151    Ok(IsNull::NotNull)
152});
153
154impl_pg_type!([u8], BYTEAOID, |this, buf| {
155    buf.put_slice(this);
156    Ok(IsNull::NotNull)
157});
158
159impl_pg_type!(Vec<u8>, BYTEAOID, |this, buf| {
160    buf.put_slice(this);
161    Ok(IsNull::NotNull)
162});
163
164impl_pg_type!(uuid::Uuid, UUIDOID, |this, buf| {
165    buf.put_slice(&this.as_bytes()[..]);
166    Ok(IsNull::NotNull)
167});
168
169impl_pg_type!(time::DateTime, TIMESTAMPTZOID, |this, buf| {
170    let (secs, nanos) = this.signed_duration_since(pg_epoch()?);
171    let micros = secs
172        .checked_mul(1_000_000)
173        .and_then(|v| v.checked_add((nanos / 1000) as i64))
174        .ok_or_else(|| Error::Encode("timestamp is out of range for timestamptz".into()))?;
175    buf.put_i64(micros);
176    Ok(IsNull::NotNull)
177});
178
179/// Serializes any `serde` value as a `jsonb` payload.
180///
181/// Returns an error if the value cannot be serialized to JSON.
182pub fn json_encode<T: serde::Serialize + ?Sized>(value: &T, buf: &mut BytesMut) -> Result<IsNull, Error> {
183    buf.put_u8(1); // jsonb format version
184    // Serialize straight into the destination buffer instead of through an
185    // intermediate `Vec` (one fewer allocation and copy per parameter).
186    serde_json::to_writer(JsonWriter(buf), value)?;
187    Ok(IsNull::NotNull)
188}
189
190/// Adapts a `BytesMut` to [`std::io::Write`] so `serde_json` can serialize into
191/// it directly.
192///
193/// The `bytes` crate is built here without its `std` feature, so `BytesMut` has
194/// no `io::Write` implementation of its own.
195struct JsonWriter<'a>(&'a mut BytesMut);
196
197impl std::io::Write for JsonWriter<'_> {
198    fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
199        self.0.extend_from_slice(data);
200        Ok(data.len())
201    }
202
203    fn flush(&mut self) -> std::io::Result<()> {
204        Ok(())
205    }
206}
207
208impl_pg_type!(serde_json::Value, JSONBOID, |this, buf| { json_encode(this, buf) });
209
210impl<T: serde::Serialize + Send + Sync> ToSql for Json<T> {
211    fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error> {
212        json_encode(&self.0, buf)
213    }
214
215    fn oid(&self) -> Oid {
216        JSONBOID
217    }
218}
219
220impl<T> StaticType for Json<T> {
221    const OID: Oid = JSONBOID;
222}
223
224impl_pg_type!(ipnetwork::IpNetwork, INETOID, |this, buf| {
225    match this {
226        ipnetwork::IpNetwork::V4(net) => {
227            buf.put_u8(2);
228            buf.put_u8(net.prefix());
229            buf.put_u8(0); // not a cidr
230            buf.put_u8(4);
231            buf.put_slice(&net.ip().octets());
232        }
233        ipnetwork::IpNetwork::V6(net) => {
234            buf.put_u8(3);
235            buf.put_u8(net.prefix());
236            buf.put_u8(0);
237            buf.put_u8(16);
238            buf.put_slice(&net.ip().octets());
239        }
240    }
241    Ok(IsNull::NotNull)
242});
243
244impl_pg_type!(Interval, INTERVALOID, |this, buf| {
245    buf.put_i64(this.micros);
246    buf.put_i32(this.days);
247    buf.put_i32(this.months);
248    Ok(IsNull::NotNull)
249});
250
251impl<T: ToSql + StaticType + ?Sized> ToSql for &T {
252    fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error> {
253        (**self).encode(buf)
254    }
255
256    fn oid(&self) -> Oid {
257        (**self).oid()
258    }
259}
260
261impl<T: StaticType + ?Sized> StaticType for &T {
262    const OID: Oid = T::OID;
263    const NULLABLE: bool = T::NULLABLE;
264}
265
266impl<T: ToSql + StaticType> ToSql for Option<T> {
267    fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error> {
268        match self {
269            Some(value) => value.encode(buf),
270            None => Ok(IsNull::Null),
271        }
272    }
273
274    fn oid(&self) -> Oid {
275        T::OID
276    }
277}
278
279impl<T: StaticType> StaticType for Option<T> {
280    const OID: Oid = T::OID;
281    const NULLABLE: bool = true;
282}
283
284impl<T: ToSql + StaticType> ToSql for [T] {
285    fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error> {
286        encode_array(self.iter(), T::OID, buf)
287    }
288
289    fn oid(&self) -> Oid {
290        array_oid(T::OID).unwrap_or(0)
291    }
292}
293
294impl<T: StaticType> StaticType for [T] {
295    const OID: Oid = array_oid_checked(T::OID);
296}
297
298impl<T: ToSql + StaticType> ToSql for Vec<T> {
299    fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error> {
300        encode_array(self.iter(), T::OID, buf)
301    }
302
303    fn oid(&self) -> Oid {
304        array_oid(T::OID).unwrap_or(0)
305    }
306}
307
308impl<T: StaticType> StaticType for Vec<T> {
309    const OID: Oid = array_oid_checked(T::OID);
310}
311
312/// Encodes a `time::DateTime` as a zone-less `timestamp`.
313#[derive(Debug, Clone, Copy)]
314pub struct Timestamp(pub time::DateTime);
315
316impl_pg_type!(Timestamp, TIMESTAMPOID, |this, buf| {
317    let (secs, nanos) = this.0.signed_duration_since(pg_epoch()?);
318    let micros = secs
319        .checked_mul(1_000_000)
320        .and_then(|v| v.checked_add((nanos / 1000) as i64))
321        .ok_or_else(|| Error::Encode("timestamp is out of range".into()))?;
322    buf.put_i64(micros);
323    Ok(IsNull::NotNull)
324});
325
326/// Encodes a `time::DateTime` as a `date` (taken in UTC).
327#[derive(Debug, Clone, Copy)]
328pub struct Date(pub time::DateTime);
329
330impl_pg_type!(Date, DATEOID, |this, buf| {
331    let (secs, _) = this.0.signed_duration_since(pg_epoch()?);
332    let days = secs.div_euclid(86_400);
333    let days = i32::try_from(days).map_err(|_| Error::Encode("date is out of range".into()))?;
334    buf.put_i32(days);
335    Ok(IsNull::NotNull)
336});