From 83520434e0b5e787ea931aa0ce6722b4fc640a78 Mon Sep 17 00:00:00 2001 From: Trezy Date: Thu, 12 Feb 2026 20:54:44 -0600 Subject: [PATCH] feat: add Jetstream listener and record indexing with sqlx migrations --- Cargo.lock | 13 ++ Cargo.toml | 3 +- migrations/20260212000000_create_records.sql | 12 ++ src/jetstream.rs | 157 +++++++++++++++++++ src/main.rs | 8 + 5 files changed, 192 insertions(+), 1 deletion(-) create mode 100644 migrations/20260212000000_create_records.sql create mode 100644 src/jetstream.rs diff --git a/Cargo.lock b/Cargo.lock index 4114fa2..4b55070 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -539,6 +539,17 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +[[package]] +name = "futures-macro" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.31" @@ -559,6 +570,7 @@ checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" dependencies = [ "futures-core", "futures-io", + "futures-macro", "futures-sink", "futures-task", "memchr", @@ -655,6 +667,7 @@ dependencies = [ "bytes", "chrono", "dotenvy", + "futures-util", "jsonwebtoken", "p256", "reqwest", diff --git a/Cargo.toml b/Cargo.toml index e9822b9..410ba48 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ edition = "2024" axum = "0.8" base64 = "0.22" dotenvy = "0.15" +futures-util = "0.3" jsonwebtoken = "9" bytes = "1" chrono = { version = "0.4", features = ["serde"] } @@ -16,7 +17,7 @@ reqwest = { version = "0.12", features = ["json"] } serde = { version = "1", features = ["derive"] } serde_json = "1" sha2 = "0.10" -sqlx = { version = "0.8", features = ["runtime-tokio", "tls-rustls", "postgres", "json", "chrono"] } +sqlx = { version = "0.8", features = ["runtime-tokio", "tls-rustls", "postgres", "json", "chrono", "migrate"] } tokio = { version = "1", features = ["full"] } tokio-tungstenite = { version = "0.26", features = ["rustls-tls-webpki-roots"] } tower-http = { version = "0.6", features = ["cors", "trace"] } diff --git a/migrations/20260212000000_create_records.sql b/migrations/20260212000000_create_records.sql new file mode 100644 index 0000000..803b819 --- /dev/null +++ b/migrations/20260212000000_create_records.sql @@ -0,0 +1,12 @@ +CREATE TABLE IF NOT EXISTS records ( + uri TEXT PRIMARY KEY, + did TEXT NOT NULL, + collection TEXT NOT NULL, + rkey TEXT NOT NULL, + record JSONB NOT NULL, + cid TEXT NOT NULL, + indexed_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX IF NOT EXISTS idx_records_did_collection ON records (did, collection); +CREATE INDEX IF NOT EXISTS idx_records_collection ON records (collection); diff --git a/src/jetstream.rs b/src/jetstream.rs new file mode 100644 index 0000000..c806290 --- /dev/null +++ b/src/jetstream.rs @@ -0,0 +1,157 @@ +use futures_util::StreamExt; +use serde::Deserialize; +use serde_json::Value; +use sqlx::PgPool; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::Arc; +use tokio_tungstenite::tungstenite::Message; + +// --------------------------------------------------------------------------- +// Jetstream event types +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +struct JetstreamEvent { + did: String, + time_us: i64, + kind: String, + commit: Option, +} + +#[derive(Deserialize)] +struct JetstreamCommit { + operation: String, + collection: String, + rkey: String, + record: Option, + cid: Option, +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +/// Spawn a background task that subscribes to the Jetstream firehose and +/// indexes `games.gamesgamesgamesgames.*` records into PostgreSQL. +pub fn spawn(db: PgPool, jetstream_url: String) { + tokio::spawn(async move { + let cursor: Arc = Arc::new(AtomicI64::new(0)); + + loop { + if let Err(e) = run(&db, &jetstream_url, &cursor).await { + tracing::warn!("jetstream disconnected: {e}"); + } + + // Back off before reconnecting. + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + tracing::info!("reconnecting to jetstream..."); + } + }); +} + +// --------------------------------------------------------------------------- +// Connection loop +// --------------------------------------------------------------------------- + +async fn run( + db: &PgPool, + jetstream_url: &str, + cursor: &Arc, +) -> Result<(), Box> { + let mut url = format!( + "{}?wantedCollections=games.gamesgamesgamesgames.*", + jetstream_url, + ); + + let last = cursor.load(Ordering::Relaxed); + if last > 0 { + // Rewind 5 seconds for gapless playback. + let rewound = last - 5_000_000; + url.push_str(&format!("&cursor={rewound}")); + tracing::info!(cursor = rewound, "resuming jetstream with cursor"); + } + + let (ws, _) = tokio_tungstenite::connect_async(&url).await?; + tracing::info!("connected to jetstream"); + + let (_, mut read) = ws.split(); + + while let Some(msg) = read.next().await { + let msg = msg?; + + let text = match msg { + Message::Text(t) => t, + Message::Close(_) => break, + _ => continue, + }; + + let event: JetstreamEvent = match serde_json::from_str(&text) { + Ok(e) => e, + Err(e) => { + tracing::debug!("skipping unparseable event: {e}"); + continue; + } + }; + + // Update cursor. + cursor.store(event.time_us, Ordering::Relaxed); + + if event.kind != "commit" { + continue; + } + + let commit = match event.commit { + Some(c) => c, + None => continue, + }; + + let uri = format!( + "at://{}/{}/{}", + event.did, commit.collection, commit.rkey, + ); + + match commit.operation.as_str() { + "create" | "update" => { + let record = match commit.record { + Some(r) => r, + None => continue, + }; + let cid = commit.cid.unwrap_or_default(); + + if let Err(e) = sqlx::query( + r#" + INSERT INTO records (uri, did, collection, rkey, record, cid) + VALUES ($1, $2, $3, $4, $5, $6) + ON CONFLICT (uri) DO UPDATE + SET record = EXCLUDED.record, + cid = EXCLUDED.cid, + indexed_at = NOW() + "#, + ) + .bind(&uri) + .bind(&event.did) + .bind(&commit.collection) + .bind(&commit.rkey) + .bind(&record) + .bind(&cid) + .execute(db) + .await + { + tracing::warn!(uri = %uri, "failed to upsert record: {e}"); + } + } + "delete" => { + if let Err(e) = sqlx::query("DELETE FROM records WHERE uri = $1") + .bind(&uri) + .execute(db) + .await + { + tracing::warn!(uri = %uri, "failed to delete record: {e}"); + } + } + _ => {} + } + } + + Ok(()) +} diff --git a/src/main.rs b/src/main.rs index a28b109..86c34f4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ mod auth; mod config; mod error; +mod jetstream; mod profile; mod repo; mod server; @@ -35,12 +36,19 @@ async fn main() { info!("connected to database"); + sqlx::migrate!() + .run(&db) + .await + .expect("failed to run migrations"); + let state = AppState { config: config.clone(), http: reqwest::Client::new(), db, }; + jetstream::spawn(state.db.clone(), config.jetstream_url.clone()); + let app = server::router(state); let addr = config.listen_addr(); -- 2.54.0