Merge branch 'add-kql-query' of https://github.com/afoucret/elasticse…

…arch into add-kql-query
afoucret · Nov 7, 2024 · 0992c3a · 0992c3a
2 parents 4fb76e8 + c752ad9
commit 0992c3a
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 10 deletions.
diff --git a/x-pack/plugin/kql/src/main/java/org/elasticsearch/xpack/kql/parser/KqlParser.java b/x-pack/plugin/kql/src/main/java/org/elasticsearch/xpack/kql/parser/KqlParser.java
@@ -14,6 +14,7 @@
 import org.antlr.v4.runtime.RecognitionException;
 import org.antlr.v4.runtime.Recognizer;
 import org.antlr.v4.runtime.atn.PredictionMode;
+import org.elasticsearch.core.Strings;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
@@ -25,10 +26,7 @@ public class KqlParser {
     private static final Logger log = LogManager.getLogger(KqlParser.class);
 
     public QueryBuilder parseKqlQuery(String kqlQuery, KqlParsingContext kqlParserContext) {
-        if (log.isTraceEnabled()) {
-            log.debug("Parsing KQL query: {}", kqlQuery);
-        }
-
+        log.trace("Parsing KQL query: {}", kqlQuery);
         return invokeParser(kqlQuery, kqlParserContext, KqlBaseParser::topLevelQuery, KqlAstBuilder::toQueryBuilder);
     }
 
@@ -53,9 +51,7 @@ private <T> T invokeParser(
 
         ParserRuleContext tree = parseFunction.apply(parser);
 
-        if (log.isTraceEnabled()) {
-            log.trace("Parse tree: {}", tree.toStringTree());
-        }
+        log.trace(() -> Strings.format("Parse tree: %s", tree.toStringTree()));
 
         return visitor.apply(new KqlAstBuilder(kqlParsingContext), tree);
     }

diff --git a/x-pack/plugin/kql/src/main/java/org/elasticsearch/xpack/kql/query/KqlQueryBuilder.java b/x-pack/plugin/kql/src/main/java/org/elasticsearch/xpack/kql/query/KqlQueryBuilder.java
@@ -154,9 +154,7 @@ protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) throw
         KqlParser parser = new KqlParser();
         QueryBuilder rewrittenQuery = parser.parseKqlQuery(query, createKqlParserContext(context));
 
-        if (log.isTraceEnabled()) {
-            log.trace("KQL query {} translated to Query DSL: {}", query, Strings.toString(rewrittenQuery));
-        }
+        log.trace(() -> Strings.format("KQL query %s translated to Query DSL: %s", query, Strings.toString(rewrittenQuery)));
 
         return rewrittenQuery;
     }

diff --git a/x-pack/plugin/kql/src/test/java/org/elasticsearch/xpack/kql/query/KqlQueryBuilderTests.java b/x-pack/plugin/kql/src/test/java/org/elasticsearch/xpack/kql/query/KqlQueryBuilderTests.java
@@ -10,6 +10,7 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.index.query.MultiMatchQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.QueryStringQueryBuilder;
 import org.elasticsearch.index.query.RangeQueryBuilder;
@@ -264,4 +265,22 @@ public void testDefaultFieldMatchQuery() throws IOException {
             assertThat(rewritenQuery.fields().keySet(), contains("mapped_object.mapped_date", "mapped_object.mapped_int"));
         }
     }
+
+    public void testQueryNameIsPreserved() throws IOException {
+        QueryRewriteContext queryRewriteContext = createQueryRewriteContext();
+        SearchExecutionContext searchExecutionContext = createSearchExecutionContext();
+
+        KqlQueryBuilder kqlQuery = new KqlQueryBuilder(generateRandomKqlQuery()).queryName(randomIdentifier());
+        QueryBuilder rewrittenQuery = rewriteQuery(kqlQuery, queryRewriteContext, searchExecutionContext);
+        assertThat(rewrittenQuery.queryName(), equalTo(kqlQuery.queryName()));
+    }
+
+    public void testQueryBoostIsPreserved() throws IOException {
+        QueryRewriteContext queryRewriteContext = createQueryRewriteContext();
+        SearchExecutionContext searchExecutionContext = createSearchExecutionContext();
+
+        KqlQueryBuilder kqlQuery = new KqlQueryBuilder(generateRandomKqlQuery()).boost(randomFloatBetween(0, Float.MAX_VALUE, true));
+        QueryBuilder rewrittenQuery = rewriteQuery(kqlQuery, queryRewriteContext, searchExecutionContext);
+        assertThat(rewrittenQuery.boost(), equalTo(kqlQuery.boost()));
+    }
 }