use a database :D

Cargo.toml

@@ -10,4 +10,5 @@ dotenv = "0.15.0"
 rand = "0.8.5"
 clap = { version = "4.5.20", features = ["derive"] }
 lamm = { git = "https://github.com/minneelyyyy/lamm", branch = "dev" }
-regex = "1.11.1"
+regex = "1.11.1"
+sqlx = { version = "0.8", features = [ "runtime-tokio-native-tls", "postgres" ] }

src/commands/gambling/balance.rs

@@ -1,4 +1,3 @@
-use super::get_user_wealth_mut;
 use crate::common::{Context, Error};
 use poise::serenity_prelude as serenity;
 
@@ -6,16 +5,15 @@ use poise::serenity_prelude as serenity;
 #[poise::command(slash_command, prefix_command)]
 pub async fn balance(ctx: Context<'_>, user: Option<serenity::User>) -> Result<(), Error> {
     let user = user.as_ref().unwrap_or(ctx.author());
-    let mut users = ctx.data().users.lock().await;
 
-    let wealth = get_user_wealth_mut(&mut users, user.id);
+    let wealth = super::get_balance(user.id, &ctx.data()).await?;
 
     ctx.reply(format!("{} **{}** token(s).",
                       if user.id == ctx.author().id {
                           "You have".to_string()
                       } else {
                           format!("{} has", user.name)
-                      }, *wealth)).await?;
+                      }, wealth)).await?;
 
     Ok(())
 }

src/commands/gambling/give.rs

@@ -1,30 +1,29 @@
 use crate::{Context, Error};
-use super::get_user_wealth_mut;
 use poise::serenity_prelude as serenity;
 
 /// Generously donate your tokens to someone else
 #[poise::command(slash_command, prefix_command)]
-pub async fn give(ctx: Context<'_>, user: serenity::User, amount: usize) -> Result<(), Error> {
+pub async fn give(ctx: Context<'_>, user: serenity::User, #[min = 1] amount: i32) -> Result<(), Error> {
     if user.bot {
-        ctx.reply("Don't waste your token(s) by giving them to a bot!").await?;
+        ctx.reply("Don't waste your tokens by giving them to a bot!").await?;
         return Ok(());
     }
 
-    let mut users = ctx.data().users.lock().await;
+    let data = ctx.data();
-    let author_wealth = get_user_wealth_mut(&mut users, ctx.author().id);
 
-    if *author_wealth < amount {
+    let author_balance = super::get_balance(ctx.author().id, &data).await?;
-        ctx.reply(format!("You only have **{}** token(s) and cannot give away **{}**.",
+
-                          *author_wealth, amount)).await?;
+    if author_balance < amount {
-        return Ok(());
+        ctx.reply(format!("You do not have a high enough balance (**{author_balance}**) to complete this transaction.")).await?;
+    } else {
+        let author_new_balance = author_balance - amount;
+        let reciever_new_balance = super::get_balance(user.id, &data).await? + amount;
+
+        super::change_balance(user.id, reciever_new_balance, &data).await?;
+        super::change_balance(ctx.author().id, author_new_balance, data).await?;
+
+        ctx.reply(format!("You've given **{}** **{}** tokens!", user.display_name(), amount)).await?;
     }
 
-    *author_wealth -= amount;
-
-    let receiver_wealth = get_user_wealth_mut(&mut users, user.id);
-    *receiver_wealth += amount;
-
-    ctx.reply(format!("You've given **{}** **{}** token(s).", user.name, amount)).await?;
-
     Ok(())
 }

src/commands/gambling/mod.rs

@@ -1,14 +1,35 @@
-use std::collections::HashMap;
-use poise::serenity_prelude::UserId;
 
 pub mod balance;
 pub mod give;
 pub mod wager;
 
-pub(self) fn get_user_wealth_mut(users: &mut HashMap<UserId, usize>, id: UserId) -> &mut usize {
+use crate::common::{Data, Error};
-    if users.get(&id).is_none() {
+use poise::serenity_prelude::UserId;
-        users.insert(id, 100);
+use sqlx::Row;
-    }
 
-    users.get_mut(&id).unwrap()
+pub async fn get_balance(id: UserId, data: &Data) -> Result<i32, Error> {
+    let db = &data.database;
+
+    let row = sqlx::query("SELECT balance FROM bank WHERE id = $1")
+        .bind(id.get() as i64)
+        .fetch_one(db).await.ok();
+
+    let balance = if let Some(row) = row {
+        row.try_get("balance")?
+    } else {
+        100
+    };
+
+    Ok(balance)
 }
+
+pub async fn change_balance(id: UserId, balance: i32, data: &Data) -> Result<(), Error> {
+    let db = &data.database;
+
+    sqlx::query("INSERT INTO bank (id, balance) VALUES ($1, $2) ON CONFLICT (id) DO UPDATE SET balance = EXCLUDED.balance")
+        .bind(id.get() as i64)
+        .bind(balance)
+        .execute(db).await?;
+
+    Ok(())
+}

src/commands/gambling/wager.rs

@@ -1,32 +1,32 @@
 use crate::common::{Context, Error};
-use super::get_user_wealth_mut;
 
 /// Put forward an amount of tokens to either lose or earn
 #[poise::command(slash_command, prefix_command)]
 pub async fn wager(
     ctx: Context<'_>,
     #[min = 1]
-    amount: usize) -> Result<(), Error>
+    amount: i32) -> Result<(), Error>
 {
-    let mut users = ctx.data().users.lock().await;
+    let data = ctx.data();
+    let mut wealth = super::get_balance(ctx.author().id, &data).await?;
 
-    let wealth = get_user_wealth_mut(&mut users, ctx.author().id);
+    if wealth < amount {
-
-    if *wealth < amount {
         ctx.reply(format!("You do not have enough tokens (**{}**) to wager this amount.",
-                          *wealth)).await?;
+                          wealth)).await?;
         return Ok(());
     }
 
     if rand::random() {
-        *wealth += amount;
+        wealth += amount;
         ctx.reply(format!("You just gained **{}** token(s)! You now have **{}**.",
-                          amount, *wealth)).await?;
+                          amount, wealth)).await?;
     } else {
-        *wealth -= amount;
+        wealth -= amount;
         ctx.reply(format!("You've lost **{}** token(s), you now have **{}**.",
-                          amount, *wealth)).await?;
+                          amount, wealth)).await?;
     }
 
+    super::change_balance(ctx.author().id, wealth, &data).await?;
+
     Ok(())
 }

src/common.rs

@@ -2,9 +2,10 @@ use std::sync::Arc;
 use tokio::sync::Mutex;
 use std::collections::HashMap;
 use poise::serenity_prelude::UserId;
+use sqlx::{Pool, Postgres};
 
 pub struct Data {
-    pub users: Arc<Mutex<HashMap<UserId, usize>>>,
+    pub database: Pool<Postgres>,
     pub mentions: Arc<Mutex<HashMap<UserId, std::time::Instant>>>,
 }
 

src/main.rs

@@ -7,13 +7,14 @@ use crate::common::{Context, Error, Data};
 use std::collections::HashMap;
 use std::env;
 use std::sync::Arc;
-use std::time::{Instant, Duration};
 
-use poise::serenity_prelude as serenity;
+use poise::serenity_prelude::{self as serenity};
 use tokio::sync::Mutex;
 
 use clap::Parser;
 
+use sqlx::postgres::PgPoolOptions;
+
 #[derive(Parser, Debug)]
 struct BotArgs {
     /// Prefix for the bot (if unspecified, the bot will not have one)
@@ -29,28 +30,9 @@ async fn event_handler(
 ) -> Result<(), Error> {
     match event {
         serenity::FullEvent::Message { new_message: message } => {
-            let mentions = ping_limit::extract_mentions(&message.content);
+            if message.author.bot { return Ok(()) }
-            let mut cooldowns = data.mentions.lock().await;
 
-            if mentions.iter()
+            ping_limit::ping_spam_yeller(ctx, &message, data).await?;
-                .filter(|&&id| id != message.author.id)
-                .any(|mention| cooldowns.get(mention).map(|t| Instant::now().duration_since(*t) < Duration::from_secs(20)).unwrap_or(false))
-            {
-                message.reply(ctx, "stop spamming!").await?;
-
-                let guild = match message.guild_id {
-                    Some(g) => g,
-                    None => return Ok(()),
-                };
-
-                let mut member = guild.member(ctx, message.author.id).await.unwrap();
-                member.disable_communication_until_datetime(ctx,
-                    serenity::Timestamp::from_unix_timestamp(serenity::Timestamp::now().unix_timestamp() + 60 as i64).unwrap()).await?;
-            }
-
-            for mention in mentions {
-                cooldowns.insert(mention, Instant::now());
-            }
         }
         _ => (),
     }
@@ -64,6 +46,7 @@ async fn main() -> Result<(), Error> {
     let args = BotArgs::parse();
 
     let token = env::var("DISCORD_BOT_TOKEN")?;
+    let database_url = env::var("DATABASE_URL")?;
     let intents = serenity::GatewayIntents::all();
 
     let framework = poise::Framework::builder()
@@ -84,8 +67,22 @@ async fn main() -> Result<(), Error> {
         .setup(|ctx, _ready, framework| {
             Box::pin(async move {
                 poise::builtins::register_globally(ctx, &framework.options().commands).await?;
+
+                let database = PgPoolOptions::new()
+                    .max_connections(4)
+                    .connect(&database_url).await?;
+
+                sqlx::query(
+                    r#"
+                    CREATE TABLE IF NOT EXISTS bank (
+                        id BIGINT PRIMARY KEY,
+                        balance INT
+                    )
+                    "#,
+                ).execute(&database).await?;
+
                 Ok(Data {
-                    users: Arc::new(Mutex::new(HashMap::new())),
+                    database,
                     mentions: Arc::new(Mutex::new(HashMap::new())),
                 })
             })

src/ping_limit.rs

@@ -1,7 +1,10 @@
+use crate::common::{Error, Data};
 
-use poise::serenity_prelude::*;
+use poise::serenity_prelude::{self as serenity, Message, UserId};
 use regex::Regex;
 
+use std::time::{Instant, Duration};
+
 pub fn extract_mentions(content: &str) -> Vec<UserId> {
     // Define the regex pattern for user mentions
     let re = Regex::new(r"<@(\d+)>").unwrap();
@@ -11,3 +14,30 @@ pub fn extract_mentions(content: &str) -> Vec<UserId> {
         .filter_map(|cap| cap.get(1).map(|id| id.as_str().parse().unwrap()))
         .collect()
 }
+
+pub async fn ping_spam_yeller(ctx: &serenity::Context, message: &Message, data: &Data) -> Result<(), Error> {
+    let mentions = extract_mentions(&message.content);
+    let mut cooldowns = data.mentions.lock().await;
+
+    if mentions.iter()
+        .filter(|&&id| id != message.author.id)
+        .any(|mention| cooldowns.get(mention).map(|t| Instant::now().duration_since(*t) < Duration::from_secs(20)).unwrap_or(false))
+    {
+        message.reply_ping(ctx, "stop spamming!").await?;
+
+        let guild = match message.guild_id {
+            Some(g) => g,
+            None => return Ok(()),
+        };
+
+        let mut member = guild.member(ctx, message.author.id).await.unwrap();
+        member.disable_communication_until_datetime(ctx,
+            serenity::Timestamp::from_unix_timestamp(serenity::Timestamp::now().unix_timestamp() + 60i64).unwrap()).await?;
+    }
+
+    for mention in mentions {
+        cooldowns.insert(mention, Instant::now());
+    }
+
+    Ok(())
+}