Skip to content

Commit

Permalink
Add kql query to the DSL (elastic#116262)
Browse files Browse the repository at this point in the history 
(cherry picked from commit e2c29f5)

# Conflicts:
#	server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java
  • Loading branch information
@afoucret
afoucret committed Nov 8, 2024
1 parent b7e96f1 commit 0875ab0
Show file tree
Hide file tree
Showing 20 changed files with 1,846 additions and 229 deletions.
1 change: 1 addition & 0 deletions server/src/main/java/org/elasticsearch/TransportVersions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,7 @@ static TransportVersion def(int id) {
public static final TransportVersion ESQL_CCS_EXEC_INFO_WITH_FAILURES = def(8_783_00_0);
public static final TransportVersion LOGSDB_TELEMETRY = def(8_784_00_0);
public static final TransportVersion LOGSDB_TELEMETRY_STATS = def(8_785_00_0);
public static final TransportVersion KQL_QUERY_ADDED = def(8_786_00_0);

/*
* STOP! READ THIS FIRST! No, really,
Expand Down
30 changes: 25 additions & 5 deletions server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@

package org.elasticsearch.rest.action.search;

import org.elasticsearch.Build;
import org.elasticsearch.common.util.set.Sets;

import java.util.Collections;
import java.util.Set;

/**
Expand All @@ -24,10 +28,26 @@ private SearchCapabilities() {}
private static final String BIT_DENSE_VECTOR_SYNTHETIC_SOURCE_CAPABILITY = "bit_dense_vector_synthetic_source";
/** Support Byte and Float with Bit dot product. */
private static final String BYTE_FLOAT_BIT_DOT_PRODUCT_CAPABILITY = "byte_float_bit_dot_product";
/** Support kql query. */
private static final String KQL_QUERY_SUPPORTED = "kql_query";

public static final Set<String> CAPABILITIES = capabilities();

private static Set<String> capabilities() {
Set<String> capabilities = Set.of(
RANGE_REGEX_INTERVAL_QUERY_CAPABILITY,
BIT_DENSE_VECTOR_SYNTHETIC_SOURCE_CAPABILITY,
BYTE_FLOAT_BIT_DOT_PRODUCT_CAPABILITY,
);

if (Build.current().isSnapshot()) {
return Collections.unmodifiableSet(Sets.union(capabilities, snapshotBuildCapabilities()));
}

return capabilities;
}

public static final Set<String> CAPABILITIES = Set.of(
RANGE_REGEX_INTERVAL_QUERY_CAPABILITY,
BIT_DENSE_VECTOR_SYNTHETIC_SOURCE_CAPABILITY,
BYTE_FLOAT_BIT_DOT_PRODUCT_CAPABILITY
);
private static Set<String> snapshotBuildCapabilities() {
return Set.of(KQL_QUERY_SUPPORTED);
}
}
18 changes: 11 additions & 7 deletions x-pack/plugin/kql/build.gradle
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import org.elasticsearch.gradle.internal.info.BuildParams

import static org.elasticsearch.gradle.util.PlatformUtils.normalize

apply plugin: 'elasticsearch.internal-es-plugin'
apply plugin: 'elasticsearch.internal-cluster-test'
apply plugin: 'elasticsearch.internal-yaml-rest-test'
apply plugin: 'elasticsearch.publish'

esplugin {
Expand All @@ -17,19 +19,21 @@ base {

dependencies {
compileOnly project(path: xpackModule('core'))
api "org.antlr:antlr4-runtime:${versions.antlr4}"
implementation "org.antlr:antlr4-runtime:${versions.antlr4}"

testImplementation "org.antlr:antlr4-runtime:${versions.antlr4}"
testImplementation project(':test:framework')
testImplementation(testArtifact(project(xpackModule('core'))))
}

/****************************************************************
* Enable QA/rest integration tests for snapshot builds only *
* TODO: Enable for all builds upon this feature release *
****************************************************************/
if (BuildParams.isSnapshotBuild()) {
addQaCheckDependencies(project)
tasks.named('yamlRestTest') {
usesDefaultDistribution()
}.configure {
/****************************************************************
* Enable QA/rest integration tests for snapshot builds only *
* TODO: Enable for all builds upon this feature release *
****************************************************************/
enabled = BuildParams.isSnapshotBuild()
}

/**********************************
Expand Down
2 changes: 1 addition & 1 deletion x-pack/plugin/kql/src/main/antlr/KqlBase.g4
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ fieldQueryValue
;

fieldName
: value=UNQUOTED_LITERAL+
: value=UNQUOTED_LITERAL
| value=QUOTED_STRING
| value=WILDCARD
;
Expand Down
1 change: 1 addition & 0 deletions x-pack/plugin/kql/src/main/java/module-info.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,4 +16,5 @@

exports org.elasticsearch.xpack.kql;
exports org.elasticsearch.xpack.kql.parser;
exports org.elasticsearch.xpack.kql.query;
}
11 changes: 11 additions & 0 deletions x-pack/plugin/kql/src/main/java/org/elasticsearch/xpack/kql/KqlPlugin.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,21 @@

package org.elasticsearch.xpack.kql;

import org.elasticsearch.Build;
import org.elasticsearch.plugins.ExtensiblePlugin;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.plugins.SearchPlugin;
import org.elasticsearch.xpack.kql.query.KqlQueryBuilder;

import java.util.List;

public class KqlPlugin extends Plugin implements SearchPlugin, ExtensiblePlugin {
@Override
public List<QuerySpec<?>> getQueries() {
if (Build.current().isSnapshot()) {
return List.of(new SearchPlugin.QuerySpec<>(KqlQueryBuilder.NAME, KqlQueryBuilder::new, KqlQueryBuilder::fromXContent));
}

return List.of();
}
}
97 changes: 69 additions & 28 deletions x-pack/plugin/kql/src/main/java/org/elasticsearch/xpack/kql/parser/KqlAstBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,36 +9,42 @@

import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.Token;
import org.elasticsearch.common.regex.Regex;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.MatchNoneQueryBuilder;
import org.elasticsearch.index.query.MultiMatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.QueryStringQueryBuilder;
import org.elasticsearch.index.query.RangeQueryBuilder;

import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;

import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
import static org.elasticsearch.xpack.kql.parser.KqlParserExecutionContext.isDateField;
import static org.elasticsearch.xpack.kql.parser.KqlParserExecutionContext.isKeywordField;
import static org.elasticsearch.xpack.kql.parser.KqlParserExecutionContext.isRuntimeField;
import static org.elasticsearch.xpack.kql.parser.KqlParsingContext.isDateField;
import static org.elasticsearch.xpack.kql.parser.KqlParsingContext.isKeywordField;
import static org.elasticsearch.xpack.kql.parser.KqlParsingContext.isRuntimeField;
import static org.elasticsearch.xpack.kql.parser.KqlParsingContext.isSearchableField;
import static org.elasticsearch.xpack.kql.parser.ParserUtils.escapeLuceneQueryString;
import static org.elasticsearch.xpack.kql.parser.ParserUtils.extractText;
import static org.elasticsearch.xpack.kql.parser.ParserUtils.hasWildcard;
import static org.elasticsearch.xpack.kql.parser.ParserUtils.typedParsing;

class KqlAstBuilder extends KqlBaseBaseVisitor<QueryBuilder> {
private final KqlParserExecutionContext kqlParserExecutionContext;
private final KqlParsingContext kqlParsingContext;

KqlAstBuilder(KqlParserExecutionContext kqlParserExecutionContext) {
this.kqlParserExecutionContext = kqlParserExecutionContext;
KqlAstBuilder(KqlParsingContext kqlParsingContext) {
this.kqlParsingContext = kqlParsingContext;
}

public QueryBuilder toQueryBuilder(ParserRuleContext ctx) {
if (ctx instanceof KqlBaseParser.TopLevelQueryContext topLeveQueryContext) {
if (topLeveQueryContext.query() != null) {
return ParserUtils.typedParsing(this, topLeveQueryContext.query(), QueryBuilder.class);
return typedParsing(this, topLeveQueryContext.query(), QueryBuilder.class);
}

return new MatchAllQueryBuilder();
Expand All @@ -59,9 +65,9 @@ public QueryBuilder visitAndBooleanQuery(KqlBaseParser.BooleanQueryContext ctx)
// TODO: KQLContext has an option to wrap the clauses into a filter instead of a must clause. Do we need it?
for (ParserRuleContext subQueryCtx : ctx.query()) {
if (subQueryCtx instanceof KqlBaseParser.BooleanQueryContext booleanSubQueryCtx && isAndQuery(booleanSubQueryCtx)) {
ParserUtils.typedParsing(this, subQueryCtx, BoolQueryBuilder.class).must().forEach(builder::must);
typedParsing(this, subQueryCtx, BoolQueryBuilder.class).must().forEach(builder::must);
} else {
builder.must(ParserUtils.typedParsing(this, subQueryCtx, QueryBuilder.class));
builder.must(typedParsing(this, subQueryCtx, QueryBuilder.class));
}
}

Expand All @@ -73,9 +79,9 @@ public QueryBuilder visitOrBooleanQuery(KqlBaseParser.BooleanQueryContext ctx) {

for (ParserRuleContext subQueryCtx : ctx.query()) {
if (subQueryCtx instanceof KqlBaseParser.BooleanQueryContext booleanSubQueryCtx && isOrQuery(booleanSubQueryCtx)) {
ParserUtils.typedParsing(this, subQueryCtx, BoolQueryBuilder.class).should().forEach(builder::should);
typedParsing(this, subQueryCtx, BoolQueryBuilder.class).should().forEach(builder::should);
} else {
builder.should(ParserUtils.typedParsing(this, subQueryCtx, QueryBuilder.class));
builder.should(typedParsing(this, subQueryCtx, QueryBuilder.class));
}
}

Expand All @@ -84,12 +90,12 @@ public QueryBuilder visitOrBooleanQuery(KqlBaseParser.BooleanQueryContext ctx) {

@Override
public QueryBuilder visitNotQuery(KqlBaseParser.NotQueryContext ctx) {
return QueryBuilders.boolQuery().mustNot(ParserUtils.typedParsing(this, ctx.simpleQuery(), QueryBuilder.class));
return QueryBuilders.boolQuery().mustNot(typedParsing(this, ctx.simpleQuery(), QueryBuilder.class));
}

@Override
public QueryBuilder visitParenthesizedQuery(KqlBaseParser.ParenthesizedQueryContext ctx) {
return ParserUtils.typedParsing(this, ctx.query(), QueryBuilder.class);
return typedParsing(this, ctx.query(), QueryBuilder.class);
}

@Override
Expand Down Expand Up @@ -121,12 +127,16 @@ public QueryBuilder visitExistsQuery(KqlBaseParser.ExistsQueryContext ctx) {
public QueryBuilder visitRangeQuery(KqlBaseParser.RangeQueryContext ctx) {
BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery().minimumShouldMatch(1);

String queryText = ParserUtils.extractText(ctx.rangeQueryValue());
String queryText = extractText(ctx.rangeQueryValue());
BiFunction<RangeQueryBuilder, String, RangeQueryBuilder> rangeOperation = rangeOperation(ctx.operator);

withFields(ctx.fieldName(), (fieldName, mappedFieldType) -> {
RangeQueryBuilder rangeQuery = rangeOperation.apply(QueryBuilders.rangeQuery(fieldName), queryText);
// TODO: add timezone for date fields

if (kqlParsingContext.timeZone() != null) {
rangeQuery.timeZone(kqlParsingContext.timeZone().getId());
}

boolQueryBuilder.should(rangeQuery);
});

Expand All @@ -135,42 +145,54 @@ public QueryBuilder visitRangeQuery(KqlBaseParser.RangeQueryContext ctx) {

@Override
public QueryBuilder visitFieldLessQuery(KqlBaseParser.FieldLessQueryContext ctx) {
String queryText = ParserUtils.extractText(ctx.fieldQueryValue());
String queryText = extractText(ctx.fieldQueryValue());

if (hasWildcard(ctx.fieldQueryValue())) {
// TODO: set default fields.
return QueryBuilders.queryStringQuery(escapeLuceneQueryString(queryText, true));
QueryStringQueryBuilder queryString = QueryBuilders.queryStringQuery(escapeLuceneQueryString(queryText, true));
if (kqlParsingContext.defaultField() != null) {
queryString.defaultField(kqlParsingContext.defaultField());
}
return queryString;
}

boolean isPhraseMatch = ctx.fieldQueryValue().QUOTED_STRING() != null;

return QueryBuilders.multiMatchQuery(queryText)
// TODO: add default fields?
MultiMatchQueryBuilder multiMatchQuery = QueryBuilders.multiMatchQuery(queryText)
.type(isPhraseMatch ? MultiMatchQueryBuilder.Type.PHRASE : MultiMatchQueryBuilder.Type.BEST_FIELDS)
.lenient(true);

if (kqlParsingContext.defaultField() != null) {
kqlParsingContext.resolveDefaultFieldNames()
.stream()
.filter(kqlParsingContext::isSearchableField)
.forEach(multiMatchQuery::field);
}

return multiMatchQuery;
}

@Override
public QueryBuilder visitFieldQuery(KqlBaseParser.FieldQueryContext ctx) {

BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery().minimumShouldMatch(1);
String queryText = ParserUtils.extractText(ctx.fieldQueryValue());
String queryText = extractText(ctx.fieldQueryValue());
boolean hasWildcard = hasWildcard(ctx.fieldQueryValue());

withFields(ctx.fieldName(), (fieldName, mappedFieldType) -> {
QueryBuilder fieldQuery = null;

if (hasWildcard && isKeywordField(mappedFieldType)) {
fieldQuery = QueryBuilders.wildcardQuery(fieldName, queryText)
.caseInsensitive(kqlParserExecutionContext.isCaseSensitive() == false);
fieldQuery = QueryBuilders.wildcardQuery(fieldName, queryText).caseInsensitive(kqlParsingContext.caseInsensitive());
} else if (hasWildcard) {
fieldQuery = QueryBuilders.queryStringQuery(escapeLuceneQueryString(queryText, true)).field(fieldName);
} else if (isDateField(mappedFieldType)) {
// TODO: add timezone
fieldQuery = QueryBuilders.rangeQuery(fieldName).gte(queryText).lte(queryText);
RangeQueryBuilder rangeFieldQuery = QueryBuilders.rangeQuery(fieldName).gte(queryText).lte(queryText);
if (kqlParsingContext.timeZone() != null) {
rangeFieldQuery.timeZone(kqlParsingContext.timeZone().getId());
}
fieldQuery = rangeFieldQuery;
} else if (isKeywordField(mappedFieldType)) {
fieldQuery = QueryBuilders.termQuery(fieldName, queryText)
.caseInsensitive(kqlParserExecutionContext.isCaseSensitive() == false);
fieldQuery = QueryBuilders.termQuery(fieldName, queryText).caseInsensitive(kqlParsingContext.caseInsensitive());
} else if (ctx.fieldQueryValue().QUOTED_STRING() != null) {
fieldQuery = QueryBuilders.matchPhraseQuery(fieldName, queryText);
} else {
Expand All @@ -194,7 +216,26 @@ private static boolean isOrQuery(KqlBaseParser.BooleanQueryContext ctx) {
}

private void withFields(KqlBaseParser.FieldNameContext ctx, BiConsumer<String, MappedFieldType> fieldConsummer) {
kqlParserExecutionContext.resolveFields(ctx).forEach(fieldDef -> fieldConsummer.accept(fieldDef.v1(), fieldDef.v2()));
assert ctx != null : "Field ctx cannot be null";
String fieldNamePattern = extractText(ctx);
Set<String> fieldNames = kqlParsingContext.resolveFieldNames(fieldNamePattern);

if (ctx.value.getType() == KqlBaseParser.QUOTED_STRING && Regex.isSimpleMatchPattern(fieldNamePattern)) {
// When using quoted string, wildcards are not expanded.
// No field can match and we can return early.
return;
}

if (ctx.value.getType() == KqlBaseParser.QUOTED_STRING) {
assert fieldNames.size() < 2 : "expecting only one matching field";
}

fieldNames.forEach(fieldName -> {
MappedFieldType fieldType = kqlParsingContext.fieldType(fieldName);
if (isSearchableField(fieldName, fieldType)) {
fieldConsummer.accept(fieldName, fieldType);
}
});
}

private QueryBuilder rewriteDisjunctionQuery(BoolQueryBuilder boolQueryBuilder) {
Expand Down
Loading

0 comments on commit 0875ab0

Please sign in to comment.