Add handling for the rest of the entity types that reference streams

Graylog2 · Nov 1, 2023 · ad8564f · ad8564f
1 parent 78adef9
commit ad8564f
Show file tree

Hide file tree

Showing 11 changed files with 44 additions and 34 deletions.
diff --git a/...n/java/org/graylog/events/contentpack/entities/AggregationEventProcessorConfigEntity.java b/...n/java/org/graylog/events/contentpack/entities/AggregationEventProcessorConfigEntity.java
@@ -41,6 +41,7 @@
 import java.util.stream.Collectors;
 
 import static org.graylog2.contentpacks.facades.StreamReferenceFacade.resolveStreamEntity;
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.resolveStreamEntityObject;
 
 @AutoValue
 @JsonTypeName(AggregationEventProcessorConfigEntity.TYPE_NAME)
@@ -129,7 +130,7 @@ public EventProcessorConfig toNativeEntity(Map<String, ValueReference> parameter
                                                Map<EntityDescriptor, Object> nativeEntities) {
         final ImmutableSet<String> streamSet = ImmutableSet.copyOf(
                 streams().stream()
-                        .map(id -> resolveStreamEntity(id, nativeEntities))
+                        .map(id -> resolveStreamEntityObject(id, nativeEntities))
                         .map(object -> {
                             if (object == null) {
                                 throw new ContentPackException("Missing Stream for event definition");
@@ -161,7 +162,7 @@ public void resolveForInstallation(EntityV1 entity,
                                        Map<EntityDescriptor, Entity> entities,
                                        MutableGraph<Entity> graph) {
         streams().stream()
-                .map(id -> (Entity) resolveStreamEntity(id, entities))
+                .map(id -> resolveStreamEntity(id, entities))
                 .filter(Objects::nonNull)
                 .forEach(stream -> graph.putEdge(entity, stream));
     }

diff --git a/...c/main/java/org/graylog/events/processor/aggregation/AggregationEventProcessorConfig.java b/...c/main/java/org/graylog/events/processor/aggregation/AggregationEventProcessorConfig.java
@@ -53,7 +53,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
-import static org.graylog2.contentpacks.facades.StreamReferenceFacade.getStreamDescriptor;
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.getStreamEntityId;
 import static org.graylog2.shared.utilities.StringUtils.f;
 
 @AutoValue
@@ -248,7 +248,7 @@ private void checkEventLimitGreaterZero(ValidationResult validationResult) {
     @Override
     public EventProcessorConfigEntity toContentPackEntity(EntityDescriptorIds entityDescriptorIds) {
         final ImmutableSet<String> streamRefs = ImmutableSet.copyOf(streams().stream()
-                .map(streamId -> getStreamDescriptor(streamId, entityDescriptorIds))
+                .map(streamId -> getStreamEntityId(streamId, entityDescriptorIds))
                 .filter(Optional::isPresent)
                 .map(Optional::get)
                 .collect(Collectors.toSet()));

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
@@ -40,7 +40,6 @@
 import org.graylog.plugins.views.search.searchfilters.model.UsesSearchFilters;
 import org.graylog2.contentpacks.ContentPackable;
 import org.graylog2.contentpacks.EntityDescriptorIds;
-import org.graylog2.contentpacks.model.ModelTypes;
 import org.graylog2.contentpacks.model.entities.QueryEntity;
 import org.graylog2.plugin.indexer.searches.timeranges.InvalidRangeParametersException;
 import org.graylog2.plugin.indexer.searches.timeranges.RelativeRange;
@@ -60,6 +59,7 @@
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.collect.ImmutableSortedSet.of;
 import static java.util.stream.Collectors.toSet;
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.getStreamEntityIdOrThrow;
 
 @AutoValue
 @JsonAutoDetect
@@ -290,8 +290,7 @@ private Filter shallowMappedFilter(EntityDescriptorIds entityDescriptorIds) {
                             .map(filter -> {
                                 if (filter.type().equals(StreamFilter.NAME)) {
                                     final StreamFilter streamFilter = (StreamFilter) filter;
-                                    final String streamId = entityDescriptorIds.
-                                            getOrThrow(streamFilter.streamId(), ModelTypes.STREAM_V1);
+                                    final String streamId = getStreamEntityIdOrThrow(streamFilter.streamId(), entityDescriptorIds);
                                     return streamFilter.toBuilder().streamId(streamId).build();
                                 }
                                 return filter;

diff --git a/graylog2-server/src/main/java/org/graylog/plugins/views/search/SearchType.java b/graylog2-server/src/main/java/org/graylog/plugins/views/search/SearchType.java
@@ -29,8 +29,6 @@
 import org.graylog2.contentpacks.ContentPackable;
 import org.graylog2.contentpacks.EntityDescriptorIds;
 import org.graylog2.contentpacks.exceptions.ContentPackException;
-import org.graylog2.contentpacks.model.ModelTypes;
-import org.graylog2.contentpacks.model.entities.EntityDescriptor;
 import org.graylog2.contentpacks.model.entities.SearchTypeEntity;
 
 import javax.annotation.Nullable;
@@ -42,6 +40,8 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.getStreamEntityId;
+
 /**
  * A search type represents parts of a query that generates a {@see Result result}.
  * <p>
@@ -243,7 +243,7 @@ public SearchTypeEntity toContentPackEntity(EntityDescriptorIds entityDescriptor
 
     default Set<String> mappedStreams(EntityDescriptorIds entityDescriptorIds) {
         return streams().stream()
-                .map(streamId -> entityDescriptorIds.get(EntityDescriptor.create(streamId, ModelTypes.STREAM_V1)))
+                .map(streamId -> getStreamEntityId(streamId, entityDescriptorIds))
                 .map(optionalStreamId -> optionalStreamId.orElseThrow(() ->
                         new ContentPackException("Did not find matching stream id")))
                 .collect(Collectors.toSet());

diff --git a/graylog2-server/src/main/java/org/graylog/plugins/views/search/views/WidgetDTO.java b/graylog2-server/src/main/java/org/graylog/plugins/views/search/views/WidgetDTO.java
@@ -27,8 +27,6 @@
 import org.graylog.plugins.views.search.searchfilters.model.UsesSearchFilters;
 import org.graylog2.contentpacks.ContentPackable;
 import org.graylog2.contentpacks.EntityDescriptorIds;
-import org.graylog2.contentpacks.model.ModelTypes;
-import org.graylog2.contentpacks.model.entities.EntityDescriptor;
 import org.graylog2.contentpacks.model.entities.WidgetEntity;
 import org.graylog2.plugin.indexer.searches.timeranges.TimeRange;
 
@@ -39,6 +37,8 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.getStreamEntityId;
+
 @AutoValue
 @JsonDeserialize(builder = WidgetDTO.Builder.class)
 @WithBeanGetter
@@ -128,8 +128,8 @@ static Builder builder() {
 
     @Override
     public WidgetEntity toContentPackEntity(EntityDescriptorIds entityDescriptorIds) {
-        Set<String> mappedStreams = streams().stream().map(streamId ->
-                entityDescriptorIds.get(EntityDescriptor.create(streamId, ModelTypes.STREAM_V1)))
+        Set<String> mappedStreams = streams().stream()
+                .map(streamId -> getStreamEntityId(streamId, entityDescriptorIds))
                 .filter(Optional::isPresent)
                 .map(Optional::get)
                 .collect(Collectors.toSet());

diff --git a/graylog2-server/src/main/java/org/graylog2/contentpacks/facades/PipelineFacade.java b/graylog2-server/src/main/java/org/graylog2/contentpacks/facades/PipelineFacade.java
@@ -64,6 +64,9 @@
 import java.util.stream.Collectors;
 
 import static java.util.Objects.requireNonNull;
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.getStreamEntityIdOrThrow;
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.resolveStreamEntity;
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.resolveStreamEntityObject;
 
 public class PipelineFacade implements EntityFacade<PipelineDao> {
     private static final Logger LOG = LoggerFactory.getLogger(PipelineFacade.class);
@@ -112,7 +115,7 @@ Entity exportNativeEntity(PipelineDao pipelineDao, EntityDescriptorIds entityDes
     private Set<ValueReference> connectedStreams(String pipelineId, EntityDescriptorIds entityDescriptorIds) {
         final Set<PipelineConnections> connections = connectionsService.loadByPipelineId(pipelineId);
         return connections.stream()
-                .map(pipelineConnections -> entityDescriptorIds.getOrThrow(pipelineConnections.streamId(), ModelTypes.STREAM_V1))
+                .map(pipelineConnections -> getStreamEntityIdOrThrow(pipelineConnections.streamId(), entityDescriptorIds))
                 .map(ValueReference::of)
                 .collect(Collectors.toSet());
     }
@@ -157,7 +160,7 @@ private NativeEntity<PipelineDao> decode(EntityV1 entity,
     private Set<Stream> connectedStreams(Set<EntityDescriptor> connectedStreamEntities, Map<EntityDescriptor, Object> nativeEntities) {
         final ImmutableSet.Builder<Stream> streams = ImmutableSet.builder();
         for (EntityDescriptor descriptor : connectedStreamEntities) {
-            final Object stream = nativeEntities.get(descriptor);
+            final Object stream = resolveStreamEntityObject(descriptor.id().id(), nativeEntities);
             if (stream instanceof Stream) {
                 streams.add((Stream) stream);
             } else {
@@ -283,7 +286,7 @@ public Graph<EntityDescriptor> resolveNativeEntity(EntityDescriptor entityDescri
             pipelineConnections.stream()
                     .map(PipelineConnections::streamId)
                     .map(ModelId::of)
-                    .map(id -> EntityDescriptor.create(id, ModelTypes.STREAM_V1))
+                    .map(id -> EntityDescriptor.create(id, ModelTypes.STREAM_REF_V1))
                     .forEach(stream -> mutableGraph.putEdge(entityDescriptor, stream));
         } catch (NotFoundException e) {
             LOG.debug("Couldn't find pipeline {}", entityDescriptor, e);
@@ -321,9 +324,7 @@ private Graph<Entity> resolveForInstallation(EntityV1 entity,
 
         pipelineEntity.connectedStreams().stream()
                 .map(valueReference -> valueReference.asString(parameters))
-                .map(ModelId::of)
-                .map(modelId -> EntityDescriptor.create(modelId, ModelTypes.STREAM_V1))
-                .map(entities::get)
+                .map(id -> resolveStreamEntity(id, entities))
                 .filter(Objects::nonNull)
                 .forEach(streamEntity -> mutableGraph.putEdge(entity, streamEntity));
 

diff --git a/graylog2-server/src/main/java/org/graylog2/contentpacks/facades/StreamReferenceFacade.java b/graylog2-server/src/main/java/org/graylog2/contentpacks/facades/StreamReferenceFacade.java
@@ -135,19 +135,28 @@ private Optional<NativeEntity<Stream>> findExisting(EntityV1 entity, Map<String,
         }
     }
 
-    public static Object resolveStreamEntity(String id, Map entities) {
+    public static Entity resolveStreamEntity(String id, Map<EntityDescriptor, Entity> entities) {
+        return (Entity) resolveStreamEntityObject(id, entities);
+    }
+
+    public static Object resolveStreamEntityObject(String id, Map entities) {
         Object streamEntity = entities.get(EntityDescriptor.create(id, ModelTypes.STREAM_V1));
         if (streamEntity == null) {
             streamEntity = entities.get(EntityDescriptor.create(id, ModelTypes.STREAM_REF_V1));
         }
         return streamEntity;
     }
 
-    public static Optional<String> getStreamDescriptor(String id, EntityDescriptorIds entityDescriptorIds) {
+    public static Optional<String> getStreamEntityId(String id, EntityDescriptorIds entityDescriptorIds) {
         Optional<String> descriptor = entityDescriptorIds.get(id, ModelTypes.STREAM_V1);
         if (descriptor.isEmpty()) {
             descriptor = entityDescriptorIds.get(id, ModelTypes.STREAM_REF_V1);
         }
         return descriptor;
     }
+
+    public static String getStreamEntityIdOrThrow(String id, EntityDescriptorIds entityDescriptorIds) {
+        return getStreamEntityId(id, entityDescriptorIds).orElseThrow(() ->
+                new ContentPackException("Couldn't find entity " + id + "/" + ModelTypes.STREAM_V1 + " or " + ModelTypes.STREAM_REF_V1));
+    }
 }
diff --git a/graylog2-server/src/main/java/org/graylog2/contentpacks/facades/ViewFacade.java b/graylog2-server/src/main/java/org/graylog2/contentpacks/facades/ViewFacade.java
@@ -58,6 +58,8 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.resolveStreamEntity;
+
 public abstract class ViewFacade implements EntityWithExcerptFacade<ViewDTO, ViewSummaryDTO> {
     private static final Logger LOG = LoggerFactory.getLogger(ViewFacade.class);
 
@@ -188,7 +190,7 @@ public Graph<EntityDescriptor> resolveNativeEntity(EntityDescriptor entityDescri
                 orElseThrow(() -> new NoSuchElementException("Could not find view with id " + modelId.id()));
         final Search search = searchDbService.get(viewSummaryDTO.searchId()).
                 orElseThrow(() -> new NoSuchElementException("Could not find search with id " + viewSummaryDTO.searchId()));
-        search.usedStreamIds().stream().map(s -> EntityDescriptor.create(s, ModelTypes.STREAM_V1))
+        search.usedStreamIds().stream().map(s -> EntityDescriptor.create(s, ModelTypes.STREAM_REF_V1))
                 .forEach(streamDescriptor -> mutableGraph.putEdge(entityDescriptor, streamDescriptor));
         return ImmutableGraph.copyOf(mutableGraph);
     }
@@ -216,8 +218,7 @@ protected Graph<Entity> resolveViewEntity(EntityV1 entity,
         final MutableGraph<Entity> mutableGraph = GraphBuilder.directed().build();
         mutableGraph.addNode(entity);
         viewEntity.search().usedStreamIds().stream()
-                .map(s -> EntityDescriptor.create(s, ModelTypes.STREAM_V1))
-                .map(entities::get)
+                .map(id -> resolveStreamEntity(id, entities))
                 .filter(Objects::nonNull)
                 .forEach(stream -> mutableGraph.putEdge(entity, stream));
         return ImmutableGraph.copyOf(mutableGraph);

diff --git a/graylog2-server/src/main/java/org/graylog2/contentpacks/model/entities/QueryEntity.java b/graylog2-server/src/main/java/org/graylog2/contentpacks/model/entities/QueryEntity.java
@@ -34,7 +34,6 @@
 import org.graylog.plugins.views.search.searchfilters.model.UsedSearchFilter;
 import org.graylog2.contentpacks.NativeEntityConverter;
 import org.graylog2.contentpacks.exceptions.ContentPackException;
-import org.graylog2.contentpacks.model.ModelTypes;
 import org.graylog2.contentpacks.model.entities.references.ValueReference;
 import org.graylog2.plugin.indexer.searches.timeranges.TimeRange;
 import org.graylog2.plugin.streams.Stream;
@@ -53,6 +52,7 @@
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.collect.ImmutableSortedSet.of;
 import static java.util.stream.Collectors.toSet;
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.resolveStreamEntityObject;
 
 @AutoValue
 @JsonAutoDetect
@@ -149,8 +149,7 @@ private Filter shallowMappedFilter(Map<EntityDescriptor, Object> nativeEntities)
                           .map(filter -> {
                               if (filter.type().matches(StreamFilter.NAME)) {
                                   final StreamFilter streamFilter = (StreamFilter) filter;
-                                  final Stream stream = (Stream) nativeEntities.get(
-                                          EntityDescriptor.create(streamFilter.streamId(), ModelTypes.STREAM_V1));
+                                  final Stream stream = (Stream) resolveStreamEntityObject(streamFilter.streamId(), nativeEntities);
                                   if (Objects.isNull(stream)) {
                                       throw new ContentPackException("Could not find matching stream id: " +
                                               streamFilter.streamId());

diff --git a/graylog2-server/src/main/java/org/graylog2/contentpacks/model/entities/SearchTypeEntity.java b/graylog2-server/src/main/java/org/graylog2/contentpacks/model/entities/SearchTypeEntity.java
@@ -27,7 +27,6 @@
 import org.graylog.plugins.views.search.timeranges.DerivedTimeRange;
 import org.graylog2.contentpacks.NativeEntityConverter;
 import org.graylog2.contentpacks.exceptions.ContentPackException;
-import org.graylog2.contentpacks.model.ModelTypes;
 import org.graylog2.contentpacks.model.entities.references.ValueReference;
 import org.graylog2.plugin.streams.Stream;
 
@@ -40,6 +39,8 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.resolveStreamEntityObject;
+
 @JsonTypeInfo(
         use = JsonTypeInfo.Id.NAME,
         include = JsonTypeInfo.As.EXISTING_PROPERTY,
@@ -193,8 +194,7 @@ public SearchType toNativeEntity(Map<String, ValueReference> parameters, Map<Ent
 
     default Set<String> mappedStreams(Map<EntityDescriptor, Object> nativeEntities) {
         return streams().stream()
-                .map(s -> EntityDescriptor.create(s, ModelTypes.STREAM_V1))
-                .map(nativeEntities::get)
+                .map(id -> resolveStreamEntityObject(id, nativeEntities))
                 .map(object -> {
                     if (object == null) {
                         throw new ContentPackException("Missing Stream for event definition");

diff --git a/graylog2-server/src/main/java/org/graylog2/contentpacks/model/entities/WidgetEntity.java b/graylog2-server/src/main/java/org/graylog2/contentpacks/model/entities/WidgetEntity.java
@@ -48,7 +48,6 @@
 import org.graylog.plugins.views.search.views.widgets.aggregation.sort.SortConfigDTO;
 import org.graylog2.contentpacks.NativeEntityConverter;
 import org.graylog2.contentpacks.exceptions.ContentPackException;
-import org.graylog2.contentpacks.model.ModelTypes;
 import org.graylog2.contentpacks.model.entities.references.ValueReference;
 import org.graylog2.plugin.indexer.searches.timeranges.TimeRange;
 import org.graylog2.plugin.streams.Stream;
@@ -64,6 +63,8 @@
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
+import static org.graylog2.contentpacks.facades.StreamReferenceFacade.resolveStreamEntityObject;
+
 @AutoValue
 @JsonDeserialize(builder = WidgetEntity.Builder.class)
 @WithBeanGetter
@@ -155,8 +156,7 @@ public WidgetDTO toNativeEntity(Map<String, ValueReference> parameters, Map<Enti
                 .filters(convertSearchFilters(this.filters()))
                 .id(this.id())
                 .streams(this.streams().stream()
-                        .map(id -> EntityDescriptor.create(id, ModelTypes.STREAM_V1))
-                        .map(nativeEntities::get)
+                        .map(id -> resolveStreamEntityObject(id, nativeEntities))
                         .map(object -> {
                             if (object == null) {
                                 throw new ContentPackException("Missing Stream for widget entity");