Skip to content

Commit

Permalink
Add support for PostgreSQL HSTORE data type (#3343)
Browse files Browse the repository at this point in the history
* Add support for PostgreSQL HSTORE data type

* Changes to make the future evolution of the API easier

* Fix clippy lints

* Add basic documentation
  • Loading branch information
KobusEllis authored Jul 19, 2024
1 parent 08e45f4 commit 4683cc3
Show file tree
Hide file tree
Showing 2 changed files with 297 additions and 0 deletions.
294 changes: 294 additions & 0 deletions sqlx-postgres/src/types/hstore.rs
Original file line number Diff line number Diff line change
@@ -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);
}
}
3 changes: 3 additions & 0 deletions sqlx-postgres/src/types/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -187,6 +188,7 @@ mod bool;
mod bytes;
mod citext;
mod float;
mod hstore;
mod int;
mod interval;
mod lquery;
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit 4683cc3

Please sign in to comment.