more refactoring

master
Edward Shen 2020-05-01 22:41:29 -04:00
parent 4b3dfcbc68
commit ae069154a0
Signed by: edward
GPG Key ID: 19182661E818369F
3 changed files with 68 additions and 54 deletions

View File

@ -2,38 +2,15 @@ use crate::commands::GENERAL_GROUP;
use crate::passive::Handler; use crate::passive::Handler;
use serenity::client::Client; use serenity::client::Client;
use serenity::framework::standard::StandardFramework; use serenity::framework::standard::StandardFramework;
use serenity::prelude::TypeMapKey;
use sqlx::sqlite::SqlitePool;
use std::env; use std::env;
use util::db::DbConnPool;
mod commands; mod commands;
mod passive; mod passive;
mod util;
pub(crate) const COMMAND_PREFIX: &str = "~"; pub(crate) const COMMAND_PREFIX: &str = "~";
pub(crate) struct DbConnPool {
pool: sqlx::pool::Pool<sqlx::sqlite::SqliteConnection>,
}
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] #[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> { async fn main() -> Result<(), Box<dyn std::error::Error>> {
dotenv::dotenv().ok(); dotenv::dotenv().ok();
@ -42,7 +19,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.configure(|c| c.prefix(COMMAND_PREFIX)) .configure(|c| c.prefix(COMMAND_PREFIX))
.group(&GENERAL_GROUP); .group(&GENERAL_GROUP);
// Login with a bot token from the environment
let mut client = let mut client =
Client::new_with_extras(&env::var("DISCORD_TOKEN").expect("token"), |extras| { Client::new_with_extras(&env::var("DISCORD_TOKEN").expect("token"), |extras| {
extras extras
@ -52,15 +28,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.await .await
.expect("Error creating client"); .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; let mut data = client.data.write().await;
data.insert::<DbConnPool>(DbConnPool { pool }); data.insert::<DbConnPool>(DbConnPool::new().await);
} }
// start listening for events by starting a single shard // start listening for events by starting a single shard
@ -70,24 +40,3 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
Ok(()) Ok(())
} }
async fn init_pool(pool: &sqlx::pool::Pool<sqlx::sqlite::SqliteConnection>) {
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();
}
}

64
src/util/db.rs Normal file
View File

@ -0,0 +1,64 @@
use serenity::prelude::TypeMapKey;
use sqlx::{
sqlite::{SqliteConnection, SqlitePool},
Pool,
};
use std::env;
type DbPool = Pool<SqliteConnection>;
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
}

1
src/util/mod.rs Normal file
View File

@ -0,0 +1 @@
pub mod db;