Adds stream categories to dashboard widgets

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191125144500_MigrateDashboardsToViewsSupport/AggregationWidget.java

@@ -40,6 +40,7 @@ public abstract class AggregationWidget implements ViewWidget {
     private static final String FIELD_TIMERANGE = "timerange";
     private static final String FIELD_QUERY = "query";
     private static final String FIELD_STREAMS = "streams";
+    private static final String FIELD_STREAM_CATEGORIES = "stream_categories";
 
     @JsonProperty(FIELD_ID)
     public abstract String id();
@@ -60,13 +61,17 @@ public abstract class AggregationWidget implements ViewWidget {
     @JsonProperty(FIELD_STREAMS)
     abstract Set<String> streams();
 
+    @JsonProperty(FIELD_STREAM_CATEGORIES)
+    abstract Set<String> streamCategories();
+
     @JsonProperty(FIELD_CONFIG)
     public abstract AggregationConfig config();
 
     public static Builder builder() {
         return new AutoValue_AggregationWidget.Builder()
                 .type(TYPE_AGGREGATION)
-                .streams(Collections.emptySet());
+                .streams(Collections.emptySet())
+                .streamCategories(Collections.emptySet());
     }
 
     public Set<SearchType> toSearchTypes(RandomUUIDProvider randomUUIDProvider) {
@@ -75,6 +80,7 @@ public Set<SearchType> toSearchTypes(RandomUUIDProvider randomUUIDProvider) {
                 .name("chart")
                 .query(query())
                 .streams(streams())
+                .streamCategories(streamCategories())
                 .timerange(timerange())
                 .rollup(config().rollup())
                 .rowGroups(config().rowPivots().stream().map(pivot -> pivot.toBucketSpec()).collect(Collectors.toList()))
@@ -122,6 +128,9 @@ public Builder query(String query) {
         @JsonProperty(FIELD_STREAMS)
         public abstract Builder streams(Set<String> streams);
 
+        @JsonProperty(FIELD_STREAM_CATEGORIES)
+        public abstract Builder streamCategories(Set<String> streamCategories);
+
         @JsonProperty(FIELD_CONFIG)
         public abstract Builder config(AggregationConfig config);
 

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191125144500_MigrateDashboardsToViewsSupport/MessageList.java

@@ -73,6 +73,8 @@ Builder query(String query) {
 
         abstract Builder streams(Set<String> streams);
 
+        abstract Builder streamCategories(Set<String> streamCategories);
+
         abstract Builder limit(int limit);
 
         abstract Builder offset(int offset);

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191125144500_MigrateDashboardsToViewsSupport/Pivot.java

@@ -92,6 +92,8 @@ Builder query(String query) {
 
         abstract Builder streams(Set<String> streams);
 
+        abstract Builder streamCategories(Set<String> streamCategories);
+
         abstract Pivot build();
     }
 }

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191125144500_MigrateDashboardsToViewsSupport/SearchType.java

@@ -34,6 +34,9 @@ public interface SearchType {
     @JsonProperty
     Set<String> streams();
 
+    @JsonProperty
+    Set<String> streamCategories();
+
     @JsonProperty
     Optional<String> name();
 }

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191203120602_MigrateSavedSearchesToViewsSupport/search/MessagesSearchType.java

@@ -22,6 +22,7 @@
 import org.graylog.plugins.views.migrations.V20191203120602_MigrateSavedSearchesToViewsSupport.view.ElasticsearchQueryString;
 import org.graylog.plugins.views.migrations.V20191203120602_MigrateSavedSearchesToViewsSupport.view.TimeRange;
 
+import javax.annotation.Nullable;
 import java.util.Collections;
 import java.util.Optional;
 import java.util.Set;
@@ -35,6 +36,7 @@ public abstract class MessagesSearchType implements SearchType {
     private static final String FIELD_TIMERANGE = "timerange";
     private static final String FIELD_QUERY = "query";
     private static final String FIELD_STREAMS = "streams";
+    private static final String FIELD_STREAM_CATEGORIES = "stream_categories";
     private static final String FIELD_NAME = "name";
     private static final String FIELD_LIMIT = "limit";
     private static final String FIELD_OFFSET = "offset";
@@ -62,6 +64,12 @@ public Set<String> streams() {
         return Collections.emptySet();
     }
 
+    @Nullable
+    @JsonProperty(FIELD_STREAM_CATEGORIES)
+    public Set<String> streamCategories() {
+        return Collections.emptySet();
+    }
+
     @JsonProperty(FIELD_NAME)
     public Optional<String> name() {
         return Optional.empty();

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191203120602_MigrateSavedSearchesToViewsSupport/search/Pivot.java

@@ -56,6 +56,11 @@ public Set<String> streams() {
         return Collections.emptySet();
     }
 
+    @Override
+    public Set<String> streamCategories() {
+        return Collections.emptySet();
+    }
+
     @JsonProperty("row_groups")
     abstract List<Time> rowGroups();
 

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191203120602_MigrateSavedSearchesToViewsSupport/search/SearchType.java

@@ -35,4 +35,7 @@ public interface SearchType {
 
     @JsonProperty
     Set<String> streams();
+
+    @JsonProperty
+    Set<String> streamCategories();
 }

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191203120602_MigrateSavedSearchesToViewsSupport/view/AggregationWidget.java

@@ -42,6 +42,7 @@ public abstract class AggregationWidget implements ViewWidget {
     private static final String FIELD_TIMERANGE = "timerange";
     private static final String FIELD_QUERY = "query";
     private static final String FIELD_STREAMS = "streams";
+    private static final String FIELD_STREAM_CATEGORIES = "stream_categories";
 
     @JsonProperty(FIELD_ID)
     public abstract String id();
@@ -74,6 +75,11 @@ Set<String> streams() {
         return Collections.emptySet();
     }
 
+    @JsonProperty(FIELD_STREAM_CATEGORIES)
+    Set<String> streamCategories() {
+        return Collections.emptySet();
+    }
+
     @JsonProperty(FIELD_CONFIG)
     public abstract AggregationConfig config();
 

graylog2-server/src/main/java/org/graylog/plugins/views/migrations/V20191203120602_MigrateSavedSearchesToViewsSupport/view/MessagesWidget.java

@@ -36,6 +36,7 @@ public abstract class MessagesWidget implements ViewWidget {
     private static final String FIELD_TIMERANGE = "timerange";
     private static final String FIELD_QUERY = "query";
     private static final String FIELD_STREAMS = "streams";
+    private static final String FIELD_STREAM_CATEGORIES = "stream_categories";
     private static final String FIELD_CONFIG = "config";
 
     @JsonProperty(FIELD_ID)
@@ -63,6 +64,11 @@ Set<String> streams() {
         return Collections.emptySet();
     }
 
+    @JsonProperty(FIELD_STREAM_CATEGORIES)
+    Set<String> streamCategories() {
+        return Collections.emptySet();
+    }
+
     @JsonProperty(FIELD_CONFIG)
     abstract MessagesWidgetConfig config();
 

graylog2-server/src/main/java/org/graylog/plugins/views/search/Search.java

@@ -172,6 +172,39 @@ public Search addStreamsToQueriesWithCategories(Function<Collection<String>, Str
         return toBuilder().queries(newQueries).build();
     }
 
+    public Search addStreamsToSearchTypesWithCategories(Function<Collection<String>, Stream<String>> categoryMappingFunction,
+                                                        StreamPermissions streamPermissions) {
+        if (!hasQuerySearchTypesWithStreamCategories()) {
+            return this;
+        }
+        final Set<Query> withStreamCategories = queries().stream()
+                .filter(q -> q.searchTypes().stream()
+                        .anyMatch(SearchType::hasStreamCategories))
+                .collect(toSet());
+        final Set<Query> withoutStreamCategories = Sets.difference(queries(), withStreamCategories);
+        final Set<Query> withMappedStreamCategories = new HashSet<>();
+
+        for (Query query : withStreamCategories) {
+            final Set<SearchType> mappedSearchTypes = new HashSet<>();
+            for (SearchType st : query.searchTypes()) {
+                if (!st.hasStreamCategories()) {
+                    mappedSearchTypes.add(st);
+                } else {
+                    final Set<String> mappedStreamIds = categoryMappingFunction.apply(st.streamCategories())
+                            .filter(streamPermissions::canReadStream)
+                            .collect(toSet());
+                    mappedStreamIds.addAll(st.streams());
+                    mappedSearchTypes.add(st.toBuilder().streams(mappedStreamIds).build());
+                }
+            }
+            withMappedStreamCategories.add(query.toBuilder().searchTypes(mappedSearchTypes).build());
+        }
+
+        final ImmutableSet<Query> newQueries = Sets.union(withMappedStreamCategories, withoutStreamCategories).immutableCopy();
+
+        return toBuilder().queries(newQueries).build();
+    }
+
     private boolean hasQueriesWithoutStreams() {
         return !queries().stream().allMatch(Query::hasStreams);
     }
@@ -180,6 +213,12 @@ private boolean hasQueriesWithStreamCategories() {
         return queries().stream().anyMatch(q -> !q.usedStreamCategories().isEmpty());
     }
 
+    private boolean hasQuerySearchTypesWithStreamCategories() {
+        return queries().stream()
+                .flatMap(q -> q.searchTypes().stream())
+                .anyMatch(SearchType::hasStreamCategories);
+    }
+
     public abstract Builder toBuilder();
 
     public static Builder builder() {

graylog2-server/src/main/java/org/graylog/plugins/views/search/SearchType.java

@@ -85,6 +85,10 @@ public interface SearchType extends ContentPackable<SearchTypeEntity>, Exportabl
     @JsonProperty
     Set<String> streams();
 
+    @Nullable
+    @JsonProperty
+    Set<String> streamCategories();
+
     SearchType applyExecutionContext(SearchTypeExecutionState executionState);
 
     default SearchType withQuery(BackendQuery query) {
@@ -110,6 +114,10 @@ default SearchType withReferenceDate(DateTime now) {
                 .orElse(this);
     }
 
+    default boolean hasStreamCategories() {
+        return streamCategories() != null && !streamCategories().isEmpty();
+    }
+
     SearchTypeBuilder toBuilder();
 
     /**
@@ -164,6 +172,10 @@ class Fallback implements SearchType {
         @JsonProperty
         private Set<String> streams;
 
+        @Nullable
+        @JsonProperty
+        private Set<String> streamCategories;
+
         @Override
         public String type() {
             return type;
@@ -204,6 +216,11 @@ public Set<String> streams() {
             return this.streams == null ? Collections.emptySet() : this.streams;
         }
 
+        @Override
+        public Set<String> streamCategories() {
+            return this.streamCategories == null ? Collections.emptySet() : this.streamCategories;
+        }
+
         @Override
         public SearchType applyExecutionContext(SearchTypeExecutionState state) {
             return this;
@@ -284,6 +301,11 @@ public SearchTypeBuilder query(@Nullable BackendQuery query) {
                 public SearchTypeBuilder filter(@Nullable Filter filter) {
                     return this;
                 }
+
+                @Override
+                public SearchTypeBuilder streams(Set<String> streams) {
+                    return this;
+                }
             };
         }
     }

graylog2-server/src/main/java/org/graylog/plugins/views/search/SearchTypeBuilder.java

@@ -22,6 +22,7 @@
 
 import javax.annotation.Nullable;
 import java.util.List;
+import java.util.Set;
 
 public interface SearchTypeBuilder {
     SearchType build();
@@ -33,4 +34,6 @@ public interface SearchTypeBuilder {
     SearchTypeBuilder query(@Nullable BackendQuery query);
 
     SearchTypeBuilder filter(@Nullable Filter filter);
+
+    SearchTypeBuilder streams(Set<String> streams);
 }

graylog2-server/src/main/java/org/graylog/plugins/views/search/engine/normalization/PluggableSearchNormalization.java

@@ -76,8 +76,9 @@ private Query normalize(final Query query,
     @Override
     public Search preValidation(Search search, SearchUser searchUser, ExecutionState executionState) {
         final Search searchWithStreams = search
-                .addStreamsToQueriesWithoutStreams(() -> searchUser.streams().loadMessageStreamsWithFallback())
-                .addStreamsToQueriesWithCategories(streamCategoryMapper, searchUser);
+                .addStreamsToQueriesWithCategories(streamCategoryMapper, searchUser)
+                .addStreamsToSearchTypesWithCategories(streamCategoryMapper, searchUser)
+                .addStreamsToQueriesWithoutStreams(() -> searchUser.streams().loadMessageStreamsWithFallback());
         final var now = referenceDateFromOverrideOrNow(executionState);
         final var normalizedSearch = searchWithStreams.applyExecutionState(firstNonNull(executionState, ExecutionState.empty()))
                 .withReferenceDate(now);

graylog2-server/src/main/java/org/graylog/plugins/views/search/permissions/UserStreams.java

@@ -58,6 +58,10 @@ public ImmutableSet<String> loadMessageStreamsWithFallback() {
         }
     }
 
+    public ImmutableSet<String> loadStreamsWithCategories(Set<String> categories) {
+        return permittedStreams.loadWithCategories(categories, streamPermissions);
+    }
+
     /**
      * If any stream IDs are provided, they will be filtered out by read permission. If none are given, we'll load
      * all available streams for the current SearchUser

graylog2-server/src/main/java/org/graylog/plugins/views/search/rest/QueryValidationResource.java

@@ -43,6 +43,7 @@
 import org.graylog2.shared.rest.resources.RestResource;
 
 import java.util.Comparator;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -90,10 +91,19 @@ private boolean containsWarmIndices(Set<ExplainResults.IndexRangeResult> searche
     }
 
     private ValidationRequest prepareRequest(ValidationRequestDTO validationRequest, SearchUser searchUser) {
+        // Combine stream IDs mapped from streamCategories and stream IDs from the validationRequest before calling
+        // readableOrAllIfEmpty to ensure all possible requested streams are added first.
+        final Set<String> streamsAndMappedCategories = validationRequest.streams().orElse(new HashSet<>());
+        if (validationRequest.streamCategories().isPresent()) {
+            streamsAndMappedCategories.addAll(searchUser.streams().loadStreamsWithCategories(validationRequest.streamCategories().get()));
+        }
+        ImmutableSet.Builder<String> allRequestedStreams = ImmutableSet.<String>builder()
+                .addAll(searchUser.streams().readableOrAllIfEmpty(streamsAndMappedCategories));
+
         final ValidationRequest.Builder q = ValidationRequest.Builder.builder()
                 .query(validationRequest.query())
                 .timerange(validationRequest.timerange().orElse(defaultTimeRange()))
-                .streams(searchUser.streams().readableOrAllIfEmpty(validationRequest.streams()))
+                .streams(allRequestedStreams.build())
                 .parameters(resolveParameters(validationRequest))
                 .validationMode(validationRequest.validationMode().toInternalRepresentation());
 

graylog2-server/src/main/java/org/graylog/plugins/views/search/rest/ValidationRequestDTO.java

@@ -38,6 +38,7 @@
 public abstract class ValidationRequestDTO {
 
     private static final String FIELD_STREAMS = "streams";
+    private static final String FIELD_STREAM_CATEGORIES = "stream_categories";
     private static final String FIELD_TIMERANGE = "timerange";
 
     @JsonProperty
@@ -58,6 +59,9 @@ public abstract class ValidationRequestDTO {
     @JsonProperty(FIELD_STREAMS)
     public abstract Optional<Set<String>> streams();
 
+    @JsonProperty(FIELD_STREAM_CATEGORIES)
+    public abstract Optional<Set<String>> streamCategories();
+
     @JsonProperty
     public abstract ImmutableSet<Parameter> parameters();
 
@@ -79,6 +83,9 @@ public abstract static class Builder {
         @JsonProperty(FIELD_STREAMS)
         public abstract Builder streams(@Nullable Set<String> streams);
 
+        @JsonProperty(FIELD_STREAM_CATEGORIES)
+        public abstract Builder streamCategories(@Nullable Set<String> streamCategories);
+
         @JsonProperty(FIELD_TIMERANGE)
         public abstract Builder timerange(@Nullable TimeRange timerange);
 

graylog2-server/src/main/java/org/graylog/plugins/views/search/searchtypes/MessageList.java

@@ -105,6 +105,7 @@ public static Builder builder() {
                 .offset(0)
                 .filters(Collections.emptyList())
                 .streams(Collections.emptySet())
+                .streamCategories(Collections.emptySet())
                 .decorators(Collections.emptyList())
                 .fields(Collections.emptyList());
     }
@@ -129,6 +130,7 @@ public abstract static class Builder implements SearchTypeBuilder {
         public static Builder createDefault() {
             return builder()
                     .filters(Collections.emptyList())
+                    .streamCategories(Collections.emptySet())
                     .streams(Collections.emptySet());
         }
 
@@ -164,6 +166,9 @@ public Builder timerange(@Nullable TimeRange timerange) {
         @JsonProperty
         public abstract Builder streams(Set<String> streams);
 
+        @JsonProperty
+        public abstract Builder streamCategories(Set<String> streamCategories);
+
         @JsonProperty
         public abstract Builder limit(int limit);