diff --git a/src/main.rs b/src/main.rs index 59fbf09..002716f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,38 +2,15 @@ use crate::commands::GENERAL_GROUP; use crate::passive::Handler; use serenity::client::Client; use serenity::framework::standard::StandardFramework; -use serenity::prelude::TypeMapKey; -use sqlx::sqlite::SqlitePool; use std::env; +use util::db::DbConnPool; mod commands; mod passive; +mod util; pub(crate) const COMMAND_PREFIX: &str = "~"; -pub(crate) struct DbConnPool { - pool: sqlx::pool::Pool, -} - -impl DbConnPool { - async fn get_heck(&self) -> i32 { - sqlx::query!("UPDATE Heck SET count = count + 1") - .execute(&self.pool) - .await - .unwrap(); - - sqlx::query!("SELECT count FROM Heck") - .fetch_one(&self.pool) - .await - .unwrap() - .count - } -} - -impl TypeMapKey for DbConnPool { - type Value = Self; -} - #[tokio::main] async fn main() -> Result<(), Box> { dotenv::dotenv().ok(); @@ -42,7 +19,6 @@ async fn main() -> Result<(), Box> { .configure(|c| c.prefix(COMMAND_PREFIX)) .group(&GENERAL_GROUP); - // Login with a bot token from the environment let mut client = Client::new_with_extras(&env::var("DISCORD_TOKEN").expect("token"), |extras| { extras @@ -52,15 +28,9 @@ async fn main() -> Result<(), Box> { .await .expect("Error creating client"); - let pool = SqlitePool::builder() - .build(&env::var("DATABASE_URL").unwrap()) - .await?; - - init_pool(&pool).await; - { let mut data = client.data.write().await; - data.insert::(DbConnPool { pool }); + data.insert::(DbConnPool::new().await); } // start listening for events by starting a single shard @@ -70,24 +40,3 @@ async fn main() -> Result<(), Box> { Ok(()) } - -async fn init_pool(pool: &sqlx::pool::Pool) { - sqlx::query!( - "CREATE TABLE IF NOT EXISTS Heck (id INTEGER PRIMARY KEY NOT NULL, count INTEGER NOT NULL)" - ) - .execute(pool) - .await - .unwrap(); - - if sqlx::query!("SELECT count FROM Heck") - .fetch_all(pool) - .await - .unwrap() - .is_empty() - { - sqlx::query!("INSERT INTO Heck VALUES (1, 0)") - .execute(pool) - .await - .unwrap(); - } -} diff --git a/src/util/db.rs b/src/util/db.rs new file mode 100644 index 0000000..831ec78 --- /dev/null +++ b/src/util/db.rs @@ -0,0 +1,64 @@ +use serenity::prelude::TypeMapKey; +use sqlx::{ + sqlite::{SqliteConnection, SqlitePool}, + Pool, +}; +use std::env; + +type DbPool = Pool; + +pub(crate) struct DbConnPool { + pool: DbPool, +} + +impl DbConnPool { + pub async fn new() -> Self { + Self { + pool: init_pool().await, + } + } + pub async fn get_heck(&self) -> i32 { + sqlx::query!("UPDATE Heck SET count = count + 1") + .execute(&self.pool) + .await + .unwrap(); + + sqlx::query!("SELECT count FROM Heck") + .fetch_one(&self.pool) + .await + .unwrap() + .count + } +} + +impl TypeMapKey for DbConnPool { + type Value = Self; +} + +async fn init_pool() -> DbPool { + let pool = SqlitePool::builder() + .build(&env::var("DATABASE_URL").unwrap()) + .await + .unwrap(); + + sqlx::query!( + "CREATE TABLE IF NOT EXISTS Heck (id INTEGER PRIMARY KEY NOT NULL, count INTEGER NOT NULL)" + ) + .execute(&pool) + .await + .unwrap(); + + if sqlx::query!("SELECT count FROM Heck") + .fetch_all(&pool) + .await + .unwrap() + .is_empty() + { + sqlx::query!("INSERT INTO Heck VALUES (1, 0)") + .execute(&pool) + .await + .unwrap(); + } + + pool +} diff --git a/src/util/mod.rs b/src/util/mod.rs new file mode 100644 index 0000000..dec1023 --- /dev/null +++ b/src/util/mod.rs @@ -0,0 +1 @@ +pub mod db;