Add stream categories to search replay (#20380)

* Add stream categories to search replay * Fix tests
Graylog2 · Sep 13, 2024 · f41bc23 · f41bc23
1 parent 1fd1e6d
commit f41bc23
Show file tree

Hide file tree

Showing 19 changed files with 130 additions and 32 deletions.
diff --git a/changelog/unreleased/pr-20110.toml b/changelog/unreleased/pr-20110.toml
@@ -2,4 +2,4 @@ type = "a"
 message = "Added categories to Streams to allow Illuminate content to be scoped to multiple products."
 
 issues = ["graylog-plugin-enterprise#7945"]
-pulls = ["20110", "20160", "20199", "20371", "20373", "graylog-plugin-enterprise#8295", "graylog-plugin-enterprise#8418", "graylog-plugin-enterprise#8437" ]
+pulls = ["20110", "20160", "20199", "20371", "20373", "20380", "graylog-plugin-enterprise#8295", "graylog-plugin-enterprise#8418", "graylog-plugin-enterprise#8437", "graylog-plugin-enterprise#8452", "graylog-plugin-enterprise#8454" ]
diff --git a/graylog2-server/src/main/java/org/graylog/plugins/views/search/Query.java b/graylog2-server/src/main/java/org/graylog/plugins/views/search/Query.java
@@ -32,6 +32,7 @@
 import org.graylog.plugins.views.search.engine.BackendQuery;
 import org.graylog.plugins.views.search.engine.EmptyTimeRange;
 import org.graylog.plugins.views.search.filter.AndFilter;
+import org.graylog.plugins.views.search.filter.OrFilter;
 import org.graylog.plugins.views.search.filter.StreamCategoryFilter;
 import org.graylog.plugins.views.search.filter.StreamFilter;
 import org.graylog.plugins.views.search.permissions.StreamPermissions;
@@ -52,6 +53,7 @@
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
@@ -251,6 +253,22 @@ public Query addStreamsToFilter(Set<String> streamIds) {
         return toBuilder().filter(newFilter).build();
     }
 
+    // This is a very specific call to set a flat OrFilter of a set of streamIds and categories when replaying a search.
+    // It is unlikely that it would need to be called outside the context of recreating a top-level search.
+    public Query orStreamAndStreamCategoryFilters(Set<String> streamIds, Set<String> categories) {
+        List<Filter> combinedFilters = new ArrayList<>();
+        if (streamIds != null && !streamIds.isEmpty()) {
+            combinedFilters.add(StreamFilter.anyIdOf(streamIds.toArray(new String[]{})));
+        }
+
+        if (categories != null && !categories.isEmpty()) {
+            categories.forEach(category -> {
+                combinedFilters.add(StreamCategoryFilter.ofCategory(category));
+            });
+        }
+        return toBuilder().filter(OrFilter.or(combinedFilters.toArray(new Filter[]{}))).build();
+    }
+
     public Query replaceStreamCategoryFilters(Function<Collection<String>, Stream<String>> categoryMappingFunction,
                                               StreamPermissions streamPermissions) {
         if (filter() == null) {

diff --git a/...rc/main/java/org/graylog/plugins/views/search/rest/scriptingapi/ScriptingApiResource.java b/...rc/main/java/org/graylog/plugins/views/search/rest/scriptingapi/ScriptingApiResource.java
@@ -119,6 +119,7 @@ public TabularResponse executeQuery(@ApiParam(name = "queryRequestSpec") @Valid
     @NoAuditEvent("Creating audit event manually in method body.")
     public TabularResponse executeQuery(@ApiParam(name = "query") @QueryParam("query") String query,
                                         @ApiParam(name = "streams") @QueryParam("streams") Set<String> streams,
+                                        @ApiParam(name = "stream_categories") @QueryParam("stream_categories") Set<String> streamCategories,
                                         @ApiParam(name = "timerange") @QueryParam("timerange") String timerangeKeyword,
                                         @ApiParam(name = "fields") @QueryParam("fields") List<String> fields,
                                         @ApiParam(name = "sort") @QueryParam("sort") String sort,
@@ -130,6 +131,7 @@ public TabularResponse executeQuery(@ApiParam(name = "query") @QueryParam("query
         try {
             MessagesRequestSpec messagesRequestSpec = queryParamsToFullRequestSpecificationMapper.simpleQueryParamsToFullRequestSpecification(query,
                     splitByComma(streams),
+                    splitByComma(streamCategories),
                     timerangeKeyword,
                     splitByComma(fields),
                     sort,
@@ -171,6 +173,7 @@ public TabularResponse executeQuery(@ApiParam(name = "searchRequestSpec") @Valid
     @NoAuditEvent("Creating audit event manually in method body.")
     public TabularResponse executeQuery(@ApiParam(name = "query") @QueryParam("query") String query,
                                         @ApiParam(name = "streams") @QueryParam("streams") Set<String> streams,
+                                        @ApiParam(name = "stream_categories") @QueryParam("stream_categories") Set<String> streamCategories,
                                         @ApiParam(name = "timerange") @QueryParam("timerange") String timerangeKeyword,
                                         @ApiParam(name = "groups") @QueryParam("groups") List<String> groups,
                                         @ApiParam(name = "metrics") @QueryParam("metrics") List<String> metrics,
@@ -179,6 +182,7 @@ public TabularResponse executeQuery(@ApiParam(name = "query") @QueryParam("query
             AggregationRequestSpec aggregationRequestSpec = queryParamsToFullRequestSpecificationMapper.simpleQueryParamsToFullRequestSpecification(
                     query,
                     StringUtils.splitByComma(streams),
+                    StringUtils.splitByComma(streamCategories),
                     timerangeKeyword,
                     splitByComma(groups),
                     splitByComma(metrics)

diff --git a/...s/views/search/rest/scriptingapi/mapping/QueryParamsToFullRequestSpecificationMapper.java b/...s/views/search/rest/scriptingapi/mapping/QueryParamsToFullRequestSpecificationMapper.java
@@ -16,6 +16,7 @@
  */
 package org.graylog.plugins.views.search.rest.scriptingapi.mapping;
 
+import jakarta.inject.Inject;
 import org.apache.commons.collections4.CollectionUtils;
 import org.graylog.plugins.views.search.rest.scriptingapi.parsing.TimerangeParser;
 import org.graylog.plugins.views.search.rest.scriptingapi.request.AggregationRequestSpec;
@@ -24,8 +25,6 @@
 import org.graylog.plugins.views.search.rest.scriptingapi.request.Metric;
 import org.graylog.plugins.views.search.searchtypes.pivot.SortSpec;
 
-import jakarta.inject.Inject;
-
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
@@ -42,6 +41,7 @@ public QueryParamsToFullRequestSpecificationMapper(final TimerangeParser timeran
 
     public MessagesRequestSpec simpleQueryParamsToFullRequestSpecification(final String query,
                                                                            final Set<String> streams,
+                                                                           final Set<String> streamCategories,
                                                                            final String timerangeKeyword,
                                                                            final List<String> fields,
                                                                            final String sort,
@@ -51,6 +51,7 @@ public MessagesRequestSpec simpleQueryParamsToFullRequestSpecification(final Str
 
         return new MessagesRequestSpec(query,
                 streams,
+                streamCategories,
                 timerangeParser.parseTimeRange(timerangeKeyword),
                 sort,
                 sortOrder,
@@ -61,6 +62,7 @@ public MessagesRequestSpec simpleQueryParamsToFullRequestSpecification(final Str
 
     public AggregationRequestSpec simpleQueryParamsToFullRequestSpecification(final String query,
                                                                               final Set<String> streams,
+                                                                              final Set<String> streamCategories,
                                                                               final String timerangeKeyword,
                                                                               List<String> groups,
                                                                               List<String> metrics) {
@@ -80,6 +82,7 @@ public AggregationRequestSpec simpleQueryParamsToFullRequestSpecification(final
         return new AggregationRequestSpec(
                 query,
                 streams,
+                streamCategories,
                 timerangeParser.parseTimeRange(timerangeKeyword),
                 groups.stream().map(Grouping::new).collect(Collectors.toList()),
                 metrics.stream().map(Metric::fromStringRepresentation).flatMap(Optional::stream).collect(Collectors.toList())

diff --git a/...aylog/plugins/views/search/rest/scriptingapi/mapping/SearchRequestSpecToSearchMapper.java b/...aylog/plugins/views/search/rest/scriptingapi/mapping/SearchRequestSpecToSearchMapper.java
@@ -17,6 +17,7 @@
 package org.graylog.plugins.views.search.rest.scriptingapi.mapping;
 
 import com.google.common.collect.ImmutableSet;
+import jakarta.inject.Inject;
 import org.graylog.plugins.views.search.Query;
 import org.graylog.plugins.views.search.Search;
 import org.graylog.plugins.views.search.SearchType;
@@ -27,25 +28,29 @@
 import org.graylog.plugins.views.search.rest.scriptingapi.request.SearchRequestSpec;
 import org.graylog2.plugin.indexer.searches.timeranges.RelativeRange;
 import org.graylog2.plugin.indexer.searches.timeranges.TimeRange;
+import org.graylog2.streams.StreamService;
 
-import jakarta.inject.Inject;
-
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.stream.Stream;
 
 public class SearchRequestSpecToSearchMapper {
 
     public static final String QUERY_ID = "scripting_api_temporary_query";
 
     private final AggregationSpecToPivotMapper pivotCreator;
     private final MessagesSpecToMessageListMapper messageListCreator;
+    private final Function<Collection<String>, Stream<String>> streamCategoryMapper;
 
     @Inject
     public SearchRequestSpecToSearchMapper(final AggregationSpecToPivotMapper pivotCreator,
-                                           final MessagesSpecToMessageListMapper messageListCreator) {
+                                           final MessagesSpecToMessageListMapper messageListCreator,
+                                           StreamService streamService) {
         this.pivotCreator = pivotCreator;
         this.messageListCreator = messageListCreator;
+        this.streamCategoryMapper = (categories) -> streamService.mapCategoriesToIds(categories).stream();
     }
 
     public Search mapToSearch(MessagesRequestSpec messagesRequestSpec, SearchUser searchUser) {
@@ -65,13 +70,17 @@ private <T extends SearchRequestSpec> Search mapToSearch(final T searchRequestSp
                 .timerange(getTimerange(searchRequestSpec))
                 .build();
 
-        if (!searchRequestSpec.streams().isEmpty()) {
-            query = query.addStreamsToFilter(new HashSet<>(searchRequestSpec.streams()));
+        if (!(searchRequestSpec.streams().isEmpty() && searchRequestSpec.streamCategories().isEmpty())) {
+            query = query.orStreamAndStreamCategoryFilters(
+                    new HashSet<>(searchRequestSpec.streams()),
+                    new HashSet<>(searchRequestSpec.streamCategories())
+            );
         }
 
         return Search.builder()
                 .queries(ImmutableSet.of(query))
                 .build()
+                .addStreamsToQueriesWithCategories(streamCategoryMapper, searchUser)
                 .addStreamsToQueriesWithoutStreams(() -> searchUser.streams().readableOrAllIfEmpty(searchRequestSpec.streams()));
     }
 

diff --git a/...va/org/graylog/plugins/views/search/rest/scriptingapi/request/AggregationRequestSpec.java b/...va/org/graylog/plugins/views/search/rest/scriptingapi/request/AggregationRequestSpec.java
@@ -19,11 +19,10 @@
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Strings;
-import org.graylog.plugins.views.search.rest.scriptingapi.response.ResponseSchemaEntry;
-import org.graylog2.plugin.indexer.searches.timeranges.TimeRange;
-
 import jakarta.validation.Valid;
 import jakarta.validation.constraints.NotEmpty;
+import org.graylog.plugins.views.search.rest.scriptingapi.response.ResponseSchemaEntry;
+import org.graylog2.plugin.indexer.searches.timeranges.TimeRange;
 
 import java.util.List;
 import java.util.Set;
@@ -32,6 +31,7 @@
 
 public record AggregationRequestSpec(@JsonProperty("query") String queryString,
                                      @JsonProperty("streams") Set<String> streams,
+                                     @JsonProperty("stream_categories") Set<String> streamCategories,
                                      @JsonProperty("timerange") TimeRange timerange,
                                      @JsonProperty("group_by") @Valid @NotEmpty List<Grouping> groupings,
                                      @JsonProperty("metrics") @Valid @NotEmpty List<Metric> metrics) implements SearchRequestSpec {
@@ -47,6 +47,9 @@ public record AggregationRequestSpec(@JsonProperty("query") String queryString,
         if (streams == null) {
             streams = Set.of();
         }
+        if (streamCategories == null) {
+            streamCategories = Set.of();
+        }
     }
 
     public boolean hasCustomSort() {

diff --git a/.../java/org/graylog/plugins/views/search/rest/scriptingapi/request/MessagesRequestSpec.java b/.../java/org/graylog/plugins/views/search/rest/scriptingapi/request/MessagesRequestSpec.java
@@ -28,6 +28,7 @@
 
 public record MessagesRequestSpec(@JsonProperty("query") String queryString,
                                   @JsonProperty("streams") Set<String> streams,
+                                  @JsonProperty("stream_categories") Set<String> streamCategories,
                                   @JsonProperty("timerange") TimeRange timerange,
                                   @JsonProperty("sort") String sort,
                                   @JsonProperty("sort_order") SortSpec.Direction sortOrder,
@@ -58,6 +59,9 @@ public record MessagesRequestSpec(@JsonProperty("query") String queryString,
         if (streams == null) {
             streams = Set.of();
         }
+        if (streamCategories == null) {
+            streamCategories = Set.of();
+        }
         if (from < 0) {
             from = DEFAULT_FROM;
         }

diff --git a/...in/java/org/graylog/plugins/views/search/rest/scriptingapi/request/SearchRequestSpec.java b/...in/java/org/graylog/plugins/views/search/rest/scriptingapi/request/SearchRequestSpec.java
@@ -31,6 +31,8 @@ public interface SearchRequestSpec {
 
     Set<String> streams();
 
+    Set<String> streamCategories();
+
     TimeRange timerange();
 
 }
diff --git a/...ews/search/rest/scriptingapi/mapping/QueryParamsToFullRequestSpecificationMapperTest.java b/...ews/search/rest/scriptingapi/mapping/QueryParamsToFullRequestSpecificationMapperTest.java
@@ -56,6 +56,7 @@ void setUp() {
     @Test
     void throwsExceptionOnNullGroups() {
         assertThrows(IllegalArgumentException.class, () -> toTest.simpleQueryParamsToFullRequestSpecification("*",
+                Set.of(),
                 Set.of(),
                 "42d",
                 null,
@@ -65,6 +66,7 @@ void throwsExceptionOnNullGroups() {
     @Test
     void throwsExceptionOnEmptyGroups() {
         assertThrows(IllegalArgumentException.class, () -> toTest.simpleQueryParamsToFullRequestSpecification("*",
+                Set.of(),
                 Set.of(),
                 "42d",
                 List.of(),
@@ -74,6 +76,7 @@ void throwsExceptionOnEmptyGroups() {
     @Test
     void throwsExceptionOnWrongMetricFormat() {
         assertThrows(IllegalArgumentException.class, () -> toTest.simpleQueryParamsToFullRequestSpecification("*",
+                Set.of(),
                 Set.of(),
                 "42d",
                 List.of("http_method"),
@@ -83,6 +86,7 @@ void throwsExceptionOnWrongMetricFormat() {
     @Test
     void usesProperDefaults() {
         AggregationRequestSpec aggregationRequestSpec = toTest.simpleQueryParamsToFullRequestSpecification(null,
+                null,
                 null,
                 null,
                 List.of("http_method"),
@@ -91,13 +95,15 @@ void usesProperDefaults() {
         assertThat(aggregationRequestSpec).isEqualTo(new AggregationRequestSpec(
                         "*",
                         Set.of(),
+                Set.of(),
                         DEFAULT_TIMERANGE,
                         List.of(new Grouping("http_method")),
                         List.of(new Metric("count", null))
                 )
         );
 
         aggregationRequestSpec = toTest.simpleQueryParamsToFullRequestSpecification(null,
+                null,
                 null,
                 null,
                 List.of("http_method"),
@@ -106,17 +112,19 @@ void usesProperDefaults() {
         assertThat(aggregationRequestSpec).isEqualTo(new AggregationRequestSpec(
                         "*",
                         Set.of(),
+                Set.of(),
                         DEFAULT_TIMERANGE,
                         List.of(new Grouping("http_method")),
                         List.of(new Metric("count", null))
 
                 )
         );
 
-        final MessagesRequestSpec messagesRequestSpec = toTest.simpleQueryParamsToFullRequestSpecification(null, null, null, null, null, null, 0, 10);
+        final MessagesRequestSpec messagesRequestSpec = toTest.simpleQueryParamsToFullRequestSpecification(null, null, null, null, null, null, null, 0, 10);
         assertThat(messagesRequestSpec).isEqualTo(new MessagesRequestSpec(
                         "*",
                         Set.of(),
+                Set.of(),
                         DEFAULT_TIMERANGE,
                         DEFAULT_SORT,
                         DEFAULT_SORT_ORDER,
@@ -134,13 +142,15 @@ void createsProperRequestSpec() {
         doReturn(KeywordRange.create("last 1 day", "UTC")).when(timerangeParser).parseTimeRange("1d");
         final AggregationRequestSpec aggregationRequestSpec = toTest.simpleQueryParamsToFullRequestSpecification("http_method:GET",
                 Set.of("000000000000000000000001"),
+                Set.of("category1"),
                 "1d",
                 List.of("http_method", "controller"),
                 List.of("avg:took_ms"));
 
         assertThat(aggregationRequestSpec).isEqualTo(new AggregationRequestSpec(
                         "http_method:GET",
                         Set.of("000000000000000000000001"),
+                Set.of("category1"),
                         KeywordRange.create("last 1 day", "UTC"),
                         List.of(new Grouping("http_method"), new Grouping("controller")),
                         List.of(new Metric("avg", "took_ms"))

diff --git a/graylog2-web-interface/src/components/search/SurroundingSearchButton.tsx b/graylog2-web-interface/src/components/search/SurroundingSearchButton.tsx
@@ -44,22 +44,23 @@ const buildFilterFields = (messageFields: {
 
 const buildSearchLink = (id: string, from: string, to: string, filterFields: {
   [key: string]: unknown;
-}, streams: string[]) => SearchLink.builder()
+}, streams: string[], streamCategories: string[]) => SearchLink.builder()
   .timerange({ type: 'absolute', from, to })
   .streams(streams)
+  .streamCategories(streamCategories)
   .filterFields(filterFields)
   .highlightedMessage(id)
   .build()
   .toURL();
 
 const searchLink = (range: string, timestamp: moment.MomentInput, id: string, messageFields: {
   [key: string]: unknown;
-}, searchConfig: Pick<SearchesConfig, 'surrounding_filter_fields'>, streams: string[]) => {
+}, searchConfig: Pick<SearchesConfig, 'surrounding_filter_fields'>, streams: string[], streamCategories: string[]) => {
   const fromTime = moment(timestamp).subtract(Number(range), 'seconds').toISOString();
   const toTime = moment(timestamp).add(Number(range), 'seconds').toISOString();
   const filterFields = buildFilterFields(messageFields, searchConfig);
 
-  return buildSearchLink(id, fromTime, toTime, filterFields, streams);
+  return buildSearchLink(id, fromTime, toTime, filterFields, streams, streamCategories);
 };
 
 type Props = {
@@ -70,7 +71,7 @@ type Props = {
 };
 
 const SurroundingSearchButton = ({ searchConfig, timestamp, id, messageFields }: Props) => {
-  const { streams } = useContext(DrilldownContext);
+  const { streams, streamCategories } = useContext(DrilldownContext);
   const timeRangeOptions = buildTimeRangeOptions(searchConfig);
   const location = useLocation();
   const sendTelemetry = useSendTelemetry();
@@ -91,7 +92,7 @@ const SurroundingSearchButton = ({ searchConfig, timestamp, id, messageFields }:
     .map((range) => (
       <MenuItem key={range}
                 onClick={() => sendEvent(range)}
-                href={searchLink(range, timestamp, id, messageFields, searchConfig, streams)}
+                href={searchLink(range, timestamp, id, messageFields, searchConfig, streams, streamCategories)}
                 target="_blank"
                 rel="noopener noreferrer">
         {timeRangeOptions[range]}