pg-instance-handler/src/actors/driven/sqlx_handler.rs

use crate::app::database_port::ForManagingDatabases;
use regex::Regex;
use sqlx::postgres::PgPoolOptions;
use tracing::{debug, info};

pub struct SQLXHandler {
    database_url: String,
}

impl SQLXHandler {
    pub fn new(database_url: String) -> Self {
        SQLXHandler { database_url }
    }

    fn make_runtime() -> Result<tokio::runtime::Runtime, String> {
        tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .map_err(|e| format!("failed to build runtime: {}", e))
    }

    fn sanitize_identifier(name: &str) -> Result<String, String> {
        // Restrict to simple PostgreSQL identifiers for safety.
        // Allows letters, numbers, and underscore; must start with a letter or underscore.
        let re = Regex::new(r"^[A-Za-z_][A-Za-z0-9_]*$").map_err(|e| e.to_string())?;
        if re.is_match(name) {
            Ok(name.to_string())
        } else {
            Err(format!("invalid database name: '{}'", name))
        }
    }
}

impl ForManagingDatabases for SQLXHandler {
    fn list_databases(&self) -> Result<Vec<String>, String> {
        info!("Listing PostgreSQL databases");
        let rt = Self::make_runtime()?;
        rt.block_on(async {
            let pool = PgPoolOptions::new()
                .max_connections(5)
                .connect(&self.database_url)
                .await
                .map_err(|e| format!("failed to connect: {}", e))?;

            let dbs: Vec<String> = sqlx::query_scalar(
                r#"SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname"#,
            )
            .fetch_all(&pool)
            .await
            .map_err(|e| format!("list databases failed: {}", e))?;

            debug!(?dbs, "Got databases");
            Ok(dbs)
        })
    }

    fn create_database(&self, name: &str) -> Result<(), String> {
        let ident = Self::sanitize_identifier(name)?;
        info!(database = %ident, "Creating database");
        let rt = Self::make_runtime()?;
        rt.block_on(async {
            let pool = PgPoolOptions::new()
                .max_connections(5)
                .connect(&self.database_url)
                .await
                .map_err(|e| format!("failed to connect: {}", e))?;

            let sql = format!("CREATE DATABASE {}", ident);
            sqlx::query(&sql)
                .execute(&pool)
                .await
                .map_err(|e| format!("create database failed: {}", e))?;
            Ok(())
        })
    }

    fn delete_database(&self, name: &str) -> Result<(), String> {
        let ident = Self::sanitize_identifier(name)?;
        info!(database = %ident, "Deleting database");
        let rt = Self::make_runtime()?;
        rt.block_on(async {
            let pool = PgPoolOptions::new()
                .max_connections(5)
                .connect(&self.database_url)
                .await
                .map_err(|e| format!("failed to connect: {}", e))?;

            // Terminate connections before dropping (best effort)
            let terminate_sql = r#"
                SELECT pg_terminate_backend(pid)
                FROM pg_stat_activity
                WHERE datname = $1 AND pid <> pg_backend_pid()
            "#;
            let _ = sqlx::query(terminate_sql).bind(&ident).execute(&pool).await;

            let sql = format!("DROP DATABASE IF EXISTS {}", ident);
            sqlx::query(&sql)
                .execute(&pool)
                .await
                .map_err(|e| format!("drop database failed: {}", e))?;
            Ok(())
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sanitize_identifier_allows_valid() {
        for ok in ["db", "_db", "db_123", "A_b1"] {
            assert!(
                SQLXHandler::sanitize_identifier(ok).is_ok(),
                "{} should be valid",
                ok
            );
        }
    }

    #[test]
    fn sanitize_identifier_rejects_invalid() {
        for bad in ["", "1db", "db-1", "db name", "db;drop", "db$", "db."] {
            assert!(
                SQLXHandler::sanitize_identifier(bad).is_err(),
                "{} should be invalid",
                bad
            );
        }
    }

    // Optional smoke test against a real Postgres instance.
    // Set TEST_ADMIN_DATABASE_URL to something like:
    //   postgres://postgres:postgres@localhost/postgres
    // Run with: cargo test -- --ignored
    #[test]
    #[ignore]
    fn postgres_smoke_create_list_delete() {
        let url = std::env::var("TEST_ADMIN_DATABASE_URL")
            .expect("set TEST_ADMIN_DATABASE_URL to an admin database URL");
        let handler = SQLXHandler::new(url);

        let db_name = format!("testdb_{}", std::process::id());

        // Ensure clean slate (ignore errors)
        let _ = handler.delete_database(&db_name);

        handler.create_database(&db_name).expect("create ok");
        let list = handler.list_databases().expect("list ok");
        assert!(list.contains(&db_name));

        handler.delete_database(&db_name).expect("drop ok");
        let list2 = handler.list_databases().expect("list ok");
        assert!(!list2.contains(&db_name));
    }
}