rusqlite/
functions.rs

1//! 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 Auxiliary 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    fn constraint_error_code() -> i32 {
79        ffi::SQLITE_CONSTRAINT_FUNCTION
80    }
81
82    if let Error::SqliteFailure(ref err, ref s) = *err {
83        ffi::sqlite3_result_error_code(ctx, err.extended_code);
84        if let Some(Ok(cstr)) = s.as_ref().map(|s| str_to_cstring(s)) {
85            ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1);
86        }
87    } else {
88        ffi::sqlite3_result_error_code(ctx, constraint_error_code());
89        if let Ok(cstr) = str_to_cstring(&err.to_string()) {
90            ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1);
91        }
92    }
93}
94
95unsafe extern "C" fn free_boxed_value<T>(p: *mut c_void) {
96    drop(Box::from_raw(p.cast::<T>()));
97}
98
99/// Context is a wrapper for the SQLite function
100/// evaluation context.
101pub struct Context<'a> {
102    ctx: *mut sqlite3_context,
103    args: &'a [*mut sqlite3_value],
104}
105
106impl Context<'_> {
107    /// Returns the number of arguments to the function.
108    #[inline]
109    #[must_use]
110    pub fn len(&self) -> usize {
111        self.args.len()
112    }
113
114    /// Returns `true` when there is no argument.
115    #[inline]
116    #[must_use]
117    pub fn is_empty(&self) -> bool {
118        self.args.is_empty()
119    }
120
121    /// Returns the `idx`th argument as a `T`.
122    ///
123    /// # Failure
124    ///
125    /// Will panic if `idx` is greater than or equal to
126    /// [`self.len()`](Context::len).
127    ///
128    /// Will return Err if the underlying SQLite type cannot be converted to a
129    /// `T`.
130    pub fn get<T: FromSql>(&self, idx: usize) -> Result<T> {
131        let arg = self.args[idx];
132        let value = unsafe { ValueRef::from_value(arg) };
133        FromSql::column_result(value).map_err(|err| match err {
134            FromSqlError::InvalidType => {
135                Error::InvalidFunctionParameterType(idx, value.data_type())
136            }
137            FromSqlError::OutOfRange(i) => Error::IntegralValueOutOfRange(idx, i),
138            FromSqlError::Other(err) => {
139                Error::FromSqlConversionFailure(idx, value.data_type(), err)
140            }
141            FromSqlError::InvalidBlobSize { .. } => {
142                Error::FromSqlConversionFailure(idx, value.data_type(), Box::new(err))
143            }
144        })
145    }
146
147    /// Returns the `idx`th argument as a `ValueRef`.
148    ///
149    /// # Failure
150    ///
151    /// Will panic if `idx` is greater than or equal to
152    /// [`self.len()`](Context::len).
153    #[inline]
154    #[must_use]
155    pub fn get_raw(&self, idx: usize) -> ValueRef<'_> {
156        let arg = self.args[idx];
157        unsafe { ValueRef::from_value(arg) }
158    }
159
160    /// Returns the subtype of `idx`th argument.
161    ///
162    /// # Failure
163    ///
164    /// Will panic if `idx` is greater than or equal to
165    /// [`self.len()`](Context::len).
166    pub fn get_subtype(&self, idx: usize) -> std::os::raw::c_uint {
167        let arg = self.args[idx];
168        unsafe { ffi::sqlite3_value_subtype(arg) }
169    }
170
171    /// Fetch or insert the auxiliary data associated with a particular
172    /// parameter. This is intended to be an easier-to-use way of fetching it
173    /// compared to calling [`get_aux`](Context::get_aux) and
174    /// [`set_aux`](Context::set_aux) separately.
175    ///
176    /// See `https://www.sqlite.org/c3ref/get_auxdata.html` for a discussion of
177    /// this feature, or the unit tests of this module for an example.
178    pub fn get_or_create_aux<T, E, F>(&self, arg: c_int, func: F) -> Result<Arc<T>>
179    where
180        T: Send + Sync + 'static,
181        E: Into<Box<dyn std::error::Error + Send + Sync + 'static>>,
182        F: FnOnce(ValueRef<'_>) -> Result<T, E>,
183    {
184        if let Some(v) = self.get_aux(arg)? {
185            Ok(v)
186        } else {
187            let vr = self.get_raw(arg as usize);
188            self.set_aux(
189                arg,
190                func(vr).map_err(|e| Error::UserFunctionError(e.into()))?,
191            )
192        }
193    }
194
195    /// Sets the auxiliary data associated with a particular parameter. See
196    /// `https://www.sqlite.org/c3ref/get_auxdata.html` for a discussion of
197    /// this feature, or the unit tests of this module for an example.
198    pub fn set_aux<T: Send + Sync + 'static>(&self, arg: c_int, value: T) -> Result<Arc<T>> {
199        let orig: Arc<T> = Arc::new(value);
200        let inner: AuxInner = orig.clone();
201        let outer = Box::new(inner);
202        let raw: *mut AuxInner = Box::into_raw(outer);
203        unsafe {
204            ffi::sqlite3_set_auxdata(
205                self.ctx,
206                arg,
207                raw.cast(),
208                Some(free_boxed_value::<AuxInner>),
209            );
210        };
211        Ok(orig)
212    }
213
214    /// Gets the auxiliary data that was associated with a given parameter via
215    /// [`set_aux`](Context::set_aux). Returns `Ok(None)` if no data has been
216    /// associated, and Ok(Some(v)) if it has. Returns an error if the
217    /// requested type does not match.
218    pub fn get_aux<T: Send + Sync + 'static>(&self, arg: c_int) -> Result<Option<Arc<T>>> {
219        let p = unsafe { ffi::sqlite3_get_auxdata(self.ctx, arg) as *const AuxInner };
220        if p.is_null() {
221            Ok(None)
222        } else {
223            let v: AuxInner = AuxInner::clone(unsafe { &*p });
224            v.downcast::<T>()
225                .map(Some)
226                .map_err(|_| Error::GetAuxWrongType)
227        }
228    }
229
230    /// Get the db connection handle via [sqlite3_context_db_handle](https://www.sqlite.org/c3ref/context_db_handle.html)
231    ///
232    /// # Safety
233    ///
234    /// This function is marked unsafe because there is a potential for other
235    /// references to the connection to be sent across threads, [see this comment](https://github.com/rusqlite/rusqlite/issues/643#issuecomment-640181213).
236    pub unsafe fn get_connection(&self) -> Result<ConnectionRef<'_>> {
237        let handle = ffi::sqlite3_context_db_handle(self.ctx);
238        Ok(ConnectionRef {
239            conn: Connection::from_handle(handle)?,
240            phantom: PhantomData,
241        })
242    }
243
244    /// Set the Subtype of an SQL function
245    pub fn set_result_subtype(&self, sub_type: std::os::raw::c_uint) {
246        unsafe { ffi::sqlite3_result_subtype(self.ctx, sub_type) };
247    }
248}
249
250/// A reference to a connection handle with a lifetime bound to something.
251pub struct ConnectionRef<'ctx> {
252    // comes from Connection::from_handle(sqlite3_context_db_handle(...))
253    // and is non-owning
254    conn: Connection,
255    phantom: PhantomData<&'ctx Context<'ctx>>,
256}
257
258impl Deref for ConnectionRef<'_> {
259    type Target = Connection;
260
261    #[inline]
262    fn deref(&self) -> &Connection {
263        &self.conn
264    }
265}
266
267type AuxInner = Arc<dyn Any + Send + Sync + 'static>;
268
269/// Aggregate is the callback interface for user-defined
270/// aggregate function.
271///
272/// `A` is the type of the aggregation context and `T` is the type of the final
273/// result. Implementations should be stateless.
274pub trait Aggregate<A, T>
275where
276    A: RefUnwindSafe + UnwindSafe,
277    T: ToSql,
278{
279    /// Initializes the aggregation context. Will be called prior to the first
280    /// call to [`step()`](Aggregate::step) to set up the context for an
281    /// invocation of the function. (Note: `init()` will not be called if
282    /// there are no rows.)
283    fn init(&self, _: &mut Context<'_>) -> Result<A>;
284
285    /// "step" function called once for each row in an aggregate group. May be
286    /// called 0 times if there are no rows.
287    fn step(&self, _: &mut Context<'_>, _: &mut A) -> Result<()>;
288
289    /// Computes and returns the final result. Will be called exactly once for
290    /// each invocation of the function. If [`step()`](Aggregate::step) was
291    /// called at least once, will be given `Some(A)` (the same `A` as was
292    /// created by [`init`](Aggregate::init) and given to
293    /// [`step`](Aggregate::step)); if [`step()`](Aggregate::step) was not
294    /// called (because the function is running against 0 rows), will be
295    /// given `None`.
296    ///
297    /// The passed context will have no arguments.
298    fn finalize(&self, _: &mut Context<'_>, _: Option<A>) -> Result<T>;
299}
300
301/// `WindowAggregate` is the callback interface for
302/// user-defined aggregate window function.
303#[cfg(feature = "window")]
304#[cfg_attr(docsrs, doc(cfg(feature = "window")))]
305pub trait WindowAggregate<A, T>: Aggregate<A, T>
306where
307    A: RefUnwindSafe + UnwindSafe,
308    T: ToSql,
309{
310    /// Returns the current value of the aggregate. Unlike xFinal, the
311    /// implementation should not delete any context.
312    fn value(&self, _: Option<&A>) -> Result<T>;
313
314    /// Removes a row from the current window.
315    fn inverse(&self, _: &mut Context<'_>, _: &mut A) -> Result<()>;
316}
317
318bitflags::bitflags! {
319    /// Function Flags.
320    /// See [sqlite3_create_function](https://sqlite.org/c3ref/create_function.html)
321    /// and [Function Flags](https://sqlite.org/c3ref/c_deterministic.html) for details.
322    #[repr(C)]
323    pub struct FunctionFlags: ::std::os::raw::c_int {
324        /// Specifies UTF-8 as the text encoding this SQL function prefers for its parameters.
325        const SQLITE_UTF8     = ffi::SQLITE_UTF8;
326        /// Specifies UTF-16 using little-endian byte order as the text encoding this SQL function prefers for its parameters.
327        const SQLITE_UTF16LE  = ffi::SQLITE_UTF16LE;
328        /// Specifies UTF-16 using big-endian byte order as the text encoding this SQL function prefers for its parameters.
329        const SQLITE_UTF16BE  = ffi::SQLITE_UTF16BE;
330        /// Specifies UTF-16 using native byte order as the text encoding this SQL function prefers for its parameters.
331        const SQLITE_UTF16    = ffi::SQLITE_UTF16;
332        /// Means that the function always gives the same output when the input parameters are the same.
333        const SQLITE_DETERMINISTIC = ffi::SQLITE_DETERMINISTIC; // 3.8.3
334        /// Means that the function may only be invoked from top-level SQL.
335        const SQLITE_DIRECTONLY    = 0x0000_0008_0000; // 3.30.0
336        /// Indicates to SQLite that a function may call `sqlite3_value_subtype()` to inspect the sub-types of its arguments.
337        const SQLITE_SUBTYPE       = 0x0000_0010_0000; // 3.30.0
338        /// Means that the function is unlikely to cause problems even if misused.
339        const SQLITE_INNOCUOUS     = 0x0000_0020_0000; // 3.31.0
340    }
341}
342
343impl Default for FunctionFlags {
344    #[inline]
345    fn default() -> FunctionFlags {
346        FunctionFlags::SQLITE_UTF8
347    }
348}
349
350impl Connection {
351    /// Attach a user-defined scalar function to
352    /// this database connection.
353    ///
354    /// `fn_name` is the name the function will be accessible from SQL.
355    /// `n_arg` is the number of arguments to the function. Use `-1` for a
356    /// variable number. If the function always returns the same value
357    /// given the same input, `deterministic` should be `true`.
358    ///
359    /// The function will remain available until the connection is closed or
360    /// until it is explicitly removed via
361    /// [`remove_function`](Connection::remove_function).
362    ///
363    /// # Example
364    ///
365    /// ```rust
366    /// # use rusqlite::{Connection, Result};
367    /// # use rusqlite::functions::FunctionFlags;
368    /// fn scalar_function_example(db: Connection) -> Result<()> {
369    ///     db.create_scalar_function(
370    ///         "halve",
371    ///         1,
372    ///         FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
373    ///         |ctx| {
374    ///             let value = ctx.get::<f64>(0)?;
375    ///             Ok(value / 2f64)
376    ///         },
377    ///     )?;
378    ///
379    ///     let six_halved: f64 = db.query_row("SELECT halve(6)", [], |r| r.get(0))?;
380    ///     assert_eq!(six_halved, 3f64);
381    ///     Ok(())
382    /// }
383    /// ```
384    ///
385    /// # Failure
386    ///
387    /// Will return Err if the function could not be attached to the connection.
388    #[inline]
389    pub fn create_scalar_function<F, T>(
390        &self,
391        fn_name: &str,
392        n_arg: c_int,
393        flags: FunctionFlags,
394        x_func: F,
395    ) -> Result<()>
396    where
397        F: FnMut(&Context<'_>) -> Result<T> + Send + UnwindSafe + 'static,
398        T: ToSql,
399    {
400        self.db
401            .borrow_mut()
402            .create_scalar_function(fn_name, n_arg, flags, x_func)
403    }
404
405    /// Attach a user-defined aggregate function to this
406    /// database connection.
407    ///
408    /// # Failure
409    ///
410    /// Will return Err if the function could not be attached to the connection.
411    #[inline]
412    pub fn create_aggregate_function<A, D, T>(
413        &self,
414        fn_name: &str,
415        n_arg: c_int,
416        flags: FunctionFlags,
417        aggr: D,
418    ) -> Result<()>
419    where
420        A: RefUnwindSafe + UnwindSafe,
421        D: Aggregate<A, T> + 'static,
422        T: ToSql,
423    {
424        self.db
425            .borrow_mut()
426            .create_aggregate_function(fn_name, n_arg, flags, aggr)
427    }
428
429    /// Attach a user-defined aggregate window function to
430    /// this database connection.
431    ///
432    /// See `https://sqlite.org/windowfunctions.html#udfwinfunc` for more
433    /// information.
434    #[cfg(feature = "window")]
435    #[cfg_attr(docsrs, doc(cfg(feature = "window")))]
436    #[inline]
437    pub fn create_window_function<A, W, T>(
438        &self,
439        fn_name: &str,
440        n_arg: c_int,
441        flags: FunctionFlags,
442        aggr: W,
443    ) -> Result<()>
444    where
445        A: RefUnwindSafe + UnwindSafe,
446        W: WindowAggregate<A, T> + 'static,
447        T: ToSql,
448    {
449        self.db
450            .borrow_mut()
451            .create_window_function(fn_name, n_arg, flags, aggr)
452    }
453
454    /// Removes a user-defined function from this
455    /// database connection.
456    ///
457    /// `fn_name` and `n_arg` should match the name and number of arguments
458    /// given to [`create_scalar_function`](Connection::create_scalar_function)
459    /// or [`create_aggregate_function`](Connection::create_aggregate_function).
460    ///
461    /// # Failure
462    ///
463    /// Will return Err if the function could not be removed.
464    #[inline]
465    pub fn remove_function(&self, fn_name: &str, n_arg: c_int) -> Result<()> {
466        self.db.borrow_mut().remove_function(fn_name, n_arg)
467    }
468}
469
470impl InnerConnection {
471    fn create_scalar_function<F, T>(
472        &mut self,
473        fn_name: &str,
474        n_arg: c_int,
475        flags: FunctionFlags,
476        x_func: F,
477    ) -> Result<()>
478    where
479        F: FnMut(&Context<'_>) -> Result<T> + Send + UnwindSafe + 'static,
480        T: ToSql,
481    {
482        unsafe extern "C" fn call_boxed_closure<F, T>(
483            ctx: *mut sqlite3_context,
484            argc: c_int,
485            argv: *mut *mut sqlite3_value,
486        ) where
487            F: FnMut(&Context<'_>) -> Result<T>,
488            T: ToSql,
489        {
490            let r = catch_unwind(|| {
491                let boxed_f: *mut F = ffi::sqlite3_user_data(ctx).cast::<F>();
492                assert!(!boxed_f.is_null(), "Internal error - null function pointer");
493                let ctx = Context {
494                    ctx,
495                    args: slice::from_raw_parts(argv, argc as usize),
496                };
497                (*boxed_f)(&ctx)
498            });
499            let t = match r {
500                Err(_) => {
501                    report_error(ctx, &Error::UnwindingPanic);
502                    return;
503                }
504                Ok(r) => r,
505            };
506            let t = t.as_ref().map(|t| ToSql::to_sql(t));
507
508            match t {
509                Ok(Ok(ref value)) => set_result(ctx, value),
510                Ok(Err(err)) => report_error(ctx, &err),
511                Err(err) => report_error(ctx, err),
512            }
513        }
514
515        let boxed_f: *mut F = Box::into_raw(Box::new(x_func));
516        let c_name = str_to_cstring(fn_name)?;
517        let r = unsafe {
518            ffi::sqlite3_create_function_v2(
519                self.db(),
520                c_name.as_ptr(),
521                n_arg,
522                flags.bits(),
523                boxed_f.cast::<c_void>(),
524                Some(call_boxed_closure::<F, T>),
525                None,
526                None,
527                Some(free_boxed_value::<F>),
528            )
529        };
530        self.decode_result(r)
531    }
532
533    fn create_aggregate_function<A, D, T>(
534        &mut self,
535        fn_name: &str,
536        n_arg: c_int,
537        flags: FunctionFlags,
538        aggr: D,
539    ) -> Result<()>
540    where
541        A: RefUnwindSafe + UnwindSafe,
542        D: Aggregate<A, T> + 'static,
543        T: ToSql,
544    {
545        let boxed_aggr: *mut D = Box::into_raw(Box::new(aggr));
546        let c_name = str_to_cstring(fn_name)?;
547        let r = unsafe {
548            ffi::sqlite3_create_function_v2(
549                self.db(),
550                c_name.as_ptr(),
551                n_arg,
552                flags.bits(),
553                boxed_aggr.cast::<c_void>(),
554                None,
555                Some(call_boxed_step::<A, D, T>),
556                Some(call_boxed_final::<A, D, T>),
557                Some(free_boxed_value::<D>),
558            )
559        };
560        self.decode_result(r)
561    }
562
563    #[cfg(feature = "window")]
564    fn create_window_function<A, W, T>(
565        &mut self,
566        fn_name: &str,
567        n_arg: c_int,
568        flags: FunctionFlags,
569        aggr: W,
570    ) -> Result<()>
571    where
572        A: RefUnwindSafe + UnwindSafe,
573        W: WindowAggregate<A, T> + 'static,
574        T: ToSql,
575    {
576        let boxed_aggr: *mut W = Box::into_raw(Box::new(aggr));
577        let c_name = str_to_cstring(fn_name)?;
578        let r = unsafe {
579            ffi::sqlite3_create_window_function(
580                self.db(),
581                c_name.as_ptr(),
582                n_arg,
583                flags.bits(),
584                boxed_aggr.cast::<c_void>(),
585                Some(call_boxed_step::<A, W, T>),
586                Some(call_boxed_final::<A, W, T>),
587                Some(call_boxed_value::<A, W, T>),
588                Some(call_boxed_inverse::<A, W, T>),
589                Some(free_boxed_value::<W>),
590            )
591        };
592        self.decode_result(r)
593    }
594
595    fn remove_function(&mut self, fn_name: &str, n_arg: c_int) -> Result<()> {
596        let c_name = str_to_cstring(fn_name)?;
597        let r = unsafe {
598            ffi::sqlite3_create_function_v2(
599                self.db(),
600                c_name.as_ptr(),
601                n_arg,
602                ffi::SQLITE_UTF8,
603                ptr::null_mut(),
604                None,
605                None,
606                None,
607                None,
608            )
609        };
610        self.decode_result(r)
611    }
612}
613
614unsafe fn aggregate_context<A>(ctx: *mut sqlite3_context, bytes: usize) -> Option<*mut *mut A> {
615    let pac = ffi::sqlite3_aggregate_context(ctx, bytes as c_int) as *mut *mut A;
616    if pac.is_null() {
617        return None;
618    }
619    Some(pac)
620}
621
622unsafe extern "C" fn call_boxed_step<A, D, T>(
623    ctx: *mut sqlite3_context,
624    argc: c_int,
625    argv: *mut *mut sqlite3_value,
626) where
627    A: RefUnwindSafe + UnwindSafe,
628    D: Aggregate<A, T>,
629    T: ToSql,
630{
631    let pac = if let Some(pac) = aggregate_context(ctx, std::mem::size_of::<*mut A>()) {
632        pac
633    } else {
634        ffi::sqlite3_result_error_nomem(ctx);
635        return;
636    };
637
638    let r = catch_unwind(|| {
639        let boxed_aggr: *mut D = ffi::sqlite3_user_data(ctx).cast::<D>();
640        assert!(
641            !boxed_aggr.is_null(),
642            "Internal error - null aggregate pointer"
643        );
644        let mut ctx = Context {
645            ctx,
646            args: slice::from_raw_parts(argv, argc as usize),
647        };
648
649        if (*pac as *mut A).is_null() {
650            *pac = Box::into_raw(Box::new((*boxed_aggr).init(&mut ctx)?));
651        }
652
653        (*boxed_aggr).step(&mut ctx, &mut **pac)
654    });
655    let r = match r {
656        Err(_) => {
657            report_error(ctx, &Error::UnwindingPanic);
658            return;
659        }
660        Ok(r) => r,
661    };
662    match r {
663        Ok(_) => {}
664        Err(err) => report_error(ctx, &err),
665    };
666}
667
668#[cfg(feature = "window")]
669unsafe extern "C" fn call_boxed_inverse<A, W, T>(
670    ctx: *mut sqlite3_context,
671    argc: c_int,
672    argv: *mut *mut sqlite3_value,
673) where
674    A: RefUnwindSafe + UnwindSafe,
675    W: WindowAggregate<A, T>,
676    T: ToSql,
677{
678    let pac = if let Some(pac) = aggregate_context(ctx, std::mem::size_of::<*mut A>()) {
679        pac
680    } else {
681        ffi::sqlite3_result_error_nomem(ctx);
682        return;
683    };
684
685    let r = catch_unwind(|| {
686        let boxed_aggr: *mut W = ffi::sqlite3_user_data(ctx).cast::<W>();
687        assert!(
688            !boxed_aggr.is_null(),
689            "Internal error - null aggregate pointer"
690        );
691        let mut ctx = Context {
692            ctx,
693            args: slice::from_raw_parts(argv, argc as usize),
694        };
695        (*boxed_aggr).inverse(&mut ctx, &mut **pac)
696    });
697    let r = match r {
698        Err(_) => {
699            report_error(ctx, &Error::UnwindingPanic);
700            return;
701        }
702        Ok(r) => r,
703    };
704    match r {
705        Ok(_) => {}
706        Err(err) => report_error(ctx, &err),
707    };
708}
709
710unsafe extern "C" fn call_boxed_final<A, D, T>(ctx: *mut sqlite3_context)
711where
712    A: RefUnwindSafe + UnwindSafe,
713    D: Aggregate<A, T>,
714    T: ToSql,
715{
716    // Within the xFinal callback, it is customary to set N=0 in calls to
717    // sqlite3_aggregate_context(C,N) so that no pointless memory allocations occur.
718    let a: Option<A> = match aggregate_context(ctx, 0) {
719        Some(pac) => {
720            if (*pac as *mut A).is_null() {
721                None
722            } else {
723                let a = Box::from_raw(*pac);
724                Some(*a)
725            }
726        }
727        None => None,
728    };
729
730    let r = catch_unwind(|| {
731        let boxed_aggr: *mut D = ffi::sqlite3_user_data(ctx).cast::<D>();
732        assert!(
733            !boxed_aggr.is_null(),
734            "Internal error - null aggregate pointer"
735        );
736        let mut ctx = Context { ctx, args: &mut [] };
737        (*boxed_aggr).finalize(&mut ctx, a)
738    });
739    let t = match r {
740        Err(_) => {
741            report_error(ctx, &Error::UnwindingPanic);
742            return;
743        }
744        Ok(r) => r,
745    };
746    let t = t.as_ref().map(|t| ToSql::to_sql(t));
747    match t {
748        Ok(Ok(ref value)) => set_result(ctx, value),
749        Ok(Err(err)) => report_error(ctx, &err),
750        Err(err) => report_error(ctx, err),
751    }
752}
753
754#[cfg(feature = "window")]
755unsafe extern "C" fn call_boxed_value<A, W, T>(ctx: *mut sqlite3_context)
756where
757    A: RefUnwindSafe + UnwindSafe,
758    W: WindowAggregate<A, T>,
759    T: ToSql,
760{
761    // Within the xValue callback, it is customary to set N=0 in calls to
762    // sqlite3_aggregate_context(C,N) so that no pointless memory allocations occur.
763    let a: Option<&A> = match aggregate_context(ctx, 0) {
764        Some(pac) => {
765            if (*pac as *mut A).is_null() {
766                None
767            } else {
768                let a = &**pac;
769                Some(a)
770            }
771        }
772        None => None,
773    };
774
775    let r = catch_unwind(|| {
776        let boxed_aggr: *mut W = ffi::sqlite3_user_data(ctx).cast::<W>();
777        assert!(
778            !boxed_aggr.is_null(),
779            "Internal error - null aggregate pointer"
780        );
781        (*boxed_aggr).value(a)
782    });
783    let t = match r {
784        Err(_) => {
785            report_error(ctx, &Error::UnwindingPanic);
786            return;
787        }
788        Ok(r) => r,
789    };
790    let t = t.as_ref().map(|t| ToSql::to_sql(t));
791    match t {
792        Ok(Ok(ref value)) => set_result(ctx, value),
793        Ok(Err(err)) => report_error(ctx, &err),
794        Err(err) => report_error(ctx, err),
795    }
796}
797
798#[cfg(test)]
799mod test {
800    use regex::Regex;
801    use std::os::raw::c_double;
802
803    #[cfg(feature = "window")]
804    use crate::functions::WindowAggregate;
805    use crate::functions::{Aggregate, Context, FunctionFlags};
806    use crate::{Connection, Error, Result};
807
808    fn half(ctx: &Context<'_>) -> Result<c_double> {
809        assert_eq!(ctx.len(), 1, "called with unexpected number of arguments");
810        let value = ctx.get::<c_double>(0)?;
811        Ok(value / 2f64)
812    }
813
814    #[test]
815    fn test_function_half() -> Result<()> {
816        let db = Connection::open_in_memory()?;
817        db.create_scalar_function(
818            "half",
819            1,
820            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
821            half,
822        )?;
823        let result: f64 = db.one_column("SELECT half(6)")?;
824
825        assert!((3f64 - result).abs() < f64::EPSILON);
826        Ok(())
827    }
828
829    #[test]
830    fn test_remove_function() -> Result<()> {
831        let db = Connection::open_in_memory()?;
832        db.create_scalar_function(
833            "half",
834            1,
835            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
836            half,
837        )?;
838        let result: f64 = db.one_column("SELECT half(6)")?;
839        assert!((3f64 - result).abs() < f64::EPSILON);
840
841        db.remove_function("half", 1)?;
842        let result: Result<f64> = db.one_column("SELECT half(6)");
843        result.unwrap_err();
844        Ok(())
845    }
846
847    // This implementation of a regexp scalar function uses SQLite's auxiliary data
848    // (https://www.sqlite.org/c3ref/get_auxdata.html) to avoid recompiling the regular
849    // expression multiple times within one query.
850    fn regexp_with_auxilliary(ctx: &Context<'_>) -> Result<bool> {
851        assert_eq!(ctx.len(), 2, "called with unexpected number of arguments");
852        type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
853        let regexp: std::sync::Arc<Regex> = ctx
854            .get_or_create_aux(0, |vr| -> Result<_, BoxError> {
855                Ok(Regex::new(vr.as_str()?)?)
856            })?;
857
858        let is_match = {
859            let text = ctx
860                .get_raw(1)
861                .as_str()
862                .map_err(|e| Error::UserFunctionError(e.into()))?;
863
864            regexp.is_match(text)
865        };
866
867        Ok(is_match)
868    }
869
870    #[test]
871    fn test_function_regexp_with_auxilliary() -> Result<()> {
872        let db = Connection::open_in_memory()?;
873        db.execute_batch(
874            "BEGIN;
875             CREATE TABLE foo (x string);
876             INSERT INTO foo VALUES ('lisa');
877             INSERT INTO foo VALUES ('lXsi');
878             INSERT INTO foo VALUES ('lisX');
879             END;",
880        )?;
881        db.create_scalar_function(
882            "regexp",
883            2,
884            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
885            regexp_with_auxilliary,
886        )?;
887
888        let result: bool = db.one_column("SELECT regexp('l.s[aeiouy]', 'lisa')")?;
889
890        assert!(result);
891
892        let result: i64 =
893            db.one_column("SELECT COUNT(*) FROM foo WHERE regexp('l.s[aeiouy]', x) == 1")?;
894
895        assert_eq!(2, result);
896        Ok(())
897    }
898
899    #[test]
900    fn test_varargs_function() -> Result<()> {
901        let db = Connection::open_in_memory()?;
902        db.create_scalar_function(
903            "my_concat",
904            -1,
905            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
906            |ctx| {
907                let mut ret = String::new();
908
909                for idx in 0..ctx.len() {
910                    let s = ctx.get::<String>(idx)?;
911                    ret.push_str(&s);
912                }
913
914                Ok(ret)
915            },
916        )?;
917
918        for &(expected, query) in &[
919            ("", "SELECT my_concat()"),
920            ("onetwo", "SELECT my_concat('one', 'two')"),
921            ("abc", "SELECT my_concat('a', 'b', 'c')"),
922        ] {
923            let result: String = db.one_column(query)?;
924            assert_eq!(expected, result);
925        }
926        Ok(())
927    }
928
929    #[test]
930    fn test_get_aux_type_checking() -> Result<()> {
931        let db = Connection::open_in_memory()?;
932        db.create_scalar_function("example", 2, FunctionFlags::default(), |ctx| {
933            if !ctx.get::<bool>(1)? {
934                ctx.set_aux::<i64>(0, 100)?;
935            } else {
936                assert_eq!(ctx.get_aux::<String>(0), Err(Error::GetAuxWrongType));
937                assert_eq!(*ctx.get_aux::<i64>(0)?.unwrap(), 100);
938            }
939            Ok(true)
940        })?;
941
942        let res: bool =
943            db.one_column("SELECT example(0, i) FROM (SELECT 0 as i UNION SELECT 1)")?;
944        // Doesn't actually matter, we'll assert in the function if there's a problem.
945        assert!(res);
946        Ok(())
947    }
948
949    struct Sum;
950    struct Count;
951
952    impl Aggregate<i64, Option<i64>> for Sum {
953        fn init(&self, _: &mut Context<'_>) -> Result<i64> {
954            Ok(0)
955        }
956
957        fn step(&self, ctx: &mut Context<'_>, sum: &mut i64) -> Result<()> {
958            *sum += ctx.get::<i64>(0)?;
959            Ok(())
960        }
961
962        fn finalize(&self, _: &mut Context<'_>, sum: Option<i64>) -> Result<Option<i64>> {
963            Ok(sum)
964        }
965    }
966
967    impl Aggregate<i64, i64> for Count {
968        fn init(&self, _: &mut Context<'_>) -> Result<i64> {
969            Ok(0)
970        }
971
972        fn step(&self, _ctx: &mut Context<'_>, sum: &mut i64) -> Result<()> {
973            *sum += 1;
974            Ok(())
975        }
976
977        fn finalize(&self, _: &mut Context<'_>, sum: Option<i64>) -> Result<i64> {
978            Ok(sum.unwrap_or(0))
979        }
980    }
981
982    #[test]
983    fn test_sum() -> Result<()> {
984        let db = Connection::open_in_memory()?;
985        db.create_aggregate_function(
986            "my_sum",
987            1,
988            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
989            Sum,
990        )?;
991
992        // sum should return NULL when given no columns (contrast with count below)
993        let no_result = "SELECT my_sum(i) FROM (SELECT 2 AS i WHERE 1 <> 1)";
994        let result: Option<i64> = db.one_column(no_result)?;
995        assert!(result.is_none());
996
997        let single_sum = "SELECT my_sum(i) FROM (SELECT 2 AS i UNION ALL SELECT 2)";
998        let result: i64 = db.one_column(single_sum)?;
999        assert_eq!(4, result);
1000
1001        let dual_sum = "SELECT my_sum(i), my_sum(j) FROM (SELECT 2 AS i, 1 AS j UNION ALL SELECT \
1002                        2, 1)";
1003        let result: (i64, i64) = db.query_row(dual_sum, [], |r| Ok((r.get(0)?, r.get(1)?)))?;
1004        assert_eq!((4, 2), result);
1005        Ok(())
1006    }
1007
1008    #[test]
1009    fn test_count() -> Result<()> {
1010        let db = Connection::open_in_memory()?;
1011        db.create_aggregate_function(
1012            "my_count",
1013            -1,
1014            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
1015            Count,
1016        )?;
1017
1018        // count should return 0 when given no columns (contrast with sum above)
1019        let no_result = "SELECT my_count(i) FROM (SELECT 2 AS i WHERE 1 <> 1)";
1020        let result: i64 = db.one_column(no_result)?;
1021        assert_eq!(result, 0);
1022
1023        let single_sum = "SELECT my_count(i) FROM (SELECT 2 AS i UNION ALL SELECT 2)";
1024        let result: i64 = db.one_column(single_sum)?;
1025        assert_eq!(2, result);
1026        Ok(())
1027    }
1028
1029    #[cfg(feature = "window")]
1030    impl WindowAggregate<i64, Option<i64>> for Sum {
1031        fn inverse(&self, ctx: &mut Context<'_>, sum: &mut i64) -> Result<()> {
1032            *sum -= ctx.get::<i64>(0)?;
1033            Ok(())
1034        }
1035
1036        fn value(&self, sum: Option<&i64>) -> Result<Option<i64>> {
1037            Ok(sum.copied())
1038        }
1039    }
1040
1041    #[test]
1042    #[cfg(feature = "window")]
1043    fn test_window() -> Result<()> {
1044        use fallible_iterator::FallibleIterator;
1045
1046        let db = Connection::open_in_memory()?;
1047        db.create_window_function(
1048            "sumint",
1049            1,
1050            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
1051            Sum,
1052        )?;
1053        db.execute_batch(
1054            "CREATE TABLE t3(x, y);
1055             INSERT INTO t3 VALUES('a', 4),
1056                     ('b', 5),
1057                     ('c', 3),
1058                     ('d', 8),
1059                     ('e', 1);",
1060        )?;
1061
1062        let mut stmt = db.prepare(
1063            "SELECT x, sumint(y) OVER (
1064                   ORDER BY x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING
1065                 ) AS sum_y
1066                 FROM t3 ORDER BY x;",
1067        )?;
1068
1069        let results: Vec<(String, i64)> = stmt
1070            .query([])?
1071            .map(|row| Ok((row.get("x")?, row.get("sum_y")?)))
1072            .collect()?;
1073        let expected = vec![
1074            ("a".to_owned(), 9),
1075            ("b".to_owned(), 12),
1076            ("c".to_owned(), 16),
1077            ("d".to_owned(), 12),
1078            ("e".to_owned(), 9),
1079        ];
1080        assert_eq!(expected, results);
1081        Ok(())
1082    }
1083}
rusqlite/functions.rs

rusqlite/
functions.rs
