Skip to main content

diesel/sqlite/connection/
functions.rs

1#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
2extern crate libsqlite3_sys as ffi;
3
4#[cfg(all(target_family = "wasm", target_os = "unknown"))]
5use sqlite_wasm_rs as ffi;
6
7use super::raw::RawConnection;
8use super::{Sqlite, SqliteAggregateFunction, SqliteBindValue, SqliteConnection};
9use crate::backend::Backend;
10use crate::deserialize::{FromSqlRow, StaticallySizedRow};
11use crate::result::{DatabaseErrorKind, Error, QueryResult};
12use crate::row::{Field, PartialRow, Row, RowIndex, RowSealed};
13use crate::serialize::{IsNull, Output, ToSql};
14use crate::sql_types::HasSqlType;
15use crate::sqlite::SqliteFunctionBehavior;
16use crate::sqlite::SqliteValue;
17use crate::sqlite::connection::bind_collector::SqliteBindValueRef;
18use crate::sqlite::connection::sqlite_value::OwnedSqliteValue;
19use alloc::boxed::Box;
20use alloc::string::ToString;
21
22pub(super) fn register<ArgsSqlType, RetSqlType, Args, Ret, F>(
23    conn: &RawConnection,
24    fn_name: &str,
25    behavior: SqliteFunctionBehavior,
26    mut f: F,
27) -> QueryResult<()>
28where
29    F: FnMut(&RawConnection, Args) -> Ret + core::panic::UnwindSafe + Send + 'static,
30    Args: FromSqlRow<ArgsSqlType, Sqlite> + StaticallySizedRow<ArgsSqlType, Sqlite>,
31    Ret: ToSql<RetSqlType, Sqlite>,
32    Sqlite: HasSqlType<RetSqlType>,
33{
34    let fields_needed = Args::FIELD_COUNT;
35    if fields_needed > 127 {
36        return Err(Error::DatabaseError(
37            DatabaseErrorKind::UnableToSendCommand,
38            Box::new("SQLite functions cannot take more than 127 parameters".to_string()),
39        ));
40    }
41
42    conn.register_sql_function(fn_name, fields_needed, behavior, move |conn, args| {
43        let args = build_sql_function_args::<ArgsSqlType, Args>(args, conn.internal_connection)?;
44
45        Ok(f(conn, args))
46    })?;
47    Ok(())
48}
49
50pub(super) fn register_noargs<RetSqlType, Ret, F>(
51    conn: &RawConnection,
52    fn_name: &str,
53    behavior: SqliteFunctionBehavior,
54    mut f: F,
55) -> QueryResult<()>
56where
57    F: FnMut() -> Ret + core::panic::UnwindSafe + Send + 'static,
58    Ret: ToSql<RetSqlType, Sqlite>,
59    Sqlite: HasSqlType<RetSqlType>,
60{
61    conn.register_sql_function(fn_name, 0, behavior, move |_, _| Ok(f()))?;
62    Ok(())
63}
64
65pub(super) fn register_aggregate<ArgsSqlType, RetSqlType, Args, Ret, A>(
66    conn: &RawConnection,
67    fn_name: &str,
68    behavior: SqliteFunctionBehavior,
69) -> QueryResult<()>
70where
71    A: SqliteAggregateFunction<Args, Output = Ret> + 'static + Send + core::panic::UnwindSafe,
72    Args: FromSqlRow<ArgsSqlType, Sqlite> + StaticallySizedRow<ArgsSqlType, Sqlite>,
73    Ret: ToSql<RetSqlType, Sqlite>,
74    Sqlite: HasSqlType<RetSqlType>,
75{
76    let fields_needed = Args::FIELD_COUNT;
77    if fields_needed > 127 {
78        return Err(Error::DatabaseError(
79            DatabaseErrorKind::UnableToSendCommand,
80            Box::new("SQLite functions cannot take more than 127 parameters".to_string()),
81        ));
82    }
83
84    conn.register_aggregate_function::<ArgsSqlType, RetSqlType, Args, Ret, A>(
85        fn_name,
86        fields_needed,
87        behavior,
88    )?;
89
90    Ok(())
91}
92
93pub(super) fn build_sql_function_args<ArgsSqlType, Args>(
94    args: &mut [*mut ffi::sqlite3_value],
95    connection: core::ptr::NonNull<ffi::sqlite3>,
96) -> Result<Args, Error>
97where
98    Args: FromSqlRow<ArgsSqlType, Sqlite>,
99{
100    let row = FunctionRow::new(args, connection);
101    Args::build_from_row(&row).map_err(Error::DeserializationError)
102}
103
104// clippy is wrong here, the let binding is required
105// for lifetime reasons
106#[allow(clippy::let_unit_value)]
107pub(super) fn process_sql_function_result<RetSqlType, Ret>(
108    result: &'_ Ret,
109) -> QueryResult<SqliteBindValueRef<'_>>
110where
111    Ret: ToSql<RetSqlType, Sqlite>,
112    Sqlite: HasSqlType<RetSqlType>,
113{
114    let mut metadata_lookup = ();
115    let value = SqliteBindValue {
116        inner: SqliteBindValueRef::Null,
117    };
118    let mut buf = Output::new(value, &mut metadata_lookup);
119    let is_null = result.to_sql(&mut buf).map_err(Error::SerializationError)?;
120
121    if let IsNull::Yes = is_null {
122        Ok(SqliteBindValueRef::Null)
123    } else {
124        Ok(buf.into_inner().inner)
125    }
126}
127
128struct FunctionRow<'a> {
129    args: &'a [Option<OwnedSqliteValue>],
130    field_count: usize,
131    connection: core::ptr::NonNull<ffi::sqlite3>,
132}
133
134impl FunctionRow<'_> {
135    #[allow(unsafe_code)] // complicated ptr cast
136    fn new(
137        args: &mut [*mut ffi::sqlite3_value],
138        connection: core::ptr::NonNull<ffi::sqlite3>,
139    ) -> Self {
140        let lengths = args.len();
141        let args = unsafe {
142            core::slice::from_raw_parts(
143                // This cast is safe because:
144                // * Casting from a pointer to an array to a pointer to the first array
145                // element is safe
146                // * Casting from a raw pointer to `NonNull<T>` is safe,
147                // because `NonNull` is #[repr(transparent)]
148                // * Casting from `NonNull<T>` to `OwnedSqliteValue` is safe,
149                // as the struct is `#[repr(transparent)]
150                // * Casting from `NonNull<T>` to `Option<NonNull<T>>` as the documentation
151                // states: "This is so that enums may use this forbidden value as a discriminant –
152                // Option<NonNull<T>> has the same size as *mut T"
153                // * The last point remains true for `OwnedSqliteValue` as `#[repr(transparent)]
154                // guarantees the same layout as the inner type
155                args as *mut [*mut ffi::sqlite3_value] as *mut ffi::sqlite3_value
156                    as *mut Option<OwnedSqliteValue>,
157                lengths,
158            )
159        };
160
161        Self {
162            field_count: lengths,
163            args,
164            connection,
165        }
166    }
167}
168
169impl RowSealed for FunctionRow<'_> {}
170
171impl<'a> Row<'a, Sqlite> for FunctionRow<'a> {
172    type Field<'f>
173        = FunctionArgument<'f>
174    where
175        'a: 'f,
176        Self: 'f;
177    type InnerPartialRow = Self;
178
179    fn field_count(&self) -> usize {
180        self.field_count
181    }
182
183    fn get<'b, I>(&'b self, idx: I) -> Option<Self::Field<'b>>
184    where
185        'a: 'b,
186        Self: crate::row::RowIndex<I>,
187    {
188        let col_idx = self.idx(idx)?;
189        Some(FunctionArgument {
190            args: self.args,
191            col_idx,
192            connection: self.connection,
193        })
194    }
195
196    fn partial_row(&self, range: core::ops::Range<usize>) -> PartialRow<'_, Self::InnerPartialRow> {
197        PartialRow::new(self, range)
198    }
199}
200
201impl RowIndex<usize> for FunctionRow<'_> {
202    fn idx(&self, idx: usize) -> Option<usize> {
203        if idx < self.field_count() {
204            Some(idx)
205        } else {
206            None
207        }
208    }
209}
210
211impl<'a> RowIndex<&'a str> for FunctionRow<'_> {
212    fn idx(&self, _idx: &'a str) -> Option<usize> {
213        None
214    }
215}
216
217struct FunctionArgument<'a> {
218    args: &'a [Option<OwnedSqliteValue>],
219    col_idx: usize,
220    connection: core::ptr::NonNull<ffi::sqlite3>,
221}
222
223impl<'a> Field<'a, Sqlite> for FunctionArgument<'a> {
224    fn field_name(&self) -> Option<&str> {
225        None
226    }
227
228    fn is_null(&self) -> bool {
229        self.value().is_none()
230    }
231
232    fn value(&self) -> Option<<Sqlite as Backend>::RawValue<'_>> {
233        SqliteValue::from_function_row(self.args, self.col_idx, self.connection)
234    }
235}
236
237impl SqliteConnection {
238    #[doc(hidden)]
239    pub fn register_sql_function<ArgsSqlType, RetSqlType, Args, Ret, F>(
240        &mut self,
241        fn_name: &str,
242        behavior: SqliteFunctionBehavior,
243        mut f: F,
244    ) -> QueryResult<()>
245    where
246        F: FnMut(Args) -> Ret + core::panic::UnwindSafe + Send + 'static,
247        Args: FromSqlRow<ArgsSqlType, Sqlite> + StaticallySizedRow<ArgsSqlType, Sqlite>,
248        Ret: ToSql<RetSqlType, Sqlite>,
249        Sqlite: HasSqlType<RetSqlType>,
250    {
251        register(&self.raw_connection, fn_name, behavior, move |_, args| {
252            f(args)
253        })
254    }
255
256    #[doc(hidden)]
257    pub fn register_noarg_sql_function<RetSqlType, Ret, F>(
258        &mut self,
259        fn_name: &str,
260        behavior: SqliteFunctionBehavior,
261        f: F,
262    ) -> QueryResult<()>
263    where
264        F: FnMut() -> Ret + core::panic::UnwindSafe + Send + 'static,
265        Ret: ToSql<RetSqlType, Sqlite>,
266        Sqlite: HasSqlType<RetSqlType>,
267    {
268        register_noargs(&self.raw_connection, fn_name, behavior, f)
269    }
270
271    #[doc(hidden)]
272    pub fn register_aggregate_function<ArgsSqlType, RetSqlType, Args, Ret, A>(
273        &mut self,
274        fn_name: &str,
275        behavior: SqliteFunctionBehavior,
276    ) -> QueryResult<()>
277    where
278        A: SqliteAggregateFunction<Args, Output = Ret> + 'static + Send + core::panic::UnwindSafe,
279        Args: FromSqlRow<ArgsSqlType, Sqlite> + StaticallySizedRow<ArgsSqlType, Sqlite>,
280        Ret: ToSql<RetSqlType, Sqlite>,
281        Sqlite: HasSqlType<RetSqlType>,
282    {
283        register_aggregate::<_, _, _, _, A>(&self.raw_connection, fn_name, behavior)
284    }
285
286    /// Register a collation function.
287    ///
288    /// `collation` must always return the same answer given the same inputs.
289    /// If `collation` panics and unwinds the stack, the process is aborted, since it is used
290    /// across a C FFI boundary, which cannot be unwound across and there is no way to
291    /// signal failures via the SQLite interface in this case..
292    ///
293    /// If the name is already registered it will be overwritten.
294    ///
295    /// This method will return an error if registering the function fails, either due to an
296    /// out-of-memory situation or because a collation with that name already exists and is
297    /// currently being used in parallel by a query.
298    ///
299    /// The collation needs to be specified when creating a table:
300    /// `CREATE TABLE my_table ( str TEXT COLLATE MY_COLLATION )`,
301    /// where `MY_COLLATION` corresponds to name passed as `collation_name`.
302    ///
303    /// # Example
304    ///
305    /// ```rust
306    /// # include!("../../doctest_setup.rs");
307    /// #
308    /// # // This test registers a Rust callback that sqlite invokes through the
309    /// # // FFI boundary, which is not supported by miri `-Zmiri-native-lib`.
310    /// # #[cfg(miri)]
311    /// # fn main() {}
312    /// # #[cfg(not(miri))]
313    /// # fn main() {
314    /// #     run_test().unwrap();
315    /// # }
316    /// #
317    /// # fn run_test() -> QueryResult<()> {
318    /// #     let mut conn = SqliteConnection::establish(":memory:").unwrap();
319    /// // sqlite NOCASE only works for ASCII characters,
320    /// // this collation allows handling UTF-8 (barring locale differences)
321    /// conn.register_collation("RUSTNOCASE", |rhs, lhs| {
322    ///     rhs.to_lowercase().cmp(&lhs.to_lowercase())
323    /// })
324    /// # }
325    /// ```
326    pub fn register_collation<F>(&mut self, collation_name: &str, collation: F) -> QueryResult<()>
327    where
328        F: Fn(&str, &str) -> core::cmp::Ordering + Send + 'static + core::panic::UnwindSafe,
329    {
330        self.raw_connection
331            .register_collation_function(collation_name, collation)
332    }
333}
334
335// miri doesn't support callback over ffi at this point
336#[cfg(all(test, not(miri)))]
337mod tests {
338    use super::*;
339    use crate::connection::SimpleConnection;
340    use crate::prelude::*;
341    use crate::sql_types::{Integer, Text};
342
343    fn connection() -> SqliteConnection {
344        SqliteConnection::establish(":memory:").unwrap()
345    }
346
347    #[declare_sql_function]
348    extern "SQL" {
349        fn fun_case(x: Text) -> Text;
350        fn my_add(x: Integer, y: Integer) -> Integer;
351        fn answer() -> Integer;
352        fn add_counter(x: Integer) -> Integer;
353
354        #[aggregate]
355        fn my_sum(expr: Integer) -> Integer;
356        #[aggregate]
357        fn range_max(expr1: Integer, expr2: Integer, expr3: Integer) -> Nullable<Integer>;
358    }
359
360    #[diesel_test_helper::test]
361    fn register_custom_function() {
362        let connection = &mut connection();
363        fun_case_utils::register_impl(connection, |x: String| {
364            x.chars()
365                .enumerate()
366                .map(|(i, c)| {
367                    if i % 2 == 0 {
368                        c.to_lowercase().to_string()
369                    } else {
370                        c.to_uppercase().to_string()
371                    }
372                })
373                .collect::<String>()
374        })
375        .unwrap();
376
377        let mapped_string = crate::select(fun_case("foobar"))
378            .get_result::<String>(connection)
379            .unwrap();
380        assert_eq!("fOoBaR", mapped_string);
381    }
382
383    #[diesel_test_helper::test]
384    fn register_multiarg_function() {
385        let connection = &mut connection();
386        my_add_utils::register_impl(connection, |x: i32, y: i32| x + y).unwrap();
387
388        let added = crate::select(my_add(1, 2)).get_result::<i32>(connection);
389        assert_eq!(Ok(3), added);
390    }
391
392    #[diesel_test_helper::test]
393    fn register_noarg_function() {
394        let connection = &mut connection();
395        answer_utils::register_impl(connection, || 42).unwrap();
396
397        let answer = crate::select(answer()).get_result::<i32>(connection);
398        assert_eq!(Ok(42), answer);
399    }
400
401    #[diesel_test_helper::test]
402    fn register_nondeterministic_noarg_function() {
403        let connection = &mut connection();
404        answer_utils::register_nondeterministic_impl(connection, || 42).unwrap();
405
406        let answer = crate::select(answer()).get_result::<i32>(connection);
407        assert_eq!(Ok(42), answer);
408    }
409
410    #[diesel_test_helper::test]
411    fn register_nondeterministic_function() {
412        let connection = &mut connection();
413        let mut y = 0;
414        add_counter_utils::register_nondeterministic_impl(connection, move |x: i32| {
415            y += 1;
416            x + y
417        })
418        .unwrap();
419
420        let added = crate::select((add_counter(1), add_counter(1), add_counter(1)))
421            .get_result::<(i32, i32, i32)>(connection);
422        assert_eq!(Ok((2, 3, 4)), added);
423    }
424
425    #[derive(Default)]
426    struct MySum {
427        sum: i32,
428    }
429
430    impl SqliteAggregateFunction<i32> for MySum {
431        type Output = i32;
432
433        fn step(&mut self, expr: i32) {
434            self.sum += expr;
435        }
436
437        fn finalize(aggregator: Option<Self>) -> Self::Output {
438            aggregator.map(|a| a.sum).unwrap_or_default()
439        }
440    }
441
442    table! {
443        my_sum_example {
444            id -> Integer,
445            value -> Integer,
446        }
447    }
448
449    #[diesel_test_helper::test]
450    fn register_aggregate_function() {
451        use self::my_sum_example::dsl::*;
452
453        let connection = &mut connection();
454        crate::sql_query(
455            "CREATE TABLE my_sum_example (id integer primary key autoincrement, value integer)",
456        )
457        .execute(connection)
458        .unwrap();
459        crate::sql_query("INSERT INTO my_sum_example (value) VALUES (1), (2), (3)")
460            .execute(connection)
461            .unwrap();
462
463        my_sum_utils::register_impl_with_behavior::<MySum, _>(
464            connection,
465            SqliteFunctionBehavior::DETERMINISTIC,
466        )
467        .unwrap();
468
469        let result = my_sum_example
470            .select(my_sum(value))
471            .get_result::<i32>(connection);
472        assert_eq!(Ok(6), result);
473    }
474
475    #[diesel_test_helper::test]
476    fn register_aggregate_function_returns_finalize_default_on_empty_set() {
477        use self::my_sum_example::dsl::*;
478
479        let connection = &mut connection();
480        crate::sql_query(
481            "CREATE TABLE my_sum_example (id integer primary key autoincrement, value integer)",
482        )
483        .execute(connection)
484        .unwrap();
485
486        my_sum_utils::register_impl_with_behavior::<MySum, _>(
487            connection,
488            SqliteFunctionBehavior::DETERMINISTIC,
489        )
490        .unwrap();
491
492        let result = my_sum_example
493            .select(my_sum(value))
494            .get_result::<i32>(connection);
495        assert_eq!(Ok(0), result);
496    }
497
498    #[derive(Default)]
499    struct RangeMax<T> {
500        max_value: Option<T>,
501    }
502
503    impl<T: Default + Ord + Copy + Clone> SqliteAggregateFunction<(T, T, T)> for RangeMax<T> {
504        type Output = Option<T>;
505
506        fn step(&mut self, (x0, x1, x2): (T, T, T)) {
507            let max = if x0 >= x1 && x0 >= x2 {
508                x0
509            } else if x1 >= x0 && x1 >= x2 {
510                x1
511            } else {
512                x2
513            };
514
515            self.max_value = match self.max_value {
516                Some(current_max_value) if max > current_max_value => Some(max),
517                None => Some(max),
518                _ => self.max_value,
519            };
520        }
521
522        fn finalize(aggregator: Option<Self>) -> Self::Output {
523            aggregator?.max_value
524        }
525    }
526
527    table! {
528        range_max_example {
529            id -> Integer,
530            value1 -> Integer,
531            value2 -> Integer,
532            value3 -> Integer,
533        }
534    }
535
536    #[diesel_test_helper::test]
537    fn register_aggregate_multiarg_function() {
538        use self::range_max_example::dsl::*;
539
540        let connection = &mut connection();
541        crate::sql_query(
542            r#"CREATE TABLE range_max_example (
543                id integer primary key autoincrement,
544                value1 integer,
545                value2 integer,
546                value3 integer
547            )"#,
548        )
549        .execute(connection)
550        .unwrap();
551        crate::sql_query(
552            "INSERT INTO range_max_example (value1, value2, value3) VALUES (3, 2, 1), (2, 2, 2)",
553        )
554        .execute(connection)
555        .unwrap();
556
557        range_max_utils::register_impl_with_behavior::<RangeMax<i32>, _, _, _>(
558            connection,
559            SqliteFunctionBehavior::DETERMINISTIC,
560        )
561        .unwrap();
562        let result = range_max_example
563            .select(range_max(value1, value2, value3))
564            .get_result::<Option<i32>>(connection)
565            .unwrap();
566        assert_eq!(Some(3), result);
567    }
568
569    table! {
570        my_collation_example {
571            id -> Integer,
572            value -> Text,
573        }
574    }
575
576    #[diesel_test_helper::test]
577    fn register_collation_function() {
578        use self::my_collation_example::dsl::*;
579
580        let connection = &mut connection();
581
582        connection
583            .register_collation("RUSTNOCASE", |rhs, lhs| {
584                rhs.to_lowercase().cmp(&lhs.to_lowercase())
585            })
586            .unwrap();
587
588        crate::sql_query(
589                "CREATE TABLE my_collation_example (id integer primary key autoincrement, value text collate RUSTNOCASE)",
590            ).execute(connection)
591            .unwrap();
592        crate::sql_query(
593            "INSERT INTO my_collation_example (value) VALUES ('foo'), ('FOo'), ('f00')",
594        )
595        .execute(connection)
596        .unwrap();
597
598        let result = my_collation_example
599            .filter(value.eq("foo"))
600            .select(value)
601            .load::<String>(connection);
602        assert_eq!(
603            Ok(&["foo".to_owned(), "FOo".to_owned()][..]),
604            result.as_ref().map(|vec| vec.as_ref())
605        );
606
607        let result = my_collation_example
608            .filter(value.eq("FOO"))
609            .select(value)
610            .load::<String>(connection);
611        assert_eq!(
612            Ok(&["foo".to_owned(), "FOo".to_owned()][..]),
613            result.as_ref().map(|vec| vec.as_ref())
614        );
615
616        let result = my_collation_example
617            .filter(value.eq("f00"))
618            .select(value)
619            .load::<String>(connection);
620        assert_eq!(
621            Ok(&["f00".to_owned()][..]),
622            result.as_ref().map(|vec| vec.as_ref())
623        );
624
625        let result = my_collation_example
626            .filter(value.eq("F00"))
627            .select(value)
628            .load::<String>(connection);
629        assert_eq!(
630            Ok(&["f00".to_owned()][..]),
631            result.as_ref().map(|vec| vec.as_ref())
632        );
633
634        let result = my_collation_example
635            .filter(value.eq("oof"))
636            .select(value)
637            .load::<String>(connection);
638        assert_eq!(Ok(&[][..]), result.as_ref().map(|vec| vec.as_ref()));
639    }
640
641    #[diesel_test_helper::test]
642    fn aggregate_function_works_with_aligned_data() {
643        #[derive(Debug, Default)]
644        #[repr(align(64))]
645        struct OverAligned;
646
647        impl SqliteAggregateFunction<i32> for OverAligned {
648            type Output = i64;
649
650            fn step(&mut self, _value: i32) {
651                let need = core::mem::align_of::<Self>();
652                let got = core::mem::align_of_val(self);
653                assert_eq!(need, got);
654            }
655
656            fn finalize(_agg: Option<Self>) -> i64 {
657                0
658            }
659        }
660        #[declare_sql_function]
661        extern "SQL" {
662            #[aggregate]
663            fn over_aligned_sum(x: Integer) -> diesel::sql_types::BigInt;
664        }
665
666        let mut conn = SqliteConnection::establish(":memory:").unwrap();
667        over_aligned_sum_utils::register_impl::<OverAligned, _>(&mut conn).unwrap();
668
669        diesel::select(over_aligned_sum(1))
670            .execute(&mut conn)
671            .unwrap();
672    }
673
674    #[diesel_test_helper::test]
675    fn sum_twice() {
676        #[derive(Default)]
677        struct Sum(i32);
678
679        impl SqliteAggregateFunction<i32> for Sum {
680            type Output = i32;
681
682            fn step(&mut self, value: i32) {
683                self.0 += value;
684            }
685
686            fn finalize(agg: Option<Self>) -> i32 {
687                agg.map(|s| s.0).unwrap_or_default()
688            }
689        }
690
691        #[declare_sql_function]
692        extern "SQL" {
693            #[aggregate]
694            fn my_sum(x: Integer) -> Integer;
695        }
696
697        let mut conn = SqliteConnection::establish(":memory:").unwrap();
698        my_sum_utils::register_impl::<Sum, _>(&mut conn).unwrap();
699
700        conn.batch_execute(
701            "
702            CREATE TABLE test(key1 INTEGER, key2 INTEGER);
703            INSERT INTO test(key1, key2) VALUES (1, 2), (2, 4), (3, 6);
704",
705        )
706        .unwrap();
707
708        table! {
709            test (key1, key2) {
710                key1 -> Integer,
711                key2 -> Integer,
712            }
713        }
714
715        let (first_res, second_res) = test::table
716            .select((my_sum(test::key1), my_sum(test::key2)))
717            .get_result::<(i32, i32)>(&mut conn)
718            .unwrap();
719
720        assert_eq!(first_res, 6);
721        assert_eq!(second_res, 12);
722
723        conn.batch_execute("DELETE FROM test").unwrap();
724        let (first_res, second_res) = test::table
725            .select((my_sum(test::key1), my_sum(test::key2)))
726            .get_result::<(i32, i32)>(&mut conn)
727            .unwrap();
728
729        assert_eq!(first_res, 0);
730        assert_eq!(second_res, 0);
731    }
732
733    // ---- DIRECTONLY / INNOCUOUS function behavior tests ----
734
735    #[declare_sql_function]
736    extern "SQL" {
737        fn directonly_fn() -> Integer;
738        fn innocuous_fn() -> Integer;
739    }
740
741    #[diesel_test_helper::test]
742    fn directonly_function_blocked_from_view() {
743        let conn = &mut connection();
744
745        // Register a DIRECTONLY function
746        directonly_fn_utils::register_impl_with_behavior(
747            conn,
748            SqliteFunctionBehavior::DIRECTONLY,
749            || 42,
750        )
751        .unwrap();
752
753        // Direct call works
754        let result = crate::select(directonly_fn()).get_result::<i32>(conn);
755        assert_eq!(Ok(42), result);
756
757        // Create a view that calls the function
758        crate::sql_query("CREATE VIEW test_view AS SELECT directonly_fn() AS val")
759            .execute(conn)
760            .unwrap();
761
762        // Disable trusted schema so DIRECTONLY is enforced from schema objects
763        conn.set_trusted_schema(false).unwrap();
764
765        // Querying the view should fail because the function is DIRECTONLY
766        let result = crate::sql_query("SELECT val FROM test_view").execute(conn);
767        assert!(result.is_err());
768    }
769
770    #[diesel_test_helper::test]
771    fn innocuous_function_allowed_from_view_with_untrusted_schema() {
772        let conn = &mut connection();
773
774        // Register an INNOCUOUS function
775        innocuous_fn_utils::register_impl_with_behavior(
776            conn,
777            SqliteFunctionBehavior::DETERMINISTIC | SqliteFunctionBehavior::INNOCUOUS,
778            || 99,
779        )
780        .unwrap();
781
782        // Create a view that calls the function
783        crate::sql_query("CREATE VIEW innocuous_view AS SELECT innocuous_fn() AS val")
784            .execute(conn)
785            .unwrap();
786
787        // Disable trusted schema
788        conn.set_trusted_schema(false).unwrap();
789
790        // Querying the view should succeed because the function is INNOCUOUS
791        let result = crate::sql_query("SELECT val FROM innocuous_view").execute(conn);
792        assert!(result.is_ok());
793    }
794}