93 lines
2.0 KiB
Rust
93 lines
2.0 KiB
Rust
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)
|
|
}
|