Skip to content

Commit

Permalink
Merge pull request #875 from halfmatthalfcat/ltree-support
Browse files Browse the repository at this point in the history
Add ltree, lquery and ltxtquery support
  • Loading branch information
sfackler authored Mar 23, 2022
2 parents a638ada + 6fae655 commit 38da7fa
Show file tree
Hide file tree
Showing 8 changed files with 272 additions and 17 deletions.
1 change: 1 addition & 0 deletions docker/sql_setup.sh
Original file line number Diff line number Diff line change
Expand Up @@ -96,4 +96,5 @@ psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" <<-EOSQL
CREATE ROLE ssl_user LOGIN;
CREATE EXTENSION hstore;
CREATE EXTENSION citext;
CREATE EXTENSION ltree;
EOSQL
2 changes: 1 addition & 1 deletion postgres-protocol/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "postgres-protocol"
version = "0.6.3"
version = "0.6.4"
authors = ["Steven Fackler <sfackler@gmail.com>"]
edition = "2018"
description = "Low level Postgres protocol APIs"
Expand Down
57 changes: 57 additions & 0 deletions postgres-protocol/src/types/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1059,3 +1059,60 @@ impl Inet {
self.netmask
}
}

/// Serializes a Postgres ltree string
#[inline]
pub fn ltree_to_sql(v: &str, buf: &mut BytesMut) {
// A version number is prepended to an ltree string per spec
buf.put_u8(1);
// Append the rest of the query
buf.put_slice(v.as_bytes());
}

/// Deserialize a Postgres ltree string
#[inline]
pub fn ltree_from_sql(buf: &[u8]) -> Result<&str, StdBox<dyn Error + Sync + Send>> {
match buf {
// Remove the version number from the front of the ltree per spec
[1u8, rest @ ..] => Ok(str::from_utf8(rest)?),
_ => Err("ltree version 1 only supported".into()),
}
}

/// Serializes a Postgres lquery string
#[inline]
pub fn lquery_to_sql(v: &str, buf: &mut BytesMut) {
// A version number is prepended to an lquery string per spec
buf.put_u8(1);
// Append the rest of the query
buf.put_slice(v.as_bytes());
}

/// Deserialize a Postgres lquery string
#[inline]
pub fn lquery_from_sql(buf: &[u8]) -> Result<&str, StdBox<dyn Error + Sync + Send>> {
match buf {
// Remove the version number from the front of the lquery per spec
[1u8, rest @ ..] => Ok(str::from_utf8(rest)?),
_ => Err("lquery version 1 only supported".into()),
}
}

/// Serializes a Postgres ltxtquery string
#[inline]
pub fn ltxtquery_to_sql(v: &str, buf: &mut BytesMut) {
// A version number is prepended to an ltxtquery string per spec
buf.put_u8(1);
// Append the rest of the query
buf.put_slice(v.as_bytes());
}

/// Deserialize a Postgres ltxtquery string
#[inline]
pub fn ltxtquery_from_sql(buf: &[u8]) -> Result<&str, StdBox<dyn Error + Sync + Send>> {
match buf {
// Remove the version number from the front of the ltxtquery per spec
[1u8, rest @ ..] => Ok(str::from_utf8(rest)?),
_ => Err("ltxtquery version 1 only supported".into()),
}
}
86 changes: 85 additions & 1 deletion postgres-protocol/src/types/test.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use bytes::BytesMut;
use bytes::{Buf, BytesMut};
use fallible_iterator::FallibleIterator;
use std::collections::HashMap;

Expand Down Expand Up @@ -156,3 +156,87 @@ fn non_null_array() {
assert_eq!(array.dimensions().collect::<Vec<_>>().unwrap(), dimensions);
assert_eq!(array.values().collect::<Vec<_>>().unwrap(), values);
}

#[test]
fn ltree_sql() {
let mut query = vec![1u8];
query.extend_from_slice("A.B.C".as_bytes());

let mut buf = BytesMut::new();

ltree_to_sql("A.B.C", &mut buf);

assert_eq!(query.as_slice(), buf.chunk());
}

#[test]
fn ltree_str() {
let mut query = vec![1u8];
query.extend_from_slice("A.B.C".as_bytes());

assert!(matches!(ltree_from_sql(query.as_slice()), Ok(_)))
}

#[test]
fn ltree_wrong_version() {
let mut query = vec![2u8];
query.extend_from_slice("A.B.C".as_bytes());

assert!(matches!(ltree_from_sql(query.as_slice()), Err(_)))
}

#[test]
fn lquery_sql() {
let mut query = vec![1u8];
query.extend_from_slice("A.B.C".as_bytes());

let mut buf = BytesMut::new();

lquery_to_sql("A.B.C", &mut buf);

assert_eq!(query.as_slice(), buf.chunk());
}

#[test]
fn lquery_str() {
let mut query = vec![1u8];
query.extend_from_slice("A.B.C".as_bytes());

assert!(matches!(lquery_from_sql(query.as_slice()), Ok(_)))
}

#[test]
fn lquery_wrong_version() {
let mut query = vec![2u8];
query.extend_from_slice("A.B.C".as_bytes());

assert!(matches!(lquery_from_sql(query.as_slice()), Err(_)))
}

#[test]
fn ltxtquery_sql() {
let mut query = vec![1u8];
query.extend_from_slice("a & b*".as_bytes());

let mut buf = BytesMut::new();

ltree_to_sql("a & b*", &mut buf);

assert_eq!(query.as_slice(), buf.chunk());
}

#[test]
fn ltxtquery_str() {
let mut query = vec![1u8];
query.extend_from_slice("a & b*".as_bytes());

assert!(matches!(ltree_from_sql(query.as_slice()), Ok(_)))
}

#[test]
fn ltxtquery_wrong_version() {
let mut query = vec![2u8];
query.extend_from_slice("a & b*".as_bytes());

assert!(matches!(ltree_from_sql(query.as_slice()), Err(_)))
}
4 changes: 2 additions & 2 deletions postgres-types/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "postgres-types"
version = "0.2.2"
version = "0.2.3"
authors = ["Steven Fackler <sfackler@gmail.com>"]
edition = "2018"
license = "MIT/Apache-2.0"
Expand Down Expand Up @@ -28,7 +28,7 @@ with-time-0_3 = ["time-03"]
[dependencies]
bytes = "1.0"
fallible-iterator = "0.2"
postgres-protocol = { version = "0.6.1", path = "../postgres-protocol" }
postgres-protocol = { version = "0.6.4", path = "../postgres-protocol" }
postgres-derive = { version = "0.4.0", optional = true, path = "../postgres-derive" }

array-init = { version = "2", optional = true }
Expand Down
46 changes: 36 additions & 10 deletions postgres-types/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -407,6 +407,7 @@ impl WrongType {
/// | `f32` | REAL |
/// | `f64` | DOUBLE PRECISION |
/// | `&str`/`String` | VARCHAR, CHAR(n), TEXT, CITEXT, NAME, UNKNOWN |
/// | | LTREE, LQUERY, LTXTQUERY |
/// | `&[u8]`/`Vec<u8>` | BYTEA |
/// | `HashMap<String, Option<String>>` | HSTORE |
/// | `SystemTime` | TIMESTAMP, TIMESTAMP WITH TIME ZONE |
Expand Down Expand Up @@ -594,8 +595,8 @@ impl<'a> FromSql<'a> for &'a [u8] {
}

impl<'a> FromSql<'a> for String {
fn from_sql(_: &Type, raw: &'a [u8]) -> Result<String, Box<dyn Error + Sync + Send>> {
types::text_from_sql(raw).map(ToString::to_string)
fn from_sql(ty: &Type, raw: &'a [u8]) -> Result<String, Box<dyn Error + Sync + Send>> {
<&str as FromSql>::from_sql(ty, raw).map(ToString::to_string)
}

fn accepts(ty: &Type) -> bool {
Expand All @@ -604,8 +605,8 @@ impl<'a> FromSql<'a> for String {
}

impl<'a> FromSql<'a> for Box<str> {
fn from_sql(_: &Type, raw: &'a [u8]) -> Result<Box<str>, Box<dyn Error + Sync + Send>> {
types::text_from_sql(raw)
fn from_sql(ty: &Type, raw: &'a [u8]) -> Result<Box<str>, Box<dyn Error + Sync + Send>> {
<&str as FromSql>::from_sql(ty, raw)
.map(ToString::to_string)
.map(String::into_boxed_str)
}
Expand All @@ -616,14 +617,26 @@ impl<'a> FromSql<'a> for Box<str> {
}

impl<'a> FromSql<'a> for &'a str {
fn from_sql(_: &Type, raw: &'a [u8]) -> Result<&'a str, Box<dyn Error + Sync + Send>> {
types::text_from_sql(raw)
fn from_sql(ty: &Type, raw: &'a [u8]) -> Result<&'a str, Box<dyn Error + Sync + Send>> {
match *ty {
ref ty if ty.name() == "ltree" => types::ltree_from_sql(raw),
ref ty if ty.name() == "lquery" => types::lquery_from_sql(raw),
ref ty if ty.name() == "ltxtquery" => types::ltxtquery_from_sql(raw),
_ => types::text_from_sql(raw),
}
}

fn accepts(ty: &Type) -> bool {
match *ty {
Type::VARCHAR | Type::TEXT | Type::BPCHAR | Type::NAME | Type::UNKNOWN => true,
ref ty if ty.name() == "citext" => true,
ref ty
if (ty.name() == "citext"
|| ty.name() == "ltree"
|| ty.name() == "lquery"
|| ty.name() == "ltxtquery") =>
{
true
}
_ => false,
}
}
Expand Down Expand Up @@ -727,6 +740,7 @@ pub enum IsNull {
/// | `f32` | REAL |
/// | `f64` | DOUBLE PRECISION |
/// | `&str`/`String` | VARCHAR, CHAR(n), TEXT, CITEXT, NAME |
/// | | LTREE, LQUERY, LTXTQUERY |
/// | `&[u8]`/`Vec<u8>` | BYTEA |
/// | `HashMap<String, Option<String>>` | HSTORE |
/// | `SystemTime` | TIMESTAMP, TIMESTAMP WITH TIME ZONE |
Expand Down Expand Up @@ -924,15 +938,27 @@ impl ToSql for Vec<u8> {
}

impl<'a> ToSql for &'a str {
fn to_sql(&self, _: &Type, w: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>> {
types::text_to_sql(*self, w);
fn to_sql(&self, ty: &Type, w: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>> {
match *ty {
ref ty if ty.name() == "ltree" => types::ltree_to_sql(*self, w),
ref ty if ty.name() == "lquery" => types::lquery_to_sql(*self, w),
ref ty if ty.name() == "ltxtquery" => types::ltxtquery_to_sql(*self, w),
_ => types::text_to_sql(*self, w),
}
Ok(IsNull::No)
}

fn accepts(ty: &Type) -> bool {
match *ty {
Type::VARCHAR | Type::TEXT | Type::BPCHAR | Type::NAME | Type::UNKNOWN => true,
ref ty if ty.name() == "citext" => true,
ref ty
if (ty.name() == "citext"
|| ty.name() == "ltree"
|| ty.name() == "lquery"
|| ty.name() == "ltxtquery") =>
{
true
}
_ => false,
}
}
Expand Down
6 changes: 3 additions & 3 deletions tokio-postgres/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "tokio-postgres"
version = "0.7.5"
version = "0.7.6"
authors = ["Steven Fackler <sfackler@gmail.com>"]
edition = "2018"
license = "MIT/Apache-2.0"
Expand Down Expand Up @@ -50,8 +50,8 @@ parking_lot = "0.12"
percent-encoding = "2.0"
pin-project-lite = "0.2"
phf = "0.10"
postgres-protocol = { version = "0.6.1", path = "../postgres-protocol" }
postgres-types = { version = "0.2.2", path = "../postgres-types" }
postgres-protocol = { version = "0.6.4", path = "../postgres-protocol" }
postgres-types = { version = "0.2.3", path = "../postgres-types" }
socket2 = "0.4"
tokio = { version = "1.0", features = ["io-util"] }
tokio-util = { version = "0.7", features = ["codec"] }
Expand Down
87 changes: 87 additions & 0 deletions tokio-postgres/tests/test/types/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -648,3 +648,90 @@ async fn inet() {
)
.await;
}

#[tokio::test]
async fn ltree() {
test_type(
"ltree",
&[(Some("b.c.d".to_owned()), "'b.c.d'"), (None, "NULL")],
)
.await;
}

#[tokio::test]
async fn ltree_any() {
test_type(
"ltree[]",
&[
(Some(vec![]), "ARRAY[]"),
(Some(vec!["a.b.c".to_string()]), "ARRAY['a.b.c']"),
(
Some(vec!["a.b.c".to_string(), "e.f.g".to_string()]),
"ARRAY['a.b.c','e.f.g']",
),
(None, "NULL"),
],
)
.await;
}

#[tokio::test]
async fn lquery() {
test_type(
"lquery",
&[
(Some("b.c.d".to_owned()), "'b.c.d'"),
(Some("b.c.*".to_owned()), "'b.c.*'"),
(Some("b.*{1,2}.d|e".to_owned()), "'b.*{1,2}.d|e'"),
(None, "NULL"),
],
)
.await;
}

#[tokio::test]
async fn lquery_any() {
test_type(
"lquery[]",
&[
(Some(vec![]), "ARRAY[]"),
(Some(vec!["b.c.*".to_string()]), "ARRAY['b.c.*']"),
(
Some(vec!["b.c.*".to_string(), "b.*{1,2}.d|e".to_string()]),
"ARRAY['b.c.*','b.*{1,2}.d|e']",
),
(None, "NULL"),
],
)
.await;
}

#[tokio::test]
async fn ltxtquery() {
test_type(
"ltxtquery",
&[
(Some("b & c & d".to_owned()), "'b & c & d'"),
(Some("b@* & !c".to_owned()), "'b@* & !c'"),
(None, "NULL"),
],
)
.await;
}

#[tokio::test]
async fn ltxtquery_any() {
test_type(
"ltxtquery[]",
&[
(Some(vec![]), "ARRAY[]"),
(Some(vec!["b & c & d".to_string()]), "ARRAY['b & c & d']"),
(
Some(vec!["b & c & d".to_string(), "b@* & !c".to_string()]),
"ARRAY['b & c & d','b@* & !c']",
),
(None, "NULL"),
],
)
.await;
}

0 comments on commit 38da7fa

Please sign in to comment.