refactor: unify message query into one method

crates/delta/src/routes/channels/message_query.rs

@@ -1,6 +1,8 @@
 use revolt_quark::{
     models::{
-        message::{BulkMessageResponse, MessageSort},
+        message::{
+            BulkMessageResponse, MessageFilter, MessageQuery, MessageSort, MessageTimePeriod,
+        },
         User,
     },
     perms, Db, Error, Permission, Ref, Result,
@@ -68,11 +70,25 @@ pub async fn req(
         sort,
         nearby,
         include_users,
-        ..
     } = options;
 
     let messages = db
-        .fetch_messages(channel.id(), limit, before, after, sort, nearby)
+        .fetch_messages(MessageQuery {
+            filter: MessageFilter {
+                channel: Some(channel.id().to_string()),
+                ..Default::default()
+            },
+            time_period: if let Some(nearby) = nearby {
+                MessageTimePeriod::Relative { nearby }
+            } else {
+                MessageTimePeriod::Absolute {
+                    before,
+                    after,
+                    sort,
+                }
+            },
+            limit,
+        })
         .await?;
 
     BulkMessageResponse::transform(db, &channel, messages, include_users)

crates/delta/src/routes/channels/message_search.rs

@@ -1,6 +1,8 @@
 use revolt_quark::{
     models::{
-        message::{BulkMessageResponse, MessageSort},
+        message::{
+            BulkMessageResponse, MessageFilter, MessageQuery, MessageSort, MessageTimePeriod,
+        },
         User,
     },
     perms, Db, Error, Permission, Ref, Result,
@@ -30,7 +32,7 @@ pub struct OptionsMessageSearch {
     after: Option<String>,
     /// Message sort direction
     ///
-    /// By default, it will be sorted by relevance.
+    /// By default, it will be sorted by latest.
     #[serde(default = "MessageSort::default")]
     sort: MessageSort,
     /// Whether to include user (and member, if server channel) objects
@@ -73,7 +75,19 @@ pub async fn req(
     } = options;
 
     let messages = db
-        .search_messages(channel.id(), &query, limit, before, after, sort)
+        .fetch_messages(MessageQuery {
+            filter: MessageFilter {
+                channel: Some(channel.id().to_string()),
+                query: Some(query),
+                ..Default::default()
+            },
+            time_period: MessageTimePeriod::Absolute {
+                before,
+                after,
+                sort: Some(sort),
+            },
+            limit,
+        })
         .await?;
 
     BulkMessageResponse::transform(db, &channel, messages, include_users)

crates/delta/src/routes/safety/report_content.rs

@@ -1,4 +1,5 @@
 use revolt_quark::events::client::EventV1;
+use revolt_quark::models::message::{MessageFilter, MessageQuery, MessageSort, MessageTimePeriod};
 use revolt_quark::models::report::{ReportStatus, ReportedContent};
 use revolt_quark::models::snapshot::{Snapshot, SnapshotContent};
 use revolt_quark::models::{Report, User};
@@ -55,26 +56,34 @@ pub async fn report_content(db: &Db, user: User, data: Json<DataReportContent>)
 
             // Collect prior context
             let prior_context = db
-                .fetch_messages(
+                .fetch_messages(MessageQuery {
-                    &message.channel,
+                    filter: MessageFilter {
-                    Some(15),
+                        channel: Some(message.channel.to_string()),
-                    Some(message.id.to_string()),
+                        ..Default::default()
-                    None,
+                    },
-                    None,
+                    limit: Some(15),
-                    None,
+                    time_period: MessageTimePeriod::Absolute {
-                )
+                        before: Some(message.id.to_string()),
+                        after: None,
+                        sort: Some(MessageSort::Latest),
+                    },
+                })
                 .await?;
 
             // Collect leading context
             let leading_context = db
-                .fetch_messages(
+                .fetch_messages(MessageQuery {
-                    &message.channel,
+                    filter: MessageFilter {
-                    Some(15),
+                        channel: Some(message.channel.to_string()),
-                    None,
+                        ..Default::default()
-                    Some(message.id.to_string()),
+                    },
-                    None,
+                    limit: Some(15),
-                    None,
+                    time_period: MessageTimePeriod::Absolute {
-                )
+                        before: None,
+                        after: Some(message.id.to_string()),
+                        sort: Some(MessageSort::Oldest),
+                    },
+                })
                 .await?;
 
             (

crates/quark/src/impl/dummy/channels/message.rs

@@ -1,4 +1,4 @@
-use crate::models::message::{AppendMessage, Message, MessageSort, PartialMessage};
+use crate::models::message::{AppendMessage, Message, MessageQuery, PartialMessage};
 use crate::{AbstractMessage, Result};
 
 use super::super::DummyDb;
@@ -41,28 +41,8 @@ impl AbstractMessage for DummyDb {
         Ok(())
     }
 
-    async fn fetch_messages(
+    async fn fetch_messages(&self, _query: MessageQuery) -> Result<Vec<Message>> {
-        &self,
+        Ok(vec![])
-        channel: &str,
-        _limit: Option<i64>,
-        _before: Option<String>,
-        _after: Option<String>,
-        _sort: Option<MessageSort>,
-        _nearby: Option<String>,
-    ) -> Result<Vec<Message>> {
-        Ok(vec![self.fetch_message(channel).await.unwrap()])
-    }
-
-    async fn search_messages(
-        &self,
-        channel: &str,
-        _query: &str,
-        _limit: Option<i64>,
-        _before: Option<String>,
-        _after: Option<String>,
-        _sort: MessageSort,
-    ) -> Result<Vec<Message>> {
-        Ok(vec![self.fetch_message(channel).await.unwrap()])
     }
 
     /// Add a new reaction to a message

crates/quark/src/impl/mongo/channels/message.rs

@@ -2,7 +2,9 @@ use bson::{to_bson, Document};
 use futures::try_join;
 use mongodb::options::FindOptions;
 
-use crate::models::message::{AppendMessage, Message, MessageSort, PartialMessage};
+use crate::models::message::{
+    AppendMessage, Message, MessageQuery, MessageSort, MessageTimePeriod, PartialMessage,
+};
 use crate::r#impl::mongo::DocumentId;
 use crate::{AbstractMessage, Error, Result};
 
@@ -139,145 +141,138 @@ impl AbstractMessage for MongoDb {
         .await
     }
 
-    async fn fetch_messages(
+    async fn fetch_messages(&self, query: MessageQuery) -> Result<Vec<Message>> {
-        &self,
+        let mut filter = doc! {};
-        channel: &str,
+
-        limit: Option<i64>,
+        // 1. Apply message filters
-        before: Option<String>,
+        if let Some(channel) = query.filter.channel {
-        after: Option<String>,
+            filter.insert("channel", channel);
-        sort: Option<MessageSort>,
+        }
-        nearby: Option<String>,
+
-    ) -> Result<Vec<Message>> {
+        if let Some(author) = query.filter.author {
-        let limit = limit.unwrap_or(50);
+            filter.insert("author", author);
-        Ok(if let Some(nearby) = nearby {
+        }
-            let (a, b) = try_join!(
+
-                self.find_with_options::<_, Message>(
+        let is_search_query = if let Some(query) = query.filter.query {
-                    COL,
+            filter.insert(
+                "$text",
+                doc! {
+                    "$search": query
+                },
+            );
+
+            true
+        } else {
+            false
+        };
+
+        // 2. Find query limit
+        let limit = query.limit.unwrap_or(50);
+
+        // 3. Apply message time period
+        match query.time_period {
+            MessageTimePeriod::Relative { nearby } => {
+                // 3.1. Prepare filters
+                let mut older_message_filter = filter.clone();
+                let mut newer_message_filter = filter;
+
+                older_message_filter.insert(
+                    "_id",
                     doc! {
-                        "channel": channel,
+                        "$lt": &nearby
-                        "_id": {
-                            "$gte": &nearby
-                        }
                     },
-                    FindOptions::builder()
+                );
-                        .limit(limit / 2 + 1)
+
-                        .sort(doc! {
+                newer_message_filter.insert(
-                            "_id": 1_i32
+                    "_id",
-                        })
-                        .build(),
-                ),
-                self.find_with_options::<_, Message>(
-                    COL,
                     doc! {
-                        "channel": channel,
+                        "$gte": &nearby
-                        "_id": {
-                            "$lt": &nearby
-                        }
                     },
+                );
+
+                // 3.2. Execute in both directions
+                let (a, b) = try_join!(
+                    self.find_with_options::<_, Message>(
+                        COL,
+                        newer_message_filter,
+                        FindOptions::builder()
+                            .limit(limit / 2 + 1)
+                            .sort(doc! {
+                                "_id": 1_i32
+                            })
+                            .build(),
+                    ),
+                    self.find_with_options::<_, Message>(
+                        COL,
+                        older_message_filter,
+                        FindOptions::builder()
+                            .limit(limit / 2)
+                            .sort(doc! {
+                                "_id": -1_i32
+                            })
+                            .build(),
+                    )
+                )?;
+
+                Ok([a, b].concat())
+            }
+            MessageTimePeriod::Absolute {
+                before,
+                after,
+                sort,
+            } => {
+                // 3.1. Apply message ID filter
+                if let Some(doc) = match (before, after) {
+                    (Some(before), Some(after)) => Some(doc! {
+                        "$lt": before,
+                        "$gt": after
+                    }),
+                    (Some(before), _) => Some(doc! {
+                        "$lt": before
+                    }),
+                    (_, Some(after)) => Some(doc! {
+                        "$gt": after
+                    }),
+                    _ => None,
+                } {
+                    filter.insert("_id", doc);
+                }
+
+                // 3.2. Execute with given message sort
+                self.find_with_options(
+                    COL,
+                    filter,
                     FindOptions::builder()
-                        .limit(limit / 2)
+                        .limit(limit)
-                        .sort(doc! {
+                        .sort(match sort.unwrap_or(MessageSort::Latest) {
-                            "_id": -1_i32
+                            // Sort by relevance, fallback to latest
+                            MessageSort::Relevance => {
+                                if is_search_query {
+                                    doc! {
+                                        "score": {
+                                            "$meta": "textScore"
+                                        }
+                                    }
+                                } else {
+                                    doc! {
+                                        "_id": -1_i32
+                                    }
+                                }
+                            }
+                            // Sort by latest first
+                            MessageSort::Latest => doc! {
+                                "_id": -1_i32
+                            },
+                            // Sort by oldest first
+                            MessageSort::Oldest => doc! {
+                                "_id": 1_i32
+                            },
                         })
                         .build(),
                 )
-            )?;
+                .await
-
-            [a, b].concat()
-        } else {
-            let mut query = doc! { "channel": channel };
-            if let Some(before) = before {
-                query.insert("_id", doc! { "$lt": before });
             }
-
-            if let Some(after) = after {
-                query.insert("_id", doc! { "$gt": after });
-            }
-
-            let sort: i32 = if let MessageSort::Latest = sort.unwrap_or(MessageSort::Latest) {
-                -1
-            } else {
-                1
-            };
-
-            self.find_with_options::<_, Message>(
-                COL,
-                query,
-                FindOptions::builder()
-                    .limit(limit)
-                    .sort(doc! {
-                        "_id": sort
-                    })
-                    .build(),
-            )
-            .await?
-        })
-    }
-
-    async fn search_messages(
-        &self,
-        channel: &str,
-        query: &str,
-        limit: Option<i64>,
-        before: Option<String>,
-        after: Option<String>,
-        sort: MessageSort,
-    ) -> Result<Vec<Message>> {
-        let limit = limit.unwrap_or(50);
-
-        let mut filter = doc! {
-            "channel": channel,
-            "$text": {
-                "$search": query
-            }
-        };
-
-        if let Some(doc) = match (before, after) {
-            (Some(before), Some(after)) => Some(doc! {
-                "lt": before,
-                "gt": after
-            }),
-            (Some(before), _) => Some(doc! {
-                "lt": before
-            }),
-            (_, Some(after)) => Some(doc! {
-                "gt": after
-            }),
-            _ => None,
-        } {
-            filter.insert("_id", doc);
         }
-
-        self.find_with_options(
-            COL,
-            filter,
-            FindOptions::builder()
-                .projection(if let MessageSort::Relevance = &sort {
-                    doc! {
-                        "score": {
-                            "$meta": "textScore"
-                        }
-                    }
-                } else {
-                    doc! {}
-                })
-                .limit(limit)
-                .sort(match &sort {
-                    MessageSort::Relevance => doc! {
-                        "score": {
-                            "$meta": "textScore"
-                        }
-                    },
-                    MessageSort::Latest => doc! {
-                        "_id": -1_i32
-                    },
-                    MessageSort::Oldest => doc! {
-                        "_id": 1_i32
-                    },
-                })
-                .build(),
-        )
-        .await
     }
 
     /// Add a new reaction to a message

crates/quark/src/impl/mongo/mod.rs

@@ -90,24 +90,25 @@ impl MongoDb {
     where
         O: Into<Option<FindOptions>>,
     {
-        Ok(self
+        let result = self.col::<T>(collection).find(projection, options).await;
-            .col::<T>(collection)
+        Ok(if cfg!(debug_assertions) {
-            .find(projection, options)
+            result.unwrap()
-            .await
+        } else {
-            .map_err(|_| Error::DatabaseError {
+            result.map_err(|_| Error::DatabaseError {
                 operation: "find",
                 with: collection,
             })?
-            .filter_map(|s| async {
+        }
-                if cfg!(debug_assertions) {
+        .filter_map(|s| async {
-                    // Hard fail on invalid documents
+            if cfg!(debug_assertions) {
-                    Some(s.unwrap())
+                // Hard fail on invalid documents
-                } else {
+                Some(s.unwrap())
-                    s.ok()
+            } else {
-                }
+                s.ok()
-            })
+            }
-            .collect::<Vec<T>>()
+        })
-            .await)
+        .collect::<Vec<T>>()
+        .await)
     }
 
     async fn find<T: DeserializeOwned + Unpin + Send + Sync>(

crates/quark/src/models/channels/message.rs

@@ -157,7 +157,7 @@ pub struct Message {
 /// # Message Sort
 ///
 /// Sort used for retrieving messages
-#[derive(Serialize, Deserialize, JsonSchema)]
+#[derive(Serialize, Deserialize, JsonSchema, Debug)]
 #[cfg_attr(feature = "rocket_impl", derive(FromFormField))]
 pub enum MessageSort {
     /// Sort by the most relevant messages
@@ -174,6 +174,56 @@ impl Default for MessageSort {
     }
 }
 
+/// # Message Time Period
+///
+/// Filter and sort messages by time
+#[derive(Serialize, Deserialize, JsonSchema)]
+#[serde(untagged)]
+pub enum MessageTimePeriod {
+    Relative {
+        /// Message id to search around
+        ///
+        /// Specifying 'nearby' ignores 'before', 'after' and 'sort'.
+        /// It will also take half of limit rounded as the limits to each side.
+        /// It also fetches the message ID specified.
+        nearby: String,
+    },
+    Absolute {
+        /// Message id before which messages should be fetched
+        before: Option<String>,
+        /// Message id after which messages should be fetched
+        after: Option<String>,
+        /// Message sort direction
+        sort: Option<MessageSort>,
+    },
+}
+
+/// # Message Filter
+#[derive(Serialize, Deserialize, JsonSchema, Default)]
+pub struct MessageFilter {
+    /// Parent channel ID
+    pub channel: Option<String>,
+    /// Message author ID
+    pub author: Option<String>,
+    /// Search query
+    pub query: Option<String>,
+}
+
+/// # Message Query
+#[derive(Serialize, Deserialize, JsonSchema)]
+pub struct MessageQuery {
+    /// Maximum number of messages to fetch
+    ///
+    /// For fetching nearby messages, this is \`(limit + 1)\`.
+    pub limit: Option<i64>,
+    /// Filter to apply
+    #[serde(flatten)]
+    pub filter: MessageFilter,
+    /// Time period to fetch
+    #[serde(flatten)]
+    pub time_period: MessageTimePeriod,
+}
+
 /// # Bulk Message Response
 ///
 /// Response used when multiple messages are fetched

crates/quark/src/traits/channels/message.rs

@@ -1,4 +1,4 @@
-use crate::models::message::{AppendMessage, Message, MessageSort, PartialMessage};
+use crate::models::message::{AppendMessage, Message, MessageQuery, PartialMessage};
 use crate::Result;
 
 #[async_trait]
@@ -21,27 +21,8 @@ pub trait AbstractMessage: Sync + Send {
     /// Delete messages from a channel by their ids and corresponding channel id
     async fn delete_messages(&self, channel: &str, ids: Vec<String>) -> Result<()>;
 
-    /// Fetch multiple messages
+    /// Fetch multiple messages by given query
-    async fn fetch_messages(
+    async fn fetch_messages(&self, query: MessageQuery) -> Result<Vec<Message>>;
-        &self,
-        channel: &str,
-        limit: Option<i64>,
-        before: Option<String>,
-        after: Option<String>,
-        sort: Option<MessageSort>,
-        nearby: Option<String>,
-    ) -> Result<Vec<Message>>;
-
-    /// Search for messages
-    async fn search_messages(
-        &self,
-        channel: &str,
-        query: &str,
-        limit: Option<i64>,
-        before: Option<String>,
-        after: Option<String>,
-        sort: MessageSort,
-    ) -> Result<Vec<Message>>;
 
     /// Add a new reaction to a message
     async fn add_reaction(&self, id: &str, emoji: &str, user: &str) -> Result<()>;