diff --git a/crates/core/database/src/events/rabbit.rs b/crates/core/database/src/events/rabbit.rs index cf33c15f..35d77427 100644 --- a/crates/core/database/src/events/rabbit.rs +++ b/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 } \ No newline at end of file diff --git a/crates/core/database/src/models/messages/model.rs b/crates/core/database/src/models/messages/model.rs index 3f9fe7e5..d40a7e4c 100644 --- a/crates/core/database/src/models/messages/model.rs +++ b/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 { - id, - channel: channel.to_string(), - append: append.into(), + EventV1::MessageAppend { + id, + channel: channel.to_string(), + append: append.into(), + } + .p(channel) + .await; } - .p(channel) - .await; Ok(()) } diff --git a/crates/core/database/src/models/messages/ops.rs b/crates/core/database/src/models/messages/ops.rs index ed542dde..b6358cfb 100644 --- a/crates/core/database/src/models/messages/ops.rs +++ b/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) -> 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>; /// Add a new reaction to a message async fn add_reaction(&self, id: &str, emoji: &str, user: &str) -> Result<()>; diff --git a/crates/core/database/src/models/messages/ops/mongodb.rs b/crates/core/database/src/models/messages/ops/mongodb.rs index a87ddbf2..a518d5b8 100644 --- a/crates/core/database/src/models/messages/ops/mongodb.rs +++ b/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> { 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::(COL) - .update_one( + self.col::(COL) + .find_one_and_update( doc! { "_id": id }, query, ) + .return_document(ReturnDocument::After) .await - .map(|_| ()) .map_err(|_| create_database_error!("update_one", COL)) } diff --git a/crates/core/database/src/models/messages/ops/reference.rs b/crates/core/database/src/models/messages/ops/reference.rs index b70f7a97..1c4c8e2e 100644 --- a/crates/core/database/src/models/messages/ops/reference.rs +++ b/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> { 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)) } diff --git a/crates/core/database/src/tasks/mod.rs b/crates/core/database/src/tasks/mod.rs index 111d5106..c305fcb4 100644 --- a/crates/core/database/src/tasks/mod.rs +++ b/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())); } } diff --git a/crates/core/database/src/tasks/process_embeds.rs b/crates/core/database/src/tasks/process_embeds.rs index d9d61d75..82c800ec 100644 --- a/crates/core/database/src/tasks/process_embeds.rs +++ b/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 { diff --git a/crates/core/search/src/client.rs b/crates/core/search/src/client.rs index 900858a7..b2d69d29 100644 --- a/crates/core/search/src/client.rs +++ b/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 - // Need to pass FileHash to this to extract metadata + if let Some(components) = terms.filters.components { + 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 diff --git a/crates/core/search/src/types.rs b/crates/core/search/src/types.rs index dc79d63a..2d9b8336 100644 --- a/crates/core/search/src/types.rs +++ b/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, pub author: Option>, @@ -31,14 +32,14 @@ pub struct SearchFilters { pub components: Option>, } -#[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, pub filters: SearchFilters, @@ -47,6 +48,12 @@ pub struct SearchTerms { pub sort: Option, } +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MetadataFile { + pub file: File, + pub metadata: Metadata +} + impl From for AuthorType { fn from(value: v0::AuthorType) -> Self { match value { diff --git a/crates/daemons/searchd/src/consumers/channel_delete.rs b/crates/daemons/searchd/src/consumers/channel_delete.rs index 170f8189..0385b578 100644 --- a/crates/daemons/searchd/src/consumers/channel_delete.rs +++ b/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, ) { let payload = serde_json::from_slice::(&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 diff --git a/crates/daemons/searchd/src/consumers/message.rs b/crates/daemons/searchd/src/consumers/message.rs index ce288100..8dccf36f 100644 --- a/crates/daemons/searchd/src/consumers/message.rs +++ b/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::(&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 diff --git a/crates/daemons/searchd/src/consumers/message_delete.rs b/crates/daemons/searchd/src/consumers/message_delete.rs index e1a1f17a..a3f70259 100644 --- a/crates/daemons/searchd/src/consumers/message_delete.rs +++ b/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, ) { let payload = serde_json::from_slice::(&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 diff --git a/crates/daemons/searchd/src/consumers/message_edit.rs b/crates/daemons/searchd/src/consumers/message_edit.rs index 8b8ed42a..f9904056 100644 --- a/crates/daemons/searchd/src/consumers/message_edit.rs +++ b/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, ) { let message = serde_json::from_slice::(&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 diff --git a/crates/daemons/searchd/src/main.rs b/crates/daemons/searchd/src/main.rs index ae0bb1d0..2097f389 100644 --- a/crates/daemons/searchd/src/main.rs +++ b/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, );