#![forbid(unsafe_code)]
#![warn(missing_docs)]
pub mod ft;
pub mod ft_filter;
pub mod query_fsm;
pub mod registry;
pub mod schema;
pub mod wire;
use std::sync::Arc;
use dynomite::embed::{CommandExtension, HsetOutcome, ServerBuilder};
use dynomite::msg::MsgType;
pub use crate::registry::{
RegistryError, TextFieldIndex, TextHit, TextRegexApproxResult, TextRegexResult, VectorRegistry,
VectorTable, VectorTableInfo,
};
pub use crate::schema::{
DistanceMetric, IndexAlgorithm, MetadataField, MetadataFieldType, VectorSchema, VectorType,
};
#[derive(Clone, Debug)]
pub struct SearchExtension {
registry: Arc<VectorRegistry>,
}
impl SearchExtension {
#[must_use]
pub fn new(registry: Arc<VectorRegistry>) -> Self {
Self { registry }
}
#[must_use]
pub fn registry(&self) -> &Arc<VectorRegistry> {
&self.registry
}
}
impl Default for SearchExtension {
fn default() -> Self {
Self {
registry: Arc::new(VectorRegistry::new()),
}
}
}
impl CommandExtension for SearchExtension {
fn handles_msg_type(&self, ty: MsgType) -> bool {
matches!(
ty,
MsgType::ReqRedisFtCreate
| MsgType::ReqRedisFtSearch
| MsgType::ReqRedisFtInfo
| MsgType::ReqRedisFtList
| MsgType::ReqRedisFtDropindex
| MsgType::ReqRedisFtRegex
| MsgType::ReqRedisFtUnknown
)
}
fn try_dispatch(&self, args: &[&[u8]]) -> Option<Vec<u8>> {
Some(crate::ft::dispatch(&self.registry, args))
}
fn try_intercept_hset(&self, args: &[&[u8]]) -> HsetOutcome {
match crate::ft::maybe_index_hset(&self.registry, args) {
Ok(Some(_)) => HsetOutcome::Absorbed,
Ok(None) => HsetOutcome::NotIndexed,
Err(e) => HsetOutcome::Error(format!("{e}")),
}
}
}
pub fn install(builder: &mut ServerBuilder) -> Arc<VectorRegistry> {
let ext = SearchExtension::default();
let registry = Arc::clone(ext.registry());
builder.set_command_extension(Arc::new(ext));
registry
}
#[must_use]
pub fn install_owned(builder: ServerBuilder) -> (ServerBuilder, Arc<VectorRegistry>) {
let ext = SearchExtension::default();
let registry = Arc::clone(ext.registry());
let builder = builder.with_command_extension(Arc::new(ext));
(builder, registry)
}