288 lines
10 KiB
Rust
288 lines
10 KiB
Rust
#![allow(clippy::unwrap_used)]
|
|
//! tidalDB + Axum: production-ready embedding patterns.
|
|
//!
|
|
//! Demonstrates:
|
|
//! - Schema defined at startup before `axum::serve`
|
|
//! - `Arc<TidalDb>` shared state across handlers
|
|
//! - Three routes: `POST /signal`, `GET /feed`, `GET /health`
|
|
//! - `TidalError` to HTTP status mapping
|
|
//! - Graceful shutdown with `db.close()`
|
|
//!
|
|
//! `TidalDb` is not `Clone`, so it must be wrapped in `Arc` for use with
|
|
//! Axum's `State<T>`, which requires `T: Clone + Send + Sync + 'static`.
|
|
//!
|
|
//! # Running
|
|
//!
|
|
//! ```bash
|
|
//! cargo run --example axum_embedding --manifest-path tidal/Cargo.toml
|
|
//! # Then:
|
|
//! # curl http://127.0.0.1:3000/health
|
|
//! # curl "http://127.0.0.1:3000/feed?user_id=1"
|
|
//! # curl -X POST http://127.0.0.1:3000/signal \
|
|
//! # -H 'Content-Type: application/json' \
|
|
//! # -d '{"entity_id": 1, "signal": "view", "weight": 1.0}'
|
|
//! ```
|
|
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use axum::extract::{Query, State};
|
|
use axum::http::StatusCode;
|
|
use axum::response::IntoResponse;
|
|
use axum::routing::{get, post};
|
|
use axum::{Json, Router};
|
|
use tidaldb::schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window};
|
|
use tidaldb::{TidalDb, TidalError};
|
|
|
|
// ── Request / Response types ────────────────────────────────────────────────
|
|
|
|
/// JSON body for `POST /signal`.
|
|
#[derive(serde::Deserialize)]
|
|
struct SignalRequest {
|
|
entity_id: u64,
|
|
signal: String,
|
|
weight: f64,
|
|
}
|
|
|
|
/// Query parameters for `GET /feed`.
|
|
#[derive(serde::Deserialize)]
|
|
struct FeedQuery {
|
|
user_id: u64,
|
|
}
|
|
|
|
/// JSON response for `GET /health`.
|
|
#[derive(serde::Serialize)]
|
|
struct HealthResponse {
|
|
status: &'static str,
|
|
items: u64,
|
|
}
|
|
|
|
/// JSON response for `GET /feed`.
|
|
#[derive(serde::Serialize)]
|
|
struct FeedResponse {
|
|
items: Vec<FeedItem>,
|
|
total_candidates: usize,
|
|
}
|
|
|
|
#[derive(serde::Serialize)]
|
|
struct FeedItem {
|
|
entity_id: u64,
|
|
score: f64,
|
|
rank: usize,
|
|
}
|
|
|
|
// ── Error mapping ───────────────────────────────────────────────────────────
|
|
|
|
/// Maps `TidalError` variants to appropriate HTTP status codes.
|
|
///
|
|
/// This is the recommended pattern for wrapping tidalDB errors in an HTTP API.
|
|
/// Each error variant maps to the semantically correct HTTP status:
|
|
///
|
|
/// - `Backpressure` / `RateLimited` -> 429 Too Many Requests (client should retry)
|
|
/// - `NotFound` -> 404 Not Found
|
|
/// - `Schema` / `Config` / `InvalidInput` -> 400 Bad Request (client's fault)
|
|
/// - `Query` -> 400 Bad Request (malformed query)
|
|
/// - `PolicyViolation` / `SessionExpired` -> 403 Forbidden
|
|
/// - Everything else -> 500 Internal Server Error
|
|
fn error_to_status(err: &TidalError) -> StatusCode {
|
|
match err {
|
|
TidalError::Backpressure { .. } | TidalError::RateLimited { .. } => {
|
|
StatusCode::TOO_MANY_REQUESTS
|
|
}
|
|
TidalError::NotFound { .. } => StatusCode::NOT_FOUND,
|
|
TidalError::Schema(_) | TidalError::Config(_) | TidalError::InvalidInput(_) => {
|
|
StatusCode::BAD_REQUEST
|
|
}
|
|
TidalError::Query(_) => StatusCode::BAD_REQUEST,
|
|
TidalError::PolicyViolation { .. } | TidalError::SessionExpired { .. } => {
|
|
StatusCode::FORBIDDEN
|
|
}
|
|
// Storage, Durability, Internal -- all server-side issues.
|
|
_ => StatusCode::INTERNAL_SERVER_ERROR,
|
|
}
|
|
}
|
|
|
|
/// Axum-compatible error wrapper that converts `TidalError` into an HTTP response.
|
|
struct AppError(TidalError);
|
|
|
|
impl IntoResponse for AppError {
|
|
fn into_response(self) -> axum::response::Response {
|
|
let status = error_to_status(&self.0);
|
|
let body = serde_json::json!({ "error": self.0.to_string() });
|
|
(status, Json(body)).into_response()
|
|
}
|
|
}
|
|
|
|
impl From<TidalError> for AppError {
|
|
fn from(err: TidalError) -> Self {
|
|
Self(err)
|
|
}
|
|
}
|
|
|
|
// ── Handlers ────────────────────────────────────────────────────────────────
|
|
|
|
/// `POST /signal` -- record an engagement signal.
|
|
///
|
|
/// Accepts a JSON body with `entity_id`, `signal` type name, and `weight`.
|
|
/// The signal is recorded into the WAL-backed signal ledger and is immediately
|
|
/// reflected in subsequent ranking queries.
|
|
async fn signal_handler(
|
|
State(db): State<Arc<TidalDb>>,
|
|
Json(req): Json<SignalRequest>,
|
|
) -> Result<StatusCode, AppError> {
|
|
db.signal(
|
|
&req.signal,
|
|
EntityId::new(req.entity_id),
|
|
req.weight,
|
|
Timestamp::now(),
|
|
)?;
|
|
Ok(StatusCode::NO_CONTENT)
|
|
}
|
|
|
|
/// `GET /feed?user_id=123` -- retrieve a ranked feed.
|
|
///
|
|
/// Executes a RETRIEVE query using the `trending` profile. The profile
|
|
/// ranks items by share + view velocity with per-creator diversity.
|
|
/// Personalization is applied when `for_user` is set.
|
|
async fn feed_handler(
|
|
State(db): State<Arc<TidalDb>>,
|
|
Query(params): Query<FeedQuery>,
|
|
) -> Result<Json<FeedResponse>, AppError> {
|
|
let query = tidaldb::query::retrieve::Retrieve::builder()
|
|
.for_user(params.user_id)
|
|
.profile("trending")
|
|
.limit(20)
|
|
.build()
|
|
.map_err(TidalError::from)?;
|
|
|
|
let results = db.retrieve(&query)?;
|
|
|
|
let items = results
|
|
.items
|
|
.iter()
|
|
.map(|r| FeedItem {
|
|
entity_id: r.entity_id.as_u64(),
|
|
score: r.score,
|
|
rank: r.rank,
|
|
})
|
|
.collect();
|
|
|
|
Ok(Json(FeedResponse {
|
|
items,
|
|
total_candidates: results.total_candidates,
|
|
}))
|
|
}
|
|
|
|
/// `GET /health` -- health check endpoint.
|
|
///
|
|
/// Returns 200 with item count if the database is operational,
|
|
/// or 500 if it has been closed or is in a degraded state.
|
|
async fn health_handler(State(db): State<Arc<TidalDb>>) -> Result<Json<HealthResponse>, AppError> {
|
|
db.health_check()?;
|
|
Ok(Json(HealthResponse {
|
|
status: "ok",
|
|
items: db.item_count(),
|
|
}))
|
|
}
|
|
|
|
// ── Main ────────────────────────────────────────────────────────────────────
|
|
|
|
/// Resolves when Ctrl+C is received.
|
|
async fn shutdown_signal() {
|
|
tokio::signal::ctrl_c()
|
|
.await
|
|
.unwrap_or_else(|e| eprintln!("ctrl-c error: {e}"));
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
tracing_subscriber::fmt()
|
|
.with_env_filter("tidaldb=info")
|
|
.init();
|
|
|
|
// ── Schema definition ───────────────────────────────────────────────
|
|
// Define the schema before opening the database. In production, this
|
|
// would match your application's signal types and text fields.
|
|
let mut schema = SchemaBuilder::new();
|
|
let _ = schema
|
|
.signal(
|
|
"view",
|
|
EntityKind::Item,
|
|
DecaySpec::Exponential {
|
|
half_life: Duration::from_secs(7 * 24 * 3600),
|
|
},
|
|
)
|
|
.windows(&[Window::OneHour, Window::TwentyFourHours, Window::AllTime])
|
|
.velocity(true)
|
|
.add();
|
|
let _ = schema
|
|
.signal(
|
|
"like",
|
|
EntityKind::Item,
|
|
DecaySpec::Exponential {
|
|
half_life: Duration::from_secs(30 * 24 * 3600),
|
|
},
|
|
)
|
|
.windows(&[Window::AllTime])
|
|
.velocity(false)
|
|
.add();
|
|
let _ = schema
|
|
.signal(
|
|
"share",
|
|
EntityKind::Item,
|
|
DecaySpec::Exponential {
|
|
half_life: Duration::from_secs(3 * 24 * 3600),
|
|
},
|
|
)
|
|
.windows(&[Window::TwentyFourHours, Window::AllTime])
|
|
.velocity(true)
|
|
.add();
|
|
let schema = schema.build()?;
|
|
|
|
// ── Database initialization ─────────────────────────────────────────
|
|
// Wrap in Arc: TidalDb is not Clone, but Axum's State<T> requires Clone.
|
|
// The Arc lets the router and main both hold references.
|
|
let db = Arc::new(TidalDb::builder().ephemeral().with_schema(schema).open()?);
|
|
|
|
// ── Router ──────────────────────────────────────────────────────────
|
|
let app = Router::new()
|
|
.route("/signal", post(signal_handler))
|
|
.route("/feed", get(feed_handler))
|
|
.route("/health", get(health_handler))
|
|
.with_state(Arc::clone(&db));
|
|
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await?;
|
|
let addr = listener.local_addr()?;
|
|
println!("listening on http://{addr}");
|
|
println!(" POST /signal -> record engagement signal");
|
|
println!(" GET /feed?user_id=123 -> ranked feed (trending profile)");
|
|
println!(" GET /health -> health check + item count");
|
|
println!("press Ctrl+C to stop");
|
|
|
|
axum::serve(listener, app)
|
|
.with_graceful_shutdown(shutdown_signal())
|
|
.await?;
|
|
|
|
// ── Graceful shutdown ───────────────────────────────────────────────
|
|
// After axum::serve returns (Ctrl+C received), we own the only remaining
|
|
// Arc reference. try_unwrap extracts the TidalDb so we can call close()
|
|
// for a clean shutdown (WAL drain, checkpoint, storage flush).
|
|
//
|
|
// If try_unwrap fails (another reference still exists), the Drop impl
|
|
// will still perform best-effort cleanup.
|
|
match Arc::try_unwrap(db) {
|
|
Ok(db) => {
|
|
println!("shutting down tidalDB...");
|
|
db.close()?;
|
|
println!("tidalDB closed cleanly.");
|
|
}
|
|
Err(arc) => {
|
|
// Drop triggers best-effort shutdown_inner().
|
|
drop(arc);
|
|
println!("tidalDB dropped (best-effort shutdown).");
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|