From 62a55a2a2e5f862dc3b1f610eea9ffc3a2fa2e02 Mon Sep 17 00:00:00 2001
From: Vladimir Kuznichenkov <kuzaxak.tech@gmail.com>
Date: Sat, 8 Jun 2024 22:22:50 +0300
Subject: [PATCH] Add missing fields to ES compatible Query API

* [_source][1] schema
* [docvalue][2] schema

I've added additional fields as `serde_json::Value` to prevent complex
struct definitions as we won't use them anyway.

[1]: https://www.elastic.co/guide/en/elasticsearch/reference/7.17/search-fields.html#source-filtering
[2]: https://www.elastic.co/guide/en/elasticsearch/reference/current/search-fields.html#docvalue-fields
---
 .../src/elastic_query_dsl/range_query.rs      |  3 +++
 .../elasticsearch_api/model/search_body.rs    | 25 +++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/quickwit/quickwit-query/src/elastic_query_dsl/range_query.rs b/quickwit/quickwit-query/src/elastic_query_dsl/range_query.rs
index 3fe74c28053..5fc1d970f78 100644
--- a/quickwit/quickwit-query/src/elastic_query_dsl/range_query.rs
+++ b/quickwit/quickwit-query/src/elastic_query_dsl/range_query.rs
@@ -40,6 +40,8 @@ pub struct RangeQueryParams {
     lte: Option<JsonLiteral>,
     #[serde(default)]
     boost: Option<NotNaNf32>,
+    #[serde(default)]
+    format: Option<JsonLiteral>,
 }
 
 pub type RangeQuery = OneFieldMap<RangeQueryParams>;
@@ -53,6 +55,7 @@ impl ConvertableToQueryAst for RangeQuery {
             lt,
             lte,
             boost,
+            format: _,
         } = self.value;
         let range_query_ast = crate::query_ast::RangeQuery {
             field,
diff --git a/quickwit/quickwit-serve/src/elasticsearch_api/model/search_body.rs b/quickwit/quickwit-serve/src/elasticsearch_api/model/search_body.rs
index 51b7296d516..2f57a2b3aeb 100644
--- a/quickwit/quickwit-serve/src/elasticsearch_api/model/search_body.rs
+++ b/quickwit/quickwit-serve/src/elasticsearch_api/model/search_body.rs
@@ -66,9 +66,28 @@ struct FieldSortParams {
     pub date_format: Option<ElasticDateFormat>,
 }
 
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
+pub struct SourceObject {
+    includes: Option<Vec<String>>,
+    excludes: Option<Vec<String>>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
+#[serde(untagged)]
+pub enum Source {
+    Bool(bool),
+    String(String),
+    List(Vec<String>),
+    Object(SourceObject),
+}
+
 #[derive(Debug, Default, Clone, Deserialize, PartialEq)]
 #[serde(deny_unknown_fields)]
 pub struct SearchBody {
+    #[serde(default)]
+    pub _source: Option<Source>,
+    #[serde(default)]
+    pub docvalue_fields: Option<serde_json::Value>,
     #[serde(default)]
     pub from: Option<u64>,
     #[serde(default)]
@@ -86,6 +105,12 @@ pub struct SearchBody {
     pub stored_fields: Option<BTreeSet<String>>,
     #[serde(default)]
     pub search_after: Vec<serde_json::Value>,
+    #[serde(default)]
+    pub script_fields: Option<serde_json::Value>,
+    #[serde(default)]
+    pub highlight: Option<serde_json::Value>,
+    #[serde(default)]
+    pub version: Option<bool>,
 }
 
 struct FieldSortVecVisitor;