Refactoring/agg types 2 (#24437)

* Removing storage-specyfic method from AggTypes

* Simplified bucket handlers

* AggTypes class not duplicated anymore in storage modules

* Get rid of agg-types record methods in bucket handlers

* Better name for AggTypes

* Series handlers simplified

* Ranaming

* Removed record method from series handler

* Pivot context init moved to a single place

* Missing license header

* Init of pivot context replaced with init on access, renaming

Author: luk-kaminski

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/AggTypes.java

@@ -18,6 +18,7 @@
 
 import com.google.common.collect.ImmutableList;
 import io.opentelemetry.instrumentation.annotations.WithSpan;
+import jakarta.inject.Inject;
 import org.graylog.plugins.views.search.Query;
 import org.graylog.plugins.views.search.SearchJob;
 import org.graylog.plugins.views.search.SearchType;
@@ -42,8 +43,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import jakarta.inject.Inject;
-
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.List;
@@ -72,10 +71,6 @@ public void doGenerateQueryPart(Query query, Pivot pivot, ESGeneratedQueryContex
         LOG.debug("Generating aggregation for {}", pivot);
         final SearchSourceBuilder searchSourceBuilder = queryContext.searchSourceBuilder(pivot);
 
-        final Map<Object, Object> contextMap = queryContext.contextMap();
-        final AggTypes aggTypes = new AggTypes();
-        contextMap.put(pivot.id(), aggTypes);
-
         var generateRollups = pivot.rollup() || (pivot.rowGroups().isEmpty() && pivot.columnGroups().isEmpty());
 
         // add global rollup series if those were requested

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/ESPivot.java

@@ -21,11 +21,9 @@
 import org.graylog.plugins.views.search.searchtypes.pivot.BucketSpecHandler;
 import org.graylog.plugins.views.search.searchtypes.pivot.Pivot;
 import org.graylog.plugins.views.search.searchtypes.pivot.PivotSort;
-import org.graylog.plugins.views.search.searchtypes.pivot.PivotSpec;
 import org.graylog.plugins.views.search.searchtypes.pivot.SeriesSort;
 import org.graylog.plugins.views.search.searchtypes.pivot.SeriesSpec;
 import org.graylog.plugins.views.search.searchtypes.pivot.SortSpec;
-import org.graylog.shaded.elasticsearch7.org.elasticsearch.search.aggregations.Aggregation;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.search.aggregations.AggregationBuilder;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.search.aggregations.AggregationBuilders;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.search.aggregations.BucketOrder;
@@ -41,14 +39,6 @@
 public abstract class ESPivotBucketSpecHandler<SPEC_TYPE extends BucketSpec>
         implements BucketSpecHandler<SPEC_TYPE, AggregationBuilder, ESGeneratedQueryContext> {
 
-    protected AggTypes aggTypes(ESGeneratedQueryContext queryContext, Pivot pivot) {
-        return (AggTypes) queryContext.contextMap().get(pivot.id());
-    }
-
-    protected void record(ESGeneratedQueryContext queryContext, Pivot pivot, PivotSpec spec, String name, Class<? extends Aggregation> aggregationClass) {
-        aggTypes(queryContext, pivot).record(spec, name);
-    }
-
     public record SortOrders(List<BucketOrder> orders, List<AggregationBuilder> sortingAggregations) {}
 
     protected SortOrders orderListForPivot(Pivot pivot, ESGeneratedQueryContext esGeneratedQueryContext, BucketOrder defaultOrder, Query query) {

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/ESPivotBucketSpecHandler.java

@@ -33,16 +33,9 @@
 public abstract class ESPivotSeriesSpecHandler<SPEC_TYPE extends SeriesSpec, AGGREGATION_RESULT extends Aggregation>
         implements SeriesSpecHandler<SPEC_TYPE, SeriesAggregationBuilder, SearchResponse, AGGREGATION_RESULT, ESSearchTypeHandler<Pivot>, ESGeneratedQueryContext> {
 
-    protected AggTypes aggTypes(ESGeneratedQueryContext queryContext, Pivot pivot) {
-        return (AggTypes) queryContext.contextMap().get(pivot.id());
-    }
-
-    protected void record(ESGeneratedQueryContext queryContext, Pivot pivot, PivotSpec spec, String name, Class<? extends Aggregation> aggregationClass) {
-        aggTypes(queryContext, pivot).record(spec, name);
-    }
-
-    public Aggregation extractAggregationFromResult(Pivot pivot, PivotSpec spec, HasAggregations aggregations, ESGeneratedQueryContext queryContext) {
-        return aggTypes(queryContext, pivot).getSubAggregation(spec, aggregations);
+    public Aggregation extractAggregationFromResult(Pivot pivot, PivotSpec spec, HasAggregations currentAggregationOrBucket, ESGeneratedQueryContext queryContext) {
+        final String aggName = queryContext.getAggNameForPivotSpecFromContext(pivot, spec);
+        return currentAggregationOrBucket.getAggregations().get(aggName);
     }
 
     @SuppressWarnings("unchecked")

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/ESPivotSeriesSpecHandler.java

@@ -57,7 +57,7 @@ public CreatedAggregations<AggregationBuilder> doCreateAggregation(Direction dir
             builder.format("date_time");
             builder.keyed(false);
 
-            record(queryContext, pivot, dateRangeBucket, name, ParsedDateRange.class);
+            queryContext.recordNameForPivotSpec(pivot, dateRangeBucket, name);
 
             if (root == null && leaf == null) {
                 root = builder;

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/buckets/ESDateRangeHandler.java

@@ -36,7 +36,7 @@ public class ESAverageHandler extends ESPivotSeriesSpecHandler<Average, Avg> {
     @Override
     public List<SeriesAggregationBuilder> doCreateAggregation(String name, Pivot pivot, Average avgSpec, ESSearchTypeHandler<Pivot> searchTypeHandler, ESGeneratedQueryContext queryContext) {
         final AvgAggregationBuilder avg = AggregationBuilders.avg(name).field(avgSpec.field());
-        record(queryContext, pivot, avgSpec, name, Avg.class);
+        queryContext.recordNameForPivotSpec(pivot, avgSpec, name);
         return List.of(SeriesAggregationBuilder.metric(avg));
     }
 

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/series/ESAverageHandler.java

@@ -35,7 +35,7 @@ public class ESCardinalityHandler extends ESPivotSeriesSpecHandler<Cardinality,
     @Override
     public List<SeriesAggregationBuilder> doCreateAggregation(String name, Pivot pivot, Cardinality cardinalitySpec, ESSearchTypeHandler<Pivot> searchTypeHandler, ESGeneratedQueryContext queryContext) {
         final CardinalityAggregationBuilder card = AggregationBuilders.cardinality(name).field(cardinalitySpec.field());
-        record(queryContext, pivot, cardinalitySpec, name, org.graylog.shaded.elasticsearch7.org.elasticsearch.search.aggregations.metrics.Cardinality.class);
+        queryContext.recordNameForPivotSpec(pivot, cardinalitySpec, name);
         return List.of(SeriesAggregationBuilder.metric(card));
     }
 

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/series/ESCardinalityHandler.java

@@ -52,7 +52,7 @@ public List<SeriesAggregationBuilder> doCreateAggregation(String name, Pivot piv
                 .map(field -> {
                     // the request was for a field count, we have to add a value_count sub aggregation
                     final ValueCountAggregationBuilder value = AggregationBuilders.count(name).field(field);
-                    record(queryContext, pivot, count, name, ValueCount.class);
+                    queryContext.recordNameForPivotSpec(pivot, count, name);
                     return List.of(SeriesAggregationBuilder.metric(value));
                 })
                 // doc_count is always present in elasticsearch's bucket aggregations, no need to add it
@@ -82,7 +82,7 @@ public Stream<Value> doHandleResult(Pivot pivot,
 
     @Override
     public Aggregation extractAggregationFromResult(Pivot pivot, PivotSpec spec, HasAggregations aggregations, ESGeneratedQueryContext queryContext) {
-        final String agg = aggTypes(queryContext, pivot).getTypes(spec);
+        final String agg = queryContext.getAggNameForPivotSpecFromContext(pivot, spec);
         if (agg == null) {
             if (aggregations instanceof MultiBucketsAggregation.Bucket) {
                 return createValueCount(((MultiBucketsAggregation.Bucket) aggregations).getDocCount());

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/series/ESCountHandler.java

@@ -49,7 +49,7 @@ public List<SeriesAggregationBuilder> doCreateAggregation(String name, Pivot piv
                         .size(1)
                         .fetchSource(latestSpec.field(), null)
                         .sort(SortBuilders.fieldSort("timestamp").order(SortOrder.DESC)));
-        record(queryContext, pivot, latestSpec, name, ParsedFilter.class);
+        queryContext.recordNameForPivotSpec(pivot, latestSpec, name);
         return List.of(SeriesAggregationBuilder.metric(latest));
     }
 

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/views/searchtypes/pivot/series/ESLatestHandler.java

@@ -35,7 +35,7 @@ public class ESMaxHandler extends ESPivotSeriesSpecHandler<Max, org.graylog.shad
     @Override
     public List<SeriesAggregationBuilder> doCreateAggregation(String name, Pivot pivot, Max maxSpec, ESSearchTypeHandler<Pivot> searchTypeHandler, ESGeneratedQueryContext queryContext) {
         final MaxAggregationBuilder max = AggregationBuilders.max(name).field(maxSpec.field());
-        record(queryContext, pivot, maxSpec, name, org.graylog.shaded.elasticsearch7.org.elasticsearch.search.aggregations.metrics.Max.class);
+        queryContext.recordNameForPivotSpec(pivot, maxSpec, name);
         return List.of(SeriesAggregationBuilder.metric(max));
     }