use axum::{
Extension,
body::Bytes,
extract::{Json, Path, Query, State},
http::{HeaderMap, StatusCode, header},
response::{Html, IntoResponse, Response},
};
use chrono::Utc;
use serde::Deserialize;
use serde_json::{Value, json};
use std::sync::atomic::Ordering;
use std::time::Instant;
use uuid::Uuid;
use crate::{
backup,
database::{
MAX_BATCH_ITEMS, execute_query_in_transaction, execute_query_with_pool,
execute_transaction_with_pool, get_or_create_pool, validate_query_permissions,
},
models::*,
};
#[derive(Debug, Deserialize, Default)]
pub struct BackupExportOptions {
#[serde(default)]
pub exclude_table: Vec<String>,
}
#[derive(Debug, Deserialize, Default)]
pub struct BackupImportOptions {
#[serde(default)]
pub clean: bool,
}
async fn backup_context(
state: &AppState,
account: &Account,
database: &str,
) -> Result<(PostgresInstance, DatabaseAccess), StatusCode> {
let access = resolve_database_access(state, account, database).await?;
let instance = state
.instances
.read()
.await
.get(&account.instance_id)
.cloned()
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
Ok((instance, access))
}
pub async fn export_database(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Path(db): Path<String>,
Query(options): Query<BackupExportOptions>,
) -> Result<Response, StatusCode> {
if !account.has_permission(Permission::Export) {
return Err(StatusCode::FORBIDDEN);
}
let (instance, access) = backup_context(&state, &account, &db).await?;
let bytes = backup::export_custom(&instance, &access, &options.exclude_table)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "application/octet-stream")
.header(
header::CONTENT_DISPOSITION,
format!("attachment; filename=\"{db}.dump\""),
)
.header("x-pg-api-backup-format", "custom")
.body(axum::body::Body::from(bytes))
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
}
pub async fn import_database(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Path(db): Path<String>,
Query(options): Query<BackupImportOptions>,
body: Bytes,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
let metadata = ResponseMetadata {
request_id: Uuid::new_v4().to_string(),
execution_time_ms: 0,
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
if !account.has_permission(Permission::Import) {
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
"IMPORT permission is required".into(),
metadata,
)));
}
let (instance, access) = backup_context(&state, &account, &db).await?;
match backup::import_custom(&instance, &access, &body, options.clean).await {
Ok(()) => Ok(Json(ApiResponse::success(
json!({"database": db, "restored": true, "clean": options.clean}),
metadata,
))),
Err(error) => Ok(Json(ApiResponse::error("BACKUP_ERROR", error, metadata))),
}
}
fn gate_error_code(msg: &str) -> &'static str {
if msg.starts_with("MULTI_STATEMENT") {
"MULTI_STATEMENT_NOT_ALLOWED"
} else if msg.starts_with("STATEMENT_NOT_ALLOWED") {
"STATEMENT_NOT_ALLOWED"
} else if msg.starts_with("STATEMENT_TIMEOUT") {
"STATEMENT_TIMEOUT"
} else {
"PERMISSION_DENIED"
}
}
pub async fn health_check() -> impl IntoResponse {
Json(json!({
"status": "healthy",
"service": env!("CARGO_PKG_NAME"),
"version": env!("CARGO_PKG_VERSION")
}))
}
pub async fn status_handler(State(state): State<AppState>) -> impl IntoResponse {
let instances = state.instances.read().await;
let connections = state.connections.len();
Json(json!({
"instances": instances.len(),
"active_connections": connections,
"status": "operational"
}))
}
pub async fn query_handler(
State(state): State<AppState>,
Extension(account): Extension<Account>,
headers: HeaderMap,
Json(request): Json<QueryRequest>,
) -> Result<Json<ApiResponse<QueryResult>>, StatusCode> {
let start = Instant::now();
let request_id = headers
.get("x-request-id")
.and_then(|v| v.to_str().ok())
.unwrap_or("unknown")
.to_string();
let mut metadata = ResponseMetadata {
request_id: request_id.clone(),
execution_time_ms: 0,
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
let query_type = request.query.split_whitespace().next().unwrap_or("UNKNOWN");
tracing::debug!(event_type = "query_received", request_id = %request_id, account_id = %account.id, account_name = %account.name, database = %request.database, query_type, "PG-API query received");
let db_access = resolve_database_access(&state, &account, &request.database).await?;
if let Err(e) =
validate_query_permissions(&request.query, &db_access.permissions, &account.role)
{
metadata.execution_time_ms = start.elapsed().as_millis();
let code = gate_error_code(&e);
return Ok(Json(ApiResponse::error(code, e, metadata)));
}
let pool_key = format!(
"{}-{}-{}",
account.instance_id, db_access.username, request.database
);
let pool = match get_or_create_pool(&state, &pool_key, &account, &db_access).await {
Ok(p) => p,
Err(e) => {
metadata.execution_time_ms = start.elapsed().as_millis();
return Ok(Json(ApiResponse::error("CONNECTION_ERROR", e, metadata)));
}
};
match execute_query_with_pool(pool, request.query, request.params, &request.options).await {
Ok(result) => {
metadata.execution_time_ms = start.elapsed().as_millis();
metadata.rows_affected = Some(result.rows_affected.unwrap_or(result.rows.len() as u64));
Ok(Json(ApiResponse::success(result, metadata)))
}
Err(e) => {
metadata.execution_time_ms = start.elapsed().as_millis();
let code = if e.starts_with("STATEMENT_TIMEOUT") {
"STATEMENT_TIMEOUT"
} else {
"QUERY_ERROR"
};
Ok(Json(ApiResponse::error(code, e, metadata)))
}
}
}
pub async fn batch_query_handler(
State(state): State<AppState>,
Extension(account): Extension<Account>,
headers: HeaderMap,
Json(requests): Json<Vec<QueryRequest>>,
) -> Result<Json<ApiResponse<Vec<QueryResult>>>, StatusCode> {
let start = Instant::now();
let request_id = Uuid::new_v4().to_string();
if requests.len() > MAX_BATCH_ITEMS {
let metadata = ResponseMetadata {
request_id,
execution_time_ms: start.elapsed().as_millis(),
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
return Ok(Json(ApiResponse::error(
"STATEMENT_NOT_ALLOWED",
format!("batch exceeds {MAX_BATCH_ITEMS} statements"),
metadata,
)));
}
if requests.is_empty() {
let metadata = ResponseMetadata {
request_id,
execution_time_ms: start.elapsed().as_millis(),
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
return Ok(Json(ApiResponse::error(
"STATEMENT_NOT_ALLOWED",
"empty batch".to_string(),
metadata,
)));
}
let mut results = Vec::new();
let mut total_rows = 0u64;
for request in requests {
match query_handler(
State(state.clone()),
Extension(account.clone()),
headers.clone(),
Json(request),
)
.await
{
Ok(Json(response)) => {
if let Some(data) = response.data {
total_rows += data.rows_affected.unwrap_or(data.rows.len() as u64);
results.push(data);
}
}
Err(_) => {
}
}
}
let metadata = ResponseMetadata {
request_id,
execution_time_ms: start.elapsed().as_millis(),
rows_affected: Some(total_rows),
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
Ok(Json(ApiResponse::success(results, metadata)))
}
pub async fn transaction_handler(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Json(requests): Json<Vec<QueryRequest>>,
) -> Result<Json<ApiResponse<Vec<QueryResult>>>, StatusCode> {
let start = Instant::now();
let request_id = Uuid::new_v4().to_string();
let mkmeta = |rows_affected: Option<u64>| ResponseMetadata {
request_id: request_id.clone(),
execution_time_ms: start.elapsed().as_millis(),
rows_affected,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
if requests.is_empty() || requests.len() > MAX_BATCH_ITEMS {
return Ok(Json(ApiResponse::error(
"STATEMENT_NOT_ALLOWED",
format!("transaction must contain 1..={MAX_BATCH_ITEMS} statements"),
mkmeta(None),
)));
}
let database = requests[0].database.clone();
if requests.iter().any(|r| r.database != database) {
return Ok(Json(ApiResponse::error(
"STATEMENT_NOT_ALLOWED",
"all transaction statements must target the same database".to_string(),
mkmeta(None),
)));
}
let db_access = resolve_database_access(&state, &account, &database).await?;
for req in &requests {
if let Err(e) =
validate_query_permissions(&req.query, &db_access.permissions, &account.role)
{
let code = gate_error_code(&e);
return Ok(Json(ApiResponse::error(code, e, mkmeta(None))));
}
}
let pool_key = format!(
"{}-{}-{}",
account.instance_id, db_access.username, database
);
let pool = match get_or_create_pool(&state, &pool_key, &account, &db_access).await {
Ok(p) => p,
Err(e) => {
return Ok(Json(ApiResponse::error(
"CONNECTION_ERROR",
e,
mkmeta(None),
)));
}
};
let stmts: Vec<(String, Vec<serde_json::Value>, QueryOptions)> = requests
.into_iter()
.map(|r| (r.query, r.params, r.options))
.collect();
match execute_transaction_with_pool(pool, stmts).await {
Ok(results) => {
let total = results
.iter()
.map(|r| r.rows_affected.unwrap_or(r.rows.len() as u64))
.sum();
Ok(Json(ApiResponse::success(results, mkmeta(Some(total)))))
}
Err(e) => {
let code = if e.starts_with("STATEMENT_TIMEOUT") {
"STATEMENT_TIMEOUT"
} else {
"QUERY_ERROR"
};
Ok(Json(ApiResponse::error(code, e, mkmeta(None))))
}
}
}
const TRANSACTION_SESSION_TTL_SECS: u64 = 900;
#[derive(Debug, Deserialize)]
pub struct TransactionSessionBeginRequest {
pub database: String,
#[serde(default)]
pub read_only: bool,
#[serde(default)]
pub isolation_level: Option<TransactionIsolationLevel>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum TransactionIsolationLevel {
ReadCommitted,
RepeatableRead,
Serializable,
}
impl TransactionIsolationLevel {
fn begin_statement(&self, read_only: bool) -> String {
let isolation = match self {
Self::ReadCommitted => "READ COMMITTED",
Self::RepeatableRead => "REPEATABLE READ",
Self::Serializable => "SERIALIZABLE",
};
format!(
"BEGIN TRANSACTION ISOLATION LEVEL {isolation}{}",
if read_only { " READ ONLY" } else { "" }
)
}
}
fn cleanup_transaction_sessions(state: &AppState) {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
state.transaction_sessions.retain(|_, session| {
now.saturating_sub(session.last_used.load(Ordering::Relaxed))
<= TRANSACTION_SESSION_TTL_SECS
});
}
fn session_metadata(account: &Account, session_id: Uuid) -> ResponseMetadata {
ResponseMetadata {
request_id: session_id.to_string(),
execution_time_ms: 0,
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
}
}
pub async fn begin_transaction_session(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Json(request): Json<TransactionSessionBeginRequest>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
cleanup_transaction_sessions(&state);
let db_access = resolve_database_access(&state, &account, &request.database).await?;
let pool_key = format!(
"{}-{}-{}",
account.instance_id, db_access.username, request.database
);
let pool = get_or_create_pool(&state, &pool_key, &account, &db_access)
.await
.map_err(|_| StatusCode::SERVICE_UNAVAILABLE)?;
let client = pool
.get()
.await
.map_err(|_| StatusCode::SERVICE_UNAVAILABLE)?;
let begin_statement = request.isolation_level.as_ref().map_or_else(
|| {
if request.read_only {
"BEGIN TRANSACTION READ ONLY".to_owned()
} else {
"BEGIN".to_owned()
}
},
|level| level.begin_statement(request.read_only),
);
client
.batch_execute(&begin_statement)
.await
.map_err(|_| StatusCode::SERVICE_UNAVAILABLE)?;
let session_id = Uuid::new_v4();
let session = std::sync::Arc::new(TransactionSession {
account_id: account.id.clone(),
database: request.database.clone(),
client: std::sync::Arc::new(tokio::sync::Mutex::new(client)),
last_used: std::sync::atomic::AtomicU64::new(0),
});
session.touch();
state.transaction_sessions.insert(session_id, session);
Ok(Json(ApiResponse::success(
json!({"session_id": session_id, "database": request.database, "read_only": request.read_only}),
session_metadata(&account, session_id),
)))
}
pub async fn query_transaction_session(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Path(session_id): Path<Uuid>,
Json(request): Json<QueryRequest>,
) -> Result<Json<ApiResponse<QueryResult>>, StatusCode> {
let Some(session) = state
.transaction_sessions
.get(&session_id)
.map(|entry| entry.clone())
else {
return Err(StatusCode::NOT_FOUND);
};
if session.account_id != account.id || session.database != request.database {
return Err(StatusCode::FORBIDDEN);
}
let db_access = resolve_database_access(&state, &account, &request.database).await?;
if let Err(error) =
validate_query_permissions(&request.query, &db_access.permissions, &account.role)
{
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
error,
session_metadata(&account, session_id),
)));
}
let client = session.client.lock().await;
let result =
execute_query_in_transaction(&client, request.query, request.params, &request.options)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
session.touch();
let mut metadata = session_metadata(&account, session_id);
metadata.rows_affected = result.rows_affected.or(Some(result.rows.len() as u64));
Ok(Json(ApiResponse::success(result, metadata)))
}
pub async fn commit_transaction_session(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Path(session_id): Path<Uuid>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
let Some(session) = state
.transaction_sessions
.get(&session_id)
.map(|entry| entry.clone())
else {
return Err(StatusCode::NOT_FOUND);
};
if session.account_id != account.id {
return Err(StatusCode::FORBIDDEN);
}
let Some((_, session)) = state.transaction_sessions.remove(&session_id) else {
return Err(StatusCode::NOT_FOUND);
};
let client = session.client.lock().await;
client
.batch_execute("COMMIT")
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
Ok(Json(ApiResponse::success(
json!({"committed": true, "session_id": session_id}),
session_metadata(&account, session_id),
)))
}
pub async fn rollback_transaction_session(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Path(session_id): Path<Uuid>,
) -> Result<StatusCode, StatusCode> {
let Some(session) = state
.transaction_sessions
.get(&session_id)
.map(|entry| entry.clone())
else {
return Err(StatusCode::NOT_FOUND);
};
if session.account_id != account.id {
return Err(StatusCode::FORBIDDEN);
}
let Some((_, session)) = state.transaction_sessions.remove(&session_id) else {
return Err(StatusCode::NOT_FOUND);
};
let client = session.client.lock().await;
client
.batch_execute("ROLLBACK")
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
Ok(StatusCode::NO_CONTENT)
}
pub async fn list_databases(
State(state): State<AppState>,
Extension(account): Extension<Account>,
) -> Result<Json<ApiResponse<Vec<String>>>, StatusCode> {
let request_id = Uuid::new_v4().to_string();
let mut databases: Vec<String> = account
.databases
.iter()
.map(|db| db.database.clone())
.collect();
for database in discover_owned_disposable_databases(&state, &account).await? {
if !databases.contains(&database) {
databases.push(database);
}
}
let metadata = ResponseMetadata {
request_id,
execution_time_ms: 0,
rows_affected: Some(databases.len() as u64),
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
Ok(Json(ApiResponse::success(databases, metadata)))
}
pub async fn create_database(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Json(request): Json<DatabaseManagementRequest>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
let request_id = Uuid::new_v4().to_string();
let metadata = ResponseMetadata {
request_id,
execution_time_ms: 0,
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
if !account.has_permission(Permission::CreateDatabase) {
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
"CREATE_DATABASE permission is required".to_string(),
metadata,
)));
}
if !is_test_database_name(&request.name) || !account.owns_database_role(&request.owner) {
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
"database name or owner is not allowed for this account".to_string(),
metadata,
)));
}
let Some(postgres_access) = account
.databases
.iter()
.find(|database| database.database == "postgres" && database.username == request.owner)
else {
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
"owner must have a postgres database access entry".to_string(),
metadata,
)));
};
let pool = match get_or_create_pool(
&state,
&format!(
"{}-{}-postgres",
account.instance_id, postgres_access.username
),
&account,
postgres_access,
)
.await
{
Ok(pool) => pool,
Err(error) => {
return Ok(Json(ApiResponse::error(
"CONNECTION_ERROR",
format!("database administration unavailable: {error}"),
metadata,
)));
}
};
let client = match pool.get().await {
Ok(client) => client,
Err(error) => {
return Ok(Json(ApiResponse::error(
"CONNECTION_ERROR",
format!("database administration unavailable: {error}"),
metadata,
)));
}
};
let statement = format!(
"CREATE DATABASE {} OWNER {}",
quote_identifier(&request.name),
quote_identifier(&request.owner)
);
match client.batch_execute(&statement).await {
Ok(()) => {
tracing::info!(account_id = %account.id, database = %request.name, owner = %request.owner, "test database created");
Ok(Json(ApiResponse::success(
json!({"name": request.name, "owner": request.owner}),
metadata,
)))
}
Err(error) => Ok(Json(ApiResponse::error(
"DATABASE_ERROR",
format!("database creation failed: {error}"),
metadata,
))),
}
}
pub async fn drop_database(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Path(name): Path<String>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
let request_id = Uuid::new_v4().to_string();
let metadata = ResponseMetadata {
request_id,
execution_time_ms: 0,
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
if !account.has_permission(Permission::DropOwnedDatabase) {
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
"DROP_OWNED_DATABASE permission is required".to_string(),
metadata,
)));
}
if !is_test_database_name(&name) {
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
"only disposable/validation databases may be dropped".to_string(),
metadata,
)));
}
let Some(postgres_access) = account
.databases
.iter()
.find(|database| database.database == "postgres")
else {
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
"account has no postgres database access entry".to_string(),
metadata,
)));
};
let pool = match get_or_create_pool(
&state,
&format!(
"{}-{}-postgres",
account.instance_id, postgres_access.username
),
&account,
postgres_access,
)
.await
{
Ok(pool) => pool,
Err(error) => {
return Ok(Json(ApiResponse::error(
"CONNECTION_ERROR",
format!("database administration unavailable: {error}"),
metadata,
)));
}
};
let client = match pool.get().await {
Ok(client) => client,
Err(error) => {
return Ok(Json(ApiResponse::error(
"CONNECTION_ERROR",
format!("database administration unavailable: {error}"),
metadata,
)));
}
};
let row = match client.query_opt(
"SELECT pg_catalog.pg_get_userbyid(datdba) FROM pg_catalog.pg_database WHERE datname = $1",
&[&name],
).await {
Ok(row) => row,
Err(error) => return Ok(Json(ApiResponse::error(
"DATABASE_ERROR",
format!("database ownership lookup failed: {error}"), metadata,
)))
};
let Some(row) = row else {
return Ok(Json(ApiResponse::error(
"NOT_FOUND",
"database not found".to_string(),
metadata,
)));
};
let owner: String = row.get(0);
if owner != postgres_access.username {
return Ok(Json(ApiResponse::error(
"PERMISSION_DENIED",
"database is not owned by the account postgres role".to_string(),
metadata,
)));
}
match client
.batch_execute(&format!(
"DROP DATABASE {} WITH (FORCE)",
quote_identifier(&name)
))
.await
{
Ok(()) => {
tracing::info!(account_id = %account.id, database = %name, owner = %owner, "test database dropped");
Ok(Json(ApiResponse::success(json!({"name": name}), metadata)))
}
Err(error) => Ok(Json(ApiResponse::error(
"DATABASE_ERROR",
format!("database deletion failed: {error}"),
metadata,
))),
}
}
#[derive(Debug, Deserialize)]
pub struct DatabaseManagementRequest {
pub name: String,
pub owner: String,
}
fn is_test_database_name(name: &str) -> bool {
let valid_chars = !name.is_empty()
&& name.len() <= 63
&& name
.bytes()
.all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_');
let reserved = matches!(name, "postgres" | "index" | "index_monitoramento");
valid_chars
&& !reserved
&& (name.ends_with("_disposable")
|| name.starts_with("identity_cutover_validation_")
|| name.starts_with("commercial_party_rbac_validation_"))
}
async fn resolve_database_access(
state: &AppState,
account: &Account,
database: &str,
) -> Result<DatabaseAccess, StatusCode> {
if let Some(access) = account.databases.iter().find(|db| db.database == database) {
return Ok(access.clone());
}
if !is_test_database_name(database) {
return Err(StatusCode::FORBIDDEN);
}
let postgres_access = account
.databases
.iter()
.find(|db| db.database == "postgres")
.ok_or(StatusCode::FORBIDDEN)?;
let pool = get_or_create_pool(
state,
&format!(
"{}-{}-postgres",
account.instance_id, postgres_access.username
),
account,
postgres_access,
)
.await
.map_err(|_| StatusCode::FORBIDDEN)?;
let client = pool.get().await.map_err(|_| StatusCode::FORBIDDEN)?;
let owner: String = client
.query_opt(
"SELECT pg_get_userbyid(datdba) FROM pg_database WHERE datname = $1",
&[&database],
)
.await
.map_err(|_| StatusCode::FORBIDDEN)?
.and_then(|row| row.try_get(0).ok())
.ok_or(StatusCode::FORBIDDEN)?;
let mut access = account
.databases
.iter()
.find(|db| db.username == owner)
.cloned()
.ok_or(StatusCode::FORBIDDEN)?;
access.database = database.to_string();
Ok(access)
}
async fn discover_owned_disposable_databases(
state: &AppState,
account: &Account,
) -> Result<Vec<String>, StatusCode> {
let postgres_access = account
.databases
.iter()
.find(|db| db.database == "postgres")
.ok_or(StatusCode::FORBIDDEN)?;
let pool = get_or_create_pool(
state,
&format!(
"{}-{}-postgres",
account.instance_id, postgres_access.username
),
account,
postgres_access,
)
.await
.map_err(|_| StatusCode::FORBIDDEN)?;
let client = pool.get().await.map_err(|_| StatusCode::FORBIDDEN)?;
let owners: Vec<String> = account
.databases
.iter()
.map(|db| db.username.clone())
.collect();
let rows = client
.query(
"SELECT datname FROM pg_database WHERE pg_get_userbyid(datdba) = ANY($1::text[])",
&[&owners],
)
.await
.map_err(|_| StatusCode::FORBIDDEN)?;
Ok(rows
.into_iter()
.filter_map(|row| row.try_get::<_, String>(0).ok())
.filter(|name| is_test_database_name(name))
.collect())
}
fn quote_identifier(identifier: &str) -> String {
format!("\"{}\"", identifier.replace('"', "\"\""))
}
pub async fn list_tables(
State(state): State<AppState>,
Extension(account): Extension<Account>,
Path(db): Path<String>,
) -> Result<Json<ApiResponse<Vec<Value>>>, StatusCode> {
let start = Instant::now();
let request_id = Uuid::new_v4().to_string();
if resolve_database_access(&state, &account, &db)
.await
.is_err()
{
let metadata = ResponseMetadata {
request_id,
execution_time_ms: start.elapsed().as_millis(),
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
return Ok(Json(ApiResponse::error(
"FORBIDDEN",
format!("No access to database: {db}"),
metadata,
)));
}
let query_req = QueryRequest {
query: "SELECT schemaname as schema, tablename as name, 'table' as type, tableowner as owner FROM pg_tables WHERE schemaname NOT IN ('pg_catalog', 'information_schema') ORDER BY schemaname, tablename".to_string(),
database: db,
params: vec![],
options: QueryOptions::default(),
};
match query_handler(
State(state),
Extension(account.clone()),
HeaderMap::new(),
Json(query_req),
)
.await
{
Ok(Json(response)) => {
if response.success {
let data = response.data.unwrap();
let metadata = ResponseMetadata {
request_id,
execution_time_ms: start.elapsed().as_millis(),
rows_affected: Some(data.rows.len() as u64),
instance_id: Some(account.instance_id),
timestamp: Utc::now(),
};
Ok(Json(ApiResponse::success(data.rows, metadata)))
} else {
Ok(Json(ApiResponse {
success: false,
data: None,
error: response.error,
metadata: response.metadata,
}))
}
}
Err(e) => Err(e),
}
}
pub async fn get_schema(
State(_state): State<AppState>,
Extension(account): Extension<Account>,
Path(_db): Path<String>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
let request_id = Uuid::new_v4().to_string();
let metadata = ResponseMetadata {
request_id,
execution_time_ms: 0,
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
Ok(Json(ApiResponse::error(
"NOT_IMPLEMENTED",
"Schema introspection is not implemented in this version".to_string(),
metadata,
)))
}
pub async fn get_account_info(
State(state): State<AppState>,
Extension(account): Extension<Account>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
let metadata = ResponseMetadata {
request_id: Uuid::new_v4().to_string(),
execution_time_ms: 0,
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
let active_connections = state
.connection_limits
.get(&account.id)
.map(|semaphore| account.max_connections as usize - semaphore.available_permits())
.unwrap_or(0);
let info = json!({
"id": account.id,
"name": account.name,
"role": account.role,
"databases": account.databases.iter().map(|db| db.database.clone()).collect::<Vec<_>>(),
"rate_limit": account.rate_limit,
"max_connections": account.max_connections,
"active_connections": active_connections,
});
Ok(Json(ApiResponse::success(info, metadata)))
}
pub async fn get_usage_stats(
Extension(account): Extension<Account>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
let metadata = ResponseMetadata {
request_id: Uuid::new_v4().to_string(),
execution_time_ms: 0,
rows_affected: None,
instance_id: Some(account.instance_id.clone()),
timestamp: Utc::now(),
};
Ok(Json(ApiResponse::error(
"NOT_IMPLEMENTED",
"Usage statistics are not implemented in this version".to_string(),
metadata,
)))
}
pub async fn serve_openapi() -> impl IntoResponse {
let openapi_content = include_str!("../openapi.json");
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "application/json")
.body(openapi_content.to_string())
.unwrap()
}
pub async fn serve_docs() -> impl IntoResponse {
Html(
r#"<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>pg-api Documentation</title>
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9/swagger-ui.css">
<style>
body {
margin: 0;
padding: 0;
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
}
.header {
background-color: #1a1a1a;
color: white;
padding: 1rem 2rem;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.header h1 {
margin: 0;
font-size: 1.5rem;
font-weight: 500;
}
.header p {
margin: 0.5rem 0 0 0;
opacity: 0.8;
font-size: 0.9rem;
}
#swagger-ui {
margin-top: 0;
}
.swagger-ui .topbar {
display: none;
}
</style>
</head>
<body>
<div class="header">
<h1>pg-api Documentation</h1>
<p>PostgreSQL API Service - RESTful API for database operations</p>
</div>
<div id="swagger-ui"></div>
<script src="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9/swagger-ui-bundle.js"></script>
<script src="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9/swagger-ui-standalone-preset.js"></script>
<script>
window.onload = function() {
window.ui = SwaggerUIBundle({
url: "/openapi.json",
dom_id: '#swagger-ui',
deepLinking: true,
presets: [
SwaggerUIBundle.presets.apis,
SwaggerUIStandalonePreset
],
plugins: [
SwaggerUIBundle.plugins.DownloadUrl
],
layout: "StandaloneLayout",
persistAuthorization: true,
tryItOutEnabled: true,
supportedSubmitMethods: ['get', 'post', 'put', 'delete', 'patch'],
onComplete: function() {
console.log("Swagger UI loaded");
}
});
};
</script>
</body>
</html>"#,
)
}
#[cfg(test)]
mod tests {
use super::{TransactionIsolationLevel, TransactionSessionBeginRequest};
#[test]
fn builds_safe_begin_statements_for_supported_isolation_levels() {
assert_eq!(
TransactionIsolationLevel::RepeatableRead.begin_statement(true),
"BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY"
);
assert_eq!(
TransactionIsolationLevel::Serializable.begin_statement(false),
"BEGIN TRANSACTION ISOLATION LEVEL SERIALIZABLE"
);
}
#[test]
fn isolation_level_is_optional_for_backward_compatible_requests() {
let request: TransactionSessionBeginRequest = serde_json::from_value(
serde_json::json!({"database": "index_monitoramento", "read_only": true}),
)
.expect("legacy session request should remain valid");
assert!(request.isolation_level.is_none());
}
}