Rust BB8

pg.rs

use std::net::IpAddr;
use std::sync::OnceLock;

use anyhow::{Context, Result};
use bb8_postgres::PostgresConnectionManager;
use bb8_postgres::bb8::{Pool, PooledConnection};
use tokio_postgres::types::ToSql;
use tokio_postgres::{NoTls, Row, ToStatement};
use postgres_native_tls::MakeTlsConnector;
use native_tls::TlsConnector;
use uuid::Uuid;
use short_uuid::ShortUuid;

static DB_MANAGER_INSTANCE: OnceLock<DBManager> = OnceLock::new();

pub type DBConnection<'a> = PooledConnection<'a, PostgresConnectionManager<MakeTlsConnector>>;
pub type DBPool = Pool<PostgresConnectionManager<MakeTlsConnector>>;

pub struct DBOptions {
    pub is_dev: bool,
    pub pg_params: String,
    pub pool_max_size: u32,
}

pub struct DBManager {
    pool: DBPool,
}

impl DBManager {
    pub async fn get() -> &'static DBManager {
        DB_MANAGER_INSTANCE.get().unwrap()
    }

    async fn new(config: DBOptions) -> Result<Self, anyhow::Error> {
        let DBOptions {
            is_dev,
            pg_params,
            pool_max_size,
        } = config;

        let manager = PostgresConnectionManager::new_from_stringlike(
            pg_params,
            MakeTlsConnector::new(TlsConnector::builder().build().unwrap())
        ).expect("Unable to build PostgresConnectionManager");

        let pool = Pool::builder()
            .max_size(pool_max_size)
            .build(manager)
            .await
            .context("Postgres error")?;

        Ok(Self { pool })
    }

    pub async fn connection(&self) -> Result<DBConnection<'_>, anyhow::Error> {
        let conn = self.pool.get().await.context("Connection error")?;
        Ok(conn)
    }

    pub async fn query<T>(
        &self,
        statement: &T,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<Vec<Row>, anyhow::Error>
    where
        T: ?Sized + ToStatement,
    {
        let conn = self.connection().await?;
        let rows = conn
            .query(statement, params)
            .await
            .context("Postgres error")?;
        Ok(rows)
    }

    pub async fn query_one<T>(
        &self,
        statement: &T,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<Row, anyhow::Error>
    where
        T: ?Sized + ToStatement,
    {
        let conn = self.connection().await?;
        let row = conn
            .query_one(statement, params)
            .await
            .context("Postgres error")?;
        Ok(row)
    }
}

#[derive(Debug)]
pub struct ApiKey {
    pub iter: i32,
    pub salt: String,
    pub hash: String,
}

pub async fn get_api_key(id: &Uuid) -> Result<Option<ApiKey>, anyhow::Error> {
    let sql = "
        select
          key_iter,
          key_salt,
          key_hash
        from api_keys
        where api_key_id = $1
              and is_active = true
              and is_deleted = false
              and revoked_at is null
              and (expires_at is null or expires_at > now())";
    let rows = DBManager::get().await.query(sql, &[id]).await?;
    if rows.is_empty() {
        Ok(None)
    } else {
        Ok(Some(ApiKey {
            iter: rows[0].get("key_iter"),
            salt: rows[0].get("key_salt"),
            hash: rows[0].get("key_hash"),
        }))
    }
}

fn hostname_to_vm_id(hostname: &str) -> String {
    if hostname.len() != 33 { //len of short UUID + .vm.vers.sh
        String::from("")
    } else {
        let short_id = ShortUuid::parse_str(&hostname[..22]);
        match short_id {
           Ok(x) => x.to_uuid().to_string(),
           _ => String::from("")
        }
    }
}

#[derive(Debug)]
pub struct Vm {
    pub ip: IpAddr,
    pub wg_public_key: String,
}

pub async fn get_vm(hostname: &str) -> Option<Vm> {
    let sql = if hostname.ends_with(".vm.vers.sh") {
        "select
           ip,
           wg_public_key
         from vms
         where vm_id = $1
         limit 1"
    } else {
        "select
           vms.ip,
           vms.wg_public_key
         from domains
         left join vms on domains.vm_id = vms.vm_id
         where domains.domain = $1
         limit 1"
    };
    let domain = if hostname.ends_with(".vm.vers.sh") {
        hostname_to_vm_id(hostname)
    } else {
        hostname.to_string()
    };
    let row = DBManager::get().await.query_one(sql, &[&domain]).await.unwrap();
    if row.is_empty() {
        None
    } else {
        Some(Vm {
            ip: row.get("ip"),
            wg_public_key: row.get("wg_public_key"),
        })
    }
}


async fn check_conn() -> Result<Option<bool>, anyhow::Error> {
    let sql = "select 1 = 1";
    let row = DBManager::get().await.query_one(sql, &[]).await.unwrap();
    Ok(Some(row.is_empty()))
}

pub async fn init(is_dev: bool, connection_string: String) -> Result<(), anyhow::Error> {
    let options = DBOptions {
        is_dev,
        pg_params: connection_string,
        pool_max_size: 8u32,
    };

    let _ = DB_MANAGER_INSTANCE.set(DBManager::new(options).await?);

    //todo finish check_conn and remove println!
    println!("About to test connection");
    check_conn().await?;

    Ok(())
}