Merge pull request diesel-rs#1955 from kaj/sg-sqlite-manage-updated-at

Rebase of diesel-rs#1871: Support `diesel_manage_updated_at` on SQLite

Author: weiznich

CHANGELOG.md

@@ -38,6 +38,11 @@ for Rust libraries in [RFC #1105](https://github.com/rust-lang/rfcs/blob/master/
 
 * `Nullable<Text>` now supports `ilike` expression on  in PostgreSQL.
 
+* `diesel_manage_updated_at('table_name')` is now available on SQLite. This
+  function can be called in your migrations to create a trigger which
+  automatically sets the `updated_at` column, unless that column was updated in
+  the query.
+
 ### Changed
 
 * Diesel's derives now require that `extern crate diesel;` be at your crate root

diesel/src/sqlite/connection/diesel_manage_updated_at.sql

@@ -0,0 +1,11 @@
+CREATE TRIGGER __diesel_manage_updated_at_{table_name}
+AFTER UPDATE ON {table_name}
+FOR EACH ROW WHEN
+  old.updated_at IS NULL AND
+  new.updated_at IS NULL OR
+  old.updated_at == new.updated_at
+BEGIN
+  UPDATE {table_name}
+  SET updated_at = CURRENT_TIMESTAMP
+  WHERE ROWID = new.ROWID;
+END

diesel/src/sqlite/connection/functions.rs

@@ -16,7 +16,7 @@ pub fn register<ArgsSqlType, RetSqlType, Args, Ret, F>(
     mut f: F,
 ) -> QueryResult<()>
 where
-    F: FnMut(Args) -> Ret + Send + 'static,
+    F: FnMut(&RawConnection, Args) -> Ret + Send + 'static,
     Args: Queryable<ArgsSqlType, Sqlite>,
     Ret: ToSql<RetSqlType, Sqlite>,
     Sqlite: HasSqlType<RetSqlType>,
@@ -29,12 +29,12 @@ where
         ));
     }
 
-    conn.register_sql_function(fn_name, fields_needed, deterministic, move |args| {
+    conn.register_sql_function(fn_name, fields_needed, deterministic, move |conn, 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 result = f(conn, args);
 
         let mut buf = Output::new(Vec::new(), &());
         let is_null = result.to_sql(&mut buf).map_err(Error::SerializationError)?;

diesel/src/sqlite/connection/mod.rs

@@ -50,11 +50,17 @@ impl Connection for SqliteConnection {
     type TransactionManager = AnsiTransactionManager;
 
     fn establish(database_url: &str) -> ConnectionResult<Self> {
-        RawConnection::establish(database_url).map(|conn| SqliteConnection {
+        use result::ConnectionError::CouldntSetupConfiguration;
+
+        let raw_connection = RawConnection::establish(database_url)?;
+        let conn = Self {
             statement_cache: StatementCache::new(),
-            raw_connection: conn,
+            raw_connection,
             transaction_manager: AnsiTransactionManager::new(),
-        })
+        };
+        conn.register_diesel_sql_functions()
+            .map_err(CouldntSetupConfiguration)?;
+        Ok(conn)
     }
 
     #[doc(hidden)]
@@ -218,15 +224,38 @@ impl SqliteConnection {
         &self,
         fn_name: &str,
         deterministic: bool,
-        f: F,
+        mut 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)
+        functions::register(
+            &self.raw_connection,
+            fn_name,
+            deterministic,
+            move |_, args| f(args),
+        )
+    }
+
+    fn register_diesel_sql_functions(&self) -> QueryResult<()> {
+        use sql_types::{Integer, Text};
+
+        functions::register::<Text, Integer, _, _, _>(
+            &self.raw_connection,
+            "diesel_manage_updated_at",
+            false,
+            |conn, table_name: String| {
+                conn.exec(&format!(
+                    include_str!("diesel_manage_updated_at.sql"),
+                    table_name = table_name
+                ))
+                .expect("Failed to create trigger");
+                0 // have to return *something*
+            },
+        )
     }
 }
 

diesel/src/sqlite/connection/raw.rs

@@ -4,7 +4,7 @@ use std::ffi::{CStr, CString};
 use std::io::{stderr, Write};
 use std::os::raw as libc;
 use std::ptr::NonNull;
-use std::{ptr, slice, str};
+use std::{mem, ptr, slice, str};
 
 use super::serialized_value::SerializedValue;
 use result::Error::DatabaseError;
@@ -72,7 +72,9 @@ impl RawConnection {
         f: F,
     ) -> QueryResult<()>
     where
-        F: FnMut(&[*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue> + Send + 'static,
+        F: FnMut(&Self, &[*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue>
+            + Send
+            + 'static,
     {
         let fn_name = CString::new(fn_name)?;
         let mut flags = ffi::SQLITE_UTF8;
@@ -143,9 +145,13 @@ extern "C" fn run_custom_function<F>(
     num_args: libc::c_int,
     value_ptr: *mut *mut ffi::sqlite3_value,
 ) where
-    F: FnMut(&[*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue> + Send + 'static,
+    F: FnMut(&RawConnection, &[*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue>
+        + Send
+        + 'static,
 {
-    const NULL_DATA_ERR: &str = "An unknown error occurred. sqlite3_user_data returned a null pointer. This should never happen.";
+    static NULL_DATA_ERR: &str = "An unknown error occurred. sqlite3_user_data returned a null pointer. This should never happen.";
+    static NULL_CONN_ERR: &str = "An unknown error occurred. sqlite3_context_db_handle returned a null pointer. This should never happen.";
+
     unsafe {
         let data_ptr = ffi::sqlite3_user_data(ctx);
         let data_ptr = data_ptr as *mut F;
@@ -162,19 +168,36 @@ extern "C" fn run_custom_function<F>(
         };
 
         let args = slice::from_raw_parts(value_ptr, num_args as _);
-        match f(args) {
+        let conn = match NonNull::new(ffi::sqlite3_context_db_handle(ctx)) {
+            Some(conn) => RawConnection {
+                internal_connection: conn,
+            },
+            None => {
+                ffi::sqlite3_result_error(
+                    ctx,
+                    NULL_DATA_ERR.as_ptr() as *const _ as *const _,
+                    NULL_DATA_ERR.len() as _,
+                );
+                return;
+            }
+        };
+        match f(&conn, args) {
             Ok(value) => value.result_of(ctx),
             Err(e) => {
                 let msg = e.to_string();
                 ffi::sqlite3_result_error(ctx, msg.as_ptr() as *const _, msg.len() as _);
             }
         }
+
+        mem::forget(conn);
     }
 }
 
 extern "C" fn destroy_boxed_fn<F>(data: *mut libc::c_void)
 where
-    F: FnMut(&[*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue> + Send + 'static,
+    F: FnMut(&RawConnection, &[*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue>
+        + Send
+        + 'static,
 {
     let ptr = data as *mut F;
     unsafe { Box::from_raw(ptr) };

diesel_tests/tests/connection.rs

@@ -1,6 +1,5 @@
-use diesel::dsl::sql;
 use diesel::*;
-use schema::connection_without_transaction;
+use schema::{connection_without_transaction, DropTable};
 
 table! {
     auto_time {
@@ -11,50 +10,66 @@ table! {
 }
 
 #[test]
-#[cfg(feature = "postgres")]
+#[cfg(any(feature = "postgres", feature = "sqlite"))]
 fn managing_updated_at_for_table() {
-    use self::auto_time::columns::*;
-    use self::auto_time::table as auto_time;
-    use diesel::pg::types::date_and_time::PgTimestamp;
+    use self::auto_time::dsl::*;
+    use chrono::NaiveDateTime;
+    use schema_dsl::*;
+    use std::{thread, time::Duration};
 
     // transactions have frozen time, so we can't use them
     let connection = connection_without_transaction();
-    connection
-        .execute(
-            "CREATE TABLE auto_time (
-        id SERIAL PRIMARY KEY,
-        n INTEGER,
-        updated_at TIMESTAMP
-    );",
-        )
-        .unwrap();
-    connection
-        .execute("SELECT diesel_manage_updated_at('auto_time');")
+    create_table(
+        "auto_time",
+        (
+            integer("id").primary_key().auto_increment(),
+            integer("n"),
+            timestamp("updated_at"),
+        ),
+    )
+    .execute(&connection)
+    .unwrap();
+    let _guard = DropTable {
+        connection: &connection,
+        table_name: "auto_time",
+    };
+    sql_query("SELECT diesel_manage_updated_at('auto_time')")
+        .execute(&connection)
         .unwrap();
 
-    connection
-        .execute("INSERT INTO auto_time (n) VALUES (2), (1), (5);")
+    insert_into(auto_time)
+        .values(&vec![n.eq(2), n.eq(1), n.eq(5)])
+        .execute(&connection)
         .unwrap();
-    let result = select(sql("COUNT(*) FROM auto_time WHERE updated_at IS NULL"))
+
+    let result = auto_time
+        .count()
+        .filter(updated_at.is_null())
         .get_result::<i64>(&connection);
     assert_eq!(Ok(3), result);
 
-    connection
-        .execute("UPDATE auto_time SET n = n + 1 WHERE true;")
+    update(auto_time)
+        .set(n.eq(n + 1))
+        .execute(&connection)
         .unwrap();
-    let result = select(sql("COUNT(*) FROM auto_time WHERE updated_at IS NULL"))
+
+    let result = auto_time
+        .count()
+        .filter(updated_at.is_null())
         .get_result::<i64>(&connection);
     assert_eq!(Ok(0), result);
 
+    if cfg!(feature = "sqlite") {
+        // SQLite only has second precision
+        thread::sleep(Duration::from_millis(1000));
+    }
+
     let query = auto_time.find(2).select(updated_at);
-    let old_time: PgTimestamp = query.first(&connection).unwrap();
+    let old_time: NaiveDateTime = query.first(&connection).unwrap();
     update(auto_time.find(2))
         .set(n.eq(0))
         .execute(&connection)
         .unwrap();
-    let new_time: PgTimestamp = query.first(&connection).unwrap();
+    let new_time: NaiveDateTime = query.first(&connection).unwrap();
     assert!(old_time < new_time);
-
-    // clean up because we aren't in a transaction
-    connection.execute("DROP TABLE auto_time;").unwrap();
 }

diesel_tests/tests/lib.rs

@@ -3,6 +3,7 @@
 
 #[macro_use]
 extern crate assert_matches;
+extern crate chrono;
 #[macro_use]
 extern crate diesel;
 #[macro_use]