1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum IsNull {
14 Null,
16 NotNull,
18}
19
20pub trait ToSql: Send + Sync {
25 fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error>;
28
29 fn oid(&self) -> Oid;
31}
32
33pub trait StaticType {
39 const OID: Oid;
41 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
65pub(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); let nulls_pos = buf.len();
73 buf.put_i32(0); buf.put_u32(element_oid);
75 let dim_pos = buf.len();
76 buf.put_i32(0); buf.put_i32(1); 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
107pub(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
179pub fn json_encode<T: serde::Serialize + ?Sized>(value: &T, buf: &mut BytesMut) -> Result<IsNull, Error> {
183 buf.put_u8(1); serde_json::to_writer(JsonWriter(buf), value)?;
187 Ok(IsNull::NotNull)
188}
189
190struct 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); 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#[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#[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});