Merge pull request #875 from halfmatthalfcat/ltree-support

Add ltree, lquery and ltxtquery support

Author: sfackler

docker/sql_setup.sh

@@ -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

postgres-protocol/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "postgres-protocol"
-version = "0.6.3"
+version = "0.6.4"
 authors = ["Steven Fackler <[email protected]>"]
 edition = "2018"
 description = "Low level Postgres protocol APIs"

postgres-protocol/src/types/mod.rs

@@ -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()),
+    }
+}

postgres-protocol/src/types/test.rs

@@ -1,4 +1,4 @@
-use bytes::BytesMut;
+use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
 use std::collections::HashMap;
 
@@ -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(_)))
+}

postgres-types/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "postgres-types"
-version = "0.2.2"
+version = "0.2.3"
 authors = ["Steven Fackler <[email protected]>"]
 edition = "2018"
 license = "MIT/Apache-2.0"
@@ -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 }

postgres-types/src/lib.rs

@@ -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           |
@@ -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 {
@@ -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)
     }
@@ -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,
         }
     }
@@ -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  |
@@ -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,
         }
     }

tokio-postgres/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tokio-postgres"
-version = "0.7.5"
+version = "0.7.6"
 authors = ["Steven Fackler <[email protected]>"]
 edition = "2018"
 license = "MIT/Apache-2.0"
@@ -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"] }

tokio-postgres/tests/test/types/mod.rs

@@ -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;
+}