feat: initial work on elasticsearch integration

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

crates/daemons/searchd/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "revolt-searchd"
+version = "0.11.5"
+edition = "2024"
+license = "AGPL-3.0-or-later"
+publish = false
+
+[dependencies]
+tokio = { version = "1.50.0", features = ["macros"] }
+revolt-database = { version = "0.11.5", path = "../../core/database" }
+revolt-config = { version = "0.11.5", path = "../../core/config" }
+revolt-search = { version = "0.11.5", path = "../../core/search" }
+amqprs = "1.7.0"
+async-trait = "0.1.89"
+log = "0.4.29"
+serde_json = "1"

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

@@ -0,0 +1,25 @@
+use amqprs::{BasicProperties, Deliver, channel::{BasicAckArguments, Channel}, consumer::AsyncConsumer};
+use async_trait::async_trait;
+use revolt_database::events::rabbit::{ChannelDeletePayload, MessageDeletePayload};
+use revolt_search::ElasticsearchClient;
+
+pub struct ChannelDeleteConsumer(pub ElasticsearchClient);
+
+#[async_trait]
+impl AsyncConsumer for ChannelDeleteConsumer {
+    async fn consume(
+        &mut self,
+        channel: &Channel,
+        deliver: Deliver,
+        _basic_properties: BasicProperties,
+        content: Vec<u8>,
+    ) {
+        let payload = serde_json::from_slice::<ChannelDeletePayload>(&content).expect("Failed to decode message");
+
+        if self.0.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

@@ -0,0 +1,26 @@
+use amqprs::{BasicProperties, Deliver, channel::{BasicAckArguments, Channel}, consumer::AsyncConsumer};
+use async_trait::async_trait;
+use revolt_database::Message;
+use revolt_search::ElasticsearchClient;
+
+pub struct MessageConsumer(pub ElasticsearchClient);
+
+#[async_trait]
+impl AsyncConsumer for MessageConsumer {
+    async fn consume(
+        &mut self,
+        channel: &Channel,
+        deliver: Deliver,
+        _basic_properties: BasicProperties,
+        content: Vec<u8>,
+    ) {
+        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() {
+            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

@@ -0,0 +1,25 @@
+use amqprs::{BasicProperties, Deliver, channel::{BasicAckArguments, Channel}, consumer::AsyncConsumer};
+use async_trait::async_trait;
+use revolt_database::events::rabbit::MessageDeletePayload;
+use revolt_search::ElasticsearchClient;
+
+pub struct MessageDeleteConsumer(pub ElasticsearchClient);
+
+#[async_trait]
+impl AsyncConsumer for MessageDeleteConsumer {
+    async fn consume(
+        &mut self,
+        channel: &Channel,
+        deliver: Deliver,
+        _basic_properties: BasicProperties,
+        content: Vec<u8>,
+    ) {
+        let payload = serde_json::from_slice::<MessageDeletePayload>(&content).expect("Failed to decode message");
+
+        if self.0.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

@@ -0,0 +1,25 @@
+use amqprs::{BasicProperties, Deliver, channel::{BasicAckArguments, Channel}, consumer::AsyncConsumer};
+use async_trait::async_trait;
+use revolt_database::Message;
+use revolt_search::ElasticsearchClient;
+
+pub struct MessageEditConsumer(pub ElasticsearchClient);
+
+#[async_trait]
+impl AsyncConsumer for MessageEditConsumer {
+    async fn consume(
+        &mut self,
+        channel: &Channel,
+        deliver: Deliver,
+        _basic_properties: BasicProperties,
+        content: Vec<u8>,
+    ) {
+        let message = serde_json::from_slice::<Message>(&content).expect("Failed to decode message");
+
+        if self.0.edit_message(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/mod.rs

@@ -0,0 +1,8 @@
+mod message;
+mod message_edit;
+mod message_delete;
+mod channel_delete;
+
+pub use message::*;
+pub use message_edit::*;
+pub use channel_delete::*;

crates/daemons/searchd/src/index.rs

@@ -0,0 +1,14 @@
+use revolt_config::capture_error;
+use revolt_database::{AMQP, Database};
+
+
+pub async fn index_existing_messages(db: Database, amqp: AMQP) {
+    let mut generator = db.fetch_all_messages().await.expect("Database query failed");
+
+    while let Some(message) = generator.next().await {
+        if let Err(e) = amqp.new_message_search(message).await {
+            log::error!("Error pushing message to RabbitMQ: {e}");
+            capture_error(&e);
+        }
+    }
+}

crates/daemons/searchd/src/main.rs

@@ -0,0 +1,151 @@
+use amqprs::{
+    channel::{
+        BasicConsumeArguments, Channel, ExchangeDeclareArguments, QueueBindArguments,
+        QueueDeclareArguments,
+    },
+    connection::{Connection, OpenConnectionArguments},
+    consumer::AsyncConsumer,
+};
+use revolt_config::{Settings, config, configure};
+use revolt_database::{AMQP, DatabaseInfo};
+use revolt_search::ElasticsearchClient;
+use tokio::{signal::ctrl_c, spawn};
+
+mod consumers;
+mod index;
+
+async fn _main() {
+    configure!(api);
+    let config = config().await;
+
+    let db = DatabaseInfo::Auto.connect().await.unwrap();
+
+    let client = ElasticsearchClient::new(
+        &config.elasticsearch.host,
+        config.elasticsearch.port,
+        config.elasticsearch.api_key.clone(),
+    );
+
+    if std::env::var("REMAKE_MESSAGES_INDEX").as_deref().is_ok_and(|v| v == "1") {
+        if client.delete_indexes().await.is_err() {
+            log::info!("Index does not existing, skipping.")
+        };
+
+        client.setup_indexes().await.unwrap();
+    };
+
+    let connection = Connection::open(&OpenConnectionArguments::new(
+        &config.rabbit.host,
+        config.rabbit.port,
+        &config.rabbit.username,
+        &config.rabbit.password,
+    ))
+    .await
+    .expect("Failed to connect to RabbitMQ");
+
+    let mut channels = Vec::new();
+
+    channels.push(
+        make_queue_and_consume(
+            &config,
+            &connection,
+            &config.elasticsearch.message_queue,
+            consumers::MessageConsumer(client.clone()),
+        )
+        .await,
+    );
+
+    channels.push(
+        make_queue_and_consume(
+            &config,
+            &connection,
+            &config.elasticsearch.message_edit_queue,
+            consumers::MessageEditConsumer(client.clone()),
+        )
+        .await,
+    );
+
+    channels.push(
+        make_queue_and_consume(
+            &config,
+            &connection,
+            &config.elasticsearch.message_delete_queue,
+            consumers::MessageConsumer(client.clone()),
+        )
+        .await,
+    );
+
+    channels.push(
+        make_queue_and_consume(
+            &config,
+            &connection,
+            &config.elasticsearch.channel_delete_queue,
+            consumers::ChannelDeleteConsumer(client.clone()),
+        )
+        .await,
+    );
+
+    let mut task = None;
+
+    if std::env::var("INDEX_ALL_MESSAGES").as_deref().is_ok_and(|v| v == "1") {
+        let amqp = AMQP::new_auto().await;
+        task = Some(spawn(index::index_existing_messages(db, amqp)));
+    }
+
+    ctrl_c().await.expect("Failed to wait for ctrl-c");
+
+    for channel in channels {
+        channel.close().await.unwrap();
+    }
+
+    if let Some(task) = task {
+        task.abort();
+    }
+}
+
+async fn make_queue_and_consume<F: AsyncConsumer + Send + 'static>(
+    config: &Settings,
+    connection: &Connection,
+    queue: &str,
+    consumer: F,
+) -> Channel {
+    let channel = connection.open_channel(None).await.unwrap();
+
+    channel
+        .exchange_declare(
+            ExchangeDeclareArguments::new(&config.elasticsearch.exchange, "direct")
+                .durable(true)
+                .finish(),
+        )
+        .await
+        .expect("Failed to declare pushd exchange");
+
+    _ = channel
+        .queue_declare(QueueDeclareArguments::new(queue).durable(true).finish())
+        .await
+        .unwrap()
+        .unwrap();
+
+    channel
+        .queue_bind(QueueBindArguments::new(
+            queue,
+            &config.elasticsearch.exchange,
+            queue,
+        ))
+        .await
+        .expect("This probably means the revolt.messages exchange does not exist in rabbitmq!");
+
+    let args = BasicConsumeArguments::new(queue, "")
+        .manual_ack(true)
+        .finish();
+
+    let tag = channel.basic_consume(consumer, args).await.unwrap();
+    log::info!("Consuming queue {queue}, tag {tag}");
+
+    channel
+}
+
+#[tokio::main]
+async fn main() {
+    _main().await
+}