Skip to content

Commit

Permalink
Split out ltree,query,txtquery protocol parsers, add tests, rust fmt
Browse files Browse the repository at this point in the history
  • Loading branch information
halfmatthalfcat committed Mar 17, 2022
1 parent 6ae60d6 commit d9d283e
Show file tree
Hide file tree
Showing 4 changed files with 274 additions and 71 deletions.
51 changes: 46 additions & 5 deletions postgres-protocol/src/types/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1060,18 +1060,59 @@ impl Inet {
}
}

/// Serializes a Postgres l{tree,query,txtquery} string
/// 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
// 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 l{tree,query,txtquery} string
/// Deserialize a Postgres ltree string
#[inline]
pub fn ltree_from_sql(buf: &[u8]) -> Result<&str, StdBox<dyn Error + Sync + Send>> {
// Remove the version number from the front of the string per spec
Ok(str::from_utf8(&buf[1..])?)
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()),
}
}
116 changes: 115 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,117 @@ 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());

let success = match ltree_from_sql(query.as_slice()) {
Ok(_) => true,
_ => false,
};

assert!(success)
}

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

let success = match ltree_from_sql(query.as_slice()) {
Err(_) => true,
_ => false,
};

assert!(success)
}

#[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());

let success = match lquery_from_sql(query.as_slice()) {
Ok(_) => true,
_ => false,
};

assert!(success)
}

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

let success = match lquery_from_sql(query.as_slice()) {
Err(_) => true,
_ => false,
};

assert!(success)
}

#[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());

let success = match ltree_from_sql(query.as_slice()) {
Ok(_) => true,
_ => false,
};

assert!(success)
}

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

let success = match ltree_from_sql(query.as_slice()) {
Err(_) => true,
_ => false,
};

assert!(success)
}
50 changes: 25 additions & 25 deletions postgres-types/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -619,24 +619,24 @@ impl<'a> FromSql<'a> for Box<str> {
impl<'a> FromSql<'a> for &'a str {
fn from_sql(ty: &Type, raw: &'a [u8]) -> Result<&'a str, Box<dyn Error + Sync + Send>> {
match *ty {
ref ty if (
ty.name() == "ltree" ||
ty.name() == "lquery" ||
ty.name() == "ltxtquery"
) => types::ltree_from_sql(raw),
_ => types::text_from_sql(raw)
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" ||
ty.name() == "ltree" ||
ty.name() == "lquery" ||
ty.name() == "ltxtquery"
) => true,
ref ty
if (ty.name() == "citext"
|| ty.name() == "ltree"
|| ty.name() == "lquery"
|| ty.name() == "ltxtquery") =>
{
true
}
_ => false,
}
}
Expand Down Expand Up @@ -939,26 +939,26 @@ impl ToSql for Vec<u8> {

impl<'a> ToSql for &'a str {
fn to_sql(&self, ty: &Type, w: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>> {
match ty {
ref ty if (
ty.name() == "ltree" ||
ty.name() == "lquery" ||
ty.name() == "ltxtquery"
) => types::ltree_to_sql(*self, w),
_ => types::text_to_sql(*self, w)
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" ||
ty.name() == "ltree" ||
ty.name() == "lquery" ||
ty.name() == "ltxtquery"
) => true,
ref ty
if (ty.name() == "citext"
|| ty.name() == "ltree"
|| ty.name() == "lquery"
|| ty.name() == "ltxtquery") =>
{
true
}
_ => false,
}
}
Expand Down
Loading

0 comments on commit d9d283e

Please sign in to comment.