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 {
    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?;

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

    debug!("Table Heck 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 VALUES (1, 0)")
            .execute(&pool)
            .await?;
    }

    Ok(pool)
}