athena_rs 0.82.2

//! Resolves a Postgres pool from either X-Athena-Client (registry) or X-JDBC-URL (direct).
use actix_web::{HttpRequest, HttpResponse};
use serde_json::json;
use sqlx::postgres::PgPoolOptions;
use sqlx::{Pool, Postgres};

use crate::api::headers::x_athena_client::x_athena_client;
use crate::api::headers::x_jdbc_url::{jdbc_to_postgres_url, x_jdbc_url};
use crate::AppState;

/// Resolves a Postgres pool from the request.
///
/// Prefers `X-JDBC-URL` when present (direct connection); otherwise uses
/// `X-Athena-Client` to look up a configured pool. Returns an error response
/// if neither is provided or the lookup/connection fails.
pub async fn resolve_postgres_pool(
    req: &HttpRequest,
    app_state: &AppState,
) -> Result<Pool<Postgres>, HttpResponse> {
    if let Some(jdbc_url) = x_jdbc_url(req) {
        return resolve_pool_from_jdbc_url(&jdbc_url, app_state).await;
    }

    let client_name = x_athena_client(req);
    if client_name.is_empty() {
        return Err(HttpResponse::BadRequest().json(json!({
            "error": "X-Athena-Client or X-JDBC-URL header is required"
        })));
    }

    match app_state.pg_registry.get_pool(&client_name) {
        Some(pool) => Ok(pool),
        None => Err(HttpResponse::BadRequest().json(json!({
            "error": format!("Postgres client '{}' is not configured", client_name)
        }))),
    }
}

async fn resolve_pool_from_jdbc_url(
    jdbc_url: &str,
    app_state: &AppState,
) -> Result<Pool<Postgres>, HttpResponse> {
    let postgres_url = match jdbc_to_postgres_url(jdbc_url) {
        Some(url) => url,
        None => {
            return Err(HttpResponse::BadRequest().json(json!({
                "error": "X-JDBC-URL must be a valid PostgreSQL JDBC URL (jdbc:postgresql://...)"
            })));
        }
    };

    if let Some(pool) = app_state.jdbc_pool_cache.get(&postgres_url).await {
        return Ok(pool);
    }

    let pool = match PgPoolOptions::new()
        .max_connections(4)
        .acquire_timeout(std::time::Duration::from_secs(10))
        .connect(&postgres_url)
        .await
    {
        Ok(p) => p,
        Err(err) => {
            return Err(HttpResponse::BadGateway().json(json!({
                "error": format!("Failed to connect to database: {}", err)
            })));
        }
    };

    app_state
        .jdbc_pool_cache
        .insert(postgres_url.clone(), pool.clone())
        .await;

    Ok(pool)
}