revolt_database/models/files/ops/
mongodb.rs

1use bson::to_document;
2use bson::Document;
3use revolt_config::report_internal_error;
4use revolt_result::Result;
5
6use crate::File;
7use crate::FileUsedFor;
8use crate::MongoDb;
9
10use super::AbstractAttachments;
11
12static COL: &str = "attachments";
13
14#[async_trait]
15impl AbstractAttachments for MongoDb {
16    /// Insert attachment into database.
17    async fn insert_attachment(&self, attachment: &File) -> Result<()> {
18        query!(self, insert_one, COL, &attachment).map(|_| ())
19    }
20
21    /// Fetch an attachment by its id.
22    async fn fetch_attachment(&self, tag: &str, file_id: &str) -> Result<File> {
23        query!(
24            self,
25            find_one,
26            COL,
27            doc! {
28                "_id": file_id,
29                "tag": tag
30            }
31        )?
32        .ok_or_else(|| create_error!(NotFound))
33    }
34
35    /// Fetch all deleted attachments.
36    async fn fetch_deleted_attachments(&self) -> Result<Vec<File>> {
37        query!(
38            self,
39            find,
40            COL,
41            doc! {
42                "deleted": true,
43                "reported": {
44                    "$ne": true
45                }
46            }
47        )
48    }
49
50    /// Fetch all dangling attachments.
51    async fn fetch_dangling_files(&self) -> Result<Vec<File>> {
52        query!(
53            self,
54            find,
55            COL,
56            doc! {
57                "used_for.type": {
58                    "$exists": 0
59                },
60                "deleted": {
61                    "$ne": true
62                }
63            }
64        )
65    }
66
67    /// Count references to a given hash.
68    async fn count_file_hash_references(&self, hash: &str) -> Result<usize> {
69        query!(
70            self,
71            count_documents,
72            COL,
73            doc! {
74                "hash": hash
75            }
76        )
77        .map(|count| count as usize)
78    }
79
80    /// Find an attachment by its details and mark it as used by a given parent.
81    async fn find_and_use_attachment(
82        &self,
83        id: &str,
84        tag: &str,
85        used_for: FileUsedFor,
86        uploader_id: String,
87    ) -> Result<File> {
88        let file = query!(
89            self,
90            find_one,
91            COL,
92            doc! {
93                "_id": id,
94                "tag": tag,
95                "used_for": {
96                    "$exists": false
97                }
98            }
99        )?
100        .ok_or_else(|| create_error!(NotFound))?;
101
102        self.col::<Document>(COL)
103            .update_one(
104                doc! {
105                    "_id": id
106                },
107                doc! {
108                    "$set": {
109                        "used_for": report_internal_error!(to_document(&used_for))?,
110                        "uploader_id": uploader_id
111                    }
112                },
113            )
114            .await
115            .map_err(|_| create_database_error!("update_one", COL))?;
116
117        Ok(file)
118    }
119
120    /// Mark an attachment as having been reported.
121    async fn mark_attachment_as_reported(&self, id: &str) -> Result<()> {
122        self.col::<Document>(COL)
123            .update_one(
124                doc! {
125                    "_id": id
126                },
127                doc! {
128                    "$set": {
129                        "reported": true
130                    }
131                },
132            )
133            .await
134            .map(|_| ())
135            .map_err(|_| create_database_error!("update_one", COL))
136    }
137
138    /// Mark an attachment as having been deleted.
139    async fn mark_attachment_as_deleted(&self, id: &str) -> Result<()> {
140        self.col::<Document>(COL)
141            .update_one(
142                doc! {
143                    "_id": id
144                },
145                doc! {
146                    "$set": {
147                        "deleted": true
148                    }
149                },
150            )
151            .await
152            .map(|_| ())
153            .map_err(|_| create_database_error!("update_one", COL))
154    }
155
156    /// Mark multiple attachments as having been deleted.
157    async fn mark_attachments_as_deleted(&self, ids: &[String]) -> Result<()> {
158        self.col::<Document>(COL)
159            .update_many(
160                doc! {
161                    "_id": {
162                        "$in": ids
163                    }
164                },
165                doc! {
166                    "$set": {
167                        "deleted": true
168                    }
169                },
170            )
171            .await
172            .map(|_| ())
173            .map_err(|_| create_database_error!("update_many", COL))
174    }
175
176    /// Delete the attachment entry.
177    async fn delete_attachment(&self, id: &str) -> Result<()> {
178        query!(self, delete_one_by_id, COL, id).map(|_| ())
179    }
180}
181
182impl MongoDb {
183    pub async fn delete_many_attachments(&self, projection: Document) -> Result<()> {
184        self.col::<Document>(COL)
185            .update_many(
186                projection,
187                doc! {
188                    "$set": {
189                        "deleted": true
190                    }
191                },
192            )
193            .await
194            .map(|_| ())
195            .map_err(|_| create_database_error!("update_many", COL))
196    }
197}