Add support for PostgreSQL HSTORE data type

sqlx-postgres/src/types/hstore.rs

@@ -0,0 +1,294 @@
+use std::{
+    collections::{btree_map, BTreeMap},
+    mem::size_of,
+    ops::{Deref, DerefMut},
+    str::from_utf8,
+};
+
+use serde::{Deserialize, Serialize};
+
+use crate::{
+    decode::Decode,
+    encode::{Encode, IsNull},
+    error::BoxDynError,
+    types::Type,
+    PgArgumentBuffer, PgTypeInfo, PgValueRef, Postgres,
+};
+
+/// Key-value support (`hstore`) for Postgres.
+///
+/// SQLx currently maps `hstore` to a `BTreeMap<String, Option<String>>` but this may be expanded in
+/// future to allow for user defined types.
+///
+/// See [the Postgres manual, Appendix F, Section 18][PG.F.18]
+///
+/// [PG.F.18]: https://www.postgresql.org/docs/current/hstore.html
+///
+/// ### Note: Requires Postgres 8.3+
+/// Introduced as a method for storing unstructured data, the `hstore` extension was first added in
+/// Postgres 8.3.
+///
+///
+/// ### Note: Extension Required
+/// The `hstore` extension is not enabled by default in Postgres. You will need to do so explicitly:
+///
+/// ```ignore
+/// CREATE EXTENSION IF NOT EXISTS hstore;
+/// ```
+///
+/// # Examples
+///
+/// ```
+/// # use sqlx_postgres::types::PgHstore;
+/// // Shows basic usage of the PgHstore type.
+/// //
+/// #[derive(Clone, Debug, Default, Eq, PartialEq)]
+/// struct UserCreate<'a> {
+///     username: &'a str,
+///     password: &'a str,
+///     additional_data: PgHstore
+/// }
+///
+/// let mut new_user = UserCreate {
+///     username: "name.surname@email.com",
+///     password: "@super_secret_1",
+///     ..Default::default()
+/// };
+///
+/// new_user.additional_data.insert("department".to_string(), Some("IT".to_string()));
+/// new_user.additional_data.insert("equipment_issued".to_string(), None);
+/// ```
+/// ```ignore
+/// query_scalar::<_, i64>(
+///     "insert into user(username, password, additional_data) values($1, $2, $3) returning id"
+/// )
+/// .bind(new_user.username)
+/// .bind(new_user.password)
+/// .bind(new_user.additional_data)
+/// .fetch_one(pg_conn)
+/// .await?;
+/// ```
+///
+/// ```
+/// # use sqlx_postgres::types::PgHstore;
+/// // PgHstore implements FromIterator to simplify construction.
+/// //
+/// let additional_data = PgHstore::from_iter([
+///     ("department".to_string(), Some("IT".to_string())),
+///     ("equipment_issued".to_string(), None),
+/// ]);
+///
+/// assert_eq!(additional_data["department"], Some("IT".to_string()));
+/// assert_eq!(additional_data["equipment_issued"], None);
+///
+/// // Also IntoIterator for ease of iteration.
+/// //
+/// for (key, value) in additional_data {
+///     println!("{key}: {value:?}");
+/// }
+/// ```
+///
+#[derive(Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize)]
+pub struct PgHstore(pub BTreeMap<String, Option<String>>);
+
+impl Deref for PgHstore {
+    type Target = BTreeMap<String, Option<String>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for PgHstore {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl FromIterator<(String, String)> for PgHstore {
+    fn from_iter<T: IntoIterator<Item = (String, String)>>(iter: T) -> Self {
+        iter.into_iter().map(|(k, v)| (k, Some(v))).collect()
+    }
+}
+
+impl FromIterator<(String, Option<String>)> for PgHstore {
+    fn from_iter<T: IntoIterator<Item = (String, Option<String>)>>(iter: T) -> Self {
+        let mut result = Self::default();
+
+        for (key, value) in iter {
+            result.0.insert(key, value);
+        }
+
+        result
+    }
+}
+
+impl IntoIterator for PgHstore {
+    type Item = (String, Option<String>);
+    type IntoIter = btree_map::IntoIter<String, Option<String>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
+
+impl Type<Postgres> for PgHstore {
+    fn type_info() -> PgTypeInfo {
+        PgTypeInfo::with_name("hstore")
+    }
+}
+
+impl<'r> Decode<'r, Postgres> for PgHstore {
+    fn decode(value: PgValueRef<'r>) -> Result<Self, BoxDynError> {
+        let mut buf = <&[u8] as Decode<Postgres>>::decode(value)?;
+        let len = read_length(&mut buf)?;
+
+        if len < 0 {
+            Err(format!("hstore, invalid entry count: {len}"))?;
+        }
+
+        let mut result = Self::default();
+
+        while !buf.is_empty() {
+            let key_len = read_length(&mut buf)?;
+            let key = read_value(&mut buf, key_len)?.ok_or("hstore, key not found")?;
+
+            let value_len = read_length(&mut buf)?;
+            let value = read_value(&mut buf, value_len)?;
+
+            result.insert(key, value);
+        }
+
+        Ok(result)
+    }
+}
+
+impl Encode<'_, Postgres> for PgHstore {
+    fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
+        buf.extend_from_slice(&i32::to_be_bytes(self.0.len() as i32));
+
+        for (key, val) in &self.0 {
+            let key_bytes = key.as_bytes();
+
+            buf.extend_from_slice(&i32::to_be_bytes(key_bytes.len() as i32));
+            buf.extend_from_slice(key_bytes);
+
+            match val {
+                Some(val) => {
+                    let val_bytes = val.as_bytes();
+
+                    buf.extend_from_slice(&i32::to_be_bytes(val_bytes.len() as i32));
+                    buf.extend_from_slice(val_bytes);
+                }
+                None => {
+                    buf.extend_from_slice(&i32::to_be_bytes(-1));
+                }
+            }
+        }
+
+        Ok(IsNull::No)
+    }
+}
+
+fn read_length(buf: &mut &[u8]) -> Result<i32, BoxDynError> {
+    let (bytes, rest) = buf.split_at(size_of::<i32>());
+
+    *buf = rest;
+
+    Ok(i32::from_be_bytes(
+        bytes
+            .try_into()
+            .map_err(|err| format!("hstore, reading length: {err}"))?,
+    ))
+}
+
+fn read_value(buf: &mut &[u8], len: i32) -> Result<Option<String>, BoxDynError> {
+    match len {
+        len if len <= 0 => Ok(None),
+        len => {
+            let (val, rest) = buf.split_at(len as usize);
+
+            *buf = rest;
+
+            Ok(Some(
+                from_utf8(val)
+                    .map_err(|err| format!("hstore, reading value: {err}"))?
+                    .to_string(),
+            ))
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::PgValueFormat;
+
+    const EMPTY: &str = "00000000";
+
+    const NAME_SURNAME_AGE: &str =
+        "0000000300000003616765ffffffff000000046e616d65000000044a6f686e000000077375726e616d6500000003446f65";
+
+    #[test]
+    fn hstore_deserialize_ok() {
+        let empty = hex::decode(EMPTY).unwrap();
+        let name_surname_age = hex::decode(NAME_SURNAME_AGE).unwrap();
+
+        let empty = PgValueRef {
+            value: Some(empty.as_slice()),
+            row: None,
+            type_info: PgTypeInfo::with_name("hstore"),
+            format: PgValueFormat::Binary,
+        };
+
+        let name_surname = PgValueRef {
+            value: Some(name_surname_age.as_slice()),
+            row: None,
+            type_info: PgTypeInfo::with_name("hstore"),
+            format: PgValueFormat::Binary,
+        };
+
+        let res_empty = PgHstore::decode(empty).unwrap();
+        let res_name_surname = PgHstore::decode(name_surname).unwrap();
+
+        assert!(res_empty.is_empty());
+        assert_eq!(res_name_surname["name"], Some("John".to_string()));
+        assert_eq!(res_name_surname["surname"], Some("Doe".to_string()));
+        assert_eq!(res_name_surname["age"], None);
+    }
+
+    #[test]
+    #[should_panic(expected = "hstore, invalid entry count: -5")]
+    fn hstore_deserialize_buffer_length_error() {
+        let buf = PgValueRef {
+            value: Some(&[255, 255, 255, 251]),
+            row: None,
+            type_info: PgTypeInfo::with_name("hstore"),
+            format: PgValueFormat::Binary,
+        };
+
+        PgHstore::decode(buf).unwrap();
+    }
+
+    #[test]
+    fn hstore_serialize_ok() {
+        let mut buff = PgArgumentBuffer::default();
+        let _ = PgHstore::from_iter::<[(String, String); 0]>([])
+            .encode_by_ref(&mut buff)
+            .unwrap();
+
+        assert_eq!(hex::encode(buff.as_slice()), EMPTY);
+
+        buff.clear();
+
+        let _ = PgHstore::from_iter([
+            ("name".to_string(), Some("John".to_string())),
+            ("surname".to_string(), Some("Doe".to_string())),
+            ("age".to_string(), None),
+        ])
+        .encode_by_ref(&mut buff)
+        .unwrap();
+
+        assert_eq!(hex::encode(buff.as_slice()), NAME_SURNAME_AGE);
+    }
+}

sqlx-postgres/src/types/mod.rs

@@ -21,6 +21,7 @@
 //! | [`PgLQuery`]                          | LQUERY                                               |
 //! | [`PgCiText`]                          | CITEXT<sup>1</sup>                                   |
 //! | [`PgCube`]                            | CUBE                                                 |
+//! | [`PgHstore`]                          | HSTORE                                               |
 //!
 //! <sup>1</sup> SQLx generally considers `CITEXT` to be compatible with `String`, `&str`, etc.,
 //! but this wrapper type is available for edge cases, such as `CITEXT[]` which Postgres
@@ -187,6 +188,7 @@ mod bool;
 mod bytes;
 mod citext;
 mod float;
+mod hstore;
 mod int;
 mod interval;
 mod lquery;
@@ -240,6 +242,7 @@ mod bit_vec;
 pub use array::PgHasArrayType;
 pub use citext::PgCiText;
 pub use cube::PgCube;
+pub use hstore::PgHstore;
 pub use interval::PgInterval;
 pub use lquery::PgLQuery;
 pub use lquery::PgLQueryLevel;