athena_rs 3.3.0

//! Client and pool resolution for API handlers.
//!
//! Use [`required_client_name`] + [`pool_for_client`] or the combined
//! [`client_name_and_pool`] when a handler needs both the client name and a Postgres pool.

use actix_web::{HttpRequest, HttpResponse};
use sqlx::postgres::PgPool;

use crate::AppState;
use crate::api::response::{bad_request, service_unavailable};
use crate::drivers::postgresql::sqlx_driver::RegisteredClient;

pub const ATHENA_CLIENT_HEADER: &str = "X-Athena-Client";

/// Returns the client name from `X-Athena-Client` or an error response.
pub fn required_client_name(req: &HttpRequest) -> Result<String, HttpResponse> {
    req.headers()
        .get("x-athena-client")
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToString::to_string)
        .ok_or_else(|| {
            bad_request(
                "Missing required header",
                format!("{} header is required", ATHENA_CLIENT_HEADER),
            )
        })
}

/// Returns the Postgres pool for the given client name or an error response.
pub fn pool_for_client(state: &AppState, client_name: &str) -> Result<PgPool, HttpResponse> {
    if let Some(pool) = state.pg_registry.get_pool(client_name) {
        return Ok(pool);
    }

    if let Some(registered_client) = state.pg_registry.registered_client(client_name) {
        return Err(unavailable_registered_client_response(
            client_name,
            &registered_client,
        ));
    }

    Err(bad_request(
        format!("Client '{}' is not available in the registry", client_name),
        "Postgres client not configured",
    ))
}

/// Resolves both client name and pool in one call (e.g. for gateway handlers).
///
/// Equivalent to `required_client_name(req).and_then(|name| {
///     pool_for_client(state, &name).map(|pool| (name, pool))
/// })`.
pub fn client_name_and_pool(
    req: &HttpRequest,
    state: &AppState,
) -> Result<(String, PgPool), HttpResponse> {
    let name: String = required_client_name(req)?;
    let pool: PgPool = pool_for_client(state, &name)?;
    Ok((name, pool))
}

pub fn registered_client_for(
    state: &AppState,
    client_name: &str,
) -> Result<RegisteredClient, HttpResponse> {
    state
        .pg_registry
        .registered_client(client_name)
        .ok_or_else(|| {
            bad_request(
                "Unknown client",
                format!("Client '{}' is not registered.", client_name),
            )
        })
}

pub fn logging_pool(state: &AppState) -> Result<PgPool, HttpResponse> {
    let Some(client_name) = state.logging_client_name.as_ref() else {
        return Err(service_unavailable(
            "Logging store unavailable",
            "No athena logging client is configured.",
        ));
    };

    state.pg_registry.get_pool(client_name).ok_or_else(|| {
        service_unavailable(
            "Logging store unavailable",
            format!("Logging client '{}' is not connected.", client_name),
        )
    })
}

pub fn auth_pool(state: &AppState) -> Result<PgPool, HttpResponse> {
    let Some(client_name) = state.gateway_auth_client_name.as_ref() else {
        return Err(service_unavailable(
            "Auth store unavailable",
            "No gateway auth client is configured.",
        ));
    };

    state.pg_registry.get_pool(client_name).ok_or_else(|| {
        service_unavailable(
            "Auth store unavailable",
            format!("Gateway auth client '{}' is not connected.", client_name),
        )
    })
}

fn unavailable_registered_client_response(
    client_name: &str,
    registered_client: &RegisteredClient,
) -> HttpResponse {
    if !registered_client.is_active {
        return bad_request(
            format!("Client '{}' is inactive", client_name),
            "Postgres client is configured but inactive",
        );
    }

    if registered_client.is_frozen {
        return bad_request(
            format!("Client '{}' is frozen", client_name),
            "Postgres client is configured but frozen",
        );
    }

    service_unavailable(
        format!(
            "Client '{}' is configured but currently unavailable",
            client_name
        ),
        "Postgres client is configured but has no active connection pool",
    )
}