revolt_database/models/safety_snapshots/
model.rs

1use revolt_models::v0::MessageSort;
2use revolt_result::Result;
3
4use crate::{Database, Message, MessageFilter, MessageQuery, MessageTimePeriod, Server, User};
5
6auto_derived!(
7    /// Snapshot of some content
8    pub struct Snapshot {
9        /// Unique Id
10        #[serde(rename = "_id")]
11        pub id: String,
12        /// Report parent Id
13        pub report_id: String,
14        /// Snapshot of content
15        pub content: SnapshotContent,
16    }
17
18    /// Enum to map into different models
19    /// that can be saved in a snapshot
20    #[serde(tag = "_type")]
21    pub enum SnapshotContent {
22        Message {
23            /// Context before the message
24            #[serde(rename = "_prior_context", default)]
25            prior_context: Vec<Message>,
26
27            /// Context after the message
28            #[serde(rename = "_leading_context", default)]
29            leading_context: Vec<Message>,
30
31            /// Message
32            #[serde(flatten)]
33            message: Message,
34        },
35        Server(Server),
36        User(User),
37    }
38);
39
40impl SnapshotContent {
41    /// Generate snapshot from a given message
42    pub async fn generate_from_message(
43        db: &Database,
44        message: Message,
45    ) -> Result<(SnapshotContent, Vec<String>)> {
46        // Collect message attachments
47        let files = message
48            .attachments
49            .as_ref()
50            .map(|attachments| attachments.iter().map(|x| x.id.to_string()).collect())
51            .unwrap_or_default();
52
53        // Collect prior context
54        let prior_context = db
55            .fetch_messages(MessageQuery {
56                filter: MessageFilter {
57                    channel: Some(message.channel.to_string()),
58                    ..Default::default()
59                },
60                limit: Some(15),
61                time_period: MessageTimePeriod::Absolute {
62                    before: Some(message.id.to_string()),
63                    after: None,
64                    sort: Some(MessageSort::Latest),
65                },
66            })
67            .await?;
68
69        // Collect leading context
70        let leading_context = db
71            .fetch_messages(MessageQuery {
72                filter: MessageFilter {
73                    channel: Some(message.channel.to_string()),
74                    ..Default::default()
75                },
76                limit: Some(15),
77                time_period: MessageTimePeriod::Absolute {
78                    before: None,
79                    after: Some(message.id.to_string()),
80                    sort: Some(MessageSort::Oldest),
81                },
82            })
83            .await?;
84
85        Ok((
86            SnapshotContent::Message {
87                message,
88                prior_context: prior_context.into_iter().collect(),
89                leading_context: leading_context.into_iter().collect(),
90            },
91            files,
92        ))
93    }
94
95    /// Generate snapshot from a given server
96    pub fn generate_from_server(server: Server) -> Result<(SnapshotContent, Vec<String>)> {
97        // Collect server's icon and banner
98        let files = [&server.icon, &server.banner]
99            .iter()
100            .filter_map(|x| x.as_ref().map(|x| x.id.to_string()))
101            .collect();
102
103        Ok((SnapshotContent::Server(server), files))
104    }
105
106    /// Generate snapshot from a given user
107    pub fn generate_from_user(user: User) -> Result<(SnapshotContent, Vec<String>)> {
108        // Collect user's avatar and profile background
109        let files = [
110            user.avatar.as_ref(),
111            user.profile
112                .as_ref()
113                .and_then(|profile| profile.background.as_ref()),
114        ]
115        .iter()
116        .filter_map(|x| x.as_ref().map(|x| x.id.to_string()))
117        .collect();
118
119        Ok((SnapshotContent::User(user), files))
120    }
121}