feat: search by message components

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

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

@@ -79,12 +79,12 @@ pub struct AckPayload {
     pub message_id: String,
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct MessageDeletePayload {
     pub message_id: String
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ChannelDeletePayload {
     pub channel_id: String
 }

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

@@ -891,19 +891,27 @@ impl Message {
     /// Append content to message
     pub async fn append(
         db: &Database,
+        amqp: Option<&AMQP>,
         id: String,
         channel: String,
         append: AppendMessage,
     ) -> Result<()> {
-        db.append_message(&id, &append).await?;
+        if let Some(message) = db.append_message(&id, &append).await? {
+            if let Some(amqp) = amqp {
+                if let Err(e) = amqp.edit_message_search(message).await {
+                    log::error!("Error pushing message to RabbitMQ: {e}");
+                    capture_error(&e);
+                }
+            }
 
-        EventV1::MessageAppend {
+            EventV1::MessageAppend {
-            id,
+                id,
-            channel: channel.to_string(),
+                channel: channel.to_string(),
-            append: append.into(),
+                append: append.into(),
+            }
+            .p(channel)
+            .await;
         }
-        .p(channel)
-        .await;
 
         Ok(())
     }

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

@@ -26,7 +26,7 @@ pub trait AbstractMessages: Sync + Send {
     async fn update_message(&self, id: &str, message: &PartialMessage, remove: Vec<FieldsMessage>) -> Result<()>;
 
     /// Append information to a given message
-    async fn append_message(&self, id: &str, append: &AppendMessage) -> Result<()>;
+    async fn append_message(&self, id: &str, append: &AppendMessage) -> Result<Option<Message>>;
 
     /// Add a new reaction to a message
     async fn add_reaction(&self, id: &str, emoji: &str, user: &str) -> Result<()>;

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

@@ -1,7 +1,7 @@
 use bson::{to_bson, Document};
 use futures::try_join;
 use futures::StreamExt;
-use mongodb::options::{FindOptions, ReadConcern};
+use mongodb::options::{FindOptions, ReadConcern, ReturnDocument};
 use revolt_models::v0::MessageSort;
 use revolt_result::Result;
 use std::collections::{HashMap, HashSet};
@@ -205,7 +205,7 @@ impl AbstractMessages for MongoDb {
     }
 
     /// Append information to a given message
-    async fn append_message(&self, id: &str, append: &AppendMessage) -> Result<()> {
+    async fn append_message(&self, id: &str, append: &AppendMessage) -> Result<Option<Message>> {
         let mut query = doc! {};
 
         if let Some(embeds) = &append.embeds {
@@ -223,18 +223,18 @@ impl AbstractMessages for MongoDb {
         }
 
         if query.is_empty() {
-            return Ok(());
+            return Ok(None);
         }
 
-        self.col::<Document>(COL)
+        self.col::<Message>(COL)
-            .update_one(
+            .find_one_and_update(
                 doc! {
                     "_id": id
                 },
                 query,
             )
+            .return_document(ReturnDocument::After)
             .await
-            .map(|_| ())
             .map_err(|_| create_database_error!("update_one", COL))
     }
 

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

@@ -211,7 +211,7 @@ impl AbstractMessages for ReferenceDb {
     }
 
     /// Append information to a given message
-    async fn append_message(&self, id: &str, append: &AppendMessage) -> Result<()> {
+    async fn append_message(&self, id: &str, append: &AppendMessage) -> Result<Option<Message>> {
         let mut messages = self.messages.lock().await;
         if let Some(message_data) = messages.get_mut(id) {
             if let Some(embeds) = &append.embeds {
@@ -222,9 +222,11 @@ impl AbstractMessages for ReferenceDb {
                         message_data.embeds = Some(embeds.clone());
                     }
                 }
-            }
 
-            Ok(())
+                Ok(Some(message_data.clone()))
+            } else {
+                Ok(None)
+            }
         } else {
             Err(create_error!(NotFound))
         }

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

@@ -19,7 +19,7 @@ pub fn start_workers(db: Database, amqp: AMQP) {
     for _ in 0..WORKER_COUNT {
         task::spawn(ack::worker(db.clone(), amqp.clone()));
         task::spawn(last_message_id::worker(db.clone()));
-        task::spawn(process_embeds::worker(db.clone()));
+        task::spawn(process_embeds::worker(db.clone(), amqp.clone()));
     }
 }
 

crates/core/database/src/tasks/process_embeds.rs

@@ -1,4 +1,4 @@
-use crate::{models::Message, AppendMessage, Database};
+use crate::{AMQP, AppendMessage, Database, models::Message};
 
 use futures::future::join_all;
 use linkify::{LinkFinder, LinkKind};
@@ -41,7 +41,7 @@ pub async fn queue(channel: String, id: String, content: String) {
 }
 
 /// Start a new worker
-pub async fn worker(db: Database) {
+pub async fn worker(db: Database, amqp: AMQP) {
     let semaphore = Arc::new(Semaphore::new(
         config().await.api.workers.max_concurrent_connections,
     ));
@@ -49,6 +49,7 @@ pub async fn worker(db: Database) {
     loop {
         let task = Q.pop().await;
         let db = db.clone();
+        let amqp = amqp.clone();
         let semaphore = semaphore.clone();
 
         spawn(async move {
@@ -64,6 +65,7 @@ pub async fn worker(db: Database) {
             if let Ok(embeds) = embeds {
                 if let Err(err) = Message::append(
                     &db,
+                    Some(&amqp),
                     task.id,
                     task.channel,
                     AppendMessage {

crates/core/search/src/client.rs

@@ -10,12 +10,12 @@ use elasticsearch::{
     indices::{IndicesCreateParts, IndicesDeleteParts},
 };
 use elasticsearch_dsl::{FieldSort, Query, Search, SearchResponse, Sort};
-use revolt_database::Message;
+use revolt_database::{Database, Message};
 use serde_json::{Map, Value, json};
 
 pub use elasticsearch;
 
-use crate::SearchTerms;
+use crate::{MessageComponent, MetadataFile, SearchTerms};
 
 #[derive(Debug)]
 pub enum Error {
@@ -87,12 +87,6 @@ impl ElasticsearchClient {
             .indices()
             .create(IndicesCreateParts::Index("messages"))
             .body(json!({
-                // "settings": {
-                //     "index": {
-                //         "sort.field": "_id",
-                //         "sort.order": ["asc", "desc"],
-                //     },
-                // },
                 "mappings": {
                     "properties": {
                         "content": {"type": "text"},
@@ -105,9 +99,13 @@ impl ElasticsearchClient {
                         "embeds": {
                             "properties": {}
                         },
-                        "attachments": {  // TODO: images, videos, files
+                        "attachments": {
                             "type": "nested",
-                            "properties": {}
+                            "properties": {
+                                "metadata.type": {
+                                    "type": "keyword"
+                                }
+                            }
                         },
                         // TODO: links
                     }
@@ -152,15 +150,27 @@ impl ElasticsearchClient {
             }
         }
 
-        // TODO: component filtering
+        if let Some(components) = terms.filters.components {
-        // Need to pass FileHash to this to extract metadata
+            let mut components_query = Query::bool();
+
+            for component in components {
+                match component {
+                    MessageComponent::Image => components_query = components_query.should(Query::term("attachments.metadata.type", "Image")),
+                    MessageComponent::Video => components_query = components_query.should(Query::term("attachments.metadata.type", "Video")),
+                    MessageComponent::File => components_query = components_query.should(Query::exists("attachments.file._id")),
+                    MessageComponent::Embed => query = query.filter(Query::exists("embeds")),
+                };
+            }
+
+            query = query.filter(Query::nested("attachments", components_query));
+        }
 
         let search = Search::new()
             .query(query)
             .stats(false)
-            .sort([Sort::FieldSort(
+            .sort(Sort::FieldSort(
                 FieldSort::new("_id".to_string()).order(terms.sort.unwrap_or_default().into()),
-            )]);
+            ));
 
         let response = self
             .inner
@@ -184,7 +194,7 @@ impl ElasticsearchClient {
         }
     }
 
-    fn create_message_source(&self, message: Message) -> Value {
+    fn create_message_source(&self, _db: &Database, message: Message) -> Value {
         let mut map = Map::new();
 
         map.insert("channel".to_string(), Value::String(message.channel));
@@ -196,9 +206,21 @@ impl ElasticsearchClient {
         }
 
         if let Some(attachments) = message.attachments {
+            let mut files = Vec::new();
+
+            for attachment in attachments {
+                // TODO: swap this out for fetching the metadata from FileHash because of deprecation
+                // let metadata = attachment.as_hash(db).await.expect("Failed to fetch FileHash").metadata;
+
+                files.push(MetadataFile {
+                    metadata: attachment.metadata.clone(),
+                    file: attachment,
+                })
+            }
+
             map.insert(
                 "attachments".to_string(),
-                serde_json::to_value(attachments).unwrap(),
+                serde_json::to_value(files).unwrap(),
             );
         }
 
@@ -240,9 +262,9 @@ impl ElasticsearchClient {
         Value::Object(map)
     }
 
-    pub async fn index_message(&self, message: Message) -> Result<(), Error> {
+    pub async fn index_message(&self, db: &Database, message: Message) -> Result<(), Error> {
         let id = message.id.clone();
-        let source = self.create_message_source(message);
+        let source = self.create_message_source(db, message);
 
         let exception = self
             .inner
@@ -260,9 +282,9 @@ impl ElasticsearchClient {
         }
     }
 
-    pub async fn edit_message(&self, message: Message) -> Result<(), Error> {
+    pub async fn edit_message(&self, db: &Database, message: Message) -> Result<(), Error> {
         let id = message.id.clone();
-        let source = self.create_message_source(message);
+        let source = self.create_message_source(db, message);
 
         let exception = self
             .inner

crates/core/search/src/types.rs

@@ -1,15 +1,16 @@
 use iso8601_timestamp::Timestamp;
 use revolt_models::v0;
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
+use revolt_database::{File, Metadata};
 
-#[derive(Copy, Clone, PartialEq, Eq, Serialize)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize)]
 pub enum AuthorType {
     User,
     // Bot,
     Webhook,
 }
 
-#[derive(Copy, Clone, PartialEq, Eq)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub enum MessageComponent {
     Image,
     Video,
@@ -18,7 +19,7 @@ pub enum MessageComponent {
     Embed,
 }
 
-#[derive(Clone, Default, PartialEq)]
+#[derive(Debug, Clone, Default, PartialEq)]
 pub struct SearchFilters {
     pub content: Option<String>,
     pub author: Option<Vec<String>>,
@@ -31,14 +32,14 @@ pub struct SearchFilters {
     pub components: Option<Vec<MessageComponent>>,
 }
 
-#[derive(Clone, Copy, Default, PartialEq, Eq)] 
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
 pub enum SortOrder {
     Asc,
     #[default]
     Desc,
 }
 
-#[derive(Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct SearchTerms {
     pub channels: Vec<String>,
     pub filters: SearchFilters,
@@ -47,6 +48,12 @@ pub struct SearchTerms {
     pub sort: Option<SortOrder>,
 }
 
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct MetadataFile {
+    pub file: File,
+    pub metadata: Metadata
+}
+
 impl From<v0::AuthorType> for AuthorType {
     fn from(value: v0::AuthorType) -> Self {
         match value {

crates/daemons/searchd/src/consumers/channel_delete.rs

@@ -1,9 +1,18 @@
 use amqprs::{BasicProperties, Deliver, channel::{BasicAckArguments, Channel}, consumer::AsyncConsumer};
 use async_trait::async_trait;
-use revolt_database::events::rabbit::ChannelDeletePayload;
+use revolt_database::{Database, events::rabbit::ChannelDeletePayload};
 use revolt_search::ElasticsearchClient;
 
-pub struct ChannelDeleteConsumer(pub ElasticsearchClient);
+pub struct ChannelDeleteConsumer {
+    client: ElasticsearchClient,
+    database: Database
+}
+
+impl ChannelDeleteConsumer {
+    pub fn new(client: ElasticsearchClient, database: Database) -> Self {
+        Self { client, database }
+    }
+}
 
 #[async_trait]
 impl AsyncConsumer for ChannelDeleteConsumer {
@@ -15,8 +24,9 @@ impl AsyncConsumer for ChannelDeleteConsumer {
         content: Vec<u8>,
     ) {
         let payload = serde_json::from_slice::<ChannelDeletePayload>(&content).expect("Failed to decode message");
+        log::debug!("Received channel delete {payload:?}");
 
-        if self.0.delete_channel(&payload.channel_id).await.is_ok() {
+        if self.client.delete_channel(&payload.channel_id).await.is_ok() {
             channel.basic_ack(BasicAckArguments::new(deliver.delivery_tag(), false)).await.expect("Failed to ack");
         } else {
             // todo requeue

crates/daemons/searchd/src/consumers/message.rs

@@ -1,9 +1,18 @@
 use amqprs::{BasicProperties, Deliver, channel::{BasicAckArguments, Channel}, consumer::AsyncConsumer};
 use async_trait::async_trait;
-use revolt_database::Message;
+use revolt_database::{Database, Message};
 use revolt_search::ElasticsearchClient;
 
-pub struct MessageConsumer(pub ElasticsearchClient);
+pub struct MessageConsumer {
+    client: ElasticsearchClient,
+    database: Database
+}
+
+impl MessageConsumer {
+    pub fn new(client: ElasticsearchClient, database: Database) -> Self {
+        Self { client, database }
+    }
+}
 
 #[async_trait]
 impl AsyncConsumer for MessageConsumer {
@@ -17,7 +26,7 @@ impl AsyncConsumer for MessageConsumer {
         let message = serde_json::from_slice::<Message>(&content).expect("Failed to decode message");
         log::debug!("Received message {message:?}");
 
-        if self.0.index_message(message).await.is_ok() {
+        if self.client.index_message(&self.database, message).await.is_ok() {
             channel.basic_ack(BasicAckArguments::new(deliver.delivery_tag(), false)).await.expect("Failed to ack");
         } else {
             // todo requeue

crates/daemons/searchd/src/consumers/message_delete.rs

@@ -1,9 +1,18 @@
 use amqprs::{BasicProperties, Deliver, channel::{BasicAckArguments, Channel}, consumer::AsyncConsumer};
 use async_trait::async_trait;
-use revolt_database::events::rabbit::MessageDeletePayload;
+use revolt_database::{Database, events::rabbit::MessageDeletePayload};
 use revolt_search::ElasticsearchClient;
 
-pub struct MessageDeleteConsumer(pub ElasticsearchClient);
+pub struct MessageDeleteConsumer {
+    client: ElasticsearchClient,
+    database: Database
+}
+
+impl MessageDeleteConsumer {
+    pub fn new(client: ElasticsearchClient, database: Database) -> Self {
+        Self { client, database }
+    }
+}
 
 #[async_trait]
 impl AsyncConsumer for MessageDeleteConsumer {
@@ -15,8 +24,9 @@ impl AsyncConsumer for MessageDeleteConsumer {
         content: Vec<u8>,
     ) {
         let payload = serde_json::from_slice::<MessageDeletePayload>(&content).expect("Failed to decode message");
+        log::debug!("Received message delete {payload:?}");
 
-        if self.0.delete_message(&payload.message_id).await.is_ok() {
+        if self.client.delete_message(&payload.message_id).await.is_ok() {
             channel.basic_ack(BasicAckArguments::new(deliver.delivery_tag(), false)).await.expect("Failed to ack");
         } else {
             // todo requeue

crates/daemons/searchd/src/consumers/message_edit.rs

@@ -1,9 +1,18 @@
 use amqprs::{BasicProperties, Deliver, channel::{BasicAckArguments, Channel}, consumer::AsyncConsumer};
 use async_trait::async_trait;
-use revolt_database::Message;
+use revolt_database::{Database, Message};
 use revolt_search::ElasticsearchClient;
 
-pub struct MessageEditConsumer(pub ElasticsearchClient);
+pub struct MessageEditConsumer {
+    client: ElasticsearchClient,
+    database: Database
+}
+
+impl MessageEditConsumer {
+    pub fn new(client: ElasticsearchClient, database: Database) -> Self {
+        Self { client, database }
+    }
+}
 
 #[async_trait]
 impl AsyncConsumer for MessageEditConsumer {
@@ -15,8 +24,9 @@ impl AsyncConsumer for MessageEditConsumer {
         content: Vec<u8>,
     ) {
         let message = serde_json::from_slice::<Message>(&content).expect("Failed to decode message");
+        log::debug!("Received edit message {message:?}");
 
-        if self.0.edit_message(message).await.is_ok() {
+        if self.client.edit_message(&self.database, message).await.is_ok() {
             channel.basic_ack(BasicAckArguments::new(deliver.delivery_tag(), false)).await.expect("Failed to ack");
         } else {
             // todo requeue

crates/daemons/searchd/src/main.rs

@@ -50,7 +50,7 @@ async fn _main() {
             &config,
             &connection,
             &config.elasticsearch.message_queue,
-            consumers::MessageConsumer(client.clone()),
+            consumers::MessageConsumer::new(client.clone(), db.clone()),
         )
         .await,
     );
@@ -60,7 +60,7 @@ async fn _main() {
             &config,
             &connection,
             &config.elasticsearch.message_edit_queue,
-            consumers::MessageEditConsumer(client.clone()),
+            consumers::MessageEditConsumer::new(client.clone(), db.clone()),
         )
         .await,
     );
@@ -70,7 +70,7 @@ async fn _main() {
             &config,
             &connection,
             &config.elasticsearch.message_delete_queue,
-            consumers::MessageDeleteConsumer(client.clone()),
+            consumers::MessageDeleteConsumer::new(client.clone(), db.clone()),
         )
         .await,
     );
@@ -80,7 +80,7 @@ async fn _main() {
             &config,
             &connection,
             &config.elasticsearch.channel_delete_queue,
-            consumers::ChannelDeleteConsumer(client.clone()),
+            consumers::ChannelDeleteConsumer::new(client.clone(), db.clone()),
         )
         .await,
     );