From 2f995db057be6bf6dfdba833804bfaa20ac72107 Mon Sep 17 00:00:00 2001 From: David Zane Date: Wed, 24 Jul 2024 14:45:23 -0700 Subject: [PATCH] Created QueryShapeService Signed-off-by: David Zane --- .../categorizer/QueryShapeService.java | 142 ++++++++++++++++++ .../categorizer/SearchQueryCategorizer.java | 121 +-------------- 2 files changed, 147 insertions(+), 116 deletions(-) create mode 100644 src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeService.java diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeService.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeService.java new file mode 100644 index 0000000..fdf8f75 --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeService.java @@ -0,0 +1,142 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.insights.core.service.categorizer; + +import org.opensearch.index.query.QueryBuilder; +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 java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; + +class QueryShapeService { + static final String TWO_SPACE_INDENT = " "; + + static String buildShape(SearchSourceBuilder source, Boolean showFields) { + StringBuilder shape = new StringBuilder(); + shape.append(buildQueryShape(source.query())); + shape.append(buildAggregationShape(source.aggregations(), showFields)); + shape.append(buildSortShape(source.sorts(), showFields)); + return shape.toString(); + } + + static String buildQueryShape(QueryBuilder queryBuilder) { + if (queryBuilder == null) { + return ""; + } + QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); + queryBuilder.visit(shapeVisitor); + return shapeVisitor.prettyPrintTree(""); + } + + static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuilder, Boolean showFields) { + if (aggregationsBuilder == null) { + return ""; + } + StringBuilder aggregationShape = recursiveAggregationShapeBuilder( + aggregationsBuilder.getAggregatorFactories(), + aggregationsBuilder.getPipelineAggregatorFactories(), + new StringBuilder(), + 0, + showFields + ); + return aggregationShape.toString(); + } + + static StringBuilder recursiveAggregationShapeBuilder( + 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 + recursiveAggregationShapeBuilder( + 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; + } + + static String buildSortShape(List> 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"); + } + return sortShape.toString(); + } +} 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 9665d86..5f394c2 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,19 +15,12 @@ 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; @@ -38,7 +31,7 @@ public final class SearchQueryCategorizer { private static final Logger logger = LogManager.getLogger(SearchQueryCategorizer.class); - public static final String TWO_SPACE_INDENT = " "; + /** * Contains all the search query counters */ @@ -91,13 +84,12 @@ public void categorize(SearchQueryRecord record) { SearchSourceBuilder source = (SearchSourceBuilder) record.getAttributes().get(Attribute.SOURCE); Map measurements = record.getMeasurements(); - QueryBuilder topLevelQueryBuilder = source.query(); - logQueryShape(topLevelQueryBuilder); - incrementQueryTypeCounters(topLevelQueryBuilder, measurements); + incrementQueryTypeCounters(source.query(), measurements); incrementQueryAggregationCounters(source.aggregations(), measurements); - logAggregationsShape(source.aggregations(), true); incrementQuerySortCounters(source.sorts(), measurements); - logSortShape(source.sorts(), true); + + String searchShape = QueryShapeService.buildShape(source, true); + logger.trace(searchShape); } private void incrementQuerySortCounters(List> sorts, Map measurements) { @@ -125,59 +117,6 @@ 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()); - } - /** * Get search query counters * @return search query counters @@ -194,54 +133,4 @@ 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; - } }