From 3dc6f5bb6fc65826cea8eae5a7608a3a9ff4ac24 Mon Sep 17 00:00:00 2001 From: Mia Date: Thu, 30 Jan 2025 20:36:06 +0000 Subject: [PATCH] blocks and follows --- consumer/src/indexer/db.rs | 63 ++++++++++ consumer/src/indexer/mod.rs | 113 ++++++++++++------ consumer/src/indexer/records.rs | 16 +++ consumer/src/indexer/types.rs | 14 ++- consumer/src/sql/follow_stats_upsert.sql | 21 ++++ .../down.sql | 3 + .../up.sql | 29 +++++ parakeet-db/src/models.rs | 30 +++++ parakeet-db/src/schema.rs | 36 ++++++ 9 files changed, 288 insertions(+), 37 deletions(-) create mode 100644 consumer/src/indexer/records.rs create mode 100644 consumer/src/sql/follow_stats_upsert.sql create mode 100644 migrations/2025-01-29-213341_follows_and_blocks/down.sql create mode 100644 migrations/2025-01-29-213341_follows_and_blocks/up.sql diff --git a/consumer/src/indexer/db.rs b/consumer/src/indexer/db.rs index ebbf4feb..b875f7f0 100644 --- a/consumer/src/indexer/db.rs +++ b/consumer/src/indexer/db.rs @@ -1,5 +1,7 @@ +use super::records; use chrono::prelude::*; use diesel::prelude::*; +use diesel::sql_types::{Array, Text}; use diesel_async::{AsyncPgConnection, RunQueryDsl}; use ipld_core::cid::Cid; use parakeet_db::{models, schema, types}; @@ -60,3 +62,64 @@ pub async fn update_repo_state( .execute(conn) .await } + +pub async fn insert_block( + conn: &mut AsyncPgConnection, + repo: &str, + at_uri: &str, + rec: records::AppBskyGraphBlock, +) -> QueryResult { + diesel::insert_into(schema::blocks::table) + .values(&models::NewBlock { + at_uri, + did: repo, + subject: &rec.subject, + created_at: rec.created_at.naive_utc(), + }) + .on_conflict_do_nothing() + .execute(conn) + .await +} + +pub async fn delete_block(conn: &mut AsyncPgConnection, at_uri: &str) -> QueryResult { + diesel::delete(schema::blocks::table) + .filter(schema::blocks::at_uri.eq(at_uri)) + .execute(conn) + .await +} + +pub async fn insert_follow( + conn: &mut AsyncPgConnection, + repo: &str, + at_uri: &str, + rec: records::AppBskyGraphFollow, +) -> QueryResult { + diesel::insert_into(schema::follows::table) + .values(&models::NewFollow { + at_uri, + did: repo, + subject: &rec.subject, + created_at: rec.created_at.naive_utc(), + }) + .on_conflict_do_nothing() + .execute(conn) + .await +} + +pub async fn delete_follow(conn: &mut AsyncPgConnection, at_uri: &str) -> QueryResult { + diesel::delete(schema::follows::table) + .filter(schema::follows::at_uri.eq(at_uri)) + .returning(schema::follows::subject) + .get_result(conn) + .await +} + +pub async fn update_follow_stats( + conn: &mut AsyncPgConnection, + subjects: &[&str], +) -> QueryResult { + diesel::sql_query(include_str!("../sql/follow_stats_upsert.sql")) + .bind::, _>(subjects) + .execute(conn) + .await +} diff --git a/consumer/src/indexer/mod.rs b/consumer/src/indexer/mod.rs index 41598a54..9cb9a835 100644 --- a/consumer/src/indexer/mod.rs +++ b/consumer/src/indexer/mod.rs @@ -1,4 +1,4 @@ -use crate::firehose::{AtpAccountEvent, AtpCommitEvent, AtpIdentityEvent, FirehoseEvent}; +use crate::firehose::{AtpAccountEvent, AtpCommitEvent, AtpIdentityEvent, CommitOp, FirehoseEvent}; use crate::indexer::types::{CollectionType, RecordTypes}; use diesel_async::pooled_connection::deadpool::Pool; use diesel_async::{AsyncConnection, AsyncPgConnection}; @@ -10,6 +10,7 @@ use tokio::sync::mpsc::Receiver; use tracing::{instrument, Instrument}; mod db; +mod records; mod types; pub async fn relay_indexer( @@ -103,48 +104,88 @@ async fn index_commit(conn: &mut AsyncPgConnection, commit: AtpCommitEvent) -> e } }; - conn.transaction::<_, diesel::result::Error, _>(|t| Box::pin(async move { - // ensure new repo rev > current rev, but only when the repo is synced - // (if repo is processing, they'll end up on the queue anyway) - if is_active && db::update_repo_state(t, &commit.repo, &commit.rev, commit.commit).await? == 0 { - tracing::warn!("Got a repo update older than the current rev. not processing ops."); - return Ok(false); + conn.transaction::<_, diesel::result::Error, _>(|t| { + Box::pin(async move { + // ensure new repo rev > current rev, but only when the repo is synced + // (if repo is processing, they'll end up on the queue anyway) + if is_active + && db::update_repo_state(t, &commit.repo, &commit.rev, commit.commit).await? == 0 + { + tracing::warn!("Got a repo update older than the current rev. not processing ops."); + return Ok(false); + } + + for op in &commit.ops { + process_op(t, &commit.repo, op, &blocks).await?; + } + + Ok(true) + }) + }) + .await?; + + Ok(()) +} + +async fn process_op( + conn: &mut AsyncPgConnection, + repo: &str, + op: &CommitOp, + blocks: &HashMap>, +) -> diesel::QueryResult<()> { + let Some((collection_raw, _)) = op.path.split_once("/") else { + tracing::warn!("op contained invalid path {}", op.path); + return Ok(()); + }; + + let collection = CollectionType::from_str(collection_raw); + if collection == CollectionType::Unsupported { + tracing::debug!("{} {collection_raw} is unsupported", op.action); + return Ok(()); + } + + let full_path = format!("at://{repo}/{}", &op.path); + + if op.action == "create" || op.action == "update" { + if op.action == "update" && !collection.can_update() { + return Ok(()); } - for op in &commit.ops { - let Some((collection_raw, rkey)) = op.path.split_once("/") else { - tracing::warn!("op contained invalid path {}", op.path); - continue; - }; + let Some(block) = op.cid.and_then(|cid| blocks.get(&cid)) else { + tracing::error!("Missing Cid or the block was not found"); + return Ok(()); + }; - let collection = CollectionType::from_str(collection_raw); - if collection == CollectionType::Unsupported { - tracing::debug!("{} {collection_raw} is unsupported", op.action); - continue; + let decoded: RecordTypes = match serde_ipld_dagcbor::from_slice(block) { + Ok(decoded) => decoded, + Err(err) => { + tracing::error!("Failed to decode record: {err}"); + return Ok(()); } + }; - if op.action == "create" || op.action == "update" { - if let Some(block) = op.cid.and_then(|cid| blocks.get(&cid)) { - let decoded: RecordTypes = match serde_ipld_dagcbor::from_slice(block) { - Ok(decoded) => decoded, - Err(err) => { - tracing::error!("Failed to decode record: {err}"); - continue; - } - }; - } else { - tracing::error!("Missing Cid or the block was not found"); - continue; - } - } else if op.action == "delete" { - // - } else { - tracing::warn!("op contained invalid action {}", op.action); + match decoded { + RecordTypes::AppBskyGraphBlock(record) => { + db::insert_block(conn, repo, &full_path, record).await?; + } + RecordTypes::AppBskyGraphFollow(record) => { + let subj = record.subject.clone(); + db::insert_follow(conn, repo, &full_path, record).await?; + db::update_follow_stats(conn, &[repo, &subj]).await?; } } - - Ok(true) - })).await?; + } else if op.action == "delete" { + match collection { + CollectionType::BskyBlock => db::delete_block(conn, &full_path).await?, + CollectionType::BskyFollow => { + let subject = db::delete_follow(conn, &full_path).await?; + db::update_follow_stats(conn, &[repo, &subject]).await? + }, + _ => unreachable!(), + }; + } else { + tracing::warn!("op contained invalid action {}", op.action); + } Ok(()) } diff --git a/consumer/src/indexer/records.rs b/consumer/src/indexer/records.rs new file mode 100644 index 00000000..844d6189 --- /dev/null +++ b/consumer/src/indexer/records.rs @@ -0,0 +1,16 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AppBskyGraphBlock { + pub subject: String, + pub created_at: DateTime, +} + +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AppBskyGraphFollow { + pub subject: String, + pub created_at: DateTime, +} diff --git a/consumer/src/indexer/types.rs b/consumer/src/indexer/types.rs index 4e150c86..9099ec13 100644 --- a/consumer/src/indexer/types.rs +++ b/consumer/src/indexer/types.rs @@ -1,23 +1,35 @@ +use super::records; use serde::{Deserialize, Serialize}; #[derive(Debug, Deserialize, Serialize)] #[serde(tag = "$type")] -pub enum RecordTypes {} +pub enum RecordTypes { + #[serde(rename = "app.bsky.graph.block")] + AppBskyGraphBlock(records::AppBskyGraphBlock), + #[serde(rename = "app.bsky.graph.follow")] + AppBskyGraphFollow(records::AppBskyGraphFollow), +} #[derive(Debug, PartialOrd, PartialEq)] pub enum CollectionType { + BskyBlock, + BskyFollow, Unsupported, } impl CollectionType { pub(crate) fn from_str(input: &str) -> CollectionType { match input { + "app.bsky.graph.block" => CollectionType::BskyBlock, + "app.bsky.graph.follow" => CollectionType::BskyFollow, _ => CollectionType::Unsupported, } } pub fn can_update(&self) -> bool { match self { + CollectionType::BskyBlock => false, + CollectionType::BskyFollow => false, CollectionType::Unsupported => false, } } diff --git a/consumer/src/sql/follow_stats_upsert.sql b/consumer/src/sql/follow_stats_upsert.sql new file mode 100644 index 00000000..d4ff7fe3 --- /dev/null +++ b/consumer/src/sql/follow_stats_upsert.sql @@ -0,0 +1,21 @@ +with new_info as ( + SELECT did, coalesce(following, 0) as following, coalesce(followers, 0) as followers + FROM actors + LEFT JOIN (SELECT did, count(*) AS following + FROM follows + GROUP BY 1) fi USING (did) + LEFT JOIN (SELECT subject as did, count(*) AS followers + FROM follows + GROUP BY 1) fe USING (did) + where actors.did = any($1) + ) + merge into follow_stats fs +using new_info ni +on ni.did = fs.did +when matched then + update + set followers=ni.followers, + following=ni.following +when not matched then + insert (did, followers, following) + values (ni.did, ni.followers, ni.following); \ No newline at end of file diff --git a/migrations/2025-01-29-213341_follows_and_blocks/down.sql b/migrations/2025-01-29-213341_follows_and_blocks/down.sql new file mode 100644 index 00000000..da2c6dc3 --- /dev/null +++ b/migrations/2025-01-29-213341_follows_and_blocks/down.sql @@ -0,0 +1,3 @@ +drop table blocks; +drop table follow_stats; +drop table follows; \ No newline at end of file diff --git a/migrations/2025-01-29-213341_follows_and_blocks/up.sql b/migrations/2025-01-29-213341_follows_and_blocks/up.sql new file mode 100644 index 00000000..829f488a --- /dev/null +++ b/migrations/2025-01-29-213341_follows_and_blocks/up.sql @@ -0,0 +1,29 @@ +create table blocks +( + at_uri text primary key, + did text not null references actors (did), + subject text not null, + created_at timestamptz not null +); + +create index blocks_did_index on blocks using hash (did); +create index blocks_subject_index on blocks using hash (subject); + +create table follows +( + at_uri text primary key, + did text not null references actors (did), + subject text not null, + created_at timestamptz not null +); + +create index follow_did_index on follows using hash (did); +create index follow_subject_index on follows using hash (subject); + +create table follow_stats +( + did text primary key, + + followers integer not null default 0, + following integer not null default 0 +); diff --git a/parakeet-db/src/models.rs b/parakeet-db/src/models.rs index 13d0142a..cc007c0e 100644 --- a/parakeet-db/src/models.rs +++ b/parakeet-db/src/models.rs @@ -25,3 +25,33 @@ pub struct NewActor<'a> { pub status: Option, pub last_indexed: Option, } + +#[derive(Insertable)] +#[diesel(table_name = crate::schema::blocks)] +#[diesel(check_for_backend(diesel::pg::Pg))] +pub struct NewBlock<'a> { + pub at_uri: &'a str, + pub did: &'a str, + pub subject: &'a str, + pub created_at: NaiveDateTime, +} + +#[derive(Insertable)] +#[diesel(table_name = crate::schema::follows)] +#[diesel(check_for_backend(diesel::pg::Pg))] +pub struct NewFollow<'a> { + pub at_uri: &'a str, + pub did: &'a str, + pub subject: &'a str, + pub created_at: NaiveDateTime, +} + +#[derive(Debug, Queryable, Selectable, Identifiable)] +#[diesel(table_name = crate::schema::follow_stats)] +#[diesel(primary_key(did))] +#[diesel(check_for_backend(diesel::pg::Pg))] +pub struct FollowStats { + pub did: String, + pub followers: i32, + pub following: i32, +} diff --git a/parakeet-db/src/schema.rs b/parakeet-db/src/schema.rs index 07ad7b5e..2e029ea8 100644 --- a/parakeet-db/src/schema.rs +++ b/parakeet-db/src/schema.rs @@ -11,3 +11,39 @@ diesel::table! { last_indexed -> Nullable, } } + +diesel::table! { + blocks (at_uri) { + at_uri -> Text, + did -> Text, + subject -> Text, + created_at -> Timestamptz, + } +} + +diesel::table! { + follow_stats (did) { + did -> Text, + followers -> Int4, + following -> Int4, + } +} + +diesel::table! { + follows (at_uri) { + at_uri -> Text, + did -> Text, + subject -> Text, + created_at -> Timestamptz, + } +} + +diesel::joinable!(blocks -> actors (did)); +diesel::joinable!(follows -> actors (did)); + +diesel::allow_tables_to_appear_in_same_query!( + actors, + blocks, + follow_stats, + follows, +); -- 2.51.2