diff --git a/crates/core/database/src/drivers/reference.rs b/crates/core/database/src/drivers/reference.rs index cd4bc1a1..f59dda16 100644 --- a/crates/core/database/src/drivers/reference.rs +++ b/crates/core/database/src/drivers/reference.rs @@ -2,7 +2,7 @@ use std::{collections::HashMap, sync::Arc}; use futures::lock::Mutex; -use crate::{Bot, User}; +use crate::{Bot, File, User}; database_derived!( /// Reference implementation @@ -10,5 +10,6 @@ database_derived!( pub struct ReferenceDb { pub bots: Arc>>, pub users: Arc>>, + pub files: Arc>>, } ); diff --git a/crates/core/database/src/models/files/mod.rs b/crates/core/database/src/models/files/mod.rs index 4a7ebf60..4d801b73 100644 --- a/crates/core/database/src/models/files/mod.rs +++ b/crates/core/database/src/models/files/mod.rs @@ -1,3 +1,5 @@ mod model; +mod ops; pub use model::*; +pub use ops::*; diff --git a/crates/core/database/src/models/files/model.rs b/crates/core/database/src/models/files/model.rs index d39e1eb3..85edcaea 100644 --- a/crates/core/database/src/models/files/model.rs +++ b/crates/core/database/src/models/files/model.rs @@ -1,3 +1,7 @@ +use crate::Database; + +use revolt_result::Result; + auto_derived_partial!( /// File pub struct File { @@ -57,3 +61,47 @@ auto_derived!( Audio, } ); + +impl File { + /// Use a file for a message attachment + pub async fn use_attachment(db: &Database, id: &str, parent: &str) -> Result { + db.find_and_use_attachment(id, "attachments", "message", parent) + .await + } + + /// Use a file for a user profile background + pub async fn use_background(db: &Database, id: &str, parent: &str) -> Result { + db.find_and_use_attachment(id, "backgrounds", "user", parent) + .await + } + + /// Use a file for a user avatar + pub async fn use_avatar(db: &Database, id: &str, parent: &str) -> Result { + db.find_and_use_attachment(id, "avatars", "user", parent) + .await + } + + /// Use a file for an icon + pub async fn use_icon(db: &Database, id: &str, parent: &str) -> Result { + db.find_and_use_attachment(id, "icons", "object", parent) + .await + } + + /// Use a file for a server icon + pub async fn use_server_icon(db: &Database, id: &str, parent: &str) -> Result { + db.find_and_use_attachment(id, "icons", "object", parent) + .await + } + + /// Use a file for a server banner + pub async fn use_banner(db: &Database, id: &str, parent: &str) -> Result { + db.find_and_use_attachment(id, "banners", "server", parent) + .await + } + + /// Use a file for an emoji + pub async fn use_emoji(db: &Database, id: &str, parent: &str) -> Result { + db.find_and_use_attachment(id, "emojis", "object", parent) + .await + } +} diff --git a/crates/core/database/src/models/files/ops.rs b/crates/core/database/src/models/files/ops.rs new file mode 100644 index 00000000..49c1e9a1 --- /dev/null +++ b/crates/core/database/src/models/files/ops.rs @@ -0,0 +1,30 @@ +use revolt_result::Result; + +use crate::File; + +mod mongodb; +mod reference; + +#[async_trait] +pub trait AbstractAttachment: Sync + Send { + /// Insert attachment into database. + async fn insert_attachment(&self, attachment: &File) -> Result<()>; + + /// Find an attachment by its details and mark it as used by a given parent. + async fn find_and_use_attachment( + &self, + id: &str, + tag: &str, + parent_type: &str, + parent_id: &str, + ) -> Result; + + /// Mark an attachment as having been reported. + async fn mark_attachment_as_reported(&self, id: &str) -> Result<()>; + + /// Mark an attachment as having been deleted. + async fn mark_attachment_as_deleted(&self, id: &str) -> Result<()>; + + /// Mark multiple attachments as having been deleted. + async fn mark_attachments_as_deleted(&self, ids: &[String]) -> Result<()>; +} diff --git a/crates/core/database/src/models/files/ops/mongodb.rs b/crates/core/database/src/models/files/ops/mongodb.rs new file mode 100644 index 00000000..de883c77 --- /dev/null +++ b/crates/core/database/src/models/files/ops/mongodb.rs @@ -0,0 +1,117 @@ +use bson::Document; +use revolt_result::Result; + +use crate::File; +use crate::MongoDb; + +use super::AbstractAttachment; + +static COL: &str = "bots"; + +#[async_trait] +impl AbstractAttachment for MongoDb { + /// Insert attachment into database. + async fn insert_attachment(&self, attachment: &File) -> Result<()> { + query!(self, insert_one, COL, &attachment).map(|_| ()) + } + + /// Find an attachment by its details and mark it as used by a given parent. + async fn find_and_use_attachment( + &self, + id: &str, + tag: &str, + parent_type: &str, + parent_id: &str, + ) -> Result { + let key = format!("{parent_type}_id"); + let file = query!( + self, + find_one, + COL, + doc! { + "_id": id, + "tag": tag, + &key: { + "$exists": false + } + } + )? + .ok_or_else(|| create_error!(NotFound))?; + + self.col::(COL) + .update_one( + doc! { + "_id": id + }, + doc! { + "$set": { + key: parent_id + } + }, + None, + ) + .await + .map_err(|_| create_database_error!("update_one", COL))?; + + Ok(file) + } + + /// Mark an attachment as having been reported. + async fn mark_attachment_as_reported(&self, id: &str) -> Result<()> { + self.col::(COL) + .update_one( + doc! { + "_id": id + }, + doc! { + "$set": { + "reported": true + } + }, + None, + ) + .await + .map(|_| ()) + .map_err(|_| create_database_error!("update_one", COL)) + } + + /// Mark an attachment as having been deleted. + async fn mark_attachment_as_deleted(&self, id: &str) -> Result<()> { + self.col::(COL) + .update_one( + doc! { + "_id": id + }, + doc! { + "$set": { + "deleted": true + } + }, + None, + ) + .await + .map(|_| ()) + .map_err(|_| create_database_error!("update_one", COL)) + } + + /// Mark multiple attachments as having been deleted. + async fn mark_attachments_as_deleted(&self, ids: &[String]) -> Result<()> { + self.col::(COL) + .update_one( + doc! { + "_id": { + "$in": ids + } + }, + doc! { + "$set": { + "deleted": true + } + }, + None, + ) + .await + .map(|_| ()) + .map_err(|_| create_database_error!("update_one", COL)) + } +} diff --git a/crates/core/database/src/models/files/ops/reference.rs b/crates/core/database/src/models/files/ops/reference.rs new file mode 100644 index 00000000..ee4d036c --- /dev/null +++ b/crates/core/database/src/models/files/ops/reference.rs @@ -0,0 +1,79 @@ +use revolt_result::Result; + +use crate::File; +use crate::ReferenceDb; + +use super::AbstractAttachment; + +#[async_trait] +impl AbstractAttachment for ReferenceDb { + /// Insert attachment into database. + async fn insert_attachment(&self, attachment: &File) -> Result<()> { + let mut attachments = self.files.lock().await; + if attachments.contains_key(&attachment.id) { + Err(create_database_error!("insert", "attachment")) + } else { + attachments.insert(attachment.id.to_string(), attachment.clone()); + Ok(()) + } + } + + /// Find an attachment by its details and mark it as used by a given parent. + async fn find_and_use_attachment( + &self, + id: &str, + _tag: &str, + _parent_type: &str, + _parent_id: &str, + ) -> Result { + let mut files = self.files.lock().await; + if let Some(file) = files.get_mut(id) { + // TODO: check tag + // TODO: set parent ID + Ok(file.clone()) + } else { + Err(create_error!(NotFound)) + } + } + + /// Mark an attachment as having been reported. + async fn mark_attachment_as_reported(&self, id: &str) -> Result<()> { + let mut files = self.files.lock().await; + if let Some(file) = files.get_mut(id) { + file.reported = Some(true); + Ok(()) + } else { + Err(create_error!(NotFound)) + } + } + + /// Mark an attachment as having been deleted. + async fn mark_attachment_as_deleted(&self, id: &str) -> Result<()> { + let mut files = self.files.lock().await; + if let Some(file) = files.get_mut(id) { + file.deleted = Some(true); + Ok(()) + } else { + Err(create_error!(NotFound)) + } + } + + /// Mark multiple attachments as having been deleted. + async fn mark_attachments_as_deleted(&self, ids: &[String]) -> Result<()> { + let mut files = self.files.lock().await; + + for id in ids { + if !files.contains_key(id) { + return Err(create_error!(NotFound)); + } + } + + for id in ids { + if let Some(file) = files.get_mut(id) { + file.reported = Some(true); + } + } + + Ok(()) + } +} diff --git a/crates/core/database/src/models/mod.rs b/crates/core/database/src/models/mod.rs index 9fb3ad3d..0560ca9f 100644 --- a/crates/core/database/src/models/mod.rs +++ b/crates/core/database/src/models/mod.rs @@ -11,7 +11,12 @@ pub use users::*; use crate::{Database, MongoDb, ReferenceDb}; pub trait AbstractDatabase: - Sync + Send + admin_migrations::AbstractMigrations + bots::AbstractBots + users::AbstractUsers + Sync + + Send + + admin_migrations::AbstractMigrations + + bots::AbstractBots + + files::AbstractAttachment + + users::AbstractUsers { }