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#[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)] 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 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 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#[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 #[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 directonly_fn_utils::register_impl_with_behavior(
747 conn,
748 SqliteFunctionBehavior::DIRECTONLY,
749 || 42,
750 )
751 .unwrap();
752
753 let result = crate::select(directonly_fn()).get_result::<i32>(conn);
755 assert_eq!(Ok(42), result);
756
757 crate::sql_query("CREATE VIEW test_view AS SELECT directonly_fn() AS val")
759 .execute(conn)
760 .unwrap();
761
762 conn.set_trusted_schema(false).unwrap();
764
765 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 innocuous_fn_utils::register_impl_with_behavior(
776 conn,
777 SqliteFunctionBehavior::DETERMINISTIC | SqliteFunctionBehavior::INNOCUOUS,
778 || 99,
779 )
780 .unwrap();
781
782 crate::sql_query("CREATE VIEW innocuous_view AS SELECT innocuous_fn() AS val")
784 .execute(conn)
785 .unwrap();
786
787 conn.set_trusted_schema(false).unwrap();
789
790 let result = crate::sql_query("SELECT val FROM innocuous_view").execute(conn);
792 assert!(result.is_ok());
793 }
794}