Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix query date parsing in the query nodes #1081

Merged
merged 5 commits into from
Sep 19, 2024
Merged

Fix query date parsing in the query nodes #1081

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2cc59e7
Merged
kyle-sammons Sep 19, 2024
f754ec5
Remove other stuff
kyle-sammons Sep 19, 2024
15310b7
Remove this too
kyle-sammons Sep 19, 2024
73bbb11
Formatting
kyle-sammons Sep 19, 2024
e82a37b
unnecessary line
ermontross Sep 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 42 additions & 103 deletions astra/src/main/java/com/slack/astra/elasticsearchApi/OpenSearchRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,12 @@
package com.slack.astra.elasticsearchApi;

import static org.apache.lucene.document.LongPoint.decodeDimension;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.slack.astra.logstore.opensearch.AstraBigArrays;
import com.slack.astra.logstore.opensearch.AstraIndexSettings;
import com.slack.astra.logstore.opensearch.OpenSearchAdapter;
import com.slack.astra.logstore.search.SearchResultUtils;
import com.slack.astra.logstore.search.aggregations.AutoDateHistogramAggBuilder;
Expand All @@ -31,32 +27,22 @@
import com.slack.astra.logstore.search.aggregations.UniqueCountAggBuilder;
import com.slack.astra.metadata.schema.LuceneFieldDef;
import com.slack.astra.proto.service.AstraSearch;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import org.apache.commons.lang3.NotImplementedException;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.IndexOrDocValuesQuery;
import org.apache.lucene.search.IndexSortSortedNumericDocValuesRangeQuery;
import org.apache.lucene.search.PointRangeQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.QueryVisitor;
import org.apache.lucene.search.TermRangeQuery;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.BytesRef;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.xcontent.json.JsonXContentParser;
import org.opensearch.core.xcontent.DeprecationHandler;
import org.opensearch.core.xcontent.NamedXContentRegistry;
import org.opensearch.index.query.AbstractQueryBuilder;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.index.query.QueryBuilderVisitor;
import org.opensearch.index.query.RangeQueryBuilder;
import org.opensearch.search.SearchModule;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -78,75 +64,42 @@ public class OpenSearchRequest {
"@timestamp", new LuceneFieldDef("@timestamp", "date", true, true, true)));
private static final Logger log = LoggerFactory.getLogger(OpenSearchRequest.class);

// Used to get the date ranges from the query
private static class DateRangeQueryVisitor extends QueryVisitor {
private static class DateRangeQueryBuilderVistor implements QueryBuilderVisitor {
private Long dateRangeStart;
private Long dateRangeEnd;

private String toString(byte[] value) {
return Long.toString(decodeDimension(value, 0));
}

private void checkTimeRange(Query query) {
if (query instanceof TermRangeQuery rangeQuery) {
BytesRef lowerTerm = rangeQuery.getLowerTerm();
BytesRef upperTerm = rangeQuery.getUpperTerm();
@Override
public void accept(QueryBuilder qb) {
if (qb instanceof RangeQueryBuilder) {
RangeQueryBuilder rangeQueryBuilder = (RangeQueryBuilder) qb;
if (!rangeQueryBuilder.fieldName().equals("@timestamp")
&& !rangeQueryBuilder.fieldName().equals("_timesinceepoch")) {
return;
}

if (rangeQuery.getField().equals("@timestamp")
|| rangeQuery.getField().equals("_timesinceepoch")) {
String lowerDateString = lowerTerm.utf8ToString();
String upperDateString = upperTerm.utf8ToString();
Object from = rangeQueryBuilder.from();
Object to = rangeQueryBuilder.to();

this.dateRangeStart = Instant.parse(lowerDateString).toEpochMilli();
this.dateRangeEnd = Instant.parse(upperDateString).toEpochMilli();
if (from instanceof Long) {
dateRangeStart = (Long) from;
} else if (from instanceof Integer) {
dateRangeStart = ((Integer) from).longValue();
} else if (from instanceof String) {
dateRangeStart = Instant.parse((String) from).toEpochMilli();
}
} else if (query
instanceof
IndexSortSortedNumericDocValuesRangeQuery indexSortSortedNumericDocValuesRangeQuery) {
Query fallbackQuery = indexSortSortedNumericDocValuesRangeQuery.getFallbackQuery();
checkTimeRange(fallbackQuery);
} else if (query instanceof IndexOrDocValuesQuery indexOrDocValuesQuery) {
Query indexQuery = indexOrDocValuesQuery.getIndexQuery();
checkTimeRange(indexQuery);
} else if (query instanceof PointRangeQuery pointRangeQuery) {

if (pointRangeQuery.getField().equals("@timestamp")
|| pointRangeQuery.getField().equals("_timesinceepoch")) {
String lowerDateString =
new String(pointRangeQuery.getLowerPoint(), StandardCharsets.UTF_8);
String upperDateString =
new String(pointRangeQuery.getUpperPoint(), StandardCharsets.UTF_8);

for (int i = 0; i < pointRangeQuery.getNumDims(); i++) {
int startOffset = pointRangeQuery.getBytesPerDim() * i;
lowerDateString =
toString(
ArrayUtil.copyOfSubArray(
pointRangeQuery.getLowerPoint(),
startOffset,
startOffset + pointRangeQuery.getBytesPerDim()));
upperDateString =
toString(
ArrayUtil.copyOfSubArray(
pointRangeQuery.getUpperPoint(),
startOffset,
startOffset + pointRangeQuery.getBytesPerDim()));
}

this.dateRangeStart = Long.valueOf(lowerDateString);
this.dateRangeEnd = Long.valueOf(upperDateString);
if (to instanceof Long) {
dateRangeEnd = (Long) to;
} else if (to instanceof Integer) {
dateRangeEnd = ((Integer) to).longValue();
} else if (to instanceof String) {
dateRangeEnd = Instant.parse((String) to).toEpochMilli();
}
}
}

@Override
public void visitLeaf(Query query) {
checkTimeRange(query);
}

@Override
public QueryVisitor getSubVisitor(BooleanClause.Occur occur, Query parent) {
checkTimeRange(parent);
public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
return this;
}
}
Expand All @@ -165,7 +118,13 @@ public List<AstraSearch.SearchRequest> parseHttpPostBody(String postBody)
JsonNode header = OM.readTree(pair.get(0));
JsonNode body = OM.readTree(pair.get(1));
String query = getQuery(body);
Optional<DateRangeQueryVisitor> dateRangeQueryVisitor = Optional.of(getDateRange(query));
DateRangeQueryBuilderVistor dateRangeQueryBuilderVistor = getDateRange(query);
Long startTimeEpochMs = null;
Long endTimeEpochMs = null;
if (dateRangeQueryBuilderVistor != null) {
startTimeEpochMs = dateRangeQueryBuilderVistor.dateRangeStart;
endTimeEpochMs = dateRangeQueryBuilderVistor.dateRangeEnd;
}

searchRequests.add(
AstraSearch.SearchRequest.newBuilder()
Expand All @@ -174,10 +133,8 @@ public List<AstraSearch.SearchRequest> parseHttpPostBody(String postBody)
.setAggregations(getAggregations(body))
.setQuery(getQuery(body))
.setSourceFieldFilter(getSourceFieldFilter(body))
.setStartTimeEpochMs(
dateRangeQueryVisitor.orElse(new DateRangeQueryVisitor()).dateRangeStart)
.setEndTimeEpochMs(
dateRangeQueryVisitor.orElse(new DateRangeQueryVisitor()).dateRangeEnd)
.setStartTimeEpochMs(startTimeEpochMs)
.setEndTimeEpochMs(endTimeEpochMs)
.build());
}
return searchRequests;
Expand Down Expand Up @@ -253,7 +210,7 @@ private static AstraSearch.SearchRequest.SourceFieldFilter getSourceFieldFilter(
return AstraSearch.SearchRequest.SourceFieldFilter.newBuilder().build();
}

private static DateRangeQueryVisitor getDateRange(String queryBody) {
private static DateRangeQueryBuilderVistor getDateRange(String queryBody) {
try {
openSearchAdapter.reloadSchema();
JsonXContentParser jsonXContentParser =
Expand All @@ -262,31 +219,13 @@ private static DateRangeQueryVisitor getDateRange(String queryBody) {
new SearchModule(Settings.EMPTY, List.of()).getNamedXContents()),
DeprecationHandler.IGNORE_DEPRECATIONS,
OM.createParser(queryBody));

QueryBuilder queryBuilder = AbstractQueryBuilder.parseInnerQueryBuilder(jsonXContentParser);
QueryShardContext queryShardContext =
new QueryShardContext(
0,
AstraIndexSettings.getInstance(),
AstraBigArrays.getInstance(),
null,
null,
openSearchAdapter.getMapperService(),
null,
null,
null,
null,
null,
null,
() -> Instant.now().toEpochMilli(),
null,
s -> false,
() -> true,
null);
Query query = queryBuilder.rewrite(queryShardContext).toQuery(queryShardContext);
DateRangeQueryVisitor dateRangeQueryVisitor = new DateRangeQueryVisitor();
query.visit(dateRangeQueryVisitor);
return dateRangeQueryVisitor;
} catch (IOException e) {
DateRangeQueryBuilderVistor dateRangeQueryBuilderVistor = new DateRangeQueryBuilderVistor();
queryBuilder.visit(dateRangeQueryBuilderVistor);
return dateRangeQueryBuilderVistor;

} catch (Exception e) {
log.error("Unable to parse date/time range from query body: {}. Error: {}", queryBody, e);
return null;
}
Expand Down
Loading