From ddcb138f249559f474d53be8f0e171aa091eac46 Mon Sep 17 00:00:00 2001 From: trinity-1686a Date: Mon, 15 Jul 2024 16:20:56 +0200 Subject: [PATCH] add api endpoint to get query ast (#5212) * add api to get quickwit ast, tantivy ast, and some query plan detail * add test for search-plan --- .../protos/quickwit/search.proto | 7 + .../src/codegen/quickwit/quickwit.search.rs | 85 ++++++ quickwit/quickwit-search/src/lib.rs | 6 +- quickwit/quickwit-search/src/root.rs | 268 +++++++++++++++--- .../src/search_response_rest.rs | 37 +++ quickwit/quickwit-search/src/service.rs | 16 +- quickwit/quickwit-serve/src/rest.rs | 7 +- .../src/search_api/grpc_adapter.rs | 11 + quickwit/quickwit-serve/src/search_api/mod.rs | 5 +- .../src/search_api/rest_handler.rs | 94 +++++- 10 files changed, 492 insertions(+), 44 deletions(-) diff --git a/quickwit/quickwit-proto/protos/quickwit/search.proto b/quickwit/quickwit-proto/protos/quickwit/search.proto index b5cc92cd5d2..d3d8954377b 100644 --- a/quickwit/quickwit-proto/protos/quickwit/search.proto +++ b/quickwit/quickwit-proto/protos/quickwit/search.proto @@ -76,6 +76,9 @@ service SearchService { rpc ListFields(ListFieldsRequest) returns (ListFieldsResponse); rpc LeafListFields(LeafListFieldsRequest) returns (ListFieldsResponse); + + // Describe how a search would be processed. + rpc SearchPlan(SearchRequest) returns (SearchPlanResponse); } /// Scroll Request @@ -298,6 +301,10 @@ message SearchResponse { optional string scroll_id = 6; } +message SearchPlanResponse { + string result = 1; +} + message SplitSearchError { // The searcherror that occurred formatted as string. string error = 1; diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs index bee455bccd3..70b7ba42191 100644 --- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs +++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs @@ -234,6 +234,13 @@ pub struct SearchResponse { #[derive(Serialize, Deserialize, utoipa::ToSchema)] #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] +pub struct SearchPlanResponse { + #[prost(string, tag = "1")] + pub result: ::prost::alloc::string::String, +} +#[derive(Serialize, Deserialize, utoipa::ToSchema)] +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct SplitSearchError { /// The searcherror that occurred formatted as string. #[prost(string, tag = "1")] @@ -1201,6 +1208,32 @@ pub mod search_service_client { ); self.inner.unary(req, path, codec).await } + /// Describe how a search would be processed. + pub async fn search_plan( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/quickwit.search.SearchService/SearchPlan", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert(GrpcMethod::new("quickwit.search.SearchService", "SearchPlan")); + self.inner.unary(req, path, codec).await + } } } /// Generated server implementations. @@ -1322,6 +1355,14 @@ pub mod search_service_server { tonic::Response, tonic::Status, >; + /// Describe how a search would be processed. + async fn search_plan( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; } #[derive(Debug)] pub struct SearchServiceServer { @@ -1940,6 +1981,50 @@ pub mod search_service_server { }; Box::pin(fut) } + "/quickwit.search.SearchService/SearchPlan" => { + #[allow(non_camel_case_types)] + struct SearchPlanSvc(pub Arc); + impl< + T: SearchService, + > tonic::server::UnaryService + for SearchPlanSvc { + type Response = super::SearchPlanResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { (*inner).search_plan(request).await }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let inner = inner.0; + let method = SearchPlanSvc(inner); + let codec = tonic::codec::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } _ => { Box::pin(async move { Ok( diff --git a/quickwit/quickwit-search/src/lib.rs b/quickwit/quickwit-search/src/lib.rs index d43bba9f507..b5e7a881c08 100644 --- a/quickwit/quickwit-search/src/lib.rs +++ b/quickwit/quickwit-search/src/lib.rs @@ -84,11 +84,11 @@ pub use crate::cluster_client::ClusterClient; pub use crate::error::{parse_grpc_error, SearchError}; use crate::fetch_docs::fetch_docs; pub use crate::root::{ - check_all_index_metadata_found, jobs_to_leaf_request, root_search, IndexMetasForLeafSearch, - SearchJob, + check_all_index_metadata_found, jobs_to_leaf_request, root_search, search_plan, + IndexMetasForLeafSearch, SearchJob, }; pub use crate::search_job_placer::{Job, SearchJobPlacer}; -pub use crate::search_response_rest::SearchResponseRest; +pub use crate::search_response_rest::{SearchPlanResponseRest, SearchResponseRest}; pub use crate::search_stream::root_search_stream; pub use crate::service::{MockSearchService, SearchService, SearchServiceImpl}; diff --git a/quickwit/quickwit-search/src/root.rs b/quickwit/quickwit-search/src/root.rs index ef23596587d..2ed4d000329 100644 --- a/quickwit/quickwit-search/src/root.rs +++ b/quickwit/quickwit-search/src/root.rs @@ -35,8 +35,8 @@ use quickwit_proto::metastore::{ }; use quickwit_proto::search::{ FetchDocsRequest, FetchDocsResponse, Hit, LeafHit, LeafRequestRef, LeafSearchRequest, - LeafSearchResponse, PartialHit, SearchRequest, SearchResponse, SnippetRequest, - SortDatetimeFormat, SortField, SortValue, SplitIdAndFooterOffsets, + LeafSearchResponse, PartialHit, SearchPlanResponse, SearchRequest, SearchResponse, + SnippetRequest, SortDatetimeFormat, SortField, SortValue, SplitIdAndFooterOffsets, }; use quickwit_proto::types::{IndexUid, SplitId}; use quickwit_query::query_ast::{ @@ -55,10 +55,11 @@ use crate::collector::{make_merge_collector, QuickwitAggregations}; use crate::find_trace_ids_collector::Span; use crate::scroll_context::{ScrollContext, ScrollKeyAndStartOffset}; use crate::search_job_placer::{group_by, group_jobs_by_index_id, Job}; +use crate::search_response_rest::StorageRequestCount; use crate::service::SearcherContext; use crate::{ extract_split_and_footer_offsets, list_relevant_splits, SearchError, SearchJobPlacer, - SearchServiceClient, + SearchPlanResponseRest, SearchServiceClient, }; /// Maximum accepted scroll TTL. @@ -1014,6 +1015,47 @@ pub fn check_all_index_metadata_found( Ok(()) } +async fn refine_and_list_matches( + metastore: &mut MetastoreServiceClient, + search_request: &mut SearchRequest, + indexes_metadata: Vec, + query_ast_resolved: QueryAst, + sort_fields_is_datetime: HashMap, + timestamp_field_opt: Option, +) -> crate::Result> { + let index_uids = indexes_metadata + .iter() + .map(|index_metadata| index_metadata.index_uid.clone()) + .collect_vec(); + search_request.query_ast = serde_json::to_string(&query_ast_resolved)?; + + // convert search_after datetime values from input datetime format to nanos. + convert_search_after_datetime_values(search_request, &sort_fields_is_datetime)?; + + // update_search_after_datetime_in_nanos(&mut search_request)?; + if let Some(timestamp_field) = ×tamp_field_opt { + refine_start_end_timestamp_from_ast( + &query_ast_resolved, + timestamp_field, + &mut search_request.start_timestamp, + &mut search_request.end_timestamp, + ); + } + let tag_filter_ast = extract_tags_from_query(query_ast_resolved); + + // TODO if search after is set, we sort by timestamp and we don't want to count all results, + // we can refine more here. Same if we sort by _shard_doc + let split_metadatas: Vec = list_relevant_splits( + index_uids, + search_request.start_timestamp, + search_request.end_timestamp, + tag_filter_ast, + metastore, + ) + .await?; + Ok(split_metadatas) +} + /// Performs a distributed search. /// 1. Sends leaf request over gRPC to multiple leaf nodes. /// 2. Merges the search results. @@ -1055,38 +1097,14 @@ pub async fn root_search( return Ok(search_response); } - let index_uids = indexes_metadata - .iter() - .map(|index_metadata| index_metadata.index_uid.clone()) - .collect_vec(); let request_metadata = validate_request_and_build_metadata(&indexes_metadata, &search_request)?; - search_request.query_ast = serde_json::to_string(&request_metadata.query_ast_resolved)?; - - // convert search_after datetime values from input datetime format to nanos. - convert_search_after_datetime_values( - &mut search_request, - &request_metadata.sort_fields_is_datetime, - )?; - - // update_search_after_datetime_in_nanos(&mut search_request)?; - if let Some(timestamp_field) = &request_metadata.timestamp_field_opt { - refine_start_end_timestamp_from_ast( - &request_metadata.query_ast_resolved, - timestamp_field, - &mut search_request.start_timestamp, - &mut search_request.end_timestamp, - ); - } - let tag_filter_ast = extract_tags_from_query(request_metadata.query_ast_resolved); - - // TODO if search after is set, we sort by timestamp and we don't want to count all results, - // we can refine more here. Same if we sort by _shard_doc - let split_metadatas: Vec = list_relevant_splits( - index_uids, - search_request.start_timestamp, - search_request.end_timestamp, - tag_filter_ast, + let split_metadatas = refine_and_list_matches( &mut metastore, + &mut search_request, + indexes_metadata, + request_metadata.query_ast_resolved, + request_metadata.sort_fields_is_datetime, + request_metadata.timestamp_field_opt, ) .await?; @@ -1103,6 +1121,104 @@ pub async fn root_search( Ok(search_response) } +/// Returns details on how a query would be executed +pub async fn search_plan( + mut search_request: SearchRequest, + mut metastore: MetastoreServiceClient, +) -> crate::Result { + let list_indexes_metadatas_request = ListIndexesMetadataRequest { + index_id_patterns: search_request.index_id_patterns.clone(), + }; + let indexes_metadata: Vec = metastore + .list_indexes_metadata(list_indexes_metadatas_request) + .await? + .deserialize_indexes_metadata() + .await?; + + check_all_index_metadata_found(&indexes_metadata[..], &search_request.index_id_patterns[..])?; + if indexes_metadata.is_empty() { + return Ok(SearchPlanResponse { + result: serde_json::to_string(&SearchPlanResponseRest { + quickwit_ast: QueryAst::MatchAll, + tantivy_ast: String::new(), + searched_splits: Vec::new(), + storage_requests: StorageRequestCount::default(), + })?, + }); + } + let doc_mapper = build_doc_mapper( + &indexes_metadata[0].index_config.doc_mapping, + &indexes_metadata[0].index_config.search_settings, + ) + .map_err(|err| SearchError::Internal(format!("failed to build doc mapper. cause: {err}")))?; + + let request_metadata = validate_request_and_build_metadata(&indexes_metadata, &search_request)?; + let split_metadatas = refine_and_list_matches( + &mut metastore, + &mut search_request, + indexes_metadata, + request_metadata.query_ast_resolved.clone(), + request_metadata.sort_fields_is_datetime, + request_metadata.timestamp_field_opt, + ) + .await?; + + let (query, mut warmup_info) = doc_mapper.query( + doc_mapper.schema(), + &request_metadata.query_ast_resolved, + true, + )?; + let merge_collector = make_merge_collector(&search_request, &Default::default())?; + warmup_info.merge(merge_collector.warmup_info()); + warmup_info.simplify(); + + let split_ids = split_metadatas + .into_iter() + .map(|split| format!("{}/{}", split.index_uid.index_id, split.split_id)) + .collect(); + // this is an upper bound, we'd need access to a hotdir for more precise results + let fieldnorm_query_count = if warmup_info.field_norms { + doc_mapper + .schema() + .fields() + .filter(|(_, entry)| entry.has_fieldnorms()) + .count() + } else { + 0 + }; + let sstable_query_count = warmup_info.term_dict_fields.len() + + warmup_info + .terms_grouped_by_field + .values() + .map(|terms: &HashMap| terms.len()) + .sum::(); + let position_query_count = warmup_info + .terms_grouped_by_field + .values() + .map(|terms: &HashMap| { + terms + .values() + .filter(|load_position| **load_position) + .count() + }) + .sum(); + Ok(SearchPlanResponse { + result: serde_json::to_string(&SearchPlanResponseRest { + quickwit_ast: request_metadata.query_ast_resolved, + tantivy_ast: format!("{query:#?}"), + searched_splits: split_ids, + storage_requests: StorageRequestCount { + footer: 1, + fastfield: warmup_info.fast_field_names.len(), + fieldnorm: fieldnorm_query_count, + sstable: sstable_query_count, + posting: sstable_query_count, + position: position_query_count, + }, + })?, + }) +} + /// Converts search after with datetime format to nanoseconds (representation in tantivy). /// If the sort field is a datetime field and no datetime format is set, the default format is /// milliseconds. @@ -3755,6 +3871,92 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_search_plan_multiple_splits() -> anyhow::Result<()> { + use quickwit_query::query_ast::{FullTextMode, FullTextParams, FullTextQuery}; + use quickwit_query::MatchAllOrNone; + + let search_request = quickwit_proto::search::SearchRequest { + index_id_patterns: vec!["test-index".to_string()], + query_ast: qast_json_helper("test-query", &["body"]), + max_hits: 10, + ..Default::default() + }; + let mut mock_metastore = MockMetastoreService::new(); + let index_metadata = IndexMetadata::for_test("test-index", "ram:///test-index"); + let index_uid = index_metadata.index_uid.clone(); + mock_metastore + .expect_list_indexes_metadata() + .returning(move |_index_ids_query| { + Ok(ListIndexesMetadataResponse::for_test(vec![ + index_metadata.clone() + ])) + }); + mock_metastore + .expect_list_splits() + .returning(move |_filter| { + let splits = vec![ + MockSplitBuilder::new("split1") + .with_index_uid(&index_uid) + .build(), + MockSplitBuilder::new("split2") + .with_index_uid(&index_uid) + .build(), + ]; + let splits_response = ListSplitsResponse::try_from_splits(splits).unwrap(); + Ok(ServiceStream::from(vec![Ok(splits_response)])) + }); + let search_response = search_plan( + search_request, + MetastoreServiceClient::from_mock(mock_metastore), + ) + .await + .unwrap(); + let response: SearchPlanResponseRest = + serde_json::from_str(&search_response.result).unwrap(); + assert_eq!( + response, + SearchPlanResponseRest { + quickwit_ast: QueryAst::FullText(FullTextQuery { + field: "body".to_string(), + text: "test-query".to_string(), + params: FullTextParams { + tokenizer: None, + mode: FullTextMode::PhraseFallbackToIntersection, + zero_terms_query: MatchAllOrNone::MatchNone, + }, + },), + tantivy_ast: r#"BooleanQuery { + subqueries: [ + ( + Must, + TermQuery(Term(field=3, type=Str, "test")), + ), + ( + Must, + TermQuery(Term(field=3, type=Str, "query")), + ), + ], + minimum_number_should_match: 0, +}"# + .to_string(), + searched_splits: vec![ + "test-index/split1".to_string(), + "test-index/split2".to_string() + ], + storage_requests: StorageRequestCount { + footer: 1, + fastfield: 0, + fieldnorm: 0, + sstable: 2, + posting: 2, + position: 0, + }, + } + ); + Ok(()) + } + #[test] fn test_extract_timestamp_range_from_ast() { use std::ops::Bound; diff --git a/quickwit/quickwit-search/src/search_response_rest.rs b/quickwit/quickwit-search/src/search_response_rest.rs index 52d5ee8d9ae..e34895ac659 100644 --- a/quickwit/quickwit-search/src/search_response_rest.rs +++ b/quickwit/quickwit-search/src/search_response_rest.rs @@ -21,6 +21,7 @@ use std::convert::TryFrom; use quickwit_common::truncate_str; use quickwit_proto::search::SearchResponse; +use quickwit_query::query_ast::QueryAst; use serde::{Deserialize, Serialize}; use serde_json::Value as JsonValue; @@ -100,3 +101,39 @@ impl TryFrom for SearchResponseRest { }) } } + +/// Details on how a query would be executed. +#[derive(Serialize, Deserialize, PartialEq, Debug, utoipa::ToSchema)] +pub struct SearchPlanResponseRest { + /// Quickwit AST of the query. + #[schema(value_type = Object)] + pub quickwit_ast: QueryAst, + /// Resolved Tantivy AST of the query, according to the latest docmapping. + /// + /// It's possible older splits actually resolve to a different ast. + pub tantivy_ast: String, + /// List of splits that would be searched by this query + pub searched_splits: Vec, + /// Requests expected for each split + #[schema(value_type = Object)] + pub storage_requests: StorageRequestCount, +} + +/// Number of expected storage requests, per request kind. +/// +/// These figures do not take in account whether the data is already cached or not. +#[derive(Serialize, Deserialize, PartialEq, Debug, Default)] +pub struct StorageRequestCount { + /// Number of split footer downloaded, always 1 + pub footer: usize, + /// Number of fastfields downloaded + pub fastfield: usize, + /// Number of fieldnorm downloaded + pub fieldnorm: usize, + /// Number of sstable dowloaded + pub sstable: usize, + /// Number of posting list downloaded + pub posting: usize, + /// Number of position list downloaded + pub position: usize, +} diff --git a/quickwit/quickwit-search/src/service.rs b/quickwit/quickwit-search/src/service.rs index 2fb03220e22..e7d0f685315 100644 --- a/quickwit/quickwit-search/src/service.rs +++ b/quickwit/quickwit-search/src/service.rs @@ -33,7 +33,8 @@ use quickwit_proto::search::{ LeafListTermsRequest, LeafListTermsResponse, LeafSearchRequest, LeafSearchResponse, LeafSearchStreamRequest, LeafSearchStreamResponse, ListFieldsRequest, ListFieldsResponse, ListTermsRequest, ListTermsResponse, PutKvRequest, ReportSplitsRequest, ReportSplitsResponse, - ScrollRequest, SearchRequest, SearchResponse, SearchStreamRequest, SnippetRequest, + ScrollRequest, SearchPlanResponse, SearchRequest, SearchResponse, SearchStreamRequest, + SnippetRequest, }; use quickwit_storage::{ MemorySizedCache, QuickwitCache, SplitCache, StorageCache, StorageResolver, @@ -50,7 +51,7 @@ use crate::list_terms::{leaf_list_terms, root_list_terms}; use crate::root::fetch_docs_phase; use crate::scroll_context::{MiniKV, ScrollContext, ScrollKeyAndStartOffset}; use crate::search_stream::{leaf_search_stream, root_search_stream}; -use crate::{fetch_docs, root_search, ClusterClient, SearchError}; +use crate::{fetch_docs, root_search, search_plan, ClusterClient, SearchError}; #[derive(Clone)] /// The search service implementation. @@ -149,6 +150,9 @@ pub trait SearchService: 'static + Send + Sync { &self, list_fields: LeafListFieldsRequest, ) -> crate::Result; + + /// Describe how a search would be processed. + async fn search_plan(&self, reqiest: SearchRequest) -> crate::Result; } impl SearchServiceImpl { @@ -352,6 +356,14 @@ impl SearchService for SearchServiceImpl { ) .await } + + async fn search_plan( + &self, + search_request: SearchRequest, + ) -> crate::Result { + let search_plan = search_plan(search_request, self.metastore.clone()).await?; + Ok(search_plan) + } } pub(crate) async fn scroll( diff --git a/quickwit/quickwit-serve/src/rest.rs b/quickwit/quickwit-serve/src/rest.rs index e2818b35928..483961a7378 100644 --- a/quickwit/quickwit-serve/src/rest.rs +++ b/quickwit/quickwit-serve/src/rest.rs @@ -48,7 +48,10 @@ use crate::metrics_api::metrics_handler; use crate::node_info_handler::node_info_handler; use crate::otlp_api::otlp_ingest_api_handlers; use crate::rest_api_response::{RestApiError, RestApiResponse}; -use crate::search_api::{search_get_handler, search_post_handler, search_stream_handler}; +use crate::search_api::{ + search_get_handler, search_plan_get_handler, search_plan_post_handler, search_post_handler, + search_stream_handler, +}; use crate::template_api::index_template_api_handlers; use crate::ui_handler::ui_handler; use crate::{BodyFormat, BuildInfo, QuickwitServices, RuntimeInfo}; @@ -234,6 +237,8 @@ fn search_routes( ) -> impl Filter + Clone { search_get_handler(search_service.clone()) .or(search_post_handler(search_service.clone())) + .or(search_plan_get_handler(search_service.clone())) + .or(search_plan_post_handler(search_service.clone())) .or(search_stream_handler(search_service)) .recover(recover_fn) } diff --git a/quickwit/quickwit-serve/src/search_api/grpc_adapter.rs b/quickwit/quickwit-serve/src/search_api/grpc_adapter.rs index 85297ca1848..38a96d960f3 100644 --- a/quickwit/quickwit-serve/src/search_api/grpc_adapter.rs +++ b/quickwit/quickwit-serve/src/search_api/grpc_adapter.rs @@ -183,4 +183,15 @@ impl grpc::SearchService for GrpcSearchAdapter { let resp = self.0.leaf_list_fields(request.into_inner()).await; convert_to_grpc_result(resp) } + + #[instrument(skip(self, request))] + async fn search_plan( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + set_parent_span_from_request_metadata(request.metadata()); + let search_request = request.into_inner(); + let search_result = self.0.search_plan(search_request).await; + convert_to_grpc_result(search_result) + } } diff --git a/quickwit/quickwit-serve/src/search_api/mod.rs b/quickwit/quickwit-serve/src/search_api/mod.rs index e351a949d27..0b0be17fbeb 100644 --- a/quickwit/quickwit-serve/src/search_api/mod.rs +++ b/quickwit/quickwit-serve/src/search_api/mod.rs @@ -23,8 +23,9 @@ mod rest_handler; pub use self::grpc_adapter::GrpcSearchAdapter; pub(crate) use self::rest_handler::{extract_index_id_patterns, extract_index_id_patterns_default}; pub use self::rest_handler::{ - search_get_handler, search_post_handler, search_request_from_api_request, - search_stream_handler, SearchApi, SearchRequestQueryString, SortBy, + search_get_handler, search_plan_get_handler, search_plan_post_handler, search_post_handler, + search_request_from_api_request, search_stream_handler, SearchApi, SearchRequestQueryString, + SortBy, }; #[cfg(test)] diff --git a/quickwit/quickwit-serve/src/search_api/rest_handler.rs b/quickwit/quickwit-serve/src/search_api/rest_handler.rs index a51fb439e8c..cf9b5d40c84 100644 --- a/quickwit/quickwit-serve/src/search_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/search_api/rest_handler.rs @@ -29,7 +29,7 @@ use quickwit_proto::search::{CountHits, OutputFormat, SortField, SortOrder}; use quickwit_proto::types::IndexId; use quickwit_proto::ServiceError; use quickwit_query::query_ast::query_ast_from_user_text; -use quickwit_search::{SearchError, SearchResponseRest, SearchService}; +use quickwit_search::{SearchError, SearchPlanResponseRest, SearchResponseRest, SearchService}; use serde::{de, Deserialize, Deserializer, Serialize, Serializer}; use serde_json::Value as JsonValue; use tracing::info; @@ -43,12 +43,19 @@ use crate::{with_arg, BodyFormat}; #[derive(utoipa::OpenApi)] #[openapi( - paths(search_get_handler, search_post_handler, search_stream_handler,), + paths( + search_get_handler, + search_post_handler, + search_stream_handler, + search_plan_get_handler, + search_plan_post_handler, + ), components(schemas( BodyFormat, OutputFormat, SearchRequestQueryString, SearchResponseRest, + SearchPlanResponseRest, SortBy, SortField, SortOrder, @@ -318,6 +325,23 @@ fn search_post_filter( .and(warp::body::json()) } +fn search_plan_get_filter( +) -> impl Filter, SearchRequestQueryString), Error = Rejection> + Clone { + warp::path!(String / "search-plan") + .and_then(extract_index_id_patterns) + .and(warp::get()) + .and(serde_qs::warp::query(serde_qs::Config::default())) +} + +fn search_plan_post_filter( +) -> impl Filter, SearchRequestQueryString), Error = Rejection> + Clone { + warp::path!(String / "search-plan") + .and_then(extract_index_id_patterns) + .and(warp::post()) + .and(warp::body::content_length_limit(1024 * 1024)) + .and(warp::body::json()) +} + async fn search( index_id_patterns: Vec, search_request: SearchRequestQueryString, @@ -329,6 +353,22 @@ async fn search( into_rest_api_response(result, body_format) } +async fn search_plan( + index_id_patterns: Vec, + search_request: SearchRequestQueryString, + search_service: Arc, +) -> impl warp::Reply { + let body_format = search_request.format; + let result: Result = async { + let plan_request = search_request_from_api_request(index_id_patterns, search_request)?; + let plan_response = search_service.search_plan(plan_request).await?; + let response = serde_json::from_str(&plan_response.result)?; + Ok(response) + } + .await; + into_rest_api_response(result, body_format) +} + #[utoipa::path( get, tag = "Search", @@ -398,6 +438,52 @@ pub fn search_stream_handler( .then(search_stream) } +#[utoipa::path( + get, + tag = "Search", + path = "/{index_id}/search-plan", + responses( + (status = 200, description = "Metadata about how a request would be executed.", body = SearchPlanResponseRest) + ), + params( + SearchRequestQueryString, + ("index_id" = String, Path, description = "The index ID to search."), + ) +)] +/// Plan Query (GET Variant) +/// +/// Parses the search request from the request query string. +pub fn search_plan_get_handler( + search_service: Arc, +) -> impl Filter + Clone { + search_plan_get_filter() + .and(with_arg(search_service)) + .then(search_plan) +} + +#[utoipa::path( + post, + tag = "Search", + path = "/{index_id}/search-plan", + request_body = SearchRequestQueryString, + responses( + (status = 200, description = "Metadata about how a request would be executed.", body = SearchPlanResponseRest) + ), + params( + ("index_id" = String, Path, description = "The index ID to search."), + ) +)] +/// Plan Query (POST Variant) +/// +/// Parses the search request from the request body. +pub fn search_plan_post_handler( + search_service: Arc, +) -> impl Filter + Clone { + search_plan_post_filter() + .and(with_arg(search_service)) + .then(search_plan) +} + /// This struct represents the search stream query passed to /// the REST API. #[derive(Deserialize, Debug, Eq, PartialEq, utoipa::IntoParams)] @@ -537,7 +623,9 @@ mod tests { let mock_search_service_in_arc = Arc::new(mock_search_service); search_get_handler(mock_search_service_in_arc.clone()) .or(search_post_handler(mock_search_service_in_arc.clone())) - .or(search_stream_handler(mock_search_service_in_arc)) + .or(search_stream_handler(mock_search_service_in_arc.clone())) + .or(search_plan_get_handler(mock_search_service_in_arc.clone())) + .or(search_plan_post_handler(mock_search_service_in_arc.clone())) .recover(recover_fn) }