discord-kurante/src/util/db.rs

use log::debug;
use serenity::prelude::TypeMapKey;
use sqlx::{
    sqlite::{SqliteConnection, SqlitePool},
    Error, Pool,
};
use std::env;

type DbPool = Pool<SqliteConnection>;

pub(crate) struct DbConnPool {
    pub pool: DbPool,
}

impl DbConnPool {
    pub async fn new() -> Result<Self, Error> {
        Ok(Self {
            pool: init_pool().await?,
        })
    }

    pub async fn get_heck(&self) -> Result<i32, Error> {
        sqlx::query!("UPDATE Heck SET count = count + 1")
            .execute(&self.pool)
            .await?;

        Ok(sqlx::query!("SELECT count FROM Heck")
            .fetch_one(&self.pool)
            .await?
            .count)
    }
}

impl TypeMapKey for DbConnPool {
    type Value = Self;
}

async fn init_pool() -> Result<DbPool, Error> {
    let pool = SqlitePool::builder()
        .build(&env::var("DATABASE_URL").unwrap())
        .await?;

    // Heck table
    sqlx::query!(
        "CREATE TABLE IF NOT EXISTS Heck (
            id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL UNIQUE,
            count INTEGER NOT NULL
        )"
    )
    .execute(&pool)
    .await?;

    debug!("Table Heck exists or was created");

    // Arknights row roll counting
    sqlx::query!(
        "CREATE TABLE IF NOT EXISTS RollCount (
            user_id TEXT PRIMARY KEY UNIQUE NOT NULL,
            count INTEGER NOT NULL
        )"
    )
    .execute(&pool)
    .await?;

    debug!("Table RollCount exists or was created");

    sqlx::query!(
        "CREATE TABLE IF NOT EXISTS OperatorCount (
            user_id TEXT NOT NULL,
            operator TEXT NOT NULL,
            count INTEGER NOT NULL,
            UNIQUE(user_id, operator)
        )"
    )
    .execute(&pool)
    .await?;

    debug!("Table OperatorCount exists or was created");

    if sqlx::query!("SELECT count FROM Heck")
        .fetch_all(&pool)
        .await?
        .is_empty()
    {
        debug!("No entries in Heck, inserting default one");
        sqlx::query!("INSERT INTO Heck (count) VALUES (0)")
            .execute(&pool)
            .await?;
    }

    Ok(pool)
}