use std::fs::File; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; use anyhow::Result; use chrono::{DateTime, Utc}; use dashmap::DashMap; use dashmap::mapref::one::Ref; use heed::types::{Str, SerdeBincode}; use heed::{EnvOpenOptions, Env, Database}; use milli::{Index, FieldsIdsMap, SearchResult, FieldId}; use serde::{Serialize, Deserialize}; use crate::data::SearchQuery; const CONTROLLER_META_FILENAME: &str = "index_controller_meta"; const INDEXES_CONTROLLER_FILENAME: &str = "indexes_db"; const INDEXES_DB_NAME: &str = "indexes_db"; pub trait UpdateStore {} pub struct IndexController { path: PathBuf, update_store: U, env: Env, indexes_db: Database>, indexes: DashMap, } #[derive(Debug, Serialize, Deserialize)] struct IndexControllerMeta { open_options: EnvOpenOptions, created_at: DateTime, } impl IndexControllerMeta { fn from_path(path: impl AsRef) -> Result> { let mut path = path.as_ref().to_path_buf(); path.push(CONTROLLER_META_FILENAME); if path.exists() { let mut file = File::open(path)?; let mut buffer = Vec::new(); let n = file.read_to_end(&mut buffer)?; let meta: IndexControllerMeta = serde_json::from_slice(&buffer[..n])?; Ok(Some(meta)) } else { Ok(None) } } fn to_path(self, path: impl AsRef) -> Result<()> { let mut path = path.as_ref().to_path_buf(); path.push(CONTROLLER_META_FILENAME); if path.exists() { Err(anyhow::anyhow!("Index controller metadata already exists")) } else { let mut file = File::create(path)?; let json = serde_json::to_vec(&self)?; file.write_all(&json)?; Ok(()) } } } #[derive(Debug, Serialize, Deserialize)] struct IndexMetadata { created_at: DateTime, open_options: EnvOpenOptions, id: String, } impl IndexMetadata { fn open_index(&self, path: impl AsRef) -> Result { // create a path in the form "db_path/indexes/index_id" let mut path = path.as_ref().to_path_buf(); path.push("indexes"); path.push(&self.id); Ok(Index::new(self.open_options, path)?) } } struct IndexView<'a, U> { txn: heed::RoTxn<'a>, index: Ref<'a, String, Index>, update_store: &'a U, } impl<'a, U: UpdateStore> IndexView<'a, U> { pub fn search(&self, search_query: SearchQuery) -> Result { let mut search = self.index.search(&self.txn); if let Some(query) = &search_query.q { search.query(query); } if let Some(offset) = search_query.offset { search.offset(offset); } let limit = search_query.limit; search.limit(limit); Ok(search.execute()?) } pub fn fields_ids_map(&self) -> Result { Ok(self.index.fields_ids_map(&self.txn)?) } pub fn fields_displayed_fields_ids(&self) -> Result>> { Ok(self.index.displayed_fields_ids(&self.txn)?) } pub fn documents(&self, ids: Vec) -> Result)>> { Ok(self.index.documents(&self.txn, ids)?) } } impl IndexController { /// Open the index controller from meta found at path, and create a new one if no meta is /// found. pub fn new(path: impl AsRef, update_store: U) -> Result { // If index controller metadata is present, we return the env, otherwise, we create a new // metadata from scratch before returning a new env. let path = path.as_ref().to_path_buf(); let env = match IndexControllerMeta::from_path(&path)? { Some(meta) => meta.open_options.open(INDEXES_CONTROLLER_FILENAME)?, None => { let open_options = EnvOpenOptions::new() .map_size(page_size::get() * 1000); let env = open_options.open(INDEXES_CONTROLLER_FILENAME)?; let created_at = Utc::now(); let meta = IndexControllerMeta { open_options: open_options.clone(), created_at }; meta.to_path(path)?; env } }; let indexes = DashMap::new(); let indexes_db = match env.open_database(Some(INDEXES_DB_NAME))? { Some(indexes_db) => indexes_db, None => env.create_database(Some(INDEXES_DB_NAME))?, }; Ok(Self { env, indexes, indexes_db, update_store, path }) } pub fn get_or_create>(&mut self, name: S) -> Result> { todo!() } /// Get an index with read access to the db. The index are lazily loaded, meaning that we first /// check for its exixtence in the indexes map, and if it doesn't exist, the index db is check /// for metadata to launch the index. pub fn get>(&self, name: S) -> Result>> { let update_store = &self.update_store; match self.indexes.get(name.as_ref()) { Some(index) => { let txn = index.read_txn()?; Ok(Some(IndexView { index, update_store, txn })) } None => { let txn = self.env.read_txn()?; match self.indexes_db.get(&txn, name.as_ref())? { Some(meta) => { let index = meta.open_index(self.path)?; self.indexes.insert(name.as_ref().to_owned(), index); // TODO: create index view match self.indexes.get(name.as_ref()) { Some(index) => { let txn = index.read_txn()?; Ok(Some(IndexView { index, txn, update_store })) } None => Ok(None) } } None => Ok(None) } } } } pub fn get_mut>(&self, name: S) -> Result>> { todo!() } pub async fn delete_index>(&self, name:S) -> Result<()> { todo!() } pub async fn list_indices(&self) -> Result> { todo!() } pub async fn rename_index(&self, old: &str, new: &str) -> Result<()> { todo!() } }