Implement support for 1-dim arrays for PostgreSQL

sqlx-core/src/postgres/types/array.rs

@@ -0,0 +1,211 @@
+/// Encoding and decoding of Postgres arrays. Documentation of the byte format can be found [here](https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/include/utils/array.h;h=7f7e744cb12bc872f628f90dad99dfdf074eb314;hb=master#l6)
+use crate::decode::Decode;
+use crate::decode::DecodeError;
+use crate::encode::Encode;
+use crate::io::{Buf, BufMut};
+use crate::postgres::database::Postgres;
+use crate::types::HasSqlType;
+use PhantomData;
+
+impl<T> Encode<Postgres> for [T]
+where
+    T: Encode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    fn encode(&self, buf: &mut Vec<u8>) {
+        let mut encoder = ArrayEncoder::new(buf);
+        for item in self {
+            encoder.push(item);
+        }
+    }
+}
+impl<T> Encode<Postgres> for Vec<T>
+where
+    [T]: Encode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    fn encode(&self, buf: &mut Vec<u8>) {
+        self.as_slice().encode(buf)
+    }
+}
+
+impl<T> Decode<Postgres> for Vec<T>
+where
+    T: Decode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    fn decode(buf: &[u8]) -> Result<Self, DecodeError> {
+        let decoder = ArrayDecoder::<T>::new(buf)?;
+        decoder.collect()
+    }
+}
+
+type Order = byteorder::BigEndian;
+
+struct ArrayDecoder<'a, T>
+where
+    T: Decode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    left: usize,
+    did_error: bool,
+
+    buf: &'a [u8],
+
+    phantom: PhantomData<T>,
+}
+
+impl<T> ArrayDecoder<'_, T>
+where
+    T: Decode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    fn new(mut buf: &[u8]) -> Result<ArrayDecoder<T>, DecodeError> {
+        let ndim = buf.get_i32::<Order>()?;
+        let dataoffset = buf.get_i32::<Order>()?;
+        let elemtype = buf.get_i32::<Order>()?;
+
+        if ndim == 0 {
+            return Ok(ArrayDecoder {
+                left: 0,
+                did_error: false,
+                buf,
+                phantom: PhantomData,
+            });
+        }
+
+        assert_eq!(ndim, 1, "only arrays of dimension 1 is supported");
+
+        let dimensions = buf.get_i32::<Order>()?;
+        let lower_bnds = buf.get_i32::<Order>()?;
+
+        assert_eq!(dataoffset, 0, "arrays with [null bitmap] is not supported");
+        assert_eq!(
+            elemtype,
+            <Postgres as HasSqlType<T>>::type_info().id.0 as i32,
+            "mismatched array element type"
+        );
+        assert_eq!(lower_bnds, 1);
+
+        Ok(ArrayDecoder {
+            left: dimensions as usize,
+            did_error: false,
+            buf,
+
+            phantom: PhantomData,
+        })
+    }
+
+    /// Decodes the next element without worring how many are left, or if it previously errored
+    fn decode_next_element(&mut self) -> Result<T, DecodeError> {
+        let len = self.buf.get_i32::<Order>()?;
+        let bytes = self.buf.get_bytes(len as usize)?;
+        Decode::decode(bytes)
+    }
+}
+
+impl<T> Iterator for ArrayDecoder<'_, T>
+where
+    T: Decode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    type Item = Result<T, DecodeError>;
+
+    fn next(&mut self) -> Option<Result<T, DecodeError>> {
+        if self.did_error || self.left == 0 {
+            return None;
+        }
+
+        self.left -= 1;
+
+        let decoded = self.decode_next_element();
+        self.did_error = decoded.is_err();
+        Some(decoded)
+    }
+}
+
+struct ArrayEncoder<'a, T>
+where
+    T: Encode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    count: usize,
+    len_start_index: usize,
+    buf: &'a mut Vec<u8>,
+
+    phantom: PhantomData<T>,
+}
+
+impl<T> ArrayEncoder<'_, T>
+where
+    T: Encode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    fn new(buf: &mut Vec<u8>) -> ArrayEncoder<T> {
+        let ty = <Postgres as HasSqlType<T>>::type_info();
+
+        // ndim
+        buf.put_i32::<Order>(1);
+        // dataoffset
+        buf.put_i32::<Order>(0);
+        // elemtype
+        buf.put_i32::<Order>(ty.id.0 as i32);
+        let len_start_index = buf.len();
+        // dimensions
+        buf.put_i32::<Order>(0);
+        // lower_bnds
+        buf.put_i32::<Order>(1);
+
+        ArrayEncoder {
+            count: 0,
+            len_start_index,
+            buf,
+
+            phantom: PhantomData,
+        }
+    }
+    fn push(&mut self, item: &T) {
+        // Allocate space for the length of the encoded elemement up front
+        let el_len_index = self.buf.len();
+        self.buf.put_i32::<Order>(0);
+
+        // Allocate the element it self
+        let el_start = self.buf.len();
+        Encode::encode(item, self.buf);
+        let el_end = self.buf.len();
+
+        // Now we know the actual length of the encoded element
+        let el_len = el_end - el_start;
+
+        // And we can now go back and update the length
+        self.buf[el_len_index..el_start].copy_from_slice(&(el_len as i32).to_be_bytes());
+
+        self.count += 1;
+    }
+    fn extend<'a, I>(&mut self, items: I)
+    where
+        I: Iterator<Item = &'a T>,
+        T: 'a,
+    {
+        for item in items {
+            self.push(item);
+        }
+    }
+    fn update_len(&mut self) {
+        const I32_SIZE: usize = std::mem::size_of::<i32>();
+
+        let size_bytes = (self.count as i32).to_be_bytes();
+
+        self.buf[self.len_start_index..self.len_start_index + I32_SIZE]
+            .copy_from_slice(&size_bytes);
+    }
+}
+impl<T> Drop for ArrayEncoder<'_, T>
+where
+    T: Encode<Postgres>,
+    Postgres: HasSqlType<T>,
+{
+    fn drop(&mut self) {
+        self.update_len();
+    }
+}

sqlx-core/src/postgres/types/bool.rs

@@ -16,6 +16,11 @@ impl HasSqlType<[bool]> for Postgres {
         PgTypeInfo::new(TypeId::ARRAY_BOOL)
     }
 }
+impl HasSqlType<Vec<bool>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Self as HasSqlType<[bool]>>::type_info()
+    }
+}
 
 impl Encode<Postgres> for bool {
     fn encode(&self, buf: &mut Vec<u8>) {

sqlx-core/src/postgres/types/bytes.rs

@@ -17,6 +17,12 @@ impl HasSqlType<[&'_ [u8]]> for Postgres {
     }
 }
 
+impl HasSqlType<Vec<&'_ [u8]>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Postgres as HasSqlType<[&'_ [u8]]>>::type_info()
+    }
+}
+
 // TODO: Do we need the [HasSqlType] here on the Vec?
 impl HasSqlType<Vec<u8>> for Postgres {
     fn type_info() -> PgTypeInfo {

sqlx-core/src/postgres/types/chrono.rs

@@ -55,6 +55,24 @@ impl HasSqlType<[NaiveDateTime]> for Postgres {
     }
 }
 
+impl HasSqlType<Vec<NaiveTime>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Postgres as HasSqlType<[NaiveTime]>>::type_info()
+    }
+}
+
+impl HasSqlType<Vec<NaiveDate>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Postgres as HasSqlType<[NaiveDate]>>::type_info()
+    }
+}
+
+impl HasSqlType<Vec<NaiveDateTime>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Postgres as HasSqlType<[NaiveDateTime]>>::type_info()
+    }
+}
+
 impl<Tz> HasSqlType<[DateTime<Tz>]> for Postgres
 where
     Tz: TimeZone,

sqlx-core/src/postgres/types/float.rs

@@ -16,6 +16,11 @@ impl HasSqlType<[f32]> for Postgres {
         PgTypeInfo::new(TypeId::ARRAY_FLOAT4)
     }
 }
+impl HasSqlType<Vec<f32>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Self as HasSqlType<[f32]>>::type_info()
+    }
+}
 
 impl Encode<Postgres> for f32 {
     fn encode(&self, buf: &mut Vec<u8>) {
@@ -42,6 +47,11 @@ impl HasSqlType<[f64]> for Postgres {
         PgTypeInfo::new(TypeId::ARRAY_FLOAT8)
     }
 }
+impl HasSqlType<Vec<f64>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Self as HasSqlType<[f64]>>::type_info()
+    }
+}
 
 impl Encode<Postgres> for f64 {
     fn encode(&self, buf: &mut Vec<u8>) {

sqlx-core/src/postgres/types/int.rs

@@ -18,6 +18,11 @@ impl HasSqlType<[i16]> for Postgres {
         PgTypeInfo::new(TypeId::ARRAY_INT2)
     }
 }
+impl HasSqlType<Vec<i16>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Self as HasSqlType<[i16]>>::type_info()
+    }
+}
 
 impl Encode<Postgres> for i16 {
     fn encode(&self, buf: &mut Vec<u8>) {
@@ -42,6 +47,11 @@ impl HasSqlType<[i32]> for Postgres {
         PgTypeInfo::new(TypeId::ARRAY_INT4)
     }
 }
+impl HasSqlType<Vec<i32>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Self as HasSqlType<[i32]>>::type_info()
+    }
+}
 
 impl Encode<Postgres> for i32 {
     fn encode(&self, buf: &mut Vec<u8>) {
@@ -66,6 +76,11 @@ impl HasSqlType<[i64]> for Postgres {
         PgTypeInfo::new(TypeId::ARRAY_INT8)
     }
 }
+impl HasSqlType<Vec<i64>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Self as HasSqlType<[i64]>>::type_info()
+    }
+}
 
 impl Encode<Postgres> for i64 {
     fn encode(&self, buf: &mut Vec<u8>) {

sqlx-core/src/postgres/types/mod.rs

@@ -1,3 +1,4 @@
+mod array;
 mod bool;
 mod bytes;
 mod float;

sqlx-core/src/postgres/types/str.rs

@@ -18,6 +18,11 @@ impl HasSqlType<[&'_ str]> for Postgres {
         PgTypeInfo::new(TypeId::ARRAY_TEXT)
     }
 }
+impl HasSqlType<Vec<&'_ str>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Self as HasSqlType<[&'_ str]>>::type_info()
+    }
+}
 
 // TODO: Do we need [HasSqlType] on String here?
 impl HasSqlType<String> for Postgres {
@@ -25,6 +30,11 @@ impl HasSqlType<String> for Postgres {
         <Self as HasSqlType<str>>::type_info()
     }
 }
+impl HasSqlType<Vec<String>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Self as HasSqlType<Vec<&'_ str>>>::type_info()
+    }
+}
 
 impl Encode<Postgres> for str {
     fn encode(&self, buf: &mut Vec<u8>) {

sqlx-core/src/postgres/types/uuid.rs

@@ -19,6 +19,12 @@ impl HasSqlType<[Uuid]> for Postgres {
     }
 }
 
+impl HasSqlType<Vec<Uuid>> for Postgres {
+    fn type_info() -> PgTypeInfo {
+        <Postgres as HasSqlType<[Uuid]>>::type_info()
+    }
+}
+
 impl Encode<Postgres> for Uuid {
     fn encode(&self, buf: &mut Vec<u8>) {
         buf.extend_from_slice(self.as_bytes());

tests/postgres-types.rs

@@ -37,6 +37,12 @@ test!(postgres_double: f64: "939399419.1225182::double precision" == 939399419.1
 
 test!(postgres_text: String: "'this is foo'" == "this is foo", "''" == "");
 
+test!(postgres_int_vec: Vec<i32>: "ARRAY[1, 2, 3]::int[]" == vec![1, 2, 3i32], "ARRAY[3, 292, 15, 2, 3]::int[]" == vec![3, 292, 15, 2, 3], "ARRAY[7, 6, 5, 4, 3, 2, 1]::int[]" == vec![7, 6, 5, 4, 3, 2, 1], "ARRAY[]::int[]" == vec![] as Vec<i32>);
+test!(postgres_string_vec: Vec<String>: "ARRAY['Hello', 'world', 'friend']::text[]" == vec!["Hello", "world", "friend"]);
+test!(postgres_bool_vec: Vec<bool>: "ARRAY[true, true, false, true]::bool[]" == vec![true, true, false, true]);
+test!(postgres_real_vec: Vec<f32>: "ARRAY[0.0, 1.0, 3.14, 1.234, -0.002, 100000.0]::real[]" == vec![0.0, 1.0, 3.14, 1.234, -0.002, 100000.0_f32]);
+test!(postgres_double_vec: Vec<f64>: "ARRAY[0.0, 1.0, 3.14, 1.234, -0.002, 100000.0]::double precision[]" == vec![0.0, 1.0, 3.14, 1.234, -0.002, 100000.0_f64]);
+
 #[cfg_attr(feature = "runtime-async-std", async_std::test)]
 #[cfg_attr(feature = "runtime-tokio", tokio::test)]
 async fn postgres_bytes() -> anyhow::Result<()> {