feat: initial work on elasticsearch integration
Signed-off-by: Zomatree <me@zomatree.live>
This commit is contained in:
Zomatree
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(())
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user