diff --git a/src/commands/gambling/balance.rs b/src/commands/gambling/balance.rs index 069b657..c867a17 100644 --- a/src/commands/gambling/balance.rs +++ b/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) -> Result<(), Error> { let user = user.as_ref().unwrap_or(ctx.author()); - let data = ctx.data(); - let mut db = data.database.lock().await; - let db = db.as_mut(); + let db = &ctx.data().database; let wealth = super::get_balance(user.id, db).await?; - ctx.reply(format!("{} **{}** token(s).", - if user.id == ctx.author().id { - "You have".to_string() - } else { - format!("{} has", user.name) - }, wealth)).await?; + common::no_ping_reply(&ctx, format!("{} **{}** token(s).", + if user.id == ctx.author().id { + "You have".to_string() + } else { + format!("{} has", user) + }, wealth) + ).await?; Ok(()) } diff --git a/src/commands/gambling/daily.rs b/src/commands/gambling/daily.rs index 0343b6b..6217ca4 100644 --- a/src/commands/gambling/daily.rs +++ b/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 db = db.as_mut(); + let user = ctx.author().id; - 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) { - Some(daily) => { - - if daily.elapsed() >= Duration::from_secs(24 * 60 * 60) { - *daily = Instant::now(); - super::add_balance(id, 50, db).await?; - ctx.reply("Added **50** credits to your account!").await?; - } else { - let until_next_daily = Duration::from_secs(24 * 60 * 60) - daily.elapsed(); - ctx.reply(format!("Your next daily will be available in **{}**.", format_duration(until_next_daily))).await?; - } - }, - None => { - dailies.insert(id.clone(), Instant::now()); - super::add_balance(id, 50, db).await?; - ctx.reply("Added **50** credits to your account!").await?; - } + let db = &data.database; + let mut tx = db.begin().await?; + + let bal = super::get_balance(user, &mut *tx).await?; + super::change_balance(user, bal + 50, &mut *tx).await?; + + tx.commit().await?; + + ctx.reply(format!("**50** tokens have been added to your balance.")).await?; + } else { + let next = Duration::from_secs(24 * 60 * 60) - last.elapsed(); + ctx.reply(format!("Your next daily will be available in **{}**.", format_duration(next))).await?; } Ok(()) diff --git a/src/commands/gambling/give.rs b/src/commands/gambling/give.rs index 85ffa86..03381fb 100644 --- a/src/commands/gambling/give.rs +++ b/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 db = data.database.lock().await; - let db = db.as_mut(); + let mut tx = ctx.data().database.begin().await?; - 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(ctx.author().id, author_new_balance, db).await?; + super::change_balance(user.id, reciever_new_balance, &mut *tx).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(()) } \ No newline at end of file diff --git a/src/commands/gambling/leaderboard.rs b/src/commands/gambling/leaderboard.rs new file mode 100644 index 0000000..e69de29 diff --git a/src/commands/gambling/mod.rs b/src/commands/gambling/mod.rs index 4642500..1129b44 100644 --- a/src/commands/gambling/mod.rs +++ b/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 { +pub async fn get_balance<'a, E>(id: UserId, db: E) -> Result +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 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(()) -} \ No newline at end of file diff --git a/src/commands/gambling/wager.rs b/src/commands/gambling/wager.rs index 3885969..8548916 100644 --- a/src/commands/gambling/wager.rs +++ b/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(); - let mut db = data.database.lock().await; - let db = db.as_mut(); + // #[min = 1] does not appear to work with prefix commands + if amount < 1 { + 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(()) } \ No newline at end of file diff --git a/src/commands/self_roles/color.rs b/src/commands/self_roles/color.rs index 16968e9..860634b 100644 --- a/src/commands/self_roles/color.rs +++ b/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 db = data.database.lock().await; - let db = db.as_mut(); + let mut tx = ctx.data().database.begin().await?; 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?; diff --git a/src/commands/self_roles/disown.rs b/src/commands/self_roles/disown.rs index f95e997..c388c65 100644 --- a/src/commands/self_roles/disown.rs +++ b/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 mut db = data.database.lock().await; - let db = db.as_mut(); + let db = &ctx.data().database; if let Some(guild) = ctx.guild_id() { if let Some(role) = super::get_user_role(ctx.author().id, guild, db).await? { diff --git a/src/commands/self_roles/mod.rs b/src/commands/self_roles/mod.rs index 46dae12..2e20ce0 100644 --- a/src/commands/self_roles/mod.rs +++ b/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, Error> { +pub async fn get_user_role<'a, E>(user: UserId, guild: GuildId, db: E) -> Result, 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) diff --git a/src/commands/self_roles/name.rs b/src/commands/self_roles/name.rs index 5564ff0..7f67cc2 100644 --- a/src/commands/self_roles/name.rs +++ b/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 mut db = data.database.lock().await; - let db = db.as_mut(); + let db = &ctx.data().database; if let Some(guild) = ctx.guild_id() { let role = match super::get_user_role(ctx.author().id, guild, db).await? { diff --git a/src/commands/self_roles/register.rs b/src/commands/self_roles/register.rs index 580790c..40d432c 100644 --- a/src/commands/self_roles/register.rs +++ b/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 db = data.database.lock().await; - let db = db.as_mut(); + let mut tx = ctx.data().database.begin().await?; 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(()) } diff --git a/src/commands/self_roles/remove.rs b/src/commands/self_roles/remove.rs index 3493dbb..f842c7d 100644 --- a/src/commands/self_roles/remove.rs +++ b/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 mut db = data.database.lock().await; - let db = db.as_mut(); + let db = &ctx.data().database; if let Some(guild) = ctx.guild_id() { match super::get_user_role(user.id, guild, db).await? { diff --git a/src/commands/self_roles/whois.rs b/src/commands/self_roles/whois.rs index 8475494..1bd6417 100644 --- a/src/commands/self_roles/whois.rs +++ b/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 mut db = data.database.lock().await; - let db = db.as_mut(); + let db = &ctx.data().database; if let Some(guild) = ctx.guild_id() { let row = match sqlx::query("SELECT userid FROM selfroles WHERE roleid = $1") diff --git a/src/common.rs b/src/common.rs index 9a7ba04..7f892f3 100644 --- a/src/common.rs +++ b/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>, + pub database: Pool, pub mentions: Arc>>, /// last time the user redeemed a daily - pub dailies: Arc>>, + pub dailies: Arc>>, } pub type Error = Box; diff --git a/src/main.rs b/src/main.rs index c040bc0..c15e1da 100644 --- a/src/main.rs +++ b/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())), }) }) })