Skip to main content

sqlite3_ext/function/
context.rs

1use super::FromUserData;
2use crate::{ffi, sqlite3_match_version, types::*, value::*, Connection};
3use sealed::sealed;
4use std::{
5    any::TypeId,
6    ffi::CString,
7    mem::{size_of, MaybeUninit},
8};
9
10#[repr(transparent)]
11pub(crate) struct InternalContext {
12    base: ffi::sqlite3_context,
13}
14
15/// Describes the run-time environment of an application-defined function.
16#[repr(transparent)]
17pub struct Context {
18    base: ffi::sqlite3_context,
19}
20
21struct AggregateContext<T> {
22    init: bool,
23    val: MaybeUninit<T>,
24}
25
26#[repr(C)]
27struct AuxData<T> {
28    type_id: TypeId,
29    val: T,
30}
31
32impl InternalContext {
33    pub unsafe fn from_ptr<'a>(base: *mut ffi::sqlite3_context) -> &'a mut Self {
34        &mut *(base as *mut Self)
35    }
36
37    pub fn as_ptr(&self) -> *mut ffi::sqlite3_context {
38        &self.base as *const ffi::sqlite3_context as _
39    }
40
41    /// # Safety
42    ///
43    /// The called must verify that Rust pointer aliasing rules are followed.
44    pub unsafe fn user_data<U>(&self) -> &mut U {
45        &mut *(ffi::sqlite3_user_data(self.as_ptr()) as *mut U)
46    }
47
48    /// Get the aggregate context, returning a mutable reference to it.
49    pub unsafe fn aggregate_context<U, F: FromUserData<U>>(&mut self) -> Result<&mut F> {
50        let ptr =
51            ffi::sqlite3_aggregate_context(self.as_ptr(), size_of::<AggregateContext<F>>() as _)
52                as *mut AggregateContext<F>;
53        if ptr.is_null() {
54            return Err(SQLITE_NOMEM);
55        }
56        let context = &mut *ptr;
57        if !context.init {
58            context.val = MaybeUninit::new(F::from_user_data(self.user_data()));
59            context.init = true;
60        }
61        Ok(context.val.assume_init_mut())
62    }
63
64    /// Try to get the aggregate context, consuming it if it is found.
65    pub unsafe fn try_aggregate_context<U, F: FromUserData<U>>(&mut self) -> Option<F> {
66        let ptr = ffi::sqlite3_aggregate_context(self.as_ptr(), 0 as _) as *mut AggregateContext<F>;
67        if ptr.is_null() {
68            return None;
69        }
70        let context = &mut *ptr;
71        if !context.init {
72            None
73        } else {
74            context.init = false;
75            Some(context.val.assume_init_read())
76        }
77    }
78}
79
80impl Context {
81    pub(crate) fn as_ptr<'a>(&self) -> *mut ffi::sqlite3_context {
82        &self.base as *const ffi::sqlite3_context as _
83    }
84
85    pub(crate) unsafe fn from_ptr<'a>(base: *mut ffi::sqlite3_context) -> &'a mut Self {
86        &mut *(base as *mut Self)
87    }
88
89    /// Return a handle to the current database.
90    pub fn db(&self) -> &Connection {
91        unsafe { Connection::from_ptr(ffi::sqlite3_context_db_handle(self.as_ptr())) }
92    }
93
94    /// Retrieve data about a function parameter that was previously set with
95    /// [set_aux_data](Context::set_aux_data).
96    ///
97    /// This method returns None if T is different from the data type that was stored
98    /// previously.
99    pub fn aux_data<T: 'static>(&self, idx: usize) -> Option<&mut T> {
100        unsafe {
101            let data = ffi::sqlite3_get_auxdata(self.as_ptr(), idx as _) as *mut AuxData<T>;
102            if data.is_null() {
103                None
104            } else {
105                let data = &mut *data;
106                if data.type_id == TypeId::of::<T>() {
107                    Some(&mut data.val)
108                } else {
109                    None
110                }
111            }
112        }
113    }
114
115    /// Set the auxiliary data associated with the corresponding function parameter.
116    ///
117    /// If some processing is necessary in order for a function parameter to be useful (for
118    /// example, compiling a regular expression), this method can be used to cache the
119    /// processed value in case it is later reused in the same query. The cached value can
120    /// be retrieved with the [aux_data](Context::aux_data) method.
121    pub fn set_aux_data<T: 'static>(&self, idx: usize, val: T) {
122        let data = Box::new(AuxData {
123            type_id: TypeId::of::<T>(),
124            val,
125        });
126        unsafe {
127            ffi::sqlite3_set_auxdata(
128                self.as_ptr(),
129                idx as _,
130                Box::into_raw(data) as _,
131                Some(ffi::drop_boxed::<AuxData<T>>),
132            )
133        };
134    }
135
136    /// Assign the given value to the result of the function. This function always returns Ok.
137    pub fn set_result(&self, val: impl ToContextResult) -> Result<()> {
138        unsafe { val.assign_to(self.as_ptr()) };
139        Ok(())
140    }
141}
142
143/// A value that can be returned from an SQL function.
144///
145/// There are several useful implementations available:
146///
147/// - For nullable values, Option\<ToContextResult\> provides an implementation.
148/// - For fallible functions, [Result]\<ToContextResult\> provides an implementation.
149/// - For arbitrary Rust objects, [PassedRef] provides an implementation.
150/// - For borrowed SQLite values, &[ValueRef] provides an implementation. Note that you have to
151///   reborrow as immutable in most cases: `&*value_ref`.
152/// - For owned types known only at run-time, [Value] provides an implementation.
153#[sealed]
154pub trait ToContextResult {
155    #[doc(hidden)]
156    unsafe fn assign_to(self, context: *mut ffi::sqlite3_context);
157}
158
159macro_rules! to_context_result {
160    ($($(#[$attr:meta])* match $ty:ty as ($ctx:ident, $val:ident) => $impl:expr),*) => {
161        $(
162        $(#[$attr])*
163        #[sealed]
164        impl ToContextResult for $ty {
165            unsafe fn assign_to(self, $ctx: *mut ffi::sqlite3_context) {
166                let $val = self;
167                $impl
168            }
169        }
170        )*
171    };
172}
173
174to_context_result! {
175    /// Assign NULL to the context result.
176    match () as (ctx, _val) => ffi::sqlite3_result_null(ctx),
177    match bool as (ctx, val) => ffi::sqlite3_result_int(ctx, val as i32),
178    match i32 as (ctx, val) => ffi::sqlite3_result_int(ctx, val),
179    match i64 as (ctx, val) => ffi::sqlite3_result_int64(ctx, val),
180    match f64 as (ctx, val) => ffi::sqlite3_result_double(ctx, val),
181    /// Assign a static string to the context result.
182    match &'static str as (ctx, val) => {
183        let val = val.as_bytes();
184        let len = val.len();
185        sqlite3_match_version! {
186            3_008_007 => ffi::sqlite3_result_text64(ctx, val.as_ptr() as _, len as _, None, ffi::SQLITE_UTF8 as _),
187            _ => ffi::sqlite3_result_text(ctx, val.as_ptr() as _, len as _, None),
188        }
189    },
190    /// Assign an owned string to the context result.
191    match String as (ctx, val) => {
192        let val = val.as_bytes();
193        let len = val.len();
194        let cstring = CString::new(val).unwrap().into_raw();
195        sqlite3_match_version! {
196            3_008_007 => ffi::sqlite3_result_text64(ctx, cstring, len as _, Some(ffi::drop_cstring), ffi::SQLITE_UTF8 as _),
197            _ => ffi::sqlite3_result_text(ctx, cstring, len as _, Some(ffi::drop_cstring)),
198        }
199    },
200    match Blob as (ctx, val) => {
201        let len = val.len();
202        sqlite3_match_version! {
203            3_008_007 => ffi::sqlite3_result_blob64(ctx, val.into_raw(), len as _, Some(ffi::drop_blob),),
204            _ => ffi::sqlite3_result_blob(ctx, val.into_raw(), len as _, Some(ffi::drop_blob)),
205        }
206    },
207    /// Sets the context error to this error.
208    match Error as (ctx, err) => {
209        match err {
210            Error::Sqlite(_, Some(desc)) => {
211                let bytes = desc.as_bytes();
212                ffi::sqlite3_result_error(ctx, bytes.as_ptr() as _, bytes.len() as _)
213            },
214            Error::Sqlite(code, None) => ffi::sqlite3_result_error_code(ctx, code),
215            Error::NoChange => (),
216            _ => {
217                let msg = format!("{err}");
218                let msg = msg.as_bytes();
219                let len = msg.len();
220                ffi::sqlite3_result_error(ctx, msg.as_ptr() as _, len as _);
221            }
222        }
223    }
224}
225
226/// Sets the context result to the contained value.
227#[sealed]
228impl<'a> ToContextResult for &'a ValueRef {
229    unsafe fn assign_to(self, ctx: *mut ffi::sqlite3_context) {
230        ffi::sqlite3_result_value(ctx, self.as_ptr())
231    }
232}
233
234/// Sets the context result to the contained value.
235#[sealed]
236impl<'a> ToContextResult for &'a mut ValueRef {
237    unsafe fn assign_to(self, ctx: *mut ffi::sqlite3_context) {
238        ffi::sqlite3_result_value(ctx, self.as_ptr())
239    }
240}
241
242/// Sets the context result to the given BLOB.
243#[sealed]
244impl<'a> ToContextResult for &'a [u8] {
245    unsafe fn assign_to(self, ctx: *mut ffi::sqlite3_context) {
246        let len = self.len();
247        sqlite3_match_version! {
248            3_008_007 => ffi::sqlite3_result_blob64(
249                ctx,
250                self.as_ptr() as _,
251                len as _,
252                ffi::sqlite_transient(),
253            ),
254            _ => ffi::sqlite3_result_blob(
255                ctx,
256                self.as_ptr() as _,
257                len as _,
258                ffi::sqlite_transient(),
259            ),
260        }
261    }
262}
263
264/// Sets the context result to the given BLOB.
265#[sealed]
266impl<'a, const N: usize> ToContextResult for &'a [u8; N] {
267    unsafe fn assign_to(self, ctx: *mut ffi::sqlite3_context) {
268        self.as_slice().assign_to(ctx);
269    }
270}
271
272/// Sets the context result to the contained value or NULL.
273#[sealed]
274impl<T: ToContextResult> ToContextResult for Option<T> {
275    unsafe fn assign_to(self, context: *mut ffi::sqlite3_context) {
276        match self {
277            Some(x) => x.assign_to(context),
278            None => ().assign_to(context),
279        }
280    }
281}
282
283/// Sets either the context result or error, depending on the result.
284#[sealed]
285impl<T: ToContextResult> ToContextResult for Result<T> {
286    unsafe fn assign_to(self, context: *mut ffi::sqlite3_context) {
287        match self {
288            Ok(x) => x.assign_to(context),
289            Err(x) => x.assign_to(context),
290        }
291    }
292}
293
294/// Sets a dynamically typed [Value] to the context result.
295#[sealed]
296impl ToContextResult for Value {
297    unsafe fn assign_to(self, context: *mut ffi::sqlite3_context) {
298        match self {
299            Value::Integer(x) => x.assign_to(context),
300            Value::Float(x) => x.assign_to(context),
301            Value::Text(x) => x.assign_to(context),
302            Value::Blob(x) => x.assign_to(context),
303            Value::Null => ().assign_to(context),
304        }
305    }
306}
307
308/// Sets an arbitrary pointer to the context result.
309#[sealed]
310impl<T: 'static + ?Sized> ToContextResult for UnsafePtr<T> {
311    unsafe fn assign_to(self, context: *mut ffi::sqlite3_context) {
312        sqlite3_match_version! {
313        3_009_000 => {
314            let subtype = self.subtype;
315            self.to_bytes().assign_to(context);
316            ffi::sqlite3_result_subtype(context, subtype as _);
317        },
318        _ => self.to_bytes().assign_to(context),
319        }
320    }
321}
322
323/// Sets the context result to NULL with this value as an associated pointer.
324#[sealed]
325impl<T: 'static> ToContextResult for PassedRef<T> {
326    unsafe fn assign_to(self, context: *mut ffi::sqlite3_context) {
327        let _ = (POINTER_TAG, context);
328        sqlite3_match_version! {
329            3_020_000 => ffi::sqlite3_result_pointer(
330                context,
331                Box::into_raw(Box::new(self)) as _,
332                POINTER_TAG,
333                Some(ffi::drop_boxed::<PassedRef<T>>),
334            ),
335            _ => (),
336        }
337    }
338}