From 8b9757787f08117cf7670c047e792940de2d951a Mon Sep 17 00:00:00 2001 From: Timothy Quilling Date: Sun, 7 Dec 2025 23:27:38 -0500 Subject: [PATCH] feat: denormalize strategy; labels --- consumer/src/db/labels.rs | 102 +++++++++++++---- consumer/src/label_indexer/mod.rs | 178 +++++++++++++++++++++++------- 2 files changed, 219 insertions(+), 61 deletions(-) diff --git a/consumer/src/db/labels.rs b/consumer/src/db/labels.rs index 698c95c2..bbc636bf 100644 --- a/consumer/src/db/labels.rs +++ b/consumer/src/db/labels.rs @@ -84,7 +84,7 @@ pub async fn maintain_label_defs( pub async fn maintain_self_labels( conn: &C, repo: &str, - cid: Option, + _cid: Option, at_uri: &str, self_labels: SelfLabels, ) -> Result { @@ -95,32 +95,90 @@ pub async fn maintain_self_labels( return Ok(0); // Actor doesn't exist yet }; - let _ = conn - .execute( - "DELETE FROM labels WHERE self_label=TRUE AND uri=$1", - &[&at_uri], + // Parse AT URI to determine if this is for a post or actor profile + // AT URI format: at://{did}/collection/rkey + // Actor profiles: at://{did}/app.bsky.actor.profile/self + // Posts: at://{did}/app.bsky.feed.post/{rkey} + let (did, collection, rkey) = parakeet_db::at_uri_util::parse_at_uri(at_uri) + .ok_or_else(|| eyre::eyre!("Invalid AT URI: {}", at_uri))?; + + // Build labels array values + // Note: Self-labels only have 'val' field, no negated/expires + let mut labels_data: Vec = Vec::new(); + let negated_vals: Vec = vec![false; self_labels.values.len()]; // All false for self-labels + let expires_vals: Vec> = vec![None; self_labels.values.len()]; // No expiration + + for label in self_labels.values { + labels_data.push(label.val.clone()); + } + + if collection == "app.bsky.actor.profile" && rkey == "self" { + // Update actor labels array + conn.execute( + "UPDATE actors + SET labels = ( + SELECT ARRAY_AGG( + ROW( + $1, + label, + NOW(), + negated, + expires::timestamptz + )::actor_label + ORDER BY idx + ) + FROM unnest($2::text[], $3::boolean[], $4::text[]) + WITH ORDINALITY AS t(label, negated, expires, idx) + ) + WHERE did = $5", + &[ + &labeler_actor_id, + &labels_data, + &negated_vals, + &expires_vals, + &did, + ], ) .await?; + } else if collection == "app.bsky.feed.post" { + // Update post labels array + // Convert TID string to INT8 for rkey lookup + let rkey_i64 = parakeet_db::tid_util::decode_tid(rkey)?; - let cid_bytes = cid.map(|c| c.to_bytes()); - - // Use ON CONFLICT DO NOTHING to handle race conditions when multiple workers - // process the same profile update concurrently. - // Since we DELETE all existing labels first, any conflict means another worker - // just inserted the same label from the same event - no need to update. - let stmt = conn - .prepare_cached( - "INSERT INTO labels (labeler_actor_id, label, uri, self_label, cid, created_at) - VALUES ($1, $2, $3, TRUE, $4, NOW()) - ON CONFLICT (labeler_actor_id, label, uri) - DO NOTHING", + conn.execute( + "UPDATE posts p + SET labels = ( + SELECT ARRAY_AGG( + ROW( + $1, + label, + NOW(), + negated, + expires::timestamptz + )::post_label + ORDER BY idx + ) + FROM unnest($2::text[], $3::boolean[], $4::text[]) + WITH ORDINALITY AS t(label, negated, expires, idx) + ) + FROM actors a + WHERE p.actor_id = a.id + AND a.did = $5 + AND p.rkey = $6", + &[ + &labeler_actor_id, + &labels_data, + &negated_vals, + &expires_vals, + &did, + &rkey_i64, + ], ) .await?; - - for label in self_labels.values { - let _ = conn - .execute(&stmt, &[&labeler_actor_id, &label.val, &at_uri, &cid_bytes]) - .await?; + } else { + // Unsupported collection type for labels + tracing::warn!("Unsupported AT URI collection for labels: {}", collection); + return Ok(0); } Ok(0) diff --git a/consumer/src/label_indexer/mod.rs b/consumer/src/label_indexer/mod.rs index 1db5650a..ee52fe0d 100644 --- a/consumer/src/label_indexer/mod.rs +++ b/consumer/src/label_indexer/mod.rs @@ -1,12 +1,9 @@ use crate::firehose::{AtpLabel, FirehoseConsumer, FirehoseEvent, FirehoseOutput}; -use futures::pin_mut; use metrics::counter; use std::time::Duration; use tokio::sync::mpsc::{channel, Sender}; use tokio::sync::watch::Receiver as WatchReceiver; use tokio::time::Instant; -use tokio_postgres::binary_copy::BinaryCopyInWriter; -use tokio_postgres::types::Type; use tracing::instrument; pub struct LabelServiceManager { @@ -81,54 +78,157 @@ impl LabelServiceManager { } async fn store_labels(conn: &mut tokio_postgres::Client, labels: &[AtpLabel]) -> eyre::Result { + use std::collections::HashMap; + + // Group labels by target (actor or post) + // Key: (did, collection, rkey), Value: Vec of labels for that target + let mut actor_labels: HashMap> = HashMap::new(); + let mut post_labels: HashMap<(String, i64), Vec<&AtpLabel>> = HashMap::new(); + + for label in labels { + // Parse AT URI to determine target + let Some((did, collection, rkey)) = parakeet_db::at_uri_util::parse_at_uri(&label.uri) else { + tracing::warn!("Invalid AT URI in label: {}", label.uri); + continue; + }; + + if collection == "app.bsky.actor.profile" && rkey == "self" { + actor_labels.entry(did.to_string()).or_default().push(label); + } else if collection == "app.bsky.feed.post" { + // Convert TID string to INT8 + let Some(rkey_i64) = parakeet_db::tid_util::decode_tid(rkey) else { + tracing::warn!("Invalid TID in label URI: {}", label.uri); + continue; + }; + post_labels.entry((did.to_string(), rkey_i64)).or_default().push(label); + } else { + tracing::warn!("Unsupported collection for labels: {}", collection); + } + } + let t = conn.transaction().await?; - let _ = t - .execute( - "CREATE TEMP TABLE label_tmp (LIKE labels INCLUDING DEFAULTS) ON COMMIT DROP", - &[], + let mut updated_count = 0u64; + + // Update actor labels + for (did, label_list) in actor_labels { + // Resolve labeler actor_id from DID (src field) + // Note: All labels in label_list should have the same src (labeler DID) + let labeler_did = &label_list[0].src; + + let labeler_actor_id_row = t + .query_opt("SELECT id FROM actors WHERE did = $1", &[labeler_did]) + .await?; + + let Some(labeler_actor_id_row) = labeler_actor_id_row else { + tracing::warn!("Labeler actor not found for DID: {}", labeler_did); + continue; + }; + let labeler_actor_id: i32 = labeler_actor_id_row.get(0); + + // Build parallel arrays for labels + let mut label_vals: Vec = Vec::new(); + let mut negated_vals: Vec = Vec::new(); + let mut expires_vals: Vec> = Vec::new(); + + for label in label_list { + label_vals.push(label.val.clone()); + negated_vals.push(label.neg.unwrap_or(false)); + expires_vals.push(label.exp.map(|dt| dt.naive_utc())); + } + + let rows_updated = t.execute( + "UPDATE actors + SET labels = ( + SELECT ARRAY_AGG( + ROW( + $1, + label, + NOW(), + negated, + expires + )::actor_label + ORDER BY idx + ) + FROM unnest($2::text[], $3::boolean[], $4::timestamp[]) + WITH ORDINALITY AS t(label, negated, expires, idx) + ) + WHERE did = $5", + &[ + &labeler_actor_id, + &label_vals, + &negated_vals, + &expires_vals, + &did, + ], ) .await?; - let sink = t.copy_in("COPY label_tmp (labeler, label, uri, cid, negated, expires, sig, created_at) FROM STDIN (FORMAT binary)").await?; - let binary_writer = BinaryCopyInWriter::new( - sink, - &[ - Type::TEXT, - Type::TEXT, - Type::TEXT, - Type::TEXT, - Type::BOOL, - Type::TIMESTAMP, - Type::BYTEA, - Type::TIMESTAMP, - ], - ); - - pin_mut!(binary_writer); + updated_count += rows_updated; + } - for label in labels { - let exp = label.exp.map(|v| v.naive_utc()); - let sig = label.sig.as_ref().map(|v| v.as_slice()); - let created = label.cts.naive_utc(); - let neg = label.neg.unwrap_or_default(); - - binary_writer - .as_mut() - .write(&[ - &label.src, &label.val, &label.uri, &label.cid, &neg, &exp, &sig, &created, - ]) + // Update post labels + for ((did, rkey_i64), label_list) in post_labels { + // Resolve labeler actor_id from DID + let labeler_did = &label_list[0].src; + + let labeler_actor_id_row = t + .query_opt("SELECT id FROM actors WHERE did = $1", &[labeler_did]) .await?; - } - let count = binary_writer.finish().await?; + let Some(labeler_actor_id_row) = labeler_actor_id_row else { + tracing::warn!("Labeler actor not found for DID: {}", labeler_did); + continue; + }; + let labeler_actor_id: i32 = labeler_actor_id_row.get(0); + + // Build parallel arrays for labels + let mut label_vals: Vec = Vec::new(); + let mut negated_vals: Vec = Vec::new(); + let mut expires_vals: Vec> = Vec::new(); + + for label in label_list { + label_vals.push(label.val.clone()); + negated_vals.push(label.neg.unwrap_or(false)); + expires_vals.push(label.exp.map(|dt| dt.naive_utc())); + } - let _ = t - .execute(include_str!("../db/sql/label_copy_upsert.sql"), &[]) + let rows_updated = t.execute( + "UPDATE posts p + SET labels = ( + SELECT ARRAY_AGG( + ROW( + $1, + label, + NOW(), + negated, + expires + )::post_label + ORDER BY idx + ) + FROM unnest($2::text[], $3::boolean[], $4::timestamp[]) + WITH ORDINALITY AS t(label, negated, expires, idx) + ) + FROM actors a + WHERE p.actor_id = a.id + AND a.did = $5 + AND p.rkey = $6", + &[ + &labeler_actor_id, + &label_vals, + &negated_vals, + &expires_vals, + &did, + &rkey_i64, + ], + ) .await?; + updated_count += rows_updated; + } + t.commit().await?; - Ok(count) + Ok(updated_count) } #[instrument(skip(resume, trigger, user_agent, db_tx, stop))] -- 2.51.2