jmug/stoatchat
1
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

feat: search for links and author type

Browse Source 
Signed-off-by: Zomatree <me@zomatree.live>
This commit is contained in:
Zomatree
2026-04-09 03:54:14 +01:00
parent b7ed543156
commit 9c62568132
27 changed files with 615 additions and 368 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
crates/core/database/src/amqp/amqp.rs
View File

@@ -326,10 +326,14 @@ impl AMQP {
.await
}
pub async fn new_message_search(&self, message: Message) -> Result<(), AMQPError> {
pub async fn new_message_search(
&self,
message: Message,
user: Option<User>,
) -> Result<(), AMQPError> {
let config = revolt_config::config().await;
let payload = to_string(&message).unwrap();
let payload = to_string(&MessageCreatePayload { message, user }).unwrap();
debug!(
"Sending new message search payload on channel {}: {}",
@@ -351,10 +355,14 @@ impl AMQP {
.await
}
pub async fn edit_message_search(&self, message: Message) -> Result<(), AMQPError> {
pub async fn edit_message_search(
&self,
message: Message,
user: Option<User>,
) -> Result<(), AMQPError> {
let config = revolt_config::config().await;
let payload = to_string(&message).unwrap();
let payload = to_string(&MessageEditPayload { message, user }).unwrap();
debug!(
"Sending edit message search payload on channel {}: {}",

20
crates/core/database/src/events/rabbit.rs
View File

@@ -3,7 +3,7 @@ use std::collections::HashMap;
use revolt_models::v0::PushNotification;
use serde::{Deserialize, Serialize};
use crate::User;
use crate::{Message, User};
#[derive(Serialize, Deserialize)]
pub struct MessageSentPayload {
@@ -81,10 +81,22 @@ pub struct AckPayload {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageDeletePayload {
pub message_id: String
pub message_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChannelDeletePayload {
pub channel_id: String
}
pub channel_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageCreatePayload {
pub message: Message,
pub user: Option<User>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageEditPayload {
pub message: Message,
pub user: Option<User>,
}

208
crates/core/database/src/models/messages/model.rs
View File

@@ -1,6 +1,6 @@
use indexmap::{IndexMap, IndexSet};
use iso8601_timestamp::Timestamp;
use revolt_config::{FeaturesLimits, capture_error, config};
use revolt_config::{capture_error, config, FeaturesLimits};
use revolt_models::v0::{
self, BulkMessageResponse, DataMessageSend, Embed, MessageAuthor, MessageFlags, MessageSort,
MessageWebhook, PushNotification, ReplyIntent, SendableEmbed, Text,
@@ -18,7 +18,7 @@ use crate::{
bulk_permissions::BulkDatabasePermissionQuery, idempotency::IdempotencyKey,
permissions::DatabasePermissionQuery,
},
Channel, Database, Emoji, File, User, AMQP,
Channel, Database, Emoji, File, Member, User, AMQP,
};
#[cfg(feature = "tasks")]
@@ -210,6 +210,13 @@ auto_derived!(
pub enum FieldsMessage {
Pinned,
}
/// Message along with the user for the author
pub struct MessageWithUser {
#[serde(flatten)]
pub message: Message,
pub user: Option<User>,
}
);
pub struct MessageFlagsValue(pub u32);
@@ -272,8 +279,8 @@ impl Message {
channel: Channel,
data: DataMessageSend,
author: MessageAuthor<'_>,
user: Option<v0::User>,
member: Option<v0::Member>,
user: Option<User>,
member: Option<Member>,
limits: FeaturesLimits,
mut idempotency: IdempotencyKey,
generate_embeds: bool,
@@ -610,10 +617,10 @@ impl Message {
/// Send a message without any notifications
pub async fn send_without_notifications(
&mut self,
&self,
db: &Database,
user: Option<v0::User>,
member: Option<v0::Member>,
user: Option<User>,
member: Option<Member>,
is_dm: bool,
generate_embeds: bool,
// This determines if this function should queue the mentions task or if somewhere else will.
@@ -623,9 +630,18 @@ impl Message {
db.insert_message(self).await?;
// Fan out events
EventV1::Message(self.clone().into_model(user, member))
.p(self.channel.to_string())
.await;
EventV1::Message(self.clone().into_model(
match user {
Some(user) => {
let is_online = revolt_presence::is_online(&user.id).await;
Some(user.into_known_static(is_online).await)
}
None => None,
},
member.map(Into::into),
))
.p(self.channel.to_string())
.await;
// Update last_message_id
#[cfg(feature = "tasks")]
@@ -673,8 +689,8 @@ impl Message {
db: &Database,
amqp: Option<&AMQP>, // this is optional mostly for tests.
author: MessageAuthor<'_>,
user: Option<v0::User>,
member: Option<v0::Member>,
user: Option<User>,
member: Option<Member>,
channel: &Channel,
generate_embeds: bool,
) -> Result<()> {
@@ -704,7 +720,17 @@ impl Message {
messages: vec![(
Some(
PushNotification::from(
self.clone().into_model(user, member),
self.clone().into_model(
match user.clone() {
Some(user) => {
let is_online =
revolt_presence::is_online(&user.id).await;
Some(user.into_known_static(is_online).await)
}
None => None,
},
member.map(Into::into),
),
Some(author),
channel.to_owned().into(),
)
@@ -727,7 +753,7 @@ impl Message {
}
if let Some(amqp) = amqp {
if let Err(e) = amqp.new_message_search(self.clone()).await {
if let Err(e) = amqp.new_message_search(self.clone(), user).await {
log::error!("Error pushing message to RabbitMQ: {e}");
capture_error(&e);
}
@@ -785,6 +811,7 @@ impl Message {
pub async fn update(
&mut self,
db: &Database,
amqp: Option<&AMQP>,
partial: PartialMessage,
remove: Vec<FieldsMessage>,
) -> Result<()> {
@@ -806,9 +833,74 @@ impl Message {
.p(self.channel.clone())
.await;
if let Some(amqp) = amqp {
if let Err(e) = amqp
.edit_message_search(self.clone(), self.fetch_author(db).await)
.await
{
log::error!("Error pushing message to RabbitMQ: {e}");
capture_error(&e);
}
}
Ok(())
}
pub async fn fetch_users(
db: &Database,
messages: &[Message],
server_id: Option<&str>,
) -> Result<(Vec<User>, Vec<Member>)> {
let user_ids = messages
.iter()
.flat_map(|m| {
let mut users = vec![m.author.clone()];
if let Some(system) = &m.system {
match system {
SystemMessage::ChannelDescriptionChanged { by } => users.push(by.clone()),
SystemMessage::ChannelIconChanged { by } => users.push(by.clone()),
SystemMessage::ChannelOwnershipChanged { from, to, .. } => {
users.push(from.clone());
users.push(to.clone())
}
SystemMessage::ChannelRenamed { by, .. } => users.push(by.clone()),
SystemMessage::UserAdded { by, id, .. }
| SystemMessage::UserRemove { by, id, .. } => {
users.push(by.clone());
users.push(id.clone());
}
SystemMessage::UserBanned { id, .. }
| SystemMessage::UserKicked { id, .. }
| SystemMessage::UserJoined { id, .. }
| SystemMessage::UserLeft { id, .. } => {
users.push(id.clone());
}
SystemMessage::Text { .. } => {}
SystemMessage::MessagePinned { by, .. } => {
users.push(by.clone());
}
SystemMessage::MessageUnpinned { by, .. } => {
users.push(by.clone());
}
SystemMessage::CallStarted { by, .. } => users.push(by.clone()),
}
}
users
})
.collect::<HashSet<String>>()
.into_iter()
.collect::<Vec<String>>();
let users = db.fetch_users(&user_ids).await?;
let members = if let Some(server_id) = server_id {
db.fetch_members(server_id, &user_ids).await?
} else {
Vec::new()
};
Ok((users, members))
}
/// Helper function to fetch many messages with users
pub async fn fetch_with_users(
db: &Database,
@@ -817,74 +909,26 @@ impl Message {
include_users: Option<bool>,
server_id: Option<&str>,
) -> Result<BulkMessageResponse> {
let messages: Vec<v0::Message> = db
.fetch_messages(query)
.await?
.into_iter()
.map(|msg| msg.into_model(None, None))
.collect();
let messages = db.fetch_messages(query).await?;
if let Some(true) = 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, perspective, &user_ids).await?;
let (users, members) = Message::fetch_users(db, &messages, server_id).await?;
Ok(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
},
messages: messages
.into_iter()
.map(|msg| msg.into_model(None, None))
.collect(),
users: User::into_mutuals(perspective, users).await,
members: Some(members.into_iter().map(Into::into).collect()),
})
} else {
Ok(BulkMessageResponse::JustMessages(messages))
Ok(BulkMessageResponse::JustMessages(
messages
.into_iter()
.map(|msg| msg.into_model(None, None))
.collect(),
))
}
}
@@ -898,7 +942,9 @@ impl Message {
) -> Result<()> {
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 {
let author = message.fetch_author(db).await;
if let Err(e) = amqp.edit_message_search(message, author).await {
log::error!("Error pushing message to RabbitMQ: {e}");
capture_error(&e);
}
@@ -1136,6 +1182,14 @@ impl Message {
FieldsMessage::Pinned => self.pinned = None,
}
}
pub async fn fetch_author(&self, db: &Database) -> Option<User> {
if self.webhook.is_some() {
None
} else {
db.fetch_user(&self.author).await.ok()
}
}
}
impl SystemMessage {

20
crates/core/database/src/models/messages/ops.rs
View File

@@ -1,8 +1,11 @@
use revolt_result::Result;
use std::collections::HashMap;
use std::time::SystemTime;
use revolt_result::Result;
use crate::{AppendMessage, FieldsMessage, Message, MessageQuery, PartialMessage, util::ChunkedDatabaseGenerator};
use crate::{
util::ChunkedDatabaseGenerator, AppendMessage, FieldsMessage, Message, MessageQuery,
MessageWithUser, PartialMessage,
};
#[cfg(feature = "mongodb")]
mod mongodb;
@@ -23,7 +26,12 @@ pub trait AbstractMessages: Sync + Send {
async fn fetch_messages_by_id(&self, ids: &[String]) -> Result<Vec<Message>>;
/// Update a given message with new information
async fn update_message(&self, id: &str, message: &PartialMessage, remove: Vec<FieldsMessage>) -> Result<()>;
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<Option<Message>>;
@@ -48,7 +56,9 @@ pub trait AbstractMessages: Sync + Send {
&self,
channels: &[String],
author: &str,
since: SystemTime
since: SystemTime,
) -> Result<HashMap<String, Vec<String>>>;
async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<Message>>;
/// Fetches all messages along with their author from every message in decending order
async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<MessageWithUser>>;
}

34
crates/core/database/src/models/messages/ops/mongodb.rs
View File

@@ -9,8 +9,8 @@ use std::time::SystemTime;
use ulid::Ulid;
use crate::{
AppendMessage, DocumentId, FieldsMessage, IntoDocumentPath, Message, MessageQuery,
MessageTimePeriod, MongoDb, PartialMessage, util::ChunkedDatabaseGenerator,
util::ChunkedDatabaseGenerator, AppendMessage, DocumentId, FieldsMessage, IntoDocumentPath,
Message, MessageQuery, MessageTimePeriod, MessageWithUser, MongoDb, PartialMessage,
};
use super::AbstractMessages;
@@ -417,7 +417,8 @@ impl AbstractMessages for MongoDb {
Ok(deleted_messages)
}
async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<Message>> {
/// Fetches all messages along with their author from every message in decending order
async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<MessageWithUser>> {
let mut session = self
.start_session()
.await
@@ -431,13 +432,34 @@ impl AbstractMessages for MongoDb {
let cursor = self
.col::<Message>(COL)
.find(doc! {})
.sort(doc ! { "_id": -1i32 })
.aggregate([
doc! {
"$lookup": {
"from": "users",
"localField": "author",
"foreignField": "_id",
"as": "user"
}
},
doc! {
"$set": {
"user": {
"$first": "$user"
}
}
},
doc! {
"$sort": {
"_id": -1
}
},
])
.with_type::<MessageWithUser>()
.session(&mut session)
.batch_size(1000)
.await
.inspect_err(|e| log::error!("{e}"))
.map_err(|_| create_database_error!("find", COL))?;
.map_err(|_| create_database_error!("aggregate", COL))?;
Ok(ChunkedDatabaseGenerator::new_mongo(session, cursor))
}

46
crates/core/database/src/models/messages/ops/reference.rs
View File

@@ -1,13 +1,13 @@
use std::collections::HashMap;
use futures::future::try_join_all;
use indexmap::IndexSet;
use revolt_result::Result;
use std::collections::HashMap;
use std::time::SystemTime;
use ulid::Ulid;
use crate::{
util::ChunkedDatabaseGenerator, AppendMessage, FieldsMessage, Message, MessageQuery,
PartialMessage, ReferenceDb,
MessageWithUser, PartialMessage, ReferenceDb,
};
use super::AbstractMessages;
@@ -64,7 +64,7 @@ impl AbstractMessages for ReferenceDb {
if let Some(pinned) = query.filter.pinned {
if message.pinned.unwrap_or_default() == pinned {
return false
return false;
}
}
@@ -195,7 +195,12 @@ impl AbstractMessages for ReferenceDb {
}
/// Update a given message with new information
async fn update_message(&self, id: &str, message: &PartialMessage, remove: Vec<FieldsMessage>) -> Result<()> {
async fn update_message(
&self,
id: &str,
message: &PartialMessage,
remove: Vec<FieldsMessage>,
) -> Result<()> {
let mut messages = self.messages.lock().await;
if let Some(message_data) = messages.get_mut(id) {
message_data.apply_options(message.to_owned());
@@ -300,7 +305,7 @@ impl AbstractMessages for ReferenceDb {
&self,
channels: &[String],
author: &str,
since: SystemTime
since: SystemTime,
) -> Result<HashMap<String, Vec<String>>> {
let threshold_ulid = Ulid::from_datetime(since).to_string();
let mut deleted_messages: HashMap<String, Vec<String>> = HashMap::new();
@@ -341,22 +346,31 @@ impl AbstractMessages for ReferenceDb {
}
// Delete the messages
self.messages
.lock()
.await
.retain(|id, message| {
let should_keep = !(message.author == author
&& channels.contains(&message.channel)
&& id.as_str() >= threshold_ulid.as_str());
should_keep
});
self.messages.lock().await.retain(|id, message| {
let should_keep = !(message.author == author
&& channels.contains(&message.channel)
&& id.as_str() >= threshold_ulid.as_str());
should_keep
});
Ok(deleted_messages)
}
async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<Message>> {
/// Fetches all messages along with their author from every message in decending order
async fn fetch_all_messages(&self) -> Result<ChunkedDatabaseGenerator<MessageWithUser>> {
let users = self.users.lock().await;
Ok(ChunkedDatabaseGenerator::new_reference(
self.messages.lock().await.values().cloned().collect(),
self.messages
.lock()
.await
.values()
.cloned()
.map(|message| MessageWithUser {
user: users.get(&message.author).cloned(),
message,
})
.collect(),
))
}
}

11
crates/core/database/src/models/users/model.rs
View File

@@ -362,6 +362,17 @@ impl User {
)
}
pub async fn into_mutuals(perspective: &User, users: Vec<User>) -> Vec<v0::User> {
let online_ids =
filter_online(&users.iter().map(|user| user.id.clone()).collect::<Vec<_>>()).await;
join_all(users.into_iter().map(|user| async {
let is_online = online_ids.contains(&user.id);
user.into_known(perspective, is_online).await
}))
.await
}
/// Find a free discriminator for a given username
pub async fn find_discriminator(
db: &Database,

27
crates/core/database/src/util/chunked.rs
View File

@@ -1,5 +1,6 @@
#[cfg(feature = "mongodb")]
use ::mongodb::{ClientSession, SessionCursor};
use revolt_result::{Result, ToRevoltError};
use serde::Deserialize;
#[derive(Debug)]
@@ -12,7 +13,7 @@ pub enum ChunkedDatabaseGenerator<T> {
},
Reference {
offset: i32,
offset: usize,
data: Vec<T>,
},
}
@@ -20,33 +21,25 @@ pub enum ChunkedDatabaseGenerator<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,
}
Self::MongoDb { session, cursor }
}
pub fn new_reference(data: Vec<T>) -> Self {
Self::Reference {
offset: 0,
data,
}
Self::Reference { offset: 0, data }
}
pub async fn next(&mut self) -> Option<T> {
pub async fn next(&mut self) -> Result<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"))
cursor.next(session).await.transpose().to_internal_error()
}
Self::Reference { offset, data } => {
if data.len() as i32 >= *offset {
None
} else {
let resp = &data[*offset as usize];
if let Some(value) = data.get(*offset) {
*offset += 1;
Some(resp.clone())
Ok(Some(value.clone()))
} else {
Ok(None)
}
}
}

34
crates/core/models/src/v0/search.rs
View File

@@ -2,44 +2,68 @@ use iso8601_timestamp::Timestamp;
use std::collections::HashSet;
auto_derived!(
/// Options for searching messages in a server or channel
pub struct DataChannelMessagesSearch {
/// Channel to search in
pub channel: Option<String>,
/// Server to search in
pub server: Option<String>,
pub offset: Option<u64>,
pub limit: Option<u64>,
/// Filter options
pub filters: Option<DataChannelMessagesSearchFilters>,
/// What index to start the search at
pub offset: Option<u64>,
/// Max amount of messages to return
pub limit: Option<u64>,
/// Sort order
pub sort: Option<SortOrder>,
pub include_users: Option<bool>,
}
/// Message search filters
pub struct DataChannelMessagesSearchFilters {
/// Message content
pub content: Option<String>,
/// Specific user
pub author: Option<HashSet<String>>,
/// Mentions a user
pub mentions: Option<HashSet<String>>,
/// Mentions a role
pub role_mentions: Option<HashSet<String>>,
/// Send before a specific date
pub before_date: Option<Timestamp>,
/// Sent after a specific date
pub after_date: Option<Timestamp>,
/// What type of user sent the message
pub author_type: Option<HashSet<AuthorType>>,
/// Whether the message is pinned or not
pub pinned: Option<bool>,
/// Require message to have a specific component type
pub components: Option<HashSet<MessageComponent>>,
}
/// Message author type
#[derive(Copy, Hash)]
pub enum AuthorType {
User,
// Bot,
Bot,
Webhook,
}
/// Message component
#[derive(Copy, Hash)]
pub enum MessageComponent {
Image,
Video,
// Link,
Link,
File,
Embed,
}
/// Message sort order
#[derive(Copy, Default)]
#[cfg_attr(feature = "serde", serde(rename_all = "lowercase"))]
pub enum SortOrder {

9
crates/core/search/Cargo.toml
View File

@@ -1,14 +1,14 @@
[package]
name = "revolt-search"
version = "0.11.5"
version = "0.12.0"
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" }
revolt-config = { version = "0.12.0", path = "../config" }
revolt-database = { version = "0.12.0", path = "../database" }
revolt-models = { version = "0.12.0", path = "../models" }
elasticsearch = "9.1.0-alpha.1"
serde_json = "1"
ulid = "1.2.1"
@@ -16,3 +16,4 @@ iso8601-timestamp = { version = "0.2.10", features = ["serde", "bson"] }
futures = "0.3.32"
elasticsearch-dsl = "0.4"
serde = "1"
linkify = "0.10.0"

135
crates/core/search/src/client.rs
View File

@@ -1,19 +1,25 @@
use std::fmt::Display;
use elasticsearch::{
BulkOperation, BulkParts, CreateParts, DeleteByQueryParts, DeleteParts, Elasticsearch, IndexParts, SearchParts, auth::Credentials, http::{
BulkOperation, BulkParts, CreateParts, DeleteByQueryParts, DeleteParts, Elasticsearch,
IndexParts, SearchParts,
auth::Credentials,
http::{
response::Exception,
transport::{SingleNodeConnectionPool, TransportBuilder},
}, indices::{IndicesCreateParts, IndicesDeleteParts}
},
indices::{IndicesCreateParts, IndicesDeleteParts},
};
use elasticsearch_dsl::{FieldSort, Query, Search, SearchResponse, Sort};
use revolt_database::{Database, Message};
use serde_json::{Map, Value, json};
use linkify::{LinkFinder, LinkKind};
use revolt_database::{Database, Message, MessageWithUser, User};
use serde_json::{Map, Value, json, to_value};
pub use elasticsearch;
use crate::{MessageComponent, MetadataFile, SearchTerms};
use crate::{AuthorType, MessageComponent, SearchTerms};
/// Elasticsearch errors
#[derive(Debug)]
pub enum Error {
Http(elasticsearch::Error),
@@ -42,6 +48,7 @@ impl Display for Error {
}
}
/// Higher level elasticsearch API more fit for our specific usecase
#[derive(Debug, Clone)]
pub struct ElasticsearchClient {
pub inner: Elasticsearch,
@@ -61,6 +68,7 @@ impl ElasticsearchClient {
Self { inner }
}
/// Delete messages index along with all documents
pub async fn delete_indexes(&self) -> Result<(), Error> {
let exception = self
.inner
@@ -78,6 +86,7 @@ impl ElasticsearchClient {
}
}
/// Create the messages index
pub async fn setup_indexes(&self) -> Result<(), Error> {
let exception = self
.inner
@@ -98,13 +107,14 @@ impl ElasticsearchClient {
},
"attachments": {
"type": "nested",
"dynamic": false,
"properties": {
"metadata.type": {
"type": "keyword"
}
}
},
// TODO: links
"has_link": { "type": "boolean" },
}
}
}))
@@ -120,6 +130,7 @@ impl ElasticsearchClient {
}
}
/// Performs a search for messages, returns a vec of message ids
pub async fn search(&self, terms: SearchTerms) -> Result<Vec<String>, Error> {
let mut query = Query::bool().filter(Query::terms("channel", terms.channels));
@@ -149,17 +160,32 @@ impl ElasticsearchClient {
if let Some(components) = terms.filters.components {
let mut components_query = Query::bool();
let mut attachments_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")),
MessageComponent::Image => {
attachments_query = attachments_query
.should(Query::term("attachments.metadata.type", "Image"))
}
MessageComponent::Video => {
attachments_query = attachments_query
.should(Query::term("attachments.metadata.type", "Video"))
}
MessageComponent::Link => {
components_query = components_query.should(Query::exists("has_link"))
}
MessageComponent::File => {
attachments_query = attachments_query.should(Query::exists("attachments"))
}
MessageComponent::Embed => {
components_query = components_query.should(Query::exists("embeds"))
}
};
}
query = query.filter(Query::nested("attachments", components_query));
query = query
.filter(components_query.should(Query::nested("attachments", attachments_query)));
}
let search = Search::new()
@@ -191,7 +217,13 @@ impl ElasticsearchClient {
}
}
fn create_message_source(&self, _db: &Database, message: Message) -> Value {
/// Creates a source for a message which can be stored and indexed into elasticsearch
fn create_message_source(
&self,
_db: &Database,
message: Message,
author: Option<User>,
) -> Value {
let mut map = Map::new();
map.insert("channel".to_string(), Value::String(message.channel));
@@ -199,25 +231,25 @@ impl ElasticsearchClient {
map.insert("author".to_string(), Value::String(message.author));
if let Some(content) = message.content {
// Is there a better way to handle this? can elasticsearch index links itself?
// Maybe in the future store the domains and be able to filter by that as well
let mut finder = LinkFinder::new();
finder.kinds(&[LinkKind::Url]);
if finder.links(&content).next().is_some() {
map.insert("has_link".to_string(), Value::Bool(true));
}
map.insert("content".to_string(), Value::String(content));
}
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,
})
}
// TODO: fetch the file metadata from FileHash because of File.metadata deprecation
// let metadata = attachment.as_hash(db).await.expect("Failed to fetch FileHash").metadata;
map.insert(
"attachments".to_string(),
serde_json::to_value(files).unwrap(),
serde_json::to_value(attachments).unwrap(),
);
}
@@ -243,31 +275,38 @@ impl ElasticsearchClient {
map.insert("pinned".to_string(), Value::Bool(pinned));
}
// TODO: bot
// This will turn bot author type to user author type if this is ran on a deleted message,
// due to the author not existing anymore so fetching will fail, this is probably niche enough
// to not really matter, might try fix in the future.
map.insert(
"author_type".to_string(),
Value::String(
if message.webhook.is_some() {
"webhook"
} else {
"user"
}
.to_string(),
),
to_value(if message.webhook.is_some() {
AuthorType::Webhook
} else if author.is_some_and(|user| user.bot.is_some()) {
AuthorType::Bot
} else {
AuthorType::User
})
.unwrap(),
);
Value::Object(map)
}
pub async fn bulk_index_messages(&self, db: &Database, messages: Vec<Message>) -> Result<(), Error> {
/// Bulk uploads and indexes messages to elasticsearch
pub async fn bulk_index_messages(
&self,
db: &Database,
messages: Vec<MessageWithUser>,
) -> Result<(), Error> {
let mut ops = Vec::<BulkOperation<Value>>::new();
for message in messages {
let id = message.id.clone();
let source = self.create_message_source(db, message);
let id = message.message.id.clone();
let source = self.create_message_source(db, message.message, message.user);
ops.push(BulkOperation::create(source).id(id).into());
};
}
let exception = self
.inner
@@ -285,9 +324,15 @@ impl ElasticsearchClient {
}
}
pub async fn index_message(&self, db: &Database, message: Message) -> Result<(), Error> {
/// Uploads and indexes a single message to elasticsearch
pub async fn index_message(
&self,
db: &Database,
message: Message,
author: Option<User>,
) -> Result<(), Error> {
let id = message.id.clone();
let source = self.create_message_source(db, message);
let source = self.create_message_source(db, message, author);
let exception = self
.inner
@@ -305,9 +350,15 @@ impl ElasticsearchClient {
}
}
pub async fn edit_message(&self, db: &Database, message: Message) -> Result<(), Error> {
/// Updates or upserts an existing message to elasticsearch
pub async fn edit_message(
&self,
db: &Database,
message: Message,
author: Option<User>,
) -> Result<(), Error> {
let id = message.id.clone();
let source = self.create_message_source(db, message);
let source = self.create_message_source(db, message, author);
let exception = self
.inner
@@ -325,6 +376,7 @@ impl ElasticsearchClient {
}
}
/// Deletes a message from elasticsearch
pub async fn delete_message(&self, message_id: &str) -> Result<(), Error> {
let exception = self
.inner
@@ -341,6 +393,7 @@ impl ElasticsearchClient {
}
}
/// Deletes all messages in a channel from elasticsearch
pub async fn delete_channel(&self, channel_id: &str) -> Result<(), Error> {
let exception = self
.inner

47
crates/core/search/src/types.rs
View File

@@ -1,39 +1,54 @@
use std::collections::HashSet;
use iso8601_timestamp::Timestamp;
use revolt_database::{File, Metadata};
use revolt_models::v0;
use serde::{Deserialize, Serialize};
use serde::Serialize;
/// Message author type
#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Hash)]
pub enum AuthorType {
User,
// Bot,
Bot,
Webhook,
}
/// Message component
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum MessageComponent {
Image,
Video,
// Link,
Link,
File,
Embed,
}
/// Message search filters
#[derive(Debug, Clone, Default, PartialEq)]
pub struct SearchFilters {
/// Message content
pub content: Option<String>,
/// Specific user
pub author: Option<HashSet<String>>,
/// Mentions a user
pub mentions: Option<HashSet<String>>,
/// Mentions a role
pub role_mentions: Option<HashSet<String>>,
/// Send before a specific date
pub before_date: Option<Timestamp>,
/// Sent after a specific date
pub after_date: Option<Timestamp>,
/// What type of user sent the message
pub author_type: Option<HashSet<AuthorType>>,
/// Whether the message is pinned or not
pub pinned: Option<bool>,
/// Require message to have a specific component type
pub components: Option<HashSet<MessageComponent>>,
}
/// Message sort order
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum SortOrder {
Asc,
@@ -41,26 +56,28 @@ pub enum SortOrder {
Desc,
}
/// Options for searching messages in a server or channel
#[derive(Debug, Clone, PartialEq)]
pub struct SearchTerms {
/// Channels to search in
pub channels: Vec<String>,
pub filters: SearchFilters,
pub offset: Option<u64>,
pub limit: Option<u64>,
pub sort: Option<SortOrder>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MetadataFile {
pub file: File,
pub metadata: Metadata,
/// Filter options
pub filters: SearchFilters,
/// What index to start the search at
pub offset: Option<u64>,
/// Max amount of messages to return
pub limit: Option<u64>,
/// Sort order
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::Bot => AuthorType::Bot,
v0::AuthorType::Webhook => AuthorType::Webhook,
}
}
@@ -71,7 +88,7 @@ impl From<v0::MessageComponent> for MessageComponent {
match value {
v0::MessageComponent::Image => MessageComponent::Image,
v0::MessageComponent::Video => MessageComponent::Video,
// v0::MessageComponent::Link => MessageComponent::Link,
v0::MessageComponent::Link => MessageComponent::Link,
v0::MessageComponent::File => MessageComponent::File,
v0::MessageComponent::Embed => MessageComponent::Embed,
}