Add badic mongo models and functions

This commit is contained in:
Victor Mignot 2022-11-19 21:19:48 -05:00
parent cfd481a86f
commit 3c11236404
No known key found for this signature in database
GPG key ID: FFE4EF056FB5E0D0
4 changed files with 250 additions and 28 deletions

View file

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

View file

@ -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::errors::ClientError;
use crate::DatabaseCredentials; use crate::DatabaseCredentials;
use super::models::{YorokobotModel, COLLECTIONS_NAMES};
/// Database client /// Database client
pub struct Client<'a> { pub struct Client {
mongo_client: Option<MongoClient>, 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 /// Create a new database client
pub fn new(credentials: &'a DatabaseCredentials) -> Client { pub fn new(credentials: DatabaseCredentials) -> Client {
return Client { return Client {
credentials, credentials,
mongo_client: None, mongo_client: None,
// database: None, database: None,
}; };
} }
/// Connect the client /// Connect the client
pub async fn connect(&mut self) -> Result<(), ClientError> { 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), Ok(c) => Some(c),
Err(e) => return Err(ClientError::Database(e)), 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(()) 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(()),
}
}
} }

View file

@ -4,13 +4,10 @@
use serde::{Deserialize, Serialize}; 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 /// Settings for a server
@ -35,3 +32,14 @@ pub struct Tag {
is_nsfw: bool, is_nsfw: bool,
subscribers: Vec<String>, 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();
}
}

View file

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