create workspace with meilisearch-error

meilisearch-http/src/data/mod.rs

@@ -0,0 +1,147 @@
+mod search;
+mod updates;
+
+pub use search::{SearchQuery, SearchResult, DEFAULT_SEARCH_LIMIT};
+
+use std::fs::create_dir_all;
+use std::ops::Deref;
+use std::sync::Arc;
+
+use sha2::Digest;
+
+use crate::index_controller::{IndexController, LocalIndexController, IndexMetadata, Settings, IndexSettings};
+use crate::option::Opt;
+
+#[derive(Clone)]
+pub struct Data {
+    inner: Arc<DataInner>,
+}
+
+impl Deref for Data {
+    type Target = DataInner;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+#[derive(Clone)]
+pub struct DataInner {
+    pub index_controller: Arc<LocalIndexController>,
+    pub api_keys: ApiKeys,
+    options: Opt,
+}
+
+#[derive(Clone)]
+pub struct ApiKeys {
+    pub public: Option<String>,
+    pub private: Option<String>,
+    pub master: Option<String>,
+}
+
+impl ApiKeys {
+    pub fn generate_missing_api_keys(&mut self) {
+        if let Some(master_key) = &self.master {
+            if self.private.is_none() {
+                let key = format!("{}-private", master_key);
+                let sha = sha2::Sha256::digest(key.as_bytes());
+                self.private = Some(format!("{:x}", sha));
+            }
+            if self.public.is_none() {
+                let key = format!("{}-public", master_key);
+                let sha = sha2::Sha256::digest(key.as_bytes());
+                self.public = Some(format!("{:x}", sha));
+            }
+        }
+    }
+}
+
+impl Data {
+    pub fn new(options: Opt) -> anyhow::Result<Data> {
+        let path = options.db_path.clone();
+        let indexer_opts = options.indexer_options.clone();
+        create_dir_all(&path)?;
+        let index_controller = LocalIndexController::new(
+            &path,
+            indexer_opts,
+            options.max_mdb_size.get_bytes(),
+            options.max_udb_size.get_bytes(),
+        )?;
+        let index_controller = Arc::new(index_controller);
+
+        let mut api_keys = ApiKeys {
+            master: options.clone().master_key,
+            private: None,
+            public: None,
+        };
+
+        api_keys.generate_missing_api_keys();
+
+        let inner = DataInner { index_controller, options, api_keys };
+        let inner = Arc::new(inner);
+
+        Ok(Data { inner })
+    }
+
+    pub fn settings<S: AsRef<str>>(&self, index_uid: S) -> anyhow::Result<Settings> {
+        let index = self.index_controller
+            .index(&index_uid)?
+            .ok_or_else(|| anyhow::anyhow!("Index {} does not exist.", index_uid.as_ref()))?;
+
+        let txn = index.read_txn()?;
+
+        let displayed_attributes = index
+            .displayed_fields(&txn)?
+            .map(|fields| fields.into_iter().map(String::from).collect())
+            .unwrap_or_else(|| vec!["*".to_string()]);
+
+        let searchable_attributes = index
+            .searchable_fields(&txn)?
+            .map(|fields| fields.into_iter().map(String::from).collect())
+            .unwrap_or_else(|| vec!["*".to_string()]);
+
+        let faceted_attributes = index
+            .faceted_fields(&txn)?
+            .into_iter()
+            .map(|(k, v)| (k, v.to_string()))
+            .collect();
+
+        Ok(Settings {
+            displayed_attributes: Some(Some(displayed_attributes)),
+            searchable_attributes: Some(Some(searchable_attributes)),
+            faceted_attributes: Some(Some(faceted_attributes)),
+            criteria: None,
+        })
+    }
+
+    pub fn list_indexes(&self) -> anyhow::Result<Vec<IndexMetadata>> {
+        self.index_controller.list_indexes()
+    }
+
+    pub fn index(&self, name: impl AsRef<str>) -> anyhow::Result<Option<IndexMetadata>> {
+        Ok(self
+            .list_indexes()?
+            .into_iter()
+            .find(|i| i.uid == name.as_ref()))
+    }
+
+    pub fn create_index(&self, name: impl AsRef<str>, primary_key: Option<impl AsRef<str>>) -> anyhow::Result<IndexMetadata> {
+        let settings = IndexSettings {
+            name: Some(name.as_ref().to_string()),
+            primary_key: primary_key.map(|s| s.as_ref().to_string()),
+        };
+
+        let meta = self.index_controller.create_index(settings)?;
+        Ok(meta)
+    }
+
+    #[inline]
+    pub fn http_payload_size_limit(&self) -> usize {
+        self.options.http_payload_size_limit.get_bytes() as usize
+    }
+
+    #[inline]
+    pub fn api_keys(&self) -> &ApiKeys {
+        &self.api_keys
+    }
+}

meilisearch-http/src/data/search.rs

@@ -0,0 +1,353 @@
+use std::collections::{HashSet, BTreeMap};
+use std::mem;
+use std::time::Instant;
+
+use anyhow::{bail, Context};
+use either::Either;
+use heed::RoTxn;
+use meilisearch_tokenizer::{Analyzer, AnalyzerConfig};
+use milli::{obkv_to_json, FacetCondition, Index, facet::FacetValue};
+use serde::{Deserialize, Serialize};
+use serde_json::{Map, Value};
+
+use super::Data;
+use crate::index_controller::IndexController;
+
+pub const DEFAULT_SEARCH_LIMIT: usize = 20;
+
+const fn default_search_limit() -> usize {
+    DEFAULT_SEARCH_LIMIT
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+#[allow(dead_code)]
+pub struct SearchQuery {
+    pub q: Option<String>,
+    pub offset: Option<usize>,
+    #[serde(default = "default_search_limit")]
+    pub limit: usize,
+    pub attributes_to_retrieve: Option<Vec<String>>,
+    pub attributes_to_crop: Option<Vec<String>>,
+    pub crop_length: Option<usize>,
+    pub attributes_to_highlight: Option<HashSet<String>>,
+    pub filters: Option<String>,
+    pub matches: Option<bool>,
+    pub facet_filters: Option<Value>,
+    pub facet_distributions: Option<Vec<String>>,
+}
+
+impl SearchQuery {
+    pub fn perform(&self, index: impl AsRef<Index>) -> anyhow::Result<SearchResult> {
+        let index = index.as_ref();
+        let before_search = Instant::now();
+        let rtxn = index.read_txn()?;
+
+        let mut search = index.search(&rtxn);
+
+        if let Some(ref query) = self.q {
+            search.query(query);
+        }
+
+        search.limit(self.limit);
+        search.offset(self.offset.unwrap_or_default());
+
+        if let Some(ref facets) = self.facet_filters {
+            if let Some(facets) = parse_facets(facets, index, &rtxn)? {
+                search.facet_condition(facets);
+            }
+        }
+
+        let milli::SearchResult {
+            documents_ids,
+            found_words,
+            candidates,
+        } = search.execute()?;
+
+        let mut documents = Vec::new();
+        let fields_ids_map = index.fields_ids_map(&rtxn)?;
+
+        let displayed_fields_ids = index.displayed_fields_ids(&rtxn)?;
+
+        let attributes_to_retrieve_ids = match self.attributes_to_retrieve {
+            Some(ref attrs) if attrs.iter().any(|f| f == "*") => None,
+            Some(ref attrs) => attrs
+                .iter()
+                .filter_map(|f| fields_ids_map.id(f))
+                .collect::<Vec<_>>()
+                .into(),
+            None => None,
+        };
+
+        let displayed_fields_ids = match (displayed_fields_ids, attributes_to_retrieve_ids) {
+            (_, Some(ids)) => ids,
+            (Some(ids), None) => ids,
+            (None, None) => fields_ids_map.iter().map(|(id, _)| id).collect(),
+        };
+
+        let stop_words = fst::Set::default();
+        let highlighter = Highlighter::new(&stop_words);
+
+        for (_id, obkv) in index.documents(&rtxn, documents_ids)? {
+            let mut object = obkv_to_json(&displayed_fields_ids, &fields_ids_map, obkv)?;
+            if let Some(ref attributes_to_highlight) = self.attributes_to_highlight {
+                highlighter.highlight_record(&mut object, &found_words, attributes_to_highlight);
+            }
+            documents.push(object);
+        }
+
+        let nb_hits = candidates.len();
+
+        let facet_distributions = match self.facet_distributions {
+            Some(ref fields) => {
+                let mut facet_distribution = index.facets_distribution(&rtxn);
+                if fields.iter().all(|f| f != "*") {
+                    facet_distribution.facets(fields);
+                }
+                Some(facet_distribution.candidates(candidates).execute()?)
+            }
+            None => None,
+        };
+
+        Ok(SearchResult {
+            hits: documents,
+            nb_hits,
+            query: self.q.clone().unwrap_or_default(),
+            limit: self.limit,
+            offset: self.offset.unwrap_or_default(),
+            processing_time_ms: before_search.elapsed().as_millis(),
+            facet_distributions,
+        })
+    }
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SearchResult {
+    hits: Vec<Map<String, Value>>,
+    nb_hits: u64,
+    query: String,
+    limit: usize,
+    offset: usize,
+    processing_time_ms: u128,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    facet_distributions: Option<BTreeMap<String, BTreeMap<FacetValue, u64>>>,
+}
+
+struct Highlighter<'a, A> {
+    analyzer: Analyzer<'a, A>,
+}
+
+impl<'a, A: AsRef<[u8]>> Highlighter<'a, A> {
+    fn new(stop_words: &'a fst::Set<A>) -> Self {
+        let analyzer = Analyzer::new(AnalyzerConfig::default_with_stopwords(stop_words));
+        Self { analyzer }
+    }
+
+    fn highlight_value(&self, value: Value, words_to_highlight: &HashSet<String>) -> Value {
+        match value {
+            Value::Null => Value::Null,
+            Value::Bool(boolean) => Value::Bool(boolean),
+            Value::Number(number) => Value::Number(number),
+            Value::String(old_string) => {
+                let mut string = String::new();
+                let analyzed = self.analyzer.analyze(&old_string);
+                for (word, token) in analyzed.reconstruct() {
+                    if token.is_word() {
+                        let to_highlight = words_to_highlight.contains(token.text());
+                        if to_highlight {
+                            string.push_str("<mark>")
+                        }
+                        string.push_str(word);
+                        if to_highlight {
+                            string.push_str("</mark>")
+                        }
+                    } else {
+                        string.push_str(word);
+                    }
+                }
+                Value::String(string)
+            }
+            Value::Array(values) => Value::Array(
+                values
+                    .into_iter()
+                    .map(|v| self.highlight_value(v, words_to_highlight))
+                    .collect(),
+            ),
+            Value::Object(object) => Value::Object(
+                object
+                    .into_iter()
+                    .map(|(k, v)| (k, self.highlight_value(v, words_to_highlight)))
+                    .collect(),
+            ),
+        }
+    }
+
+    fn highlight_record(
+        &self,
+        object: &mut Map<String, Value>,
+        words_to_highlight: &HashSet<String>,
+        attributes_to_highlight: &HashSet<String>,
+    ) {
+        // TODO do we need to create a string for element that are not and needs to be highlight?
+        for (key, value) in object.iter_mut() {
+            if attributes_to_highlight.contains(key) {
+                let old_value = mem::take(value);
+                *value = self.highlight_value(old_value, words_to_highlight);
+            }
+        }
+    }
+}
+
+impl Data {
+    pub fn search<S: AsRef<str>>(
+        &self,
+        index: S,
+        search_query: SearchQuery,
+    ) -> anyhow::Result<SearchResult> {
+        match self.index_controller.index(&index)? {
+            Some(index) => Ok(search_query.perform(index)?),
+            None => bail!("index {:?} doesn't exists", index.as_ref()),
+        }
+    }
+
+    pub async fn retrieve_documents<S>(
+        &self,
+        index: impl AsRef<str> + Send + Sync + 'static,
+        offset: usize,
+        limit: usize,
+        attributes_to_retrieve: Option<Vec<S>>,
+    ) -> anyhow::Result<Vec<Map<String, Value>>>
+    where
+        S: AsRef<str> + Send + Sync + 'static,
+    {
+        let index_controller = self.index_controller.clone();
+        let documents: anyhow::Result<_> = tokio::task::spawn_blocking(move || {
+            let index = index_controller
+                .index(&index)?
+                .with_context(|| format!("Index {:?} doesn't exist", index.as_ref()))?;
+
+            let txn = index.read_txn()?;
+
+            let fields_ids_map = index.fields_ids_map(&txn)?;
+
+            let attributes_to_retrieve_ids = match attributes_to_retrieve {
+                Some(attrs) => attrs
+                    .iter()
+                    .filter_map(|f| fields_ids_map.id(f.as_ref()))
+                    .collect::<Vec<_>>(),
+                None => fields_ids_map.iter().map(|(id, _)| id).collect(),
+            };
+
+            let iter = index.documents.range(&txn, &(..))?.skip(offset).take(limit);
+
+            let mut documents = Vec::new();
+
+            for entry in iter {
+                let (_id, obkv) = entry?;
+                let object = obkv_to_json(&attributes_to_retrieve_ids, &fields_ids_map, obkv)?;
+                documents.push(object);
+            }
+
+            Ok(documents)
+        })
+        .await?;
+        documents
+    }
+
+    pub async fn retrieve_document<S>(
+        &self,
+        index: impl AsRef<str> + Sync + Send + 'static,
+        document_id: impl AsRef<str> + Sync + Send + 'static,
+        attributes_to_retrieve: Option<Vec<S>>,
+    ) -> anyhow::Result<Map<String, Value>>
+    where
+        S: AsRef<str> + Sync + Send + 'static,
+    {
+        let index_controller = self.index_controller.clone();
+        let document: anyhow::Result<_> = tokio::task::spawn_blocking(move || {
+            let index = index_controller
+                .index(&index)?
+                .with_context(|| format!("Index {:?} doesn't exist", index.as_ref()))?;
+            let txn = index.read_txn()?;
+
+            let fields_ids_map = index.fields_ids_map(&txn)?;
+
+            let attributes_to_retrieve_ids = match attributes_to_retrieve {
+                Some(attrs) => attrs
+                    .iter()
+                    .filter_map(|f| fields_ids_map.id(f.as_ref()))
+                    .collect::<Vec<_>>(),
+                None => fields_ids_map.iter().map(|(id, _)| id).collect(),
+            };
+
+            let internal_id = index
+                .external_documents_ids(&txn)?
+                .get(document_id.as_ref().as_bytes())
+                .with_context(|| format!("Document with id {} not found", document_id.as_ref()))?;
+
+            let document = index
+                .documents(&txn, std::iter::once(internal_id))?
+                .into_iter()
+                .next()
+                .map(|(_, d)| d);
+
+            match document {
+                Some(document) => Ok(obkv_to_json(
+                    &attributes_to_retrieve_ids,
+                    &fields_ids_map,
+                    document,
+                )?),
+                None => bail!("Document with id {} not found", document_id.as_ref()),
+            }
+        })
+        .await?;
+        document
+    }
+}
+
+fn parse_facets_array(
+    txn: &RoTxn,
+    index: &Index,
+    arr: &Vec<Value>,
+) -> anyhow::Result<Option<FacetCondition>> {
+    let mut ands = Vec::new();
+    for value in arr {
+        match value {
+            Value::String(s) => ands.push(Either::Right(s.clone())),
+            Value::Array(arr) => {
+                let mut ors = Vec::new();
+                for value in arr {
+                    match value {
+                        Value::String(s) => ors.push(s.clone()),
+                        v => bail!("Invalid facet expression, expected String, found: {:?}", v),
+                    }
+                }
+                ands.push(Either::Left(ors));
+            }
+            v => bail!(
+                "Invalid facet expression, expected String or [String], found: {:?}",
+                v
+            ),
+        }
+    }
+
+    FacetCondition::from_array(txn, index, ands)
+}
+
+fn parse_facets(
+    facets: &Value,
+    index: &Index,
+    txn: &RoTxn,
+) -> anyhow::Result<Option<FacetCondition>> {
+    match facets {
+        // Disabled for now
+        //Value::String(expr) => Ok(Some(FacetCondition::from_str(txn, index, expr)?)),
+        Value::Array(arr) => parse_facets_array(txn, index, arr),
+        v => bail!(
+            "Invalid facet expression, expected Array, found: {:?}",
+            v
+        ),
+    }
+}
+

meilisearch-http/src/data/updates.rs

@@ -0,0 +1,115 @@
+use std::ops::Deref;
+
+use async_compression::tokio_02::write::GzipEncoder;
+use futures_util::stream::StreamExt;
+use milli::update::{IndexDocumentsMethod, UpdateFormat};
+use tokio::io::AsyncWriteExt;
+
+use crate::index_controller::UpdateStatus;
+use crate::index_controller::{IndexController, Settings, IndexSettings, IndexMetadata};
+use super::Data;
+
+impl Data {
+    pub async fn add_documents<B, E>(
+        &self,
+        index: impl AsRef<str> + Send + Sync + 'static,
+        method: IndexDocumentsMethod,
+        format: UpdateFormat,
+        mut stream: impl futures::Stream<Item=Result<B, E>> + Unpin,
+        primary_key: Option<String>,
+    ) -> anyhow::Result<UpdateStatus>
+    where
+        B: Deref<Target = [u8]>,
+        E: std::error::Error + Send + Sync + 'static,
+    {
+        let file = tokio::task::spawn_blocking(tempfile::tempfile).await?;
+        let file = tokio::fs::File::from_std(file?);
+        let mut encoder = GzipEncoder::new(file);
+
+        let mut empty_update = true;
+        while let Some(result) = stream.next().await {
+            empty_update = false;
+            let bytes = &*result?;
+            encoder.write_all(&bytes[..]).await?;
+        }
+
+        encoder.shutdown().await?;
+        let mut file = encoder.into_inner();
+        file.sync_all().await?;
+        let file = file.into_std().await;
+
+        let index_controller = self.index_controller.clone();
+        let update = tokio::task::spawn_blocking(move ||{
+            let mmap;
+            let bytes = if empty_update {
+                &[][..]
+            } else {
+                mmap = unsafe { memmap::Mmap::map(&file)? };
+                &mmap
+            };
+            index_controller.add_documents(index, method, format, &bytes, primary_key)
+        }).await??;
+        Ok(update.into())
+    }
+
+    pub async fn update_settings(
+        &self,
+        index: impl AsRef<str> + Send + Sync + 'static,
+        settings: Settings
+    ) -> anyhow::Result<UpdateStatus> {
+        let index_controller = self.index_controller.clone();
+        let update = tokio::task::spawn_blocking(move || index_controller.update_settings(index, settings)).await??;
+        Ok(update.into())
+    }
+
+    pub async fn clear_documents(
+        &self,
+        index: impl AsRef<str> + Sync + Send + 'static,
+    ) -> anyhow::Result<UpdateStatus> {
+        let index_controller = self.index_controller.clone();
+        let update = tokio::task::spawn_blocking(move || index_controller.clear_documents(index)).await??;
+        Ok(update.into())
+    }
+
+    pub async fn delete_documents(
+        &self,
+        index: impl AsRef<str> + Sync + Send + 'static,
+        document_ids: Vec<String>,
+    ) -> anyhow::Result<UpdateStatus> {
+        let index_controller = self.index_controller.clone();
+        let update = tokio::task::spawn_blocking(move || index_controller.delete_documents(index, document_ids)).await??;
+        Ok(update.into())
+    }
+
+    pub async fn delete_index(
+        &self,
+        index: impl AsRef<str> + Send + Sync + 'static,
+    ) -> anyhow::Result<()> {
+        let index_controller = self.index_controller.clone();
+        tokio::task::spawn_blocking(move || { index_controller.delete_index(index) }).await??;
+        Ok(())
+    }
+
+    #[inline]
+    pub fn get_update_status(&self, index: impl AsRef<str>, uid: u64) -> anyhow::Result<Option<UpdateStatus>> {
+        self.index_controller.update_status(index, uid)
+    }
+
+    pub fn get_updates_status(&self, index: impl AsRef<str>) -> anyhow::Result<Vec<UpdateStatus>> {
+        self.index_controller.all_update_status(index)
+    }
+
+    pub fn update_index(
+        &self,
+        name: impl AsRef<str>,
+        primary_key: Option<impl AsRef<str>>,
+        new_name: Option<impl AsRef<str>>
+    ) -> anyhow::Result<IndexMetadata> {
+        let settings = IndexSettings {
+            name: new_name.map(|s| s.as_ref().to_string()),
+            primary_key: primary_key.map(|s| s.as_ref().to_string()),
+        };
+
+        self.index_controller.update_index(name, settings)
+    }
+}