Skip to main content

postgresql/
array.rs

1//! Zero-allocation binding of lazy iterators as PostgreSQL arrays.
2
3use std::sync::Mutex;
4
5use bytes::BytesMut;
6use postgresql_protocol::oid::{Oid, array_oid, array_oid_checked};
7
8use crate::{
9    encode::{IsNull, StaticType, ToSql, encode_array},
10    error::Error,
11};
12
13/// A single-use adapter that binds an iterator as a PostgreSQL array.
14///
15/// The iterator is consumed on the first encode. Wrapping a `Vec` or slice is
16/// unnecessary — those already implement [`ToSql`] as arrays.
17pub struct Array<I> {
18    inner: Mutex<Option<I>>,
19}
20
21impl<I> Array<I> {
22    /// Wraps an iterator.
23    pub fn new(iter: I) -> Self {
24        Array {
25            inner: Mutex::new(Some(iter)),
26        }
27    }
28}
29
30impl<I, T> ToSql for Array<I>
31where
32    I: Iterator<Item = T> + Send,
33    T: ToSql + StaticType,
34{
35    fn encode(&self, buf: &mut BytesMut) -> Result<IsNull, Error> {
36        let mut guard = self
37            .inner
38            .lock()
39            .map_err(|_| Error::Encode("array iterator mutex was poisoned".into()))?;
40        let iter = guard
41            .take()
42            .ok_or_else(|| Error::Encode("array iterator has already been consumed".into()))?;
43        encode_array(iter, T::OID, buf)
44    }
45
46    fn oid(&self) -> Oid {
47        array_oid(T::OID).unwrap_or(0)
48    }
49}
50
51impl<I, T> StaticType for Array<I>
52where
53    I: Iterator<Item = T> + Send,
54    T: ToSql + StaticType,
55{
56    const OID: Oid = array_oid_checked(T::OID);
57}
58
59/// Binds a lazy iterator as a one-dimensional PostgreSQL array.
60///
61/// Use this inside a query macro argument for `UNNEST` / `= ANY`:
62///
63/// ```ignore
64/// # async fn run(db: postgresql::Connection, uuids: Vec<uuid::Uuid>) -> Result<(), postgresql::Error> {
65/// # use postgresql::query;
66/// let rows = query!(
67///     "SELECT id, name FROM users WHERE id = ANY($1::uuid[])",
68///     postgresql::array(uuids.iter()))
69/// .fetch_all::<_, (uuid::Uuid, String)>(&db)
70/// .await?;
71/// # Ok(())
72/// # }
73/// ```
74pub fn array<I: IntoIterator>(iter: I) -> Array<I::IntoIter>
75where
76    I::Item: ToSql + StaticType,
77{
78    Array::new(iter.into_iter())
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use crate::decode::FromSql;
85
86    fn decode_array(bytes: &[u8]) -> Vec<i32> {
87        Vec::<i32>::from_sql(1007, bytes).unwrap()
88    }
89
90    #[test]
91    fn encodes_like_vec() {
92        let values = vec![1i32, 2, 3];
93        let mut from_vec = BytesMut::new();
94        values.encode(&mut from_vec).unwrap();
95
96        let mut from_iter = BytesMut::new();
97        array(values.iter()).encode(&mut from_iter).unwrap();
98
99        assert_eq!(&from_vec[..], &from_iter[..]);
100        assert_eq!(decode_array(&from_iter), values);
101    }
102
103    #[test]
104    fn encodes_empty() {
105        let empty: Vec<i32> = Vec::new();
106        let mut buf = BytesMut::new();
107        array(empty.iter()).encode(&mut buf).unwrap();
108        assert_eq!(decode_array(&buf), Vec::<i32>::new());
109        assert_eq!(buf.len(), 20);
110    }
111
112    #[test]
113    fn second_encode_errors() {
114        let values = vec![1i32];
115        let adapter = array(values.iter());
116        let mut buf = BytesMut::new();
117        adapter.encode(&mut buf).unwrap();
118        assert!(adapter.encode(&mut buf).is_err());
119    }
120}