Skip to main content

rusqlite/
functions.rs

1//! `feature = "functions"` Create or redefine SQL functions.
2//!
3//! # Example
4//!
5//! Adding a `regexp` function to a connection in which compiled regular
6//! expressions are cached in a `HashMap`. For an alternative implementation
7//! that uses SQLite's [Function Auxilliary Data](https://www.sqlite.org/c3ref/get_auxdata.html) interface
8//! to avoid recompiling regular expressions, see the unit tests for this
9//! module.
10//!
11//! ```rust
12//! use regex::Regex;
13//! use rusqlite::functions::FunctionFlags;
14//! use rusqlite::{Connection, Error, Result};
15//! use std::sync::Arc;
16//! type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
17//!
18//! fn add_regexp_function(db: &Connection) -> Result<()> {
19//!     db.create_scalar_function(
20//!         "regexp",
21//!         2,
22//!         FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
23//!         move |ctx| {
24//!             assert_eq!(ctx.len(), 2, "called with unexpected number of arguments");
25//!             let regexp: Arc<Regex> = ctx.get_or_create_aux(0, |vr| -> Result<_, BoxError> {
26//!                 Ok(Regex::new(vr.as_str()?)?)
27//!             })?;
28//!             let is_match = {
29//!                 let text = ctx
30//!                     .get_raw(1)
31//!                     .as_str()
32//!                     .map_err(|e| Error::UserFunctionError(e.into()))?;
33//!
34//!                 regexp.is_match(text)
35//!             };
36//!
37//!             Ok(is_match)
38//!         },
39//!     )
40//! }
41//!
42//! fn main() -> Result<()> {
43//!     let db = Connection::open_in_memory()?;
44//!     add_regexp_function(&db)?;
45//!
46//!     let is_match: bool =
47//!         db.query_row("SELECT regexp('[aeiou]*', 'aaaaeeeiii')", [], |row| {
48//!             row.get(0)
49//!         })?;
50//!
51//!     assert!(is_match);
52//!     Ok(())
53//! }
54//! ```
55use 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    // Extended constraint error codes were added in SQLite 3.7.16. We don't have
75    // an explicit feature check for that, and this doesn't really warrant one.
76    // We'll use the extended code if we're on the bundled version (since it's
77    // at least 3.17.0) and the normal constraint error code if not.
78    #[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
107/// `feature = "functions"` Context is a wrapper for the SQLite function
108/// evaluation context.
109pub struct Context<'a> {
110    ctx: *mut sqlite3_context,
111    args: &'a [*mut sqlite3_value],
112}
113
114impl Context<'_> {
115    /// Returns the number of arguments to the function.
116    #[inline]
117    pub fn len(&self) -> usize {
118        self.args.len()
119    }
120
121    /// Returns `true` when there is no argument.
122    #[inline]
123    pub fn is_empty(&self) -> bool {
124        self.args.is_empty()
125    }
126
127    /// Returns the `idx`th argument as a `T`.
128    ///
129    /// # Failure
130    ///
131    /// Will panic if `idx` is greater than or equal to [`self.len()`](Context::len).
132    ///
133    /// Will return Err if the underlying SQLite type cannot be converted to a
134    /// `T`.
135    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    /// Returns the `idx`th argument as a `ValueRef`.
158    ///
159    /// # Failure
160    ///
161    /// Will panic if `idx` is greater than or equal to [`self.len()`](Context::len).
162    #[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    /// Fetch or insert the auxilliary data associated with a particular
169    /// parameter. This is intended to be an easier-to-use way of fetching it
170    /// compared to calling [`get_aux`](Context::get_aux) and [`set_aux`](Context::set_aux) separately.
171    ///
172    /// See `https://www.sqlite.org/c3ref/get_auxdata.html` for a discussion of
173    /// this feature, or the unit tests of this module for an example.
174    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    /// Sets the auxilliary data associated with a particular parameter. See
192    /// `https://www.sqlite.org/c3ref/get_auxdata.html` for a discussion of
193    /// this feature, or the unit tests of this module for an example.
194    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    /// Gets the auxilliary data that was associated with a given parameter via
211    /// [`set_aux`](Context::set_aux). Returns `Ok(None)` if no data has been associated, and
212    /// Ok(Some(v)) if it has. Returns an error if the requested type does not
213    /// match.
214    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    /// Get the db connection handle via [sqlite3_context_db_handle](https://www.sqlite.org/c3ref/context_db_handle.html)
227    ///
228    /// # Safety
229    ///
230    /// This function is marked unsafe because there is a potential for other
231    /// references to the connection to be sent across threads, [see this comment](https://github.com/rusqlite/rusqlite/issues/643#issuecomment-640181213).
232    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
241/// A reference to a connection handle with a lifetime bound to something.
242pub struct ConnectionRef<'ctx> {
243    // comes from Connection::from_handle(sqlite3_context_db_handle(...))
244    // and is non-owning
245    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
260/// `feature = "functions"` Aggregate is the callback interface for user-defined
261/// aggregate function.
262///
263/// `A` is the type of the aggregation context and `T` is the type of the final
264/// result. Implementations should be stateless.
265pub trait Aggregate<A, T>
266where
267    A: RefUnwindSafe + UnwindSafe,
268    T: ToSql,
269{
270    /// Initializes the aggregation context. Will be called prior to the first
271    /// call to [`step()`](Aggregate::step) to set up the context for an invocation of the
272    /// function. (Note: `init()` will not be called if there are no rows.)
273    fn init(&self, _: &mut Context<'_>) -> Result<A>;
274
275    /// "step" function called once for each row in an aggregate group. May be
276    /// called 0 times if there are no rows.
277    fn step(&self, _: &mut Context<'_>, _: &mut A) -> Result<()>;
278
279    /// Computes and returns the final result. Will be called exactly once for
280    /// each invocation of the function. If [`step()`](Aggregate::step) was called at least
281    /// once, will be given `Some(A)` (the same `A` as was created by
282    /// [`init`](Aggregate::init) and given to [`step`](Aggregate::step)); if [`step()`](Aggregate::step) was not called (because
283    /// the function is running against 0 rows), will be given `None`.
284    ///
285    /// The passed context will have no arguments.
286    fn finalize(&self, _: &mut Context<'_>, _: Option<A>) -> Result<T>;
287}
288
289/// `feature = "window"` WindowAggregate is the callback interface for
290/// user-defined aggregate window function.
291#[cfg(feature = "window")]
292pub trait WindowAggregate<A, T>: Aggregate<A, T>
293where
294    A: RefUnwindSafe + UnwindSafe,
295    T: ToSql,
296{
297    /// Returns the current value of the aggregate. Unlike xFinal, the
298    /// implementation should not delete any context.
299    fn value(&self, _: Option<&A>) -> Result<T>;
300
301    /// Removes a row from the current window.
302    fn inverse(&self, _: &mut Context<'_>, _: &mut A) -> Result<()>;
303}
304
305bitflags::bitflags! {
306    /// Function Flags.
307    /// See [sqlite3_create_function](https://sqlite.org/c3ref/create_function.html)
308    /// and [Function Flags](https://sqlite.org/c3ref/c_deterministic.html) for details.
309    #[repr(C)]
310    pub struct FunctionFlags: ::std::os::raw::c_int {
311        /// Specifies UTF-8 as the text encoding this SQL function prefers for its parameters.
312        const SQLITE_UTF8     = ffi::SQLITE_UTF8;
313        /// Specifies UTF-16 using little-endian byte order as the text encoding this SQL function prefers for its parameters.
314        const SQLITE_UTF16LE  = ffi::SQLITE_UTF16LE;
315        /// Specifies UTF-16 using big-endian byte order as the text encoding this SQL function prefers for its parameters.
316        const SQLITE_UTF16BE  = ffi::SQLITE_UTF16BE;
317        /// Specifies UTF-16 using native byte order as the text encoding this SQL function prefers for its parameters.
318        const SQLITE_UTF16    = ffi::SQLITE_UTF16;
319        /// Means that the function always gives the same output when the input parameters are the same.
320        const SQLITE_DETERMINISTIC = ffi::SQLITE_DETERMINISTIC;
321        /// Means that the function may only be invoked from top-level SQL.
322        const SQLITE_DIRECTONLY    = 0x0000_0008_0000; // 3.30.0
323        /// Indicates to SQLite that a function may call `sqlite3_value_subtype()` to inspect the sub-types of its arguments.
324        const SQLITE_SUBTYPE       = 0x0000_0010_0000; // 3.30.0
325        /// Means that the function is unlikely to cause problems even if misused.
326        const SQLITE_INNOCUOUS     = 0x0000_0020_0000; // 3.31.0
327    }
328}
329
330impl Default for FunctionFlags {
331    #[inline]
332    fn default() -> FunctionFlags {
333        FunctionFlags::SQLITE_UTF8
334    }
335}
336
337impl Connection {
338    /// `feature = "functions"` Attach a user-defined scalar function to
339    /// this database connection.
340    ///
341    /// `fn_name` is the name the function will be accessible from SQL.
342    /// `n_arg` is the number of arguments to the function. Use `-1` for a
343    /// variable number. If the function always returns the same value
344    /// given the same input, `deterministic` should be `true`.
345    ///
346    /// The function will remain available until the connection is closed or
347    /// until it is explicitly removed via [`remove_function`](Connection::remove_function).
348    ///
349    /// # Example
350    ///
351    /// ```rust
352    /// # use rusqlite::{Connection, Result};
353    /// # use rusqlite::functions::FunctionFlags;
354    /// fn scalar_function_example(db: Connection) -> Result<()> {
355    ///     db.create_scalar_function(
356    ///         "halve",
357    ///         1,
358    ///         FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
359    ///         |ctx| {
360    ///             let value = ctx.get::<f64>(0)?;
361    ///             Ok(value / 2f64)
362    ///         },
363    ///     )?;
364    ///
365    ///     let six_halved: f64 = db.query_row("SELECT halve(6)", [], |r| r.get(0))?;
366    ///     assert_eq!(six_halved, 3f64);
367    ///     Ok(())
368    /// }
369    /// ```
370    ///
371    /// # Failure
372    ///
373    /// Will return Err if the function could not be attached to the connection.
374    #[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    /// `feature = "functions"` Attach a user-defined aggregate function to this
392    /// database connection.
393    ///
394    /// # Failure
395    ///
396    /// Will return Err if the function could not be attached to the connection.
397    #[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    /// `feature = "window"` Attach a user-defined aggregate window function to
416    /// this database connection.
417    ///
418    /// See `https://sqlite.org/windowfunctions.html#udfwinfunc` for more
419    /// information.
420    #[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    /// `feature = "functions"` Removes a user-defined function from this
440    /// database connection.
441    ///
442    /// `fn_name` and `n_arg` should match the name and number of arguments
443    /// given to [`create_scalar_function`](Connection::create_scalar_function) or [`create_aggregate_function`](Connection::create_aggregate_function).
444    ///
445    /// # Failure
446    ///
447    /// Will return Err if the function could not be removed.
448    #[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    // Within the xFinal callback, it is customary to set N=0 in calls to
703    // sqlite3_aggregate_context(C,N) so that no pointless memory allocations occur.
704    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    // Within the xValue callback, it is customary to set N=0 in calls to
748    // sqlite3_aggregate_context(C,N) so that no pointless memory allocations occur.
749    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    // This implementation of a regexp scalar function uses SQLite's auxilliary data
835    // (https://www.sqlite.org/c3ref/get_auxdata.html) to avoid recompiling the regular
836    // expression multiple times within one query.
837    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        // Doesn't actually matter, we'll assert in the function if there's a problem.
939        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        // sum should return NULL when given no columns (contrast with count below)
987        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        // count should return 0 when given no columns (contrast with sum above)
1013        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}