lakehouse/crates/storaged/src/federation_service.rs

/// HTTP surface for federation operator endpoints.
///
/// - `GET /buckets` — configured buckets with reachability status
/// - `GET /errors` — recent bucket op failures (filterable)
/// - `GET /bucket-health` — aggregated errors-per-bucket in the last 5 minutes
///
/// Mounted under `/storage` by the gateway. `/storage/health` is already
/// claimed by the existing storage router for service liveness, so the
/// aggregated bucket health lives at `/storage/bucket-health`.

use axum::{
    Json, Router,
    extract::{Path, Query, State},
    http::StatusCode,
    response::IntoResponse,
    routing::{delete, get, post, put},
};
use chrono::{DateTime, Utc};
use serde::Deserialize;
use shared::config::BucketConfig;
use std::sync::Arc;

use crate::error_journal::BucketErrorEvent;
use crate::registry::{BucketInfo, BucketRegistry};

pub fn router(registry: Arc<BucketRegistry>) -> Router {
    Router::new()
        .route("/buckets", get(list_buckets).post(add_bucket))
        .route("/buckets/{name}", delete(remove_bucket))
        .route("/errors", get(list_errors))
        .route("/errors/flush", post(flush_errors))
        .route("/errors/compact", post(compact_errors))
        .route("/bucket-health", get(get_health))
        .route("/buckets/{bucket}/objects/{*key}", put(put_bucket_object))
        .route("/buckets/{bucket}/objects/{*key}", get(get_bucket_object))
        .with_state(registry)
}

/// Provision + register a bucket at runtime. Body is a `BucketConfig`.
///
/// Federation layer 2: profile buckets (`profile:alice`) and tenant
/// buckets can be added after startup without a service restart.
async fn add_bucket(
    State(reg): State<Arc<BucketRegistry>>,
    Json(bc): Json<BucketConfig>,
) -> impl IntoResponse {
    // Safety net: local backends must land somewhere under the configured
    // profile_root (for profile:*) or be otherwise explicitly rooted.
    // Refuse empty / missing root to avoid "bucket created at ./" surprises.
    if bc.backend == "local" {
        match bc.root.as_deref() {
            Some(r) if !r.trim().is_empty() => {}
            _ => return Err((StatusCode::BAD_REQUEST,
                "local bucket requires a non-empty 'root' path".to_string())),
        }
    }
    match reg.add_bucket(bc).await {
        Ok(info) => Ok((StatusCode::CREATED, Json(info))),
        Err(e) => {
            let code = if e.contains("already registered") {
                StatusCode::CONFLICT
            } else {
                StatusCode::BAD_REQUEST
            };
            Err((code, e))
        }
    }
}

/// Unregister a bucket. Refused for primary / rescue / unknown names.
async fn remove_bucket(
    State(reg): State<Arc<BucketRegistry>>,
    Path(name): Path<String>,
) -> impl IntoResponse {
    match reg.remove_bucket(&name) {
        Ok(()) => Ok((StatusCode::OK, format!("unregistered: {name}"))),
        Err(e) => {
            let code = if e.contains("cannot remove") {
                StatusCode::FORBIDDEN
            } else if e.contains("not registered") {
                StatusCode::NOT_FOUND
            } else {
                StatusCode::BAD_REQUEST
            };
            Err((code, e))
        }
    }
}

async fn list_buckets(State(reg): State<Arc<BucketRegistry>>) -> Json<Vec<BucketInfo>> {
    Json(reg.list().await)
}

#[derive(Deserialize)]
struct ErrorQuery {
    bucket: Option<String>,
    since: Option<DateTime<Utc>>,
    #[serde(default = "default_limit")]
    limit: usize,
}

fn default_limit() -> usize { 50 }

async fn list_errors(
    State(reg): State<Arc<BucketRegistry>>,
    Query(q): Query<ErrorQuery>,
) -> Json<Vec<BucketErrorEvent>> {
    Json(reg.journal().filter(q.bucket.as_deref(), q.since, q.limit).await)
}

#[derive(Deserialize)]
struct HealthQuery {
    #[serde(default = "default_period")]
    minutes: i64,
}

fn default_period() -> i64 { 5 }

async fn get_health(
    State(reg): State<Arc<BucketRegistry>>,
    Query(q): Query<HealthQuery>,
) -> impl IntoResponse {
    Json(reg.journal().health(q.minutes).await)
}

/// Bucket-aware write. Hard-fails if the target bucket is unreachable;
/// never falls back to rescue for writes.
async fn put_bucket_object(
    State(reg): State<Arc<BucketRegistry>>,
    Path((bucket, key)): Path<(String, String)>,
    body: bytes::Bytes,
) -> impl IntoResponse {
    match reg.write_smart(&bucket, &key, body).await {
        Ok(()) => (StatusCode::CREATED, format!("stored: {bucket}/{key}")).into_response(),
        Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e).into_response(),
    }
}

/// Bucket-aware read. Falls through to the rescue bucket if the target
/// is unreachable or the key is missing; emits observability headers so
/// callers can detect the fallback.
async fn flush_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
    match reg.journal().flush().await {
        Ok(()) => (StatusCode::OK, "flushed").into_response(),
        Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
    }
}

async fn compact_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
    match reg.journal().compact().await {
        Ok(stats) => Json(stats).into_response(),
        Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
    }
}

async fn get_bucket_object(
    State(reg): State<Arc<BucketRegistry>>,
    Path((bucket, key)): Path<(String, String)>,
) -> impl IntoResponse {
    match reg.read_smart(&bucket, &key).await {
        Ok(outcome) => {
            let mut resp = outcome.data.into_response();
            if outcome.rescued {
                let headers = resp.headers_mut();
                headers.insert("x-lakehouse-rescue-used", "true".parse().unwrap());
                headers.insert(
                    "x-lakehouse-original-bucket",
                    outcome.original_bucket.parse().unwrap(),
                );
                headers.insert(
                    "x-lakehouse-served-by",
                    outcome.served_by.parse().unwrap(),
                );
            }
            resp
        }
        Err(e) => (StatusCode::NOT_FOUND, e).into_response(),
    }
}