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, } #[derive(Debug, Deserialize)] pub struct CreateTableRequest { pub name: String, pub columns: Vec, } #[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, pub sample_rows: Vec, } pub async fn list_tables( State(state): State, Extension(_claims): Extension, ) -> Result>, 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::("table_name").unwrap_or_default(), column_count: r.try_get::("column_count").unwrap_or(0), }) .collect(); Ok(Json(tables)) } pub async fn get_table_preview( State(state): State, Extension(_claims): Extension, Path(name): Path, ) -> Result, 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 = col_rows .into_iter() .map(|r| ColumnInfo { column_name: r.try_get::("column_name").unwrap_or_default(), data_type: r.try_get::("data_type").unwrap_or_default(), is_nullable: r.try_get::("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::(0).unwrap_or(0)) .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let sample_sql = format!("SELECT * FROM {} LIMIT 10", name); let sample_rows: Vec = 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, Extension(_claims): Extension, Json(body): Json, ) -> Result, 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, Extension(_claims): Extension, Path(name): Path, ) -> Result, 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")); } }