use crate::error::{Error, Result};
use crate::search::ast::{
explain_search_query, explain_search_query_cli, parse_search_query_with_params,
};
use crate::search::conf::{FtCreate, FtInfo, FtSearch, SearchResult, SuggestionItem};
use crate::search::index::InvertedIndex;
use crate::search::meta::SearchIndexSchema;
use crate::search::sug::SuggestionDict;
use hipstr::HipStr;
use rapidhash::RapidHashMap;
#[derive(Debug, Clone, Default)]
pub struct SearchIndexManager {
pub indexes: RapidHashMap<HipStr<'static>, (SearchIndexSchema, InvertedIndex)>,
pub aliases: RapidHashMap<HipStr<'static>, HipStr<'static>>,
pub suggestions: RapidHashMap<HipStr<'static>, SuggestionDict>,
pub configs: RapidHashMap<String, String>,
}
impl SearchIndexManager {
pub fn new() -> Self {
let mut configs = RapidHashMap::default();
configs.insert("TIMEOUT".to_string(), "500".to_string());
configs.insert("DEFAULT_DIALECT".to_string(), "2".to_string());
configs.insert("MINPREFIX".to_string(), "2".to_string());
configs.insert("MAXEXPANSIONS".to_string(), "200".to_string());
Self {
indexes: RapidHashMap::default(),
aliases: RapidHashMap::default(),
suggestions: RapidHashMap::default(),
configs,
}
}
pub fn create_index(&mut self, schema: SearchIndexSchema) -> Result<()> {
let name = schema.name.clone();
if self.indexes.contains_key(&name) {
return Err(Error::invalid_data(format!(
"Index '{name}' already exists"
)));
}
self.indexes.insert(name, (schema, InvertedIndex::new()));
Ok(())
}
pub fn create_index_from_opts(&mut self, opts: FtCreate) -> Result<()> {
self.create_index(SearchIndexSchema::from(opts))
}
pub fn drop_index(&mut self, index_name: &str, dd: bool) -> Result<Vec<HipStr<'static>>> {
let resolved = self.resolve_index_name(index_name).to_string();
let name_to_remove = HipStr::from(resolved);
let removed = self.indexes.remove(&name_to_remove);
if let Some((_, inverted)) = removed {
self.aliases
.retain(|_, target| target.as_str() != name_to_remove.as_str());
if dd {
let doc_ids: Vec<HipStr<'static>> = inverted.docs.keys().cloned().collect();
Ok(doc_ids)
} else {
Ok(Vec::new())
}
} else {
Err(Error::invalid_data(format!("Unknown index '{index_name}'")))
}
}
pub fn get_index(&self, index_name: &str) -> Option<&(SearchIndexSchema, InvertedIndex)> {
let resolved = self.resolve_index_name(index_name);
self.indexes.get(resolved)
}
pub fn get_index_mut(
&mut self,
index_name: &str,
) -> Option<&mut (SearchIndexSchema, InvertedIndex)> {
let resolved = self.resolve_index_name(index_name);
let key = HipStr::from(resolved);
self.indexes.get_mut(&key)
}
pub fn list_indexes(&self) -> Vec<String> {
let mut list: Vec<String> = self.indexes.keys().map(|k| k.to_string()).collect();
list.sort();
list
}
pub fn resolve_index_name<'a>(&'a self, alias_or_name: &'a str) -> &'a str {
if let Some(target) = self.aliases.get(alias_or_name) {
target.as_str()
} else {
alias_or_name
}
}
pub fn add_alias(&mut self, alias: &str, index_name: &str) -> Result<()> {
let alias_key = HipStr::from(alias);
if self.aliases.contains_key(&alias_key) {
return Err(Error::invalid_data(format!(
"Alias '{alias}' already exists"
)));
}
if !self.indexes.contains_key(index_name) {
return Err(Error::invalid_data(format!(
"Index '{index_name}' not found"
)));
}
self.aliases.insert(alias_key, HipStr::from(index_name));
Ok(())
}
pub fn del_alias(&mut self, alias: &str) -> Result<()> {
let alias_key = HipStr::from(alias);
if self.aliases.remove(&alias_key).is_none() {
return Err(Error::invalid_data(format!("Alias '{alias}' not found")));
}
Ok(())
}
pub fn update_alias(&mut self, alias: &str, index_name: &str) -> Result<()> {
if !self.indexes.contains_key(index_name) {
return Err(Error::invalid_data(format!(
"Index '{index_name}' not found"
)));
}
self.aliases
.insert(HipStr::from(alias), HipStr::from(index_name));
Ok(())
}
pub fn tag_vals(&self, index_name: &str, field_name: &str) -> Result<Vec<String>> {
if let Some((_, idx)) = self.get_index(index_name) {
Ok(idx.tag_vals(field_name))
} else {
Err(Error::invalid_data(format!(
"Index '{index_name}' not found"
)))
}
}
pub fn search(&self, index_name: &str, query: &str, opts: &FtSearch) -> Result<SearchResult> {
if let Some((schema, idx)) = self.get_index(index_name) {
idx.search(schema, query, opts)
} else {
Err(Error::invalid_data(format!(
"Index '{index_name}' not found"
)))
}
}
pub fn explain(&self, query: &str, opts: &FtSearch) -> String {
let ast = parse_search_query_with_params(query, &opts.params);
explain_search_query(&ast)
}
pub fn explain_cli(&self, query: &str, opts: &FtSearch) -> String {
let ast = parse_search_query_with_params(query, &opts.params);
explain_search_query_cli(&ast)
}
pub fn info(&self, index_name: &str) -> Result<FtInfo> {
if let Some((schema, idx)) = self.get_index(index_name) {
Ok(idx.info(schema))
} else {
Err(Error::invalid_data(format!(
"Index '{index_name}' not found"
)))
}
}
pub fn config_get(&self, option: &str) -> Result<String> {
let opt_upper = option.to_ascii_uppercase();
if let Some(val) = self.configs.get(&opt_upper) {
Ok(val.clone())
} else {
Err(Error::invalid_data(format!(
"No such configuration option '{option}'"
)))
}
}
pub fn config_set(&mut self, option: &str, value: &str) -> Result<()> {
let opt_upper = option.to_ascii_uppercase();
self.configs.insert(opt_upper, value.to_string());
Ok(())
}
pub fn config_help(&self, option: &str) -> Result<String> {
let opt_upper = option.to_ascii_uppercase();
match opt_upper.as_str() {
"TIMEOUT" => Ok("Query execution timeout in milliseconds".to_string()),
"DEFAULT_DIALECT" => Ok("Default RediSearch dialect version".to_string()),
"MINPREFIX" => Ok("Minimum prefix length for wildcard expansion".to_string()),
"MAXEXPANSIONS" => Ok("Maximum number of prefix expansions".to_string()),
_ => Ok(format!("Help information for {option}")),
}
}
pub fn sug_add(
&mut self,
key: &str,
string: &str,
score: f64,
incr: bool,
payload: Option<String>,
) -> usize {
let dict = self.suggestions.entry(HipStr::from(key)).or_default();
dict.sug_add(string, score, incr, payload)
}
pub fn sug_get(
&self,
key: &str,
prefix: &str,
fuzzy: bool,
withscores: bool,
withpayloads: bool,
max: Option<usize>,
) -> Vec<SuggestionItem> {
if let Some(dict) = self.suggestions.get(key) {
dict.sug_get(prefix, fuzzy, withscores, withpayloads, max)
} else {
Vec::new()
}
}
pub fn sug_del(&mut self, key: &str, string: &str) -> bool {
if let Some(dict) = self.suggestions.get_mut(key) {
dict.sug_del(string)
} else {
false
}
}
pub fn sug_len(&self, key: &str) -> usize {
if let Some(dict) = self.suggestions.get(key) {
dict.sug_len()
} else {
0
}
}
}