save category

src/model.rs

@@ -2,9 +2,10 @@ use color_eyre::Result;
 use mongodb::{
     bson::doc,
     options::{ClientOptions, Credential},
-    Client, Database,
+    Client, Collection, Database,
 };
 use serde::{Deserialize, Serialize};
+use tracing::info;
 
 use crate::{lawsuit::Lawsuit, WrapErr};
 
@@ -12,7 +13,7 @@ use crate::{lawsuit::Lawsuit, WrapErr};
 pub struct State {
     pub guild_id: String,
     pub lawsuits: Vec<Lawsuit>,
-    pub justice_category: Option<String>,
+    pub court_category: Option<String>,
     pub court_rooms: Vec<CourtRoom>,
 }
 
@@ -50,13 +51,21 @@ impl Mongo {
         Ok(Self { db })
     }
 
-    pub async fn find_state(&self, guild_id: &str) -> Result<Option<State>> {
-        let collection = self.db.collection("state");
-        let state = collection
-            .find_one(doc! {"guild_id": guild_id  }, None)
+    pub async fn find_or_insert_state(&self, guild_id: &str) -> Result<State> {
+        let coll = self.state_coll();
+        let state = coll
+            .find_one(doc! {"guild_id": &guild_id  }, None)
             .await
             .wrap_err("find state")?;
 
+        let state = match state {
+            Some(state) => state,
+            None => {
+                info!(%guild_id, "No state found for guild, creating new state");
+                self.new_state(guild_id.to_owned()).await?
+            }
+        };
+
         Ok(state)
     }
 
@@ -64,15 +73,31 @@ impl Mongo {
         let state = State {
             guild_id,
             lawsuits: vec![],
-            justice_category: None,
+            court_category: None,
             court_rooms: vec![],
         };
 
-        let collection = self.db.collection::<State>("state");
-        collection
-            .insert_one(&state, None)
+        let coll = self.db.collection::<State>("state");
+        coll.insert_one(&state, None)
             .await
             .wrap_err("insert state")?;
         Ok(state)
     }
+
+    pub async fn set_court_category(&self, guild_id: &str, category: &str) -> Result<()> {
+        let _ = self.find_or_insert_state(guild_id).await?;
+        let coll = self.state_coll();
+        coll.update_one(
+            doc! {"guild_id": &guild_id  },
+            doc! {"$set": { "court_category": category }},
+            None,
+        )
+        .await
+        .wrap_err("update court category")?;
+        Ok(())
+    }
+
+    fn state_coll(&self) -> Collection<State> {
+        self.db.collection("state")
+    }
 }