Add badic mongo models and functions

src/client.rs

@@ -26,12 +26,12 @@ use serenity::{prelude::GatewayIntents, Client as DiscordClient};
 ///
 /// # }
 /// ```
-pub struct Client<'a> {
+pub struct Client {
     /// The Serenity Discord Client
     discord_client: DiscordClient,
 
     /// The database client
-    database_client: DatabaseClient<'a>,
+    database_client: DatabaseClient,
 }
 
 /// Yorokobot connection credentials
@@ -40,10 +40,10 @@ pub struct ClientCredentials<'a> {
     pub discord_token: &'a String,
 
     /// MongoDB connection string.
-    pub db_credentials: &'a DatabaseCredentials,
+    pub db_credentials: DatabaseCredentials,
 }
 
-impl<'a> Client<'a> {
+impl<'a> Client {
     /// Create a Yorokobot client
     pub async fn new(credentials: ClientCredentials<'a>) -> Result<Client, ClientError> {
         let discord_client = match DiscordClient::builder(

src/database/client.rs

@@ -1,44 +1,258 @@
-use mongodb::Client as MongoClient;
+use std::collections::HashSet;
 
+use futures::TryStreamExt;
+use mongodb::{
+    bson::{doc, from_bson, Bson, Document},
+    Client as MongoClient, Collection, Database,
+};
+use serde::{Deserialize, Serialize};
+
+use crate::environment::get_env_variable;
 use crate::errors::ClientError;
 use crate::DatabaseCredentials;
 
+use super::models::{YorokobotModel, COLLECTIONS_NAMES};
+
 /// Database client
-pub struct Client<'a> {
+pub struct Client {
     mongo_client: Option<MongoClient>,
-    // database: Option<Database>,
+    database: Option<Database>,
-    credentials: &'a DatabaseCredentials,
+    credentials: DatabaseCredentials,
 }
 
-impl<'a> Client<'a> {
+impl Client {
     /// Create a new database client
-    pub fn new(credentials: &'a DatabaseCredentials) -> Client {
+    pub fn new(credentials: DatabaseCredentials) -> Client {
         return Client {
             credentials,
             mongo_client: None,
-            // database: None,
+            database: None,
         };
     }
 
     /// Connect the client
     pub async fn connect(&mut self) -> Result<(), ClientError> {
-        self.mongo_client = match MongoClient::with_options(self.credentials.clone()) {
+        self.mongo_client = match MongoClient::with_options(self.credentials.to_owned()) {
             Ok(c) => Some(c),
             Err(e) => return Err(ClientError::Database(e)),
         };
 
-        if let None = self.mongo_client.as_ref().unwrap().default_database() {
+        self.database = Some(
-            // TODO:
+            self.mongo_client
-            // Implement an Environment Variable catcher to wrap std::env::var()
+                .as_ref()
-            // As we often call it and always have to use a match control flow
+                .unwrap()
+                .database(get_env_variable("MONGO_DEFAULT_DB").as_str()),
+        );
 
-            // TODO:
+        // TODO:
-            // Complete error kind to be more specific.
+        // Complete error kind to be more specific.
-            // Ex: DatabaseConnection
+        // Ex: DatabaseConnection
 
-            todo!();
+        self.check_init_error().await;
-        }
 
         Ok(())
     }
+
+    async fn check_init_error(&mut self) {
+        self.check_collections_presence().await;
+    }
+
+    async fn check_collections_presence(&mut self) {
+        let mut missing_collections: Vec<&str> = vec![];
+        let collections: HashSet<String> = match self
+            .database
+            .as_ref()
+            .unwrap()
+            .list_collection_names(None)
+            .await
+        {
+            Ok(n) => n.into_iter().collect(),
+            Err(e) => panic!("Could not list collections: {e}"),
+        };
+
+        for col in COLLECTIONS_NAMES {
+            if !collections.contains(col) {
+                missing_collections.push(col);
+            }
+        }
+
+        if missing_collections.len() != 0 {
+            panic!(
+                "Missing the following the following collections: {}",
+                missing_collections.join(", ")
+            );
+        }
+    }
+
+    fn get_collection<T: YorokobotModel>(&self) -> Collection<Document> {
+        self.database
+            .as_ref()
+            .expect("Could not retrieve database")
+            .collection(&T::get_collection_name())
+    }
+
+    fn get_typed_collection<T: YorokobotModel>(&self) -> Collection<T> {
+        self.database
+            .as_ref()
+            .expect("Could not retrieve database")
+            .collection::<T>(&T::get_collection_name())
+    }
+
+    #[allow(dead_code)]
+    pub async fn get_by_id<T: YorokobotModel + for<'de> Deserialize<'de>>(&self, id: &str) -> T {
+        self.get_one(doc! {"_id": id}).await
+    }
+
+    #[allow(dead_code)]
+    pub async fn get_one<T: YorokobotModel + for<'de> Deserialize<'de>>(
+        &self,
+        filter: Document,
+    ) -> T {
+        let result = self
+            .get_collection::<T>()
+            .find_one(filter, None)
+            .await
+            .expect("Could not issue request")
+            .expect("Could not find matching data");
+
+        return from_bson(Bson::Document(result)).expect("Could not deserialize data");
+    }
+
+    #[allow(dead_code)]
+    pub async fn get_all<T: YorokobotModel + for<'de> Deserialize<'de>>(
+        &self,
+        filter: Option<Document>,
+    ) {
+        let mut result: Vec<T> = vec![];
+
+        let mut cursor = match filter {
+            Some(f) => self.get_collection::<T>().find(f, None).await,
+            None => self.get_collection::<T>().find(doc! {}, None).await,
+        }
+        .expect("Could not issue request");
+
+        while let Some(document) = cursor.try_next().await.expect("Could not fetch results") {
+            result
+                .push(from_bson(Bson::Document(document)).expect("Could not deserialize document"));
+        }
+    }
+
+    #[allow(dead_code)]
+    // TODO: Set true error handling
+    pub async fn insert_one<T: YorokobotModel + Serialize>(&self, document: T) -> Result<(), ()> {
+        match self
+            .get_typed_collection::<T>()
+            .insert_one(document, None)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(_) => Err(()),
+        }
+    }
+
+    #[allow(dead_code)]
+    // TODO: Set true error handling
+    pub async fn insert_many<T: YorokobotModel + Serialize>(
+        &self,
+        documents: Vec<T>,
+    ) -> Result<(), ()> {
+        match self
+            .get_typed_collection::<T>()
+            .insert_many(documents, None)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(_) => Err(()),
+        }
+    }
+
+    #[allow(dead_code)]
+    // TODO: Set true error handling
+    pub async fn delete_one<T: YorokobotModel>(&self, document: Document) -> Result<(), ()> {
+        match self
+            .get_typed_collection::<T>()
+            .delete_one(document, None)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(_) => Err(()),
+        }
+    }
+
+    #[allow(dead_code)]
+    // TODO: Set true error handling
+    pub async fn delete_by_id<T: YorokobotModel>(&self, id: &str) -> Result<(), ()> {
+        match self
+            .get_typed_collection::<T>()
+            .delete_one(doc! {"_id": id}, None)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(_) => Err(()),
+        }
+    }
+
+    #[allow(dead_code)]
+    // TODO: Set true error handling
+    pub async fn delete_many<T: YorokobotModel>(&self, document: Document) -> Result<(), ()> {
+        match self
+            .get_typed_collection::<T>()
+            .delete_many(document, None)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(_) => Err(()),
+        }
+    }
+
+    #[allow(dead_code)]
+    //TODO: Set true error handling
+    pub async fn update_one<T: YorokobotModel>(
+        &self,
+        document: Document,
+        update: Document,
+    ) -> Result<(), ()> {
+        match self
+            .get_typed_collection::<T>()
+            .update_one(document, update, None)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(_) => Err(()),
+        }
+    }
+
+    #[allow(dead_code)]
+    //TODO: Set true error handling
+    pub async fn update_by_id<T: YorokobotModel>(
+        &self,
+        document_id: &str,
+        update: Document,
+    ) -> Result<(), ()> {
+        match self
+            .get_typed_collection::<T>()
+            .update_one(doc! {"_id": document_id}, update, None)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(_) => Err(()),
+        }
+    }
+
+    #[allow(dead_code)]
+    //TODO: Set true error handling
+    pub async fn update_many<T: YorokobotModel>(
+        &self,
+        document: Document,
+        update: Document,
+    ) -> Result<(), ()> {
+        match self
+            .get_typed_collection::<T>()
+            .update_many(document, update, None)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(_) => Err(()),
+        }
+    }
 }

src/database/models.rs

@@ -4,13 +4,10 @@
 
 use serde::{Deserialize, Serialize};
 
-/// All the models within Mongo COllections
+pub const COLLECTIONS_NAMES: [&str; 2] = ["guilds", "tags"];
-pub enum CollectionModels {
-    /// Discord Guild
-    Guild(Guild),
 
-    /// Yorokobot tags
+pub trait YorokobotModel {
-    Tag(Tag),
+    fn get_collection_name() -> String;
 }
 
 /// Settings for a server
@@ -35,3 +32,14 @@ pub struct Tag {
     is_nsfw: bool,
     subscribers: Vec<String>,
 }
+
+impl YorokobotModel for Guild {
+    fn get_collection_name() -> String {
+        return "guilds".to_string();
+    }
+}
+impl YorokobotModel for Tag {
+    fn get_collection_name() -> String {
+        return "traits".to_string();
+    }
+}

src/main.rs

@@ -34,7 +34,7 @@ async fn main() -> std::process::ExitCode {
 
     let credentials = ClientCredentials {
         discord_token: &discord_token,
-        db_credentials: &db_credentials,
+        db_credentials: db_credentials,
     };
 
     let mut client = match Client::new(credentials).await {