athena_rs 0.83.0

use actix_web::{HttpRequest, HttpResponse, Responder, get, web};
use serde::Serialize;
use serde_json::json;
use sqlx::postgres::{PgPool, PgRow};
use sqlx::{Pool, Postgres, Row};
use web::{Data, Query};

use crate::AppState;
use crate::api::auth::authorize_static_admin_key;
use crate::api::response::processed_error;
use crate::drivers::postgresql::sqlx_driver::PostgresClientRegistry;
use crate::error::sqlx_parser::process_sqlx_error_with_context;
use crate::parser::query_builder::sanitize_identifier;

#[derive(Serialize)]
/// Payload returned by `/schema/clients`.
struct SchemaClients {
    clients: Vec<String>,
}

#[derive(Serialize)]
/// Lightweight table metadata (schema + table name) returned by `/schema/tables`.
pub struct SchemaTable {
    table_schema: String,
    table_name: String,
}

#[derive(Serialize)]
/// Column metadata returned by `/schema/columns`.
struct SchemaColumn {
    column_name: String,
    data_type: Option<String>,
    column_default: Option<String>,
    is_nullable: Option<String>,
}

#[derive(Serialize)]
/// Migration record returned by `/schema/migrations` (Supabase-style schema_migrations).
struct SchemaMigration {
    version: Option<String>,
    name: Option<String>,
}

/// Extracts the `X-Athena-Client` header or returns a 400 error response.
fn client_header(req: &HttpRequest) -> Result<String, HttpResponse> {
    req.headers()
        .get("x-athena-client")
        .and_then(|value| value.to_str().ok())
        .filter(|value| !value.is_empty())
        .map(ToString::to_string)
        .ok_or_else(|| {
            HttpResponse::BadRequest().json(json!({
                "error": "Missing X-Athena-Client header"
            }))
        })
}

/// Resolves the Postgres pool for the requested client or returns `400 Bad Request`.
fn pool_for_client(
    client_name: &str,
    registry: &PostgresClientRegistry,
) -> Result<PgPool, HttpResponse> {
    registry.get_pool(client_name).ok_or_else(|| {
        HttpResponse::BadRequest().json(json!({
            "error": format!("Unknown Postgres client '{}'", client_name)
        }))
    })
}

#[get("/schema/clients")]
/// Returns an array of configured Postgres clients for selection by callers.
///
/// Protected by ATHENA_KEY_12. Use `?api_key=`, `Authorization: Bearer <key>`, or
/// `X-Athena-Key: <key>`.
async fn schema_clients(req: HttpRequest, app_state: Data<AppState>) -> impl Responder {
    if let Err(resp) = validate_athena_key_12(&req) {
        return resp;
    }
    let clients: Vec<String> = app_state.pg_registry.list_clients();
    HttpResponse::Ok().json(SchemaClients { clients })
}

/// Extracts and validates ATHENA_KEY_12 for protected endpoints.
fn validate_athena_key_12(req: &HttpRequest) -> Result<(), HttpResponse> {
    authorize_static_admin_key(req)
}

#[get("/clients")]
/// Returns an array of configured Athena/Postgres clients.
///
/// Protected by ATHENA_KEY_12. Use `?api_key=`, `Authorization: Bearer <key>`, or
/// `X-Athena-Key: <key>`.
async fn list_clients_protected(req: HttpRequest, app_state: Data<AppState>) -> impl Responder {
    if let Err(resp) = validate_athena_key_12(&req) {
        return resp;
    }
    let clients: Vec<String> = app_state.pg_registry.list_clients();
    HttpResponse::Ok().json(SchemaClients { clients })
}

#[get("/schema/tables")]
/// Lists all tables visible to the supplied `X-Athena-Client`.
///
/// Requires `X-Athena-Client` to look up the appropriate Postgres pool before querying.
///
/// The handler queries `information_schema.tables`, excludes Postgres system schemas,
/// and returns `{ "tables": [ { table_schema, table_name }, ... ] }`.
async fn schema_tables(req: HttpRequest, app_state: Data<AppState>) -> impl Responder {
    let client_name: String = match client_header(&req) {
        Ok(value) => value,
        Err(resp) => return resp,
    };

    let pool: Pool<Postgres> = match pool_for_client(&client_name, &app_state.pg_registry) {
        Ok(pool) => pool,
        Err(resp) => return resp,
    };

    let rows: Vec<PgRow> = match sqlx::query(
        r#"
        SELECT table_schema, table_name
        FROM information_schema.tables
        WHERE table_type = 'BASE TABLE'
          AND table_schema NOT IN ('information_schema', 'pg_catalog')
        ORDER BY table_schema, table_name
        "#,
    )
    .fetch_all(&pool)
    .await
    {
        Ok(rows) => rows,
        Err(err) => {
            return HttpResponse::InternalServerError().json(json!({
                "error": format!("Failed to fetch tables: {}", err)
            }));
        }
    };

    let tables: Vec<SchemaTable> = rows
        .into_iter()
        .map(|row| SchemaTable {
            table_schema: row.get("table_schema"),
            table_name: row.get("table_name"),
        })
        .collect();

    HttpResponse::Ok().json(json!({ "tables": tables }))
}

#[derive(serde::Deserialize)]
/// Query parameters accepted by `/schema/columns`.
struct ColumnQuery {
    table_name: String,
}

#[get("/schema/columns")]
/// Returns column metadata for `table_name` using `X-Athena-Client`.
///
/// Requires `X-Athena-Client`; the handler also validates and sanitizes the
/// `table_name` query parameter before querying `information_schema.columns`.
///
/// Validates `table_name`, fetches data from `information_schema.columns`, and
/// returns `{ "columns": [ ... ] }`.
async fn schema_columns(
    req: HttpRequest,
    app_state: Data<AppState>,
    query: Query<ColumnQuery>,
) -> impl Responder {
    let client_name: String = match client_header(&req) {
        Ok(value) => value,
        Err(resp) => return resp,
    };

    let pool: Pool<Postgres> = match pool_for_client(&client_name, &app_state.pg_registry) {
        Ok(pool) => pool,
        Err(resp) => return resp,
    };

    let table_identifier: String = match sanitize_identifier(&query.table_name) {
        Some(value) => value,
        None => {
            return HttpResponse::BadRequest().json(json!({
                "error": "Invalid table_name parameter"
            }));
        }
    };

    let rows: Vec<PgRow> = match sqlx::query(
        r#"
        SELECT column_name, data_type, column_default, is_nullable
        FROM information_schema.columns
        WHERE table_name = $1
        ORDER BY ordinal_position
        "#,
    )
    .bind(table_identifier)
    .fetch_all(&pool)
    .await
    {
        Ok(rows) => rows,
        Err(err) => {
            return HttpResponse::InternalServerError().json(json!({
                "error": format!("Failed to fetch columns: {}", err)
            }));
        }
    };

    let columns: Vec<SchemaColumn> = rows
        .into_iter()
        .map(|row| SchemaColumn {
            column_name: row.get("column_name"),
            data_type: row.get("data_type"),
            column_default: row.get("column_default"),
            is_nullable: row.get("is_nullable"),
        })
        .collect();

    HttpResponse::Ok().json(json!({ "columns": columns }))
}

#[get("/schema/migrations")]
/// Lists migration records from `supabase_migrations.schema_migrations` when available.
///
/// Requires `X-Athena-Client`. If the migrations table does not exist (e.g. non-Supabase
/// or fresh database), returns `{ "migrations": [], "message": "..." }` as a graceful
/// fallback instead of an error.
async fn schema_migrations(req: HttpRequest, app_state: Data<AppState>) -> impl Responder {
    let client_name: String = match client_header(&req) {
        Ok(value) => value,
        Err(resp) => return resp,
    };

    let pool: Pool<Postgres> = match pool_for_client(&client_name, &app_state.pg_registry) {
        Ok(pool) => pool,
        Err(resp) => return resp,
    };

    let result = sqlx::query(
        r#"
        SELECT version, name
        FROM supabase_migrations.schema_migrations
        ORDER BY version
        "#,
    )
    .fetch_all(&pool)
    .await;

    match result {
        Ok(rows) => {
            let migrations: Vec<SchemaMigration> = rows
                .into_iter()
                .map(|row: PgRow| SchemaMigration {
                    version: row.get("version"),
                    name: row.get("name"),
                })
                .collect();
            HttpResponse::Ok().json(json!({ "migrations": migrations }))
        }
        Err(err) => {
            // Graceful fallback: when the migrations table/schema does not exist, return empty
            let code: Option<String> = err
                .as_database_error()
                .and_then(|db_err| db_err.code().map(|c| c.to_string()));
            if matches!(code.as_deref(), Some("42P01") | Some("3F000")) {
                return HttpResponse::Ok().json(json!({
                    "migrations": [],
                    "message": "Migration tracking is not available for this database."
                }));
            }
            let processed = process_sqlx_error_with_context(&err, Some("schema_migrations"));
            processed_error(processed)
        }
    }
}

/// Registers the schema endpoints with the Actix router.
pub fn services(cfg: &mut web::ServiceConfig) {
    cfg.service(schema_clients)
        .service(list_clients_protected)
        .service(schema_tables)
        .service(schema_columns)
        .service(schema_migrations);
}