feat: initial work on elasticsearch integration
Signed-off-by: Zomatree <me@zomatree.live>
This commit is contained in:
@@ -323,3 +323,14 @@ proxy = ""
|
||||
pushd = ""
|
||||
crond = ""
|
||||
gifbox = ""
|
||||
|
||||
[elasticsearch]
|
||||
host = "http://elasticsearch"
|
||||
port = 9200
|
||||
api_key = ""
|
||||
|
||||
exchange = "revolt.messages"
|
||||
message_queue = "messages.message"
|
||||
message_edit_queue = "messages.message_edit"
|
||||
message_delete_queue = "messages.message_delete"
|
||||
channel_delete_queue = "messages.channel_delete"
|
||||
@@ -413,6 +413,19 @@ pub struct Sentry {
|
||||
pub gifbox: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
pub struct Elasticsearch {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub api_key: String,
|
||||
|
||||
pub exchange: String,
|
||||
pub message_queue: String,
|
||||
pub message_edit_queue: String,
|
||||
pub message_delete_queue: String,
|
||||
pub channel_delete_queue: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
pub struct Settings {
|
||||
pub database: Database,
|
||||
@@ -425,6 +438,7 @@ pub struct Settings {
|
||||
pub sentry: Sentry,
|
||||
pub production: bool,
|
||||
pub disable_events_dont_use: bool,
|
||||
pub elasticsearch: Elasticsearch,
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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;
|
||||
@@ -339,7 +339,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<()> {
|
||||
@@ -376,7 +376,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()),
|
||||
@@ -669,7 +669,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,
|
||||
@@ -701,7 +701,7 @@ impl Channel {
|
||||
.into_message(id.to_string())
|
||||
.send(
|
||||
db,
|
||||
Some(amqp),
|
||||
amqp,
|
||||
MessageAuthor::System {
|
||||
username: name,
|
||||
avatar: None,
|
||||
@@ -714,7 +714,7 @@ impl Channel {
|
||||
.await
|
||||
.ok();
|
||||
} else {
|
||||
return self.delete(db).await;
|
||||
return self.delete(db, amqp).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -741,7 +741,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()),
|
||||
@@ -763,13 +763,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(())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
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,
|
||||
@@ -671,7 +671,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>,
|
||||
@@ -726,6 +726,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(())
|
||||
}
|
||||
|
||||
@@ -994,9 +1001,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();
|
||||
|
||||
@@ -1006,12 +1014,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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::collections::HashMap;
|
||||
use std::time::SystemTime;
|
||||
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;
|
||||
@@ -50,4 +50,5 @@ pub trait AbstractMessages: Sync + Send {
|
||||
author: &str,
|
||||
since: SystemTime
|
||||
) -> Result<HashMap<String, Vec<String>>>;
|
||||
async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<Message>>;
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use bson::{to_bson, Document};
|
||||
use futures::try_join;
|
||||
use futures::StreamExt;
|
||||
use mongodb::options::FindOptions;
|
||||
use mongodb::options::{FindOptions, ReadConcern};
|
||||
use revolt_models::v0::MessageSort;
|
||||
use revolt_result::Result;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
@@ -10,7 +10,7 @@ use ulid::Ulid;
|
||||
|
||||
use crate::{
|
||||
AppendMessage, DocumentId, FieldsMessage, IntoDocumentPath, Message, MessageQuery,
|
||||
MessageTimePeriod, MongoDb, PartialMessage,
|
||||
MessageTimePeriod, MongoDb, PartialMessage, util::ChunkedDatabaseGenerator,
|
||||
};
|
||||
|
||||
use super::AbstractMessages;
|
||||
@@ -416,6 +416,31 @@ impl AbstractMessages for MongoDb {
|
||||
|
||||
Ok(deleted_messages)
|
||||
}
|
||||
|
||||
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 {
|
||||
|
||||
@@ -4,7 +4,11 @@ use indexmap::IndexSet;
|
||||
use revolt_result::Result;
|
||||
use std::time::SystemTime;
|
||||
use ulid::Ulid;
|
||||
use crate::{AppendMessage, FieldsMessage, Message, MessageQuery, PartialMessage, ReferenceDb};
|
||||
|
||||
use crate::{
|
||||
util::ChunkedDatabaseGenerator, AppendMessage, FieldsMessage, Message, MessageQuery,
|
||||
PartialMessage, ReferenceDb,
|
||||
};
|
||||
|
||||
use super::AbstractMessages;
|
||||
|
||||
@@ -347,4 +351,10 @@ impl AbstractMessages for ReferenceDb {
|
||||
|
||||
Ok(deleted_messages)
|
||||
}
|
||||
|
||||
async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<Message>> {
|
||||
Ok(ChunkedDatabaseGenerator::new_reference(
|
||||
self.messages.lock().await.values().cloned().collect(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
54
crates/core/database/src/util/chunked.rs
Normal file
54
crates/core/database/src/util/chunked.rs
Normal file
@@ -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())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -14,6 +14,7 @@ mod server_members;
|
||||
mod servers;
|
||||
mod user_settings;
|
||||
mod users;
|
||||
mod search;
|
||||
|
||||
pub use bots::*;
|
||||
pub use channel_invites::*;
|
||||
@@ -31,3 +32,4 @@ pub use server_members::*;
|
||||
pub use servers::*;
|
||||
pub use user_settings::*;
|
||||
pub use users::*;
|
||||
pub use search::*;
|
||||
|
||||
49
crates/core/models/src/v0/search.rs
Normal file
49
crates/core/models/src/v0/search.rs
Normal file
@@ -0,0 +1,49 @@
|
||||
use iso8601_timestamp::Timestamp;
|
||||
|
||||
auto_derived!(
|
||||
pub struct DataChannelMessagesSearch {
|
||||
pub channel: Option<String>,
|
||||
pub server: Option<String>,
|
||||
pub offset: Option<u64>,
|
||||
pub limit: Option<u64>,
|
||||
pub filters: Option<DataChannelMessagesSearchFilters>,
|
||||
pub sort: Option<SortOrder>,
|
||||
pub include_users: Option<bool>,
|
||||
}
|
||||
|
||||
pub struct DataChannelMessagesSearchFilters {
|
||||
pub content: Option<String>,
|
||||
pub author: Option<Vec<String>>,
|
||||
pub mentions: Option<Vec<String>>,
|
||||
pub role_mentions: Option<Vec<String>>,
|
||||
pub before_date: Option<Timestamp>,
|
||||
pub after_date: Option<Timestamp>,
|
||||
pub author_type: Option<Vec<AuthorType>>,
|
||||
pub pinned: Option<bool>,
|
||||
pub components: Option<Vec<MessageComponent>>,
|
||||
}
|
||||
|
||||
#[derive(Copy)]
|
||||
pub enum AuthorType {
|
||||
User,
|
||||
// Bot,
|
||||
Webhook,
|
||||
}
|
||||
|
||||
#[derive(Copy)]
|
||||
pub enum MessageComponent {
|
||||
Image,
|
||||
Video,
|
||||
// Link,
|
||||
File,
|
||||
Embed,
|
||||
}
|
||||
|
||||
#[derive(Copy, Default)]
|
||||
#[cfg_attr(feature = "serde", serde(rename_all = "lowercase"))]
|
||||
pub enum SortOrder {
|
||||
Asc,
|
||||
#[default]
|
||||
Desc,
|
||||
}
|
||||
);
|
||||
18
crates/core/search/Cargo.toml
Normal file
18
crates/core/search/Cargo.toml
Normal file
@@ -0,0 +1,18 @@
|
||||
[package]
|
||||
name = "revolt-search"
|
||||
version = "0.11.5"
|
||||
edition = "2024"
|
||||
license = "MIT"
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
revolt-config = { version = "0.11.5", path = "../config" }
|
||||
revolt-database = { version = "0.11.5", path = "../database" }
|
||||
revolt-models = { version = "0.11.5", path = "../models" }
|
||||
elasticsearch = "9.1.0-alpha.1"
|
||||
serde_json = "1"
|
||||
ulid = "1.2.1"
|
||||
iso8601-timestamp = { version = "0.2.10", features = ["serde", "bson"] }
|
||||
futures = "0.3.32"
|
||||
elasticsearch-dsl = "0.4"
|
||||
serde = "1"
|
||||
315
crates/core/search/src/client.rs
Normal file
315
crates/core/search/src/client.rs
Normal file
@@ -0,0 +1,315 @@
|
||||
use std::fmt::Display;
|
||||
|
||||
use elasticsearch::{
|
||||
CreateParts, DeleteByQueryParts, DeleteParts, Elasticsearch, IndexParts, SearchParts,
|
||||
auth::Credentials,
|
||||
http::{
|
||||
response::Exception,
|
||||
transport::{SingleNodeConnectionPool, TransportBuilder},
|
||||
},
|
||||
indices::{IndicesCreateParts, IndicesDeleteParts},
|
||||
};
|
||||
use elasticsearch_dsl::{FieldSort, Query, Search, SearchResponse, Sort};
|
||||
use revolt_database::Message;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
pub use elasticsearch;
|
||||
|
||||
use crate::SearchTerms;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Error {
|
||||
Http(elasticsearch::Error),
|
||||
Exception(Exception),
|
||||
}
|
||||
|
||||
impl From<elasticsearch::Error> for Error {
|
||||
fn from(value: elasticsearch::Error) -> Self {
|
||||
Self::Http(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Exception> for Error {
|
||||
fn from(value: Exception) -> Self {
|
||||
Self::Exception(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for Error {}
|
||||
impl Display for Error {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Error::Http(error) => write!(f, "Http error: {error}"),
|
||||
Error::Exception(exception) => write!(f, "Elasticsearch error: {exception:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ElasticsearchClient {
|
||||
pub inner: Elasticsearch,
|
||||
}
|
||||
|
||||
impl ElasticsearchClient {
|
||||
pub fn new(host: &str, port: u16, key: String) -> Self {
|
||||
let pool =
|
||||
SingleNodeConnectionPool::new(format!("{host}:{port}").as_str().try_into().unwrap());
|
||||
let transport = TransportBuilder::new(pool)
|
||||
.auth(Credentials::EncodedApiKey(key))
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let inner = Elasticsearch::new(transport);
|
||||
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
pub async fn delete_indexes(&self) -> Result<(), Error> {
|
||||
let exception = self
|
||||
.inner
|
||||
.indices()
|
||||
.delete(IndicesDeleteParts::Index(&["messages"]))
|
||||
.send()
|
||||
.await?
|
||||
.exception()
|
||||
.await?;
|
||||
|
||||
if let Some(exception) = exception {
|
||||
Err(exception.into())
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn setup_indexes(&self) -> Result<(), Error> {
|
||||
let exception = self
|
||||
.inner
|
||||
.indices()
|
||||
.create(IndicesCreateParts::Index("messages"))
|
||||
.body(json!({
|
||||
// "settings": {
|
||||
// "index": {
|
||||
// "sort.field": "_id",
|
||||
// "sort.order": ["asc", "desc"],
|
||||
// },
|
||||
// },
|
||||
"mappings": {
|
||||
"properties": {
|
||||
"content": {"type": "text"},
|
||||
"author": {"type": "keyword"},
|
||||
"author_type": {"type": "keyword"},
|
||||
"channel": {"type": "keyword"},
|
||||
"mentions": {"type": "keyword"},
|
||||
"role_mentions": {"type": "keyword"},
|
||||
"pinned": {"type": "boolean"},
|
||||
"embeds": {
|
||||
"properties": {}
|
||||
},
|
||||
"attachments": { // TODO: images, videos, files
|
||||
"type": "nested",
|
||||
"properties": {}
|
||||
},
|
||||
// TODO: links
|
||||
}
|
||||
}
|
||||
}))
|
||||
.send()
|
||||
.await?
|
||||
.exception()
|
||||
.await?;
|
||||
|
||||
if let Some(exception) = exception {
|
||||
Err(exception.into())
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn search(&self, terms: SearchTerms) -> Result<Vec<String>, Error> {
|
||||
let mut query = Query::bool().filter(Query::terms("channel", terms.channels));
|
||||
|
||||
if let Some(content) = terms.filters.content {
|
||||
query = query.filter(Query::r#match("content", content))
|
||||
}
|
||||
|
||||
if let Some(author) = terms.filters.author {
|
||||
query = query.filter(Query::terms("author", author))
|
||||
}
|
||||
|
||||
if let Some(mentions) = terms.filters.mentions {
|
||||
query = query.filter(Query::terms("mentions", mentions))
|
||||
}
|
||||
|
||||
if let Some(author_type) = terms.filters.author_type {
|
||||
query = query.filter(Query::terms("author_type", author_type))
|
||||
}
|
||||
|
||||
if let Some(pinned) = terms.filters.pinned {
|
||||
if pinned {
|
||||
query = query.filter(Query::exists("pinned"))
|
||||
} else {
|
||||
query = query.filter(Query::bool().must_not(Query::exists("pinned")))
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: component filtering
|
||||
// Need to pass FileHash to this to extract metadata
|
||||
|
||||
let search = Search::new()
|
||||
.query(query)
|
||||
.stats(false)
|
||||
.sort([Sort::FieldSort(
|
||||
FieldSort::new("_id".to_string()).order(terms.sort.unwrap_or_default().into()),
|
||||
)]);
|
||||
|
||||
let response = self
|
||||
.inner
|
||||
.search(SearchParts::Index(&["messages"]))
|
||||
.stored_fields(&[])
|
||||
.body(search)
|
||||
.size(terms.limit.unwrap_or(100) as i64)
|
||||
.from(terms.offset.unwrap_or(0) as i64)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if response.status_code().is_success() {
|
||||
let messages = response.json::<SearchResponse>().await?;
|
||||
Ok(messages.hits.hits.into_iter().map(|hit| hit.id).collect())
|
||||
} else {
|
||||
Err(response
|
||||
.exception()
|
||||
.await?
|
||||
.expect("No exception with error response.")
|
||||
.into())
|
||||
}
|
||||
}
|
||||
|
||||
fn create_message_source(&self, message: Message) -> Value {
|
||||
let mut map = Map::new();
|
||||
|
||||
map.insert("channel".to_string(), Value::String(message.channel));
|
||||
|
||||
map.insert("author".to_string(), Value::String(message.author));
|
||||
|
||||
if let Some(content) = message.content {
|
||||
map.insert("content".to_string(), Value::String(content));
|
||||
}
|
||||
|
||||
if let Some(attachments) = message.attachments {
|
||||
map.insert(
|
||||
"attachments".to_string(),
|
||||
serde_json::to_value(attachments).unwrap(),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(embeds) = message.embeds {
|
||||
map.insert("embeds".to_string(), serde_json::to_value(embeds).unwrap());
|
||||
}
|
||||
|
||||
if let Some(mentions) = message.mentions {
|
||||
map.insert(
|
||||
"mentions".to_string(),
|
||||
serde_json::to_value(mentions).unwrap(),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(role_mentions) = message.role_mentions {
|
||||
map.insert(
|
||||
"role_mentions".to_string(),
|
||||
serde_json::to_value(role_mentions).unwrap(),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(pinned) = message.pinned {
|
||||
map.insert("pinned".to_string(), Value::Bool(pinned));
|
||||
}
|
||||
|
||||
// TODO: bot
|
||||
map.insert(
|
||||
"author_type".to_string(),
|
||||
Value::String(
|
||||
if message.webhook.is_some() {
|
||||
"webhook"
|
||||
} else {
|
||||
"user"
|
||||
}
|
||||
.to_string(),
|
||||
),
|
||||
);
|
||||
|
||||
Value::Object(map)
|
||||
}
|
||||
|
||||
pub async fn index_message(&self, message: Message) -> Result<(), Error> {
|
||||
let id = message.id.clone();
|
||||
let source = self.create_message_source(message);
|
||||
|
||||
let exception = self
|
||||
.inner
|
||||
.create(CreateParts::IndexId("messages", &id))
|
||||
.body(source)
|
||||
.send()
|
||||
.await?
|
||||
.exception()
|
||||
.await?;
|
||||
|
||||
if let Some(exception) = exception {
|
||||
Err(exception.into())
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn edit_message(&self, message: Message) -> Result<(), Error> {
|
||||
let id = message.id.clone();
|
||||
let source = self.create_message_source(message);
|
||||
|
||||
let exception = self
|
||||
.inner
|
||||
.index(IndexParts::IndexId("messages", &id))
|
||||
.body(source)
|
||||
.send()
|
||||
.await?
|
||||
.exception()
|
||||
.await?;
|
||||
|
||||
if let Some(exception) = exception {
|
||||
Err(exception.into())
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn delete_message(&self, message_id: &str) -> Result<(), Error> {
|
||||
let exception = self
|
||||
.inner
|
||||
.delete(DeleteParts::IndexId("messages", message_id))
|
||||
.send()
|
||||
.await?
|
||||
.exception()
|
||||
.await?;
|
||||
|
||||
if let Some(exception) = exception {
|
||||
Err(exception.into())
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn delete_channel(&self, channel_id: &str) -> Result<(), Error> {
|
||||
let exception = self
|
||||
.inner
|
||||
.delete_by_query(DeleteByQueryParts::Index(&["messages"]))
|
||||
.body(Search::new().query(Query::term("channel", channel_id)))
|
||||
.send()
|
||||
.await?
|
||||
.exception()
|
||||
.await?;
|
||||
|
||||
if let Some(exception) = exception {
|
||||
Err(exception.into())
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
5
crates/core/search/src/lib.rs
Normal file
5
crates/core/search/src/lib.rs
Normal file
@@ -0,0 +1,5 @@
|
||||
mod client;
|
||||
mod types;
|
||||
|
||||
pub use client::*;
|
||||
pub use types::*;
|
||||
108
crates/core/search/src/types.rs
Normal file
108
crates/core/search/src/types.rs
Normal file
@@ -0,0 +1,108 @@
|
||||
use iso8601_timestamp::Timestamp;
|
||||
use revolt_models::v0;
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Copy, Clone, PartialEq, Eq, Serialize)]
|
||||
pub enum AuthorType {
|
||||
User,
|
||||
// Bot,
|
||||
Webhook,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, PartialEq, Eq)]
|
||||
pub enum MessageComponent {
|
||||
Image,
|
||||
Video,
|
||||
// Link,
|
||||
File,
|
||||
Embed,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, PartialEq)]
|
||||
pub struct SearchFilters {
|
||||
pub content: Option<String>,
|
||||
pub author: Option<Vec<String>>,
|
||||
pub mentions: Option<Vec<String>>,
|
||||
pub role_mentions: Option<Vec<String>>,
|
||||
pub before_date: Option<Timestamp>,
|
||||
pub after_date: Option<Timestamp>,
|
||||
pub author_type: Option<Vec<AuthorType>>,
|
||||
pub pinned: Option<bool>,
|
||||
pub components: Option<Vec<MessageComponent>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub enum SortOrder {
|
||||
Asc,
|
||||
#[default]
|
||||
Desc,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct SearchTerms {
|
||||
pub channels: Vec<String>,
|
||||
pub filters: SearchFilters,
|
||||
pub offset: Option<u64>,
|
||||
pub limit: Option<u64>,
|
||||
pub sort: Option<SortOrder>,
|
||||
}
|
||||
|
||||
impl From<v0::AuthorType> for AuthorType {
|
||||
fn from(value: v0::AuthorType) -> Self {
|
||||
match value {
|
||||
v0::AuthorType::User => AuthorType::User,
|
||||
// v0::AuthorType::Bot => AuthorType::Bot,
|
||||
v0::AuthorType::Webhook => AuthorType::Webhook,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<v0::MessageComponent> for MessageComponent {
|
||||
fn from(value: v0::MessageComponent) -> Self {
|
||||
match value {
|
||||
v0::MessageComponent::Image => MessageComponent::Image,
|
||||
v0::MessageComponent::Video => MessageComponent::Video,
|
||||
// v0::MessageComponent::Link => MessageComponent::Link,
|
||||
v0::MessageComponent::File => MessageComponent::File,
|
||||
v0::MessageComponent::Embed => MessageComponent::Embed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<v0::SortOrder> for SortOrder {
|
||||
fn from(value: v0::SortOrder) -> Self {
|
||||
match value {
|
||||
v0::SortOrder::Asc => SortOrder::Asc,
|
||||
v0::SortOrder::Desc => SortOrder::Desc,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<v0::DataChannelMessagesSearchFilters> for SearchFilters {
|
||||
fn from(value: v0::DataChannelMessagesSearchFilters) -> Self {
|
||||
Self {
|
||||
content: value.content,
|
||||
author: value.author,
|
||||
mentions: value.mentions,
|
||||
role_mentions: value.role_mentions,
|
||||
before_date: value.before_date,
|
||||
after_date: value.after_date,
|
||||
author_type: value
|
||||
.author_type
|
||||
.map(|types| types.into_iter().map(Into::into).collect()),
|
||||
pinned: value.pinned,
|
||||
components: value
|
||||
.components
|
||||
.map(|types| types.into_iter().map(Into::into).collect()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SortOrder> for elasticsearch_dsl::SortOrder {
|
||||
fn from(value: SortOrder) -> Self {
|
||||
match value {
|
||||
SortOrder::Asc => elasticsearch_dsl::SortOrder::Asc,
|
||||
SortOrder::Desc => elasticsearch_dsl::SortOrder::Desc,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user