From 0cdd2735a676dbb9459120716ac58102bc647ecc Mon Sep 17 00:00:00 2001 From: bluepython508 <16466646+bluepython508@users.noreply.github.com> Date: Sun, 3 Nov 2024 22:18:10 +0000 Subject: [PATCH] Persist configuration with rusqlite --- .envrc | 1 + .gitignore | 1 + Cargo.lock | 163 ++++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 1 + flake.nix | 4 ++ src/main.rs | 89 ++++++++++++++++++---------- 6 files changed, 229 insertions(+), 30 deletions(-) diff --git a/.envrc b/.envrc index 9f972e3..738b564 100644 --- a/.envrc +++ b/.envrc @@ -4,3 +4,4 @@ export DISCORD_TOKEN_FILE=$(mktemp) cat >$DISCORD_TOKEN_FILE <>>); +struct Data(tokio::sync::Mutex); + +impl Data { + fn open(path: impl AsRef) -> Result { + let conn = rusqlite::Connection::open(path)?; + conn.execute_batch("CREATE TABLE IF NOT EXISTS guilds(guild TEXT PRIMARY KEY NOT NULL, town_square TEXT NOT NULL, cottages TEXT NOT NULL);")?; + Ok(Self(tokio::sync::Mutex::new(conn))) + } + + async fn get(&self, guild: serenity::GuildId) -> Result { + self.0 + .lock() + .await + .prepare_cached("SELECT guild, town_square, cottages FROM guilds WHERE guild = ?;")? + .query_row(rusqlite::params![guild.to_string()], |row| { + Ok(GuildData { + guild: row.get_unwrap::<_, String>("guild").parse().unwrap(), + town_square: row.get_unwrap::<_, String>("town_square").parse().unwrap(), + cottages: row.get_unwrap::<_, String>("cottages").parse().unwrap(), + }) + }) + .optional()? + .ok_or_eyre("This bot has not been configured in this server!") + } + + async fn insert( + &self, + GuildData { + guild, + town_square, + cottages, + }: GuildData, + ) -> Result<()> { + self.0 + .lock() + .await + .prepare_cached("INSERT INTO guilds (guild, town_square, cottages) VALUES (?, ?, ?)")? + .execute(rusqlite::params![ + guild.to_string(), + town_square.to_string(), + cottages.to_string() + ])?; + Ok(()) + } +} type Context<'a> = poise::Context<'a, Data, Error>; @@ -75,14 +117,8 @@ async fn dusk(ctx: Context<'_>) -> Result<()> { let GuildData { town_square, cottages, - } = ctx - .data() - .0 - .read() - .expect("No poisoned locks") - .get(&guild) - .copied() - .ok_or_eyre("This bot hasn't been configured for this server")?; + .. + } = ctx.data().get(guild).await?; run(ctx, town_square, cottages).await?; @@ -97,14 +133,8 @@ async fn dawn(ctx: Context<'_>) -> Result<()> { let GuildData { town_square, cottages, - } = ctx - .data() - .0 - .read() - .expect("No poisoned locks") - .get(&guild) - .copied() - .ok_or_eyre("This bot hasn't been configured for this server")?; + .. + } = ctx.data().get(guild).await?; run(ctx, cottages, town_square).await?; @@ -121,13 +151,11 @@ async fn configure( .guild_id() .ok_or_eyre("This bot only works in servers")?; - ctx.data().0.write().expect("No poisoned locks").insert( + ctx.data().insert(GuildData { guild, - GuildData { - town_square, - cottages, - }, - ); + town_square, + cottages, + }).await?; ctx.reply("Configured successfully!").await?; Ok(()) } @@ -160,6 +188,7 @@ async fn main() -> Result<()> { std::env::var_os("DISCORD_TOKEN_FILE").ok_or_eyre("A Discord token is required")?, ) .wrap_err("A Discord token is required")?; + let dbfile = std::env::var_os("DB_FILE").ok_or_eyre("A database file is required")?; let intents = serenity::GatewayIntents::non_privileged(); let framework = poise::Framework::builder() @@ -173,7 +202,7 @@ async fn main() -> Result<()> { for guild in &ready.guilds { register_commands(ctx, framework.options(), guild.id).await?; } - Ok(Data::default()) + Data::open(dbfile) }) }) .build();