feat: initial work on elasticsearch integration

Signed-off-by: Zomatree <me@zomatree.live>
This commit is contained in:
Zomatree
2026-03-19 23:04:37 +00:00
parent d9643ebd8d
commit 8009b3bf53
49 changed files with 1526 additions and 128 deletions

View File

@@ -84,6 +84,7 @@ revolt-presence = { path = "../core/presence" }
revolt-result = { path = "../core/result", features = ["rocket", "okapi"] }
revolt-permissions = { path = "../core/permissions", features = ["schemas"] }
revolt-ratelimits = { path = "../core/ratelimits", features = ["rocket"] }
revolt-search = { path = "../core/search" }
# voice
livekit-api = "0.4.4"

View File

@@ -12,6 +12,7 @@ use revolt_config::config;
use revolt_database::events::client::EventV1;
use revolt_database::AMQP;
use revolt_ratelimits::rocket as ratelimiter;
use revolt_search::ElasticsearchClient;
use rocket::{Build, Rocket};
use rocket_cors::{AllowedOrigins, CorsOptions};
use rocket_prometheus::PrometheusMetrics;
@@ -24,8 +25,8 @@ use amqprs::{
};
use async_std::channel::unbounded;
use authifier::AuthifierEvent;
use rocket::data::ToByteUnit;
use revolt_database::voice::VoiceClient;
use rocket::data::ToByteUnit;
pub async fn web() -> Rocket<Build> {
// Get settings
@@ -135,6 +136,15 @@ pub async fn web() -> Rocket<Build> {
.await
.expect("Failed to declare exchange");
channel
.exchange_declare(
ExchangeDeclareArguments::new(&config.elasticsearch.exchange, "direct")
.durable(true)
.finish(),
)
.await
.expect("Failed to declare exchange");
let amqp = AMQP::new(connection, channel);
// Launch background task workers
@@ -147,6 +157,13 @@ pub async fn web() -> Rocket<Build> {
// Ratelimits
let ratelimits = ratelimiter::RatelimitStorage::new(util::ratelimits::DeltaRatelimits);
// Search
let elasticsearch = ElasticsearchClient::new(
&config.elasticsearch.host,
config.elasticsearch.port,
config.elasticsearch.api_key.clone(),
);
routes::mount(config, rocket)
.attach(prometheus.clone())
.mount("/metrics", prometheus)
@@ -160,6 +177,7 @@ pub async fn web() -> Rocket<Build> {
.manage(cors.clone())
.manage(voice_client)
.manage(ratelimits)
.manage(elasticsearch)
.attach(ratelimiter::RatelimitFairing)
.attach(cors)
.configure(rocket::Config {

View File

@@ -56,7 +56,7 @@ pub async fn invite_bot(
.throw_if_lacking_channel_permission(ChannelPermission::InviteOthers)?;
channel
.add_user_to_group(db, amqp, &bot_user, &user.id)
.add_user_to_group(db, Some(amqp), &bot_user, &user.id)
.await
.map(|_| EmptyResponse)
}

View File

@@ -50,7 +50,7 @@ pub async fn delete(
channel
.remove_user_from_group(
db,
amqp,
Some(amqp),
&user,
None,
options.leave_silently.unwrap_or_default(),
@@ -65,7 +65,7 @@ pub async fn delete(
}
Channel::TextChannel { .. } => {
permissions.throw_if_lacking_channel_permission(ChannelPermission::ManageChannel)?;
channel.delete(db).await?;
channel.delete(db, Some(amqp)).await?;
delete_voice_channel(voice_client, &UserVoiceChannel::from_channel(&channel)).await?;
}

View File

@@ -39,7 +39,7 @@ pub async fn add_member(
}
channel
.add_user_to_group(db, amqp, &member, &user.id)
.add_user_to_group(db, Some(amqp), &member, &user.id)
.await
.map(|_| EmptyResponse)
}

View File

@@ -46,7 +46,7 @@ pub async fn remove_member(
}
channel
.remove_user_from_group(db, amqp, &member, Some(&user.id), false)
.remove_user_from_group(db, Some(amqp), &member, Some(&user.id), false)
.await?;
} else {
return Err(create_error!(InvalidOperation));

View File

@@ -1,6 +1,6 @@
use revolt_config::capture_error;
use revolt_database::{
util::{permissions::DatabasePermissionQuery, reference::Reference},
Database, User,
AMQP, Database, User, util::{permissions::DatabasePermissionQuery, reference::Reference}
};
use revolt_permissions::{calculate_channel_permissions, ChannelPermission};
use revolt_result::Result;
@@ -14,6 +14,7 @@ use rocket_empty::EmptyResponse;
#[delete("/<target>/messages/<msg>", rank = 2)]
pub async fn delete(
db: &State<Database>,
amqp: &State<AMQP>,
user: User,
target: Reference<'_>,
msg: Reference<'_>,
@@ -28,5 +29,5 @@ pub async fn delete(
.throw_if_lacking_channel_permission(ChannelPermission::ManageMessages)?;
}
message.delete(db).await.map(|_| EmptyResponse)
message.delete(db, Some(amqp)).await.map(|_| EmptyResponse)
}

View File

@@ -1,8 +1,7 @@
use iso8601_timestamp::Timestamp;
use revolt_config::capture_error;
use revolt_database::{
tasks,
util::{permissions::DatabasePermissionQuery, reference::Reference},
Database, Message, PartialMessage, User,
AMQP, Database, Message, PartialMessage, User, tasks, util::{permissions::DatabasePermissionQuery, reference::Reference}
};
use revolt_models::v0::{self, Embed};
use revolt_permissions::{calculate_channel_permissions, ChannelPermission};
@@ -17,6 +16,7 @@ use validator::Validate;
#[patch("/<target>/messages/<msg>", data = "<edit>")]
pub async fn edit(
db: &State<Database>,
amqp: &State<AMQP>,
user: User,
target: Reference<'_>,
msg: Reference<'_>,
@@ -96,5 +96,10 @@ pub async fn edit(
}
}
if let Err(e) = amqp.edit_message_search(message.clone()).await {
log::error!("Error pushing message to RabbitMQ: {e}");
capture_error(&e);
}
Ok(Json(message.into_model(None, None)))
}

View File

@@ -1,3 +1,4 @@
use revolt_config::capture_error;
use revolt_database::{
util::{permissions::DatabasePermissionQuery, reference::Reference},
Channel, Database, PartialMessage, SystemMessage, User, AMQP,
@@ -46,6 +47,11 @@ pub async fn message_pin(
)
.await?;
if let Err(e) = amqp.edit_message_search(message.clone()).await {
log::error!("Error pushing message to RabbitMQ: {e}");
capture_error(&e);
}
SystemMessage::MessagePinned {
id: message.id.clone(),
by: user.id.clone(),

View File

@@ -10,6 +10,7 @@ use validator::Validate;
/// # Search for Messages
///
/// DEPRECATED: use /search instead.
/// This route searches for messages within the given parameters.
#[openapi(tag = "Messaging")]
#[post("/<target>/search", data = "<options>")]

View File

@@ -420,7 +420,7 @@ mod test {
// Delete the message
message
.clone()
.delete(&harness.db)
.delete(&harness.db, None)
.await
.expect("Failed to delete message");

View File

@@ -1,3 +1,4 @@
use revolt_config::capture_error;
use revolt_database::{
util::{permissions::DatabasePermissionQuery, reference::Reference},
Channel, Database, FieldsMessage, PartialMessage, SystemMessage, User, AMQP,
@@ -39,6 +40,11 @@ pub async fn message_unpin(
.update(db, PartialMessage::default(), vec![FieldsMessage::Pinned])
.await?;
if let Err(e) = amqp.edit_message_search(message.clone()).await {
log::error!("Error pushing message to RabbitMQ: {e}");
capture_error(&e);
}
SystemMessage::MessageUnpinned {
id: message.id.clone(),
by: user.id.clone(),

View File

@@ -35,7 +35,7 @@ pub async fn join(
channel, creator, ..
} => {
let mut channel = db.fetch_channel(channel).await?;
channel.add_user_to_group(db, amqp, &user, creator).await?;
channel.add_user_to_group(db, Some(amqp), &user, creator).await?;
if let Channel::Group { recipients, .. } = &channel {
Ok(Json(InviteJoinResponse::Group {
users: User::fetch_many_ids_as_mutuals(db, &user, recipients).await?,

View File

@@ -17,6 +17,7 @@ mod servers;
mod sync;
mod users;
mod webhooks;
mod search;
pub fn mount(config: Settings, mut rocket: Rocket<Build>) -> Rocket<Build> {
let settings = OpenApiSettings::default();
@@ -40,7 +41,8 @@ pub fn mount(config: Settings, mut rocket: Rocket<Build>) -> Rocket<Build> {
"/policy" => policy::routes(),
"/push" => push::routes(),
"/sync" => sync::routes(),
"/webhooks" => webhooks::routes()
"/webhooks" => webhooks::routes(),
"/search" => search::routes(),
};
} else {
mount_endpoints_and_merged_docs! {
@@ -60,7 +62,8 @@ pub fn mount(config: Settings, mut rocket: Rocket<Build>) -> Rocket<Build> {
"/onboard" => onboard::routes(),
"/policy" => policy::routes(),
"/push" => push::routes(),
"/sync" => sync::routes()
"/sync" => sync::routes(),
"/search" => search::routes(),
};
}
@@ -82,7 +85,8 @@ pub fn mount(config: Settings, mut rocket: Rocket<Build>) -> Rocket<Build> {
"/onboard" => onboard::routes(),
"/push" => push::routes(),
"/sync" => sync::routes(),
"/webhooks" => webhooks::routes()
"/webhooks" => webhooks::routes(),
"/search" => search::routes(),
};
} else {
mount_endpoints_and_merged_docs! {
@@ -101,7 +105,8 @@ pub fn mount(config: Settings, mut rocket: Rocket<Build>) -> Rocket<Build> {
"/auth/mfa" => rocket_authifier::routes::mfa::routes(),
"/onboard" => onboard::routes(),
"/push" => push::routes(),
"/sync" => sync::routes()
"/sync" => sync::routes(),
"/search" => search::routes(),
};
}

View File

@@ -0,0 +1,10 @@
use revolt_rocket_okapi::revolt_okapi::openapi3::OpenApi;
use rocket::Route;
mod search_messages;
pub fn routes() -> (Vec<Route>, OpenApi) {
openapi_get_routes_spec![
search_messages::search_messages
]
}

View File

@@ -0,0 +1,144 @@
use std::collections::HashSet;
use revolt_database::{util::permissions::DatabasePermissionQuery, Database, User};
use revolt_models::v0;
use revolt_permissions::{calculate_channel_permissions, ChannelPermission};
use revolt_result::{create_error, Result, ToRevoltError};
use revolt_search::{ElasticsearchClient, SearchTerms};
use rocket::{serde::json::Json, State};
#[openapi(tag = "Search")]
#[post("/", data = "<data>")]
pub async fn search_messages(
db: &State<Database>,
es: &State<ElasticsearchClient>,
user: User,
data: Json<v0::DataChannelMessagesSearch>,
) -> Result<Json<v0::BulkMessageResponse>> {
let data = data.into_inner();
if (data.server.is_some() && data.channel.is_some())
|| (data.server.is_none() && data.channel.is_none())
{
return Err(create_error!(InvalidOperation));
}
let (server_id, channels) = if let Some(server_id) = data.server {
let server = db.fetch_server(&server_id).await?;
let server_channels = db.fetch_channels(&server.channels).await?;
let mut channels = Vec::new();
let query = DatabasePermissionQuery::new(db, &user).server(&server);
for channel in server_channels {
let perms = calculate_channel_permissions(&mut query.clone().channel(&channel)).await;
if perms.has(ChannelPermission::ReadMessageHistory as u64) {
channels.push(channel.id().to_string())
}
}
(Some(server_id.clone()), channels)
} else {
let channel_id = data.channel.unwrap();
let channel = db.fetch_channel(&channel_id).await?;
let mut query = DatabasePermissionQuery::new(db, &user).channel(&channel);
let perms = calculate_channel_permissions(&mut query).await;
if !perms.has(ChannelPermission::ReadMessageHistory as u64) {
return Err(create_error!(InvalidOperation));
}
(channel.server().map(ToString::to_string), vec![channel_id])
};
let terms = SearchTerms {
channels,
filters: data.filters.map(Into::into).unwrap_or_default(),
offset: data.offset,
limit: data.limit.map(|limit| limit.min(100)),
sort: data.sort.map(Into::into),
};
let message_ids = es.search(terms).await.to_internal_error()?;
let mut messages = db
.fetch_messages_by_id(&message_ids)
.await?
.into_iter()
.map(|msg| msg.into_model(None, None))
.collect::<Vec<_>>();
// Fetching the messages looses the order so resort
let sort_order = data.sort.unwrap_or_default();
messages.sort_by(|a, b| match sort_order {
v0::SortOrder::Asc => a.id.cmp(&b.id),
v0::SortOrder::Desc => b.id.cmp(&a.id),
});
// TODO: abstract the user and member fetching logic out
let response = if let Some(true) = data.include_users {
let user_ids = messages
.iter()
.flat_map(|m| {
let mut users = vec![m.author.clone()];
if let Some(system) = &m.system {
match system {
v0::SystemMessage::ChannelDescriptionChanged { by } => {
users.push(by.clone())
}
v0::SystemMessage::ChannelIconChanged { by } => users.push(by.clone()),
v0::SystemMessage::ChannelOwnershipChanged { from, to, .. } => {
users.push(from.clone());
users.push(to.clone())
}
v0::SystemMessage::ChannelRenamed { by, .. } => users.push(by.clone()),
v0::SystemMessage::UserAdded { by, id, .. }
| v0::SystemMessage::UserRemove { by, id, .. } => {
users.push(by.clone());
users.push(id.clone());
}
v0::SystemMessage::UserBanned { id, .. }
| v0::SystemMessage::UserKicked { id, .. }
| v0::SystemMessage::UserJoined { id, .. }
| v0::SystemMessage::UserLeft { id, .. } => {
users.push(id.clone());
}
v0::SystemMessage::Text { .. } => {}
v0::SystemMessage::MessagePinned { by, .. } => {
users.push(by.clone());
}
v0::SystemMessage::MessageUnpinned { by, .. } => {
users.push(by.clone());
}
v0::SystemMessage::CallStarted { by, .. } => users.push(by.clone()),
}
}
users
})
.collect::<HashSet<String>>()
.into_iter()
.collect::<Vec<String>>();
let users = User::fetch_many_ids_as_mutuals(db, &user, &user_ids).await?;
v0::BulkMessageResponse::MessagesAndUsers {
messages,
users,
members: if let Some(server_id) = server_id {
Some(
db.fetch_members(&server_id, &user_ids)
.await?
.into_iter()
.map(Into::into)
.collect(),
)
} else {
None
},
}
} else {
v0::BulkMessageResponse::JustMessages(messages)
};
Ok(Json(response))
}

View File

@@ -4,7 +4,7 @@ use revolt_database::{
delete_voice_channel, get_user_voice_channel_in_server, remove_user_from_voice_channel,
UserVoiceChannel, VoiceClient,
},
Database, RemovalIntention, User,
AMQP, Database, RemovalIntention, User,
};
use revolt_models::v0;
use revolt_result::Result;
@@ -19,6 +19,7 @@ use rocket_empty::EmptyResponse;
#[delete("/<target>?<options..>")]
pub async fn delete(
db: &State<Database>,
amqp: &State<AMQP>,
voice_client: &State<VoiceClient>,
user: User,
target: Reference<'_>,
@@ -39,7 +40,7 @@ pub async fn delete(
.await?;
}
server.delete(db).await
server.delete(db, Some(amqp)).await
} else {
if let Some(channel_id) = get_user_voice_channel_in_server(&user.id, &server.id).await? {
remove_user_from_voice_channel(

View File

@@ -22,6 +22,6 @@ pub async fn add(
return Err(create_error!(IsBot));
}
user.add_friend(db, amqp, &mut target).await?;
user.add_friend(db, Some(amqp), &mut target).await?;
Ok(Json(target.into(db, &user).await))
}

View File

@@ -23,7 +23,7 @@ pub async fn send_friend_request(
return Err(create_error!(IsBot));
}
user.add_friend(db, amqp, &mut target).await?;
user.add_friend(db, Some(amqp), &mut target).await?;
Ok(Json(target.into(db, &user).await))
} else {
Err(create_error!(InvalidProperty))