Research-Stack/4-Infrastructure/rds_probe/src/db.rs

use sqlx::{postgres::PgPoolOptions, PgPool};
use std::env;

pub struct DbConfig {
    pub host: String,
    pub port: u16,
    pub user: String,
    pub password: String,
    pub dbname: String,
}

impl DbConfig {
    pub fn from_env() -> Self {
        let host = env::var("RDS_HOST")
            .unwrap_or_else(|_| "localhost".to_string());
        let user = env::var("RDS_USER").unwrap_or_else(|_| "postgres".to_string());
        let password = env::var("RDS_PASSWORD").unwrap_or_else(|_| "".to_string());
        
        Self {
            host,
            port: env::var("RDS_PORT")
                .unwrap_or_else(|_| "5432".to_string())
                .parse()
                .unwrap_or(5432),
            user,
            password,
            dbname: env::var("RDS_DB").unwrap_or_else(|_| "postgres".to_string()),
        }
    }

    pub fn dsn(&self) -> String {
        format!(
            "postgres://{}:{}@{}:{}/{}?sslmode=prefer",
            self.user, 
            urlencoding(&self.password),
            self.host, 
            self.port, 
            self.dbname
        )
    }
}

fn urlencoding(s: &str) -> String {
    let mut encoded = String::new();
    for c in s.chars() {
        match c {
            'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' | '.' | '~' => {
                encoded.push(c);
            }
            _ => {
                encoded.push('%');
                encoded.push_str(&format!("{:02X}", c as u8));
            }
        }
    }
    encoded
}

pub async fn create_pool(config: &DbConfig) -> Result<PgPool, sqlx::Error> {
    PgPoolOptions::new()
        .max_connections(1)
        .acquire_timeout(std::time::Duration::from_secs(30))
        .connect(&config.dsn()).await
}

pub async fn get_pool() -> Result<PgPool, Box<dyn std::error::Error>> {
    let config = DbConfig::from_env();
    let pool = create_pool(&config).await?;
    Ok(pool)
}