diff --git a/migrations/20260312000000_create_record_refs.sql b/migrations/20260312000000_create_record_refs.sql new file mode 100644 --- /dev/null +++ b/migrations/20260312000000_create_record_refs.sql @@ -0,0 +1,9 @@ +CREATE TABLE IF NOT EXISTS record_refs ( + source_uri TEXT NOT NULL REFERENCES records(uri) ON DELETE CASCADE, + target_uri TEXT NOT NULL, + collection TEXT NOT NULL, + PRIMARY KEY (source_uri, target_uri) +); + +CREATE INDEX IF NOT EXISTS idx_record_refs_target ON record_refs (target_uri, collection); +CREATE INDEX IF NOT EXISTS idx_records_created_at_uri ON records (created_at DESC, uri DESC); diff --git a/src/lib.rs b/src/lib.rs --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ pub mod event_log; pub mod lexicon; pub mod lua; pub mod profile; +pub mod record_refs; pub mod repo; pub mod resolve; pub mod server; diff --git a/src/lua/db_api.rs b/src/lua/db_api.rs --- a/src/lua/db_api.rs +++ b/src/lua/db_api.rs @@ -1,3 +1,4 @@ +use base64::{Engine, engine::general_purpose::STANDARD as BASE64}; use mlua::{Lua, LuaSerdeExt, Result as LuaResult}; use serde_json::{Value, json}; use sqlx::{Column, Row}; @@ -5,11 +6,24 @@ use std::sync::Arc; use crate::AppState; +/// Encode a cursor from created_at timestamp and uri. +fn encode_cursor(created_at: &str, uri: &str) -> String { + BASE64.encode(format!("{created_at}|{uri}")) +} + +/// Decode a cursor into (created_at, uri). Returns None if invalid. +fn decode_cursor(cursor: &str) -> Option<(String, String)> { + let decoded = BASE64.decode(cursor).ok()?; + let s = String::from_utf8(decoded).ok()?; + let (ts, uri) = s.split_once('|')?; + Some((ts.to_string(), uri.to_string())) +} + /// Register the `db` table with database query functions. pub fn register_db_api(lua: &Lua, state: Arc) -> LuaResult<()> { let db_table = lua.create_table()?; - // db.query({ collection, did?, limit?, offset?, sort?, sortDirection? }) -> { records, cursor? } + // db.query({ collection, did?, limit?, offset?, cursor?, 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(); @@ -17,9 +31,9 @@ async move { let collection: String = opts.get("collection")?; let did: Option = opts.get("did").ok(); let limit: i64 = opts.get::("limit").unwrap_or(20).min(100); - let offset: i64 = opts.get::("offset").unwrap_or(0); let sort: Option = opts.get("sort").ok(); let sort_direction: Option = opts.get("sortDirection").ok(); + let cursor_str: Option = opts.get("cursor").ok(); // Validate sort field name to prevent SQL injection if let Some(ref field) = sort { @@ -42,63 +56,170 @@ ))); } }; - 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}") + let result_table = lua.create_table()?; + + if let Some(ref sort_field) = sort { + // Custom sort: use OFFSET/LIMIT with base64-encoded offset cursor + let offset: i64 = if let Some(ref cursor) = cursor_str { + BASE64.decode(cursor).ok() + .and_then(|b| String::from_utf8(b).ok()) + .and_then(|s| s.parse::().ok()) + .unwrap_or(0) + } else { + opts.get::("offset").unwrap_or(0) + }; + + let top_level_columns = ["indexed_at", "did", "uri"]; + let order_expr = if top_level_columns.contains(&sort_field.as_str()) { + format!("{sort_field} {direction}") + } else { + format!("record->'value'->>'{sort_field}' {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; + + if has_next { + let next_offset = offset + limit; + result_table.set("cursor", BASE64.encode(next_offset.to_string()))?; } - 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}")))? + let records: Vec = 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 = records + .iter() + .map(|r| lua.to_value(r)) + .collect::>()?; + let records_table = lua.create_sequence_from(record_values)?; + records_table.set_metatable(Some(lua.array_metatable()))?; + result_table.set("records", records_table)?; } 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}")))? - }; + // Cursor-based pagination on (created_at, uri) + let cursor_parts = cursor_str.as_ref().and_then(|c| decode_cursor(c)); + + type RowType = (String, String, Value, chrono::DateTime); - let has_next = rows.len() as i64 == limit; - let records: Vec = rows - .into_iter() - .map(|(uri, _did, mut record)| { - if let Some(obj) = record.as_object_mut() { - obj.insert("uri".to_string(), json!(uri)); + let rows_raw: Vec = match (&did, &cursor_parts) { + (Some(did), Some((cursor_ts, cursor_uri))) => { + let ts: chrono::DateTime = cursor_ts.parse() + .map_err(|e| mlua::Error::runtime(format!("invalid cursor timestamp: {e}")))?; + sqlx::query_as( + "SELECT uri, did, record, created_at FROM records \ + WHERE collection = $1 AND did = $2 AND (created_at, uri) < ($3, $4) \ + ORDER BY created_at DESC, uri DESC \ + LIMIT $5", + ) + .bind(&collection) + .bind(did) + .bind(ts) + .bind(cursor_uri) + .bind(limit) + .fetch_all(&state.db) + .await + .map_err(|e| mlua::Error::runtime(format!("DB query failed: {e}")))? } - record - }) - .collect(); + (Some(did), None) => { + sqlx::query_as( + "SELECT uri, did, record, created_at FROM records \ + WHERE collection = $1 AND did = $2 \ + ORDER BY created_at DESC, uri DESC \ + LIMIT $3", + ) + .bind(&collection) + .bind(did) + .bind(limit) + .fetch_all(&state.db) + .await + .map_err(|e| mlua::Error::runtime(format!("DB query failed: {e}")))? + } + (None, Some((cursor_ts, cursor_uri))) => { + let ts: chrono::DateTime = cursor_ts.parse() + .map_err(|e| mlua::Error::runtime(format!("invalid cursor timestamp: {e}")))?; + sqlx::query_as( + "SELECT uri, did, record, created_at FROM records \ + WHERE collection = $1 AND (created_at, uri) < ($2, $3) \ + ORDER BY created_at DESC, uri DESC \ + LIMIT $4", + ) + .bind(&collection) + .bind(ts) + .bind(cursor_uri) + .bind(limit) + .fetch_all(&state.db) + .await + .map_err(|e| mlua::Error::runtime(format!("DB query failed: {e}")))? + } + (None, None) => { + sqlx::query_as( + "SELECT uri, did, record, created_at FROM records \ + WHERE collection = $1 \ + ORDER BY created_at DESC, uri DESC \ + LIMIT $2", + ) + .bind(&collection) + .bind(limit) + .fetch_all(&state.db) + .await + .map_err(|e| mlua::Error::runtime(format!("DB query failed: {e}")))? + } + }; - let record_values: Vec = records - .iter() - .map(|r| lua.to_value(r)) - .collect::>()?; - let records_table = lua.create_sequence_from(record_values)?; - records_table.set_metatable(Some(lua.array_metatable()))?; + let has_next = rows_raw.len() as i64 == limit; + + if has_next + && let Some((last_uri, _, _, last_created_at)) = rows_raw.last() + { + let cursor = encode_cursor(&last_created_at.to_rfc3339(), last_uri); + result_table.set("cursor", cursor)?; + } + + let records: Vec = rows_raw + .into_iter() + .map(|(uri, _did, mut record, _created_at)| { + if let Some(obj) = record.as_object_mut() { + obj.insert("uri".to_string(), json!(uri)); + } + record + }) + .collect(); - 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)?; + let record_values: Vec = records + .iter() + .map(|r| lua.to_value(r)) + .collect::>()?; + let records_table = lua.create_sequence_from(record_values)?; + records_table.set_metatable(Some(lua.array_metatable()))?; + result_table.set("records", records_table)?; } Ok(mlua::Value::Table(result_table)) @@ -213,8 +334,8 @@ } })?; 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. + // db.backlinks({ collection, uri, did?, limit?, cursor? }) -> { records, cursor? } + // Find records in `collection` that reference the given AT URI via record_refs. let state_backlinks = state.clone(); let backlinks_fn = lua.create_async_function(move |lua, opts: mlua::Table| { let state = state_backlinks.clone(); @@ -223,46 +344,97 @@ let collection: String = opts.get("collection")?; let uri: String = opts.get("uri")?; let did: Option = opts.get("did").ok(); let limit: i64 = opts.get::("limit").unwrap_or(20).min(100); - let offset: i64 = opts.get::("offset").unwrap_or(0); + let cursor_str: Option = opts.get("cursor").ok(); + + let cursor_parts = cursor_str.as_ref().and_then(|c| decode_cursor(c)); + + type RowType = (String, String, Value, chrono::DateTime); - 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", + let rows_raw: Vec = match (&did, &cursor_parts) { + (Some(did), Some((cursor_ts, cursor_uri))) => { + let ts: chrono::DateTime = cursor_ts.parse().map_err(|e| { + mlua::Error::runtime(format!("invalid cursor timestamp: {e}")) + })?; + sqlx::query_as( + "SELECT r.uri, r.did, r.record, r.created_at FROM records r \ + INNER JOIN record_refs ref ON ref.source_uri = r.uri \ + WHERE ref.target_uri = $1 AND ref.collection = $2 AND r.did = $3 \ + AND (r.created_at, r.uri) < ($4, $5) \ + ORDER BY r.created_at DESC, r.uri DESC \ + LIMIT $6", + ) + .bind(&uri) + .bind(&collection) + .bind(did) + .bind(ts) + .bind(cursor_uri) + .bind(limit) + .fetch_all(&state.db) + .await + .map_err(|e| mlua::Error::runtime(format!("DB backlinks failed: {e}")))? + } + (Some(did), None) => sqlx::query_as( + "SELECT r.uri, r.did, r.record, r.created_at FROM records r \ + INNER JOIN record_refs ref ON ref.source_uri = r.uri \ + WHERE ref.target_uri = $1 AND ref.collection = $2 AND r.did = $3 \ + ORDER BY r.created_at DESC, r.uri DESC \ + LIMIT $4", ) - .bind(&collection) .bind(&uri) + .bind(&collection) .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", + .map_err(|e| mlua::Error::runtime(format!("DB backlinks failed: {e}")))?, + (None, Some((cursor_ts, cursor_uri))) => { + let ts: chrono::DateTime = cursor_ts.parse().map_err(|e| { + mlua::Error::runtime(format!("invalid cursor timestamp: {e}")) + })?; + sqlx::query_as( + "SELECT r.uri, r.did, r.record, r.created_at FROM records r \ + INNER JOIN record_refs ref ON ref.source_uri = r.uri \ + WHERE ref.target_uri = $1 AND ref.collection = $2 \ + AND (r.created_at, r.uri) < ($3, $4) \ + ORDER BY r.created_at DESC, r.uri DESC \ + LIMIT $5", + ) + .bind(&uri) + .bind(&collection) + .bind(ts) + .bind(cursor_uri) + .bind(limit) + .fetch_all(&state.db) + .await + .map_err(|e| mlua::Error::runtime(format!("DB backlinks failed: {e}")))? + } + (None, None) => sqlx::query_as( + "SELECT r.uri, r.did, r.record, r.created_at FROM records r \ + INNER JOIN record_refs ref ON ref.source_uri = r.uri \ + WHERE ref.target_uri = $1 AND ref.collection = $2 \ + ORDER BY r.created_at DESC, r.uri DESC \ + LIMIT $3", ) - .bind(&collection) .bind(&uri) + .bind(&collection) .bind(limit) - .bind(offset) .fetch_all(&state.db) .await - .map_err(|e| mlua::Error::runtime(format!("DB backlinks failed: {e}")))? + .map_err(|e| mlua::Error::runtime(format!("DB backlinks failed: {e}")))?, }; - let has_next = rows.len() as i64 == limit; - let records: Vec = rows + let has_next = rows_raw.len() as i64 == limit; + + let result_table = lua.create_table()?; + + if has_next && let Some((last_uri, _, _, last_created_at)) = rows_raw.last() { + let cursor = encode_cursor(&last_created_at.to_rfc3339(), last_uri); + result_table.set("cursor", cursor)?; + } + + let records: Vec = rows_raw .into_iter() - .map(|(uri, _did, mut record)| { + .map(|(uri, _did, mut record, _created_at)| { if let Some(obj) = record.as_object_mut() { obj.insert("uri".to_string(), json!(uri)); } @@ -276,13 +448,7 @@ .map(|r| lua.to_value(r)) .collect::>()?; 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)) } @@ -495,6 +661,26 @@ assert!( err.contains("invalid sortDirection"), "expected sortDirection error, got: {err}" ); + } + + #[test] + fn cursor_round_trip() { + let encoded = super::encode_cursor("2026-03-12T10:00:00Z", "at://did:plc:abc/col/rkey"); + let (ts, uri) = super::decode_cursor(&encoded).unwrap(); + assert_eq!(ts, "2026-03-12T10:00:00Z"); + assert_eq!(uri, "at://did:plc:abc/col/rkey"); + } + + #[test] + fn decode_invalid_cursor_returns_none() { + assert!(super::decode_cursor("not-valid-base64!!!").is_none()); + } + + #[test] + fn decode_cursor_missing_pipe_returns_none() { + use base64::{Engine, engine::general_purpose::STANDARD as BASE64}; + let encoded = BASE64.encode("no-pipe-here"); + assert!(super::decode_cursor(&encoded).is_none()); } #[tokio::test] diff --git a/src/lua/record.rs b/src/lua/record.rs --- a/src/lua/record.rs +++ b/src/lua/record.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use crate::AppState; use crate::auth::Claims; +use crate::record_refs::sync_refs; use crate::repo::{self, AtpSession}; use super::tid::generate_tid; @@ -112,6 +113,8 @@ .bind(cid) .execute(&state.db) .await; + let _ = sync_refs(&state.db, uri, &collection, &data).await; + result } else { // CREATE @@ -169,6 +172,8 @@ .bind(&data) .bind(cid) .execute(&state.db) .await; + + let _ = sync_refs(&state.db, uri, &collection, &data).await; } result @@ -514,6 +519,8 @@ .bind(cid) .execute(&state.db) .await; + let _ = sync_refs(&state.db, uri.as_str(), &collection, &data).await; + Ok(result) } else { let mut pds_body = json!({ @@ -576,6 +583,8 @@ .bind(&data) .bind(cid) .execute(&state.db) .await; + + let _ = sync_refs(&state.db, uri, &collection, &data).await; } Ok(result) diff --git a/src/main.rs b/src/main.rs --- a/src/main.rs +++ b/src/main.rs @@ -30,6 +30,36 @@ .run(&db) .await .expect("failed to run migrations"); + // Backfill record_refs if empty (first run after upgrade) + { + let count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM record_refs") + .fetch_one(&db) + .await + .expect("failed to count record_refs"); + + if count.0 == 0 { + info!("backfilling record_refs table..."); + let records: Vec<(String, String, serde_json::Value)> = + sqlx::query_as("SELECT uri, collection, record FROM records") + .fetch_all(&db) + .await + .expect("failed to fetch records for backfill"); + + let total = records.len(); + for (i, (uri, collection, record)) in records.iter().enumerate() { + if let Err(e) = + happyview::record_refs::sync_refs(&db, uri, collection, record).await + { + warn!(uri = uri.as_str(), "failed to backfill refs: {e}"); + } + if (i + 1) % 10000 == 0 { + info!("backfill progress: {}/{}", i + 1, total); + } + } + info!("backfill complete: processed {total} records"); + } + } + let lexicons = LexiconRegistry::new(); lexicons .load_from_db(&db) diff --git a/src/record_refs.rs b/src/record_refs.rs new file mode 100644 --- /dev/null +++ b/src/record_refs.rs @@ -0,0 +1,110 @@ +use serde_json::Value; +use std::collections::HashSet; + +/// Recursively walk a JSON value and collect all string values starting with "at://". +pub fn extract_at_uris(value: &Value) -> HashSet { + let mut uris = HashSet::new(); + collect_at_uris(value, &mut uris); + uris +} + +fn collect_at_uris(value: &Value, uris: &mut HashSet) { + match value { + Value::String(s) => { + if s.starts_with("at://") { + uris.insert(s.clone()); + } + } + Value::Array(arr) => { + for item in arr { + collect_at_uris(item, uris); + } + } + Value::Object(obj) => { + for v in obj.values() { + collect_at_uris(v, uris); + } + } + _ => {} + } +} + +/// Update record_refs for a given source record. +/// Deletes old refs and inserts new ones. +pub async fn sync_refs( + db: &sqlx::PgPool, + source_uri: &str, + collection: &str, + record: &Value, +) -> Result<(), sqlx::Error> { + let uris = extract_at_uris(record); + + // Delete existing refs for this source + sqlx::query("DELETE FROM record_refs WHERE source_uri = $1") + .bind(source_uri) + .execute(db) + .await?; + + // Insert new refs + for target_uri in &uris { + sqlx::query( + "INSERT INTO record_refs (source_uri, target_uri, collection) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING" + ) + .bind(source_uri) + .bind(target_uri) + .bind(collection) + .execute(db) + .await?; + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn extracts_top_level_uri() { + let val = json!({"subject": "at://did:plc:abc/com.example/123"}); + let uris = extract_at_uris(&val); + assert_eq!(uris.len(), 1); + assert!(uris.contains("at://did:plc:abc/com.example/123")); + } + + #[test] + fn extracts_nested_uri() { + let val = json!({"outer": {"inner": "at://did:plc:abc/col/rkey"}}); + let uris = extract_at_uris(&val); + assert_eq!(uris.len(), 1); + assert!(uris.contains("at://did:plc:abc/col/rkey")); + } + + #[test] + fn extracts_uris_from_arrays() { + let val = json!({"refs": ["at://did:plc:a/col/1", "at://did:plc:b/col/2"]}); + let uris = extract_at_uris(&val); + assert_eq!(uris.len(), 2); + } + + #[test] + fn ignores_non_at_strings() { + let val = json!({"url": "https://example.com", "name": "test"}); + let uris = extract_at_uris(&val); + assert!(uris.is_empty()); + } + + #[test] + fn empty_object_returns_empty() { + let uris = extract_at_uris(&json!({})); + assert!(uris.is_empty()); + } + + #[test] + fn deduplicates_repeated_uris() { + let val = json!({"a": "at://did:plc:x/c/1", "b": "at://did:plc:x/c/1"}); + let uris = extract_at_uris(&val); + assert_eq!(uris.len(), 1); + } +} diff --git a/src/tap.rs b/src/tap.rs --- a/src/tap.rs +++ b/src/tap.rs @@ -540,6 +540,10 @@ .execute(db) .await { Ok(_) => { + let _ = + crate::record_refs::sync_refs(db, &uri, &record.collection, &rec_to_store) + .await; + log_event( db, EventLog { diff --git a/src/xrpc/procedure.rs b/src/xrpc/procedure.rs --- a/src/xrpc/procedure.rs +++ b/src/xrpc/procedure.rs @@ -6,6 +6,7 @@ use crate::AppState; use crate::auth::Claims; use crate::error::AppError; use crate::lexicon::ProcedureAction; +use crate::record_refs::sync_refs; use crate::repo; pub(super) async fn handle_procedure( @@ -103,6 +104,8 @@ .bind(&record) .bind(cid) .execute(&state.db) .await; + + let _ = sync_refs(&state.db, uri, collection, &record).await; } Ok(( @@ -183,6 +186,8 @@ .bind(&record) .bind(cid) .execute(&state.db) .await; + + let _ = sync_refs(&state.db, uri, collection, &record).await; Ok(( StatusCode::OK,