diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java index 32a444e5..9665d86e 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java +++ b/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> sorts, Map measurements) { @@ -117,14 +126,56 @@ private void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder, Map> sortBuilderList, Boolean showFields) { + if (sortBuilderList == null || sortBuilderList.isEmpty()) { + return; + } + StringBuilder sortShape = new StringBuilder(); + sortShape.append("sort:\n"); + + List 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 aggregationBuilders, + Collection 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 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; + } }