Add search "EXPLAIN" endpoint

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/ElasticsearchBackend.java

@@ -16,11 +16,13 @@
  */
 package org.graylog.storage.elasticsearch7.views;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import io.opentelemetry.instrumentation.annotations.WithSpan;
 import jakarta.inject.Inject;
 import jakarta.inject.Named;
 import jakarta.inject.Provider;
+import org.graylog.plugins.views.search.ExplainResults;
 import org.graylog.plugins.views.search.Filter;
 import org.graylog.plugins.views.search.GlobalOverride;
 import org.graylog.plugins.views.search.Query;
@@ -218,6 +220,26 @@ public Optional<QueryBuilder> generateFilterClause(Filter filter) {
         return Optional.empty();
     }
 
+    @Override
+    public ExplainResults.QueryExplainResult doExplain(SearchJob job, Query query, ESGeneratedQueryContext queryContext) {
+        final ImmutableMap.Builder<String, ExplainResults.ExplainResult> builder = ImmutableMap.builder();
+        final Map<String, SearchSourceBuilder> searchTypeQueries = queryContext.searchTypeQueries();
+
+        final DateTime nowUTCSharedBetweenSearchTypes = Tools.nowUTC();
+
+        query.searchTypes().forEach(s -> {
+            final Set<ExplainResults.IndexRangeResult> indicesForQuery = indexLookup.indexRangesForStreamsInTimeRange(
+                            query.effectiveStreams(s), query.effectiveTimeRange(s, nowUTCSharedBetweenSearchTypes))
+                    .stream().map(ExplainResults.IndexRangeResult::fromIndexRange).collect(Collectors.toSet());
+
+            final var queryString = searchTypeQueries.get(s.id()).toString();
+
+            builder.put(s.id(), new ExplainResults.ExplainResult(queryString, indicesForQuery));
+        });
+
+        return new ExplainResults.QueryExplainResult(builder.build());
+    }
+
     @WithSpan
     @Override
     public QueryResult doRun(SearchJob job, Query query, ESGeneratedQueryContext queryContext) {

graylog-storage-elasticsearch7/src/test/java/org/graylog/storage/elasticsearch7/views/ElasticsearchBackendTest.java

@@ -19,6 +19,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
+import com.jayway.jsonpath.JsonPath;
 import jakarta.inject.Provider;
 import org.graylog.plugins.views.search.LegacyDecoratorProcessor;
 import org.graylog.plugins.views.search.Query;
@@ -28,49 +29,75 @@
 import org.graylog.plugins.views.search.SearchType;
 import org.graylog.plugins.views.search.elasticsearch.ElasticsearchQueryString;
 import org.graylog.plugins.views.search.elasticsearch.IndexLookup;
+import org.graylog.plugins.views.search.engine.GeneratedQueryContext;
 import org.graylog.plugins.views.search.engine.monitoring.collection.NoOpStatsCollector;
 import org.graylog.plugins.views.search.searchfilters.db.UsedSearchFiltersToQueryStringsMapper;
 import org.graylog.plugins.views.search.searchfilters.model.InlineQueryStringSearchFilter;
 import org.graylog.plugins.views.search.searchfilters.model.ReferencedQueryStringSearchFilter;
 import org.graylog.plugins.views.search.searchfilters.model.UsedSearchFilter;
 import org.graylog.plugins.views.search.searchtypes.MessageList;
+import org.graylog.plugins.views.search.searchtypes.pivot.Pivot;
+import org.graylog.plugins.views.search.searchtypes.pivot.buckets.AutoInterval;
+import org.graylog.plugins.views.search.searchtypes.pivot.buckets.Time;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.index.query.BoolQueryBuilder;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.index.query.QueryBuilder;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.index.query.QueryStringQueryBuilder;
 import org.graylog.storage.elasticsearch7.views.searchtypes.ESMessageList;
 import org.graylog.storage.elasticsearch7.views.searchtypes.ESSearchTypeHandler;
+import org.graylog.storage.elasticsearch7.views.searchtypes.pivot.ESPivot;
+import org.graylog.storage.elasticsearch7.views.searchtypes.pivot.EffectiveTimeRangeExtractor;
+import org.graylog.storage.elasticsearch7.views.searchtypes.pivot.buckets.ESTimeHandler;
+import org.graylog.testing.jsonpath.JsonPathAssert;
+import org.graylog2.indexer.ranges.MongoIndexRange;
 import org.graylog2.indexer.results.TestResultMessageFactory;
+import org.graylog2.plugin.indexer.searches.timeranges.AbsoluteRange;
 import org.graylog2.plugin.indexer.searches.timeranges.RelativeRange;
+import org.graylog2.plugin.indexer.searches.timeranges.TimeRange;
+import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.graylog2.plugin.Tools.nowUTC;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anySet;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 
 public class ElasticsearchBackendTest {
+    @Rule
+    public final MockitoRule mockitoRule = MockitoJUnit.rule();
     private ElasticsearchBackend backend;
     private UsedSearchFiltersToQueryStringsMapper usedSearchFiltersToQueryStringsMapper;
 
+    @Mock
+    private IndexLookup indexLookup;
+
     @Before
     public void setup() {
         Map<String, Provider<ESSearchTypeHandler<? extends SearchType>>> handlers = Maps.newHashMap();
         handlers.put(MessageList.NAME, () -> new ESMessageList(new LegacyDecoratorProcessor.Fake(),
                 new TestResultMessageFactory(), false));
+        handlers.put(Pivot.NAME, () -> new ESPivot(Map.of(Time.NAME, new ESTimeHandler()), Map.of(), new EffectiveTimeRangeExtractor()));
 
         usedSearchFiltersToQueryStringsMapper = mock(UsedSearchFiltersToQueryStringsMapper.class);
         doReturn(Collections.emptySet()).when(usedSearchFiltersToQueryStringsMapper).map(any());
         backend = new ElasticsearchBackend(handlers,
                 null,
-                mock(IndexLookup.class),
+                indexLookup,
                 ViewsUtils.createTestContextFactory(),
                 usedSearchFiltersToQueryStringsMapper,
                 new NoOpStatsCollector<>(),
@@ -143,6 +170,66 @@ public void generatedContextHasQueryThatIncludesSearchFilters() {
                 .contains("method:GET");
     }
 
+    @Test
+    public void testExplain() {
+        when(indexLookup.indexRangesForStreamsInTimeRange(anySet(), any())).thenAnswer(a -> {
+            if (a.getArgument(1, TimeRange.class).getFrom().getYear() < 2024) {
+                return Set.of(
+                        MongoIndexRange.create("graylog_0", nowUTC(), nowUTC(), nowUTC(), 0),
+                        MongoIndexRange.create("graylog_1", nowUTC(), nowUTC(), nowUTC(), 0),
+                        MongoIndexRange.create("graylog_warm_2", nowUTC(), nowUTC(), nowUTC(), 0)
+                );
+            }
+            return Set.of(MongoIndexRange.create("graylog_0", nowUTC(), nowUTC(), nowUTC(), 0));
+        });
+
+        final Query query = Query.builder()
+                .id("query1")
+                .query(ElasticsearchQueryString.of("needle"))
+                .searchTypes(Set.of(
+                                MessageList.builder()
+                                        .id("messagelist-1")
+                                        .build(),
+                                Pivot.builder()
+                                        .id("pivot-1")
+                                        .rowGroups(Time.builder().field("source").interval(AutoInterval.create()).build())
+                                        .timerange(AbsoluteRange.create(DateTime.parse("2016-05-19T00:00:00.000Z"), DateTime.parse("2022-01-09T00:00:00.000Z")))
+                                        .series()
+                                        .rollup(false)
+                                        .build()
+                        )
+                )
+                .timerange(RelativeRange.create(300))
+                .build();
+        final Search search = Search.builder().queries(ImmutableSet.of(query)).build();
+        final SearchJob job = new SearchJob("deadbeef", search, "admin", "test-node-id");
+        final GeneratedQueryContext generatedQueryContext = createContext(query);
+
+        var explainResult = backend.explain(job, query, generatedQueryContext);
+        assertThat(explainResult.searchTypes()).isNotNull();
+        assertThat(explainResult.searchTypes().get("messagelist-1")).satisfies(ml -> {
+            assertThat(ml).isNotNull();
+
+            assertThat(ml.searchedIndexRanges()).hasSize(1);
+            assertThat(ml.searchedIndexRanges()).allMatch(r -> r.indexName().equals("graylog_0"));
+
+
+            var ctx = JsonPath.parse(ml.queryString());
+            JsonPathAssert.assertThat(ctx).jsonPathAsString("$.query.bool.must[0].bool.filter[0].query_string.query").isEqualTo("needle");
+        });
+
+        assertThat(explainResult.searchTypes().get("pivot-1")).satisfies(ml -> {
+            assertThat(ml).isNotNull();
+            assertThat(ml.searchedIndexRanges()).hasSize(3);
+            assertThat(ml.searchedIndexRanges()).anyMatch(r -> r.indexName().equals("graylog_0") && !r.isWarmTiered());
+            assertThat(ml.searchedIndexRanges()).anyMatch(r -> r.indexName().equals("graylog_warm_2") && r.isWarmTiered());
+
+            var ctx = JsonPath.parse(ml.queryString());
+            JsonPathAssert.assertThat(ctx).jsonPathAsString("$.query.bool.must[0].bool.filter[0].query_string.query").isEqualTo("needle");
+            JsonPathAssert.assertThat(ctx).jsonPathAsString("$.aggregations.agg.date_histogram.field").isEqualTo("source");
+        });
+    }
+
     private ESGeneratedQueryContext createContext(Query query) {
         return backend.generate(query, Collections.emptySet(), DateTimeZone.UTC);
     }

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/views/OpenSearchBackend.java

@@ -16,11 +16,13 @@
  */
 package org.graylog.storage.opensearch2.views;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import io.opentelemetry.instrumentation.annotations.WithSpan;
 import jakarta.inject.Inject;
 import jakarta.inject.Named;
 import jakarta.inject.Provider;
+import org.graylog.plugins.views.search.ExplainResults;
 import org.graylog.plugins.views.search.Filter;
 import org.graylog.plugins.views.search.GlobalOverride;
 import org.graylog.plugins.views.search.Query;
@@ -170,7 +172,7 @@ public OSGeneratedQueryContext generate(Query query, Set<SearchError> validation
 
                     if (effectiveStreamIds.stream().noneMatch(s -> s.startsWith(Stream.DATASTREAM_PREFIX))) {
                         searchTypeOverrides
-                            .must(QueryBuilders.termsQuery(Message.FIELD_STREAMS, effectiveStreamIds));
+                                .must(QueryBuilders.termsQuery(Message.FIELD_STREAMS, effectiveStreamIds));
                     }
 
                     searchType.query().ifPresent(searchTypeQuery -> {
@@ -221,6 +223,26 @@ public Optional<QueryBuilder> generateFilterClause(Filter filter) {
         return Optional.empty();
     }
 
+    @Override
+    public ExplainResults.QueryExplainResult doExplain(SearchJob job, Query query, OSGeneratedQueryContext queryContext) {
+        final ImmutableMap.Builder<String, ExplainResults.ExplainResult> builder = ImmutableMap.builder();
+        final Map<String, SearchSourceBuilder> searchTypeQueries = queryContext.searchTypeQueries();
+
+        final DateTime nowUTCSharedBetweenSearchTypes = Tools.nowUTC();
+
+        query.searchTypes().forEach(s -> {
+            final Set<ExplainResults.IndexRangeResult> indicesForQuery = indexLookup.indexRangesForStreamsInTimeRange(
+                            query.effectiveStreams(s), query.effectiveTimeRange(s, nowUTCSharedBetweenSearchTypes))
+                    .stream().map(ExplainResults.IndexRangeResult::fromIndexRange).collect(Collectors.toSet());
+
+            final var queryString = searchTypeQueries.get(s.id()).toString();
+
+            builder.put(s.id(), new ExplainResults.ExplainResult(queryString, indicesForQuery));
+        });
+
+        return new ExplainResults.QueryExplainResult(builder.build());
+    }
+
     @Override
     @WithSpan
     public QueryResult doRun(SearchJob job, Query query, OSGeneratedQueryContext queryContext) {