Sqlite Collation Support

sqlx-core/src/sqlite/connection/collation.rs

@@ -0,0 +1,82 @@
+use std::cmp::Ordering;
+use std::ffi::CString;
+use std::os::raw::{c_char, c_int, c_void};
+use std::panic::{catch_unwind, UnwindSafe};
+use std::slice;
+
+use libsqlite3_sys::{sqlite3_create_collation_v2, SQLITE_OK, SQLITE_UTF8};
+
+use crate::error::Error;
+use crate::sqlite::connection::handle::ConnectionHandle;
+use crate::sqlite::SqliteError;
+
+unsafe extern "C" fn free_boxed_value<T>(p: *mut c_void) {
+    drop(Box::from_raw(p as *mut T));
+}
+
+pub(crate) fn create_collation<F>(
+    handle: &ConnectionHandle,
+    name: &str,
+    collation: F,
+) -> Result<(), Error>
+where
+    F: Fn(&str, &str) -> Ordering + Send + Sync + UnwindSafe + 'static,
+{
+    unsafe extern "C" fn call_boxed_closure<C>(
+        arg1: *mut c_void,
+        arg2: c_int,
+        arg3: *const c_void,
+        arg4: c_int,
+        arg5: *const c_void,
+    ) -> c_int
+    where
+        C: Fn(&str, &str) -> Ordering,
+    {
+        let r = catch_unwind(|| {
+            let boxed_f: *mut C = arg1 as *mut C;
+            assert!(!boxed_f.is_null(), "Internal error - null function pointer");
+            let s1 = {
+                let c_slice = slice::from_raw_parts(arg3 as *const u8, arg2 as usize);
+                String::from_utf8_lossy(c_slice)
+            };
+            let s2 = {
+                let c_slice = slice::from_raw_parts(arg5 as *const u8, arg4 as usize);
+                String::from_utf8_lossy(c_slice)
+            };
+            (*boxed_f)(s1.as_ref(), s2.as_ref())
+        });
+        let t = match r {
+            Err(_) => {
+                return -1; // FIXME How ?
+            }
+            Ok(r) => r,
+        };
+
+        match t {
+            Ordering::Less => -1,
+            Ordering::Equal => 0,
+            Ordering::Greater => 1,
+        }
+    }
+
+    let boxed_f: *mut F = Box::into_raw(Box::new(collation));
+    let c_name =
+        CString::new(name).map_err(|_| err_protocol!("Invalid collation name: {}", name))?;
+    let flags = SQLITE_UTF8;
+    let r = unsafe {
+        sqlite3_create_collation_v2(
+            handle.as_ptr(),
+            c_name.as_ptr(),
+            flags,
+            boxed_f as *mut c_void,
+            Some(call_boxed_closure::<F>),
+            Some(free_boxed_value::<F>),
+        )
+    };
+
+    if r == SQLITE_OK {
+        Ok(())
+    } else {
+        Err(Error::Database(Box::new(SqliteError::new(handle.as_ptr()))))
+    }
+}

sqlx-core/src/sqlite/connection/mod.rs

@@ -1,4 +1,6 @@
+use std::cmp::Ordering;
 use std::fmt::{self, Debug, Formatter};
+use std::panic::UnwindSafe;
 use std::sync::Arc;
 
 use futures_core::future::BoxFuture;
@@ -14,6 +16,7 @@ use crate::sqlite::connection::establish::establish;
 use crate::sqlite::statement::{SqliteStatement, StatementWorker};
 use crate::sqlite::{Sqlite, SqliteConnectOptions};
 
+mod collation;
 mod establish;
 mod executor;
 mod handle;
@@ -40,6 +43,15 @@ impl SqliteConnection {
     pub fn as_raw_handle(&mut self) -> *mut sqlite3 {
         self.handle.as_ptr()
     }
+
+    pub fn create_collation(
+        &mut self,
+        name: &str,
+        collation: impl Fn(&str, &str) -> Ordering + Send + Sync + UnwindSafe + 'static,
+    ) -> &mut SqliteConnection {
+        collation::create_collation(&self.handle, name, collation);
+        self
+    }
 }
 
 impl Debug for SqliteConnection {

tests/sqlite/sqlite.rs

@@ -1,6 +1,7 @@
 use futures::TryStreamExt;
 use sqlx::{
-    query, sqlite::Sqlite, Connect, Connection, Executor, Row, SqliteConnection, SqlitePool,
+    query, sqlite::Sqlite, sqlite::SqliteRow, Connect, Connection, Executor, Row, SqliteConnection,
+    SqlitePool,
 };
 use sqlx_test::new;
 
@@ -270,6 +271,37 @@ SELECT id, text FROM _sqlx_test;
     Ok(())
 }
 
+#[sqlx_macros::test]
+async fn it_supports_collations() -> anyhow::Result<()> {
+    let mut conn = new::<Sqlite>().await?;
+
+    conn.create_collation("test_collation", |l, r| l.cmp(r).reverse());
+
+    let _ = conn
+        .execute(
+            r#"
+CREATE TEMPORARY TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL COLLATE test_collation)
+            "#,
+        )
+        .await?;
+
+    sqlx::query("INSERT INTO users (name) VALUES (?)")
+        .bind("a")
+        .execute(&mut conn)
+        .await?;
+    sqlx::query("INSERT INTO users (name) VALUES (?)")
+        .bind("b")
+        .execute(&mut conn)
+        .await?;
+
+    let row: SqliteRow = conn
+        .fetch_one("SELECT name FROM users ORDER BY name ASC")
+        .await?;
+    let name: &str = row.try_get(0)?;
+
+    assert_eq!(name, "b");
+}
+
 #[sqlx_macros::test]
 async fn it_caches_statements() -> anyhow::Result<()> {
     let mut conn = new::<Sqlite>().await?;