use std::collections::HashMap; use std::ops::Deref; use futures::StreamExt; use mongodb::bson::{doc, to_document, Document}; use mongodb::error::Result; use mongodb::options::{FindOneOptions, FindOptions}; use mongodb::results::{DeleteResult, InsertOneResult, UpdateResult}; use serde::de::DeserializeOwned; use serde::Serialize; database_derived!( #[cfg(feature = "mongodb")] /// MongoDB implementation pub struct MongoDb(pub ::mongodb::Client, pub String); ); impl Deref for MongoDb { type Target = mongodb::Client; fn deref(&self) -> &Self::Target { &self.0 } } #[allow(dead_code)] impl MongoDb { /// Get the Revolt database pub fn db(&self) -> mongodb::Database { self.database(&self.1) } /// Get a collection by its name pub fn col(&self, collection: &str) -> mongodb::Collection { self.db().collection(collection) } /// Insert one document into a collection pub async fn insert_one( &self, collection: &'static str, document: T, ) -> Result { self.col::(collection).insert_one(document, None).await } /// Count documents by projection pub async fn count_documents( &self, collection: &'static str, projection: Document, ) -> Result { self.col::(collection) .count_documents(projection, None) .await } /// Find multiple documents in a collection with options pub 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? .filter_map(|s| async { if cfg!(debug_assertions) { // Hard fail on invalid documents Some(s.unwrap()) } else { s.ok() } }) .collect::>() .await) } /// Find multiple documents in a collection pub async fn find( &self, collection: &'static str, projection: Document, ) -> Result> { self.find_with_options(collection, projection, None).await } /// Find one document with options pub 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 } /// Find one document pub async fn find_one( &self, collection: &'static str, projection: Document, ) -> Result> { self.find_one_with_options(collection, projection, None) .await } /// Find one document by its ID pub async fn find_one_by_id( &self, collection: &'static str, id: &str, ) -> Result> { self.find_one( collection, doc! { "_id": id }, ) .await } /// Update one document given a projection, partial document, and list of paths to unset pub 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) }? }; self.col::(collection) .update_one(projection, query, None) .await } /// Update one document given an ID, partial document, and list of paths to unset pub 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 } /// Delete one document by the given projection pub async fn delete_one( &self, collection: &'static str, projection: Document, ) -> Result { self.col::(collection) .delete_one(projection, None) .await } /// Delete one document by the given ID pub async fn delete_one_by_id( &self, collection: &'static str, id: &str, ) -> Result { self.delete_one( collection, doc! { "_id": id }, ) .await } } /// Just a string ID struct #[derive(Deserialize)] pub struct DocumentId { #[serde(rename = "_id")] pub id: String, } pub trait IntoDocumentPath: Send + Sync { /// Create JSON key path fn as_path(&self) -> Option<&'static str>; } /// Prefix keys on an arbitrary object pub fn prefix_keys(t: &T, prefix: &str) -> HashMap { let v: String = serde_json::to_string(t).unwrap(); let v: HashMap = serde_json::from_str(&v).unwrap(); v.into_iter() .filter(|(_k, v)| !v.is_null()) .map(|(k, v)| (prefix.to_owned() + &k, v)) .collect() }