diff --git a/migrations/20240304155340_initial_migration.sql b/migrations/20240304155340_initial_migration.sql index ad70dbe..049aab9 100644 --- a/migrations/20240304155340_initial_migration.sql +++ b/migrations/20240304155340_initial_migration.sql @@ -1,10 +1,16 @@ -- Add migration script here CREATE TABLE user_stats ( - user_id BIGINT UNSIGNED NOT NULL, - guild_id BIGINT UNSIGNED NOT NULL, - experience_points INTEGER UNSIGNED NOT NULL, - level INTEGER UNSIGNED NOT NULL, - PRIMARY KEY (user_id, guild_id) + user_id BIGINT UNSIGNED NOT NULL, + guild_id BIGINT UNSIGNED NOT NULL, + experience_points INTEGER UNSIGNED NOT NULL, + level INTEGER UNSIGNED NOT NULL, + PRIMARY KEY (user_id, guild_id) ); CREATE INDEX idx_user_id ON user_stats (user_id); CREATE INDEX idx_guild_id ON user_stats (guild_id); + +CREATE TABLE bot_mentions ( + mentions BIGINT UNSIGNED NOT NULL +); +CREATE INDEX idx_mentions ON bot_mentions (mentions); +INSERT INTO bot_mentions (mentions) VALUES (0); diff --git a/src/commands/level_cmds/level.rs b/src/commands/level_cmds/level.rs index 450f0af..4e96620 100644 --- a/src/commands/level_cmds/level.rs +++ b/src/commands/level_cmds/level.rs @@ -47,10 +47,10 @@ pub async fn level( }; let level = level_xp_and_rank_row .1 - .get::(DATABASE_COLUMNS[&Level]); + .get::(LEVELS_TABLE[&Level]); let xp = level_xp_and_rank_row .1 - .get::(DATABASE_COLUMNS[&ExperiencePoints]); + .get::(LEVELS_TABLE[&ExperiencePoints]); let avatar = target_replied_user.face().replace(".webp", ".png"); let username = &target_replied_user.name; diff --git a/src/commands/level_cmds/mod.rs b/src/commands/level_cmds/mod.rs index dac923c..4f6e67e 100644 --- a/src/commands/level_cmds/mod.rs +++ b/src/commands/level_cmds/mod.rs @@ -1,11 +1,12 @@ /// dependencies for the commands. use crate::commands::cmd_utils::get_replied_user; -use crate::data::bot_data::{DATABASE_COLUMNS, DATABASE_FILENAME}; use crate::data::command_data::{Context, Error}; -use crate::data::database_interactions::{ - connect_to_db, fetch_top_nine_levels_in_guild, fetch_user_level_and_rank, +use crate::data::database::{DATABASE_FILENAME, LEVELS_TABLE}; +use crate::database::{ + connect_to_db, + level_system::{fetch_top_nine_levels_in_guild, fetch_user_level_and_rank}, }; -use crate::enums::schemas::DatabaseSchema::*; +use crate::enums::schemas::LevelsSchema::*; use ::serenity::futures::future::try_join_all; use poise::serenity_prelude as serenity; use rayon::prelude::*; diff --git a/src/commands/level_cmds/toplevels.rs b/src/commands/level_cmds/toplevels.rs index 9e99c23..79715b6 100644 --- a/src/commands/level_cmds/toplevels.rs +++ b/src/commands/level_cmds/toplevels.rs @@ -25,7 +25,7 @@ pub async fn toplevels(ctx: Context<'_>) -> Result<(), Error> { ctx.defer().await?; let user_ids: Vec = level_and_xp_rows .par_iter() - .map(|row| row.get::(DATABASE_COLUMNS[&UserId]) as u64) + .map(|row| row.get::(LEVELS_TABLE[&UserId]) as u64) .collect(); let users = try_join_all( user_ids @@ -38,8 +38,8 @@ pub async fn toplevels(ctx: Context<'_>) -> Result<(), Error> { for (counter, (row, user)) in level_and_xp_rows.iter().zip(users.iter()).enumerate() { let (level, xp) = ( - row.get::(DATABASE_COLUMNS[&Level]), - row.get::(DATABASE_COLUMNS[&ExperiencePoints]), + row.get::(LEVELS_TABLE[&Level]), + row.get::(LEVELS_TABLE[&ExperiencePoints]), ); fields.push(( diff --git a/src/commands/level_logic.rs b/src/commands/level_logic.rs index a5af994..0abf000 100644 --- a/src/commands/level_logic.rs +++ b/src/commands/level_logic.rs @@ -1,10 +1,10 @@ use std::collections::HashMap; -use crate::enums::schemas::DatabaseSchema; +use crate::enums::schemas::LevelsSchema; /// Set the leveling condition and return the updated level with reset xp if true. -pub async fn update_level(experience: i32, level: i32) -> HashMap { - use crate::enums::schemas::DatabaseSchema as DbSch; +pub async fn update_level(experience: i32, level: i32) -> HashMap { + use crate::enums::schemas::LevelsSchema as DbSch; let update_level = if experience >= level * 100 { level + 1 } else { diff --git a/src/data/bot_data.rs b/src/data/bot_data.rs index 886db2d..f6a7272 100644 --- a/src/data/bot_data.rs +++ b/src/data/bot_data.rs @@ -1,14 +1,8 @@ -use crate::{ - enums::schemas::DatabaseSchema, utils::string_manipulation::upper_lowercase_permutations, -}; -use std::collections::HashMap; +use crate::utils::string_manipulation::upper_lowercase_permutations; use lazy_static::lazy_static; use regex::Regex; -pub const DATABASE_FILENAME: &str = "database/bot_database.sqlite"; -pub const DATABASE_USERS: &str = "user_stats"; - pub const DEFAULT_XP: i64 = 0; pub const DEFAULT_LEVEL: i64 = 1; @@ -17,18 +11,7 @@ lazy_static! { pub(crate) static ref BOT_TOKEN: String = std::env::var("BOT_TOKEN").expect("Expected a token in the dotenv file."); pub(crate) static ref START_TIME: std::time::Instant = std::time::Instant::now(); - #[derive(Debug)] - pub(crate) static ref DATABASE_COLUMNS: HashMap = { - use crate::enums::schemas::DatabaseSchema as DbSch; - HashMap::from([ - (DbSch::UserId, "user_id"), - (DbSch::GuildId, "guild_id"), - (DbSch::ExperiencePoints, "experience_points"), - (DbSch::Level, "level"), - (DbSch::LastQueryTimestamp, "last_query_timestamp") - ]) - }; pub(crate) static ref XP_COOLDOWN_NUMBER_SECS: i64 = 60; pub(crate) static ref BOT_PREFIXES: Vec = { let mut temp = vec![]; diff --git a/src/data/command_data.rs b/src/data/command_data.rs index 9fcce5a..e2d2ad7 100644 --- a/src/data/command_data.rs +++ b/src/data/command_data.rs @@ -1,13 +1,11 @@ use poise::serenity_prelude as serenity; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use std::sync::atomic::AtomicU32; use std::sync::Arc; #[derive(Debug)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct Data { - pub hutao_mentions: AtomicU32, pub bot_user: Arc, pub bot_avatar: Arc, } diff --git a/src/data/database.rs b/src/data/database.rs new file mode 100644 index 0000000..551784c --- /dev/null +++ b/src/data/database.rs @@ -0,0 +1,31 @@ +use crate::enums::schemas::{LevelsSchema, MentionsSchema}; +use std::collections::HashMap; + +use lazy_static::lazy_static; + +pub const DATABASE_FILENAME: &str = "database/bot_database.sqlite"; +pub const DATABASE_USERS: &str = "user_stats"; +pub const MENTIONS_TABLE_NAME: &str = "bot_mentions"; + +lazy_static! { + #[derive(Debug)] + pub(crate) static ref LEVELS_TABLE: HashMap = { + use crate::enums::schemas::LevelsSchema as DbSch; + + HashMap::from([ + (DbSch::UserId, "user_id"), + (DbSch::GuildId, "guild_id"), + (DbSch::ExperiencePoints, "experience_points"), + (DbSch::Level, "level"), + (DbSch::LastQueryTimestamp, "last_query_timestamp") + ]) + }; + #[derive(Debug)] + pub(crate) static ref MENTIONS_TABLE: HashMap = { + use crate::enums::schemas::MentionsSchema as DbSch; + + HashMap::from([ + (DbSch::Mentions, "mentions"), + ]) + }; +} diff --git a/src/data/mod.rs b/src/data/mod.rs index c170c4c..2e60088 100644 --- a/src/data/mod.rs +++ b/src/data/mod.rs @@ -1,5 +1,5 @@ pub mod bot_data; pub mod command_data; -pub mod database_interactions; +pub mod database; pub mod embed_media; pub mod user_data; diff --git a/src/database/bot_mentions.rs b/src/database/bot_mentions.rs new file mode 100644 index 0000000..c2ed2b3 --- /dev/null +++ b/src/database/bot_mentions.rs @@ -0,0 +1,45 @@ +use crate::data::command_data::Error; +use crate::data::database::{MENTIONS_TABLE, MENTIONS_TABLE_NAME}; +use crate::enums::schemas::MentionsSchema; +use sqlx::sqlite::SqliteQueryResult; +use sqlx::Row; +use sqlx::SqlitePool; + +pub async fn fetch_mentions(db: &SqlitePool) -> Result { + let query = format!( + "SELECT `{}` FROM `{}`", + MENTIONS_TABLE[&MentionsSchema::Mentions], + MENTIONS_TABLE_NAME, + ); + let sql = sqlx::query(&query).fetch_optional(db).await?; + + let row = match sql { + Some(row) => row, + None => return Err(format!("Couldn't find a row to select (SQL: {})", query).into()), + }; + + let queried_mentions = row.get::(MENTIONS_TABLE[&MentionsSchema::Mentions]); + + Ok(queried_mentions) +} + +pub async fn update_mentions( + db: &SqlitePool, + updated_mentions: i64, +) -> Result { + let query = format!( + "UPDATE `{}` SET `{}` = ?", + MENTIONS_TABLE_NAME, + MENTIONS_TABLE[&MentionsSchema::Mentions] + ); + + sqlx::query(&query).bind(updated_mentions).execute(db).await +} + +pub async fn add_mentions(db: &SqlitePool, n: i64) -> Result { + let fetched_mentions = fetch_mentions(db).await?; + + update_mentions(db, fetched_mentions + n).await?; + + fetch_mentions(db).await +} diff --git a/src/data/database_interactions.rs b/src/database/level_system.rs similarity index 78% rename from src/data/database_interactions.rs rename to src/database/level_system.rs index de248c2..fcff9f2 100644 --- a/src/data/database_interactions.rs +++ b/src/database/level_system.rs @@ -1,33 +1,21 @@ use crate::commands::level_logic::update_level; -use crate::enums::schemas::DatabaseSchema::*; +use crate::enums::schemas::LevelsSchema::*; use poise::serenity_prelude as serenity; use serenity::User; use sqlx::sqlite::SqliteRow; use sqlx::Row; /// https://stackoverflow.com/questions/72763578/how-to-create-a-sqlite-database-with-rust-sqlx -use sqlx::{sqlite::SqliteConnectOptions, sqlite::SqlitePoolOptions, Error, SqlitePool}; -use std::{future::Future, path::Path}; +use sqlx::{Error, SqlitePool}; -use crate::data::bot_data::{DATABASE_COLUMNS, DEFAULT_LEVEL, DEFAULT_XP, XP_COOLDOWN_NUMBER_SECS}; -use crate::data::user_data::USER_COOLDOWNS; - -use super::bot_data::DATABASE_USERS; +use crate::data::bot_data::{DEFAULT_LEVEL, DEFAULT_XP, XP_COOLDOWN_NUMBER_SECS}; +use crate::data::database::{DATABASE_USERS, LEVELS_TABLE}; -/// Used to establish the database connection with its predetermined parameters. -pub async fn connect_to_db( - filename: impl AsRef, -) -> impl Future> { - SqlitePoolOptions::new().connect_with( - SqliteConnectOptions::new() - .filename(filename) - .create_if_missing(true), - ) -} +use crate::data::user_data::USER_COOLDOWNS; /// Adds a new database user with the schema from `crate::data:bot_data.rs`. /// That's the reason why the function isn't public. async fn add_user_if_not_exists( - db: SqlitePool, + db: &SqlitePool, user: &User, guild_id: serenity::GuildId, ) -> Result<(), Error> { @@ -39,10 +27,10 @@ async fn add_user_if_not_exists( "INSERT INTO `{}` (`{}`, `{}`, `{}`, `{}`) VALUES (?, ?, ?, ?)", DATABASE_USERS, - DATABASE_COLUMNS[&UserId], - DATABASE_COLUMNS[&GuildId], - DATABASE_COLUMNS[&ExperiencePoints], - DATABASE_COLUMNS[&Level], + LEVELS_TABLE[&UserId], + LEVELS_TABLE[&GuildId], + LEVELS_TABLE[&ExperiencePoints], + LEVELS_TABLE[&Level], ); sqlx::query(&query) @@ -50,7 +38,7 @@ async fn add_user_if_not_exists( .bind(guild_id.to_string()) .bind(DEFAULT_XP) .bind(DEFAULT_LEVEL) - .execute(&db) + .execute(db) .await?; Ok(()) @@ -66,14 +54,14 @@ pub async fn fetch_user_level( "SELECT `{}`, `{}`, `{}` FROM `{}` WHERE `{}` = ? AND `{}` = ?", - DATABASE_COLUMNS[&UserId], - DATABASE_COLUMNS[&ExperiencePoints], - DATABASE_COLUMNS[&Level], + LEVELS_TABLE[&UserId], + LEVELS_TABLE[&ExperiencePoints], + LEVELS_TABLE[&Level], // DATABASE_USERS, // - DATABASE_COLUMNS[&UserId], - DATABASE_COLUMNS[&GuildId] + LEVELS_TABLE[&UserId], + LEVELS_TABLE[&GuildId] ) .as_str(), ) @@ -126,19 +114,19 @@ pub async fn fetch_top_nine_levels_in_guild( WHERE `{}` = ? ORDER BY {} DESC, {} DESC LIMIT 9", - DATABASE_COLUMNS[&UserId], - DATABASE_COLUMNS[&UserId], + LEVELS_TABLE[&UserId], + LEVELS_TABLE[&UserId], // - DATABASE_COLUMNS[&ExperiencePoints], - DATABASE_COLUMNS[&ExperiencePoints], + LEVELS_TABLE[&ExperiencePoints], + LEVELS_TABLE[&ExperiencePoints], // - DATABASE_COLUMNS[&Level], - DATABASE_COLUMNS[&Level], + LEVELS_TABLE[&Level], + LEVELS_TABLE[&Level], // DATABASE_USERS, - DATABASE_COLUMNS[&GuildId], - DATABASE_COLUMNS[&Level], - DATABASE_COLUMNS[&ExperiencePoints], + LEVELS_TABLE[&GuildId], + LEVELS_TABLE[&Level], + LEVELS_TABLE[&ExperiencePoints], ) .as_str(), ) @@ -156,7 +144,7 @@ pub async fn fetch_top_nine_levels_in_guild( /// Additionally, we directly use the guild_id instead of the event as the /// parameter for add_user_if_not_exists() in order to save computing resources. pub async fn add_or_update_db_user( - db: SqlitePool, + db: &SqlitePool, message: &serenity::Message, ctx: &serenity::Context, obtained_xp: i32, @@ -203,7 +191,7 @@ pub async fn add_or_update_db_user( } // First we need to check if there's some user_id+guild_id pair that matches - let level_query: Option = fetch_user_level(&db, user, guild_id).await?; + let level_query: Option = fetch_user_level(db, user, guild_id).await?; let query_row = match level_query { Some(row) => row, @@ -214,8 +202,8 @@ pub async fn add_or_update_db_user( } }; - let queried_level = query_row.get::(DATABASE_COLUMNS[&Level]); - let queried_experience_points = query_row.get::(DATABASE_COLUMNS[&ExperiencePoints]); + let queried_level = query_row.get::(LEVELS_TABLE[&Level]); + let queried_experience_points = query_row.get::(LEVELS_TABLE[&ExperiencePoints]); let added_experience_points = queried_experience_points + obtained_xp; let update = update_level(added_experience_points, queried_level).await; @@ -259,11 +247,11 @@ pub async fn add_or_update_db_user( WHERE `{}` = ? AND `{}` = ?", DATABASE_USERS, // - DATABASE_COLUMNS[&ExperiencePoints], - DATABASE_COLUMNS[&Level], + LEVELS_TABLE[&ExperiencePoints], + LEVELS_TABLE[&Level], // - DATABASE_COLUMNS[&UserId], - DATABASE_COLUMNS[&GuildId], + LEVELS_TABLE[&UserId], + LEVELS_TABLE[&GuildId], ); sqlx::query(&query) @@ -271,7 +259,7 @@ pub async fn add_or_update_db_user( .bind(updated_level) .bind(user.id.to_string()) .bind(guild_id.to_string()) - .execute(&db) + .execute(db) .await?; Ok(()) diff --git a/src/database/mod.rs b/src/database/mod.rs new file mode 100644 index 0000000..6c40357 --- /dev/null +++ b/src/database/mod.rs @@ -0,0 +1,17 @@ +use std::{future::Future, path::Path}; + +use sqlx::{sqlite::SqliteConnectOptions, sqlite::SqlitePoolOptions, SqlitePool}; + +pub mod bot_mentions; +pub mod level_system; + +/// Used to establish the database connection with its predetermined parameters. +pub async fn connect_to_db( + filename: impl AsRef, +) -> impl Future> { + SqlitePoolOptions::new().connect_with( + SqliteConnectOptions::new() + .filename(filename) + .create_if_missing(true), + ) +} diff --git a/src/enums/schemas.rs b/src/enums/schemas.rs index 4922ef6..3726b5e 100644 --- a/src/enums/schemas.rs +++ b/src/enums/schemas.rs @@ -3,10 +3,16 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Eq, PartialEq, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] -pub enum DatabaseSchema { +pub enum LevelsSchema { UserId, GuildId, Level, ExperiencePoints, LastQueryTimestamp, } + +#[derive(Debug, Eq, PartialEq, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum MentionsSchema { + Mentions, +} diff --git a/src/event_handler/handler.rs b/src/event_handler/handler.rs index 66ec671..5a52d15 100644 --- a/src/event_handler/handler.rs +++ b/src/event_handler/handler.rs @@ -9,10 +9,11 @@ use crate::data::bot_data::START_TIME; use crate::{ data::{ - bot_data::DATABASE_FILENAME, command_data::{Data, Error}, - database_interactions::*, + database::DATABASE_FILENAME, }, + database::connect_to_db, + database::level_system::*, utils::{replies::handle_replies, string_manipulation::remove_emojis_and_embeds_from_str}, }; use poise::serenity_prelude as serenity; @@ -22,7 +23,7 @@ pub async fn event_handler( ctx: &serenity::Context, event: &serenity::FullEvent, _framework: poise::FrameworkContext<'_, Data, Error>, - data: &Data, + _data: &Data, ) -> Result<(), Error> { match event { #[cfg(feature = "debug")] @@ -57,25 +58,30 @@ pub async fn event_handler( let text_patterns = ["hutao", "hu tao"]; let lowercase_msg = msg.to_lowercase(); let trimmed_emojis = remove_emojis_and_embeds_from_str(&lowercase_msg); - match text_patterns - .iter() - .any(|text| trimmed_emojis.contains(text)) - { - true => handle_replies(ctx, new_message, data, msg).await?, - false => { - #[cfg(feature = "debug")] - println!("Msg: {} has an emoji!", msg); - } - }; let db = connect_to_db(DATABASE_FILENAME.to_string()).await; match db.await { - Ok(pool) => { + Ok(ref pool) => { let obtained_xp: i32 = rand::thread_rng().gen_range(5..=15); #[cfg(feature = "debug")] println!("Connected to the database: {pool:?}"); + match text_patterns + .iter() + .any(|text| trimmed_emojis.contains(text)) + { + true => handle_replies(pool, ctx, new_message, &trimmed_emojis).await?, + false => { + #[cfg(feature = "debug")] + println!( + "Msg: {} has an emoji or doesn't contain: [{}]", + msg, + text_patterns.join(" / ") + ); + } + }; + let status = add_or_update_db_user(pool, new_message, ctx, obtained_xp).await; diff --git a/src/lib.rs b/src/lib.rs index 348f5ec..f0a5fc7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ pub mod commands; pub mod data; +pub mod database; pub mod enums; pub mod event_handler; pub mod extra_threads; diff --git a/src/main.rs b/src/main.rs index 71106c4..cebe9fa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,13 @@ mod commands; mod data; +mod database; mod enums; mod event_handler; mod extra_threads; mod tests; mod utils; +use data::command_data::Error; use data::{ bot_data::{BOT_PREFIXES, BOT_TOKEN, START_TIME}, command_data::Data, @@ -13,10 +15,10 @@ use data::{ use event_handler::handler::event_handler; use extra_threads::xp_command_cooldown::periodically_clean_users_on_diff_thread; use poise::serenity_prelude as serenity; -use std::sync::{atomic::AtomicU32, Arc}; +use std::sync::Arc; #[tokio::main] -async fn main() { +async fn main() -> Result<(), Error> { let _ = START_TIME.elapsed().as_secs(); // Dummy data to get the time elapsing started dotenv::dotenv().ok(); @@ -88,7 +90,6 @@ async fn main() { Box::pin(async move { poise::builtins::register_globally(ctx, &framework.options().commands).await?; Ok(Data { - hutao_mentions: AtomicU32::new(1), // It's better to clone the bot user once when it starts rather than do http // requests for the serenity::CurrentUser on every comman invocation. bot_user: Arc::from(ready.user.clone()), @@ -113,4 +114,6 @@ async fn main() { .await; client.unwrap().start().await.unwrap(); + + Ok(()) } diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 793a4e9..ed824a8 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -9,9 +9,9 @@ fn sized_send_sunc_unpin() {} fn normal_types() { use crate::data::command_data::Data; use crate::enums::command_enums::EmbedType; - use crate::enums::schemas::DatabaseSchema; + use crate::enums::schemas::LevelsSchema; sized_send_sunc_unpin::(); sized_send_sunc_unpin::(); - sized_send_sunc_unpin::(); + sized_send_sunc_unpin::(); } diff --git a/src/utils/replies.rs b/src/utils/replies.rs index d749954..72fd329 100644 --- a/src/utils/replies.rs +++ b/src/utils/replies.rs @@ -1,11 +1,11 @@ -use crate::data::command_data::{Data, Error}; +use crate::{data::command_data::Error, database::bot_mentions::add_mentions}; use poise::serenity_prelude as serenity; -use std::sync::atomic::Ordering; +use sqlx::SqlitePool; pub async fn handle_replies( + db: &SqlitePool, ctx: &serenity::Context, new_message: &serenity::Message, - data: &Data, msg: &str, ) -> Result<(), Error> { if msg @@ -15,10 +15,12 @@ pub async fn handle_replies( .join("") == "damnhutaomains" { - data.hutao_mentions.fetch_add(1, Ordering::SeqCst); + add_mentions(db, 1).await?; new_message.reply(ctx, "Any last words?").await?; } else if msg.contains("hutao") || msg.contains("hu tao") { - let mentions = data.hutao_mentions.fetch_add(1, Ordering::SeqCst); + let mentions = add_mentions(db, 1).await?; + #[cfg(feature = "debug")] + println!("Mentions: {}", mentions); new_message .reply(ctx, format!("Hu Tao has been mentioned {} times", mentions)) .await?;