mlua_batteries_sqlite/kv.rs
1//! `std.kv` — SQLite-backed key-value store for Lua scripts.
2//!
3//! Storage lives in a SQLite database supplied by the host (one shared
4//! connection), in a dedicated `__kv` table:
5//!
6//! ```sql
7//! CREATE TABLE __kv (
8//! ns TEXT NOT NULL,
9//! key TEXT NOT NULL,
10//! value TEXT NOT NULL, -- JSON-serialized Lua value
11//! PRIMARY KEY (ns, key)
12//! ) WITHOUT ROWID;
13//! ```
14//!
15//! Trade-offs vs. a JSON-file-per-namespace implementation:
16//! - Per-key updates (no whole-namespace rewrite on every set).
17//! - Durability + atomicity delegated to SQLite's WAL journal.
18//! - Cross-process writes arbitrated by `busy_timeout`.
19//!
20//! # Wiring contract
21//!
22//! Symmetric with [`crate::sql`]. The host opens a `rusqlite::Connection`
23//! (typically a database file dedicated to KV scratch state, kept separate
24//! from the `std.sql` user database so backup / WAL / page-cache lifecycles
25//! do not collide) and passes it as `Arc<Mutex<_>>` plus its
26//! `InterruptHandle`. Cancellation and per-query timeout are inherited
27//! from the [`crate::sql::SqlConfig`] in `lua.app_data`; the `rusqlite`
28//! those types come from is [`crate::rusqlite`].
29
30use std::sync::Arc;
31
32use mlua::prelude::*;
33
34use mlua_batteries::json::{
35 array_metatable, json_to_lua_preserving_null, lua_to_json_preserving_null,
36};
37
38use crate::sql::{lock_conn, race_timeout, sql_query_timeout, SqlConfig};
39use crate::sqlite_backend::rusqlite::{self, Connection, InterruptHandle, OptionalExtension};
40
41// ---------------------------------------------------------------------------
42// Helpers
43// ---------------------------------------------------------------------------
44
45/// Validate a namespace string.
46///
47/// Namespaces were originally used as filenames, so `/`, `\`, `..`, `\0` were
48/// rejected for path-traversal safety. Even though storage is now a SQL table
49/// (and namespaces are just column values), we keep the same validation so
50/// that existing Lua scripts and tests see identical semantics.
51fn validate_ns(ns: &str) -> Result<(), String> {
52 if ns.is_empty() {
53 return Err(format!("Invalid namespace: '{ns}'"));
54 }
55 if ns.contains('/') || ns.contains('\\') || ns.contains('\0') || ns.contains("..") {
56 return Err(format!("Invalid namespace: '{ns}'"));
57 }
58 Ok(())
59}
60
61/// Create the `__kv` table if it is not there yet.
62///
63/// [`register`] / [`register_with`] run this on the supplied connection
64/// before exposing `std.kv`, so hosts do not have to. It is public for
65/// hosts that prefer to own schema setup themselves — right after opening
66/// the connection, before it is wrapped in `Arc<Mutex<_>>`:
67///
68/// ```rust,ignore
69/// let conn = Connection::open(kv_path)?;
70/// mlua_batteries_sqlite::kv::init_schema(&conn)?;
71/// ```
72///
73/// The DDL is `CREATE TABLE IF NOT EXISTS`, so running it twice is harmless.
74pub fn init_schema(conn: &Connection) -> Result<(), rusqlite::Error> {
75 conn.execute_batch(
76 "CREATE TABLE IF NOT EXISTS __kv (\n ns TEXT NOT NULL,\n key TEXT NOT NULL,\n value TEXT NOT NULL,\n PRIMARY KEY (ns, key)\n ) WITHOUT ROWID;",
77 )
78}
79
80// ---------------------------------------------------------------------------
81// Registration
82// ---------------------------------------------------------------------------
83
84/// Register `std.kv` with default [`SqlConfig`] (only used if `std.sql` was
85/// not registered first; otherwise the existing config is preserved).
86pub fn register(
87 lua: &Lua,
88 conn: Arc<std::sync::Mutex<Connection>>,
89 interrupt: Arc<InterruptHandle>,
90) -> LuaResult<()> {
91 register_with(lua, conn, interrupt, SqlConfig::default())
92}
93
94/// Register `std.kv` with caller-provided [`SqlConfig`].
95///
96/// If `std.sql` was registered earlier with a `SqlConfig`, the same slot
97/// in `lua.app_data` is overwritten — pass an identical config from the
98/// host to keep `sql` and `kv` in sync (the typical case).
99pub fn register_with(
100 lua: &Lua,
101 conn: Arc<std::sync::Mutex<Connection>>,
102 interrupt: Arc<InterruptHandle>,
103 cfg: SqlConfig,
104) -> LuaResult<()> {
105 lua.set_app_data::<SqlConfig>(cfg);
106
107 // One-time schema init on the supplied connection.
108 {
109 let guard = lock_conn(&conn);
110 init_schema(&guard).map_err(|e| LuaError::external(format!("kv schema init: {e}")))?;
111 }
112
113 let kv_tbl = lua.create_table()?;
114
115 // ── std.kv.get ────────────────────────────────────────────────────────
116 {
117 let conn = Arc::clone(&conn);
118 let interrupt = Arc::clone(&interrupt);
119 kv_tbl.set(
120 "get",
121 lua.create_async_function(move |lua, (ns, key): (String, String)| {
122 let conn = Arc::clone(&conn);
123 let interrupt = Arc::clone(&interrupt);
124 let ns_check = validate_ns(&ns).map_err(LuaError::external);
125 async move {
126 ns_check?;
127 let fut = tokio::task::spawn_blocking(move || {
128 let guard = lock_conn(&conn);
129 guard
130 .query_row(
131 "SELECT value FROM __kv WHERE ns = ?1 AND key = ?2",
132 rusqlite::params![ns, key],
133 |row| row.get::<_, String>(0),
134 )
135 .optional()
136 .map_err(|e| format!("kv.get sql error: {e}"))
137 });
138 let timeout = sql_query_timeout(&lua);
139 let row = race_timeout(fut, timeout, &interrupt, "kv.get").await?;
140 match row {
141 None => Ok(LuaValue::Nil),
142 Some(s) => {
143 let v: serde_json::Value = serde_json::from_str(&s).map_err(|e| {
144 LuaError::external(format!("kv.get json parse: {e}"))
145 })?;
146 json_to_lua_preserving_null(&lua, v)
147 }
148 }
149 }
150 })?,
151 )?;
152 }
153
154 // ── std.kv.set ────────────────────────────────────────────────────────
155 {
156 let conn = Arc::clone(&conn);
157 let interrupt = Arc::clone(&interrupt);
158 kv_tbl.set(
159 "set",
160 lua.create_async_function(move |lua, (ns, key, value): (String, String, LuaValue)| {
161 let conn = Arc::clone(&conn);
162 let interrupt = Arc::clone(&interrupt);
163 // Serialize synchronously on the Lua thread (LuaValue is
164 // !Send, so it can't cross the spawn_blocking boundary).
165 let ns_check = validate_ns(&ns).map_err(LuaError::external);
166 let json_result = lua_to_json_preserving_null(value).and_then(|v| {
167 serde_json::to_string(&v)
168 .map_err(|e| LuaError::external(format!("kv.set serialize: {e}")))
169 });
170 async move {
171 ns_check?;
172 let json_str = json_result?;
173 let fut = tokio::task::spawn_blocking(move || {
174 let guard = lock_conn(&conn);
175 guard
176 .execute(
177 "INSERT INTO __kv (ns, key, value) VALUES (?1, ?2, ?3) \
178 ON CONFLICT(ns, key) DO UPDATE SET value = excluded.value",
179 rusqlite::params![ns, key, json_str],
180 )
181 .map(|_| ())
182 .map_err(|e| format!("kv.set sql error: {e}"))
183 });
184 let timeout = sql_query_timeout(&lua);
185 race_timeout(fut, timeout, &interrupt, "kv.set").await
186 }
187 })?,
188 )?;
189 }
190
191 // ── std.kv.delete ─────────────────────────────────────────────────────
192 {
193 let conn = Arc::clone(&conn);
194 let interrupt = Arc::clone(&interrupt);
195 kv_tbl.set(
196 "delete",
197 lua.create_async_function(move |lua, (ns, key): (String, String)| {
198 let conn = Arc::clone(&conn);
199 let interrupt = Arc::clone(&interrupt);
200 let ns_check = validate_ns(&ns).map_err(LuaError::external);
201 async move {
202 ns_check?;
203 let fut = tokio::task::spawn_blocking(move || {
204 let guard = lock_conn(&conn);
205 guard
206 .execute(
207 "DELETE FROM __kv WHERE ns = ?1 AND key = ?2",
208 rusqlite::params![ns, key],
209 )
210 .map(|n| n > 0)
211 .map_err(|e| format!("kv.delete sql error: {e}"))
212 });
213 let timeout = sql_query_timeout(&lua);
214 race_timeout(fut, timeout, &interrupt, "kv.delete").await
215 }
216 })?,
217 )?;
218 }
219
220 // ── std.kv.list ───────────────────────────────────────────────────────
221 {
222 let conn = Arc::clone(&conn);
223 let interrupt = Arc::clone(&interrupt);
224 kv_tbl.set(
225 "list",
226 lua.create_async_function(move |lua, (ns, prefix): (String, Option<String>)| {
227 let conn = Arc::clone(&conn);
228 let interrupt = Arc::clone(&interrupt);
229 let ns_check = validate_ns(&ns).map_err(LuaError::external);
230 async move {
231 ns_check?;
232 let fut = tokio::task::spawn_blocking(move || {
233 let guard = lock_conn(&conn);
234 let mut stmt = guard
235 .prepare("SELECT key FROM __kv WHERE ns = ?1 ORDER BY key")
236 .map_err(|e| format!("kv.list prepare: {e}"))?;
237 let keys: Vec<String> = stmt
238 .query_map(rusqlite::params![ns], |row| row.get::<_, String>(0))
239 .map_err(|e| format!("kv.list query: {e}"))?
240 .collect::<Result<_, _>>()
241 .map_err(|e| format!("kv.list row: {e}"))?;
242 Ok::<_, String>(keys)
243 });
244 let timeout = sql_query_timeout(&lua);
245 let keys = race_timeout(fut, timeout, &interrupt, "kv.list").await?;
246
247 let tbl = lua.create_table()?;
248 let mut idx = 1usize;
249 for k in keys {
250 let include = prefix.as_deref().map_or(true, |p| k.starts_with(p));
251 if include {
252 tbl.set(idx, k.as_str())?;
253 idx += 1;
254 }
255 }
256 if idx == 1 {
257 // Nothing matched: tag the empty list with the shared
258 // `__jsontype = "array"` metatable so that
259 // `json.encode(kv.list(...))` renders `[]`, not `{}`.
260 tbl.set_metatable(Some(array_metatable(&lua)?))?;
261 }
262 Ok(LuaValue::Table(tbl))
263 }
264 })?,
265 )?;
266 }
267
268 let std_ns: LuaTable = lua.globals().get("std")?;
269 std_ns.set("kv", kv_tbl)?;
270 Ok(())
271}