Skip to content

Commit 4683cc3

Browse files
authored
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
1 parent 08e45f4 commit 4683cc3

File tree

2 files changed

+297
-0
lines changed

2 files changed

+297
-0
lines changed

sqlx-postgres/src/types/hstore.rs

+294
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,294 @@
1+
use std::{
2+
collections::{btree_map, BTreeMap},
3+
mem::size_of,
4+
ops::{Deref, DerefMut},
5+
str::from_utf8,
6+
};
7+
8+
use serde::{Deserialize, Serialize};
9+
10+
use crate::{
11+
decode::Decode,
12+
encode::{Encode, IsNull},
13+
error::BoxDynError,
14+
types::Type,
15+
PgArgumentBuffer, PgTypeInfo, PgValueRef, Postgres,
16+
};
17+
18+
/// Key-value support (`hstore`) for Postgres.
19+
///
20+
/// SQLx currently maps `hstore` to a `BTreeMap<String, Option<String>>` but this may be expanded in
21+
/// future to allow for user defined types.
22+
///
23+
/// See [the Postgres manual, Appendix F, Section 18][PG.F.18]
24+
///
25+
/// [PG.F.18]: https://www.postgresql.org/docs/current/hstore.html
26+
///
27+
/// ### Note: Requires Postgres 8.3+
28+
/// Introduced as a method for storing unstructured data, the `hstore` extension was first added in
29+
/// Postgres 8.3.
30+
///
31+
///
32+
/// ### Note: Extension Required
33+
/// The `hstore` extension is not enabled by default in Postgres. You will need to do so explicitly:
34+
///
35+
/// ```ignore
36+
/// CREATE EXTENSION IF NOT EXISTS hstore;
37+
/// ```
38+
///
39+
/// # Examples
40+
///
41+
/// ```
42+
/// # use sqlx_postgres::types::PgHstore;
43+
/// // Shows basic usage of the PgHstore type.
44+
/// //
45+
/// #[derive(Clone, Debug, Default, Eq, PartialEq)]
46+
/// struct UserCreate<'a> {
47+
/// username: &'a str,
48+
/// password: &'a str,
49+
/// additional_data: PgHstore
50+
/// }
51+
///
52+
/// let mut new_user = UserCreate {
53+
/// username: "[email protected]",
54+
/// password: "@super_secret_1",
55+
/// ..Default::default()
56+
/// };
57+
///
58+
/// new_user.additional_data.insert("department".to_string(), Some("IT".to_string()));
59+
/// new_user.additional_data.insert("equipment_issued".to_string(), None);
60+
/// ```
61+
/// ```ignore
62+
/// query_scalar::<_, i64>(
63+
/// "insert into user(username, password, additional_data) values($1, $2, $3) returning id"
64+
/// )
65+
/// .bind(new_user.username)
66+
/// .bind(new_user.password)
67+
/// .bind(new_user.additional_data)
68+
/// .fetch_one(pg_conn)
69+
/// .await?;
70+
/// ```
71+
///
72+
/// ```
73+
/// # use sqlx_postgres::types::PgHstore;
74+
/// // PgHstore implements FromIterator to simplify construction.
75+
/// //
76+
/// let additional_data = PgHstore::from_iter([
77+
/// ("department".to_string(), Some("IT".to_string())),
78+
/// ("equipment_issued".to_string(), None),
79+
/// ]);
80+
///
81+
/// assert_eq!(additional_data["department"], Some("IT".to_string()));
82+
/// assert_eq!(additional_data["equipment_issued"], None);
83+
///
84+
/// // Also IntoIterator for ease of iteration.
85+
/// //
86+
/// for (key, value) in additional_data {
87+
/// println!("{key}: {value:?}");
88+
/// }
89+
/// ```
90+
///
91+
#[derive(Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize)]
92+
pub struct PgHstore(pub BTreeMap<String, Option<String>>);
93+
94+
impl Deref for PgHstore {
95+
type Target = BTreeMap<String, Option<String>>;
96+
97+
fn deref(&self) -> &Self::Target {
98+
&self.0
99+
}
100+
}
101+
102+
impl DerefMut for PgHstore {
103+
fn deref_mut(&mut self) -> &mut Self::Target {
104+
&mut self.0
105+
}
106+
}
107+
108+
impl FromIterator<(String, String)> for PgHstore {
109+
fn from_iter<T: IntoIterator<Item = (String, String)>>(iter: T) -> Self {
110+
iter.into_iter().map(|(k, v)| (k, Some(v))).collect()
111+
}
112+
}
113+
114+
impl FromIterator<(String, Option<String>)> for PgHstore {
115+
fn from_iter<T: IntoIterator<Item = (String, Option<String>)>>(iter: T) -> Self {
116+
let mut result = Self::default();
117+
118+
for (key, value) in iter {
119+
result.0.insert(key, value);
120+
}
121+
122+
result
123+
}
124+
}
125+
126+
impl IntoIterator for PgHstore {
127+
type Item = (String, Option<String>);
128+
type IntoIter = btree_map::IntoIter<String, Option<String>>;
129+
130+
fn into_iter(self) -> Self::IntoIter {
131+
self.0.into_iter()
132+
}
133+
}
134+
135+
impl Type<Postgres> for PgHstore {
136+
fn type_info() -> PgTypeInfo {
137+
PgTypeInfo::with_name("hstore")
138+
}
139+
}
140+
141+
impl<'r> Decode<'r, Postgres> for PgHstore {
142+
fn decode(value: PgValueRef<'r>) -> Result<Self, BoxDynError> {
143+
let mut buf = <&[u8] as Decode<Postgres>>::decode(value)?;
144+
let len = read_length(&mut buf)?;
145+
146+
if len < 0 {
147+
Err(format!("hstore, invalid entry count: {len}"))?;
148+
}
149+
150+
let mut result = Self::default();
151+
152+
while !buf.is_empty() {
153+
let key_len = read_length(&mut buf)?;
154+
let key = read_value(&mut buf, key_len)?.ok_or("hstore, key not found")?;
155+
156+
let value_len = read_length(&mut buf)?;
157+
let value = read_value(&mut buf, value_len)?;
158+
159+
result.insert(key, value);
160+
}
161+
162+
Ok(result)
163+
}
164+
}
165+
166+
impl Encode<'_, Postgres> for PgHstore {
167+
fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
168+
buf.extend_from_slice(&i32::to_be_bytes(self.0.len() as i32));
169+
170+
for (key, val) in &self.0 {
171+
let key_bytes = key.as_bytes();
172+
173+
buf.extend_from_slice(&i32::to_be_bytes(key_bytes.len() as i32));
174+
buf.extend_from_slice(key_bytes);
175+
176+
match val {
177+
Some(val) => {
178+
let val_bytes = val.as_bytes();
179+
180+
buf.extend_from_slice(&i32::to_be_bytes(val_bytes.len() as i32));
181+
buf.extend_from_slice(val_bytes);
182+
}
183+
None => {
184+
buf.extend_from_slice(&i32::to_be_bytes(-1));
185+
}
186+
}
187+
}
188+
189+
Ok(IsNull::No)
190+
}
191+
}
192+
193+
fn read_length(buf: &mut &[u8]) -> Result<i32, BoxDynError> {
194+
let (bytes, rest) = buf.split_at(size_of::<i32>());
195+
196+
*buf = rest;
197+
198+
Ok(i32::from_be_bytes(
199+
bytes
200+
.try_into()
201+
.map_err(|err| format!("hstore, reading length: {err}"))?,
202+
))
203+
}
204+
205+
fn read_value(buf: &mut &[u8], len: i32) -> Result<Option<String>, BoxDynError> {
206+
match len {
207+
len if len <= 0 => Ok(None),
208+
len => {
209+
let (val, rest) = buf.split_at(len as usize);
210+
211+
*buf = rest;
212+
213+
Ok(Some(
214+
from_utf8(val)
215+
.map_err(|err| format!("hstore, reading value: {err}"))?
216+
.to_string(),
217+
))
218+
}
219+
}
220+
}
221+
222+
#[cfg(test)]
223+
mod test {
224+
use super::*;
225+
use crate::PgValueFormat;
226+
227+
const EMPTY: &str = "00000000";
228+
229+
const NAME_SURNAME_AGE: &str =
230+
"0000000300000003616765ffffffff000000046e616d65000000044a6f686e000000077375726e616d6500000003446f65";
231+
232+
#[test]
233+
fn hstore_deserialize_ok() {
234+
let empty = hex::decode(EMPTY).unwrap();
235+
let name_surname_age = hex::decode(NAME_SURNAME_AGE).unwrap();
236+
237+
let empty = PgValueRef {
238+
value: Some(empty.as_slice()),
239+
row: None,
240+
type_info: PgTypeInfo::with_name("hstore"),
241+
format: PgValueFormat::Binary,
242+
};
243+
244+
let name_surname = PgValueRef {
245+
value: Some(name_surname_age.as_slice()),
246+
row: None,
247+
type_info: PgTypeInfo::with_name("hstore"),
248+
format: PgValueFormat::Binary,
249+
};
250+
251+
let res_empty = PgHstore::decode(empty).unwrap();
252+
let res_name_surname = PgHstore::decode(name_surname).unwrap();
253+
254+
assert!(res_empty.is_empty());
255+
assert_eq!(res_name_surname["name"], Some("John".to_string()));
256+
assert_eq!(res_name_surname["surname"], Some("Doe".to_string()));
257+
assert_eq!(res_name_surname["age"], None);
258+
}
259+
260+
#[test]
261+
#[should_panic(expected = "hstore, invalid entry count: -5")]
262+
fn hstore_deserialize_buffer_length_error() {
263+
let buf = PgValueRef {
264+
value: Some(&[255, 255, 255, 251]),
265+
row: None,
266+
type_info: PgTypeInfo::with_name("hstore"),
267+
format: PgValueFormat::Binary,
268+
};
269+
270+
PgHstore::decode(buf).unwrap();
271+
}
272+
273+
#[test]
274+
fn hstore_serialize_ok() {
275+
let mut buff = PgArgumentBuffer::default();
276+
let _ = PgHstore::from_iter::<[(String, String); 0]>([])
277+
.encode_by_ref(&mut buff)
278+
.unwrap();
279+
280+
assert_eq!(hex::encode(buff.as_slice()), EMPTY);
281+
282+
buff.clear();
283+
284+
let _ = PgHstore::from_iter([
285+
("name".to_string(), Some("John".to_string())),
286+
("surname".to_string(), Some("Doe".to_string())),
287+
("age".to_string(), None),
288+
])
289+
.encode_by_ref(&mut buff)
290+
.unwrap();
291+
292+
assert_eq!(hex::encode(buff.as_slice()), NAME_SURNAME_AGE);
293+
}
294+
}

sqlx-postgres/src/types/mod.rs

+3
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
//! | [`PgLQuery`] | LQUERY |
2222
//! | [`PgCiText`] | CITEXT<sup>1</sup> |
2323
//! | [`PgCube`] | CUBE |
24+
//! | [`PgHstore`] | HSTORE |
2425
//!
2526
//! <sup>1</sup> SQLx generally considers `CITEXT` to be compatible with `String`, `&str`, etc.,
2627
//! but this wrapper type is available for edge cases, such as `CITEXT[]` which Postgres
@@ -187,6 +188,7 @@ mod bool;
187188
mod bytes;
188189
mod citext;
189190
mod float;
191+
mod hstore;
190192
mod int;
191193
mod interval;
192194
mod lquery;
@@ -240,6 +242,7 @@ mod bit_vec;
240242
pub use array::PgHasArrayType;
241243
pub use citext::PgCiText;
242244
pub use cube::PgCube;
245+
pub use hstore::PgHstore;
243246
pub use interval::PgInterval;
244247
pub use lquery::PgLQuery;
245248
pub use lquery::PgLQueryLevel;

0 commit comments

Comments
 (0)