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/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/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/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/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 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; 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"); +} 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..a11da95 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 @@ -1436,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();