feat: initial work on elasticsearch integration

Signed-off-by: Zomatree <me@zomatree.live>

crates/core/database/src/amqp/amqp.rs

@@ -1,7 +1,7 @@
 use std::collections::HashSet;
 
 use crate::events::rabbit::*;
-use crate::User;
+use crate::{Message, User};
 use amqprs::channel::{BasicPublishArguments, ExchangeDeclareArguments};
 use amqprs::connection::OpenConnectionArguments;
 use amqprs::{channel::Channel, connection::Connection, error::Error as AMQPError};
@@ -52,6 +52,15 @@ impl AMQP {
             .await
             .expect("Failed to declare exchange");
 
+        channel
+            .exchange_declare(
+                ExchangeDeclareArguments::new(&config.elasticsearch.exchange, "direct")
+                    .durable(true)
+                    .finish(),
+            )
+            .await
+            .expect("Failed to declare exchange");
+
         AMQP::new(connection, channel)
     }
 
@@ -316,4 +325,104 @@ impl AMQP {
             )
             .await
     }
+
+    pub async fn new_message_search(&self, message: Message) -> Result<(), AMQPError> {
+        let config = revolt_config::config().await;
+
+        let payload = to_string(&message).unwrap();
+
+        debug!(
+            "Sending new message search payload on channel {}: {}",
+            config.elasticsearch.message_queue, payload
+        );
+
+        self.channel
+            .basic_publish(
+                BasicProperties::default()
+                    .with_content_type("application/json")
+                    .with_persistence(true)
+                    .finish(),
+                payload.into(),
+                BasicPublishArguments::new(
+                    &config.elasticsearch.exchange,
+                    &config.elasticsearch.message_queue,
+                ),
+            )
+            .await
+    }
+
+    pub async fn edit_message_search(&self, message: Message) -> Result<(), AMQPError> {
+        let config = revolt_config::config().await;
+
+        let payload = to_string(&message).unwrap();
+
+        debug!(
+            "Sending edit message search payload on channel {}: {}",
+            config.elasticsearch.message_edit_queue, payload
+        );
+
+        self.channel
+            .basic_publish(
+                BasicProperties::default()
+                    .with_content_type("application/json")
+                    .with_persistence(true)
+                    .finish(),
+                payload.into(),
+                BasicPublishArguments::new(
+                    &config.elasticsearch.exchange,
+                    &config.elasticsearch.message_edit_queue,
+                ),
+            )
+            .await
+    }
+
+    pub async fn delete_message_search(&self, message_id: String) -> Result<(), AMQPError> {
+        let config = revolt_config::config().await;
+
+        let payload = to_string(&MessageDeletePayload { message_id }).unwrap();
+
+        debug!(
+            "Sending delete message search payload on channel {}: {}",
+            config.elasticsearch.message_delete_queue, payload
+        );
+
+        self.channel
+            .basic_publish(
+                BasicProperties::default()
+                    .with_content_type("application/json")
+                    .with_persistence(true)
+                    .finish(),
+                payload.into(),
+                BasicPublishArguments::new(
+                    &config.elasticsearch.exchange,
+                    &config.elasticsearch.message_delete_queue,
+                ),
+            )
+            .await
+    }
+
+    pub async fn delete_channel_search(&self, channel_id: String) -> Result<(), AMQPError> {
+        let config = revolt_config::config().await;
+
+        let payload = to_string(&ChannelDeletePayload { channel_id }).unwrap();
+
+        debug!(
+            "Sending delete channel search payload on channel {}: {}",
+            config.elasticsearch.channel_delete_queue, payload
+        );
+
+        self.channel
+            .basic_publish(
+                BasicProperties::default()
+                    .with_content_type("application/json")
+                    .with_persistence(true)
+                    .finish(),
+                payload.into(),
+                BasicPublishArguments::new(
+                    &config.elasticsearch.exchange,
+                    &config.elasticsearch.channel_delete_queue,
+                ),
+            )
+            .await
+    }
 }

crates/core/database/src/events/rabbit.rs

@@ -78,3 +78,13 @@ pub struct AckPayload {
     pub channel_id: String,
     pub message_id: String,
 }
+
+#[derive(Serialize, Deserialize)]
+pub struct MessageDeletePayload {
+    pub message_id: String
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct ChannelDeletePayload {
+    pub channel_id: String
+}

crates/core/database/src/models/channels/model.rs

@@ -1,7 +1,7 @@
 #![allow(deprecated)]
 use std::{borrow::Cow, collections::HashMap};
 
-use revolt_config::config;
+use revolt_config::{capture_error, config};
 use revolt_models::v0::{self, MessageAuthor};
 use revolt_permissions::OverrideField;
 use revolt_result::Result;
@@ -331,7 +331,7 @@ impl Channel {
     pub async fn add_user_to_group(
         &mut self,
         db: &Database,
-        amqp: &AMQP,
+        amqp: Option<&AMQP>,
         user: &User,
         by_id: &str,
     ) -> Result<()> {
@@ -368,7 +368,7 @@ impl Channel {
                 .into_message(id.to_string())
                 .send(
                     db,
-                    Some(amqp),
+                    amqp,
                     MessageAuthor::System {
                         username: &user.username,
                         avatar: user.avatar.as_ref().map(|file| file.id.as_ref()),
@@ -661,7 +661,7 @@ impl Channel {
     pub async fn remove_user_from_group(
         &self,
         db: &Database,
-        amqp: &AMQP,
+        amqp: Option<&AMQP>,
         user: &User,
         by_id: Option<&str>,
         silent: bool,
@@ -693,7 +693,7 @@ impl Channel {
                         .into_message(id.to_string())
                         .send(
                             db,
-                            Some(amqp),
+                            amqp,
                             MessageAuthor::System {
                                 username: name,
                                 avatar: None,
@@ -706,7 +706,7 @@ impl Channel {
                         .await
                         .ok();
                     } else {
-                        return self.delete(db).await;
+                        return self.delete(db, amqp).await;
                     }
                 }
 
@@ -733,7 +733,7 @@ impl Channel {
                     .into_message(id.to_string())
                     .send(
                         db,
-                        Some(amqp),
+                        amqp,
                         MessageAuthor::System {
                             username: &user.username,
                             avatar: user.avatar.as_ref().map(|file| file.id.as_ref()),
@@ -755,13 +755,22 @@ impl Channel {
     }
 
     /// Delete a channel
-    pub async fn delete(&self, db: &Database) -> Result<()> {
+    pub async fn delete(&self, db: &Database, amqp: Option<&AMQP>) -> Result<()> {
         let id = self.id().to_string();
         EventV1::ChannelDelete { id: id.clone() }.p(id).await;
         // TODO: missing functionality:
         // - group invites
         // - channels list / categories list on server
-        db.delete_channel(self).await
+        db.delete_channel(self).await?;
+
+        if let Some(amqp) = amqp {
+            if let Err(e) = amqp.delete_channel_search(self.id().to_string()).await {
+                log::error!("Error pushing message to RabbitMQ: {e}");
+                capture_error(&e);
+            }
+        }
+
+        Ok(())
     }
 }
 

crates/core/database/src/models/messages/model.rs

@@ -2,7 +2,7 @@ use std::{collections::HashSet, hash::RandomState};
 
 use indexmap::{IndexMap, IndexSet};
 use iso8601_timestamp::Timestamp;
-use revolt_config::{config, FeaturesLimits};
+use revolt_config::{FeaturesLimits, capture_error, config};
 use revolt_models::v0::{
     self, BulkMessageResponse, DataMessageSend, Embed, MessageAuthor, MessageFlags, MessageSort,
     MessageWebhook, PushNotification, ReplyIntent, SendableEmbed, Text,
@@ -670,7 +670,7 @@ impl Message {
     pub async fn send(
         &mut self,
         db: &Database,
-        _amqp: Option<&AMQP>, // this is optional mostly for tests.
+        amqp: Option<&AMQP>, // this is optional mostly for tests.
         author: MessageAuthor<'_>,
         user: Option<v0::User>,
         member: Option<v0::Member>,
@@ -720,6 +720,13 @@ impl Message {
             .await;
         }
 
+        if let Some(amqp) = amqp {
+            if let Err(e) = amqp.new_message_search(self.clone()).await {
+                log::error!("Error pushing message to RabbitMQ: {e}");
+                capture_error(&e);
+            }
+        }
+
         Ok(())
     }
 
@@ -988,9 +995,10 @@ impl Message {
     }
 
     /// Delete a message
-    pub async fn delete(self, db: &Database) -> Result<()> {
+    pub async fn delete(&self, db: &Database, amqp: Option<&AMQP>) -> Result<()> {
         let file_ids: Vec<String> = self
             .attachments
+            .as_ref()
             .map(|files| files.iter().map(|file| file.id.to_string()).collect())
             .unwrap_or_default();
 
@@ -1000,12 +1008,20 @@ impl Message {
 
         db.delete_message(&self.id).await?;
 
+        if let Some(amqp) = amqp {
+            if let Err(e) = amqp.delete_message_search(self.id.clone()).await {
+                log::error!("Error pushing message to RabbitMQ: {e}");
+                capture_error(&e);
+            }
+        }
+
         EventV1::MessageDelete {
-            id: self.id,
+            id: self.id.clone(),
             channel: self.channel.clone(),
         }
-        .p(self.channel)
+        .p(self.channel.clone())
         .await;
+
         Ok(())
     }
 

crates/core/database/src/models/messages/ops.rs

@@ -1,6 +1,6 @@
 use revolt_result::Result;
 
-use crate::{AppendMessage, FieldsMessage, Message, MessageQuery, PartialMessage};
+use crate::{AppendMessage, FieldsMessage, Message, MessageQuery, PartialMessage, util::ChunkedDatabaseGenerator};
 
 #[cfg(feature = "mongodb")]
 mod mongodb;
@@ -40,4 +40,6 @@ pub trait AbstractMessages: Sync + Send {
 
     /// Delete messages from a channel by their ids and corresponding channel id
     async fn delete_messages(&self, channel: &str, ids: &[String]) -> Result<()>;
+
+    async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<Message>>;
 }

crates/core/database/src/models/messages/ops/mongodb.rs

@@ -1,12 +1,12 @@
 use bson::{to_bson, Document};
 use futures::try_join;
-use mongodb::options::FindOptions;
+use mongodb::options::{FindOptions, ReadConcern};
 use revolt_models::v0::MessageSort;
 use revolt_result::Result;
 
 use crate::{
     AppendMessage, DocumentId, FieldsMessage, IntoDocumentPath, Message, MessageQuery,
-    MessageTimePeriod, MongoDb, PartialMessage,
+    MessageTimePeriod, MongoDb, PartialMessage, util::ChunkedDatabaseGenerator,
 };
 
 use super::AbstractMessages;
@@ -306,6 +306,31 @@ impl AbstractMessages for MongoDb {
             .map(|_| ())
             .map_err(|_| create_database_error!("delete_many", COL))
     }
+
+    async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<Message>> {
+        let mut session = self
+            .start_session()
+            .await
+            .map_err(|_| create_database_error!("start_session", COL))?;
+
+        session
+            .start_transaction()
+            .read_concern(ReadConcern::snapshot())
+            .await
+            .map_err(|_| create_database_error!("start_transaction", COL))?;
+
+        let cursor = self
+            .col::<Message>(COL)
+            .find(doc! {})
+            .sort(doc ! { "_id": -1i32 })
+            .session(&mut session)
+            .batch_size(1000)
+            .await
+            .inspect_err(|e| log::error!("{e}"))
+            .map_err(|_| create_database_error!("find", COL))?;
+
+        Ok(ChunkedDatabaseGenerator::new_mongo(session, cursor))
+    }
 }
 
 impl IntoDocumentPath for FieldsMessage {

crates/core/database/src/models/messages/ops/reference.rs

@@ -2,7 +2,10 @@ use futures::future::try_join_all;
 use indexmap::IndexSet;
 use revolt_result::Result;
 
-use crate::{AppendMessage, FieldsMessage, Message, MessageQuery, PartialMessage, ReferenceDb};
+use crate::{
+    util::ChunkedDatabaseGenerator, AppendMessage, FieldsMessage, Message, MessageQuery,
+    PartialMessage, ReferenceDb,
+};
 
 use super::AbstractMessages;
 
@@ -286,4 +289,10 @@ impl AbstractMessages for ReferenceDb {
 
         Ok(())
     }
+
+    async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<Message>> {
+        Ok(ChunkedDatabaseGenerator::new_reference(
+            self.messages.lock().await.values().cloned().collect(),
+        ))
+    }
 }

crates/core/database/src/models/servers/model.rs

@@ -1,11 +1,12 @@
 use std::collections::{HashMap, HashSet};
 
+use revolt_config::capture_error;
 use revolt_models::v0::{self, DataCreateServerChannel};
 use revolt_permissions::{OverrideField, DEFAULT_PERMISSION_SERVER};
 use revolt_result::Result;
 use ulid::Ulid;
 
-use crate::{events::client::EventV1, Channel, Database, File, User};
+use crate::{AMQP, Channel, Database, File, User, events::client::EventV1};
 
 auto_derived_partial!(
     /// Server
@@ -210,7 +211,16 @@ impl Server {
     }
 
     /// Delete a server
-    pub async fn delete(self, db: &Database) -> Result<()> {
+    pub async fn delete(self, db: &Database, amqp: Option<&AMQP>) -> Result<()> {
+        if let Some(amqp) = amqp {
+            for channel_id in self.channels {
+                if let Err(e) = amqp.delete_channel_search(channel_id).await {
+                    log::error!("Error pushing message to RabbitMQ: {e}");
+                    capture_error(&e);
+                }
+            }
+        }
+
         EventV1::ServerDelete {
             id: self.id.clone(),
         }

crates/core/database/src/models/users/model.rs

@@ -510,7 +510,7 @@ impl User {
     pub async fn add_friend(
         &mut self,
         db: &Database,
-        amqp: &AMQP,
+        amqp: Option<&AMQP>,
         target: &mut User,
     ) -> Result<()> {
         match self.relationship_with(&target.id) {
@@ -520,8 +520,10 @@ impl User {
             RelationshipStatus::Blocked => Err(create_error!(Blocked)),
             RelationshipStatus::BlockedOther => Err(create_error!(BlockedByOther)),
             RelationshipStatus::Incoming => {
-                // Accept incoming friend request
-                _ = amqp.friend_request_accepted(self, target).await;
+                if let Some(amqp) = amqp {
+                    // Accept incoming friend request
+                    _ = amqp.friend_request_accepted(self, target).await;
+                };
 
                 self.apply_relationship(
                     db,
@@ -551,7 +553,9 @@ impl User {
                     }));
                 }
 
-                _ = amqp.friend_request_received(target, self).await;
+                if let Some(amqp) = amqp {
+                    _ = amqp.friend_request_received(target, self).await;
+                };
 
                 // Send the friend request
                 self.apply_relationship(

crates/core/database/src/util/chunked.rs

@@ -0,0 +1,54 @@
+#[cfg(feature = "mongodb")]
+use ::mongodb::{ClientSession, SessionCursor};
+use serde::Deserialize;
+
+#[derive(Debug)]
+#[allow(clippy::large_enum_variant)]
+pub enum ChunkedDatabaseGenerator<T> {
+    #[cfg(feature = "mongodb")]
+    MongoDb {
+        session: ClientSession,
+        cursor: SessionCursor<T>,
+    },
+
+    Reference {
+        offset: i32,
+        data: Vec<T>,
+    },
+}
+
+impl<T: for<'d> Deserialize<'d> + Clone> ChunkedDatabaseGenerator<T> {
+    #[cfg(feature = "mongodb")]
+    pub fn new_mongo(session: ClientSession, cursor: SessionCursor<T>) -> Self {
+        Self::MongoDb {
+            session,
+            cursor,
+        }
+    }
+
+    pub fn new_reference(data: Vec<T>) -> Self {
+        Self::Reference {
+            offset: 0,
+            data,
+        }
+    }
+
+    pub async fn next(&mut self) -> Option<T> {
+        match self {
+            #[cfg(feature = "mongodb")]
+            Self::MongoDb { session, cursor } => {
+                let value = cursor.next(session).await;
+                value.map(|val| val.expect("Failed to fetch the next message"))
+            }
+            Self::Reference { offset, data } => {
+                if data.len() as i32 >= *offset {
+                    None
+                } else {
+                    let resp = &data[*offset as usize];
+                    *offset += 1;
+                    Some(resp.clone())
+                }
+            }
+        }
+    }
+}

crates/core/database/src/util/mod.rs

@@ -5,5 +5,7 @@ pub mod idempotency;
 pub mod permissions;
 pub mod reference;
 pub mod test_fixtures;
+pub mod chunked;
 
 pub use funcs::*;
+pub use chunked::ChunkedDatabaseGenerator;