datatrash/src/db.rs

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

pub async fn setup_db() -> PgPool {
    let conn_url = &get_db_url();
    log::info!("Using Connection string {}", conn_url);

    let pool = PgPoolOptions::new()
        .max_connections(5)
        .acquire_timeout(5.std_seconds())
        .connect(conn_url)
        .await
        .expect("could not create db pool");

    for query in include_str!("../init-db.sql").split_inclusive(';') {
        sqlx::query(query)
            .execute(&pool)
            .await
            .expect("could not initialize database schema");
    }

    pool
}

fn get_db_url() -> String {
    if let Ok(database_url) = env::var("DATABASE_URL") {
        return database_url;
    }

    let auth = if let Ok(user) = env::var("DATABASE_USER") {
        if let Ok(pass) = env::var("DATABASE_PASS") {
            format!("{}:{}@", user, pass)
        } else {
            format!("{}@", user)
        }
    } else {
        String::new()
    };

    format!(
        "postgresql://{auth}{host}/{name}",
        auth = auth,
        host = env::var("DATABASE_HOST").unwrap_or_else(|_| "localhost".to_string()),
        name = env::var("DATABASE_NAME").unwrap_or_else(|_| "datatrash".to_string())
    )
}