minneelyyyy/bigbirb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

better database handling, rwlock daily, some minor other improvements

Browse Source
This commit is contained in:
minneelyyyy
2024-12-09 14:07:24 -05:00
parent 4d8848094c
commit 0889297f2a
15 changed files with 86 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/commands/gambling/balance.rs
View File

@@ -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 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).",
common::no_ping_reply(&ctx, format!("{} **{}** token(s).",
if user.id == ctx.author().id {
"You have".to_string()
} else {
format!("{} has", user.name)
}, wealth)).await?;
format!("{} has", user)
}, wealth)
).await?;
Ok(())
}

34
src/commands/gambling/daily.rs
View File

@@ -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) => {
let db = &data.database;
let mut tx = db.begin().await?;
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?;
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 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 next = Duration::from_secs(24 * 60 * 60) - last.elapsed();
ctx.reply(format!("Your next daily will be available in **{}**.", format_duration(next))).await?;
}
Ok(())

14
src/commands/gambling/give.rs
View File

@@ -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(())
}

0
src/commands/gambling/leaderboard.rs Normal file
View File

19
src/commands/gambling/mod.rs
View File

@@ -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(())
}

17
src/commands/gambling/wager.rs
View File

@@ -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(())
}

8
src/commands/self_roles/color.rs
View File

@@ -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?;

4
src/commands/self_roles/disown.rs
View File

@@ -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? {

7
src/commands/self_roles/mod.rs
View File

@@ -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)

4
src/commands/self_roles/name.rs
View File

@@ -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? {

15
src/commands/self_roles/register.rs
View File

@@ -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(())
}

4
src/commands/self_roles/remove.rs
View File

@@ -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? {

4
src/commands/self_roles/whois.rs
View File

@@ -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")

8
src/common.rs
View File

@@ -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>;

16
src/main.rs
View File

@@ -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())),
})
})
})