1use std::any::Any;
56use std::marker::PhantomData;
57use std::ops::Deref;
58use std::os::raw::{c_int, c_void};
59use std::panic::{catch_unwind, RefUnwindSafe, UnwindSafe};
60use std::ptr;
61use std::slice;
62use std::sync::Arc;
63
64use crate::ffi;
65use crate::ffi::sqlite3_context;
66use crate::ffi::sqlite3_value;
67
68use crate::context::set_result;
69use crate::types::{FromSql, FromSqlError, ToSql, ValueRef};
70
71use crate::{str_to_cstring, Connection, Error, InnerConnection, Result};
72
73unsafe fn report_error(ctx: *mut sqlite3_context, err: &Error) {
74 #[cfg(feature = "modern_sqlite")]
79 fn constraint_error_code() -> i32 {
80 ffi::SQLITE_CONSTRAINT_FUNCTION
81 }
82 #[cfg(not(feature = "modern_sqlite"))]
83 fn constraint_error_code() -> i32 {
84 ffi::SQLITE_CONSTRAINT
85 }
86
87 match *err {
88 Error::SqliteFailure(ref err, ref s) => {
89 ffi::sqlite3_result_error_code(ctx, err.extended_code);
90 if let Some(Ok(cstr)) = s.as_ref().map(|s| str_to_cstring(s)) {
91 ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1);
92 }
93 }
94 _ => {
95 ffi::sqlite3_result_error_code(ctx, constraint_error_code());
96 if let Ok(cstr) = str_to_cstring(&err.to_string()) {
97 ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1);
98 }
99 }
100 }
101}
102
103unsafe extern "C" fn free_boxed_value<T>(p: *mut c_void) {
104 drop(Box::from_raw(p as *mut T));
105}
106
107pub struct Context<'a> {
110 ctx: *mut sqlite3_context,
111 args: &'a [*mut sqlite3_value],
112}
113
114impl Context<'_> {
115 #[inline]
117 pub fn len(&self) -> usize {
118 self.args.len()
119 }
120
121 #[inline]
123 pub fn is_empty(&self) -> bool {
124 self.args.is_empty()
125 }
126
127 pub fn get<T: FromSql>(&self, idx: usize) -> Result<T> {
136 let arg = self.args[idx];
137 let value = unsafe { ValueRef::from_value(arg) };
138 FromSql::column_result(value).map_err(|err| match err {
139 FromSqlError::InvalidType => {
140 Error::InvalidFunctionParameterType(idx, value.data_type())
141 }
142 FromSqlError::OutOfRange(i) => Error::IntegralValueOutOfRange(idx, i),
143 FromSqlError::Other(err) => {
144 Error::FromSqlConversionFailure(idx, value.data_type(), err)
145 }
146 #[cfg(feature = "i128_blob")]
147 FromSqlError::InvalidI128Size(_) => {
148 Error::FromSqlConversionFailure(idx, value.data_type(), Box::new(err))
149 }
150 #[cfg(feature = "uuid")]
151 FromSqlError::InvalidUuidSize(_) => {
152 Error::FromSqlConversionFailure(idx, value.data_type(), Box::new(err))
153 }
154 })
155 }
156
157 #[inline]
163 pub fn get_raw(&self, idx: usize) -> ValueRef<'_> {
164 let arg = self.args[idx];
165 unsafe { ValueRef::from_value(arg) }
166 }
167
168 pub fn get_or_create_aux<T, E, F>(&self, arg: c_int, func: F) -> Result<Arc<T>>
175 where
176 T: Send + Sync + 'static,
177 E: Into<Box<dyn std::error::Error + Send + Sync + 'static>>,
178 F: FnOnce(ValueRef<'_>) -> Result<T, E>,
179 {
180 if let Some(v) = self.get_aux(arg)? {
181 Ok(v)
182 } else {
183 let vr = self.get_raw(arg as usize);
184 self.set_aux(
185 arg,
186 func(vr).map_err(|e| Error::UserFunctionError(e.into()))?,
187 )
188 }
189 }
190
191 pub fn set_aux<T: Send + Sync + 'static>(&self, arg: c_int, value: T) -> Result<Arc<T>> {
195 let orig: Arc<T> = Arc::new(value);
196 let inner: AuxInner = orig.clone();
197 let outer = Box::new(inner);
198 let raw: *mut AuxInner = Box::into_raw(outer);
199 unsafe {
200 ffi::sqlite3_set_auxdata(
201 self.ctx,
202 arg,
203 raw as *mut _,
204 Some(free_boxed_value::<AuxInner>),
205 )
206 };
207 Ok(orig)
208 }
209
210 pub fn get_aux<T: Send + Sync + 'static>(&self, arg: c_int) -> Result<Option<Arc<T>>> {
215 let p = unsafe { ffi::sqlite3_get_auxdata(self.ctx, arg) as *const AuxInner };
216 if p.is_null() {
217 Ok(None)
218 } else {
219 let v: AuxInner = AuxInner::clone(unsafe { &*p });
220 v.downcast::<T>()
221 .map(Some)
222 .map_err(|_| Error::GetAuxWrongType)
223 }
224 }
225
226 pub unsafe fn get_connection(&self) -> Result<ConnectionRef<'_>> {
233 let handle = ffi::sqlite3_context_db_handle(self.ctx);
234 Ok(ConnectionRef {
235 conn: Connection::from_handle(handle)?,
236 phantom: PhantomData,
237 })
238 }
239}
240
241pub struct ConnectionRef<'ctx> {
243 conn: Connection,
246 phantom: PhantomData<&'ctx Context<'ctx>>,
247}
248
249impl Deref for ConnectionRef<'_> {
250 type Target = Connection;
251
252 #[inline]
253 fn deref(&self) -> &Connection {
254 &self.conn
255 }
256}
257
258type AuxInner = Arc<dyn Any + Send + Sync + 'static>;
259
260pub trait Aggregate<A, T>
266where
267 A: RefUnwindSafe + UnwindSafe,
268 T: ToSql,
269{
270 fn init(&self, _: &mut Context<'_>) -> Result<A>;
274
275 fn step(&self, _: &mut Context<'_>, _: &mut A) -> Result<()>;
278
279 fn finalize(&self, _: &mut Context<'_>, _: Option<A>) -> Result<T>;
287}
288
289#[cfg(feature = "window")]
292pub trait WindowAggregate<A, T>: Aggregate<A, T>
293where
294 A: RefUnwindSafe + UnwindSafe,
295 T: ToSql,
296{
297 fn value(&self, _: Option<&A>) -> Result<T>;
300
301 fn inverse(&self, _: &mut Context<'_>, _: &mut A) -> Result<()>;
303}
304
305bitflags::bitflags! {
306 #[repr(C)]
310 pub struct FunctionFlags: ::std::os::raw::c_int {
311 const SQLITE_UTF8 = ffi::SQLITE_UTF8;
313 const SQLITE_UTF16LE = ffi::SQLITE_UTF16LE;
315 const SQLITE_UTF16BE = ffi::SQLITE_UTF16BE;
317 const SQLITE_UTF16 = ffi::SQLITE_UTF16;
319 const SQLITE_DETERMINISTIC = ffi::SQLITE_DETERMINISTIC;
321 const SQLITE_DIRECTONLY = 0x0000_0008_0000; const SQLITE_SUBTYPE = 0x0000_0010_0000; const SQLITE_INNOCUOUS = 0x0000_0020_0000; }
328}
329
330impl Default for FunctionFlags {
331 #[inline]
332 fn default() -> FunctionFlags {
333 FunctionFlags::SQLITE_UTF8
334 }
335}
336
337impl Connection {
338 #[inline]
375 pub fn create_scalar_function<'c, F, T>(
376 &'c self,
377 fn_name: &str,
378 n_arg: c_int,
379 flags: FunctionFlags,
380 x_func: F,
381 ) -> Result<()>
382 where
383 F: FnMut(&Context<'_>) -> Result<T> + Send + UnwindSafe + 'c,
384 T: ToSql,
385 {
386 self.db
387 .borrow_mut()
388 .create_scalar_function(fn_name, n_arg, flags, x_func)
389 }
390
391 #[inline]
398 pub fn create_aggregate_function<A, D, T>(
399 &self,
400 fn_name: &str,
401 n_arg: c_int,
402 flags: FunctionFlags,
403 aggr: D,
404 ) -> Result<()>
405 where
406 A: RefUnwindSafe + UnwindSafe,
407 D: Aggregate<A, T>,
408 T: ToSql,
409 {
410 self.db
411 .borrow_mut()
412 .create_aggregate_function(fn_name, n_arg, flags, aggr)
413 }
414
415 #[cfg(feature = "window")]
421 #[inline]
422 pub fn create_window_function<A, W, T>(
423 &self,
424 fn_name: &str,
425 n_arg: c_int,
426 flags: FunctionFlags,
427 aggr: W,
428 ) -> Result<()>
429 where
430 A: RefUnwindSafe + UnwindSafe,
431 W: WindowAggregate<A, T>,
432 T: ToSql,
433 {
434 self.db
435 .borrow_mut()
436 .create_window_function(fn_name, n_arg, flags, aggr)
437 }
438
439 #[inline]
449 pub fn remove_function(&self, fn_name: &str, n_arg: c_int) -> Result<()> {
450 self.db.borrow_mut().remove_function(fn_name, n_arg)
451 }
452}
453
454impl InnerConnection {
455 fn create_scalar_function<'c, F, T>(
456 &'c mut self,
457 fn_name: &str,
458 n_arg: c_int,
459 flags: FunctionFlags,
460 x_func: F,
461 ) -> Result<()>
462 where
463 F: FnMut(&Context<'_>) -> Result<T> + Send + UnwindSafe + 'c,
464 T: ToSql,
465 {
466 unsafe extern "C" fn call_boxed_closure<F, T>(
467 ctx: *mut sqlite3_context,
468 argc: c_int,
469 argv: *mut *mut sqlite3_value,
470 ) where
471 F: FnMut(&Context<'_>) -> Result<T>,
472 T: ToSql,
473 {
474 let r = catch_unwind(|| {
475 let boxed_f: *mut F = ffi::sqlite3_user_data(ctx) as *mut F;
476 assert!(!boxed_f.is_null(), "Internal error - null function pointer");
477 let ctx = Context {
478 ctx,
479 args: slice::from_raw_parts(argv, argc as usize),
480 };
481 (*boxed_f)(&ctx)
482 });
483 let t = match r {
484 Err(_) => {
485 report_error(ctx, &Error::UnwindingPanic);
486 return;
487 }
488 Ok(r) => r,
489 };
490 let t = t.as_ref().map(|t| ToSql::to_sql(t));
491
492 match t {
493 Ok(Ok(ref value)) => set_result(ctx, value),
494 Ok(Err(err)) => report_error(ctx, &err),
495 Err(err) => report_error(ctx, err),
496 }
497 }
498
499 let boxed_f: *mut F = Box::into_raw(Box::new(x_func));
500 let c_name = str_to_cstring(fn_name)?;
501 let r = unsafe {
502 ffi::sqlite3_create_function_v2(
503 self.db(),
504 c_name.as_ptr(),
505 n_arg,
506 flags.bits(),
507 boxed_f as *mut c_void,
508 Some(call_boxed_closure::<F, T>),
509 None,
510 None,
511 Some(free_boxed_value::<F>),
512 )
513 };
514 self.decode_result(r)
515 }
516
517 fn create_aggregate_function<A, D, T>(
518 &mut self,
519 fn_name: &str,
520 n_arg: c_int,
521 flags: FunctionFlags,
522 aggr: D,
523 ) -> Result<()>
524 where
525 A: RefUnwindSafe + UnwindSafe,
526 D: Aggregate<A, T>,
527 T: ToSql,
528 {
529 let boxed_aggr: *mut D = Box::into_raw(Box::new(aggr));
530 let c_name = str_to_cstring(fn_name)?;
531 let r = unsafe {
532 ffi::sqlite3_create_function_v2(
533 self.db(),
534 c_name.as_ptr(),
535 n_arg,
536 flags.bits(),
537 boxed_aggr as *mut c_void,
538 None,
539 Some(call_boxed_step::<A, D, T>),
540 Some(call_boxed_final::<A, D, T>),
541 Some(free_boxed_value::<D>),
542 )
543 };
544 self.decode_result(r)
545 }
546
547 #[cfg(feature = "window")]
548 fn create_window_function<A, W, T>(
549 &mut self,
550 fn_name: &str,
551 n_arg: c_int,
552 flags: FunctionFlags,
553 aggr: W,
554 ) -> Result<()>
555 where
556 A: RefUnwindSafe + UnwindSafe,
557 W: WindowAggregate<A, T>,
558 T: ToSql,
559 {
560 let boxed_aggr: *mut W = Box::into_raw(Box::new(aggr));
561 let c_name = str_to_cstring(fn_name)?;
562 let r = unsafe {
563 ffi::sqlite3_create_window_function(
564 self.db(),
565 c_name.as_ptr(),
566 n_arg,
567 flags.bits(),
568 boxed_aggr as *mut c_void,
569 Some(call_boxed_step::<A, W, T>),
570 Some(call_boxed_final::<A, W, T>),
571 Some(call_boxed_value::<A, W, T>),
572 Some(call_boxed_inverse::<A, W, T>),
573 Some(free_boxed_value::<W>),
574 )
575 };
576 self.decode_result(r)
577 }
578
579 fn remove_function(&mut self, fn_name: &str, n_arg: c_int) -> Result<()> {
580 let c_name = str_to_cstring(fn_name)?;
581 let r = unsafe {
582 ffi::sqlite3_create_function_v2(
583 self.db(),
584 c_name.as_ptr(),
585 n_arg,
586 ffi::SQLITE_UTF8,
587 ptr::null_mut(),
588 None,
589 None,
590 None,
591 None,
592 )
593 };
594 self.decode_result(r)
595 }
596}
597
598unsafe fn aggregate_context<A>(ctx: *mut sqlite3_context, bytes: usize) -> Option<*mut *mut A> {
599 let pac = ffi::sqlite3_aggregate_context(ctx, bytes as c_int) as *mut *mut A;
600 if pac.is_null() {
601 return None;
602 }
603 Some(pac)
604}
605
606unsafe extern "C" fn call_boxed_step<A, D, T>(
607 ctx: *mut sqlite3_context,
608 argc: c_int,
609 argv: *mut *mut sqlite3_value,
610) where
611 A: RefUnwindSafe + UnwindSafe,
612 D: Aggregate<A, T>,
613 T: ToSql,
614{
615 let pac = match aggregate_context(ctx, ::std::mem::size_of::<*mut A>()) {
616 Some(pac) => pac,
617 None => {
618 ffi::sqlite3_result_error_nomem(ctx);
619 return;
620 }
621 };
622
623 let r = catch_unwind(|| {
624 let boxed_aggr: *mut D = ffi::sqlite3_user_data(ctx) as *mut D;
625 assert!(
626 !boxed_aggr.is_null(),
627 "Internal error - null aggregate pointer"
628 );
629 let mut ctx = Context {
630 ctx,
631 args: slice::from_raw_parts(argv, argc as usize),
632 };
633
634 if (*pac as *mut A).is_null() {
635 *pac = Box::into_raw(Box::new((*boxed_aggr).init(&mut ctx)?));
636 }
637
638 (*boxed_aggr).step(&mut ctx, &mut **pac)
639 });
640 let r = match r {
641 Err(_) => {
642 report_error(ctx, &Error::UnwindingPanic);
643 return;
644 }
645 Ok(r) => r,
646 };
647 match r {
648 Ok(_) => {}
649 Err(err) => report_error(ctx, &err),
650 };
651}
652
653#[cfg(feature = "window")]
654unsafe extern "C" fn call_boxed_inverse<A, W, T>(
655 ctx: *mut sqlite3_context,
656 argc: c_int,
657 argv: *mut *mut sqlite3_value,
658) where
659 A: RefUnwindSafe + UnwindSafe,
660 W: WindowAggregate<A, T>,
661 T: ToSql,
662{
663 let pac = match aggregate_context(ctx, ::std::mem::size_of::<*mut A>()) {
664 Some(pac) => pac,
665 None => {
666 ffi::sqlite3_result_error_nomem(ctx);
667 return;
668 }
669 };
670
671 let r = catch_unwind(|| {
672 let boxed_aggr: *mut W = ffi::sqlite3_user_data(ctx) as *mut W;
673 assert!(
674 !boxed_aggr.is_null(),
675 "Internal error - null aggregate pointer"
676 );
677 let mut ctx = Context {
678 ctx,
679 args: slice::from_raw_parts(argv, argc as usize),
680 };
681 (*boxed_aggr).inverse(&mut ctx, &mut **pac)
682 });
683 let r = match r {
684 Err(_) => {
685 report_error(ctx, &Error::UnwindingPanic);
686 return;
687 }
688 Ok(r) => r,
689 };
690 match r {
691 Ok(_) => {}
692 Err(err) => report_error(ctx, &err),
693 };
694}
695
696unsafe extern "C" fn call_boxed_final<A, D, T>(ctx: *mut sqlite3_context)
697where
698 A: RefUnwindSafe + UnwindSafe,
699 D: Aggregate<A, T>,
700 T: ToSql,
701{
702 let a: Option<A> = match aggregate_context(ctx, 0) {
705 Some(pac) => {
706 if (*pac as *mut A).is_null() {
707 None
708 } else {
709 let a = Box::from_raw(*pac);
710 Some(*a)
711 }
712 }
713 None => None,
714 };
715
716 let r = catch_unwind(|| {
717 let boxed_aggr: *mut D = ffi::sqlite3_user_data(ctx) as *mut D;
718 assert!(
719 !boxed_aggr.is_null(),
720 "Internal error - null aggregate pointer"
721 );
722 let mut ctx = Context { ctx, args: &mut [] };
723 (*boxed_aggr).finalize(&mut ctx, a)
724 });
725 let t = match r {
726 Err(_) => {
727 report_error(ctx, &Error::UnwindingPanic);
728 return;
729 }
730 Ok(r) => r,
731 };
732 let t = t.as_ref().map(|t| ToSql::to_sql(t));
733 match t {
734 Ok(Ok(ref value)) => set_result(ctx, value),
735 Ok(Err(err)) => report_error(ctx, &err),
736 Err(err) => report_error(ctx, err),
737 }
738}
739
740#[cfg(feature = "window")]
741unsafe extern "C" fn call_boxed_value<A, W, T>(ctx: *mut sqlite3_context)
742where
743 A: RefUnwindSafe + UnwindSafe,
744 W: WindowAggregate<A, T>,
745 T: ToSql,
746{
747 let a: Option<&A> = match aggregate_context(ctx, 0) {
750 Some(pac) => {
751 if (*pac as *mut A).is_null() {
752 None
753 } else {
754 let a = &**pac;
755 Some(a)
756 }
757 }
758 None => None,
759 };
760
761 let r = catch_unwind(|| {
762 let boxed_aggr: *mut W = ffi::sqlite3_user_data(ctx) as *mut W;
763 assert!(
764 !boxed_aggr.is_null(),
765 "Internal error - null aggregate pointer"
766 );
767 (*boxed_aggr).value(a)
768 });
769 let t = match r {
770 Err(_) => {
771 report_error(ctx, &Error::UnwindingPanic);
772 return;
773 }
774 Ok(r) => r,
775 };
776 let t = t.as_ref().map(|t| ToSql::to_sql(t));
777 match t {
778 Ok(Ok(ref value)) => set_result(ctx, value),
779 Ok(Err(err)) => report_error(ctx, &err),
780 Err(err) => report_error(ctx, err),
781 }
782}
783
784#[cfg(test)]
785mod test {
786 use regex::Regex;
787 use std::f64::EPSILON;
788 use std::os::raw::c_double;
789
790 #[cfg(feature = "window")]
791 use crate::functions::WindowAggregate;
792 use crate::functions::{Aggregate, Context, FunctionFlags};
793 use crate::{Connection, Error, Result};
794
795 fn half(ctx: &Context<'_>) -> Result<c_double> {
796 assert_eq!(ctx.len(), 1, "called with unexpected number of arguments");
797 let value = ctx.get::<c_double>(0)?;
798 Ok(value / 2f64)
799 }
800
801 #[test]
802 fn test_function_half() -> Result<()> {
803 let db = Connection::open_in_memory()?;
804 db.create_scalar_function(
805 "half",
806 1,
807 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
808 half,
809 )?;
810 let result: Result<f64> = db.query_row("SELECT half(6)", [], |r| r.get(0));
811
812 assert!((3f64 - result?).abs() < EPSILON);
813 Ok(())
814 }
815
816 #[test]
817 fn test_remove_function() -> Result<()> {
818 let db = Connection::open_in_memory()?;
819 db.create_scalar_function(
820 "half",
821 1,
822 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
823 half,
824 )?;
825 let result: Result<f64> = db.query_row("SELECT half(6)", [], |r| r.get(0));
826 assert!((3f64 - result?).abs() < EPSILON);
827
828 db.remove_function("half", 1)?;
829 let result: Result<f64> = db.query_row("SELECT half(6)", [], |r| r.get(0));
830 assert!(result.is_err());
831 Ok(())
832 }
833
834 fn regexp_with_auxilliary(ctx: &Context<'_>) -> Result<bool> {
838 assert_eq!(ctx.len(), 2, "called with unexpected number of arguments");
839 type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
840 let regexp: std::sync::Arc<Regex> = ctx
841 .get_or_create_aux(0, |vr| -> Result<_, BoxError> {
842 Ok(Regex::new(vr.as_str()?)?)
843 })?;
844
845 let is_match = {
846 let text = ctx
847 .get_raw(1)
848 .as_str()
849 .map_err(|e| Error::UserFunctionError(e.into()))?;
850
851 regexp.is_match(text)
852 };
853
854 Ok(is_match)
855 }
856
857 #[test]
858 fn test_function_regexp_with_auxilliary() -> Result<()> {
859 let db = Connection::open_in_memory()?;
860 db.execute_batch(
861 "BEGIN;
862 CREATE TABLE foo (x string);
863 INSERT INTO foo VALUES ('lisa');
864 INSERT INTO foo VALUES ('lXsi');
865 INSERT INTO foo VALUES ('lisX');
866 END;",
867 )?;
868 db.create_scalar_function(
869 "regexp",
870 2,
871 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
872 regexp_with_auxilliary,
873 )?;
874
875 let result: Result<bool> =
876 db.query_row("SELECT regexp('l.s[aeiouy]', 'lisa')", [], |r| r.get(0));
877
878 assert_eq!(true, result?);
879
880 let result: Result<i64> = db.query_row(
881 "SELECT COUNT(*) FROM foo WHERE regexp('l.s[aeiouy]', x) == 1",
882 [],
883 |r| r.get(0),
884 );
885
886 assert_eq!(2, result?);
887 Ok(())
888 }
889
890 #[test]
891 fn test_varargs_function() -> Result<()> {
892 let db = Connection::open_in_memory()?;
893 db.create_scalar_function(
894 "my_concat",
895 -1,
896 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
897 |ctx| {
898 let mut ret = String::new();
899
900 for idx in 0..ctx.len() {
901 let s = ctx.get::<String>(idx)?;
902 ret.push_str(&s);
903 }
904
905 Ok(ret)
906 },
907 )?;
908
909 for &(expected, query) in &[
910 ("", "SELECT my_concat()"),
911 ("onetwo", "SELECT my_concat('one', 'two')"),
912 ("abc", "SELECT my_concat('a', 'b', 'c')"),
913 ] {
914 let result: String = db.query_row(query, [], |r| r.get(0))?;
915 assert_eq!(expected, result);
916 }
917 Ok(())
918 }
919
920 #[test]
921 fn test_get_aux_type_checking() -> Result<()> {
922 let db = Connection::open_in_memory()?;
923 db.create_scalar_function("example", 2, FunctionFlags::default(), |ctx| {
924 if !ctx.get::<bool>(1)? {
925 ctx.set_aux::<i64>(0, 100)?;
926 } else {
927 assert_eq!(ctx.get_aux::<String>(0), Err(Error::GetAuxWrongType));
928 assert_eq!(*ctx.get_aux::<i64>(0)?.unwrap(), 100);
929 }
930 Ok(true)
931 })?;
932
933 let res: bool = db.query_row(
934 "SELECT example(0, i) FROM (SELECT 0 as i UNION SELECT 1)",
935 [],
936 |r| r.get(0),
937 )?;
938 assert!(res);
940 Ok(())
941 }
942
943 struct Sum;
944 struct Count;
945
946 impl Aggregate<i64, Option<i64>> for Sum {
947 fn init(&self, _: &mut Context<'_>) -> Result<i64> {
948 Ok(0)
949 }
950
951 fn step(&self, ctx: &mut Context<'_>, sum: &mut i64) -> Result<()> {
952 *sum += ctx.get::<i64>(0)?;
953 Ok(())
954 }
955
956 fn finalize(&self, _: &mut Context<'_>, sum: Option<i64>) -> Result<Option<i64>> {
957 Ok(sum)
958 }
959 }
960
961 impl Aggregate<i64, i64> for Count {
962 fn init(&self, _: &mut Context<'_>) -> Result<i64> {
963 Ok(0)
964 }
965
966 fn step(&self, _ctx: &mut Context<'_>, sum: &mut i64) -> Result<()> {
967 *sum += 1;
968 Ok(())
969 }
970
971 fn finalize(&self, _: &mut Context<'_>, sum: Option<i64>) -> Result<i64> {
972 Ok(sum.unwrap_or(0))
973 }
974 }
975
976 #[test]
977 fn test_sum() -> Result<()> {
978 let db = Connection::open_in_memory()?;
979 db.create_aggregate_function(
980 "my_sum",
981 1,
982 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
983 Sum,
984 )?;
985
986 let no_result = "SELECT my_sum(i) FROM (SELECT 2 AS i WHERE 1 <> 1)";
988 let result: Option<i64> = db.query_row(no_result, [], |r| r.get(0))?;
989 assert!(result.is_none());
990
991 let single_sum = "SELECT my_sum(i) FROM (SELECT 2 AS i UNION ALL SELECT 2)";
992 let result: i64 = db.query_row(single_sum, [], |r| r.get(0))?;
993 assert_eq!(4, result);
994
995 let dual_sum = "SELECT my_sum(i), my_sum(j) FROM (SELECT 2 AS i, 1 AS j UNION ALL SELECT \
996 2, 1)";
997 let result: (i64, i64) = db.query_row(dual_sum, [], |r| Ok((r.get(0)?, r.get(1)?)))?;
998 assert_eq!((4, 2), result);
999 Ok(())
1000 }
1001
1002 #[test]
1003 fn test_count() -> Result<()> {
1004 let db = Connection::open_in_memory()?;
1005 db.create_aggregate_function(
1006 "my_count",
1007 -1,
1008 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
1009 Count,
1010 )?;
1011
1012 let no_result = "SELECT my_count(i) FROM (SELECT 2 AS i WHERE 1 <> 1)";
1014 let result: i64 = db.query_row(no_result, [], |r| r.get(0))?;
1015 assert_eq!(result, 0);
1016
1017 let single_sum = "SELECT my_count(i) FROM (SELECT 2 AS i UNION ALL SELECT 2)";
1018 let result: i64 = db.query_row(single_sum, [], |r| r.get(0))?;
1019 assert_eq!(2, result);
1020 Ok(())
1021 }
1022
1023 #[cfg(feature = "window")]
1024 impl WindowAggregate<i64, Option<i64>> for Sum {
1025 fn inverse(&self, ctx: &mut Context<'_>, sum: &mut i64) -> Result<()> {
1026 *sum -= ctx.get::<i64>(0)?;
1027 Ok(())
1028 }
1029
1030 fn value(&self, sum: Option<&i64>) -> Result<Option<i64>> {
1031 Ok(sum.copied())
1032 }
1033 }
1034
1035 #[test]
1036 #[cfg(feature = "window")]
1037 fn test_window() -> Result<()> {
1038 use fallible_iterator::FallibleIterator;
1039
1040 let db = Connection::open_in_memory()?;
1041 db.create_window_function(
1042 "sumint",
1043 1,
1044 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
1045 Sum,
1046 )?;
1047 db.execute_batch(
1048 "CREATE TABLE t3(x, y);
1049 INSERT INTO t3 VALUES('a', 4),
1050 ('b', 5),
1051 ('c', 3),
1052 ('d', 8),
1053 ('e', 1);",
1054 )?;
1055
1056 let mut stmt = db.prepare(
1057 "SELECT x, sumint(y) OVER (
1058 ORDER BY x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING
1059 ) AS sum_y
1060 FROM t3 ORDER BY x;",
1061 )?;
1062
1063 let results: Vec<(String, i64)> = stmt
1064 .query([])?
1065 .map(|row| Ok((row.get("x")?, row.get("sum_y")?)))
1066 .collect()?;
1067 let expected = vec![
1068 ("a".to_owned(), 9),
1069 ("b".to_owned(), 12),
1070 ("c".to_owned(), 16),
1071 ("d".to_owned(), 12),
1072 ("e".to_owned(), 9),
1073 ];
1074 assert_eq!(expected, results);
1075 Ok(())
1076 }
1077}