A lexicon-driven AppView for ATProto.
happyview.dev
backfill
firehose
jetstream
atproto
appview
oauth
lexicon
feat: add `db.raw` for complex queries (read-only)
+112
-1
+112
-1
src/lua/db_api.rs
+112
-1
src/lua/db_api.rs
···
1
1
use mlua::{Lua, LuaSerdeExt, Result as LuaResult};
2
2
use serde_json::{Value, json};
3
+
use sqlx::{Column, Row};
3
4
use std::sync::Arc;
4
5
5
6
use crate::AppState;
···
152
153
db_table.set("search", search_fn)?;
153
154
154
155
// db.count(collection, did?) -> integer
155
-
let state_count = state;
156
+
let state_count = state.clone();
156
157
let count_fn =
157
158
lua.create_async_function(move |_, (collection, did): (String, Option<String>)| {
158
159
let state = state_count.clone();
···
177
178
}
178
179
})?;
179
180
db_table.set("count", count_fn)?;
181
+
182
+
// db.raw(sql, params?) -> rows[]
183
+
// Read-only: only SELECT statements are allowed.
184
+
let state_raw = state;
185
+
let raw_fn =
186
+
lua.create_async_function(move |lua, (sql, params): (String, Option<mlua::Table>)| {
187
+
let state = state_raw.clone();
188
+
async move {
189
+
// Only allow SELECT statements
190
+
let trimmed = sql.trim_start().to_uppercase();
191
+
if !trimmed.starts_with("SELECT") {
192
+
return Err(mlua::Error::runtime("db.raw only supports SELECT queries"));
193
+
}
194
+
195
+
// Build query with dynamic parameter binding
196
+
let mut query = sqlx::query(&sql);
197
+
if let Some(ref params_table) = params {
198
+
for value in params_table.sequence_values::<mlua::Value>() {
199
+
let value = value?;
200
+
query = match value {
201
+
mlua::Value::String(s) => query.bind(s.to_str()?.to_string()),
202
+
mlua::Value::Integer(n) => query.bind(n),
203
+
mlua::Value::Number(n) => query.bind(n),
204
+
mlua::Value::Boolean(b) => query.bind(b),
205
+
mlua::Value::Nil => query.bind(Option::<String>::None),
206
+
other => {
207
+
return Err(mlua::Error::runtime(format!(
208
+
"unsupported parameter type: {}",
209
+
other.type_name()
210
+
)));
211
+
}
212
+
};
213
+
}
214
+
}
215
+
216
+
let rows = query
217
+
.fetch_all(&state.db)
218
+
.await
219
+
.map_err(|e| mlua::Error::runtime(format!("db.raw query failed: {e}")))?;
220
+
221
+
// Convert rows to Lua tables
222
+
let mut lua_rows: Vec<mlua::Value> = Vec::with_capacity(rows.len());
223
+
for row in &rows {
224
+
let row_table = lua.create_table()?;
225
+
for col in row.columns() {
226
+
let name = col.name();
227
+
let type_name = col.type_info().to_string();
228
+
let lua_val: mlua::Value = match type_name.as_str() {
229
+
"TEXT" | "VARCHAR" | "CHAR" | "NAME" | "BPCHAR" => {
230
+
match row.try_get::<Option<String>, _>(name) {
231
+
Ok(Some(s)) => mlua::Value::String(lua.create_string(&s)?),
232
+
_ => mlua::Value::Nil,
233
+
}
234
+
}
235
+
"INT4" | "INT2" | "SERIAL" => {
236
+
match row.try_get::<Option<i32>, _>(name) {
237
+
Ok(Some(n)) => mlua::Value::Integer(n as i64),
238
+
_ => mlua::Value::Nil,
239
+
}
240
+
}
241
+
"INT8" | "BIGSERIAL" | "BIGINT" => {
242
+
match row.try_get::<Option<i64>, _>(name) {
243
+
Ok(Some(n)) => mlua::Value::Integer(n),
244
+
_ => mlua::Value::Nil,
245
+
}
246
+
}
247
+
"FLOAT4" => match row.try_get::<Option<f32>, _>(name) {
248
+
Ok(Some(n)) => mlua::Value::Number(n as f64),
249
+
_ => mlua::Value::Nil,
250
+
},
251
+
"FLOAT8" | "NUMERIC" => match row.try_get::<Option<f64>, _>(name) {
252
+
Ok(Some(n)) => mlua::Value::Number(n),
253
+
_ => mlua::Value::Nil,
254
+
},
255
+
"BOOL" => match row.try_get::<Option<bool>, _>(name) {
256
+
Ok(Some(b)) => mlua::Value::Boolean(b),
257
+
_ => mlua::Value::Nil,
258
+
},
259
+
"JSON" | "JSONB" => match row.try_get::<Option<Value>, _>(name) {
260
+
Ok(Some(v)) => lua.to_value(&v)?,
261
+
_ => mlua::Value::Nil,
262
+
},
263
+
"TIMESTAMPTZ" | "TIMESTAMP" => {
264
+
match row.try_get::<Option<chrono::DateTime<chrono::Utc>>, _>(name)
265
+
{
266
+
Ok(Some(dt)) => {
267
+
mlua::Value::String(lua.create_string(dt.to_rfc3339())?)
268
+
}
269
+
_ => mlua::Value::Nil,
270
+
}
271
+
}
272
+
_ => {
273
+
// Fall back to trying as a string
274
+
match row.try_get::<Option<String>, _>(name) {
275
+
Ok(Some(s)) => mlua::Value::String(lua.create_string(&s)?),
276
+
_ => mlua::Value::Nil,
277
+
}
278
+
}
279
+
};
280
+
row_table.set(name, lua_val)?;
281
+
}
282
+
lua_rows.push(mlua::Value::Table(row_table));
283
+
}
284
+
285
+
let result = lua.create_sequence_from(lua_rows)?;
286
+
result.set_metatable(Some(lua.array_metatable()))?;
287
+
Ok(result)
288
+
}
289
+
})?;
290
+
db_table.set("raw", raw_fn)?;
180
291
181
292
lua.globals().set("db", db_table)?;
182
293
Ok(())