use log::debug; use serenity::prelude::TypeMapKey; use sqlx::{ sqlite::{SqliteConnection, SqlitePool}, Error, Pool, }; use std::env; type DbPool = Pool; pub(crate) struct DbConnPool { pub pool: DbPool, } impl DbConnPool { pub async fn new() -> Result { Ok(Self { pool: init_pool().await?, }) } pub async fn get_heck(&self) -> Result { 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 { 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) }