use std::ops::Deref; use bson::{to_document, Document}; use futures::StreamExt; use mongodb::{ options::{FindOneOptions, FindOptions}, results::{DeleteResult, InsertOneResult, UpdateResult}, }; use rocket::serde::DeserializeOwned; use serde::{Deserialize, Serialize}; use crate::{util::manipulation::prefix_keys, AbstractDatabase, Error, Result}; pub mod admin { pub mod migrations; } pub mod media { pub mod attachment; pub mod emoji; } pub mod channels { pub mod channel; pub mod channel_invite; pub mod channel_unread; pub mod message; } pub mod servers { pub mod server; pub mod server_ban; pub mod server_member; } pub mod users { pub mod bot; pub mod user; pub mod user_settings; } #[derive(Debug, Clone)] pub struct MongoDb(pub mongodb::Client); impl AbstractDatabase for MongoDb {} impl Deref for MongoDb { type Target = mongodb::Client; fn deref(&self) -> &Self::Target { &self.0 } } impl MongoDb { pub fn db(&self) -> mongodb::Database { self.database("revolt") } pub fn col(&self, collection: &str) -> mongodb::Collection { self.db().collection(collection) } async fn insert_one( &self, collection: &'static str, document: T, ) -> Result { self.col::(collection) .insert_one(document, None) .await .map_err(|_| Error::DatabaseError { operation: "insert_one", with: collection, }) } async fn find_with_options( &self, collection: &'static str, projection: Document, options: O, ) -> Result> where O: Into>, { Ok(self .col::(collection) .find(projection, options) .await .map_err(|_| Error::DatabaseError { operation: "find", with: collection, })? .filter_map(|s| async { s.ok() }) .collect::>() .await) } async fn find( &self, collection: &'static str, projection: Document, ) -> Result> { self.find_with_options(collection, projection, None).await } async fn find_one_with_options( &self, collection: &'static str, projection: Document, options: O, ) -> Result where O: Into>, { self.col::(collection) .find_one(projection, options) .await .map_err(|_| Error::DatabaseError { operation: "find_one", with: collection, })? .ok_or(Error::NotFound) } async fn find_one( &self, collection: &'static str, projection: Document, ) -> Result { self.find_one_with_options(collection, projection, None) .await } async fn find_one_by_id( &self, collection: &'static str, id: &str, ) -> Result { self.find_one( collection, doc! { "_id": id }, ) .await } async fn update_one( &self, collection: &'static str, projection: Document, partial: T, remove: Vec<&dyn IntoDocumentPath>, prefix: P, ) -> Result where P: Into>, { let prefix = prefix.into(); let mut unset = doc! {}; for field in remove { if let Some(path) = field.as_path() { if let Some(prefix) = &prefix { unset.insert(prefix.to_owned() + path, 1_i32); } else { unset.insert(path, 1_i32); } } } let query = doc! { "$unset": unset, "$set": if let Some(prefix) = &prefix { to_document(&prefix_keys(&partial, prefix)) } else { to_document(&partial) }.map_err(|_| Error::DatabaseError { operation: "to_document", with: collection })? }; self.col::(collection) .update_one(projection, query, None) .await .map_err(|_| Error::DatabaseError { operation: "update_one", with: collection, }) } async fn update_one_by_id( &self, collection: &'static str, id: &str, partial: T, remove: Vec<&dyn IntoDocumentPath>, prefix: P, ) -> Result where P: Into>, { self.update_one( collection, doc! { "_id": id }, partial, remove, prefix, ) .await } async fn delete_one( &self, collection: &'static str, projection: Document, ) -> Result { self.col::(collection) .delete_one(projection, None) .await .map_err(|_| Error::DatabaseError { operation: "delete_one", with: collection, }) } async fn delete_one_by_id(&self, collection: &'static str, id: &str) -> Result { self.delete_one( collection, doc! { "_id": id }, ) .await } } #[derive(Deserialize)] pub struct DocumentId { #[serde(rename = "_id")] pub id: String, } pub trait IntoDocumentPath: Send + Sync { fn as_path(&self) -> Option<&'static str>; }