Merge pull request #18 from meilisearch/more-tenant

More tenant token, v1.48 edition

Author: dureuill

crates/meilisearch/src/routes/chats/chat_completions.rs

@@ -24,7 +24,7 @@ use index_scheduler::filter::{
     filter_into_index_filter, filters_into_index_filters_unchecked, parse_filter,
 };
 use index_scheduler::IndexScheduler;
-use meilisearch_auth::AuthController;
+use meilisearch_auth::{AuthController, IndexSearchRules};
 use meilisearch_types::error::{Code, ResponseError};
 use meilisearch_types::features::{
     ChatCompletionPrompts as DbChatCompletionPrompts,
@@ -35,7 +35,7 @@ use meilisearch_types::keys::actions;
 use meilisearch_types::milli::index::ChatConfig;
 use meilisearch_types::milli::progress::Progress;
 use meilisearch_types::milli::{
-    all_obkv_to_json, obkv_to_json, OrderBy, PatternMatch, TotalProcessingTimeStep,
+    all_obkv_to_json, obkv_to_json, Filter, OrderBy, PatternMatch, TotalProcessingTimeStep,
 };
 use meilisearch_types::{Document, Index};
 use serde::Deserialize;
@@ -239,18 +239,29 @@ fn setup_search_tool(
     let mut index_uids = Vec::new();
     let mut function_description = prompts.search_description.clone();
     let mut filter_description = prompts.search_filter_param.clone();
+    let progress = &Default::default();
     index_scheduler.try_for_each_index::<_, ()>(|name, index| {
         // Make sure to skip unauthorized indexes
         if !filters.is_index_authorized(name) {
             return Ok(());
         }
+        let search_rules = filters.get_index_search_rules(name);
 
         let rtxn = index.read_txn()?;
         let chat_config = index.chat_config(&rtxn)?;
         let index_description = chat_config.description;
         let _ = writeln!(&mut function_description, "\n\n - {name}: {index_description}\n");
         index_uids.push(name.to_string());
-        let facet_distributions = format_facet_distributions(index, &rtxn, 10).unwrap(); // TODO do not unwrap
+        let facet_distributions = format_facet_distributions(
+            index_scheduler,
+            index,
+            &rtxn,
+            10,
+            search_rules,
+            name,
+            progress,
+        )
+        .unwrap(); // TODO do not unwrap
         let _ = writeln!(&mut filter_description, "\n## Facet distributions of the {name} index");
         let _ = writeln!(&mut filter_description, "{facet_distributions}");
 
@@ -956,11 +967,27 @@ struct SearchInIndexParameters {
 }
 
 fn format_facet_distributions(
+    index_scheduler: &IndexScheduler,
     index: &Index,
     rtxn: &RoTxn,
     max_values_per_facet: usize,
-) -> meilisearch_types::milli::Result<String> {
-    let universe = index.documents_ids(rtxn)?;
+    search_rules: Option<IndexSearchRules>,
+    index_uid: &str,
+    progress: &Progress,
+) -> index_scheduler::Result<String> {
+    let from_milli = |err| index_scheduler::Error::from_milli(err, Some(index_uid.to_string()));
+    let universe = 'filter: {
+        let Some(search_rules) = search_rules else { break 'filter index.documents_ids(rtxn)? };
+        let Some(filter) = search_rules.filter else {
+            break 'filter index.documents_ids(rtxn)?;
+        };
+        let Some(filter) = Filter::from_json(&filter).map_err(from_milli)? else {
+            break 'filter index.documents_ids(rtxn)?;
+        };
+        let filter =
+            filter_into_index_filter(filter, index, rtxn, index_scheduler, progress, index_uid)?;
+        filter.evaluate(rtxn, index).map_err(from_milli)?
+    };
     let rules = index.filterable_attributes_rules(rtxn)?;
     let fields_ids_map = index.fields_ids_map(rtxn)?;
     let filterable_attributes = fields_ids_map
@@ -972,7 +999,8 @@ fn format_facet_distributions(
         .max_values_per_facet(max_values_per_facet)
         .candidates(universe)
         .facets(filterable_attributes)
-        .execute()?;
+        .execute()
+        .map_err(from_milli)?;
 
     let mut output = String::new();
     for (facet_name, entries) in facets_distribution {

crates/meilisearch/src/routes/indexes/similar.rs

@@ -207,16 +207,19 @@ async fn similar(
     mut query: SimilarQuery,
 ) -> Result<SimilarResult, ResponseError> {
     let progress = Progress::default();
+
+    let search_rules = index_scheduler.filters().get_index_search_rules(&index_uid);
+
     // Tenant token search_rules.
-    if let Some(search_rules) = index_scheduler.filters().get_index_search_rules(&index_uid) {
+    if let Some(search_rules) = search_rules.clone() {
         add_search_rules(&mut query.filter, search_rules);
     }
 
     let progress_clone = progress.clone();
     let result = tokio::task::spawn_blocking(move || {
         let _step = progress_clone.update_progress_scoped(TotalProcessingTimeStep::Search);
 
-        perform_similar(&index_scheduler, index_uid, query, &progress_clone)
+        perform_similar(&index_scheduler, index_uid, query, &progress_clone, search_rules)
     })
     .await;
 

crates/meilisearch/src/search/mod.rs

@@ -2,12 +2,16 @@ use core::fmt;
 use std::cmp::min;
 use std::collections::{BTreeMap, BTreeSet, HashSet};
 use std::ops::Not as _;
+use std::rc::Rc;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Instant;
 
 use deserr::Deserr;
-use index_scheduler::filter::{filter_into_index_filter, parse_filter};
+use index_scheduler::filter::{
+    filter_into_index_filter, filters_into_index_filters, parse_filter,
+    retrieve_foreign_keys_settings, SourceIndexUid,
+};
 use index_scheduler::{IndexScheduler, RoFeatures};
 use indexmap::IndexMap;
 use meilisearch_auth::IndexSearchRules;
@@ -23,8 +27,8 @@ use meilisearch_types::milli::score_details::{ScoreDetails, ScoringStrategy};
 use meilisearch_types::milli::vector::parsed_vectors::ExplicitVectors;
 use meilisearch_types::milli::vector::Embedder;
 use meilisearch_types::milli::{
-    AttributeState, Deadline, FacetValueHit, IndexFilter, InternalError, OrderBy, PatternMatch,
-    SearchForFacetValues, SearchStep,
+    filtered_universe, AttributeState, Deadline, FacetValueHit, Filter, IndexFilter, InternalError,
+    OrderBy, PatternMatch, SearchForFacetValues, SearchStep,
 };
 use meilisearch_types::network::Network;
 use meilisearch_types::settings::DEFAULT_PAGINATION_MAX_TOTAL_HITS;
@@ -2481,6 +2485,7 @@ pub fn perform_similar(
     index_uid: IndexUid,
     query: SimilarQuery,
     progress: &Progress,
+    search_rules: Option<IndexSearchRules>,
 ) -> Result<SimilarResult, ResponseError> {
     let before_search = Instant::now();
     let features = index_scheduler.features();
@@ -2512,6 +2517,30 @@ pub fn perform_similar(
         Route::Similar,
     )?;
 
+    let docid_filter = search_rules.and_then(|search_rules| search_rules.filter);
+    let docid_filter = docid_filter
+        .as_ref()
+        .map(|docid_filter| {
+            parse_filter(
+                docid_filter,
+                Code::InvalidSimilarFilter,
+                features,
+                Some(index_uid.as_str()),
+            )
+        })
+        .transpose()?
+        .flatten();
+
+    let candidates_filter = filter
+        .as_ref()
+        .and_then(|filter| {
+            parse_filter(filter, Code::InvalidSimilarFilter, features, None).transpose()
+        })
+        .transpose()?;
+
+    let (docid_filter, candidates_filter) =
+        extract_filters(index_scheduler, index_uid, progress, docid_filter, candidates_filter)?;
+
     let id: ExternalDocumentId = id.try_into().map_err(|error| {
         let msg = format!("Invalid value at `.id`: {error}");
         ResponseError::from_msg(msg, Code::InvalidSimilarId)
@@ -2526,6 +2555,14 @@ pub fn perform_similar(
         ));
     };
 
+    let docid_universe = filtered_universe(&index, &rtxn, &docid_filter, progress)?;
+    if docid_universe.contains(internal_id).not() {
+        return Err(ResponseError::from_msg(
+            MeilisearchHttpError::DocumentNotFound(id.into_inner()).to_string(),
+            Code::NotFoundSimilarId,
+        ));
+    }
+
     let mut similar = milli::Similar::new(
         internal_id,
         offset,
@@ -2538,18 +2575,8 @@ pub fn perform_similar(
         progress,
     );
 
-    if let Some(ref filter) = filter {
-        if let Some(filter) = parse_filter(filter, Code::InvalidSimilarFilter, features, None)? {
-            let filter = filter_into_index_filter(
-                filter,
-                &index,
-                &rtxn,
-                index_scheduler,
-                progress,
-                &index_uid,
-            )?;
-            similar.filter(filter);
-        }
+    if let Some(filter) = candidates_filter {
+        similar.filter(filter);
     }
 
     if let Some(ranking_score_threshold) = ranking_score_threshold {
@@ -2618,6 +2645,33 @@ pub fn perform_similar(
     Ok(result)
 }
 
+fn extract_filters<'a>(
+    index_scheduler: &IndexScheduler,
+    index_uid: IndexUid,
+    progress: &Progress,
+    docid_filter: Option<Filter<'a>>,
+    candidates_filter: Option<Filter<'a>>,
+) -> Result<(Option<IndexFilter<'a>>, Option<IndexFilter<'a>>), ResponseError> {
+    let source_index_uid = SourceIndexUid(Rc::from(&*index_uid));
+    let foreign_keys_settings =
+        retrieve_foreign_keys_settings(index_scheduler, std::iter::once(&source_index_uid))?;
+    let (docid_filter, candidates_filter) = match filters_into_index_filters(
+        vec![
+            (source_index_uid.clone(), docid_filter),
+            (source_index_uid.clone(), candidates_filter),
+        ],
+        &foreign_keys_settings,
+        index_scheduler,
+        progress,
+    )?
+    .as_mut_slice()
+    {
+        [docid_filter, candidates_filter] => (docid_filter.take(), candidates_filter.take()),
+        _ => unreachable!(),
+    };
+    Ok((docid_filter, candidates_filter))
+}
+
 pub fn insert_geo_distance(sorts: &[String], document: &mut Document) {
     lazy_static::lazy_static! {
         static ref GEO_REGEX: Regex =