add custom prefixes for servers to set

src/commands/mod.rs

@@ -6,6 +6,7 @@ mod yeehaw;
 mod gambling;
 mod eval;
 mod self_roles;
+mod settings;
 
 use crate::common::{Data, Error};
 
@@ -22,5 +23,6 @@ pub fn commands() -> Vec<Command<Data, Error>> {
         gambling::shop::buy(),
         eval::eval(),
         self_roles::role(),
+        settings::setting(),
     ]
 }

src/commands/settings.rs

@@ -0,0 +1,30 @@
+use crate::common::{Context, Error};
+
+#[poise::command(prefix_command, slash_command, required_permissions = "MANAGE_GUILD")]
+async fn prefix(ctx: Context<'_>, prefix: String) -> Result<(), Error> {
+    let guild = match ctx.guild_id() {
+        Some(g) => g,
+        None => {
+            ctx.reply("This command must be ran within a guild.").await?;
+            return Ok(());
+        }
+    };
+
+    let mut tx = ctx.data().database.begin().await?;
+
+    sqlx::query("INSERT INTO settings (guildid, prefix) VALUES ($1, $2) ON CONFLICT (guildid) DO UPDATE SET prefix = EXCLUDED.prefix")
+        .bind(guild.get() as i64)
+        .bind(&prefix)
+        .execute(&mut *tx).await?;
+
+    tx.commit().await?;
+
+    ctx.reply(format!("This server's custom prefix has been updated to `{prefix}`.")).await?;
+
+    Ok(())
+}
+
+#[poise::command(prefix_command, slash_command, subcommands("prefix"), subcommand_required)]
+pub async fn setting(_ctx: Context<'_>) -> Result<(), Error> {
+    Ok(())
+}

src/common.rs

@@ -3,6 +3,7 @@ use sqlx::{Pool, Postgres};
 
 pub struct Data {
     pub database: Pool<Postgres>,
+    pub prefix: Option<String>,
 }
 
 pub type Error = Box<dyn std::error::Error + Send + Sync>;

src/main.rs

@@ -8,9 +8,11 @@ use std::env;
 use std::sync::Arc;
 
 use poise::serenity_prelude::{self as serenity};
+use poise::PartialContext;
 use clap::Parser;
 
 use sqlx::postgres::PgPoolOptions;
+use sqlx::Row;
 
 #[derive(Parser, Debug)]
 struct BotArgs {
@@ -35,6 +37,22 @@ async fn event_handler(
     Ok(())
 }
 
+async fn get_prefix(ctx: PartialContext<'_, Data, Error>) -> Result<Option<String>, Error> {
+    let guild = match ctx.guild_id {
+        Some(guild) => guild,
+        None => return Ok(None),
+    };
+
+    let db = &ctx.data.database;
+
+    let prefix = sqlx::query("SELECT prefix FROM settings WHERE guildid = $1")
+        .bind(guild.get() as i64)
+        .fetch_one(db).await.ok()
+            .map(|x| x.get(0)).unwrap_or(ctx.data.prefix.clone());
+
+    Ok(prefix)
+}
+
 #[tokio::main]
 async fn main() -> Result<(), Error> {
     dotenv::dotenv().ok();
@@ -42,16 +60,17 @@ async fn main() -> Result<(), Error> {
 
     let token = env::var("DISCORD_BOT_TOKEN")?;
     let database_url = env::var("DATABASE_URL")?;
-    let intents =
+    let intents = args.prefix.clone().map(|_|
         serenity::GatewayIntents::GUILD_MESSAGES
-        | serenity::GatewayIntents::DIRECT_MESSAGES
+            | serenity::GatewayIntents::DIRECT_MESSAGES
-        | serenity::GatewayIntents::MESSAGE_CONTENT;
+            | serenity::GatewayIntents::MESSAGE_CONTENT)
+        .unwrap_or(serenity::GatewayIntents::empty());
 
     let framework = poise::Framework::builder()
         .options(poise::FrameworkOptions {
             commands: commands::commands(),
             prefix_options: poise::PrefixFrameworkOptions {
-                prefix: args.prefix,
+                dynamic_prefix: Some(|ctx| Box::pin(get_prefix(ctx))),
                 edit_tracker: Some(Arc::new(
                     poise::EditTracker::for_timespan(std::time::Duration::from_secs(10)))),
                 case_insensitive_commands: true,
@@ -122,9 +141,18 @@ async fn main() -> Result<(), Error> {
                     "#
                 ).execute(&database).await?;
 
+                sqlx::query(
+                    r#"
+                    CREATE TABLE IF NOT EXISTS settings (
+                        guildid BIGINT NOT NULL PRIMARY KEY,
+                        prefix TEXT
+                    )
+                    "#
+                ).execute(&database).await?;
+
                 println!("Bot is ready!");
 
-                Ok(Data { database })
+                Ok(Data { database, prefix: args.prefix })
             })
         })
         .build();