Merge pull request diesel-rs#1691 from diesel-rs/sg-sqlite-custom-functions

Support custom functions on SQLite

Author: sgrif

CHANGELOG.md

@@ -23,6 +23,9 @@ for Rust libraries in [RFC #1105](https://github.com/rust-lang/rfcs/blob/master/
 * Added `sqlite-bundled` feature to `diesel_cli` to make installing on
   some platforms easier.
 
+* Custom SQL functions can now be used with SQLite. See [the
+  docs][sql-function-sqlite-1-3-0] for details.
+
 * All functions and operators provided by Diesel can now be used with numeric
   operators if the SQL type supports it.
 

diesel/src/expression/functions/mod.rs

@@ -260,10 +260,94 @@ macro_rules! __diesel_sql_function_body {
                     }
                 }
             }
+
+            __diesel_sqlite_register_fn! {
+                type_args = ($($type_args)*),
+                fn_name = $fn_name,
+                args = ($($arg_name,)+),
+                sql_args = ($($arg_type,)+),
+                ret = $return_type,
+            }
         }
     }
 }
 
+#[macro_export]
+#[doc(hidden)]
+#[cfg(feature = "sqlite")]
+macro_rules! __diesel_sqlite_register_fn {
+    // We can't handle generic functions for SQLite
+    (
+        type_args = ($($type_args:tt)+),
+        $($rest:tt)*
+    ) => {
+    };
+
+    (
+        type_args = (),
+        fn_name = $fn_name:ident,
+        args = ($($args:ident,)+),
+        sql_args = $sql_args:ty,
+        ret = $ret:ty,
+    ) => {
+        #[allow(dead_code)]
+        /// Registers an implementation for this function on the given connection
+        ///
+        /// This function must be called for every `SqliteConnection` before
+        /// this SQL function can be used on SQLite. The implementation must be
+        /// deterministic (returns the same result given the same arguments). If
+        /// the function is nondeterministic, call
+        /// `register_nondeterministic_impl` instead.
+        pub fn register_impl<F, Ret, $($args,)+>(
+            conn: &$crate::SqliteConnection,
+            f: F,
+        ) -> $crate::QueryResult<()>
+        where
+            F: Fn($($args,)+) -> Ret + Send + 'static,
+            ($($args,)+): $crate::deserialize::Queryable<$sql_args, $crate::sqlite::Sqlite>,
+            Ret: $crate::serialize::ToSql<$ret, $crate::sqlite::Sqlite>,
+        {
+            conn.register_sql_function::<$sql_args, $ret, _, _, _>(
+                stringify!($fn_name),
+                true,
+                move |($($args,)+)| f($($args),+),
+            )
+        }
+
+        #[allow(dead_code)]
+        /// Registers an implementation for this function on the given connection
+        ///
+        /// This function must be called for every `SqliteConnection` before
+        /// this SQL function can be used on SQLite.
+        /// `register_nondeterministic_impl` should only be used if your
+        /// function can return different results with the same arguments (e.g.
+        /// `random`). If your function is deterministic, you should call
+        /// `register_impl` instead.
+        pub fn register_nondeterministic_impl<F, Ret, $($args,)+>(
+            conn: &$crate::SqliteConnection,
+            mut f: F,
+        ) -> $crate::QueryResult<()>
+        where
+            F: FnMut($($args,)+) -> Ret + Send + 'static,
+            ($($args,)+): $crate::deserialize::Queryable<$sql_args, $crate::sqlite::Sqlite>,
+            Ret: $crate::serialize::ToSql<$ret, $crate::sqlite::Sqlite>,
+        {
+            conn.register_sql_function::<$sql_args, $ret, _, _, _>(
+                stringify!($fn_name),
+                false,
+                move |($($args,)+)| f($($args),+),
+            )
+        }
+    };
+}
+
+#[macro_export]
+#[doc(hidden)]
+#[cfg(not(feature = "sqlite"))]
+macro_rules! __diesel_sqlite_register_fn {
+    ($($token:tt)*) => {};
+}
+
 #[macro_export]
 /// Declare a sql function for use in your code.
 ///
@@ -325,7 +409,7 @@ macro_rules! __diesel_sql_function_body {
 /// Most attributes given to this macro will be put on the generated function
 /// (including doc comments).
 ///
-/// # Example
+/// # Adding Doc Comments
 ///
 /// ```no_run
 /// # #[macro_use] extern crate diesel;
@@ -350,6 +434,8 @@ macro_rules! __diesel_sql_function_body {
 /// # }
 /// ```
 ///
+/// # Special Attributes
+///
 /// There are a handful of special attributes that Diesel will recognize. They
 /// are:
 ///
@@ -382,6 +468,49 @@ macro_rules! __diesel_sql_function_body {
 /// crates.select(sum(id));
 /// # }
 /// ```
+///
+/// # Use with SQLite
+///
+/// On most backends, the implementation of the function is defined in a
+/// migration using `CREATE FUNCTION`. On SQLite, the function is implemented in
+/// Rust instead. You must call `register_impl` or
+/// `register_nondeterministic_impl` with every connection before you can use
+/// the function.
+///
+/// These functions will only be generated if the `sqlite` feature is enabled,
+/// and the function is not generic. Generic functions and variadic functions
+/// are not supported on SQLite.
+///
+/// ```rust
+/// # #[macro_use] extern crate diesel;
+/// # use diesel::*;
+/// #
+/// # #[cfg(feature = "sqlite")]
+/// # fn main() {
+/// #     run_test().unwrap();
+/// # }
+/// #
+/// # #[cfg(not(feature = "sqlite"))]
+/// # fn main() {
+/// # }
+/// #
+/// use diesel::sql_types::{Integer, Double};
+/// sql_function!(fn add_mul(x: Integer, y: Integer, z: Double) -> Double);
+///
+/// # #[cfg(feature = "sqlite")]
+/// # fn run_test() -> Result<(), Box<::std::error::Error>> {
+/// let connection = SqliteConnection::establish(":memory:")?;
+///
+/// add_mul::register_impl(&connection, |x: i32, y: i32, z: f64| {
+///     (x + y) as f64 * z
+/// })?;
+///
+/// let result = select(add_mul(1, 2, 1.5))
+///     .get_result::<f64>(&connection)?;
+/// assert_eq!(4.5, result);
+/// #     Ok(())
+/// # }
+/// ```
 macro_rules! sql_function {
     ($(#$meta:tt)* fn $fn_name:ident $args:tt $(;)*) => {
         sql_function!($(#[$meta])* fn $fn_name $args -> ());

diesel/src/sqlite/connection/functions.rs

@@ -0,0 +1,73 @@
+extern crate libsqlite3_sys as ffi;
+
+use deserialize::{FromSqlRow, Queryable};
+use result::{DatabaseErrorKind, Error, QueryResult};
+use row::Row;
+use serialize::{IsNull, Output, ToSql};
+use sql_types::HasSqlType;
+use super::raw::RawConnection;
+use super::serialized_value::SerializedValue;
+use super::{Sqlite, SqliteValue};
+
+pub fn register<ArgsSqlType, RetSqlType, Args, Ret, F>(
+    conn: &RawConnection,
+    fn_name: &str,
+    deterministic: bool,
+    mut f: F,
+) -> QueryResult<()>
+where
+    F: FnMut(Args) -> Ret + Send + 'static,
+    Args: Queryable<ArgsSqlType, Sqlite>,
+    Ret: ToSql<RetSqlType, Sqlite>,
+    Sqlite: HasSqlType<RetSqlType>,
+{
+    let fields_needed = Args::Row::FIELDS_NEEDED;
+    if fields_needed > 127 {
+        return Err(Error::DatabaseError(
+            DatabaseErrorKind::UnableToSendCommand,
+            Box::new("SQLite functions cannot take more than 127 parameters".to_string()),
+        ));
+    }
+
+    conn.register_sql_function(fn_name, fields_needed, deterministic, move |args| {
+        let mut row = FunctionRow { args };
+        let args_row = Args::Row::build_from_row(&mut row).map_err(Error::DeserializationError)?;
+        let args = Args::build(args_row);
+
+        let result = f(args);
+
+        let mut buf = Output::new(Vec::new(), &());
+        let is_null = result.to_sql(&mut buf).map_err(Error::SerializationError)?;
+
+        let bytes = if let IsNull::Yes = is_null {
+            None
+        } else {
+            Some(buf.into_inner())
+        };
+
+        Ok(SerializedValue {
+            ty: Sqlite::metadata(&()),
+            data: bytes,
+        })
+    })?;
+    Ok(())
+}
+
+struct FunctionRow<'a> {
+    args: &'a [*mut ffi::sqlite3_value],
+}
+
+impl<'a> Row<Sqlite> for FunctionRow<'a> {
+    fn take(&mut self) -> Option<&SqliteValue> {
+        self.args.split_first().and_then(|(&first, rest)| {
+            self.args = rest;
+            unsafe { SqliteValue::new(first) }
+        })
+    }
+
+    fn next_is_null(&self, count: usize) -> bool {
+        self.args[..count]
+            .iter()
+            .all(|&p| unsafe { SqliteValue::new(p) }.is_none())
+    }
+}

diesel/src/sqlite/connection/mod.rs

@@ -1,7 +1,9 @@
 extern crate libsqlite3_sys as ffi;
 
+mod functions;
 #[doc(hidden)]
 pub mod raw;
+mod serialized_value;
 mod stmt;
 mod statement_iterator;
 mod sqlite_value;
@@ -18,6 +20,7 @@ use result::*;
 use self::raw::RawConnection;
 use self::statement_iterator::*;
 use self::stmt::{Statement, StatementUse};
+use serialize::ToSql;
 use sql_types::HasSqlType;
 use sqlite::Sqlite;
 
@@ -208,6 +211,22 @@ impl SqliteConnection {
             Statement::prepare(&self.raw_connection, sql)
         })
     }
+
+    #[doc(hidden)]
+    pub fn register_sql_function<ArgsSqlType, RetSqlType, Args, Ret, F>(
+        &self,
+        fn_name: &str,
+        deterministic: bool,
+        f: F,
+    ) -> QueryResult<()>
+    where
+        F: FnMut(Args) -> Ret + Send + 'static,
+        Args: Queryable<ArgsSqlType, Sqlite>,
+        Ret: ToSql<RetSqlType, Sqlite>,
+        Sqlite: HasSqlType<RetSqlType>,
+    {
+        functions::register(&self.raw_connection, fn_name, deterministic, f)
+    }
 }
 
 fn error_message(err_code: libc::c_int) -> &'static str {
@@ -269,4 +288,56 @@ mod tests {
         assert_eq!(Ok(true), query.get_result(&connection));
         assert_eq!(1, connection.statement_cache.len());
     }
+
+    use sql_types::Text;
+    sql_function!(fn fun_case(x: Text) -> Text);
+
+    #[test]
+    fn register_custom_function() {
+        let connection = SqliteConnection::establish(":memory:").unwrap();
+        fun_case::register_impl(&connection, |x: String| {
+            x.chars()
+                .enumerate()
+                .map(|(i, c)| {
+                    if i % 2 == 0 {
+                        c.to_lowercase().to_string()
+                    } else {
+                        c.to_uppercase().to_string()
+                    }
+                })
+                .collect::<String>()
+        }).unwrap();
+
+        let mapped_string = ::select(fun_case("foobar"))
+            .get_result::<String>(&connection)
+            .unwrap();
+        assert_eq!("fOoBaR", mapped_string);
+    }
+
+    sql_function!(fn my_add(x: Integer, y: Integer) -> Integer);
+
+    #[test]
+    fn register_multiarg_function() {
+        let connection = SqliteConnection::establish(":memory:").unwrap();
+        my_add::register_impl(&connection, |x: i32, y: i32| x + y).unwrap();
+
+        let added = ::select(my_add(1, 2)).get_result::<i32>(&connection);
+        assert_eq!(Ok(3), added);
+    }
+
+    sql_function!(fn add_counter(x: Integer) -> Integer);
+
+    #[test]
+    fn register_nondeterministic_function() {
+        let connection = SqliteConnection::establish(":memory:").unwrap();
+        let mut y = 0;
+        add_counter::register_nondeterministic_impl(&connection, move |x: i32| {
+            y += 1;
+            x + y
+        }).unwrap();
+
+        let added = ::select((add_counter(1), add_counter(1), add_counter(1)))
+            .get_result::<(i32, i32, i32)>(&connection);
+        assert_eq!(Ok((2, 3, 4)), added);
+    }
 }