feat: search for links and author type

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

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"

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

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,
}