From 89f116723903549f4daa8c9b901477a430c5b8b4 Mon Sep 17 00:00:00 2001 From: Paul Makles Date: Sat, 11 Mar 2023 16:37:48 +0000 Subject: [PATCH] refactor: unify message query into one method --- .../src/routes/channels/message_query.rs | 22 +- .../src/routes/channels/message_search.rs | 20 +- .../delta/src/routes/safety/report_content.rs | 41 +-- .../quark/src/impl/dummy/channels/message.rs | 26 +- .../quark/src/impl/mongo/channels/message.rs | 255 +++++++++--------- crates/quark/src/impl/mongo/mod.rs | 31 +-- crates/quark/src/models/channels/message.rs | 52 +++- crates/quark/src/traits/channels/message.rs | 25 +- 8 files changed, 259 insertions(+), 213 deletions(-) diff --git a/crates/delta/src/routes/channels/message_query.rs b/crates/delta/src/routes/channels/message_query.rs index a6dfe267..59ff9af8 100644 --- a/crates/delta/src/routes/channels/message_query.rs +++ b/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) diff --git a/crates/delta/src/routes/channels/message_search.rs b/crates/delta/src/routes/channels/message_search.rs index 0db35568..bb229444 100644 --- a/crates/delta/src/routes/channels/message_search.rs +++ b/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, /// 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) diff --git a/crates/delta/src/routes/safety/report_content.rs b/crates/delta/src/routes/safety/report_content.rs index bdb03336..5a7060b7 100644 --- a/crates/delta/src/routes/safety/report_content.rs +++ b/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) // Collect prior context let prior_context = db - .fetch_messages( - &message.channel, - Some(15), - Some(message.id.to_string()), - None, - None, - None, - ) + .fetch_messages(MessageQuery { + filter: MessageFilter { + channel: Some(message.channel.to_string()), + ..Default::default() + }, + limit: Some(15), + 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( - &message.channel, - Some(15), - None, - Some(message.id.to_string()), - None, - None, - ) + .fetch_messages(MessageQuery { + filter: MessageFilter { + channel: Some(message.channel.to_string()), + ..Default::default() + }, + limit: Some(15), + time_period: MessageTimePeriod::Absolute { + before: None, + after: Some(message.id.to_string()), + sort: Some(MessageSort::Oldest), + }, + }) .await?; ( diff --git a/crates/quark/src/impl/dummy/channels/message.rs b/crates/quark/src/impl/dummy/channels/message.rs index a01a32a1..c88ce11f 100644 --- a/crates/quark/src/impl/dummy/channels/message.rs +++ b/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( - &self, - channel: &str, - _limit: Option, - _before: Option, - _after: Option, - _sort: Option, - _nearby: Option, - ) -> Result> { - Ok(vec![self.fetch_message(channel).await.unwrap()]) - } - - async fn search_messages( - &self, - channel: &str, - _query: &str, - _limit: Option, - _before: Option, - _after: Option, - _sort: MessageSort, - ) -> Result> { - Ok(vec![self.fetch_message(channel).await.unwrap()]) + async fn fetch_messages(&self, _query: MessageQuery) -> Result> { + Ok(vec![]) } /// Add a new reaction to a message diff --git a/crates/quark/src/impl/mongo/channels/message.rs b/crates/quark/src/impl/mongo/channels/message.rs index 0a1db531..79ae1916 100644 --- a/crates/quark/src/impl/mongo/channels/message.rs +++ b/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( - &self, - channel: &str, - limit: Option, - before: Option, - after: Option, - sort: Option, - nearby: Option, - ) -> Result> { - let limit = limit.unwrap_or(50); - Ok(if let Some(nearby) = nearby { - let (a, b) = try_join!( - self.find_with_options::<_, Message>( - COL, + async fn fetch_messages(&self, query: MessageQuery) -> Result> { + let mut filter = doc! {}; + + // 1. Apply message filters + if let Some(channel) = query.filter.channel { + filter.insert("channel", channel); + } + + if let Some(author) = query.filter.author { + filter.insert("author", author); + } + + let is_search_query = if let Some(query) = query.filter.query { + 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, - "_id": { - "$gte": &nearby - } + "$lt": &nearby }, - FindOptions::builder() - .limit(limit / 2 + 1) - .sort(doc! { - "_id": 1_i32 - }) - .build(), - ), - self.find_with_options::<_, Message>( - COL, + ); + + newer_message_filter.insert( + "_id", doc! { - "channel": channel, - "_id": { - "$lt": &nearby - } + "$gte": &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) - .sort(doc! { - "_id": -1_i32 + .limit(limit) + .sort(match sort.unwrap_or(MessageSort::Latest) { + // 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(), ) - )?; - - [a, b].concat() - } else { - let mut query = doc! { "channel": channel }; - if let Some(before) = before { - query.insert("_id", doc! { "$lt": before }); + .await } - - 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, - before: Option, - after: Option, - sort: MessageSort, - ) -> Result> { - 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 diff --git a/crates/quark/src/impl/mongo/mod.rs b/crates/quark/src/impl/mongo/mod.rs index 3e7e9796..582d3d6e 100644 --- a/crates/quark/src/impl/mongo/mod.rs +++ b/crates/quark/src/impl/mongo/mod.rs @@ -90,24 +90,25 @@ impl MongoDb { where O: Into>, { - Ok(self - .col::(collection) - .find(projection, options) - .await - .map_err(|_| Error::DatabaseError { + let result = self.col::(collection).find(projection, options).await; + Ok(if cfg!(debug_assertions) { + result.unwrap() + } else { + result.map_err(|_| Error::DatabaseError { operation: "find", with: collection, })? - .filter_map(|s| async { - if cfg!(debug_assertions) { - // Hard fail on invalid documents - Some(s.unwrap()) - } else { - s.ok() - } - }) - .collect::>() - .await) + } + .filter_map(|s| async { + if cfg!(debug_assertions) { + // Hard fail on invalid documents + Some(s.unwrap()) + } else { + s.ok() + } + }) + .collect::>() + .await) } async fn find( diff --git a/crates/quark/src/models/channels/message.rs b/crates/quark/src/models/channels/message.rs index 2eb91733..71e89f16 100644 --- a/crates/quark/src/models/channels/message.rs +++ b/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, + /// Message id after which messages should be fetched + after: Option, + /// Message sort direction + sort: Option, + }, +} + +/// # Message Filter +#[derive(Serialize, Deserialize, JsonSchema, Default)] +pub struct MessageFilter { + /// Parent channel ID + pub channel: Option, + /// Message author ID + pub author: Option, + /// Search query + pub query: Option, +} + +/// # 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, + /// 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 diff --git a/crates/quark/src/traits/channels/message.rs b/crates/quark/src/traits/channels/message.rs index 6f3f4f96..c408de73 100644 --- a/crates/quark/src/traits/channels/message.rs +++ b/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) -> Result<()>; - /// Fetch multiple messages - async fn fetch_messages( - &self, - channel: &str, - limit: Option, - before: Option, - after: Option, - sort: Option, - nearby: Option, - ) -> Result>; - - /// Search for messages - async fn search_messages( - &self, - channel: &str, - query: &str, - limit: Option, - before: Option, - after: Option, - sort: MessageSort, - ) -> Result>; + /// Fetch multiple messages by given query + async fn fetch_messages(&self, query: MessageQuery) -> Result>; /// Add a new reaction to a message async fn add_reaction(&self, id: &str, emoji: &str, user: &str) -> Result<()>;