better database handling, rwlock daily, some minor other improvements

src/commands/gambling/balance.rs

@@ -1,22 +1,22 @@
-use crate::common::{Context, Error};
+
+use crate::common::{self, Context, Error};
 use poise::serenity_prelude as serenity;
 
 /// Tells you what your or someone else's balance is
 #[poise::command(slash_command, prefix_command, aliases("bal"))]
 pub async fn balance(ctx: Context<'_>, user: Option<serenity::User>) -> Result<(), Error> {
     let user = user.as_ref().unwrap_or(ctx.author());
-    let data = ctx.data();
+    let db = &ctx.data().database;
-    let mut db = data.database.lock().await;
-    let db = db.as_mut();
 
     let wealth = super::get_balance(user.id, db).await?;
 
-    ctx.reply(format!("{} **{}** token(s).",
+    common::no_ping_reply(&ctx, format!("{} **{}** token(s).",
         if user.id == ctx.author().id {
             "You have".to_string()
         } else {
-                          format!("{} has", user.name)
+            format!("{} has", user)
-                      }, wealth)).await?;
+        }, wealth)
+    ).await?;
 
     Ok(())
 }

src/commands/gambling/daily.rs

@@ -15,30 +15,26 @@ fn format_duration(duration: Duration) -> String {
 #[poise::command(slash_command, prefix_command)]
 pub async fn daily(ctx: Context<'_>) -> Result<(), Error> {
     let data = ctx.data();
-    let mut db = data.database.lock().await;
+    let user = ctx.author().id;
-    let db = db.as_mut();
 
-    let id = ctx.author().id;
+    let day_ago = Instant::now() - Duration::from_secs(24 * 60 * 60);
+    let last = *data.dailies.read().await.get(&user).unwrap_or(&day_ago);
 
-    let mut dailies = data.dailies.lock().await;
+    if last <= day_ago {
+        data.dailies.write().await.insert(user, Instant::now());
 
-    match dailies.get_mut(&id) {
+        let db = &data.database;
-        Some(daily) => {
+        let mut tx = db.begin().await?;
 
-            if daily.elapsed() >= Duration::from_secs(24 * 60 * 60) {
+        let bal = super::get_balance(user, &mut *tx).await?;
-                *daily = Instant::now();
+        super::change_balance(user, bal + 50, &mut *tx).await?;
-                super::add_balance(id, 50, db).await?;
+
-                ctx.reply("Added **50** credits to your account!").await?;
+        tx.commit().await?;
+
+        ctx.reply(format!("**50** tokens have been added to your balance.")).await?;
     } else {
-                let until_next_daily = Duration::from_secs(24 * 60 * 60) - daily.elapsed();
+        let next = Duration::from_secs(24 * 60 * 60) - last.elapsed();
-                ctx.reply(format!("Your next daily will be available in **{}**.", format_duration(until_next_daily))).await?;
+        ctx.reply(format!("Your next daily will be available in **{}**.", format_duration(next))).await?;
-            }
-        },
-        None => {
-            dailies.insert(id.clone(), Instant::now());
-            super::add_balance(id, 50, db).await?;
-            ctx.reply("Added **50** credits to your account!").await?;
-        }
     }
 
     Ok(())

src/commands/gambling/give.rs

@@ -9,23 +9,23 @@ pub async fn give(ctx: Context<'_>, user: serenity::User, #[min = 1] amount: i32
         return Ok(());
     }
 
-    let data = ctx.data();
+    let mut tx = ctx.data().database.begin().await?;
-    let mut db = data.database.lock().await;
-    let db = db.as_mut();
 
-    let author_balance = super::get_balance(ctx.author().id, db).await?;
+    let author_balance = super::get_balance(ctx.author().id, &mut *tx).await?;
 
     if author_balance < amount {
         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, db).await? + amount;
+        let reciever_new_balance = super::get_balance(user.id, &mut *tx).await? + amount;
 
-        super::change_balance(user.id, reciever_new_balance, db).await?;
+        super::change_balance(user.id, reciever_new_balance, &mut *tx).await?;
-        super::change_balance(ctx.author().id, author_new_balance, db).await?;
+        super::change_balance(ctx.author().id, author_new_balance, &mut *tx).await?;
 
         ctx.reply(format!("You've given **{}** **{}** tokens!", user.display_name(), amount)).await?;
     }
 
+    tx.commit().await?;
+
     Ok(())
 }

src/commands/gambling/mod.rs

@@ -3,12 +3,16 @@ pub mod balance;
 pub mod give;
 pub mod wager;
 pub mod daily;
+pub mod leaderboard;
 
 use crate::common::Error;
 use poise::serenity_prelude::UserId;
-use sqlx::{Row, PgConnection};
+use sqlx::{Row, PgExecutor};
 
-pub async fn get_balance(id: UserId, db: &mut PgConnection) -> Result<i32, Error> {
+pub async fn get_balance<'a, E>(id: UserId, db: E) -> Result<i32, Error>
+where
+    E: PgExecutor<'a>,
+{
     let row = sqlx::query("SELECT balance FROM bank WHERE id = $1")
         .bind(id.get() as i64)
         .fetch_one(db).await.ok();
@@ -22,7 +26,10 @@ pub async fn get_balance(id: UserId, db: &mut PgConnection) -> Result<i32, Error
     Ok(balance)
 }
 
-pub async fn change_balance(id: UserId, balance: i32, db: &mut PgConnection) -> Result<(), Error> {
+pub async fn change_balance<'a, E>(id: UserId, balance: i32, db: E) -> Result<(), Error>
+where
+    E: PgExecutor<'a>,
+{
     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)
@@ -30,9 +37,3 @@ pub async fn change_balance(id: UserId, balance: i32, db: &mut PgConnection) ->
 
     Ok(())
 }
-
-pub async fn add_balance(id: UserId, amount: i32, db: &mut PgConnection) -> Result<(), Error> {
-    let balance = get_balance(id, db).await?;
-    change_balance(id, balance + amount, db).await?;
-    Ok(())
-}

src/commands/gambling/wager.rs

@@ -4,14 +4,17 @@ use crate::common::{Context, Error};
 #[poise::command(slash_command, prefix_command)]
 pub async fn wager(
     ctx: Context<'_>,
-    #[min = 1]
     amount: i32) -> Result<(), Error>
 {
-    let data = ctx.data();
+    // #[min = 1] does not appear to work with prefix commands
-    let mut db = data.database.lock().await;
+    if amount < 1 {
-    let db = db.as_mut();
+        ctx.reply("you cannot wager less than 1 token.").await?;
+        return Ok(());
+    }
 
-    let mut wealth = super::get_balance(ctx.author().id, db).await?;
+    let mut tx = ctx.data().database.begin().await?;
+
+    let mut wealth = super::get_balance(ctx.author().id, &mut *tx).await?;
 
     if wealth < amount {
         ctx.reply(format!("You do not have enough tokens (**{}**) to wager this amount.",
@@ -29,7 +32,9 @@ pub async fn wager(
                           amount, wealth)).await?;
     }
 
-    super::change_balance(ctx.author().id, wealth, db).await?;
+    super::change_balance(ctx.author().id, wealth, &mut *tx).await?;
+
+    tx.commit().await?;
 
     Ok(())
 }

src/commands/self_roles/color.rs

@@ -42,12 +42,10 @@ pub async fn color(ctx: Context<'_>, color: String) -> Result<(), Error> {
         }
     };
 
-    let data = ctx.data();
+    let mut tx = ctx.data().database.begin().await?;
-    let mut db = data.database.lock().await;
-    let db = db.as_mut();
 
     if let Some(guild) = ctx.guild_id() {
-        match super::get_user_role(ctx.author().id, guild, db).await? {
+        match super::get_user_role(ctx.author().id, guild, &mut *tx).await? {
             Some(role) => {
                 guild.edit_role(ctx, role, EditRole::new().colour(color)).await?;
                 let role = guild.role(ctx, role).await?;
@@ -66,7 +64,7 @@ pub async fn color(ctx: Context<'_>, color: String) -> Result<(), Error> {
                     .bind(ctx.author().id.get() as i64)
                     .bind(role.id.get() as i64)
                     .bind(guild.get() as i64)
-                    .execute(db).await?;
+                    .execute(&mut *tx).await?;
 
                 let member = guild.member(ctx, ctx.author().id).await?;
                 member.add_role(ctx, role.clone()).await?;

src/commands/self_roles/disown.rs

@@ -4,9 +4,7 @@ use crate::common::{Context, Error};
 /// Remove and delete your personal role
 #[poise::command(slash_command, prefix_command)]
 pub async fn disown(ctx: Context<'_>) -> Result<(), Error> {
-    let data = ctx.data();
+    let db = &ctx.data().database;
-    let mut db = data.database.lock().await;
-    let db = db.as_mut();
 
     if let Some(guild) = ctx.guild_id() {
         if let Some(role) = super::get_user_role(ctx.author().id, guild, db).await? {

src/commands/self_roles/mod.rs

@@ -1,6 +1,6 @@
 
 use crate::common::{Context, Error};
-use sqlx::{PgConnection, Row};
+use sqlx::{PgExecutor, Row};
 use poise::serenity_prelude::{RoleId, UserId, GuildId};
 
 mod register;
@@ -26,7 +26,10 @@ pub async fn role(_ctx: Context<'_>) -> Result<(), Error> {
     Ok(())
 }
 
-pub async fn get_user_role(user: UserId, guild: GuildId, db: &mut PgConnection) -> Result<Option<RoleId>, Error> {
+pub async fn get_user_role<'a, E>(user: UserId, guild: GuildId, db: E) -> Result<Option<RoleId>, Error>
+where
+    E: PgExecutor<'a>,
+{
     match sqlx::query("SELECT roleid FROM selfroles WHERE userid = $1 AND guildid = $2")
         .bind(user.get() as i64)
         .bind(guild.get() as i64)

src/commands/self_roles/name.rs

@@ -6,9 +6,7 @@ use poise::serenity_prelude::EditRole;
 /// Change the name of your personal role
 #[poise::command(slash_command, prefix_command)]
 pub async fn name(ctx: Context<'_>, name: String) -> Result<(), Error> {
-    let data = ctx.data();
+    let db = &ctx.data().database;
-    let mut db = data.database.lock().await;
-    let db = db.as_mut();
 
     if let Some(guild) = ctx.guild_id() {
         let role = match super::get_user_role(ctx.author().id, guild, db).await? {

src/commands/self_roles/register.rs

@@ -6,34 +6,31 @@ use poise::serenity_prelude as serenity;
 /// Register an existing role as a user's custom role
 #[poise::command(slash_command, prefix_command, required_permissions = "MANAGE_ROLES")]
 pub async fn register(ctx: Context<'_>, user: serenity::User, role: serenity::Role) -> Result<(), Error> {
-    let data = ctx.data();
+    let mut tx = ctx.data().database.begin().await?;
-    let mut db = data.database.lock().await;
-    let db = db.as_mut();
 
     if let Some(guild) = ctx.guild_id() {
-        match super::get_user_role(user.id, guild, db).await? {
+        match super::get_user_role(user.id, guild, &mut *tx).await? {
             Some(role) => {
                 let role = guild.role(ctx, role).await?;
                 common::no_ping_reply(&ctx, format!("{} already has the role {} registered.", user, role)).await?;
-                Ok(())
             },
             None => {
                 sqlx::query("INSERT INTO selfroles (userid, guildid, roleid) VALUES ($1, $2, $3)")
                     .bind(user.id.get() as i64)
                     .bind(guild.get() as i64)
                     .bind(role.id.get() as i64)
-                    .execute(db).await?;
+                    .execute(&mut *tx).await?;
 
                 let member = guild.member(ctx, user.id).await?;
                 member.add_role(ctx, role.id).await?;
         
                 common::no_ping_reply(&ctx, format!("{} now has {} set as their personal role.", user, role)).await?;
-        
-                Ok(())
             }
         }
     } else {
         ctx.reply("This command can only be run in a guild!").await?;
-        Ok(())
     }
+
+    tx.commit().await?;
+    Ok(())
 }

src/commands/self_roles/remove.rs

@@ -6,9 +6,7 @@ use poise::serenity_prelude as serenity;
 /// force remove someone's role (this does not delete the role)
 #[poise::command(slash_command, prefix_command, required_permissions = "MANAGE_ROLES")]
 pub async fn remove(ctx: Context<'_>, user: serenity::User) -> Result<(), Error> {
-    let data = ctx.data();
+    let db = &ctx.data().database;
-    let mut db = data.database.lock().await;
-    let db = db.as_mut();
 
     if let Some(guild) = ctx.guild_id() {
         match super::get_user_role(user.id, guild, db).await? {

src/commands/self_roles/whois.rs

@@ -8,9 +8,7 @@ use sqlx::Row;
 /// Let you know who is the owner of a role.
 #[poise::command(slash_command, prefix_command)]
 pub async fn whois(ctx: Context<'_>, role: serenity::Role) -> Result<(), Error> {
-    let data = ctx.data();
+    let db = &ctx.data().database;
-    let mut db = data.database.lock().await;
-    let db = db.as_mut();
 
     if let Some(guild) = ctx.guild_id() {
         let row = match sqlx::query("SELECT userid FROM selfroles WHERE roleid = $1")

src/common.rs

@@ -1,15 +1,15 @@
 use std::sync::Arc;
-use tokio::sync::Mutex;
+use tokio::sync::{Mutex, RwLock};
 use std::collections::HashMap;
 use poise::{serenity_prelude::UserId, ReplyHandle};
-use sqlx::PgConnection;
+use sqlx::{Pool, Postgres};
 
 pub struct Data {
-    pub database: Arc<Mutex<PgConnection>>,
+    pub database: Pool<Postgres>,
     pub mentions: Arc<Mutex<HashMap<UserId, std::time::Instant>>>,
 
     /// last time the user redeemed a daily
-    pub dailies: Arc<Mutex<HashMap<UserId, std::time::Instant>>>,
+    pub dailies: Arc<RwLock<HashMap<UserId, std::time::Instant>>>,
 }
 
 pub type Error = Box<dyn std::error::Error + Send + Sync>;

src/main.rs

@@ -9,11 +9,11 @@ use std::env;
 use std::sync::Arc;
 
 use poise::serenity_prelude::{self as serenity};
-use tokio::sync::Mutex;
+use tokio::sync::{Mutex, RwLock};
 
 use clap::Parser;
 
-use sqlx::{PgConnection, Connection};
+use sqlx::postgres::PgPoolOptions;
 
 #[derive(Parser, Debug)]
 struct BotArgs {
@@ -68,7 +68,9 @@ async fn main() -> Result<(), Error> {
             Box::pin(async move {
                 poise::builtins::register_globally(ctx, &framework.options().commands).await?;
 
-                let mut database = PgConnection::connect(&database_url).await?;
+                let database = PgPoolOptions::new()
+                    .max_connections(5)
+                    .connect(&database_url).await?;
 
                 sqlx::query(
                     r#"
@@ -77,7 +79,7 @@ async fn main() -> Result<(), Error> {
                         balance INT
                     )
                     "#,
-                ).execute(&mut database).await?;
+                ).execute(&database).await?;
 
                 sqlx::query(
                     r#"
@@ -88,14 +90,14 @@ async fn main() -> Result<(), Error> {
                         UNIQUE (userid, guildid)
                     )
                     "#,
-                ).execute(&mut database).await?;
+                ).execute(&database).await?;
 
                 println!("Bot is ready!");
 
                 Ok(Data {
-                    database: Arc::new(Mutex::new(database)),
+                    database,
                     mentions: Arc::new(Mutex::new(HashMap::new())),
-                    dailies: Arc::new(Mutex::new(HashMap::new())),
+                    dailies: Arc::new(RwLock::new(HashMap::new())),
                 })
             })
         })