Basic support for ltree (#1696)

* support ltree

* add default and push to PgLTree

* add more derived for ltree

* fix copy/paste

* Update sqlx-core/src/error.rs

Co-authored-by: Paolo Barbolini <paolo@paolo565.org>

* PR fixes

* ltree with name instead of OID

* custom ltree errors

* add pop ot PgLTree

* do not hide ltree behind feature flag

* bytes() instead of chars()

* apply extend_display suggestion

* add more functions to PgLTree

* fix IntoIter

* resolve PR annotation

* add tests

* remove code from arguments

* fix array

* fix setup isse

* fix(postgres): disable `ltree` tests on Postgres 9.6

Co-authored-by: Bastian Schubert <bastian.schubert@crosscard.com>
Co-authored-by: Paolo Barbolini <paolo@paolo565.org>
Co-authored-by: Austin Bonander <austin@launchbadge.com>

.github/workflows/sqlx.yml

@@ -208,6 +208,10 @@ jobs:
           key: ${{ runner.os }}-postgres-${{ matrix.runtime }}-${{ matrix.tls }}-${{ hashFiles('**/Cargo.lock') }}
 
       - uses: actions-rs/cargo@v1
+        env:
+          # FIXME: needed to disable `ltree` tests in Postgres 9.6
+          # but `PgLTree` should just fall back to text format
+          RUSTFLAGS: --cfg postgres_${{ matrix.postgres }}
         with:
           command: build
           args: >
@@ -225,6 +229,9 @@ jobs:
             --features any,postgres,macros,all-types,runtime-${{ matrix.runtime }}-${{ matrix.tls }}
         env:
           DATABASE_URL: postgres://postgres:password@localhost:5432/sqlx
+          # FIXME: needed to disable `ltree` tests in Postgres 9.6
+          # but `PgLTree` should just fall back to text format
+          RUSTFLAGS: --cfg postgres_${{ matrix.postgres }}
 
       - uses: actions-rs/cargo@v1
         with:
@@ -234,6 +241,9 @@ jobs:
             --features any,postgres,macros,migrate,all-types,runtime-${{ matrix.runtime }}-${{ matrix.tls }}
         env:
           DATABASE_URL: postgres://postgres:password@localhost:5432/sqlx?sslmode=verify-ca&sslrootcert=.%2Ftests%2Fcerts%2Fca.crt
+          # FIXME: needed to disable `ltree` tests in Postgres 9.6
+          # but `PgLTree` should just fall back to text format
+          RUSTFLAGS: --cfg postgres_${{ matrix.postgres }}
 
   mysql:
     name: MySQL

sqlx-core/src/postgres/types/ltree.rs

@@ -0,0 +1,172 @@
+use crate::decode::Decode;
+use crate::encode::{Encode, IsNull};
+use crate::error::BoxDynError;
+use crate::postgres::{
+    PgArgumentBuffer, PgHasArrayType, PgTypeInfo, PgValueFormat, PgValueRef, Postgres,
+};
+use crate::types::Type;
+use std::fmt::{self, Display, Formatter};
+use std::io::Write;
+use std::ops::Deref;
+use std::str::FromStr;
+
+/// Represents ltree specific errors
+#[derive(Debug, thiserror::Error)]
+#[non_exhaustive]
+pub enum PgLTreeParseError {
+    /// LTree labels can only contain [A-Za-z0-9_]
+    #[error("ltree label cotains invalid characters")]
+    InvalidLtreeLabel,
+
+    /// LTree version not supported
+    #[error("ltree version not supported")]
+    InvalidLtreeVersion,
+}
+
+/// Container for a Label Tree (`ltree`) in Postgres.
+///
+/// See https://www.postgresql.org/docs/current/ltree.html
+///
+/// ### Note: Requires Postgres 13+
+///
+/// This integration requires that the `ltree` type support the binary format in the Postgres
+/// wire protocol, which only became available in Postgres 13.
+/// ([Postgres 13.0 Release Notes, Additional Modules][https://www.postgresql.org/docs/13/release-13.html#id-1.11.6.11.5.14])
+///
+/// Ideally, SQLx's Postgres driver should support falling back to text format for types
+/// which don't have `typsend` and `typrecv` entries in `pg_type`, but that work still needs
+/// to be done.
+///
+/// ### Note: Extension Required
+/// The `ltree` extension is not enabled by default in Postgres. You will need to do so explicitly:
+///
+/// ```ignore
+/// CREATE EXTENSION IF NOT EXISTS "ltree";
+/// ```
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct PgLTree {
+    labels: Vec<String>,
+}
+
+impl PgLTree {
+    /// creates default/empty ltree
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// creates ltree from a [Vec<String>] without checking labels
+    pub fn new_unchecked(labels: Vec<String>) -> Self {
+        Self { labels }
+    }
+
+    /// creates ltree from an iterator with checking labels
+    pub fn from_iter<I, S>(labels: I) -> Result<Self, PgLTreeParseError>
+    where
+        S: Into<String>,
+        I: IntoIterator<Item = S>,
+    {
+        let mut ltree = Self::default();
+        for label in labels {
+            ltree.push(label.into())?;
+        }
+        Ok(ltree)
+    }
+
+    /// push a label to ltree
+    pub fn push(&mut self, label: String) -> Result<(), PgLTreeParseError> {
+        if label.len() <= 256
+            && label
+                .bytes()
+                .all(|c| c.is_ascii_alphabetic() || c.is_ascii_digit() || c == b'_')
+        {
+            self.labels.push(label);
+            Ok(())
+        } else {
+            Err(PgLTreeParseError::InvalidLtreeLabel)
+        }
+    }
+
+    /// pop a label from ltree
+    pub fn pop(&mut self) -> Option<String> {
+        self.labels.pop()
+    }
+}
+
+impl IntoIterator for PgLTree {
+    type Item = String;
+    type IntoIter = std::vec::IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.labels.into_iter()
+    }
+}
+
+impl FromStr for PgLTree {
+    type Err = PgLTreeParseError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(Self {
+            labels: s.split('.').map(|s| s.to_owned()).collect(),
+        })
+    }
+}
+
+impl Display for PgLTree {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        let mut iter = self.labels.iter();
+        if let Some(label) = iter.next() {
+            write!(f, "{}", label)?;
+            for label in iter {
+                write!(f, ".{}", label)?;
+            }
+        }
+        Ok(())
+    }
+}
+
+impl Deref for PgLTree {
+    type Target = [String];
+
+    fn deref(&self) -> &Self::Target {
+        &self.labels
+    }
+}
+
+impl Type<Postgres> for PgLTree {
+    fn type_info() -> PgTypeInfo {
+        // Since `ltree` is enabled by an extension, it does not have a stable OID.
+        PgTypeInfo::with_name("ltree")
+    }
+}
+
+impl PgHasArrayType for PgLTree {
+    fn array_type_info() -> PgTypeInfo {
+        PgTypeInfo::with_name("_ltree")
+    }
+}
+
+impl Encode<'_, Postgres> for PgLTree {
+    fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> IsNull {
+        buf.extend(1i8.to_le_bytes());
+        write!(buf, "{}", self)
+            .expect("Display implementation panicked while writing to PgArgumentBuffer");
+
+        IsNull::No
+    }
+}
+
+impl<'r> Decode<'r, Postgres> for PgLTree {
+    fn decode(value: PgValueRef<'r>) -> Result<Self, BoxDynError> {
+        match value.format() {
+            PgValueFormat::Binary => {
+                let bytes = value.as_bytes()?;
+                let version = i8::from_le_bytes([bytes[0]; 1]);
+                if version != 1 {
+                    return Err(Box::new(PgLTreeParseError::InvalidLtreeVersion));
+                }
+                Ok(Self::from_str(std::str::from_utf8(&bytes[1..])?)?)
+            }
+            PgValueFormat::Text => Ok(Self::from_str(value.as_str()?)?),
+        }
+    }
+}

sqlx-core/src/postgres/types/mod.rs

@@ -168,6 +168,7 @@ mod bytes;
 mod float;
 mod int;
 mod interval;
+mod ltree;
 mod money;
 mod range;
 mod record;
@@ -210,6 +211,8 @@ mod bit_vec;
 
 pub use array::PgHasArrayType;
 pub use interval::PgInterval;
+pub use ltree::PgLTree;
+pub use ltree::PgLTreeParseError;
 pub use money::PgMoney;
 pub use range::PgRange;
 

sqlx-macros/src/database/postgres.rs

@@ -18,6 +18,8 @@ impl_database_ext! {
 
         sqlx::postgres::types::PgMoney,
 
+        sqlx::postgres::types::PgLTree,
+
         #[cfg(feature = "uuid")]
         sqlx::types::Uuid,
 

sqlx-test/src/lib.rs

@@ -158,6 +158,7 @@ macro_rules! __test_prepared_type {
 
                 $(
                     let query = format!($sql, $text);
+                    println!("{query}");
 
                     let row = sqlx::query(&query)
                         .bind($value)

tests/postgres/setup.sql

@@ -1,3 +1,6 @@
+-- https://www.postgresql.org/docs/current/ltree.html
+CREATE EXTENSION IF NOT EXISTS ltree;
+
 -- https://www.postgresql.org/docs/current/sql-createtype.html
 CREATE TYPE status AS ENUM ('new', 'open', 'closed');
 

tests/postgres/types.rs

@@ -1,12 +1,11 @@
 extern crate time_ as time;
 
 use std::ops::Bound;
-#[cfg(feature = "decimal")]
-use std::str::FromStr;
 
 use sqlx::postgres::types::{PgInterval, PgMoney, PgRange};
 use sqlx::postgres::Postgres;
 use sqlx_test::{test_decode_type, test_prepared_type, test_type};
+use std::str::FromStr;
 
 test_type!(null<Option<i16>>(Postgres,
     "NULL::int2" == None::<i16>
@@ -513,3 +512,22 @@ test_prepared_type!(money<PgMoney>(Postgres, "123.45::money" == PgMoney(12345)))
 test_prepared_type!(money_vec<Vec<PgMoney>>(Postgres,
     "array[123.45,420.00,666.66]::money[]" == vec![PgMoney(12345), PgMoney(42000), PgMoney(66666)],
 ));
+
+// FIXME: needed to disable `ltree` tests in Postgres 9.6
+// but `PgLTree` should just fall back to text format
+#[cfg(postgres_14)]
+test_type!(ltree<sqlx::postgres::types::PgLTree>(Postgres,
+    "'Foo.Bar.Baz.Quux'::ltree" == sqlx::postgres::types::PgLTree::from_str("Foo.Bar.Baz.Quux").unwrap(),
+    "'Alpha.Beta.Delta.Gamma'::ltree" == sqlx::postgres::types::PgLTree::from_iter(["Alpha", "Beta", "Delta", "Gamma"]).unwrap(),
+));
+
+// FIXME: needed to disable `ltree` tests in Postgres 9.6
+// but `PgLTree` should just fall back to text format
+#[cfg(postgres_14)]
+test_type!(ltree_vec<Vec<sqlx::postgres::types::PgLTree>>(Postgres,
+    "array['Foo.Bar.Baz.Quux', 'Alpha.Beta.Delta.Gamma']::ltree[]" ==
+        vec![
+            sqlx::postgres::types::PgLTree::from_str("Foo.Bar.Baz.Quux").unwrap(),
+            sqlx::postgres::types::PgLTree::from_iter(["Alpha", "Beta", "Delta", "Gamma"]).unwrap()
+        ]
+));