Persist configuration with rusqlite

src/main.rs

@@ -1,24 +1,66 @@
-use std::{
-    collections::BTreeMap,
-    sync::{Arc, RwLock},
-};
+use std::path::Path;
 
 use eyre::{bail, Context as _, Error, OptionExt, Result};
 use futures::FutureExt;
 use poise::serenity_prelude as serenity;
 use rand::seq::SliceRandom;
+use rusqlite::OptionalExtension;
 
-// TODO: persist configuration
 // TODO: buttons?
+// TODO: shuffle channels as well as users/instead of users?
 
 #[derive(Copy, Clone)]
 struct GuildData {
+    guild: serenity::GuildId,
     town_square: serenity::ChannelId,
     cottages: serenity::ChannelId,
 }
 
-#[derive(Default)]
-struct Data(Arc<RwLock<BTreeMap<serenity::GuildId, GuildData>>>);
+struct Data(tokio::sync::Mutex<rusqlite::Connection>);
+
+impl Data {
+    fn open(path: impl AsRef<Path>) -> Result<Self> {
+        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<GuildData> {
+        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();