use crate::RankedMap; use chrono::{DateTime, Utc}; use heed::types::{ByteSlice, OwnedType, SerdeBincode, Str}; use heed::Result as ZResult; use meilisearch_schema::Schema; use std::collections::HashMap; use std::sync::Arc; const CREATED_AT_KEY: &str = "created-at"; const CUSTOMS_KEY: &str = "customs-key"; const FIELDS_FREQUENCY_KEY: &str = "fields-frequency"; const NAME_KEY: &str = "name"; const NUMBER_OF_DOCUMENTS_KEY: &str = "number-of-documents"; const RANKED_MAP_KEY: &str = "ranked-map"; const SCHEMA_KEY: &str = "schema"; const STOP_WORDS_KEY: &str = "stop-words"; const SYNONYMS_KEY: &str = "synonyms"; const UPDATED_AT_KEY: &str = "updated-at"; const WORDS_KEY: &str = "words"; pub type FreqsMap = HashMap; type SerdeFreqsMap = SerdeBincode; type SerdeDatetime = SerdeBincode>; #[derive(Copy, Clone)] pub struct Main { pub(crate) main: heed::PolyDatabase, } impl Main { pub fn clear(self, writer: &mut heed::RwTxn) -> ZResult<()> { self.main.clear(writer) } pub fn put_name(self, writer: &mut heed::RwTxn, name: &str) -> ZResult<()> { self.main.put::(writer, NAME_KEY, name) } pub fn name(self, reader: &heed::RoTxn) -> ZResult> { Ok(self .main .get::(reader, NAME_KEY)? .map(|name| name.to_owned())) } pub fn put_created_at(self, writer: &mut heed::RwTxn) -> ZResult<()> { self.main .put::(writer, CREATED_AT_KEY, &Utc::now()) } pub fn created_at(self, reader: &heed::RoTxn) -> ZResult>> { self.main.get::(reader, CREATED_AT_KEY) } pub fn put_updated_at(self, writer: &mut heed::RwTxn) -> ZResult<()> { self.main .put::(writer, UPDATED_AT_KEY, &Utc::now()) } pub fn updated_at(self, reader: &heed::RoTxn) -> ZResult>> { self.main.get::(reader, UPDATED_AT_KEY) } pub fn put_words_fst(self, writer: &mut heed::RwTxn, fst: &fst::Set) -> ZResult<()> { let bytes = fst.as_fst().as_bytes(); self.main.put::(writer, WORDS_KEY, bytes) } pub fn words_fst(self, reader: &heed::RoTxn) -> ZResult> { match self.main.get::(reader, WORDS_KEY)? { Some(bytes) => { let len = bytes.len(); let bytes = Arc::new(bytes.to_owned()); let fst = fst::raw::Fst::from_shared_bytes(bytes, 0, len).unwrap(); Ok(Some(fst::Set::from(fst))) } None => Ok(None), } } pub fn put_schema(self, writer: &mut heed::RwTxn, schema: &Schema) -> ZResult<()> { self.main .put::>(writer, SCHEMA_KEY, schema) } pub fn schema(self, reader: &heed::RoTxn) -> ZResult> { self.main .get::>(reader, SCHEMA_KEY) } pub fn put_ranked_map(self, writer: &mut heed::RwTxn, ranked_map: &RankedMap) -> ZResult<()> { self.main .put::>(writer, RANKED_MAP_KEY, &ranked_map) } pub fn ranked_map(self, reader: &heed::RoTxn) -> ZResult> { self.main .get::>(reader, RANKED_MAP_KEY) } pub fn put_synonyms_fst(self, writer: &mut heed::RwTxn, fst: &fst::Set) -> ZResult<()> { let bytes = fst.as_fst().as_bytes(); self.main.put::(writer, SYNONYMS_KEY, bytes) } pub fn synonyms_fst(self, reader: &heed::RoTxn) -> ZResult> { match self.main.get::(reader, SYNONYMS_KEY)? { Some(bytes) => { let len = bytes.len(); let bytes = Arc::new(bytes.to_owned()); let fst = fst::raw::Fst::from_shared_bytes(bytes, 0, len).unwrap(); Ok(Some(fst::Set::from(fst))) } None => Ok(None), } } pub fn put_stop_words_fst(self, writer: &mut heed::RwTxn, fst: &fst::Set) -> ZResult<()> { let bytes = fst.as_fst().as_bytes(); self.main .put::(writer, STOP_WORDS_KEY, bytes) } pub fn stop_words_fst(self, reader: &heed::RoTxn) -> ZResult> { match self.main.get::(reader, STOP_WORDS_KEY)? { Some(bytes) => { let len = bytes.len(); let bytes = Arc::new(bytes.to_owned()); let fst = fst::raw::Fst::from_shared_bytes(bytes, 0, len).unwrap(); Ok(Some(fst::Set::from(fst))) } None => Ok(None), } } pub fn put_number_of_documents(self, writer: &mut heed::RwTxn, f: F) -> ZResult where F: Fn(u64) -> u64, { let new = self.number_of_documents(writer).map(f)?; self.main .put::>(writer, NUMBER_OF_DOCUMENTS_KEY, &new)?; Ok(new) } pub fn number_of_documents(self, reader: &heed::RoTxn) -> ZResult { match self .main .get::>(reader, NUMBER_OF_DOCUMENTS_KEY)? { Some(value) => Ok(value), None => Ok(0), } } pub fn put_fields_frequency( self, writer: &mut heed::RwTxn, fields_frequency: &FreqsMap, ) -> ZResult<()> { self.main .put::(writer, FIELDS_FREQUENCY_KEY, fields_frequency) } pub fn fields_frequency(&self, reader: &heed::RoTxn) -> ZResult> { match self .main .get::(reader, FIELDS_FREQUENCY_KEY)? { Some(freqs) => Ok(Some(freqs)), None => Ok(None), } } pub fn put_customs(self, writer: &mut heed::RwTxn, customs: &[u8]) -> ZResult<()> { self.main .put::(writer, CUSTOMS_KEY, customs) } pub fn customs<'txn>(self, reader: &'txn heed::RoTxn) -> ZResult> { self.main.get::(reader, CUSTOMS_KEY) } }