From b8e6caaa41fc191fd1ef7484789aca24749110ad Mon Sep 17 00:00:00 2001 From: Trezy Date: Sun, 15 Mar 2026 09:20:06 -0500 Subject: [PATCH] feat: add support for labelers --- ...314100000_create_labeler_subscriptions.sql | 7 + migrations/20260314100001_create_labels.sql | 11 + src/admin/labelers.rs | 150 ++++++ src/admin/mod.rs | 6 + src/admin/permissions.rs | 20 +- src/admin/types.rs | 23 + src/aip.rs | 2 + src/labeler.rs | 440 ++++++++++++++++++ src/lib.rs | 2 + src/lua/atproto_api.rs | 144 ++++++ src/lua/db_api.rs | 2 + src/lua/execute.rs | 2 + src/lua/http_api.rs | 2 + src/main.rs | 7 +- src/tap.rs | 2 + tests/common/app.rs | 2 + tests/common/db.rs | 2 +- tests/e2e_labelers.rs | 387 +++++++++++++++ tests/lua_atproto_api.rs | 386 +++++++++++++++ tests/lua_db_api.rs | 2 + 20 files changed, 1595 insertions(+), 4 deletions(-) create mode 100644 migrations/20260314100000_create_labeler_subscriptions.sql create mode 100644 migrations/20260314100001_create_labels.sql create mode 100644 src/admin/labelers.rs create mode 100644 src/labeler.rs create mode 100644 tests/e2e_labelers.rs create mode 100644 tests/lua_atproto_api.rs diff --git a/migrations/20260314100000_create_labeler_subscriptions.sql b/migrations/20260314100000_create_labeler_subscriptions.sql new file mode 100644 index 0000000..7f40edf --- /dev/null +++ b/migrations/20260314100000_create_labeler_subscriptions.sql @@ -0,0 +1,7 @@ +CREATE TABLE labeler_subscriptions ( + did TEXT PRIMARY KEY, + cursor BIGINT, + status TEXT NOT NULL DEFAULT 'active', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); diff --git a/migrations/20260314100001_create_labels.sql b/migrations/20260314100001_create_labels.sql new file mode 100644 index 0000000..43dbd50 --- /dev/null +++ b/migrations/20260314100001_create_labels.sql @@ -0,0 +1,11 @@ +CREATE TABLE labels ( + src TEXT NOT NULL, + uri TEXT NOT NULL, + val TEXT NOT NULL, + cts TIMESTAMPTZ NOT NULL, + exp TIMESTAMPTZ, + PRIMARY KEY (src, uri, val) +); + +CREATE INDEX idx_labels_uri ON labels (uri); +CREATE INDEX idx_labels_exp ON labels (exp) WHERE exp IS NOT NULL; diff --git a/src/admin/labelers.rs b/src/admin/labelers.rs new file mode 100644 index 0000000..a7f8c78 --- /dev/null +++ b/src/admin/labelers.rs @@ -0,0 +1,150 @@ +use axum::Json; +use axum::extract::{Path, State}; +use axum::http::StatusCode; + +use crate::AppState; +use crate::error::AppError; +use crate::event_log::{EventLog, Severity, log_event}; + +use super::auth::UserAuth; +use super::permissions::Permission; +use super::types::{AddLabelerBody, LabelerSummary, UpdateLabelerBody}; + +/// GET /admin/labelers — list all labeler subscriptions. +pub(super) async fn list( + State(state): State, + auth: UserAuth, +) -> Result>, AppError> { + auth.require(Permission::LabelersRead).await?; + + let labelers: Vec = sqlx::query_as( + "SELECT did, status, cursor, created_at, updated_at FROM labeler_subscriptions ORDER BY created_at", + ) + .fetch_all(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to list labeler subscriptions: {e}")))?; + + Ok(Json(labelers)) +} + +/// POST /admin/labelers — add a labeler subscription. +pub(super) async fn add( + State(state): State, + auth: UserAuth, + Json(body): Json, +) -> Result { + auth.require(Permission::LabelersCreate).await?; + + sqlx::query( + r#" + INSERT INTO labeler_subscriptions (did) + VALUES ($1) + ON CONFLICT (did) DO UPDATE SET status = 'active', updated_at = NOW() + "#, + ) + .bind(&body.did) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to add labeler subscription: {e}")))?; + + // Notify the labeler consumer to pick up the new subscription. + let _ = state.labeler_subscriptions_tx.send(()); + + log_event( + &state.db, + EventLog { + event_type: "labeler.added".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(body.did.clone()), + detail: serde_json::json!({}), + }, + ) + .await; + + Ok(StatusCode::CREATED) +} + +/// PATCH /admin/labelers/{did} — update labeler status (active/paused). +pub(super) async fn update( + State(state): State, + auth: UserAuth, + Path(did): Path, + Json(body): Json, +) -> Result { + auth.require(Permission::LabelersCreate).await?; + + let result = sqlx::query( + "UPDATE labeler_subscriptions SET status = $1, updated_at = NOW() WHERE did = $2", + ) + .bind(&body.status) + .bind(&did) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to update labeler subscription: {e}")))?; + + if result.rows_affected() == 0 { + return Err(AppError::NotFound(format!( + "labeler subscription '{did}' not found" + ))); + } + + let _ = state.labeler_subscriptions_tx.send(()); + + log_event( + &state.db, + EventLog { + event_type: "labeler.updated".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(did), + detail: serde_json::json!({ "status": body.status }), + }, + ) + .await; + + Ok(StatusCode::NO_CONTENT) +} + +/// DELETE /admin/labelers/{did} — remove a labeler subscription and its labels. +pub(super) async fn delete( + State(state): State, + auth: UserAuth, + Path(did): Path, +) -> Result { + auth.require(Permission::LabelersDelete).await?; + + let result = sqlx::query("DELETE FROM labeler_subscriptions WHERE did = $1") + .bind(&did) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to delete labeler subscription: {e}")))?; + + if result.rows_affected() == 0 { + return Err(AppError::NotFound(format!( + "labeler subscription '{did}' not found" + ))); + } + + // Also remove all labels from this labeler. + let _ = sqlx::query("DELETE FROM labels WHERE src = $1") + .bind(&did) + .execute(&state.db) + .await; + + let _ = state.labeler_subscriptions_tx.send(()); + + log_event( + &state.db, + EventLog { + event_type: "labeler.deleted".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(did), + detail: serde_json::json!({}), + }, + ) + .await; + + Ok(StatusCode::NO_CONTENT) +} diff --git a/src/admin/mod.rs b/src/admin/mod.rs index 34dd5dc..2eef7eb 100644 --- a/src/admin/mod.rs +++ b/src/admin/mod.rs @@ -2,6 +2,7 @@ mod api_keys; pub(crate) mod auth; mod backfill; mod events; +mod labelers; mod lexicons; mod network_lexicons; pub(crate) mod permissions; @@ -62,4 +63,9 @@ pub fn admin_routes(_state: AppState) -> Router { post(script_variables::upsert).get(script_variables::list), ) .route("/script-variables/{key}", delete(script_variables::delete)) + .route("/labelers", post(labelers::add).get(labelers::list)) + .route( + "/labelers/{did}", + patch(labelers::update).delete(labelers::delete), + ) } diff --git a/src/admin/permissions.rs b/src/admin/permissions.rs index cb09a57..ee0edf1 100644 --- a/src/admin/permissions.rs +++ b/src/admin/permissions.rs @@ -2,7 +2,7 @@ use std::collections::HashSet; use serde::{Deserialize, Serialize}; -/// All 20 permissions in the system. +/// All 23 permissions in the system. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub enum Permission { #[serde(rename = "lexicons:create")] @@ -52,6 +52,13 @@ pub enum Permission { #[serde(rename = "events:read")] EventsRead, + + #[serde(rename = "labelers:create")] + LabelersCreate, + #[serde(rename = "labelers:read")] + LabelersRead, + #[serde(rename = "labelers:delete")] + LabelersDelete, } impl Permission { @@ -78,10 +85,13 @@ impl Permission { Self::BackfillRead => "backfill:read", Self::StatsRead => "stats:read", Self::EventsRead => "events:read", + Self::LabelersCreate => "labelers:create", + Self::LabelersRead => "labelers:read", + Self::LabelersDelete => "labelers:delete", } } - /// All 20 permissions. + /// All 23 permissions. pub fn all() -> HashSet { HashSet::from([ Self::LexiconsCreate, @@ -104,6 +114,9 @@ impl Permission { Self::BackfillRead, Self::StatsRead, Self::EventsRead, + Self::LabelersCreate, + Self::LabelersRead, + Self::LabelersDelete, ]) } } @@ -145,6 +158,9 @@ impl Template { perms.insert(Permission::ScriptVariablesCreate); perms.insert(Permission::ScriptVariablesDelete); perms.insert(Permission::RecordsDelete); + perms.insert(Permission::LabelersCreate); + perms.insert(Permission::LabelersRead); + perms.insert(Permission::LabelersDelete); perms } Self::FullAccess => Permission::all(), diff --git a/src/admin/types.rs b/src/admin/types.rs index 87b1a33..6026d67 100644 --- a/src/admin/types.rs +++ b/src/admin/types.rs @@ -181,3 +181,26 @@ pub(super) struct UpsertScriptVariableBody { pub(super) key: String, pub(super) value: String, } + +// --------------------------------------------------------------------------- +// Labeler subscription types +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +pub(super) struct AddLabelerBody { + pub(super) did: String, +} + +#[derive(Serialize, sqlx::FromRow)] +pub(super) struct LabelerSummary { + pub(super) did: String, + pub(super) status: String, + pub(super) cursor: Option, + pub(super) created_at: chrono::DateTime, + pub(super) updated_at: chrono::DateTime, +} + +#[derive(Deserialize)] +pub(super) struct UpdateLabelerBody { + pub(super) status: String, +} diff --git a/src/aip.rs b/src/aip.rs index 7941edd..84c7b7f 100644 --- a/src/aip.rs +++ b/src/aip.rs @@ -111,12 +111,14 @@ mod tests { event_log_retention_days: 30, }; let (tx, _) = watch::channel(vec![]); + let (labeler_tx, _) = watch::channel(()); AppState { config, http: reqwest::Client::new(), db: sqlx::PgPool::connect_lazy("postgres://localhost/fake").unwrap(), lexicons: crate::lexicon::LexiconRegistry::new(), collections_tx: tx, + labeler_subscriptions_tx: labeler_tx, } } diff --git a/src/labeler.rs b/src/labeler.rs new file mode 100644 index 0000000..1451db0 --- /dev/null +++ b/src/labeler.rs @@ -0,0 +1,440 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use futures_util::StreamExt; +use serde::Deserialize; +use tokio::sync::watch; +use tokio::task::JoinHandle; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; + +use crate::AppState; +use crate::event_log::{EventLog, Severity, log_event}; +use crate::profile; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +struct SubscribeLabelsMessage { + seq: i64, + labels: Vec