From 48cfe3080d912a426738d2434fac51a9a4cddd86 Mon Sep 17 00:00:00 2001 From: Ethan Smith Date: Wed, 19 Aug 2026 20:14:14 +0000 Subject: [PATCH 1/5] feat: apply registered scalar functions to every connection `register_function` records a Rust scalar function in a process-global set. Both connection pools install that whole set on each connection they open, through `after_connect`, so a function is present before any caller can use the connection and remains present as the pools drop idle connections and open new ones. Registration validates the name, the arity, and the name-and-arity pair against the set, so a bad registration fails at the call site. SQLite owns one strong reference to each registration per connection and releases it through the `xDestroy` callback when the connection closes. A handler error becomes the statement's SQLite error, carrying the handler's message. A handler panic becomes an error naming the function, because an unwind across the FFI boundary would abort the process. --- Cargo.lock | 1 + crates/sqlx-sqlite-conn-mgr/Cargo.toml | 4 + crates/sqlx-sqlite-conn-mgr/src/database.rs | 18 + crates/sqlx-sqlite-conn-mgr/src/error.rs | 50 + crates/sqlx-sqlite-conn-mgr/src/functions.rs | 1038 ++++++++++++++++++ crates/sqlx-sqlite-conn-mgr/src/lib.rs | 8 + 6 files changed, 1119 insertions(+) create mode 100644 crates/sqlx-sqlite-conn-mgr/src/functions.rs diff --git a/Cargo.lock b/Cargo.lock index 2e329dd..a2585f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3615,6 +3615,7 @@ dependencies = [ name = "sqlx-sqlite-conn-mgr" version = "0.9.0" dependencies = [ + "libsqlite3-sys", "parking_lot", "serde", "sqlx", diff --git a/crates/sqlx-sqlite-conn-mgr/Cargo.toml b/crates/sqlx-sqlite-conn-mgr/Cargo.toml index 81deaec..2f22296 100644 --- a/crates/sqlx-sqlite-conn-mgr/Cargo.toml +++ b/crates/sqlx-sqlite-conn-mgr/Cargo.toml @@ -19,6 +19,10 @@ tokio = { version = "1.49.0", features = ["full"] } tracing = { version = "0.1.44", default-features = false, features = ["std", "release_max_level_off"] } serde = { version = "1.0.228", features = ["derive"] } parking_lot = "0.12.3" +# Registers consumer-supplied scalar functions on every pooled connection. +# Pinned to the version sqlx-sqlite links: only one version of a `links` +# crate can exist in a dependency graph. +libsqlite3-sys = "0.35.0" [dev-dependencies] tempfile = "3.24.0" diff --git a/crates/sqlx-sqlite-conn-mgr/src/database.rs b/crates/sqlx-sqlite-conn-mgr/src/database.rs index 642c127..b94a3a0 100644 --- a/crates/sqlx-sqlite-conn-mgr/src/database.rs +++ b/crates/sqlx-sqlite-conn-mgr/src/database.rs @@ -3,6 +3,7 @@ use crate::Result; use crate::config::SqliteDatabaseConfig; use crate::error::Error; +use crate::functions; use crate::observer_slot::ObserverSlot; use crate::registry::{get_or_open_database, is_memory_database, uncache_database}; use crate::write_guard::WriteGuard; @@ -180,12 +181,27 @@ impl SqliteDatabase { .read_only(true) .optimize_on_close(true, OPTIMIZE_ANALYSIS_LIMIT); + // The functions this database serves: the set registered before this connect. + // Both pools install this same snapshot on every connection they ever open, so + // one database never serves connections with differing function sets. A + // registration made after this call applies to the next database to connect. + let registered_functions = functions::snapshot(); + + // SQLite keeps a scalar function in per-connection state, so every connection + // either pool opens must have the whole snapshot. `after_connect` is the only + // hook that can do this. Both pools open connections lazily and drop them after + // the idle timeout, so a one-time pass over the connections open at startup will + // miss every connection opened later. sqlx runs this hook before it hands the + // connection to a caller. If the hook returns an error, sqlx discards the + // connection, so no caller ever receives a connection that lacks a registered + // function. let read_pool = SqlitePoolOptions::new() .max_connections(config.max_read_connections) .min_connections(0) .idle_timeout(Some(std::time::Duration::from_secs( config.idle_timeout_secs, ))) + .after_connect(functions::install_hook(®istered_functions)) .connect_with(read_options) .await?; @@ -213,6 +229,8 @@ impl SqliteDatabase { .idle_timeout(Some(std::time::Duration::from_secs( config.idle_timeout_secs, ))) + // Registered scalar functions. See the read pool above for why this hook is here. + .after_connect(functions::install_hook(®istered_functions)) .after_release(|conn, _meta| { Box::pin(async move { match sqlx::query("ROLLBACK").execute(&mut *conn).await { diff --git a/crates/sqlx-sqlite-conn-mgr/src/error.rs b/crates/sqlx-sqlite-conn-mgr/src/error.rs index 3a11fbd..b3068e3 100644 --- a/crates/sqlx-sqlite-conn-mgr/src/error.rs +++ b/crates/sqlx-sqlite-conn-mgr/src/error.rs @@ -40,6 +40,56 @@ pub enum Error { )] DuplicateAttachedDatabase(String), + /// Invalid name for a scalar function. See `functions::validate_name` for the rule set + /// this message must match. + #[error( + "Invalid function name '{0}': must be non-empty, at most 255 bytes long, and contain no interior NUL byte" + )] + InvalidFunctionName(String), + + /// Invalid argument count for a scalar function. In SQLite, a negative arity means a + /// variadic function. This crate does not offer that option. The upper bound is the + /// linked library's `SQLITE_MAX_FUNCTION_ARG`, enforced by the probe as + /// `Error::FunctionRefused`. + #[error("Invalid function arity {0}: must not be negative")] + InvalidFunctionArity(i32), + + /// A scalar function with this name and argument count is already registered. This + /// crate compares names case-insensitively, to match how SQLite resolves them. + /// Without this check, a second registration replaces the first instead of the two + /// coexisting. + #[error("A function named '{name}' taking {arity} argument(s) is already registered")] + DuplicateFunction { name: String, arity: i32 }, + + /// The linked SQLite library already defines this name as a built-in with the same + /// argument count, or as a variadic built-in. SQLite silently replaces a built-in with + /// an application-defined function, so every statement on the connection that names the + /// built-in reaches the handler instead, including one inside an existing view, index, + /// or constraint. See `functions::shadows_builtin`. + #[error( + "A function named '{name}' taking {arity} argument(s) replaces a SQLite built-in function. Choose a name no built-in uses." + )] + ShadowsBuiltinFunction { name: String, arity: i32 }, + + /// The in-memory connection that checks a registration failed to open or query, so + /// this crate cannot tell whether the linked library accepts the function. The failure + /// belongs to the probe, not to the function. See `functions::probe`. + #[error( + "Cannot check whether SQLite accepts a function named '{name}': the probe failed: {reason}" + )] + ProbeConnectionFailed { name: String, reason: String }, + + /// The linked SQLite library refused the registration. SQLite validates the name, the + /// argument count, and the encoding flags. The library's own limits decide the + /// outcome: an arity above its `SQLITE_MAX_FUNCTION_ARG` fails here. The `reason` + /// field gives the message SQLite returned. + #[error("SQLite refused a function named '{name}' taking {arity} argument(s): {reason}")] + FunctionRefused { + name: String, + arity: i32, + reason: String, + }, + /// Two attached-database specs used the same schema alias. Compared /// case-insensitively, matching SQLite's own schema namespace - a spec named `"x"` /// and one named `"X"` collide at `ATTACH` even though they compare unequal as diff --git a/crates/sqlx-sqlite-conn-mgr/src/functions.rs b/crates/sqlx-sqlite-conn-mgr/src/functions.rs new file mode 100644 index 0000000..9ec7a8a --- /dev/null +++ b/crates/sqlx-sqlite-conn-mgr/src/functions.rs @@ -0,0 +1,1038 @@ +//! Consumer-supplied scalar SQL functions, applied to every pooled connection. +//! +//! This module doc is the canonical statement of the scalar-function contract. The +//! READMEs, the plugin builder docs, and the CHANGELOG point here. +//! +//! # The contract +//! +//! A function belongs to a *connection*, not to a database file. SQLite keeps +//! user-defined functions in per-connection state, so every connection that runs a query +//! naming one must have it registered. [`register_function`] records a function in a +//! process-global set. [`crate::SqliteDatabase::connect`] snapshots that set, and the +//! pools' `after_connect` hook installs the snapshot on each connection they open. A +//! function is present before any caller receives a connection, and it stays present as +//! the pools drop idle connections and open new ones. +//! +//! Register a function before the `connect` call for the database that must serve it. +//! A database applies the functions registered before its `connect`, for every +//! connection it ever opens. A later registration applies to the next database, never to +//! part of an open one. Registration is process-global, and a function stays registered +//! for the life of the process. [`register_or_replace_function`] puts a different handler +//! under a name already registered, for the databases that connect after it. +//! +//! [`ScalarFunction::invocation_scope`] decides where SQLite accepts a call. +//! [`InvocationScope::DirectOnly`] confines the function to top-level SQL, so SQLite +//! refuses a call from inside a schema object: a view, a trigger, a CHECK constraint, a +//! DEFAULT clause, an expression index, a partial index, or a generated column. The error +//! reads `unsafe use of ()`. Direct-only keeps a handler out of the schema of a +//! database this application did not write, such as an attached file from sync or import. +//! It also stops a caller from naming the function in an expression index or a generated +//! column, which writes the name into the schema and leaves the table unreadable by any +//! process that opens the file without registering the function. The other two scopes +//! accept both of those risks in exchange for a function that a view or an index can +//! call. +//! +//! The handler is `Send + Sync + 'static` because sqlx runs every SQLite connection on +//! its own thread. Every connection of every database shares one handler. The handler +//! receives every argument, including SQL NULL, and decides what NULL means. Returning +//! [`SqlValue::Null`] produces NULL. A returned [`FunctionError`] fails the statement, +//! and the caller receives its message. A handler blocks its connection for as long as it +//! runs, and the write pool holds a single connection. +//! +//! If the application unwinds on panic (the Rust default), a panic inside the handler +//! becomes an error naming the function. Under `panic = "abort"`, the process aborts, and +//! nothing here can prevent that. +//! +//! Each of these fails at the [`register_function`] call, not later as a connection +//! error: +//! +//! - an empty name, a name over 255 bytes, or a name holding a NUL byte +//! - a negative arity +//! - a name a SQLite built-in already uses +//! - a name and arity pair already registered (names compare case-insensitively, +//! matching how SQLite resolves them) +//! - a function the linked SQLite library refuses, such as an arity above its +//! `SQLITE_MAX_FUNCTION_ARG` +//! +//! Aggregate functions, window functions, collations, and virtual tables are not +//! supported. + +use std::borrow::Cow; +use std::ffi::{CStr, CString, c_char, c_int, c_void}; +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::sync::{Arc, OnceLock}; + +use libsqlite3_sys as ffi; +use parking_lot::RwLock; +use sqlx::sqlite::SqliteConnection; +use tracing::{error, trace}; + +use crate::Result; +use crate::error::Error; + +/// SQLite's own limit on the length of a function name, in bytes. +/// See . +const MAX_NAME_BYTES: usize = 255; + +/// The registered functions, in registration order. +static FUNCTIONS: OnceLock>>> = OnceLock::new(); + +fn functions() -> &'static RwLock>> { + FUNCTIONS.get_or_init(|| RwLock::new(Vec::new())) +} + +/// The set of functions a database captures at connect time. +/// +/// Both of a database's pools install this exact set on every connection they open, so +/// one database never serves connections with differing function sets. +pub(crate) type FunctionSet = Arc>>; + +/// The registered set at this moment. Called once per [`crate::SqliteDatabase::connect`]. +pub(crate) fn snapshot() -> FunctionSet { + Arc::new(functions().read().clone()) +} + +/// An owned value returned from a scalar function, or read from a raw SQLite value. +/// +/// The five variants are SQLite's five storage classes. Returning [`SqlValue::Null`] +/// from a handler produces SQL NULL. Handlers receive their arguments as the borrowed +/// [`SqlValueRef`] instead. +#[derive(Debug, Clone, PartialEq)] +pub enum SqlValue { + Null, + Integer(i64), + Real(f64), + Text(String), + Blob(Vec), +} + +/// A borrowed value passed to a scalar function. +/// +/// Text and blob variants point into SQLite's own buffers, so a handler reads its +/// arguments without a copy. [`SqlValueRef::to_owned`] converts one into a [`SqlValue`]. +#[derive(Debug, Clone, PartialEq)] +pub enum SqlValueRef<'a> { + Null, + Integer(i64), + Real(f64), + Text(Cow<'a, str>), + Blob(&'a [u8]), +} + +impl SqlValue { + /// Whether this value is SQL NULL. + pub fn is_null(&self) -> bool { + matches!(self, SqlValue::Null) + } + + /// This value as an integer, when it holds one. + pub fn as_integer(&self) -> Option { + match self { + SqlValue::Integer(value) => Some(*value), + _ => None, + } + } + + /// This value as a float, when it holds one. + pub fn as_real(&self) -> Option { + match self { + SqlValue::Real(value) => Some(*value), + _ => None, + } + } + + /// This value as a string reference, when it holds text. + pub fn as_text(&self) -> Option<&str> { + match self { + SqlValue::Text(text) => Some(text), + _ => None, + } + } + + /// This value as a blob reference, when it holds one. + pub fn as_blob(&self) -> Option<&[u8]> { + match self { + SqlValue::Blob(bytes) => Some(bytes), + _ => None, + } + } +} + +impl SqlValueRef<'_> { + /// Copy this borrowed value into an owned [`SqlValue`]. + pub fn to_owned(&self) -> SqlValue { + match self { + SqlValueRef::Null => SqlValue::Null, + SqlValueRef::Integer(value) => SqlValue::Integer(*value), + SqlValueRef::Real(value) => SqlValue::Real(*value), + SqlValueRef::Text(text) => SqlValue::Text(text.to_string()), + SqlValueRef::Blob(bytes) => SqlValue::Blob(bytes.to_vec()), + } + } +} + +/// A scalar function's failure. +/// +/// The message becomes the SQLite error for the statement that invoked the function. See +/// [`ScalarFunction`] for what the caller sees. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("{message}")] +pub struct FunctionError { + message: String, +} + +impl FunctionError { + /// Build an error from `message`. + pub fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } + + /// The message this error reports to the caller. + pub fn message(&self) -> &str { + &self.message + } +} + +/// The closure that computes a scalar function's result. +/// +/// The bound is `Send + Sync + 'static` because sqlx runs every SQLite connection on its +/// own thread. Every connection of every database shares one handler: up to +/// `max_read_connections` readers plus one writer per database. +pub type ScalarHandler = + Arc]) -> std::result::Result + Send + Sync>; + +/// Where SQLite accepts a call to a registered function. +/// +/// A function a schema object can call runs whenever anything reads through that object, +/// including an object in a database file this application did not write. Pick +/// [`InvocationScope::DirectOnly`] unless the function has a reason to appear in a view, a +/// trigger, or an index. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum InvocationScope { + /// Top-level SQL only, through the `SQLITE_DIRECTONLY` flag. SQLite refuses a call + /// from inside a view, a trigger, a CHECK constraint, a DEFAULT clause, an expression + /// index, a partial index, or a generated column, with the error + /// `unsafe use of ()`. + #[default] + DirectOnly, + + /// Schema objects as well, on a connection that trusts the schema. SQLite trusts the + /// schema unless a caller runs `PRAGMA trusted_schema = OFF`. This scope sets neither + /// `SQLITE_DIRECTONLY` nor `SQLITE_INNOCUOUS`. + Schema, + + /// Schema objects as well, on every connection, including one that turns trusted + /// schema off. This scope sets `SQLITE_INNOCUOUS`, which promises SQLite that the + /// handler reads nothing but its arguments and changes nothing but its result. A + /// handler that reads a file, opens a socket, or touches process state breaks that + /// promise, and SQLite has no way to detect it. + InnocuousSchema, +} + +/// A scalar SQL function to register on every connection. +/// +/// # Example +/// +/// ```no_run +/// use std::sync::Arc; +/// use sqlx_sqlite_conn_mgr::{ +/// FunctionError, InvocationScope, ScalarFunction, SqlValue, SqlValueRef, +/// register_function, +/// }; +/// +/// # fn example() -> sqlx_sqlite_conn_mgr::Result<()> { +/// register_function(ScalarFunction { +/// name: "normalize_for_search".into(), +/// arity: 1, +/// deterministic: true, +/// invocation_scope: InvocationScope::DirectOnly, +/// // SQLite's own lower() folds ASCII only, so Unicode case folding needs Rust. +/// handler: Arc::new(|args: &[SqlValueRef]| match &args[0] { +/// SqlValueRef::Text(text) => Ok(SqlValue::Text(text.to_lowercase())), +/// SqlValueRef::Null => Ok(SqlValue::Null), +/// _ => Err(FunctionError::new("normalize_for_search expects text")), +/// }), +/// })?; +/// # Ok(()) +/// # } +/// ``` +pub struct ScalarFunction { + /// The SQL identifier, for example `normalize_for_search`. The name is at most 255 + /// bytes long, with no NUL byte. SQLite resolves it case-insensitively. + pub name: String, + + /// The exact number of arguments the function takes. Two registrations can share a + /// name when their arities differ. SQLite overloads a name this way. + pub arity: i32, + + /// Whether SQLite can treat the result as stable for the same inputs. SQLite uses this + /// to cache the result and reorder calls. Set this field only for a handler that reads + /// nothing but its arguments. + pub deterministic: bool, + + /// Where SQLite accepts a call to this function. [`InvocationScope::DirectOnly`] is + /// the safe choice, and the one a caller with no view, trigger, or index to serve + /// wants. + pub invocation_scope: InvocationScope, + + /// The closure that computes the result. A returned error fails the statement and + /// reports its message to the caller. If the application unwinds on panic (the + /// default), a panic inside the closure becomes an error naming the function. Under + /// `panic = "abort"`, the process aborts, and nothing here can prevent that. + pub handler: ScalarHandler, +} + +/// One registered function, with everything the FFI layer needs precomputed. +pub(crate) struct Registration { + /// Kept beside `c_name` for the duplicate check and for error messages, so neither + /// path must convert a `CString` back to a `str`. + name: String, + c_name: CString, + arity: i32, + flags: c_int, + handler: ScalarHandler, +} + +/// Register a scalar function on every connection the crate opens for databases that +/// connect from now on. +/// +/// See the [module documentation](self) for the full contract. This function validates +/// the name and the arity, and asks the linked SQLite library to accept the function on a +/// throwaway connection. A bad registration fails here, not later as a connection error. +/// Registration is process-global: the function applies to every database whose `connect` +/// runs after this call, and no caller can remove it. +/// +/// # Errors +/// +/// Returns: +/// +/// - [`Error::InvalidFunctionName`] for an empty name, a name over 255 bytes, or a name +/// holding an interior NUL byte. +/// - [`Error::InvalidFunctionArity`] for a negative arity. SQLite reads a negative arity +/// as variadic, and this crate does not offer variadic functions. +/// - [`Error::DuplicateFunction`] when a registration already has the same name and +/// arity. This crate compares names case-insensitively, to match how SQLite resolves +/// them. Without this check, a second registration replaces the first instead of the +/// two coexisting. +/// - [`Error::ShadowsBuiltinFunction`] when a SQLite built-in already uses the name with +/// this argument count, or variadically. SQLite silently replaces a built-in with an +/// application-defined function, so every statement on the connection that names the +/// built-in reaches the handler instead, including one inside a schema object this +/// application did not write. +/// - [`Error::FunctionRefused`] when the linked library refuses the function for a reason +/// of its own, such as an arity above its `SQLITE_MAX_FUNCTION_ARG`. +/// - [`Error::ProbeConnectionFailed`] when the in-memory connection that carries these +/// checks fails to open or query. Nothing judged the function in that case. +/// +pub fn register_function(function: ScalarFunction) -> Result<()> { + store(function, OnDuplicate::Reject) +} + +/// Register a scalar function, replacing the one already registered under the same name +/// and arity. +/// +/// Use this where a repeat registration is the caller's intent, such as a setup path that +/// can run twice. [`register_function`] rejects the repeat instead, because a rebuilt +/// handler is a new closure and the registry cannot tell an intended repeat from two call +/// sites that disagree about one name. +/// +/// The replacement reaches the databases that connect after this call. A database that +/// already connected holds the set it captured, and it serves that handler for as long as +/// the database lives. The path registry holds a weak reference, so a path whose last +/// handle drops connects fresh and captures the current set. This is the rule every +/// registration follows: a change applies to the next database to connect, never to part +/// of an open one. +/// +/// # Errors +/// +/// Returns every error [`register_function`] returns except [`Error::DuplicateFunction`], +/// which is the case this function accepts. +pub fn register_or_replace_function(function: ScalarFunction) -> Result<()> { + store(function, OnDuplicate::Replace) +} + +/// What to do with a registration that already holds the name and arity. +enum OnDuplicate { + Reject, + Replace, +} + +/// Validate, probe, and store a registration. The two public entry points differ only in +/// what they do about a name and arity already registered. +fn store(function: ScalarFunction, on_duplicate: OnDuplicate) -> Result<()> { + let ScalarFunction { + name, + arity, + deterministic, + invocation_scope, + handler, + } = function; + + let c_name = validate_name(&name)?; + + if arity < 0 { + return Err(Error::InvalidFunctionArity(arity)); + } + + // The connection is opened as UTF-8, so text crosses the boundary as UTF-8 in both + // directions. + // + // The scope is the caller's, because only the caller knows whether a view or an index + // must call the function. `DirectOnly` is the default the enum documents: a registered + // function applies to every database this crate serves, including one attached from an + // untrusted file whose views and triggers this application did not write. + let mut flags = ffi::SQLITE_UTF8; + + match invocation_scope { + InvocationScope::DirectOnly => flags |= ffi::SQLITE_DIRECTONLY, + InvocationScope::Schema => {} + InvocationScope::InnocuousSchema => flags |= ffi::SQLITE_INNOCUOUS, + } + + if deterministic { + flags |= ffi::SQLITE_DETERMINISTIC; + } + + let registration = Arc::new(Registration { + name, + c_name, + arity, + flags, + handler, + }); + + // Held across the match, the probe, and the store, so two threads registering the same + // name cannot both pass the check. The probe opens a connection of its own, and + // registration belongs to startup, so the wait this adds costs nothing. + let mut registered = functions().write(); + + let matched = registered.iter().position(|existing| { + existing.arity == registration.arity && existing.name.eq_ignore_ascii_case(®istration.name) + }); + + // Before the probe: a rejected duplicate is the caller's own doing, and the probe + // cannot report it. Probing first hides a duplicate behind `FunctionRefused` whenever + // the linked library also rejects the arity. + if matched.is_some() && matches!(on_duplicate, OnDuplicate::Reject) { + return Err(Error::DuplicateFunction { + name: registration.name.clone(), + arity, + }); + } + + probe(®istration)?; + + // A replacement takes the position it replaces, so the set stays in registration order. + match matched { + Some(index) => registered[index] = registration, + None => registered.push(registration), + } + + Ok(()) +} + +/// Validate a function name and convert it to the NUL-terminated form SQLite needs. +/// +/// `CString::new` enforces the no-interior-NUL rule, so the conversion and the last +/// validation rule are the same step. +fn validate_name(name: &str) -> Result { + if name.is_empty() || name.len() > MAX_NAME_BYTES { + return Err(Error::InvalidFunctionName(name.to_string())); + } + + CString::new(name).map_err(|_| Error::InvalidFunctionName(name.to_string())) +} + +/// Check `registration` against the linked SQLite library on a throwaway in-memory +/// connection: the name must not replace a built-in, and the library must accept the +/// registration. +/// +/// The check runs here because a rejection at connect time tells the caller nothing +/// useful. sqlx discards the `after_connect` error. Once the acquire deadline passes, +/// sqlx reports `PoolTimedOut` instead, and that error names neither the function nor the +/// reason. A returned error from this function names both, at the call that registered +/// the function. +/// +/// The name-length rule this crate applies first comes from SQLite's documented API +/// contract. The linked library is the platform's own, built with limits this crate +/// cannot read directly. The probe asks that library directly, using a connection built +/// the same way as every pooled connection. +fn probe(registration: &Arc) -> Result<()> { + let mut db: *mut ffi::sqlite3 = std::ptr::null_mut(); + + // SAFETY: `db` is a valid out pointer, the filename is a NUL-terminated literal, and a + // null VFS name selects the default VFS. + let open_code = unsafe { + ffi::sqlite3_open_v2( + c":memory:".as_ptr(), + &mut db, + ffi::SQLITE_OPEN_READWRITE | ffi::SQLITE_OPEN_CREATE, + std::ptr::null(), + ) + }; + + let outcome = if open_code == ffi::SQLITE_OK { + // SAFETY: `db` points at the connection opened above, which nothing else can reach. + unsafe { probe_on(db, registration) } + } else { + // A failed open means this crate cannot run the checks at all: the platform is out + // of memory, its default VFS forbids the connection, or the library dropped + // `:memory:` support. Naming that a refusal would tell the caller to change a + // function that nothing has judged yet. + Err(Error::ProbeConnectionFailed { + name: registration.name.clone(), + reason: describe(open_code), + }) + }; + + // SAFETY: `db` is either null, which `sqlite3_close` accepts, or the handle from the + // call above. A failed open still produces a handle to close. `probe_on` finalizes the + // one statement it prepares, so the connection has nothing left to finalize. + unsafe { ffi::sqlite3_close(db) }; + + outcome +} + +/// The checks [`probe`] runs once its connection is open. +/// +/// # Safety +/// +/// `db` must point to an open sqlite3 connection nothing else is using. +unsafe fn probe_on(db: *mut ffi::sqlite3, registration: &Arc) -> Result<()> { + // SAFETY: `db` is open and exclusively ours (caller's guarantee). + if unsafe { shadows_builtin(db, registration) }? { + return Err(Error::ShadowsBuiltinFunction { + name: registration.name.clone(), + arity: registration.arity, + }); + } + + // SAFETY: as above. + let create_code = unsafe { create_function(db, registration) }; + + if create_code != ffi::SQLITE_OK { + return Err(Error::FunctionRefused { + name: registration.name.clone(), + arity: registration.arity, + reason: describe(create_code), + }); + } + + Ok(()) +} + +/// Whether the linked library already defines this name as a built-in the registration +/// replaces: a function with the same argument count, or a variadic one. +/// +/// `sqlite3_create_function_v2` silently replaces a built-in, so every statement that +/// names it reaches the handler instead, whatever the registration's +/// [`InvocationScope`]. The probe connection holds nothing but built-ins, so a match here +/// is a built-in. `lower()` folds ASCII only, which matches how SQLite resolves function +/// names. +/// +/// # Safety +/// +/// `db` must point to an open sqlite3 connection nothing else is using. +unsafe fn shadows_builtin(db: *mut ffi::sqlite3, registration: &Registration) -> Result { + // A negative `narg` marks a variadic built-in (the exact negative value encodes the + // minimum argument count in newer SQLite versions, such as -3 for `max`). An + // exact-arity application function takes precedence over a variadic built-in for calls + // with that arity, so a variadic match is a replacement for those calls too. + const SQL: &CStr = c"SELECT 1 FROM pragma_function_list \ + WHERE lower(name) = lower(?1) AND (narg = ?2 OR narg < 0) LIMIT 1"; + + let probe_failure = |code: c_int| Error::ProbeConnectionFailed { + name: registration.name.clone(), + reason: describe(code), + }; + + let mut stmt: *mut ffi::sqlite3_stmt = std::ptr::null_mut(); + + // SAFETY: `db` is open (caller's guarantee), the SQL is NUL-terminated, and `stmt` is + // a valid out pointer. + let prepare_code = + unsafe { ffi::sqlite3_prepare_v2(db, SQL.as_ptr(), -1, &mut stmt, std::ptr::null_mut()) }; + + if prepare_code != ffi::SQLITE_OK { + // A failed prepare leaves `stmt` null, so there is nothing to finalize. + return Err(probe_failure(prepare_code)); + } + + // SAFETY: `stmt` is the statement prepared above. `c_name` outlives it, so + // `SQLITE_STATIC` applies. Its length fits in `c_int` because `validate_name` caps + // names at 255 bytes. + let mut code = unsafe { + ffi::sqlite3_bind_text( + stmt, + 1, + registration.c_name.as_ptr(), + registration.c_name.as_bytes().len() as c_int, + ffi::SQLITE_STATIC(), + ) + }; + + if code == ffi::SQLITE_OK { + // SAFETY: as above. + code = unsafe { ffi::sqlite3_bind_int(stmt, 2, registration.arity as c_int) }; + } + + let step_code = if code == ffi::SQLITE_OK { + // SAFETY: `stmt` is the prepared statement with both parameters bound. + unsafe { ffi::sqlite3_step(stmt) } + } else { + code + }; + + // SAFETY: `stmt` came from the successful prepare above. + unsafe { ffi::sqlite3_finalize(stmt) }; + + match step_code { + ffi::SQLITE_ROW => Ok(true), + ffi::SQLITE_DONE => Ok(false), + code => Err(probe_failure(code)), + } +} + +/// SQLite's own description of a result code, followed by the numeric code. +fn describe(code: c_int) -> String { + // SAFETY: sqlite3_errstr returns a static, NUL-terminated string for any code, + // including an unrecognized one. + let text = unsafe { CStr::from_ptr(ffi::sqlite3_errstr(code)) }.to_string_lossy(); + + format!("{text} (code {code})") +} + +/// The boxed future sqlx's `after_connect` expects, borrowing the connection it +/// configures. +type HookFuture<'c> = + std::pin::Pin> + Send + 'c>>; + +/// Build the `after_connect` hook that installs `functions` on each connection a pool +/// opens. Both of a database's pools use one snapshot, so their connections stay uniform. +pub(crate) fn install_hook( + functions: &FunctionSet, +) -> impl for<'c> Fn(&'c mut SqliteConnection, sqlx::pool::PoolConnectionMetadata) -> HookFuture<'c> ++ Send ++ Sync ++ 'static { + let functions = Arc::clone(functions); + + move |conn, _meta| { + let functions = Arc::clone(&functions); + Box::pin(async move { apply_all(&functions, conn).await }) + } +} + +/// Install `functions` on a newly opened connection. +/// +/// [`install_hook`] runs this from the pools' `after_connect` hook, before sqlx hands the +/// connection to a caller, with the set the database captured at connect time. If this +/// function returns an error, sqlx closes the connection, so no caller ever gets one with +/// only some functions installed. sqlx then retries the connect until the acquire +/// deadline passes, and hands the caller `sqlx::Error::PoolTimedOut` instead, without the +/// error this function returned. [`register_function`] probes every function against the +/// same library before accepting it. A failure here means the connection ran out of +/// memory, not that the library refused a function. +async fn apply_all( + functions: &[Arc], + conn: &mut SqliteConnection, +) -> std::result::Result<(), sqlx::Error> { + // An empty set never calls `lock_handle()` and never touches FFI. + if functions.is_empty() { + return Ok(()); + } + + let mut handle = conn.lock_handle().await?; + let db = handle.as_raw_handle().as_ptr(); + + for registration in functions { + // SAFETY: `db` comes from the handle we hold the lock on, so it points at an open + // connection nothing else is using for the duration of this call. + let code = unsafe { create_function(db, registration) }; + + if code != ffi::SQLITE_OK { + let message = format!( + "failed to register SQL function '{}' taking {} argument(s): {}", + registration.name, + registration.arity, + describe(code) + ); + + error!("{message}"); + + return Err(sqlx::Error::Configuration(message.into())); + } + } + + trace!( + count = functions.len(), + "Registered scalar functions on a new connection" + ); + + Ok(()) +} + +/// Register one function on a raw connection handle, returning SQLite's result code. +/// +/// # Safety +/// +/// `db` must point to an open sqlite3 connection, and the caller must hold exclusive +/// access to that connection for the duration of the call. +unsafe fn create_function(db: *mut ffi::sqlite3, registration: &Arc) -> c_int { + // One strong reference per (connection, function) pair. SQLite hands the pointer back + // to `dispatch` as user data, and passes it to `release_registration` when the + // function is deleted: that happens when the connection closes, or when the same name + // and arity is registered again on it. Both pools drop idle connections and open + // fresh ones, so the reference must be scoped to the connection: one that outlives its + // connection leaks, once per connection that replaces it. + let user_data = Arc::into_raw(Arc::clone(registration)) as *mut c_void; + + // SAFETY: `db` is an open connection (the caller's guarantee). `c_name` is + // NUL-terminated, and it outlives the registration because the Arc above holds it. + // `user_data` stays valid until `release_registration` runs. + unsafe { + ffi::sqlite3_create_function_v2( + db, + registration.c_name.as_ptr(), + registration.arity as c_int, + registration.flags, + user_data, + Some(dispatch), + // xStep and xFinal belong to an aggregate, not a scalar function. An aggregate + // takes this call with xFunc unset, so it cannot reuse `dispatch`. + None, + None, + Some(release_registration), + ) + } + + // No manual release on the failure path: SQLite invokes xDestroy even when this call + // fails, so a manual release here duplicates the drop. +} + +/// The `xFunc` every registration shares. Invokes the handler and writes its outcome into +/// the SQLite call context. +/// +/// # Safety +/// +/// SQLite calls this with the context and argument vector of a live function invocation, +/// on the thread that owns the connection. +unsafe extern "C" fn dispatch( + ctx: *mut ffi::sqlite3_context, + argc: c_int, + argv: *mut *mut ffi::sqlite3_value, +) { + // SAFETY: `ctx` belongs to a live invocation, so the user data is the pointer + // `create_function` handed SQLite for this function. + let user_data = unsafe { ffi::sqlite3_user_data(ctx) }; + + if user_data.is_null() || argv.is_null() { + return; + } + + // A borrow, not an owned Arc. SQLite keeps the reference alive until + // `release_registration` runs. A connection runs one statement at a time on its own + // thread, so this call cannot race the release. + // + // SAFETY: the pointer came from `Arc::into_raw` on an `Arc` and the + // strong count SQLite holds has not been released yet. + let registration = unsafe { &*(user_data as *const Registration) }; + + // An unwind out of an `extern "C"` function aborts the process, so everything that can + // panic runs inside here. That includes argument conversion, because a lossy decode of + // invalid UTF-8 allocates. + let outcome = catch_unwind(AssertUnwindSafe(|| { + // SAFETY: SQLite guarantees `argv` holds `argc` valid value pointers for the + // duration of this call, and the handler runs within it. The null check above is + // what keeps this sound at an `argc` of zero, where `from_raw_parts` still requires + // a non-null pointer. + let argv = unsafe { std::slice::from_raw_parts(argv, argc.max(0) as usize) }; + + let mut args = Vec::with_capacity(argv.len()); + + // Each argument borrows from `argv`, which is local to this closure, so the compiler + // keeps every borrow inside the invocation. + for value in argv { + // SAFETY: `value` points at one of the pointers SQLite passed, valid for this + // call. + args.push(unsafe { SqlValueRef::from_raw(value) }); + } + + (registration.handler)(&args) + })); + + match outcome { + // SAFETY for all three arms: `ctx` is the context of the invocation still in + // progress. These calls require nothing else. + Ok(Ok(value)) => unsafe { value.write_to(ctx) }, + Ok(Err(error)) => unsafe { result_error(ctx, error.message()) }, + Err(_) => unsafe { + // Allocating after a caught panic is safe here: Rust aborts on allocation + // failure rather than unwinding, so `format!` has no panic path of its own. + result_error(ctx, &format!("{}: handler panicked", registration.name)) + }, + } +} + +/// Release the strong reference [`create_function`] handed to SQLite. +/// +/// # Safety +/// +/// SQLite calls this once per `sqlite3_create_function_v2` call, successful or not, with +/// the `user_data` pointer from that call. +unsafe extern "C" fn release_registration(user_data: *mut c_void) { + if user_data.is_null() { + return; + } + + // Dropping the last reference drops the consumer's handler closure. This call is + // guarded because an unguarded panic in that closure's `Drop` unwinds out of this + // `extern "C"` function and aborts the process. + let _ = catch_unwind(AssertUnwindSafe(|| { + // SAFETY: the pointer came from `Arc::into_raw` in `create_function`, and SQLite + // hands each pointer back exactly once. + unsafe { Arc::decrement_strong_count(user_data as *const Registration) }; + })); +} + +/// Report `message` as the error for the function call in progress. +/// +/// # Safety +/// +/// `ctx` must be the context of a live function invocation. +unsafe fn result_error(ctx: *mut ffi::sqlite3_context, message: &str) { + // SQLite copies the message out, so pointing at Rust-owned memory is safe. A message + // longer than `i32::MAX` cannot fit in the length SQLite expects, so this function + // truncates it instead of dropping it. Nothing produces a message that long in + // practice. + let len = c_int::try_from(message.len()).unwrap_or(c_int::MAX); + + // SAFETY: `ctx` belongs to a live invocation, and `message` is valid for `len` bytes. + unsafe { ffi::sqlite3_result_error(ctx, message.as_ptr() as *const c_char, len) }; +} + +impl<'a> SqlValueRef<'a> { + /// Read a SQLite argument as a borrowed value pointing into SQLite's own buffers. + /// + /// Text is decoded lossily rather than rejected, to remove any risk of a panic + /// crossing the FFI boundary. The connection is UTF-8, so invalid UTF-8 does not occur + /// here in practice, and the decode borrows rather than allocates. + /// + /// # Safety + /// + /// `value` must point at a valid `sqlite3_value` pointer. The returned borrow lives as + /// long as that reference, so a caller holding the reference no longer than the + /// invocation cannot let the borrow escape it. + unsafe fn from_raw(value: &'a *mut ffi::sqlite3_value) -> SqlValueRef<'a> { + let value = *value; + + if value.is_null() { + return SqlValueRef::Null; + } + + // SAFETY for every call below: `value` is valid until the invocation returns + // (caller's guarantee). This code calls `sqlite3_value_bytes` after the matching + // `_text`/`_blob` accessor, in the order SQLite documents. + match unsafe { ffi::sqlite3_value_type(value) } { + ffi::SQLITE_INTEGER => SqlValueRef::Integer(unsafe { ffi::sqlite3_value_int64(value) }), + ffi::SQLITE_FLOAT => SqlValueRef::Real(unsafe { ffi::sqlite3_value_double(value) }), + ffi::SQLITE_TEXT => { + let ptr = unsafe { ffi::sqlite3_value_text(value) }; + let len = unsafe { ffi::sqlite3_value_bytes(value) }; + + if ptr.is_null() || len <= 0 { + return SqlValueRef::Text(Cow::Borrowed("")); + } + + // SAFETY: SQLite reports `len` readable bytes at `ptr`. + let bytes = unsafe { std::slice::from_raw_parts(ptr, len as usize) }; + SqlValueRef::Text(String::from_utf8_lossy(bytes)) + } + ffi::SQLITE_BLOB => { + let ptr = unsafe { ffi::sqlite3_value_blob(value) }; + let len = unsafe { ffi::sqlite3_value_bytes(value) }; + + if ptr.is_null() || len <= 0 { + return SqlValueRef::Blob(&[]); + } + + // SAFETY: SQLite reports `len` readable bytes at `ptr`. + SqlValueRef::Blob(unsafe { std::slice::from_raw_parts(ptr as *const u8, len as usize) }) + } + // SQLITE_NULL, and any type a future SQLite adds. + _ => SqlValueRef::Null, + } + } +} + +impl SqlValue { + /// Write this value as the result of the function call in progress. + /// + /// This function hands text and blob results over with `SQLITE_TRANSIENT`, so SQLite + /// copies them before this value is dropped. A text or blob result over `i32::MAX` + /// bytes fails the statement with `SQLITE_TOOBIG`. + /// + /// # Safety + /// + /// `ctx` must be the context of a live function invocation. + unsafe fn write_to(self, ctx: *mut ffi::sqlite3_context) { + // SAFETY for every call below: `ctx` belongs to a live invocation, and each pointer + // passed is valid for the length passed with it. + match self { + SqlValue::Null => unsafe { ffi::sqlite3_result_null(ctx) }, + SqlValue::Integer(value) => unsafe { ffi::sqlite3_result_int64(ctx, value) }, + SqlValue::Real(value) => unsafe { ffi::sqlite3_result_double(ctx, value) }, + SqlValue::Text(text) => match c_int::try_from(text.len()) { + Ok(len) => unsafe { + ffi::sqlite3_result_text( + ctx, + text.as_ptr() as *const c_char, + len, + ffi::SQLITE_TRANSIENT(), + ) + }, + // A value this large exceeds what SQLite holds. `SQLITE_MAX_LENGTH` bounds + // it far lower, and the 64-bit result setters reject the same bound. + Err(_) => unsafe { ffi::sqlite3_result_error_toobig(ctx) }, + }, + SqlValue::Blob(bytes) => match c_int::try_from(bytes.len()) { + Ok(len) => unsafe { + ffi::sqlite3_result_blob( + ctx, + bytes.as_ptr() as *const c_void, + len, + ffi::SQLITE_TRANSIENT(), + ) + }, + Err(_) => unsafe { ffi::sqlite3_result_error_toobig(ctx) }, + }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Every test in this binary shares the process-global registry, so each one uses a + /// name of its own. + fn function(name: &str, arity: i32) -> ScalarFunction { + ScalarFunction { + name: name.to_string(), + arity, + deterministic: true, + invocation_scope: InvocationScope::DirectOnly, + handler: Arc::new(|_args| Ok(SqlValue::Null)), + } + } + + #[test] + fn rejects_an_empty_name() { + assert!(matches!( + register_function(function("", 1)), + Err(Error::InvalidFunctionName(_)) + )); + } + + #[test] + fn rejects_a_name_over_255_bytes() { + assert!(matches!( + register_function(function(&"x".repeat(256), 1)), + Err(Error::InvalidFunctionName(_)) + )); + } + + #[test] + fn rejects_a_name_holding_a_nul_byte() { + assert!(matches!( + register_function(function("unit_n\0ul", 1)), + Err(Error::InvalidFunctionName(_)) + )); + } + + #[test] + fn rejects_a_negative_arity() { + assert!(matches!( + register_function(function("unit_variadic", -1)), + Err(Error::InvalidFunctionArity(-1)) + )); + } + + /// The linked library, not this crate, decides the arity ceiling. The bundled SQLite + /// caps `SQLITE_MAX_FUNCTION_ARG` at 1000, so 1001 reaches the probe and the library + /// refuses it there. + #[test] + fn refuses_an_arity_above_what_the_library_accepts() { + assert!(matches!( + register_function(function("unit_too_many_args", 1001)), + Err(Error::FunctionRefused { arity: 1001, .. }) + )); + } + + #[test] + fn refuses_a_name_that_replaces_a_builtin() { + // Exact-arity match against a built-in: lower() takes one argument. + assert!(matches!( + register_function(function("lower", 1)), + Err(Error::ShadowsBuiltinFunction { arity: 1, .. }) + )); + + // Case-insensitive, matching how SQLite resolves names. + assert!(matches!( + register_function(function("LOWER", 1)), + Err(Error::ShadowsBuiltinFunction { arity: 1, .. }) + )); + + // A variadic built-in matches every arity: max() accepts any argument count. + assert!(matches!( + register_function(function("max", 3)), + Err(Error::ShadowsBuiltinFunction { arity: 3, .. }) + )); + } + + /// A consumer that builds a name at runtime, such as one carrying a version suffix, + /// hands over the `String` it built. + #[test] + fn accepts_a_name_built_at_runtime() { + let name = format!("unit_runtime_{}", 3 + 4); + + register_function(function(&name, 1)).expect("registration"); + } + + #[test] + fn rejects_a_duplicate_name_and_arity_whatever_the_case() { + register_function(function("unit_dup", 1)).expect("first registration"); + + assert!(matches!( + register_function(function("unit_dup", 1)), + Err(Error::DuplicateFunction { .. }) + )); + + assert!(matches!( + register_function(function("UNIT_DUP", 1)), + Err(Error::DuplicateFunction { .. }) + )); + + // A different arity is an overload. SQLite supports overloading a name this way, so + // this registers. + register_function(function("unit_dup", 2)).expect("overload by arity"); + } + + #[test] + fn replaces_a_name_and_arity_the_plain_registration_rejects() { + register_function(function("unit_replace", 1)).expect("first registration"); + + assert!(matches!( + register_function(function("unit_replace", 1)), + Err(Error::DuplicateFunction { .. }) + )); + + register_or_replace_function(function("unit_replace", 1)).expect("replacement"); + + // Case-insensitively, matching how SQLite resolves names and how the rejection above + // compares them. + register_or_replace_function(function("UNIT_REPLACE", 1)).expect("replacement"); + } +} diff --git a/crates/sqlx-sqlite-conn-mgr/src/lib.rs b/crates/sqlx-sqlite-conn-mgr/src/lib.rs index 37a3d1a..41c0fb5 100644 --- a/crates/sqlx-sqlite-conn-mgr/src/lib.rs +++ b/crates/sqlx-sqlite-conn-mgr/src/lib.rs @@ -8,6 +8,7 @@ //! - **[`SqliteDatabase`]**: Main database type with separate read and write connection pools //! - **[`SqliteDatabaseConfig`]**: Configuration for connection pool settings //! - **[`WriteGuard`]**: RAII guard ensuring exclusive write access +//! - **[`ScalarFunction`]**: A scalar SQL function to register on every connection //! - **[`Migrator`]**: Re-exported from sqlx for running database migrations //! - **[`Error`]**: Error type for database operations //! @@ -17,6 +18,8 @@ //! - **Lazy WAL mode**: Write-Ahead Logging enabled automatically on first write //! - **Exclusive writes**: Single-connection write pool enforces serialized write access //! - **Concurrent reads**: Multiple readers can query simultaneously via the read pool +//! - **Scalar functions**: [`register_function`] applies a consumer's Rust function to +//! every connection of every database whose `connect` runs after the registration //! //! ## Usage //! @@ -64,6 +67,7 @@ mod attached; mod config; mod database; mod error; +pub mod functions; mod observer_slot; mod registry; mod write_guard; @@ -76,6 +80,10 @@ pub use attached::{ pub use config::SqliteDatabaseConfig; pub use database::SqliteDatabase; pub use error::Error; +pub use functions::{ + FunctionError, InvocationScope, ScalarFunction, ScalarHandler, SqlValue, SqlValueRef, + register_function, register_or_replace_function, +}; pub use observer_slot::ObserverSlot; pub use write_guard::WriteGuard; From 7e1665854b3e16cebdbe9ae471d2ea090181f304 Mon Sep 17 00:00:00 2001 From: Ethan Smith Date: Wed, 19 Aug 2026 20:15:44 +0000 Subject: [PATCH 2/5] test: cover scalar functions across the connection lifecycle Asserts what a consumer depends on: a registered function resolves on a read connection and on the write connection, on every connection a growing read pool opens, and on the connection that replaces one the idle reaper dropped. Also covers the two failure paths. A handler error carries its message to the caller, and a handler panic produces an error naming the function while leaving the pool usable. --- .../tests/function_tests.rs | 498 ++++++++++++++++++ 1 file changed, 498 insertions(+) create mode 100644 crates/sqlx-sqlite-conn-mgr/tests/function_tests.rs diff --git a/crates/sqlx-sqlite-conn-mgr/tests/function_tests.rs b/crates/sqlx-sqlite-conn-mgr/tests/function_tests.rs new file mode 100644 index 0000000..c21da71 --- /dev/null +++ b/crates/sqlx-sqlite-conn-mgr/tests/function_tests.rs @@ -0,0 +1,498 @@ +//! A scalar function applies to every connection that either pool opens. +//! +//! Every test in this binary shares the process-global function registry, and each +//! database captures the registered set at its `connect`. `REGISTERED` therefore holds +//! every function this binary needs, under a name per test, and `open` forces it before +//! the first database connects. + +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; + +use sqlx::sqlite::{Sqlite, SqliteRow}; +use sqlx::{Executor, FromRow}; +use sqlx_sqlite_conn_mgr::{ + FunctionError, InvocationScope, ScalarFunction, ScalarHandler, SqlValue, SqlValueRef, + SqliteDatabase, SqliteDatabaseConfig, register_function, register_or_replace_function, +}; +use tempfile::TempDir; + +/// A handler that returns whichever value it received. +fn echo() -> ScalarHandler { + Arc::new(|args| Ok(args[0].to_owned())) +} + +fn register(name: &str, arity: i32, handler: ScalarHandler) { + register_function(ScalarFunction { + name: name.to_string(), + arity, + deterministic: true, + invocation_scope: InvocationScope::DirectOnly, + handler, + }) + .expect("registration"); +} + +/// Every function this binary uses, one per test, registered before any database opens. +static REGISTERED: LazyLock<()> = LazyLock::new(|| { + register( + "fn_read_write", + 1, + Arc::new(|args| match &args[0] { + SqlValueRef::Text(text) => Ok(SqlValue::Text(format!("{text}!"))), + _ => Err(FunctionError::new("expected text")), + }), + ); + register("fn_pool_growth", 1, echo()); + register("fn_after_idle", 1, echo()); + register( + "fn_handler_error", + 1, + Arc::new(|_args| Err(FunctionError::new("payload is not decodable"))), + ); + register("fn_panics", 1, Arc::new(|_args| panic!("handler exploded"))); + register("fn_round_trip", 1, echo()); + register("fn_direct_only", 1, echo()); + register( + "fn_no_args", + 0, + Arc::new(|_args| Ok(SqlValue::Text("no arguments".into()))), + ); + register( + "fn_pair", + 2, + Arc::new(|args| match (&args[0], &args[1]) { + (SqlValueRef::Text(first), SqlValueRef::Text(second)) => { + Ok(SqlValue::Text(format!("{first}|{second}"))) + } + _ => Err(FunctionError::new("expected two text arguments")), + }), + ); + + // One name at two arities. Each handler names its own arity, so a call proves which + // registration SQLite resolved. + register( + "fn_over", + 1, + Arc::new(|_args| Ok(SqlValue::Text("one argument".into()))), + ); + register( + "fn_over", + 2, + Arc::new(|_args| Ok(SqlValue::Text("two arguments".into()))), + ); + + register_function(ScalarFunction { + name: "fn_in_schema".into(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::Schema, + handler: echo(), + }) + .expect("registration"); + + register_function(ScalarFunction { + name: "fn_trusted".into(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::Schema, + handler: echo(), + }) + .expect("registration"); + + register_function(ScalarFunction { + name: "fn_innocuous".into(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::InnocuousSchema, + handler: echo(), + }) + .expect("registration"); +}); + +async fn open(dir: &TempDir, config: Option) -> Arc { + LazyLock::force(®ISTERED); + + SqliteDatabase::connect(dir.path().join("test.db"), config) + .await + .expect("connect") +} + +/// Runs `sql` on any executor and returns its one column, panicking with the SQL itself. +async fn scalar<'e, 'c: 'e, T, E>(executor: E, sql: &'static str) -> T +where + T: Send + Unpin, + (T,): Send + Unpin + for<'r> FromRow<'r, SqliteRow>, + E: 'e + Executor<'c, Database = Sqlite>, +{ + sqlx::query_scalar(sql) + .fetch_one(executor) + .await + .expect(sql) +} + +/// Runs `sql` expecting it to fail, and returns the error it failed with. +async fn scalar_error<'e, 'c: 'e, E>(executor: E, sql: &'static str) -> sqlx::Error +where + E: 'e + Executor<'c, Database = Sqlite>, +{ + sqlx::query_scalar::<_, String>(sql) + .fetch_one(executor) + .await + .expect_err(sql) +} + +#[tokio::test] +async fn resolves_on_a_read_connection_and_on_the_write_connection() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + + let from_reader: String = scalar(db.read_pool().unwrap(), "SELECT fn_read_write('read')").await; + assert_eq!(from_reader, "read!"); + + let mut writer = db.acquire_writer().await.unwrap(); + let from_writer: String = scalar(&mut *writer, "SELECT fn_read_write('write')").await; + assert_eq!(from_writer, "write!"); +} + +#[tokio::test] +async fn resolves_on_every_connection_once_the_read_pool_grows() { + let dir = TempDir::new().unwrap(); + let config = SqliteDatabaseConfig { + max_read_connections: 3, + ..Default::default() + }; + let db = open(&dir, Some(config)).await; + + // The test holds each connection until the loop ends. The pool must therefore open a + // new connection each iteration, rather than hand back one already open. Every query + // below runs on a newly opened connection. + let mut held = Vec::new(); + + for _ in 0..3 { + let mut conn = db.read_pool().unwrap().acquire().await.expect("acquire"); + + let value: i64 = scalar(&mut *conn, "SELECT fn_pool_growth(7)").await; + assert_eq!(value, 7); + + held.push(conn); + } + + assert_eq!(db.read_pool().unwrap().size(), 3); +} + +#[tokio::test] +async fn resolves_on_a_connection_that_replaces_an_idle_one() { + let dir = TempDir::new().unwrap(); + let config = SqliteDatabaseConfig { + idle_timeout_secs: 1, + ..Default::default() + }; + let db = open(&dir, Some(config)).await; + let pool = db.read_pool().unwrap(); + + // `connect_with` opens and tests one connection, so the pool starts with one. + assert_eq!(pool.size(), 1); + + // Wait for sqlx's reaper to drop it. The reaper's period is the idle timeout, so this + // asserts on the pool being empty rather than on elapsed time. + let deadline = Instant::now() + Duration::from_secs(10); + while pool.size() > 0 && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert_eq!(pool.size(), 0, "the idle connection was never reaped"); + + let value: String = scalar(pool, "SELECT fn_after_idle('fresh')").await; + assert_eq!(value, "fresh"); +} + +#[tokio::test] +async fn a_handler_error_fails_the_statement_and_carries_its_message() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + + let error = scalar_error(db.read_pool().unwrap(), "SELECT fn_handler_error('x')").await; + + assert!( + error.to_string().contains("payload is not decodable"), + "error did not include the handler's message: {error}" + ); +} + +#[tokio::test] +async fn a_panicking_handler_produces_an_error_and_leaves_the_pool_usable() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + let pool = db.read_pool().unwrap(); + + // The panic hook still prints to stderr, so this test is noisy by design. + let error = scalar_error(pool, "SELECT fn_panics('x')").await; + + assert!( + error.to_string().contains("fn_panics: handler panicked"), + "error did not name the panicking function: {error}" + ); + + // The pool serves queries after a handler panic, rather than holding a poisoned + // connection. + let survivor: i64 = scalar(pool, "SELECT 1").await; + assert_eq!(survivor, 1); +} + +#[tokio::test] +async fn every_storage_class_round_trips_and_null_passes_through() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + let pool = db.read_pool().unwrap(); + + let null: Option = scalar(pool, "SELECT fn_round_trip(NULL)").await; + assert_eq!(null, None); + + let integer: i64 = scalar(pool, "SELECT fn_round_trip(42)").await; + assert_eq!(integer, 42); + + let real: f64 = scalar(pool, "SELECT fn_round_trip(1.5)").await; + assert_eq!(real, 1.5); + + let text: String = scalar(pool, "SELECT fn_round_trip('abc')").await; + assert_eq!(text, "abc"); + + let blob: Vec = scalar(pool, "SELECT fn_round_trip(x'0102ff')").await; + assert_eq!(blob, vec![1, 2, 255]); + + // The storage class survives the round trip, rather than everything arriving as text. + // sqlx coerces a text value into an i64, so the typed reads above pass either way. + let classes: Vec = sqlx::query_scalar( + "SELECT typeof(fn_round_trip(42)) \ + UNION ALL SELECT typeof(fn_round_trip(1.5)) \ + UNION ALL SELECT typeof(fn_round_trip('abc')) \ + UNION ALL SELECT typeof(fn_round_trip(x'01'))", + ) + .fetch_all(pool) + .await + .unwrap(); + assert_eq!(classes, vec!["integer", "real", "text", "blob"]); +} + +/// A zero-argument function drives `dispatch` with an empty argument vector, which no +/// other test in this binary reaches. +#[tokio::test] +async fn a_zero_argument_function_resolves() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + + let value: String = scalar(db.read_pool().unwrap(), "SELECT fn_no_args()").await; + assert_eq!(value, "no arguments"); +} + +/// Two calls with the arguments swapped. One call cannot tell a correct argument order +/// from a reversed one. +#[tokio::test] +async fn a_two_argument_function_receives_its_arguments_in_order() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + let pool = db.read_pool().unwrap(); + + let forward: String = scalar(pool, "SELECT fn_pair('a', 'b')").await; + assert_eq!(forward, "a|b"); + + let reversed: String = scalar(pool, "SELECT fn_pair('b', 'a')").await; + assert_eq!(reversed, "b|a"); +} + +/// SQLite overloads a name by argument count, and the registry holds one entry per name +/// and arity pair. Each call reaches the handler registered for its own arity. +#[tokio::test] +async fn one_name_registered_at_two_arities_resolves_at_both() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + let pool = db.read_pool().unwrap(); + + let one: String = scalar(pool, "SELECT fn_over('x')").await; + assert_eq!(one, "one argument"); + + let two: String = scalar(pool, "SELECT fn_over('x', 'y')").await; + assert_eq!(two, "two arguments"); +} + +#[tokio::test] +async fn a_registered_function_runs_from_top_level_sql_but_not_from_a_schema_object() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + let mut writer = db.acquire_writer().await.unwrap(); + + let direct: String = scalar(&mut *writer, "SELECT fn_direct_only('x')").await; + assert_eq!(direct, "x"); + + sqlx::query("CREATE TABLE doc (title TEXT)") + .execute(&mut *writer) + .await + .unwrap(); + sqlx::query("INSERT INTO doc (title) VALUES ('x')") + .execute(&mut *writer) + .await + .unwrap(); + sqlx::query("CREATE VIEW doc_view AS SELECT fn_direct_only(title) AS title FROM doc") + .execute(&mut *writer) + .await + .unwrap(); + + // A view in an attached file can name a registered function, so SQLite must refuse the + // call rather than run a handler this application never meant to expose there. + let from_view = scalar_error(&mut *writer, "SELECT title FROM doc_view").await; + assert!( + from_view + .to_string() + .contains("unsafe use of fn_direct_only"), + "a view calling the function was not refused: {from_view}" + ); + + // An expression index would write the function name into the schema, leaving the table + // unreadable by any process that opens the file without registering the function. + let index = sqlx::query("CREATE INDEX doc_title ON doc (fn_direct_only(title))") + .execute(&mut *writer) + .await + .expect_err("an expression index naming the function was accepted"); + assert!( + index.to_string().contains("unsafe use of fn_direct_only"), + "an expression index naming the function was not refused: {index}" + ); +} + +/// A replacement follows the same rule a first registration does: it reaches the databases +/// that connect after it, and the database already open keeps the handler it captured. +#[tokio::test] +async fn a_replacement_reaches_the_next_database_and_leaves_an_open_one_alone() { + let replaced = |result: &'static str| ScalarFunction { + name: "fn_replaced".into(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::DirectOnly, + handler: Arc::new(move |_args| Ok(SqlValue::Text(result.to_string()))), + }; + + register_function(replaced("first")).expect("registration"); + + let earlier_dir = TempDir::new().unwrap(); + let earlier = SqliteDatabase::connect(earlier_dir.path().join("earlier.db"), None) + .await + .expect("connect"); + + register_or_replace_function(replaced("second")).expect("replacement"); + + let later_dir = TempDir::new().unwrap(); + let later = SqliteDatabase::connect(later_dir.path().join("later.db"), None) + .await + .expect("connect"); + + let from_earlier: String = scalar(earlier.read_pool().unwrap(), "SELECT fn_replaced('x')").await; + assert_eq!(from_earlier, "first"); + + let from_later: String = scalar(later.read_pool().unwrap(), "SELECT fn_replaced('x')").await; + assert_eq!(from_later, "second"); +} + +/// `InvocationScope::Schema` is what a consumer picks to build a view over a function. +#[tokio::test] +async fn a_schema_scoped_function_runs_from_inside_a_view() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + let mut writer = db.acquire_writer().await.unwrap(); + + sqlx::query("CREATE TABLE note (body TEXT)") + .execute(&mut *writer) + .await + .unwrap(); + sqlx::query("INSERT INTO note (body) VALUES ('kept')") + .execute(&mut *writer) + .await + .unwrap(); + sqlx::query("CREATE VIEW note_view AS SELECT fn_in_schema(body) AS body FROM note") + .execute(&mut *writer) + .await + .unwrap(); + + let through_view: String = scalar(&mut *writer, "SELECT body FROM note_view").await; + assert_eq!(through_view, "kept"); + + // The view belongs to the schema, so a read connection resolves the function too. + let from_reader: String = scalar(db.read_pool().unwrap(), "SELECT body FROM note_view").await; + assert_eq!(from_reader, "kept"); +} + +/// `SQLITE_INNOCUOUS` is what separates [`InvocationScope::InnocuousSchema`] from +/// [`InvocationScope::Schema`]. Both run from a view on a connection that trusts the +/// schema, so the test that tells them apart needs a connection that does not. +/// `fn_trusted` is the control: without the flag, SQLite refuses the same call. +#[tokio::test] +async fn only_an_innocuous_function_runs_from_a_view_when_the_schema_is_untrusted() { + let dir = TempDir::new().unwrap(); + let db = open(&dir, None).await; + + // The views are created while the schema is still trusted, and the writer is released + // before the reads below. + { + let mut writer = db.acquire_writer().await.unwrap(); + + for sql in [ + "CREATE TABLE memo (body TEXT)", + "INSERT INTO memo (body) VALUES ('held')", + "CREATE VIEW trusted_view AS SELECT fn_trusted(body) AS body FROM memo", + "CREATE VIEW innocuous_view AS SELECT fn_innocuous(body) AS body FROM memo", + ] { + sqlx::query(sql).execute(&mut *writer).await.expect(sql); + } + } + + // `PRAGMA trusted_schema` belongs to one connection, and the read pool hands out + // whichever connection it has. So the test holds the connection it set the pragma on + // for both reads. + let mut conn = db.read_pool().unwrap().acquire().await.expect("acquire"); + sqlx::query("PRAGMA trusted_schema = OFF") + .execute(&mut *conn) + .await + .expect("trusted_schema = OFF"); + + let refused = scalar_error(&mut *conn, "SELECT body FROM trusted_view").await; + assert!( + refused.to_string().contains("unsafe use of fn_trusted"), + "a schema-scoped function was not refused on an untrusted schema: {refused}" + ); + + let accepted: String = scalar(&mut *conn, "SELECT body FROM innocuous_view").await; + assert_eq!(accepted, "held"); +} + +/// A database captures the registered set at its `connect`. A registration made after +/// that point applies to the next database to connect, and to no connection of the one +/// already open, however the pools grow or replace connections. +#[tokio::test] +async fn a_registration_after_a_connect_applies_to_the_next_database_only() { + let dir = TempDir::new().unwrap(); + let earlier = open(&dir, None).await; + + register_function(ScalarFunction { + name: "fn_after_connect".into(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::DirectOnly, + handler: echo(), + }) + .expect("registration"); + + let unresolved = + scalar_error(earlier.read_pool().unwrap(), "SELECT fn_after_connect('x')").await; + assert!( + unresolved + .to_string() + .contains("no such function: fn_after_connect"), + "the earlier database must not resolve the function: {unresolved}" + ); + + let later_dir = TempDir::new().unwrap(); + let later = SqliteDatabase::connect(later_dir.path().join("later.db"), None) + .await + .expect("connect"); + + let resolved: String = scalar(later.read_pool().unwrap(), "SELECT fn_after_connect('x')").await; + assert_eq!(resolved, "x"); +} From 32e7defce938735bb3fdd9e30200891b1fb24324 Mon Sep 17 00:00:00 2001 From: Ethan Smith Date: Wed, 19 Aug 2026 20:18:30 +0000 Subject: [PATCH 3/5] feat: expose scalar function registration on the plugin builder `Builder::register_function` and `SetupRegistrar::register_function` register a function for every database the plugin serves. The registrar variant exists for a handler that closes over the `app` instance; the builder variant covers everything else. Registration happens at the call rather than at `build()`, so a validation error names the offending function where a developer wrote it. The toolkit and the plugin both re-export the function types, so a consumer does not depend on the connection manager directly. --- crates/sqlx-sqlite-toolkit/src/lib.rs | 4 +- .../tests/function_tests.rs | 171 ++++++++++++++++++ src/lib.rs | 102 ++++++++++- 3 files changed, 275 insertions(+), 2 deletions(-) create mode 100644 crates/sqlx-sqlite-toolkit/tests/function_tests.rs diff --git a/crates/sqlx-sqlite-toolkit/src/lib.rs b/crates/sqlx-sqlite-toolkit/src/lib.rs index 6746ded..0f24380 100644 --- a/crates/sqlx-sqlite-toolkit/src/lib.rs +++ b/crates/sqlx-sqlite-toolkit/src/lib.rs @@ -55,5 +55,7 @@ pub use wrapper::{ // Re-export commonly used types from dependencies pub use sqlx_sqlite_conn_mgr::{ - AttachedMode, AttachedSpec, Migrator, SqliteDatabase, SqliteDatabaseConfig, + AttachedMode, AttachedSpec, FunctionError, InvocationScope, Migrator, ScalarFunction, + ScalarHandler, SqlValue, SqlValueRef, SqliteDatabase, SqliteDatabaseConfig, register_function, + register_or_replace_function, }; diff --git a/crates/sqlx-sqlite-toolkit/tests/function_tests.rs b/crates/sqlx-sqlite-toolkit/tests/function_tests.rs new file mode 100644 index 0000000..8f21525 --- /dev/null +++ b/crates/sqlx-sqlite-toolkit/tests/function_tests.rs @@ -0,0 +1,171 @@ +//! Scalar functions resolve on the paths the toolkit routes queries through. +//! +//! `sqlx-sqlite-conn-mgr/tests/function_tests.rs` covers the connection-level behavior. +//! These tests cover what the toolkit adds on top: transactions, `INSERT ... SELECT`, and a +//! query naming an attached database. +//! +//! Every test in this binary shares the process-global function registry, and each +//! database captures the registered set at its `connect`. `REGISTERED` therefore holds +//! every function this binary needs, under a name per test, and `database` forces it +//! before the first database connects. + +use std::sync::{Arc, LazyLock}; + +use serde_json::{Value, json}; +use sqlx_sqlite_conn_mgr::{AttachedMode, AttachedSpec}; +use sqlx_sqlite_toolkit::{ + DatabaseWrapper, FunctionError, InvocationScope, ScalarFunction, ScalarHandler, SqlValue, + SqlValueRef, register_function, +}; +use tempfile::TempDir; + +fn register(name: &str, handler: ScalarHandler) { + register_function(ScalarFunction { + name: name.to_string(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::DirectOnly, + handler, + }) + .expect("registration"); +} + +/// A handler that upper-cases text and passes NULL through. +fn shout() -> ScalarHandler { + Arc::new(|args| match &args[0] { + SqlValueRef::Text(text) => Ok(SqlValue::Text(text.to_uppercase())), + SqlValueRef::Null => Ok(SqlValue::Null), + _ => Err(FunctionError::new("expected text")), + }) +} + +/// Every function this binary uses, one per test, registered before any database opens. +static REGISTERED: LazyLock<()> = LazyLock::new(|| { + register("fn_tx_shout", shout()); + register("fn_attached_shout", shout()); + register( + "fn_tx_error", + Arc::new(|_args| Err(FunctionError::new("payload is not decodable"))), + ); +}); + +async fn database(name: &str) -> (DatabaseWrapper, TempDir) { + LazyLock::force(®ISTERED); + + let temp = TempDir::new().expect("temp dir"); + let db = DatabaseWrapper::connect(&temp.path().join(name), None) + .await + .expect("connect"); + + (db, temp) +} + +/// Runs `sql` for its effect, panicking with the SQL itself. +async fn exec(db: &DatabaseWrapper, sql: &str) { + db.execute(sql.into(), vec![]).await.expect(sql); +} + +/// Runs `sql` and returns column `name` from every row, in row order. +async fn column(db: &DatabaseWrapper, sql: &str, name: &str) -> Vec { + db.fetch_all(sql.into(), vec![]) + .await + .expect(sql) + .iter() + .map(|row| row.get(name).expect(name).clone()) + .collect() +} + +#[tokio::test] +async fn resolves_inside_a_transaction_and_inside_an_insert_select() { + let (db, _temp) = database("tx.db").await; + + exec(&db, "CREATE TABLE source (body TEXT)").await; + exec(&db, "CREATE TABLE target (body TEXT)").await; + exec(&db, "INSERT INTO source (body) VALUES ('alpha'), ('beta')").await; + + // Both statements run on the write connection inside one transaction: the first calls + // the function in a VALUES list, the second inside an INSERT ... SELECT. + db.execute_transaction(vec![ + ( + "INSERT INTO target (body) VALUES (fn_tx_shout('gamma'))", + vec![], + ), + ( + "INSERT INTO target (body) SELECT fn_tx_shout(body) FROM source ORDER BY body", + vec![], + ), + ]) + .execute() + .await + .expect("transaction"); + + assert_eq!( + column(&db, "SELECT body FROM target ORDER BY body", "body").await, + vec![json!("ALPHA"), json!("BETA"), json!("GAMMA")] + ); +} + +#[tokio::test] +async fn resolves_in_a_query_naming_an_attached_database() { + // Each database captures the registered set at its `connect`, so the registration must + // exist before either connect below. + LazyLock::force(®ISTERED); + + let temp = TempDir::new().expect("temp dir"); + let main = DatabaseWrapper::connect(&temp.path().join("main.db"), None) + .await + .expect("connect main"); + let other = DatabaseWrapper::connect(&temp.path().join("other.db"), None) + .await + .expect("connect other"); + + exec(&other, "CREATE TABLE logs (msg TEXT)").await; + exec(&other, "INSERT INTO logs (msg) VALUES ('stored')").await; + + let rows = main + .fetch_all( + "SELECT fn_attached_shout(msg) AS shouted FROM other.logs".into(), + vec![], + ) + .attach(vec![AttachedSpec { + database: Arc::clone(other.inner()), + schema_name: "other".to_string(), + mode: AttachedMode::ReadOnly, + }]) + .await + .expect("attached read"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get("shouted"), Some(&json!("STORED"))); +} + +#[tokio::test] +async fn a_handler_error_in_a_transaction_leaves_no_transaction_open() { + let (db, _temp) = database("tx_error.db").await; + + exec(&db, "CREATE TABLE t (body TEXT)").await; + + let error = db + .execute_transaction(vec![ + ("INSERT INTO t (body) VALUES ('before')", vec![]), + ("INSERT INTO t (body) VALUES (fn_tx_error('x'))", vec![]), + ]) + .execute() + .await + .expect_err("the transaction must fail"); + + assert!( + error.to_string().contains("payload is not decodable"), + "error did not include the handler's message: {error}" + ); + + // The write pool holds one connection. If the failure leaves a transaction open, this + // write will fail with "cannot start a transaction within a transaction". + exec(&db, "INSERT INTO t (body) VALUES ('after')").await; + + // The failed transaction rolled back, so only the later write survives. + assert_eq!( + column(&db, "SELECT body FROM t", "body").await, + vec![json!("after")] + ); +} diff --git a/src/lib.rs b/src/lib.rs index 59a04cb..a6efb85 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,7 +16,8 @@ mod validate; pub use error::{Error, Result}; pub use sqlx_sqlite_conn_mgr::{ - AttachedMode, AttachedSpec, Migrator as SqliteMigrator, SqliteDatabaseConfig, + AttachedMode, AttachedSpec, FunctionError, InvocationScope, Migrator as SqliteMigrator, + ScalarFunction, ScalarHandler, SqlValue, SqlValueRef, SqliteDatabaseConfig, }; pub use sqlx_sqlite_toolkit::{ ActiveInterruptibleTransactions, ActiveRegularTransactions, DatabaseWrapper, @@ -318,6 +319,24 @@ impl SetupRegistrar { .insert(key.to_string(), validated_database_info(path, migrator)?); Ok(()) } + + /// Register a scalar SQL function under the same contract as + /// [`Builder::register_function`]. The plugin's setup hook runs before any database + /// loads, so a function registered here reaches every database the plugin serves. + /// + /// Register here when the handler closes over something from the `app` instance. A + /// handler that needs nothing from `app` belongs on the builder itself. + pub fn register_function(&mut self, function: ScalarFunction) -> Result<()> { + sqlx_sqlite_conn_mgr::register_function(function)?; + Ok(()) + } + + /// Register a scalar SQL function, replacing the one already registered under the same + /// name and arity. See [`Builder::register_or_replace_function`]. + pub fn register_or_replace_function(&mut self, function: ScalarFunction) -> Result<()> { + sqlx_sqlite_conn_mgr::register_or_replace_function(function)?; + Ok(()) + } } /// Closure type for the deferred [`Builder::on_setup`] hook. @@ -413,6 +432,87 @@ impl Builder { Ok(self) } + /// Register a scalar SQL function, callable by name from every query the plugin serves. + /// + /// The full contract lives in the [`sqlx_sqlite_conn_mgr::functions`] module + /// documentation. The plugin-level facts: + /// + /// - Register before the first `load` or `connect`. Each database captures the + /// registered set when it connects, so a function registered here reaches every + /// database the plugin serves. + /// - A returned [`FunctionError`] fails the statement and reports its message to the + /// caller as error code `SQLITE_1`. Inside a transaction, the plugin's usual + /// rollback behavior applies. + /// - Registration is process-global and takes effect at this call, not at + /// [`build`](Self::build). A validation error appears at the call site. A builder + /// discarded without `build()` still leaves its functions registered. + /// - When the handler needs the `app` instance, use + /// [`SetupRegistrar::register_function`] instead. + /// - [`ScalarFunction::invocation_scope`] decides whether a schema object can call the + /// function. [`InvocationScope::DirectOnly`] confines it to top-level SQL. + /// - The plugin does not support aggregate functions, window functions, collations, + /// virtual tables, or functions defined in JavaScript. + /// + /// # Errors + /// + /// Returns `Err` for an invalid name or arity, a name a SQLite built-in already uses, + /// a name and arity pair already registered, or a function the linked SQLite library + /// refuses. See [`sqlx_sqlite_conn_mgr::register_function`] for the + /// variant-by-variant list. + /// + /// # Example + /// + /// ```no_run + /// use std::sync::Arc; + /// use tauri_plugin_sqlite::{ + /// Builder, FunctionError, InvocationScope, ScalarFunction, SqlValue, SqlValueRef, + /// }; + /// + /// # fn example() -> tauri_plugin_sqlite::Result<()> { + /// Builder::::new() + /// .register_function(ScalarFunction { + /// name: "normalize_for_search".into(), + /// arity: 1, + /// deterministic: true, + /// invocation_scope: InvocationScope::DirectOnly, + /// // SQLite's own lower() folds ASCII only, so Unicode case folding needs Rust. + /// handler: Arc::new(|args: &[SqlValueRef]| match &args[0] { + /// SqlValueRef::Text(text) => Ok(SqlValue::Text(text.to_lowercase())), + /// SqlValueRef::Null => Ok(SqlValue::Null), + /// _ => Err(FunctionError::new("normalize_for_search expects text")), + /// }), + /// })? + /// .build()?; + /// # Ok(()) + /// # } + /// ``` + pub fn register_function(self, function: ScalarFunction) -> Result { + sqlx_sqlite_conn_mgr::register_function(function)?; + + Ok(self) + } + + /// Register a scalar SQL function, replacing the one already registered under the same + /// name and arity. + /// + /// Use this where a repeat registration is the caller's intent, such as a setup path a + /// test or a restart can run twice. [`register_function`](Self::register_function) + /// rejects the repeat instead. Everything else about the two calls matches, including + /// the validation and the connect-time rule: the replacement reaches the databases that + /// load after it, and a database already loaded keeps the handler it captured. See + /// [`sqlx_sqlite_conn_mgr::register_or_replace_function`]. + /// + /// # Errors + /// + /// Returns `Err` for every reason + /// [`register_function`](Self::register_function) does, except a name and arity + /// already registered, which is the case this call accepts. + pub fn register_or_replace_function(self, function: ScalarFunction) -> Result { + sqlx_sqlite_conn_mgr::register_or_replace_function(function)?; + + Ok(self) + } + /// Set the timeout for interruptible transactions. /// /// If an interruptible transaction exceeds this duration, it will be automatically From 9a86ecb91c32fc1be38183fd60ba4eb47b10ec24 Mon Sep 17 00:00:00 2001 From: Ethan Smith Date: Wed, 19 Aug 2026 20:20:17 +0000 Subject: [PATCH 4/5] docs: document scalar function registration Both READMEs carry the same contract: a function belongs to a connection rather than to a database file, registration precedes the first open, and the handler runs on the connection's own thread. Also states what the mechanism cannot promise. A handler panic becomes an error only where the application unwinds on panic, and aggregate functions, window functions, collations, virtual tables, and functions defined in JavaScript are out of scope. --- CHANGELOG.md | 5 ++ README.md | 78 +++++++++++++++++++++++++++ crates/sqlx-sqlite-conn-mgr/README.md | 54 +++++++++++++++++++ 3 files changed, 137 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d45f962..0ba3047 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -56,9 +56,14 @@ Databases must be registered on the Rust side with a stable key before they can - Path validation module (`validate.rs`) with canonicalization at registration time. - Parent directory auto-creation during registration validation for file paths. - CI check that committed `api-iife.js` matches a fresh Rollup build. +- Scalar SQL functions written in Rust. Any query, from the frontend or from Rust, can call a registered function by name. Each database captures the registered set when it connects, so registration must precede the first `load`. See [Scalar Functions](README.md#scalar-functions) for the full contract. Aggregate functions, window functions, collations, virtual tables, and functions defined in JavaScript are out of scope. #### Rust API +- **`Builder::register_function()`** and **`SetupRegistrar::register_function()`**: register a scalar SQL function. Use the registrar variant when the handler closes over the `app` instance. Use the builder variant otherwise. Both validate the name, the arity, the name-and-arity pair, and a collision with a SQLite built-in at the call. A bad registration produces an error there, not later when a connection opens. +- **`Builder::register_or_replace_function()`** and **`SetupRegistrar::register_or_replace_function()`**: the same registration, accepting a name and arity already registered. For a setup path that runs twice, where `register_function()` reports the repeat as an error. +- **`InvocationScope`**, the required **`ScalarFunction::invocation_scope`** field: `DirectOnly` confines a function to top-level SQL through `SQLITE_DIRECTONLY`, and `Schema` and `InnocuousSchema` let a view, a trigger, or an index call it. Pick `DirectOnly` unless a schema object must call the function: a handler reachable from the schema of a database this application did not write runs on any read through that object. +- **`sqlx_sqlite_conn_mgr::register_function()`** and **`register_or_replace_function()`** with **`ScalarFunction`**, **`InvocationScope`**, **`ScalarHandler`**, **`SqlValue`**, **`SqlValueRef`**, and **`FunctionError`**: the registry the plugin builds on, usable without Tauri. `SqliteDatabase::connect` snapshots the registered set, and both of the database's pools install that snapshot through `SqlitePoolOptions::after_connect`. This keeps a function present even on a connection the idle reaper replaced, and it keeps one database's connections uniform: a registration after a `connect` applies to the next database, never to part of an open one. Handlers receive borrowed `SqlValueRef` arguments and return an owned `SqlValue`, and `ScalarFunction::name` is an owned `String`, so a name built at runtime needs no leak. The types are re-exported from `sqlx_sqlite_toolkit`; the plugin re-exports the types and exposes registration through the builder and the registrar alone. - **`sqlx_sqlite_toolkit::AttachedWriterGuard`** and **`DatabaseWrapper::acquire_writer_with_attached()`** — the attached-writer acquisition path that routes through the observer. - **`sqlx_sqlite_conn_mgr::ObserverSlot`** (with `SqliteDatabase::observer_slot()`) — the database-scoped observation slot underlying the database-wide sharing described above. Populated only through `get_or_init()` / `get_or_init_with()`, both of which reuse rather than replace, so the slot cannot come to hold two different concrete types from safe code. `get_or_init_with()` additionally runs a caller-supplied merge callback under the same write lock that decided to reuse. - **`sqlx_sqlite_observer::ObservableSqliteDatabase::acquire_writer_with_attached()`** and **`ObservableWriteGuard::detach_all()`**. diff --git a/README.md b/README.md index 1b63923..8894d4a 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,8 @@ SQLite database interface for Tauri applications using * **Migration Support**: SQLx's migration framework * **Resource Management**: Proper cleanup on application exit * **Optional Change Notifications**: SQLite hooks for reactive change notifications + * **Scalar Functions in Rust**: register a Rust function and call it by name from + any query the plugin serves ## Architecture @@ -941,6 +943,82 @@ async fn example(app: tauri::AppHandle) -> tauri_plugin_sqlite::R } ``` +### Scalar Functions + +Register a Rust function once and call it by name from any query the plugin serves. This +puts work that needs native code inside the query, so SQLite applies it while it scans. +The alternative reads every candidate row out of the database and applies the same logic +in the caller. + +```rust +use std::sync::Arc; +use tauri_plugin_sqlite::{ + Builder, FunctionError, InvocationScope, ScalarFunction, SqlValue, SqlValueRef, +}; + +Builder::new() + .register_function(ScalarFunction { + name: "normalize_for_search".into(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::DirectOnly, + // SQLite's own lower() folds ASCII only, so Unicode case folding needs Rust. + handler: Arc::new(|args: &[SqlValueRef]| match &args[0] { + SqlValueRef::Text(text) => Ok(SqlValue::Text(text.to_lowercase())), + SqlValueRef::Null => Ok(SqlValue::Null), + _ => Err(FunctionError::new("normalize_for_search expects text")), + }), + })? + .build()?; +``` + +Any query can then call it, from the frontend or from Rust: + +```sql +SELECT d.Title AS title + FROM Document d + WHERE normalize_for_search(d.Title) LIKE $1 +``` + +Register a function inside `on_setup` instead when its handler needs the `app` instance: + +```rust +Builder::new() + .on_setup(|app, reg| { + reg.register_database(MAIN_DB_KEY, app.path().app_data_dir()?.join("main.db"), None)?; + reg.register_function(ScalarFunction { /* ... */ })?; + Ok(()) + }) + .build()?; +``` + +The full contract lives in the `sqlx_sqlite_conn_mgr::functions` module documentation +(crates/sqlx-sqlite-conn-mgr/src/functions.rs). The short version: + + * Register before the first `load` or `connect`. Each database captures the registered + set when it connects, for every connection it ever opens. + * `invocation_scope` decides where SQLite accepts a call. + `InvocationScope::DirectOnly` confines the function to top-level SQL, so SQLite + refuses a call from inside a view, a trigger, a CHECK constraint, an expression + index, or any other schema object. `InvocationScope::Schema` and + `InvocationScope::InnocuousSchema` accept a call from a schema object, which is what + a view over the function needs. + * Validation happens at the `register_function` call: an invalid name or arity, a name + a SQLite built-in already uses, a duplicate name and arity pair, and a function the + linked SQLite library refuses all fail there. + * `register_or_replace_function` accepts a name and arity already registered, for a + setup path that runs twice. The replacement reaches the databases that load after + it, and a database already loaded keeps the handler it captured. + * The handler receives borrowed arguments (`SqlValueRef`), including SQL NULL, and + returns an owned `SqlValue`. A returned `FunctionError` fails the statement with + error code `SQLITE_1` and the handler's message. + * A panicking handler produces an error naming the function and leaves the pool + usable, under `panic = "unwind"` (the Rust default). Under `panic = "abort"`, the + process aborts. + +Aggregate functions, window functions, collations, virtual tables, and functions defined +in JavaScript are not supported. + ### Basic Operations ```rust diff --git a/crates/sqlx-sqlite-conn-mgr/README.md b/crates/sqlx-sqlite-conn-mgr/README.md index 2b2c087..74b3d49 100644 --- a/crates/sqlx-sqlite-conn-mgr/README.md +++ b/crates/sqlx-sqlite-conn-mgr/README.md @@ -15,6 +15,8 @@ project needing SQLx connection management. * **WAL mode**: Enabled on first `acquire_writer()` call * **Idle timeout**: Connections close after 30s inactivity (configurable) * **No perpetual caching**: Zero minimum connections (prevents idle thread sprawl) + * **Scalar functions**: Either pool applies a registered Rust function to every + connection it opens, for every database this crate serves Delegates to SQLx's `SqlitePoolOptions` and `SqliteConnectOptions` wherever possible — minimal wrapper logic. @@ -154,6 +156,36 @@ async fn example() -> Result<(), sqlx_sqlite_conn_mgr::Error> { > circumvents the connection manager's policies and will result in > unpredictable behavior, including potential deadlocks. +### Scalar Functions + +`register_function` records a Rust function. Every database whose `connect` runs after the +registration keeps that function available on every connection either of its pools opens. + +```rust +use std::sync::Arc; +use sqlx_sqlite_conn_mgr::{ + FunctionError, InvocationScope, ScalarFunction, SqlValue, SqlValueRef, register_function, +}; + +register_function(ScalarFunction { + name: "normalize_for_search".into(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::DirectOnly, + // SQLite's own lower() folds ASCII only, so Unicode case folding needs Rust. + handler: Arc::new(|args: &[SqlValueRef]| match &args[0] { + SqlValueRef::Text(text) => Ok(SqlValue::Text(text.to_lowercase())), + SqlValueRef::Null => Ok(SqlValue::Null), + _ => Err(FunctionError::new("normalize_for_search expects text")), + }), +})?; +``` + +The full contract (connection-scoped registration, connect-time snapshots, the invocation +scope, panic behavior, and the validation list) lives in the `functions` module +documentation (src/functions.rs). Aggregate functions, window +functions, collations, and virtual tables are not supported. + ## API Reference ### `SqliteDatabase` @@ -167,6 +199,28 @@ async fn example() -> Result<(), sqlx_sqlite_conn_mgr::Error> { | `close()` | Close and remove from cache | | `remove()` | Close and delete database files (.db, .db-wal, .db-shm) | +### `register_function` + +| Function | Description | +| -------- | ----------- | +| `register_function(function)` | Register a `ScalarFunction` for every database that connects after the call | +| `register_or_replace_function(function)` | The same, replacing the function already registered under the name and arity | + +Each of these fails at the call: + + * an empty name + * a name over 255 bytes + * a name holding a NUL byte + * a negative arity + * a name a SQLite built-in already uses + * a name and arity pair already registered, unless the call is + `register_or_replace_function` + * a function the linked SQLite library refuses + +The registry compares names case-insensitively. This matches how SQLite resolves them. The +call also registers the function on a temporary in-memory connection. A library that +refuses the function there returns an error naming the function and its reason. + ### `WriteGuard` RAII guard for exclusive write access. Derefs to `SqliteConnection`. Connection From 93a8d3fdf876185cb50469a9b27f120ececf15a7 Mon Sep 17 00:00:00 2001 From: Ethan Smith Date: Wed, 19 Aug 2026 20:27:38 +0000 Subject: [PATCH 5/5] test: cover the plugin's own path to a registered function Registers a function on the builder, then reads a column through it with the `fetch_all` command, which is the path a consumer's frontend takes. The layers underneath have their own tests; this one asserts that the builder entry point and the pools meet. --- src/lib.rs | 63 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index a6efb85..a11da95 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1536,6 +1536,69 @@ mod tests { app } + /// A function registered on the builder resolves in a query the plugin serves. This test + /// covers that path end to end. Each layer below has its own tests in its own crate. + /// + /// The registry is process-global, so `plugin_shout` belongs to this test alone. A + /// second test in this binary needs a name of its own. + #[test] + fn registered_function_resolves_through_the_plugin() { + let temp_dir = tempfile::tempdir().unwrap(); + let db_path = validate::validate_database_path(temp_dir.path().join("main.db")).unwrap(); + + let mut plugin = Builder::::new() + .register_function(ScalarFunction { + name: "plugin_shout".into(), + arity: 1, + deterministic: true, + invocation_scope: InvocationScope::DirectOnly, + handler: Arc::new(|args: &[SqlValueRef]| match &args[0] { + SqlValueRef::Text(text) => Ok(SqlValue::Text(text.to_uppercase())), + _ => Err(FunctionError::new("plugin_shout expects text")), + }), + }) + .unwrap() + .register_database("MAIN", &db_path, None) + .unwrap() + .build() + .unwrap(); + + let app = mock_app(); + plugin + .initialize(app.handle(), serde_json::Value::default()) + .expect("plugin init should succeed"); + + tauri::async_runtime::block_on(async { + load_and_create_test_table(&app, "MAIN").await; + + commands::execute( + app.state::(), + "MAIN".to_string(), + "INSERT INTO test (val) VALUES ('quiet')".to_string(), + vec![], + None, + ) + .await + .expect("insert should succeed"); + + let rows = commands::fetch_all( + app.state::(), + "MAIN".to_string(), + "SELECT plugin_shout(val) AS shouted FROM test".to_string(), + vec![], + None, + ) + .await + .expect("the query must resolve the registered function"); + + assert_eq!(rows.len(), 1); + assert_eq!( + rows[0].get("shouted").and_then(|value| value.as_str()), + Some("QUIET") + ); + }); + } + #[tokio::test] async fn test_connect_to_database_registered_key() { let temp_dir = tempfile::tempdir().unwrap();