Mercury/src/routes/admin/tables.rs

use axum::{
    extract::{Extension, Path, State},
    http::StatusCode,
    Json,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sqlx::Row;

use crate::{auth::Claims, routes::crud::pg_row_to_json, state::AppState};

const ALLOWED_TYPES: &[&str] = &[
    "TEXT",
    "INTEGER",
    "BIGINT",
    "SMALLINT",
    "BOOLEAN",
    "NUMERIC",
    "FLOAT4",
    "FLOAT8",
    "UUID",
    "TIMESTAMPTZ",
    "DATE",
    "JSONB",
];

const PROTECTED_TABLES: &[&str] = &["users", "blacklist", "api_keys", "queries", "permissions"];

fn is_protected(name: &str) -> bool {
    let lower = name.to_lowercase();
    PROTECTED_TABLES.iter().any(|&t| t == lower)
}

#[derive(Debug, Deserialize)]
pub struct ColumnDef {
    pub name: String,
    pub col_type: String,
    pub nullable: Option<bool>,
}

#[derive(Debug, Deserialize)]
pub struct CreateTableRequest {
    pub name: String,
    pub columns: Vec<ColumnDef>,
}

#[derive(Debug, Serialize)]
pub struct TableInfo {
    pub table_name: String,
    pub column_count: i64,
}

#[derive(Debug, Serialize)]
pub struct ColumnInfo {
    pub column_name: String,
    pub data_type: String,
    pub is_nullable: String,
}

#[derive(Debug, Serialize)]
pub struct TablePreview {
    pub table_name: String,
    pub row_count: i64,
    pub columns: Vec<ColumnInfo>,
    pub sample_rows: Vec<Value>,
}

pub async fn list_tables(
    State(state): State<AppState>,
    Extension(_claims): Extension<Claims>,
) -> Result<Json<Vec<TableInfo>>, StatusCode> {
    let rows = sqlx::query(
        r#"
        SELECT t.table_name,
               COUNT(c.column_name)::bigint AS column_count
        FROM information_schema.tables t
        LEFT JOIN information_schema.columns c
               ON c.table_schema = t.table_schema AND c.table_name = t.table_name
        WHERE t.table_schema = 'public' AND t.table_type = 'BASE TABLE'
        GROUP BY t.table_name
        ORDER BY t.table_name
        "#,
    )
    .fetch_all(&state.pool)
    .await
    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let tables = rows
        .into_iter()
        .map(|r| TableInfo {
            table_name: r.try_get::<String, _>("table_name").unwrap_or_default(),
            column_count: r.try_get::<i64, _>("column_count").unwrap_or(0),
        })
        .collect();

    Ok(Json(tables))
}

pub async fn get_table_preview(
    State(state): State<AppState>,
    Extension(_claims): Extension<Claims>,
    Path(name): Path<String>,
) -> Result<Json<TablePreview>, StatusCode> {
    if !crate::routes::is_valid_identifier(&name) {
        return Err(StatusCode::BAD_REQUEST);
    }

    let col_rows = sqlx::query(
        "SELECT column_name, data_type, is_nullable \
         FROM information_schema.columns \
         WHERE table_schema = 'public' AND table_name = $1 \
         ORDER BY ordinal_position",
    )
    .bind(&name)
    .fetch_all(&state.pool)
    .await
    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    if col_rows.is_empty() {
        return Err(StatusCode::NOT_FOUND);
    }

    let columns: Vec<ColumnInfo> = col_rows
        .into_iter()
        .map(|r| ColumnInfo {
            column_name: r.try_get::<String, _>("column_name").unwrap_or_default(),
            data_type: r.try_get::<String, _>("data_type").unwrap_or_default(),
            is_nullable: r.try_get::<String, _>("is_nullable").unwrap_or_default(),
        })
        .collect();

    let count_sql = format!("SELECT COUNT(*)::bigint FROM {}", name);
    let row_count: i64 = sqlx::query(&count_sql)
        .fetch_one(&state.pool)
        .await
        .map(|r| r.try_get::<i64, _>(0).unwrap_or(0))
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let sample_sql = format!("SELECT * FROM {} LIMIT 10", name);
    let sample_rows: Vec<Value> = sqlx::query(&sample_sql)
        .fetch_all(&state.pool)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .into_iter()
        .map(pg_row_to_json)
        .collect();

    Ok(Json(TablePreview {
        table_name: name,
        row_count,
        columns,
        sample_rows,
    }))
}

pub async fn create_table(
    State(state): State<AppState>,
    Extension(_claims): Extension<Claims>,
    Json(body): Json<CreateTableRequest>,
) -> Result<Json<TableInfo>, StatusCode> {
    if !crate::routes::is_valid_identifier(&body.name) {
        return Err(StatusCode::BAD_REQUEST);
    }
    if is_protected(&body.name) {
        return Err(StatusCode::FORBIDDEN);
    }

    if body.columns.is_empty() {
        return Err(StatusCode::BAD_REQUEST);
    }

    let mut col_defs = vec!["id SERIAL PRIMARY KEY".to_string()];

    for col in &body.columns {
        if !crate::routes::is_valid_identifier(&col.name) {
            return Err(StatusCode::BAD_REQUEST);
        }

        let upper_type = col.col_type.to_uppercase();
        if !ALLOWED_TYPES.contains(&upper_type.as_str()) {
            return Err(StatusCode::UNPROCESSABLE_ENTITY);
        }

        let nullable = col.nullable.unwrap_or(true);
        let null_clause = if nullable { "" } else { " NOT NULL" };
        col_defs.push(format!("{} {}{}", col.name, upper_type, null_clause));
    }

    let sql = format!("CREATE TABLE {} ({})", body.name, col_defs.join(", "));

    sqlx::query(&sql).execute(&state.pool).await.map_err(|e| {
        tracing::error!("create table error: {}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    Ok(Json(TableInfo {
        table_name: body.name,
        column_count: body.columns.len() as i64 + 1,
    }))
}

pub async fn drop_table(
    State(state): State<AppState>,
    Extension(_claims): Extension<Claims>,
    Path(name): Path<String>,
) -> Result<Json<Value>, StatusCode> {
    if !crate::routes::is_valid_identifier(&name) {
        return Err(StatusCode::BAD_REQUEST);
    }
    if is_protected(&name) {
        return Err(StatusCode::FORBIDDEN);
    }

    let sql = format!("DROP TABLE IF EXISTS {}", name);
    sqlx::query(&sql).execute(&state.pool).await.map_err(|e| {
        tracing::error!("drop table error: {}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    Ok(Json(serde_json::json!({ "dropped": true })))
}

#[cfg(test)]
mod tests {
    use super::{is_protected, PROTECTED_TABLES};

    #[test]
    fn test_protected_tables_list() {
        assert!(PROTECTED_TABLES.contains(&"users"));
        assert!(PROTECTED_TABLES.contains(&"blacklist"));
        assert!(PROTECTED_TABLES.contains(&"api_keys"));
        assert!(PROTECTED_TABLES.contains(&"queries"));
        assert!(PROTECTED_TABLES.contains(&"permissions"));
    }

    #[test]
    fn test_is_protected() {
        assert!(is_protected("users"));
        assert!(is_protected("USERS"));
        assert!(!is_protected("orders"));
    }
}