From 94dbfaa4266b8d284992e48299e68fddc5efc446 Mon Sep 17 00:00:00 2001 From: Trezy Date: Thu, 12 Feb 2026 23:31:52 -0600 Subject: [PATCH] feat: add lexicon storage and admin API --- migrations/20260214000000_create_lexicons.sql | 8 + src/admin.rs | 281 ++++++++++++++++++ src/config.rs | 2 + src/error.rs | 3 + src/lexicon.rs | 182 ++++++++++++ src/main.rs | 11 + src/server.rs | 2 + 7 files changed, 489 insertions(+) create mode 100644 migrations/20260214000000_create_lexicons.sql create mode 100644 src/admin.rs create mode 100644 src/lexicon.rs diff --git a/migrations/20260214000000_create_lexicons.sql b/migrations/20260214000000_create_lexicons.sql new file mode 100644 index 0000000..fa3ceb7 --- /dev/null +++ b/migrations/20260214000000_create_lexicons.sql @@ -0,0 +1,8 @@ +CREATE TABLE IF NOT EXISTS lexicons ( + id TEXT PRIMARY KEY, + revision INT NOT NULL DEFAULT 1, + lexicon_json JSONB NOT NULL, + backfill BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); diff --git a/src/admin.rs b/src/admin.rs new file mode 100644 index 0000000..51637c2 --- /dev/null +++ b/src/admin.rs @@ -0,0 +1,281 @@ +use axum::extract::{Path, State}; +use axum::http::StatusCode; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::error::AppError; +use crate::lexicon::ParsedLexicon; +use crate::AppState; + +// --------------------------------------------------------------------------- +// Admin auth middleware +// --------------------------------------------------------------------------- + +/// Axum middleware layer that rejects requests without a valid admin secret. +pub fn admin_routes(_state: AppState) -> Router { + Router::new() + .route("/lexicons", post(upload_lexicon).get(list_lexicons)) + .route("/lexicons/{id}", get(get_lexicon).delete(delete_lexicon)) + .route("/stats", get(stats)) +} + +/// Extract and validate the admin Bearer token from request headers. +fn extract_admin_token(headers: &axum::http::HeaderMap, secret: &Option) -> Result<(), AppError> { + let secret = secret + .as_ref() + .ok_or_else(|| AppError::Auth("admin API is not configured".into()))?; + + let header = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .ok_or_else(|| AppError::Auth("missing Authorization header".into()))?; + + let token = header + .strip_prefix("Bearer ") + .ok_or_else(|| AppError::Auth("invalid Authorization scheme".into()))?; + + if token != secret { + return Err(AppError::Auth("invalid admin secret".into())); + } + + Ok(()) +} + +/// Axum extractor for admin auth — validates Bearer token against ADMIN_SECRET. +pub struct AdminAuth; + +impl axum::extract::FromRequestParts for AdminAuth { + type Rejection = AppError; + + async fn from_request_parts( + parts: &mut axum::http::request::Parts, + state: &AppState, + ) -> Result { + extract_admin_token(&parts.headers, &state.config.admin_secret)?; + Ok(AdminAuth) + } +} + +// --------------------------------------------------------------------------- +// Request / response types +// --------------------------------------------------------------------------- + +#[derive(Serialize)] +struct LexiconSummary { + id: String, + revision: i32, + lexicon_type: String, + backfill: bool, + created_at: chrono::DateTime, + updated_at: chrono::DateTime, +} + +#[derive(Deserialize)] +struct UploadLexiconBody { + lexicon_json: Value, + #[serde(default = "default_backfill")] + backfill: bool, +} + +fn default_backfill() -> bool { + true +} + +#[derive(Serialize)] +struct StatsResponse { + total_records: i64, + collections: Vec, +} + +#[derive(Serialize)] +struct CollectionStat { + collection: String, + count: i64, +} + +// --------------------------------------------------------------------------- +// Handlers +// --------------------------------------------------------------------------- + +/// POST /admin/lexicons — upload (upsert) a lexicon. +async fn upload_lexicon( + State(state): State, + _admin: AdminAuth, + Json(body): Json, +) -> Result<(StatusCode, Json), AppError> { + // Validate basic structure + let lexicon_version = body + .lexicon_json + .get("lexicon") + .and_then(|v| v.as_i64()) + .ok_or_else(|| AppError::BadRequest("lexicon JSON must have a numeric 'lexicon' field".into()))?; + + if lexicon_version != 1 { + return Err(AppError::BadRequest(format!( + "unsupported lexicon version: {lexicon_version}" + ))); + } + + let id = body + .lexicon_json + .get("id") + .and_then(|v| v.as_str()) + .ok_or_else(|| AppError::BadRequest("lexicon JSON must have a string 'id' field".into()))? + .to_string(); + + // Validate it parses correctly + ParsedLexicon::parse(body.lexicon_json.clone(), 1) + .map_err(|e| AppError::BadRequest(format!("failed to parse lexicon: {e}")))?; + + // Upsert into database + let row: (i32,) = sqlx::query_as( + r#" + INSERT INTO lexicons (id, lexicon_json, backfill) + VALUES ($1, $2, $3) + ON CONFLICT (id) DO UPDATE SET + lexicon_json = EXCLUDED.lexicon_json, + backfill = EXCLUDED.backfill, + revision = lexicons.revision + 1, + updated_at = NOW() + RETURNING revision + "#, + ) + .bind(&id) + .bind(&body.lexicon_json) + .bind(body.backfill) + .fetch_one(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to upsert lexicon: {e}")))?; + + let revision = row.0; + + // Update in-memory registry with correct revision + let parsed = ParsedLexicon::parse(body.lexicon_json, revision) + .map_err(|e| AppError::Internal(format!("failed to re-parse lexicon: {e}")))?; + state.lexicons.upsert(parsed).await; + + let status = if revision == 1 { + StatusCode::CREATED + } else { + StatusCode::OK + }; + + Ok(( + status, + Json(serde_json::json!({ + "id": id, + "revision": revision, + })), + )) +} + +/// GET /admin/lexicons — list all lexicons. +async fn list_lexicons( + State(state): State, + _admin: AdminAuth, +) -> Result>, AppError> { + let rows: Vec<(String, i32, Value, bool, chrono::DateTime, chrono::DateTime)> = + sqlx::query_as( + "SELECT id, revision, lexicon_json, backfill, created_at, updated_at FROM lexicons ORDER BY id", + ) + .fetch_all(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to list lexicons: {e}")))?; + + let summaries: Vec = rows + .into_iter() + .map(|(id, revision, json, backfill, created_at, updated_at)| { + let lexicon_type = ParsedLexicon::parse(json, revision) + .map(|p| format!("{:?}", p.lexicon_type).to_lowercase()) + .unwrap_or_else(|_| "unknown".into()); + + LexiconSummary { + id, + revision, + lexicon_type, + backfill, + created_at, + updated_at, + } + }) + .collect(); + + Ok(Json(summaries)) +} + +/// GET /admin/lexicons/:id — get a single lexicon. +async fn get_lexicon( + State(state): State, + _admin: AdminAuth, + Path(id): Path, +) -> Result, AppError> { + let row: Option<(String, i32, Value, bool, chrono::DateTime, chrono::DateTime)> = + sqlx::query_as( + "SELECT id, revision, lexicon_json, backfill, created_at, updated_at FROM lexicons WHERE id = $1", + ) + .bind(&id) + .fetch_optional(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to get lexicon: {e}")))?; + + let (id, revision, lexicon_json, backfill, created_at, updated_at) = + row.ok_or_else(|| AppError::NotFound(format!("lexicon '{id}' not found")))?; + + Ok(Json(serde_json::json!({ + "id": id, + "revision": revision, + "lexicon_json": lexicon_json, + "backfill": backfill, + "created_at": created_at, + "updated_at": updated_at, + }))) +} + +/// DELETE /admin/lexicons/:id — remove a lexicon. +async fn delete_lexicon( + State(state): State, + _admin: AdminAuth, + Path(id): Path, +) -> Result { + let result = sqlx::query("DELETE FROM lexicons WHERE id = $1") + .bind(&id) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to delete lexicon: {e}")))?; + + if result.rows_affected() == 0 { + return Err(AppError::NotFound(format!("lexicon '{id}' not found"))); + } + + state.lexicons.remove(&id).await; + + Ok(StatusCode::NO_CONTENT) +} + +/// GET /admin/stats — system statistics. +async fn stats( + State(state): State, + _admin: AdminAuth, +) -> Result, AppError> { + let total: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM records") + .fetch_one(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to count records: {e}")))?; + + let collections: Vec<(String, i64)> = sqlx::query_as( + "SELECT collection, COUNT(*) FROM records GROUP BY collection ORDER BY collection", + ) + .fetch_all(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to count by collection: {e}")))?; + + Ok(Json(StatsResponse { + total_records: total.0, + collections: collections + .into_iter() + .map(|(collection, count)| CollectionStat { collection, count }) + .collect(), + })) +} diff --git a/src/config.rs b/src/config.rs index 2e460d7..e8fbd8b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -8,6 +8,7 @@ pub struct Config { pub database_url: String, pub aip_url: String, pub jetstream_url: String, + pub admin_secret: Option, } impl Config { @@ -24,6 +25,7 @@ impl Config { .expect("AIP_URL must be set"), jetstream_url: env::var("JETSTREAM_URL") .unwrap_or_else(|_| "wss://jetstream2.us-west.bsky.network/subscribe".into()), + admin_secret: env::var("ADMIN_SECRET").ok(), } } diff --git a/src/error.rs b/src/error.rs index 7668c41..471f942 100644 --- a/src/error.rs +++ b/src/error.rs @@ -5,6 +5,7 @@ use bytes::Bytes; #[derive(Debug)] pub enum AppError { Auth(String), + BadRequest(String), Internal(String), NotFound(String), PdsError(StatusCode, Bytes), @@ -14,6 +15,7 @@ impl std::fmt::Display for AppError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { AppError::Auth(msg) => write!(f, "auth error: {msg}"), + AppError::BadRequest(msg) => write!(f, "bad request: {msg}"), AppError::Internal(msg) => write!(f, "internal error: {msg}"), AppError::NotFound(msg) => write!(f, "not found: {msg}"), AppError::PdsError(status, _) => write!(f, "PDS error: {status}"), @@ -33,6 +35,7 @@ impl IntoResponse for AppError { other => { let (status, message) = match &other { AppError::Auth(msg) => (StatusCode::UNAUTHORIZED, msg.clone()), + AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()), AppError::Internal(msg) => { tracing::error!("{msg}"); ( diff --git a/src/lexicon.rs b/src/lexicon.rs new file mode 100644 index 0000000..5ab9d18 --- /dev/null +++ b/src/lexicon.rs @@ -0,0 +1,182 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::RwLock; +use tracing::{info, warn}; + +/// The type of a lexicon's `main` definition. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum LexiconType { + Record, + Query, + Procedure, + /// Lexicons with no `main` def or a non-endpoint main type (token, object, string, etc.). + Definitions, +} + +/// Metadata extracted from a raw lexicon JSON document. +#[derive(Debug, Clone)] +pub struct ParsedLexicon { + /// The NSID, e.g. "games.gamesgamesgamesgames.game". + pub id: String, + /// What kind of endpoint this lexicon defines. + pub lexicon_type: LexiconType, + /// For records: the `key` field (e.g. "tid", "any", "literal:self"). + pub record_key: Option, + /// For queries: the parameters schema from `defs.main.parameters`. + pub parameters: Option, + /// For procedures: the input schema from `defs.main.input`. + pub input: Option, + /// For queries and procedures: the output schema from `defs.main.output`. + pub output: Option, + /// For records: the record schema from `defs.main.record`. + pub record_schema: Option, + /// The entire raw lexicon JSON. + pub raw: Value, + /// Database revision number. + pub revision: i32, +} + +impl ParsedLexicon { + /// Parse a lexicon JSON document into a `ParsedLexicon`. + pub fn parse(raw: Value, revision: i32) -> Result { + let id = raw + .get("id") + .and_then(|v| v.as_str()) + .ok_or("lexicon JSON missing 'id' field")? + .to_string(); + + let main_def = raw.get("defs").and_then(|d| d.get("main")); + + let main_type_str = main_def + .and_then(|m| m.get("type")) + .and_then(|t| t.as_str()); + + let lexicon_type = match main_type_str { + Some("record") => LexiconType::Record, + Some("query") => LexiconType::Query, + Some("procedure") => LexiconType::Procedure, + _ => LexiconType::Definitions, + }; + + let record_key = main_def + .and_then(|m| m.get("key")) + .and_then(|k| k.as_str()) + .map(|s| s.to_string()); + + let parameters = main_def.and_then(|m| m.get("parameters")).cloned(); + let input = main_def.and_then(|m| m.get("input")).cloned(); + let output = main_def.and_then(|m| m.get("output")).cloned(); + let record_schema = main_def.and_then(|m| m.get("record")).cloned(); + + Ok(Self { + id, + lexicon_type, + record_key, + parameters, + input, + output, + record_schema, + raw, + revision, + }) + } +} + +/// In-memory cache of all parsed lexicons, keyed by NSID. +#[derive(Debug, Clone)] +pub struct LexiconRegistry { + inner: Arc>>, +} + +impl LexiconRegistry { + pub fn new() -> Self { + Self { + inner: Arc::new(RwLock::new(HashMap::new())), + } + } + + /// Load all lexicons from the database, replacing any existing entries. + pub async fn load_from_db(&self, db: &sqlx::PgPool) -> Result<(), String> { + let rows: Vec<(String, Value, i32)> = + sqlx::query_as("SELECT id, lexicon_json, revision FROM lexicons") + .fetch_all(db) + .await + .map_err(|e| format!("failed to load lexicons: {e}"))?; + + let mut inner = self.inner.write().await; + inner.clear(); + + let mut loaded = 0u32; + for (id, json, revision) in rows { + match ParsedLexicon::parse(json, revision) { + Ok(parsed) => { + inner.insert(id, parsed); + loaded += 1; + } + Err(e) => { + warn!(%id, "failed to parse lexicon: {e}"); + } + } + } + + info!(count = loaded, "loaded lexicons into registry"); + Ok(()) + } + + /// Insert or update a single lexicon in the registry. + pub async fn upsert(&self, parsed: ParsedLexicon) { + let mut inner = self.inner.write().await; + inner.insert(parsed.id.clone(), parsed); + } + + /// Remove a lexicon from the registry by NSID. + pub async fn remove(&self, id: &str) -> bool { + let mut inner = self.inner.write().await; + inner.remove(id).is_some() + } + + /// Get a single lexicon by NSID. + pub async fn get(&self, id: &str) -> Option { + let inner = self.inner.read().await; + inner.get(id).cloned() + } + + /// Return NSIDs of all record-type lexicons. + pub async fn get_record_collections(&self) -> Vec { + let inner = self.inner.read().await; + inner + .values() + .filter(|lex| lex.lexicon_type == LexiconType::Record) + .map(|lex| lex.id.clone()) + .collect() + } + + /// Return NSIDs of all query-type lexicons. + pub async fn get_queries(&self) -> Vec { + let inner = self.inner.read().await; + inner + .values() + .filter(|lex| lex.lexicon_type == LexiconType::Query) + .map(|lex| lex.id.clone()) + .collect() + } + + /// Return NSIDs of all procedure-type lexicons. + pub async fn get_procedures(&self) -> Vec { + let inner = self.inner.read().await; + inner + .values() + .filter(|lex| lex.lexicon_type == LexiconType::Procedure) + .map(|lex| lex.id.clone()) + .collect() + } + + /// Return the total count of registered lexicons. + pub async fn count(&self) -> usize { + let inner = self.inner.read().await; + inner.len() + } +} diff --git a/src/main.rs b/src/main.rs index 86c34f4..ad91323 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,12 +1,15 @@ +mod admin; mod auth; mod config; mod error; mod jetstream; +mod lexicon; mod profile; mod repo; mod server; use config::Config; +use lexicon::LexiconRegistry; use tracing::info; #[derive(Clone)] @@ -14,6 +17,7 @@ pub struct AppState { pub config: Config, pub http: reqwest::Client, pub db: sqlx::PgPool, + pub lexicons: LexiconRegistry, } #[tokio::main] @@ -41,10 +45,17 @@ async fn main() { .await .expect("failed to run migrations"); + let lexicons = LexiconRegistry::new(); + lexicons + .load_from_db(&db) + .await + .expect("failed to load lexicons"); + let state = AppState { config: config.clone(), http: reqwest::Client::new(), db, + lexicons, }; jetstream::spawn(state.db.clone(), config.jetstream_url.clone()); diff --git a/src/server.rs b/src/server.rs index 95a8e26..aa4605b 100644 --- a/src/server.rs +++ b/src/server.rs @@ -4,6 +4,7 @@ use axum::{Json, Router}; use tower_http::cors::CorsLayer; use tower_http::trace::TraceLayer; +use crate::admin; use crate::auth::Claims; use crate::error::AppError; use crate::profile; @@ -13,6 +14,7 @@ use crate::AppState; pub fn router(state: AppState) -> Router { Router::new() .route("/health", get(health)) + .nest("/admin", admin::admin_routes(state.clone())) .route("/xrpc/app.bsky.actor.getProfile", get(get_profile)) .route( "/xrpc/games.gamesgamesgamesgames.createGame", -- 2.54.0