Something went wrong. Try again.
A lexicon-driven AppView for ATProto.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517use mlua::{Lua, LuaSerdeExt, Result as LuaResult};use serde_json::{Value, json};use sqlx::{Column, Row};use std::sync::Arc;
use crate::AppState;
/// Register the `db` table with database query functions.pub fn register_db_api(lua: &Lua, state: Arc<AppState>) -> LuaResult<()> { let db_table = lua.create_table()?;
// db.query({ collection, did?, limit?, offset?, sort?, sortDirection? }) -> { records, cursor? } let state_query = state.clone(); let query_fn = lua.create_async_function(move |lua, opts: mlua::Table| { let state = state_query.clone(); async move { let collection: String = opts.get("collection")?; let did: Option<String> = opts.get("did").ok(); let limit: i64 = opts.get::<i64>("limit").unwrap_or(20).min(100); let offset: i64 = opts.get::<i64>("offset").unwrap_or(0); let sort: Option<String> = opts.get("sort").ok(); let sort_direction: Option<String> = opts.get("sortDirection").ok();
// Validate sort field name to prevent SQL injection if let Some(ref field) = sort { let valid = field.chars().all(|c| c.is_ascii_alphanumeric() || c == '_'); if !valid || field.is_empty() { return Err(mlua::Error::runtime( "invalid sort field: only alphanumeric characters and underscores are allowed", )); } }
let direction = match sort_direction.as_deref() { Some("asc") => "ASC", Some("desc") => "DESC", None => "DESC", Some(other) => { return Err(mlua::Error::runtime(format!( "invalid sortDirection '{other}': must be 'asc' or 'desc'" ))); } };
let top_level_columns = ["indexed_at", "did", "uri"]; let order_expr = match &sort { Some(field) if top_level_columns.contains(&field.as_str()) => { format!("{field} {direction}") } Some(field) => { format!("record->'value'->>'{field}' {direction}") } None => format!("indexed_at {direction}"), };
let rows: Vec<(String, String, Value)> = if let Some(ref did) = did { sqlx::query_as( &format!("SELECT uri, did, record FROM records WHERE collection = $1 AND did = $2 ORDER BY {order_expr} LIMIT $3 OFFSET $4"), ) .bind(&collection) .bind(did) .bind(limit) .bind(offset) .fetch_all(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("DB query failed: {e}")))? } else { sqlx::query_as( &format!("SELECT uri, did, record FROM records WHERE collection = $1 ORDER BY {order_expr} LIMIT $2 OFFSET $3"), ) .bind(&collection) .bind(limit) .bind(offset) .fetch_all(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("DB query failed: {e}")))? };
let has_next = rows.len() as i64 == limit; let records: Vec<Value> = rows .into_iter() .map(|(uri, _did, mut record)| { if let Some(obj) = record.as_object_mut() { obj.insert("uri".to_string(), json!(uri)); } record }) .collect();
let record_values: Vec<mlua::Value> = records .iter() .map(|r| lua.to_value(r)) .collect::<LuaResult<_>>()?; let records_table = lua.create_sequence_from(record_values)?; records_table.set_metatable(Some(lua.array_metatable()))?;
let result_table = lua.create_table()?; result_table.set("records", records_table)?; if has_next { let next_cursor = (offset + limit).to_string(); result_table.set("cursor", next_cursor)?; }
Ok(mlua::Value::Table(result_table)) } })?; db_table.set("query", query_fn)?;
// db.get(uri) -> record table or nil let state_get = state.clone(); let get_fn = lua.create_async_function(move |lua, uri: String| { let state = state_get.clone(); async move { let row: Option<(Value,)> = sqlx::query_as("SELECT record FROM records WHERE uri = $1") .bind(&uri) .fetch_optional(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("DB query failed: {e}")))?;
match row { Some((mut record,)) => { if let Some(obj) = record.as_object_mut() { obj.insert("uri".to_string(), json!(uri)); } lua.to_value(&record) } None => Ok(mlua::Value::Nil), } } })?; db_table.set("get", get_fn)?;
// db.search({ collection, field, query, limit? }) -> { records } let state_search = state.clone(); let search_fn = lua.create_async_function(move |lua, opts: mlua::Table| { let state = state_search.clone(); async move { let collection: String = opts.get("collection")?; let field: String = opts.get("field")?; let query: String = opts.get("query")?; let limit: i64 = opts.get::<i64>("limit").unwrap_or(10).min(100);
let rows: Vec<(String, String, Value)> = sqlx::query_as( "SELECT uri, did, record FROM records \ WHERE collection = $1 \ AND record->>$2 ILIKE '%' || $3 || '%' \ ORDER BY \ CASE \ WHEN LOWER(record->>$2) = LOWER($3) THEN 0 \ WHEN LOWER(record->>$2) LIKE LOWER($3) || '%' THEN 1 \ ELSE 2 \ END, \ record->>$2 \ LIMIT $4", ) .bind(&collection) .bind(&field) .bind(&query) .bind(limit) .fetch_all(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("DB search failed: {e}")))?;
let records: Vec<Value> = rows .into_iter() .map(|(uri, _did, mut record)| { if let Some(obj) = record.as_object_mut() { obj.insert("uri".to_string(), json!(uri)); } record }) .collect();
let record_values: Vec<mlua::Value> = records .iter() .map(|r| lua.to_value(r)) .collect::<LuaResult<_>>()?; let records_table = lua.create_sequence_from(record_values)?; records_table.set_metatable(Some(lua.array_metatable()))?;
let result_table = lua.create_table()?; result_table.set("records", records_table)?;
Ok(mlua::Value::Table(result_table)) } })?; db_table.set("search", search_fn)?;
// db.count(collection, did?) -> integer let state_count = state.clone(); let count_fn = lua.create_async_function(move |_, (collection, did): (String, Option<String>)| { let state = state_count.clone(); async move { let count: (i64,) = if let Some(ref did) = did { sqlx::query_as( "SELECT COUNT(*) FROM records WHERE collection = $1 AND did = $2", ) .bind(&collection) .bind(did) .fetch_one(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("DB count failed: {e}")))? } else { sqlx::query_as("SELECT COUNT(*) FROM records WHERE collection = $1") .bind(&collection) .fetch_one(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("DB count failed: {e}")))? }; Ok(count.0) } })?; db_table.set("count", count_fn)?;
// db.backlinks({ collection, uri, did?, limit?, offset? }) -> { records, cursor? } // Find records in `collection` whose JSONB contains the given AT URI. let state_backlinks = state.clone(); let backlinks_fn = lua.create_async_function(move |lua, opts: mlua::Table| { let state = state_backlinks.clone(); async move { let collection: String = opts.get("collection")?; let uri: String = opts.get("uri")?; let did: Option<String> = opts.get("did").ok(); let limit: i64 = opts.get::<i64>("limit").unwrap_or(20).min(100); let offset: i64 = opts.get::<i64>("offset").unwrap_or(0);
let rows: Vec<(String, String, Value)> = if let Some(ref did) = did { sqlx::query_as( "SELECT uri, did, record FROM records \ WHERE collection = $1 \ AND record::text LIKE '%' || $2 || '%' \ AND did = $3 \ ORDER BY indexed_at DESC \ LIMIT $4 OFFSET $5", ) .bind(&collection) .bind(&uri) .bind(did) .bind(limit) .bind(offset) .fetch_all(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("DB backlinks failed: {e}")))? } else { sqlx::query_as( "SELECT uri, did, record FROM records \ WHERE collection = $1 \ AND record::text LIKE '%' || $2 || '%' \ ORDER BY indexed_at DESC \ LIMIT $3 OFFSET $4", ) .bind(&collection) .bind(&uri) .bind(limit) .bind(offset) .fetch_all(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("DB backlinks failed: {e}")))? };
let has_next = rows.len() as i64 == limit; let records: Vec<Value> = rows .into_iter() .map(|(uri, _did, mut record)| { if let Some(obj) = record.as_object_mut() { obj.insert("uri".to_string(), json!(uri)); } record }) .collect();
let record_values: Vec<mlua::Value> = records .iter() .map(|r| lua.to_value(r)) .collect::<LuaResult<_>>()?; let records_table = lua.create_sequence_from(record_values)?; records_table.set_metatable(Some(lua.array_metatable()))?;
let result_table = lua.create_table()?; result_table.set("records", records_table)?; if has_next { let next_cursor = (offset + limit).to_string(); result_table.set("cursor", next_cursor)?; }
Ok(mlua::Value::Table(result_table)) } })?; db_table.set("backlinks", backlinks_fn)?;
// db.raw(sql, params?) -> rows[] let state_raw = state; let raw_fn = lua.create_async_function(move |lua, (sql, params): (String, Option<mlua::Table>)| { let state = state_raw.clone(); async move { // Build query with dynamic parameter binding let mut query = sqlx::query(&sql); if let Some(ref params_table) = params { for value in params_table.sequence_values::<mlua::Value>() { let value = value?; query = match value { mlua::Value::String(s) => query.bind(s.to_str()?.to_string()), mlua::Value::Integer(n) => query.bind(n), mlua::Value::Number(n) => query.bind(n), mlua::Value::Boolean(b) => query.bind(b), mlua::Value::Nil => query.bind(Option::<String>::None), other => { return Err(mlua::Error::runtime(format!( "unsupported parameter type: {}", other.type_name() ))); } }; } }
let rows = query .fetch_all(&state.db) .await .map_err(|e| mlua::Error::runtime(format!("db.raw query failed: {e}")))?;
// Convert rows to Lua tables let mut lua_rows: Vec<mlua::Value> = Vec::with_capacity(rows.len()); for row in &rows { let row_table = lua.create_table()?; for col in row.columns() { let name = col.name(); let type_name = col.type_info().to_string(); let lua_val: mlua::Value = match type_name.as_str() { "TEXT" | "VARCHAR" | "CHAR" | "NAME" | "BPCHAR" => { match row.try_get::<Option<String>, _>(name) { Ok(Some(s)) => mlua::Value::String(lua.create_string(&s)?), _ => mlua::Value::Nil, } } "INT4" | "INT2" | "SERIAL" => { match row.try_get::<Option<i32>, _>(name) { Ok(Some(n)) => mlua::Value::Integer(n as i64), _ => mlua::Value::Nil, } } "INT8" | "BIGSERIAL" | "BIGINT" => { match row.try_get::<Option<i64>, _>(name) { Ok(Some(n)) => mlua::Value::Integer(n), _ => mlua::Value::Nil, } } "FLOAT4" => match row.try_get::<Option<f32>, _>(name) { Ok(Some(n)) => mlua::Value::Number(n as f64), _ => mlua::Value::Nil, }, "FLOAT8" | "NUMERIC" => match row.try_get::<Option<f64>, _>(name) { Ok(Some(n)) => mlua::Value::Number(n), _ => mlua::Value::Nil, }, "BOOL" => match row.try_get::<Option<bool>, _>(name) { Ok(Some(b)) => mlua::Value::Boolean(b), _ => mlua::Value::Nil, }, "JSON" | "JSONB" => match row.try_get::<Option<Value>, _>(name) { Ok(Some(v)) => lua.to_value(&v)?, _ => mlua::Value::Nil, }, "TIMESTAMPTZ" | "TIMESTAMP" => { match row.try_get::<Option<chrono::DateTime<chrono::Utc>>, _>(name) { Ok(Some(dt)) => { mlua::Value::String(lua.create_string(dt.to_rfc3339())?) } _ => mlua::Value::Nil, } } _ => { // Fall back to trying as a string match row.try_get::<Option<String>, _>(name) { Ok(Some(s)) => mlua::Value::String(lua.create_string(&s)?), _ => mlua::Value::Nil, } } }; row_table.set(name, lua_val)?; } lua_rows.push(mlua::Value::Table(row_table)); }
let result = lua.create_sequence_from(lua_rows)?; result.set_metatable(Some(lua.array_metatable()))?; Ok(result) } })?; db_table.set("raw", raw_fn)?;
lua.globals().set("db", db_table)?; Ok(())}
#[cfg(test)]mod tests { use super::*; use crate::config::Config; use crate::lexicon::LexiconRegistry; use tokio::sync::watch;
fn test_state() -> AppState { let config = Config { host: "127.0.0.1".into(), port: 3000, database_url: String::new(), aip_url: String::new(), aip_public_url: String::new(), tap_url: String::new(), tap_admin_password: None, relay_url: String::new(), plc_url: String::new(), static_dir: String::new(), event_log_retention_days: 30, }; let (tx, _) = watch::channel(vec![]); AppState { config, http: reqwest::Client::new(), db: sqlx::PgPool::connect_lazy("postgres://localhost/fake").unwrap(), lexicons: LexiconRegistry::new(), collections_tx: tx, } }
fn setup(state: &AppState) -> Lua { let lua = Lua::new(); register_db_api(&lua, Arc::new(state.clone())).unwrap(); lua }
#[tokio::test] async fn raw_allows_non_select() { let state = test_state(); let lua = setup(&state); let result: Result<mlua::Value, _> = lua .load(r#"return db.raw("DELETE FROM records")"#) .eval_async() .await; // Should fail with a DB connection error, NOT a validation error assert!(result.is_err()); let err = result.unwrap_err().to_string(); assert!( !err.contains("only supports SELECT"), "should have passed validation but got: {err}" ); }
#[tokio::test] async fn raw_allows_select() { let state = test_state(); let lua = setup(&state); let result: Result<mlua::Value, _> = lua.load(r#"return db.raw("SELECT 1")"#).eval_async().await; // Should fail with a DB connection error, NOT a validation error assert!(result.is_err()); let err = result.unwrap_err().to_string(); assert!( !err.contains("only supports SELECT"), "should have passed validation but got: {err}" ); }
#[tokio::test] async fn query_rejects_invalid_sort_field() { let state = test_state(); let lua = setup(&state); let result: Result<mlua::Value, _> = lua .load(r#"return db.query({ collection = "test", sort = "name; DROP TABLE" })"#) .eval_async() .await; assert!(result.is_err()); let err = result.unwrap_err().to_string(); assert!( err.contains("invalid sort field"), "expected sort field error, got: {err}" ); }
#[tokio::test] async fn query_rejects_invalid_sort_direction() { let state = test_state(); let lua = setup(&state); let result: Result<mlua::Value, _> = lua .load(r#"return db.query({ collection = "test", sortDirection = "sideways" })"#) .eval_async() .await; assert!(result.is_err()); let err = result.unwrap_err().to_string(); assert!( err.contains("invalid sortDirection"), "expected sortDirection error, got: {err}" ); }
#[tokio::test] async fn query_accepts_valid_sort_direction() { let state = test_state(); let lua = setup(&state); let result: Result<mlua::Value, _> = lua .load(r#"return db.query({ collection = "test", sortDirection = "asc" })"#) .eval_async() .await; // Should fail with a DB connection error, NOT a validation error assert!(result.is_err()); let err = result.unwrap_err().to_string(); assert!( !err.contains("invalid sortDirection"), "should have passed validation but got: {err}" ); }}