From 4683cc34e54188cee98de099816b081e2e64aabf Mon Sep 17 00:00:00 2001 From: Kobus Ellis Date: Fri, 19 Jul 2024 23:33:16 +0200 Subject: [PATCH] Add support for PostgreSQL HSTORE data type (#3343) * Add support for PostgreSQL HSTORE data type * Changes to make the future evolution of the API easier * Fix clippy lints * Add basic documentation --- sqlx-postgres/src/types/hstore.rs | 294 ++++++++++++++++++++++++++++++ sqlx-postgres/src/types/mod.rs | 3 + 2 files changed, 297 insertions(+) create mode 100644 sqlx-postgres/src/types/hstore.rs diff --git a/sqlx-postgres/src/types/hstore.rs b/sqlx-postgres/src/types/hstore.rs new file mode 100644 index 0000000000..0dcc71943c --- /dev/null +++ b/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>` 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>); + +impl Deref for PgHstore { + type Target = BTreeMap>; + + 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>(iter: T) -> Self { + iter.into_iter().map(|(k, v)| (k, Some(v))).collect() + } +} + +impl FromIterator<(String, Option)> for PgHstore { + fn from_iter)>>(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); + type IntoIter = btree_map::IntoIter>; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl Type for PgHstore { + fn type_info() -> PgTypeInfo { + PgTypeInfo::with_name("hstore") + } +} + +impl<'r> Decode<'r, Postgres> for PgHstore { + fn decode(value: PgValueRef<'r>) -> Result { + let mut buf = <&[u8] as Decode>::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 { + 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 { + let (bytes, rest) = buf.split_at(size_of::()); + + *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, 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); + } +} diff --git a/sqlx-postgres/src/types/mod.rs b/sqlx-postgres/src/types/mod.rs index 3219480cb7..846f1b731d 100644 --- a/sqlx-postgres/src/types/mod.rs +++ b/sqlx-postgres/src/types/mod.rs @@ -21,6 +21,7 @@ //! | [`PgLQuery`] | LQUERY | //! | [`PgCiText`] | CITEXT1 | //! | [`PgCube`] | CUBE | +//! | [`PgHstore`] | HSTORE | //! //! 1 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;