update v0.1.2, update self roles code

Cargo.lock

@@ -202,7 +202,7 @@ dependencies = [
 
 [[package]]
 name = "bot"
-version = "0.1.1"
+version = "0.1.2"
 dependencies = [
  "clap",
  "dotenv",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bot"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 
 [dependencies]

src/commands/mod.rs

@@ -5,7 +5,7 @@ mod dox;
 mod yeehaw;
 mod gambling;
 mod eval;
-mod self_roles;
+pub mod self_roles;
 mod settings;
 
 use crate::common::{Data, Error, Context};

src/commands/self_roles/color.rs

@@ -1,11 +1,11 @@
 
-use crate::common::{Context, Error};
+use crate::common::{self, Context, Error};
 
 use once_cell::sync::Lazy;
 use std::collections::HashMap;
 
 use hex_color::HexColor;
-use poise::serenity_prelude::{colours, Color, EditRole};
+use poise::serenity_prelude::{colours, Color, EditRole, Permissions};
 
 static COLORS: Lazy<HashMap<&'static str, Color>> = Lazy::new(|| {
     HashMap::from([
@@ -69,42 +69,29 @@ pub async fn color(ctx: Context<'_>, #[autocomplete = "autocomplete_colors"] col
         Color::from_rgb(rgb.r, rgb.g, rgb.b)
     };
 
-    if let Some(guild) = ctx.guild_id() {
+    let guild = if let Some(guild) = ctx.guild_id() {
+        guild
+    } else {
+        ctx.reply("This command can only be run inside of a guild.").await?;
+        return Ok(());
+    };
+
+    let user = ctx.author();
+
     let mut tx = ctx.data().database.begin().await?;
 
-        match super::get_user_role(ctx.author().id, guild, &mut *tx).await? {
+    if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? {
-            Some(role) => {
-                guild.edit_role(ctx, role, EditRole::new().colour(color)).await?;
         let role = guild.role(ctx, role).await?;
-
+        guild.edit_role(ctx, role.id, EditRole::new().colour(color)).await?;
-                ctx.reply(format!("{}'s color has been updated!", role)).await?;
+        common::no_ping_reply(&ctx, format!("{}'s color has been updated.", role)).await?;
-
-                Ok(())
-            },
-            None => {
-                let role = guild.create_role(ctx,
-                    EditRole::new()
-                        .name(format!("#{:06x}", color.0))
-                        .colour(color)).await?;
-
-                sqlx::query("INSERT INTO selfroles (userid, roleid, guildid) VALUES ($1, $2, $3)")
-                    .bind(ctx.author().id.get() as i64)
-                    .bind(role.id.get() as i64)
-                    .bind(guild.get() as i64)
-                    .execute(&mut *tx).await?;
-                
-                tx.commit().await?;
-
-                let member = guild.member(ctx, ctx.author().id).await?;
-                member.add_role(ctx, role.clone()).await?;
-
-                ctx.reply(format!("You have been given the {} role!", role)).await?;
-
-                return Ok(());
-            }
-        }
     } else {
-        ctx.reply("This command must be run within a server.").await?;
+        let role = guild.create_role(ctx, EditRole::new().colour(color).name(user.name.clone()).permissions(Permissions::empty())).await?;
+        super::update_user_role(user.id, guild, role.id, &mut *tx).await?;
+        let member = guild.member(ctx, user).await?;
+        member.add_role(ctx, role.id).await?;
+        tx.commit().await?;
+        common::no_ping_reply(&ctx, format!("{} has been given the new role {}.", user, role)).await?;
+    }
+
     Ok(())
 }
-}

src/commands/self_roles/disown.rs

@@ -4,25 +4,25 @@ 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 db = &ctx.data().database;
+    let guild = if let Some(guild) = ctx.guild_id() {
+        guild
+    } else {
+        ctx.reply("This command can only be run inside of a guild.").await?;
+        return Ok(());
+    };
 
-    if let Some(guild) = ctx.guild_id() {
+    let user = ctx.author();
-        if let Some(role) = super::get_user_role(ctx.author().id, guild, db).await? {
+
+    let mut tx = ctx.data().database.begin().await?;
+
+    if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? {
         guild.delete_role(ctx, role).await?;
-
+        super::remove_user_role(user.id, guild, &mut *tx).await?;
-            sqlx::query("DELETE FROM selfroles WHERE roleid = $1")
+        tx.commit().await?;
-                .bind(role.get() as i64)
+        ctx.reply("Your self role has been deleted.").await?;
-                .execute(db).await?;
-
-            ctx.reply("Your role has been successfully removed.").await?;
-
-            Ok(())
     } else {
-            ctx.reply("You do not currently have a personal role to remove.").await?;
+        ctx.reply("You currently have no self role to delete.").await?;
+    }
+
     Ok(())
 }
-    } else {
-        ctx.reply("This command must be called within a server.").await?;
-        Ok(())
-    }
-}

src/commands/self_roles/mod.rs

@@ -26,9 +26,45 @@ pub async fn role(_ctx: Context<'_>) -> Result<(), Error> {
     Ok(())
 }
 
+/// Remove a row concerning a user's self role from the database
+pub async fn remove_user_role<'a, E>(user: UserId, guild: GuildId, db: E) -> Result<(), Error>
+    where E: PgExecutor<'a>,
+{
+    sqlx::query("DELETE FROM selfroles WHERE userid = $1 AND guildid = $2")
+        .bind(user.get() as i64)
+        .bind(guild.get() as i64)
+        .execute(db).await?;
+
+    Ok(())
+}
+
+pub async fn remove_role<'a, E>(role: RoleId, guild: GuildId, db: E) -> Result<(), Error>
+where E: PgExecutor<'a>
+{
+    sqlx::query("DELETE FROM selfroles WHERE roleid = $1 AND guildid = $2")
+        .bind(role.get() as i64)
+        .bind(guild.get() as i64)
+        .execute(db).await?;
+
+    Ok(())
+}
+
+/// Replace a user's custom role with a new one
+pub async fn update_user_role<'a, E>(user: UserId, guild: GuildId, role: RoleId, db: E) -> Result<(), Error>
+    where E: PgExecutor<'a>,
+{
+    sqlx::query("INSERT INTO selfroles (userid, guildid, roleid) VALUES($1, $2, $3) ON CONFLICT (userid, guildid) DO UPDATE SET roleid = EXCLUDED.roleid")
+        .bind(user.get() as i64)
+        .bind(guild.get() as i64)
+        .bind(role.get() as i64)
+        .execute(db).await?;
+
+    Ok(())
+}
+
+/// Get a user's personal role id from the database
 pub async fn get_user_role<'a, E>(user: UserId, guild: GuildId, db: E) -> Result<Option<RoleId>, Error>
-where
+    where E: PgExecutor<'a>,
-    E: PgExecutor<'a>,
 {
     match sqlx::query("SELECT roleid FROM selfroles WHERE userid = $1 AND guildid = $2")
         .bind(user.get() as i64)

src/commands/self_roles/name.rs

@@ -1,44 +1,34 @@
 
-use crate::common::{Context, Error};
+use crate::common::{self, Context, Error};
 
-use poise::serenity_prelude::EditRole;
+use poise::serenity_prelude::{EditRole, Permissions};
 
 /// Change the name of your personal role
 #[poise::command(slash_command, prefix_command)]
 pub async fn name(ctx: Context<'_>, name: String) -> Result<(), Error> {
-    if let Some(guild) = ctx.guild_id() {
+    let guild = if let Some(guild) = ctx.guild_id() {
-        let mut tx = ctx.data().database.begin().await?;
+        guild
-
+    } else {
-        let role = match super::get_user_role(ctx.author().id, guild, &mut *tx).await? {
+        ctx.reply("This command can only be run inside of a guild.").await?;
-            Some(role) => role,
-            None => {
-                let role = guild.create_role(ctx, EditRole::new().name(name)).await?;
-
-                sqlx::query("INSERT INTO selfroles (userid, roleid, guildid) VALUES ($1, $2, $3)")
-                    .bind(ctx.author().id.get() as i64)
-                    .bind(role.id.get() as i64)
-                    .bind(guild.get() as i64)
-                    .execute(&mut *tx).await?;
-
-                tx.commit().await?;
-
-                let member = guild.member(ctx, ctx.author().id).await?;
-                member.add_role(ctx, role.clone()).await?;
-
-                ctx.reply(format!("You've been given the {} role!", role)).await?;
-
         return Ok(());
-            }
     };
 
-        guild.edit_role(ctx, role, EditRole::new().name(name)).await?;
+    let user = ctx.author();
+
+    let mut tx = ctx.data().database.begin().await?;
+
+    if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? {
         let role = guild.role(ctx, role).await?;
-
+        guild.edit_role(ctx, role.id, EditRole::new().name(name)).await?;
-        ctx.reply(format!("{} has been updated.", role)).await?;
+        common::no_ping_reply(&ctx, format!("{} has been updated.", role)).await?;
-
-        Ok(())
     } else {
-        ctx.reply("This command must be run within a server.").await?;
+        let role = guild.create_role(ctx, EditRole::new().name(name).permissions(Permissions::empty())).await?;
+        super::update_user_role(user.id, guild, role.id, &mut *tx).await?;
+        let member = guild.member(ctx, user).await?;
+        member.add_role(ctx, role.id).await?;
+        tx.commit().await?;
+        common::no_ping_reply(&ctx, format!("{} has been given the new role {}.", user, role)).await?;
+    }
+
     Ok(())
 }
-}

src/commands/self_roles/register.rs

@@ -3,35 +3,29 @@ use crate::common::{self, Context, Error};
 
 use poise::serenity_prelude as serenity;
 
-/// Register an existing role as a user's custom role
+/// Register an existing role as a user's custom role. This deletes their current self 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> {
-    if let Some(guild) = ctx.guild_id() {
+    let guild = if let Some(guild) = ctx.guild_id() {
+        guild
+    } else {
+        ctx.reply("This command can only be run inside of a guild.").await?;
+        return Ok(());
+    };
+
     let mut tx = ctx.data().database.begin().await?;
 
-        match super::get_user_role(user.id, guild, &mut *tx).await? {
+    if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? {
-            Some(role) => {
+        guild.delete_role(ctx, role).await?;
-                let role = guild.role(ctx, role).await?;
+    }
-                common::no_ping_reply(&ctx, format!("{} already has the role {} registered.", user, role)).await?;
-            },
-            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(&mut *tx).await?;
 
-                tx.commit().await?;
+    let member = guild.member(ctx, user).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?;
+    super::update_user_role(member.user.id, guild, role.id, &mut *tx).await?;
-            }
+    tx.commit().await?;
-        }
+
-    } else {
+    common::no_ping_reply(&ctx, format!("{} has been set as {}'s self role.", role, member.user)).await?;
-        ctx.reply("This command can only be run in a guild!").await?;
-    }
 
     Ok(())
 }

src/commands/self_roles/remove.rs

@@ -3,35 +3,25 @@ use crate::common::{self, Context, Error};
 
 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 db = &ctx.data().database;
+    let guild = if let Some(guild) = ctx.guild_id() {
-
+        guild
-    if let Some(guild) = ctx.guild_id() {
-        match super::get_user_role(user.id, guild, db).await? {
-            Some(role) => {
-                sqlx::query("DELETE FROM selfroles WHERE userid = $1 AND guildid = $2")
-                    .bind(user.id.get() as i64)
-                    .bind(guild.get() as i64)
-                    .execute(db).await?;
-
-                let member = guild.member(ctx, user.id).await?;
-                
-                member.remove_role(ctx, role).await?;
-                let role = guild.role(ctx, role).await?;
-
-                common::no_ping_reply(&ctx, format!("{} has had {} removed.", user, role)).await?;
-
-                Ok(())
-            },
-            None => {
-                common::no_ping_reply(&ctx, format!("{} does not have a personal role to remove.", user)).await?;
-                Ok(())
-            }
-        }
     } else {
-        ctx.reply("This command can only be run in a guild!").await?;
+        ctx.reply("This command can only be run inside of a guild.").await?;
+        return Ok(());
+    };
+
+    let mut tx = ctx.data().database.begin().await?;
+
+    if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? {
+        guild.delete_role(ctx, role).await?;
+        super::remove_role(role, guild, &mut *tx).await?;
+        tx.commit().await?;
+        common::no_ping_reply(&ctx, format!("{}'s self role has been deleted.", user)).await?;
+    } else {
+        common::no_ping_reply(&ctx, format!("{} has no self role.", user)).await?;
+    }
+
     Ok(())
 }
-}

src/commands/self_roles/whois.rs

@@ -11,11 +11,11 @@ pub async fn whois(ctx: Context<'_>, role: serenity::Role) -> Result<(), Error>
     let db = &ctx.data().database;
 
     if let Some(guild) = ctx.guild_id() {
-        let row = match sqlx::query("SELECT userid FROM selfroles WHERE roleid = $1")
+        let user = match sqlx::query("SELECT userid FROM selfroles WHERE roleid = $1")
             .bind(role.id.get() as i64)
             .fetch_one(db).await
         {
-            Ok(row) => row,
+            Ok(row) => UserId::new(row.try_get::<i64, usize>(0)? as u64),
             Err(sqlx::Error::RowNotFound) => {
                 ctx.reply("This role is not owned by anyone.").await?;
                 return Ok(());
@@ -23,10 +23,6 @@ pub async fn whois(ctx: Context<'_>, role: serenity::Role) -> Result<(), Error>
             Err(e) => return Err(Box::new(e)),
         };
 
-        let user: i64 = row.try_get(0)?;
-
-        let user = UserId::new(user as u64);
-
         let member = guild.member(ctx, user).await?;
 
         common::no_ping_reply(&ctx, format!("{} owns this role.", member)).await?;

src/main.rs

@@ -2,6 +2,7 @@ mod commands;
 
 pub mod common;
 pub mod inventory;
+use crate::commands::self_roles;
 use crate::common::{Context, Error, Data};
 
 use std::env;
@@ -22,15 +23,31 @@ struct BotArgs {
 }
 
 async fn event_handler(
-    _ctx: &serenity::Context,
+    ctx: &serenity::Context,
     event: &serenity::FullEvent,
     _framework: poise::FrameworkContext<'_, Data, Error>,
-    _data: &Data,
+    data: &Data,
 ) -> Result<(), Error> {
     match event {
         serenity::FullEvent::Message { new_message: message } => {
             if message.author.bot { return Ok(()) }
         }
+        serenity::FullEvent::GuildMemberRemoval { guild_id, user, .. } => {
+            let mut tx = data.database.begin().await?;
+
+            if let Some(role) = self_roles::get_user_role(user.id, *guild_id, &mut *tx).await? {
+                guild_id.delete_role(ctx, role).await?;
+                self_roles::remove_role(role, *guild_id, &mut *tx).await?;
+                tx.commit().await?;
+            }
+        },
+        serenity::FullEvent::GuildRoleDelete { guild_id, removed_role_id, .. } => {
+            let mut tx = data.database.begin().await?;
+
+            self_roles::remove_role(*removed_role_id, *guild_id, &mut *tx).await?;
+
+            tx.commit().await?;
+        }
         _ => (),
     }
 
@@ -64,7 +81,9 @@ async fn main() -> Result<(), Error> {
         serenity::GatewayIntents::GUILD_MESSAGES
             | serenity::GatewayIntents::DIRECT_MESSAGES
             | serenity::GatewayIntents::MESSAGE_CONTENT)
-        .unwrap_or(serenity::GatewayIntents::empty());
+        .unwrap_or(serenity::GatewayIntents::empty())
+            | serenity::GatewayIntents::GUILD_MEMBERS
+            | serenity::GatewayIntents::GUILDS;
 
     let framework = poise::Framework::builder()
         .options(poise::FrameworkOptions {