forked from launchbadge/sqlx
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add support for PostgreSQL HSTORE data type (launchbadge#3343)
* 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
1 parent
ea30592
commit 749ae5c
Showing
2 changed files
with
297 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters