You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
439 lines
18 KiB
439 lines
18 KiB
use std::sync::{Arc, Mutex};
|
|
|
|
use mlua::prelude::LuaResult;
|
|
use mlua::{Lua, Table as LuaTable, UserData, UserDataMethods, Value as LuaValue};
|
|
use rusqlite::types::Value;
|
|
|
|
use super::utils::is_null;
|
|
|
|
const SQLITE_LUA: &str = include_str!("../../lua/stdlib/sqlite.lua");
|
|
|
|
/// Shared, lockable handle to a SQLite connection.
|
|
///
|
|
/// `Option` allows an explicit `close()` (set to `None`); any later use then
|
|
/// reports a clear "connection is closed" error. `Arc<Mutex<>>` is required
|
|
/// because mlua's async `UserData` must be `Send + Sync + 'static`, and it lets
|
|
/// the blocking SQLite work run on a `spawn_blocking` thread.
|
|
type ConnHandle = Arc<Mutex<Option<rusqlite::Connection>>>;
|
|
|
|
struct SqliteConn(ConnHandle);
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Errors
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Error type carried back across the `spawn_blocking` boundary. Unlike
|
|
/// `mlua::Error` it is trivially `Send + Sync`, so it travels cleanly between
|
|
/// threads before being surfaced to Lua via `mlua::Error::external`.
|
|
#[derive(Debug)]
|
|
enum DbError {
|
|
Closed,
|
|
Sqlite(rusqlite::Error),
|
|
Other(String),
|
|
}
|
|
|
|
impl std::fmt::Display for DbError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
DbError::Closed => write!(f, "sqlite: connection is closed"),
|
|
DbError::Sqlite(e) => write!(f, "sqlite: {e}"),
|
|
DbError::Other(s) => write!(f, "{s}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for DbError {}
|
|
|
|
impl From<rusqlite::Error> for DbError {
|
|
fn from(e: rusqlite::Error) -> Self {
|
|
DbError::Sqlite(e)
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Parameter binding
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Parameters bound to a statement, already converted to a `Send` form so they
|
|
/// can cross into the blocking thread.
|
|
enum BoundParams {
|
|
None,
|
|
/// Positional `?` parameters, in order.
|
|
Positional(Vec<Value>),
|
|
/// Named `:name` / `@name` / `$name` parameters (sigil stripped from the key).
|
|
Named(Vec<(String, Value)>),
|
|
}
|
|
|
|
/// Convert an optional Lua params table into `BoundParams`.
|
|
///
|
|
/// A table with any string key is treated as named; a table with only integer
|
|
/// keys is positional. Mixing the two is an error, as is a non-contiguous
|
|
/// positional array. Runs on the Lua side (before `spawn_blocking`) because it
|
|
/// touches non-`Send` Lua values.
|
|
fn convert_params(params: Option<LuaTable>) -> LuaResult<BoundParams> {
|
|
let table = match params {
|
|
None => return Ok(BoundParams::None),
|
|
Some(t) => t,
|
|
};
|
|
|
|
let mut positional: Vec<(i64, Value)> = Vec::new();
|
|
let mut named: Vec<(String, Value)> = Vec::new();
|
|
|
|
for pair in table.pairs::<LuaValue, LuaValue>() {
|
|
let (key, value) = pair?;
|
|
let value = lua_value_to_sqlite(value)?;
|
|
match key {
|
|
LuaValue::Integer(i) => positional.push((i, value)),
|
|
LuaValue::Number(n) if n.fract() == 0.0 => positional.push((n as i64, value)),
|
|
LuaValue::String(s) => named.push((s.to_str()?.to_owned(), value)),
|
|
other => {
|
|
return Err(mlua::Error::external(format!(
|
|
"sqlite: unsupported parameter key of type '{}'",
|
|
other.type_name()
|
|
)))
|
|
}
|
|
}
|
|
}
|
|
|
|
match (positional.is_empty(), named.is_empty()) {
|
|
(true, true) => Ok(BoundParams::None),
|
|
(false, false) => Err(mlua::Error::external(
|
|
"sqlite: parameter table mixes positional (array) and named keys",
|
|
)),
|
|
(true, false) => Ok(BoundParams::Named(named)),
|
|
(false, true) => {
|
|
// Require a gapless 1..=n sequence so the `?` placeholders line up.
|
|
positional.sort_by_key(|(i, _)| *i);
|
|
for (offset, (index, _)) in positional.iter().enumerate() {
|
|
if *index != offset as i64 + 1 {
|
|
return Err(mlua::Error::external(
|
|
"sqlite: positional parameters must form a contiguous array starting at index 1",
|
|
));
|
|
}
|
|
}
|
|
Ok(BoundParams::Positional(
|
|
positional.into_iter().map(|(_, v)| v).collect(),
|
|
))
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Map a Lua value to the SQLite value it binds as.
|
|
fn lua_value_to_sqlite(value: LuaValue) -> LuaResult<Value> {
|
|
// `utils.NULL` is the only practical way to bind a SQL NULL, since a literal
|
|
// `nil` cannot live inside a Lua table.
|
|
if is_null(&value) {
|
|
return Ok(Value::Null);
|
|
}
|
|
match value {
|
|
LuaValue::Nil => Ok(Value::Null),
|
|
LuaValue::Boolean(b) => Ok(Value::Integer(b as i64)),
|
|
LuaValue::Integer(i) => Ok(Value::Integer(i)),
|
|
LuaValue::Number(f) => Ok(Value::Real(f)),
|
|
// Lua strings are byte strings. Valid UTF-8 binds as TEXT; anything else
|
|
// (raw bytes, embedded NULs) binds as a BLOB so binary data round-trips
|
|
// instead of failing the UTF-8 conversion.
|
|
LuaValue::String(s) => {
|
|
let bytes = s.as_bytes();
|
|
match std::str::from_utf8(&bytes) {
|
|
Ok(text) => Ok(Value::Text(text.to_owned())),
|
|
Err(_) => Ok(Value::Blob(bytes.to_vec())),
|
|
}
|
|
}
|
|
other => Err(mlua::Error::external(format!(
|
|
"sqlite: cannot bind value of type '{}'",
|
|
other.type_name()
|
|
))),
|
|
}
|
|
}
|
|
|
|
/// Bind already-converted parameters onto a prepared statement.
|
|
fn bind_params(stmt: &mut rusqlite::CachedStatement, params: &BoundParams) -> Result<(), DbError> {
|
|
match params {
|
|
BoundParams::None => Ok(()),
|
|
BoundParams::Positional(values) => {
|
|
let expected = stmt.parameter_count();
|
|
if values.len() != expected {
|
|
return Err(DbError::Other(format!(
|
|
"sqlite: statement expects {expected} positional parameter(s) but {} were given",
|
|
values.len()
|
|
)));
|
|
}
|
|
for (offset, value) in values.iter().enumerate() {
|
|
stmt.raw_bind_parameter(offset + 1, value)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
BoundParams::Named(map) => {
|
|
for index in 1..=stmt.parameter_count() {
|
|
// `parameter_name` returns the placeholder with its sigil, e.g. ":id";
|
|
// `None` means an anonymous `?`, which can't be filled by name.
|
|
let name = stmt
|
|
.parameter_name(index)
|
|
.ok_or_else(|| {
|
|
DbError::Other(format!(
|
|
"sqlite: statement has a positional placeholder at position {index} but named parameters were given"
|
|
))
|
|
})?
|
|
.to_string();
|
|
let key = &name[1..]; // strip the leading sigil
|
|
let value = map
|
|
.iter()
|
|
.find(|(k, _)| k == key)
|
|
.map(|(_, v)| v)
|
|
.ok_or_else(|| {
|
|
DbError::Other(format!("sqlite: missing value for named parameter '{name}'"))
|
|
})?;
|
|
stmt.raw_bind_parameter(index, value)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// SQLite value -> Lua
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn sqlite_value_to_lua(lua: &Lua, value: Value) -> LuaResult<LuaValue> {
|
|
Ok(match value {
|
|
Value::Null => LuaValue::Nil,
|
|
Value::Integer(i) => LuaValue::Integer(i),
|
|
Value::Real(f) => LuaValue::Number(f),
|
|
Value::Text(s) => LuaValue::String(lua.create_string(&s)?),
|
|
Value::Blob(b) => LuaValue::String(lua.create_string(&b)?),
|
|
})
|
|
}
|
|
|
|
/// One database row as (column name, value) pairs.
|
|
type Record = Vec<(String, Value)>;
|
|
|
|
fn record_to_lua(lua: &Lua, record: Record) -> LuaResult<LuaTable> {
|
|
let table = lua.create_table_with_capacity(0, record.len())?;
|
|
for (name, value) in record {
|
|
// A SQL NULL maps to Lua nil; setting a key to nil leaves it absent,
|
|
// which reads back as nil — exactly what we want.
|
|
if matches!(value, Value::Null) {
|
|
continue;
|
|
}
|
|
table.raw_set(name, sqlite_value_to_lua(lua, value)?)?;
|
|
}
|
|
Ok(table)
|
|
}
|
|
|
|
fn rows_to_lua(lua: &Lua, rows: Vec<Record>) -> LuaResult<LuaTable> {
|
|
let array = lua.create_table_with_capacity(rows.len(), 0)?;
|
|
for (i, record) in rows.into_iter().enumerate() {
|
|
array.raw_set(i + 1, record_to_lua(lua, record)?)?;
|
|
}
|
|
Ok(array)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Blocking workers (run on the tokio blocking pool)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Reject SQL that is empty/comment-only or contains more than one statement.
|
|
///
|
|
/// `execute`/`query`/`queryOne` compile a single statement via `prepare_cached`,
|
|
/// which silently ignores anything after the first — so `execute("A; B")` would
|
|
/// run only `A` and drop `B` with no error, and `execute("")` surfaces SQLite's
|
|
/// bare "not an error". `Batch` uses SQLite's own tokenizer to count statements
|
|
/// (skipping comments/whitespace), giving a clear message instead. Use
|
|
/// `executeBatch` to run multiple statements on purpose.
|
|
fn ensure_single_statement(conn: &rusqlite::Connection, sql: &str) -> Result<(), DbError> {
|
|
let mut batch = rusqlite::Batch::new(conn, sql);
|
|
let mut count = 0usize;
|
|
while batch.next()?.is_some() {
|
|
count += 1;
|
|
if count > 1 {
|
|
return Err(DbError::Other(
|
|
"sqlite: expected a single SQL statement (use executeBatch for multiple)".to_string(),
|
|
));
|
|
}
|
|
}
|
|
if count == 0 {
|
|
return Err(DbError::Other("sqlite: no SQL statement to run".to_string()));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Run an INSERT/UPDATE/DELETE-style statement, returning rows changed.
|
|
async fn run_execute(handle: ConnHandle, sql: String, params: BoundParams) -> Result<usize, DbError> {
|
|
tokio::task::spawn_blocking(move || -> Result<usize, DbError> {
|
|
let guard = handle.lock().unwrap_or_else(|e| e.into_inner());
|
|
let conn = guard.as_ref().ok_or(DbError::Closed)?;
|
|
ensure_single_statement(conn, &sql)?;
|
|
let mut stmt = conn.prepare_cached(&sql)?;
|
|
bind_params(&mut stmt, ¶ms)?;
|
|
Ok(stmt.raw_execute()?)
|
|
})
|
|
.await
|
|
.unwrap_or_else(|e| Err(DbError::Other(format!("sqlite: background task failed: {e}"))))
|
|
}
|
|
|
|
/// Run every statement in `sql` in order, with no parameter binding.
|
|
async fn run_execute_batch(handle: ConnHandle, sql: String) -> Result<(), DbError> {
|
|
tokio::task::spawn_blocking(move || -> Result<(), DbError> {
|
|
let guard = handle.lock().unwrap_or_else(|e| e.into_inner());
|
|
let conn = guard.as_ref().ok_or(DbError::Closed)?;
|
|
conn.execute_batch(&sql)?;
|
|
Ok(())
|
|
})
|
|
.await
|
|
.unwrap_or_else(|e| Err(DbError::Other(format!("sqlite: background task failed: {e}"))))
|
|
}
|
|
|
|
/// Run a SELECT-style statement, materializing every row.
|
|
async fn run_query(handle: ConnHandle, sql: String, params: BoundParams) -> Result<Vec<Record>, DbError> {
|
|
tokio::task::spawn_blocking(move || -> Result<Vec<Record>, DbError> {
|
|
let guard = handle.lock().unwrap_or_else(|e| e.into_inner());
|
|
let conn = guard.as_ref().ok_or(DbError::Closed)?;
|
|
ensure_single_statement(conn, &sql)?;
|
|
let mut stmt = conn.prepare_cached(&sql)?;
|
|
// Column names must be captured before the mutable borrows below.
|
|
let columns: Vec<String> = stmt.column_names().into_iter().map(String::from).collect();
|
|
bind_params(&mut stmt, ¶ms)?;
|
|
|
|
let mut out = Vec::new();
|
|
let mut rows = stmt.raw_query();
|
|
while let Some(row) = rows.next()? {
|
|
let mut record = Vec::with_capacity(columns.len());
|
|
for (i, name) in columns.iter().enumerate() {
|
|
record.push((name.clone(), row.get::<_, Value>(i)?));
|
|
}
|
|
out.push(record);
|
|
}
|
|
Ok(out)
|
|
})
|
|
.await
|
|
.unwrap_or_else(|e| Err(DbError::Other(format!("sqlite: background task failed: {e}"))))
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// UserData methods
|
|
// ---------------------------------------------------------------------------
|
|
|
|
impl UserData for SqliteConn {
|
|
fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
|
|
// I/O-bound methods run on a blocking thread so SQLite never stalls the
|
|
// tokio event loop.
|
|
methods.add_async_method("execute", |_, this, (sql, params): (String, Option<LuaTable>)| {
|
|
let handle = this.0.clone();
|
|
let bound = convert_params(params);
|
|
async move {
|
|
let changed = run_execute(handle, sql, bound?).await.map_err(mlua::Error::external)?;
|
|
Ok(changed as i64)
|
|
}
|
|
});
|
|
|
|
methods.add_async_method("executeBatch", |_, this, sql: String| {
|
|
let handle = this.0.clone();
|
|
async move {
|
|
run_execute_batch(handle, sql).await.map_err(mlua::Error::external)?;
|
|
Ok(())
|
|
}
|
|
});
|
|
|
|
methods.add_async_method("query", |lua, this, (sql, params): (String, Option<LuaTable>)| {
|
|
let handle = this.0.clone();
|
|
let bound = convert_params(params);
|
|
async move {
|
|
let rows = run_query(handle, sql, bound?).await.map_err(mlua::Error::external)?;
|
|
rows_to_lua(&lua, rows)
|
|
}
|
|
});
|
|
|
|
methods.add_async_method("queryOne", |lua, this, (sql, params): (String, Option<LuaTable>)| {
|
|
let handle = this.0.clone();
|
|
let bound = convert_params(params);
|
|
async move {
|
|
let mut rows = run_query(handle, sql, bound?).await.map_err(mlua::Error::external)?;
|
|
if rows.is_empty() {
|
|
Ok(LuaValue::Nil)
|
|
} else {
|
|
Ok(LuaValue::Table(record_to_lua(&lua, rows.swap_remove(0))?))
|
|
}
|
|
}
|
|
});
|
|
|
|
// These only read in-memory connection state, but they still contend on
|
|
// the same mutex that a running query holds for its whole duration. Doing
|
|
// the lock on the blocking pool (not inline) keeps the tokio event loop
|
|
// responsive: a `close()` next to a long query no longer freezes every
|
|
// other coroutine and timer.
|
|
methods.add_async_method("lastInsertRowid", |_, this, ()| {
|
|
let handle = this.0.clone();
|
|
async move {
|
|
tokio::task::spawn_blocking(move || -> Result<i64, DbError> {
|
|
let guard = handle.lock().unwrap_or_else(|e| e.into_inner());
|
|
Ok(guard.as_ref().ok_or(DbError::Closed)?.last_insert_rowid())
|
|
})
|
|
.await
|
|
.unwrap_or_else(|e| Err(DbError::Other(format!("sqlite: background task failed: {e}"))))
|
|
.map_err(mlua::Error::external)
|
|
}
|
|
});
|
|
|
|
methods.add_async_method("changes", |_, this, ()| {
|
|
let handle = this.0.clone();
|
|
async move {
|
|
tokio::task::spawn_blocking(move || -> Result<i64, DbError> {
|
|
let guard = handle.lock().unwrap_or_else(|e| e.into_inner());
|
|
Ok(guard.as_ref().ok_or(DbError::Closed)?.changes() as i64)
|
|
})
|
|
.await
|
|
.unwrap_or_else(|e| Err(DbError::Other(format!("sqlite: background task failed: {e}"))))
|
|
.map_err(mlua::Error::external)
|
|
}
|
|
});
|
|
|
|
methods.add_async_method("close", |_, this, ()| {
|
|
let handle = this.0.clone();
|
|
async move {
|
|
// Dropping the connection releases the SQLite handle; later use errors.
|
|
// Done on the blocking pool so it waits behind an in-flight query
|
|
// there rather than stalling the event loop.
|
|
tokio::task::spawn_blocking(move || {
|
|
*handle.lock().unwrap_or_else(|e| e.into_inner()) = None;
|
|
})
|
|
.await
|
|
.map_err(|e| mlua::Error::external(format!("sqlite: background task failed: {e}")))
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Installation
|
|
// ---------------------------------------------------------------------------
|
|
|
|
pub(super) fn install(lua: &Lua) -> LuaResult<()> {
|
|
let sqlite = lua.create_table()?;
|
|
|
|
sqlite.set(
|
|
"connect",
|
|
lua.create_async_function(|_, path: String| async move {
|
|
// Opening a file-backed database touches disk, so do it on the
|
|
// blocking pool rather than the tokio event loop.
|
|
let conn = tokio::task::spawn_blocking(move || rusqlite::Connection::open(&path))
|
|
.await
|
|
.map_err(|e| mlua::Error::external(format!("sqlite.connect: background task failed: {e}")))?
|
|
.map_err(|e| mlua::Error::external(format!("sqlite.connect: {e}")))?;
|
|
Ok(SqliteConn(Arc::new(Mutex::new(Some(conn)))))
|
|
})?,
|
|
)?;
|
|
|
|
lua.globals().set("sqlite", sqlite)?;
|
|
|
|
// The Lua layer adds the transaction helpers (begin/commit/rollback/
|
|
// transaction) onto the connection's shared method table. mlua hides the
|
|
// metatable from Lua (`__metatable` is set), so we reach the `__index`
|
|
// method table from Rust and hand it to the Lua chunk. Every connection of
|
|
// this type shares that one table, so the helpers apply to all of them. A
|
|
// closed (None) sample handle opens no database and is cheap to discard.
|
|
let init: mlua::Function = lua.load(SQLITE_LUA).set_name("@[stdlib/sqlite]").eval()?;
|
|
let sample = lua.create_userdata(SqliteConn(Arc::new(Mutex::new(None))))?;
|
|
let methods: LuaTable = sample.metatable()?.get("__index")?;
|
|
init.call::<()>(methods)
|
|
}
|
|
|