Query shape for agg & sort

Signed-off-by: David Zane <davizane@amazon.com>

src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java

@@ -15,12 +15,19 @@
 import org.opensearch.plugin.insights.rules.model.Attribute;
 import org.opensearch.plugin.insights.rules.model.MetricType;
 import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
+import org.opensearch.search.aggregations.AggregationBuilder;
 import org.opensearch.search.aggregations.AggregatorFactories;
+import org.opensearch.search.aggregations.PipelineAggregationBuilder;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.sort.FieldSortBuilder;
 import org.opensearch.search.sort.SortBuilder;
 import org.opensearch.telemetry.metrics.MetricsRegistry;
 import org.opensearch.telemetry.metrics.tags.Tags;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -30,8 +37,8 @@
  */
 public final class SearchQueryCategorizer {
 
-    private static final Logger log = LogManager.getLogger(SearchQueryCategorizer.class);
-
+    private static final Logger logger = LogManager.getLogger(SearchQueryCategorizer.class);
+    public static final String TWO_SPACE_INDENT = "  ";
     /**
      * Contains all the search query counters
      */
@@ -88,7 +95,9 @@ public void categorize(SearchQueryRecord record) {
         logQueryShape(topLevelQueryBuilder);
         incrementQueryTypeCounters(topLevelQueryBuilder, measurements);
         incrementQueryAggregationCounters(source.aggregations(), measurements);
+        logAggregationsShape(source.aggregations(), true);
         incrementQuerySortCounters(source.sorts(), measurements);
+        logSortShape(source.sorts(), true);
     }
 
     private void incrementQuerySortCounters(List<SortBuilder<?>> sorts, Map<MetricType, Number> measurements) {
@@ -117,14 +126,56 @@ private void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder, Map<M
     }
 
     private void logQueryShape(QueryBuilder topLevelQueryBuilder) {
-        if (log.isTraceEnabled()) {
+        if (logger.isTraceEnabled()) {
             if (topLevelQueryBuilder == null) {
                 return;
             }
             QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
             topLevelQueryBuilder.visit(shapeVisitor);
-            log.trace("Query shape : {}", shapeVisitor.prettyPrintTree("  "));
+            logger.trace(shapeVisitor.prettyPrintTree(TWO_SPACE_INDENT));
+        }
+    }
+
+    private void logAggregationsShape(AggregatorFactories.Builder aggregationsBuilder, Boolean showFields) {
+        if (aggregationsBuilder == null) {
+            return;
+        }
+        StringBuilder aggregationShape = getAggregationShape(
+            aggregationsBuilder.getAggregatorFactories(),
+            aggregationsBuilder.getPipelineAggregatorFactories(),
+            new StringBuilder(),
+            0,
+            showFields
+        );
+        logger.trace(aggregationShape.toString());
+    }
+
+    private void logSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showFields) {
+        if (sortBuilderList == null || sortBuilderList.isEmpty()) {
+            return;
+        }
+        StringBuilder sortShape = new StringBuilder();
+        sortShape.append("sort:\n");
+
+        List<String> shapeStrings = new ArrayList<>();
+        for (SortBuilder<?> sortBuilder : sortBuilderList) {
+            StringBuilder stringBuilder = new StringBuilder();
+            stringBuilder.append(TWO_SPACE_INDENT).append(sortBuilder.order());
+            if (showFields) {
+                if (sortBuilder instanceof FieldSortBuilder) {
+                    stringBuilder.append(" [").append(((FieldSortBuilder) sortBuilder).getFieldName()).append("]");
+                } else {
+                    stringBuilder.append(" []");
+                }
+            }
+            shapeStrings.add(stringBuilder.toString());
+        }
+
+        Collections.sort(shapeStrings);
+        for (String line : shapeStrings) {
+            sortShape.append(line).append("\n");
         }
+        logger.trace(sortShape.toString());
     }
 
     /**
@@ -143,4 +194,54 @@ public void reset() {
             instance = null;
         }
     }
+
+    public StringBuilder getAggregationShape(
+        Collection<AggregationBuilder> aggregationBuilders,
+        Collection<PipelineAggregationBuilder> pipelineAggregations,
+        StringBuilder outputBuilder,
+        int indentCount,
+        Boolean showFields
+    ) {
+        String baseIndent = TWO_SPACE_INDENT.repeat(indentCount);
+
+        //// Normal Aggregations ////
+        if (aggregationBuilders.isEmpty() == false) {
+            outputBuilder.append(baseIndent).append("aggregation:").append("\n");
+        }
+        List<String> aggShapeStrings = new ArrayList<>();
+        for (AggregationBuilder agg : aggregationBuilders) {
+            StringBuilder stringBuilder = new StringBuilder();
+            stringBuilder.append(baseIndent).append(TWO_SPACE_INDENT).append(agg.getType());
+            if (showFields) {
+                if (agg instanceof ValuesSourceAggregationBuilder) {
+                    stringBuilder.append(" [").append(((ValuesSourceAggregationBuilder<?>) agg).field()).append("]");
+                } else {
+                    stringBuilder.append(" []");
+                }
+            }
+            stringBuilder.append("\n");
+
+            if (agg.getSubAggregations().isEmpty() == false) {
+                // Recursive call on sub-aggregations
+                stringBuilder = getAggregationShape(agg.getSubAggregations(), agg.getPipelineAggregations(), stringBuilder, indentCount + 2, showFields);
+            }
+            aggShapeStrings.add(stringBuilder.toString());
+        }
+
+        // Sort aggregations
+        Collections.sort(aggShapeStrings);
+        for (String shapeString : aggShapeStrings) {
+            outputBuilder.append(shapeString);
+        }
+
+        //// Pipeline Aggregation (cannot have sub-aggregations) ////
+        if (pipelineAggregations.isEmpty() == false) {
+            outputBuilder.append(baseIndent).append("pipeline aggregation:").append("\n");
+        }
+        for (PipelineAggregationBuilder pipelineAgg : pipelineAggregations) {
+            outputBuilder.append(baseIndent).append(TWO_SPACE_INDENT).append(pipelineAgg.getName()).append("\n");
+        }
+
+        return outputBuilder;
+    }
 }