From c332bd99b7a71382235ded2f0192f88e43cf5606 Mon Sep 17 00:00:00 2001 From: Siddhant Deshmukh Date: Mon, 15 Jul 2024 14:51:50 -0700 Subject: [PATCH] Move query categorization changes to plugin Signed-off-by: Siddhant Deshmukh --- .../plugin/insights/QueryInsightsPlugin.java | 20 +- .../core/listener/QueryInsightsListener.java | 16 +- .../core/service/QueryInsightsService.java | 60 ++++- .../categorizer/QueryShapeVisitor.java | 102 ++++++++ .../SearchQueryAggregationCategorizer.java | 63 +++++ .../categorizer/SearchQueryCategorizer.java | 132 ++++++++++ .../SearchQueryCategorizingVisitor.java | 39 +++ .../categorizer/SearchQueryCounters.java | 134 ++++++++++ .../service/categorizer/package-info.java | 12 + .../settings/QueryCategorizationSettings.java | 32 +++ .../insights/QueryInsightsPluginTests.java | 8 +- .../insights/QueryInsightsTestUtils.java | 13 +- .../listener/QueryInsightsListenerTests.java | 7 +- .../service/QueryInsightsServiceTests.java | 26 +- .../categorizor/QueryShapeVisitorTests.java | 38 +++ .../SearchQueryCategorizerTests.java | 247 ++++++++++++++++++ 16 files changed, 924 insertions(+), 25 deletions(-) create mode 100644 src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java create mode 100644 src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryAggregationCategorizer.java create mode 100644 src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java create mode 100644 src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizingVisitor.java create mode 100644 src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java create mode 100644 src/main/java/org/opensearch/plugin/insights/core/service/categorizer/package-info.java create mode 100644 src/main/java/org/opensearch/plugin/insights/settings/QueryCategorizationSettings.java create mode 100644 src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeVisitorTests.java create mode 100644 src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java diff --git a/src/main/java/org/opensearch/plugin/insights/QueryInsightsPlugin.java b/src/main/java/org/opensearch/plugin/insights/QueryInsightsPlugin.java index bba67643..ba33bbba 100644 --- a/src/main/java/org/opensearch/plugin/insights/QueryInsightsPlugin.java +++ b/src/main/java/org/opensearch/plugin/insights/QueryInsightsPlugin.java @@ -30,13 +30,17 @@ import org.opensearch.plugin.insights.rules.action.top_queries.TopQueriesAction; import org.opensearch.plugin.insights.rules.resthandler.top_queries.RestTopQueriesAction; import org.opensearch.plugin.insights.rules.transport.top_queries.TransportTopQueriesAction; +import org.opensearch.plugin.insights.settings.QueryCategorizationSettings; import org.opensearch.plugin.insights.settings.QueryInsightsSettings; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; +import org.opensearch.plugins.TelemetryAwarePlugin; import org.opensearch.repositories.RepositoriesService; import org.opensearch.rest.RestController; import org.opensearch.rest.RestHandler; import org.opensearch.script.ScriptService; +import org.opensearch.telemetry.metrics.MetricsRegistry; +import org.opensearch.telemetry.tracing.Tracer; import org.opensearch.threadpool.ExecutorBuilder; import org.opensearch.threadpool.ScalingExecutorBuilder; import org.opensearch.threadpool.ThreadPool; @@ -49,7 +53,7 @@ /** * Plugin class for Query Insights. */ -public class QueryInsightsPlugin extends Plugin implements ActionPlugin { +public class QueryInsightsPlugin extends Plugin implements ActionPlugin, TelemetryAwarePlugin { /** * Default constructor */ @@ -67,10 +71,17 @@ public Collection createComponents( final NodeEnvironment nodeEnvironment, final NamedWriteableRegistry namedWriteableRegistry, final IndexNameExpressionResolver indexNameExpressionResolver, - final Supplier repositoriesServiceSupplier + final Supplier repositoriesServiceSupplier, + final Tracer tracer, + final MetricsRegistry metricsRegistry ) { // create top n queries service - final QueryInsightsService queryInsightsService = new QueryInsightsService(clusterService.getClusterSettings(), threadPool, client); + final QueryInsightsService queryInsightsService = new QueryInsightsService( + clusterService.getClusterSettings(), + threadPool, + client, + metricsRegistry + ); return List.of(queryInsightsService, new QueryInsightsListener(clusterService, queryInsightsService)); } @@ -119,7 +130,8 @@ public List> getSettings() { QueryInsightsSettings.TOP_N_MEMORY_QUERIES_ENABLED, QueryInsightsSettings.TOP_N_MEMORY_QUERIES_SIZE, QueryInsightsSettings.TOP_N_MEMORY_QUERIES_WINDOW_SIZE, - QueryInsightsSettings.TOP_N_MEMORY_EXPORTER_SETTINGS + QueryInsightsSettings.TOP_N_MEMORY_EXPORTER_SETTINGS, + QueryCategorizationSettings.SEARCH_QUERY_METRICS_ENABLED_SETTING ); } } diff --git a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java index 8f186769..679e2b96 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java +++ b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java @@ -94,18 +94,20 @@ public QueryInsightsListener(final ClusterService clusterService, final QueryIns * @param enabled boolean */ public void setEnableTopQueries(final MetricType metricType, final boolean enabled) { - boolean isAllMetricsDisabled = !queryInsightsService.isEnabled(); - this.queryInsightsService.enableCollection(metricType, enabled); + boolean isInsightsServiceDisabled = !queryInsightsService.isEnabled(); + if (!enabled) { - // disable QueryInsightsListener only if all metrics collections are disabled now. - if (!queryInsightsService.isEnabled()) { + // disable QueryInsightsListener only if all metrics collections are disabled now + // and search query metrics is disabled. + if (isInsightsServiceDisabled) { super.setEnabled(false); this.queryInsightsService.stop(); } } else { super.setEnabled(true); - // restart QueryInsightsListener only if none of metrics collections is enabled before. - if (isAllMetricsDisabled) { + // restart QueryInsightsListener only if none of metrics collections is enabled before and + // search query metrics is disabled before. + if (isInsightsServiceDisabled) { this.queryInsightsService.stop(); this.queryInsightsService.start(); } @@ -176,7 +178,7 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final } Map attributes = new HashMap<>(); attributes.put(Attribute.SEARCH_TYPE, request.searchType().toString().toLowerCase(Locale.ROOT)); - attributes.put(Attribute.SOURCE, request.source().toString(FORMAT_PARAMS)); + attributes.put(Attribute.SOURCE, request.source()); attributes.put(Attribute.TOTAL_SHARDS, context.getNumShards()); attributes.put(Attribute.INDICES, request.indices()); attributes.put(Attribute.PHASE_LATENCY_MAP, searchRequestContext.phaseTookMap()); diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java b/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java index 33343aa6..40e63f78 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java @@ -8,6 +8,8 @@ package org.opensearch.plugin.insights.core.service; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.client.Client; import org.opensearch.common.inject.Inject; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; @@ -15,9 +17,11 @@ import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; import org.opensearch.plugin.insights.core.exporter.QueryInsightsExporterFactory; +import org.opensearch.plugin.insights.core.service.categorizer.SearchQueryCategorizer; import org.opensearch.plugin.insights.rules.model.MetricType; import org.opensearch.plugin.insights.rules.model.SearchQueryRecord; import org.opensearch.plugin.insights.settings.QueryInsightsSettings; +import org.opensearch.telemetry.metrics.MetricsRegistry; import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; @@ -29,6 +33,7 @@ import java.util.Map; import java.util.concurrent.LinkedBlockingQueue; +import static org.opensearch.plugin.insights.settings.QueryCategorizationSettings.SEARCH_QUERY_METRICS_ENABLED_SETTING; import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.getExporterSettings; /** @@ -36,6 +41,9 @@ * information related to search queries */ public class QueryInsightsService extends AbstractLifecycleComponent { + + private static final Logger logger = LogManager.getLogger(QueryInsightsService.class); + /** * The internal OpenSearch thread pool that execute async processing and exporting tasks */ @@ -67,15 +75,25 @@ public class QueryInsightsService extends AbstractLifecycleComponent { */ final QueryInsightsExporterFactory queryInsightsExporterFactory; + private volatile boolean searchQueryMetricsEnabled; + + private SearchQueryCategorizer searchQueryCategorizer; + /** * Constructor of the QueryInsightsService * * @param clusterSettings OpenSearch cluster level settings * @param threadPool The OpenSearch thread pool to run async tasks * @param client OS client + * @param metricsRegistry Opentelemetry Metrics registry */ @Inject - public QueryInsightsService(final ClusterSettings clusterSettings, final ThreadPool threadPool, final Client client) { + public QueryInsightsService( + final ClusterSettings clusterSettings, + final ThreadPool threadPool, + final Client client, + final MetricsRegistry metricsRegistry + ) { enableCollect = new HashMap<>(); queryRecordsQueue = new LinkedBlockingQueue<>(QueryInsightsSettings.QUERY_RECORD_QUEUE_CAPACITY); this.threadPool = threadPool; @@ -93,6 +111,10 @@ public QueryInsightsService(final ClusterSettings clusterSettings, final ThreadP (settings -> validateExporterConfig(type, settings)) ); } + + this.searchQueryMetricsEnabled = clusterSettings.get(SEARCH_QUERY_METRICS_ENABLED_SETTING); + this.searchQueryCategorizer = SearchQueryCategorizer.getInstance(metricsRegistry); + clusterSettings.addSettingsUpdateConsumer(SEARCH_QUERY_METRICS_ENABLED_SETTING, this::setSearchQueryMetricsEnabled); } /** @@ -133,6 +155,14 @@ public void drainRecords() { topQueriesServices.get(metricType).consumeRecords(records); } } + + if (searchQueryMetricsEnabled) { + try { + searchQueryCategorizer.consumeRecords(records); + } catch (Exception e) { + logger.error("Error while trying to categorize the queries.", e); + } + } } /** @@ -166,7 +196,7 @@ public boolean isCollectionEnabled(final MetricType metricType) { } /** - * Check if query insights service is enabled + * Check if any feature of Query Insights service is enabled, right now includes Top N and Categorization. * * @return if query insights service is enabled */ @@ -176,7 +206,7 @@ public boolean isEnabled() { return true; } } - return false; + return this.searchQueryMetricsEnabled; } /** @@ -239,6 +269,30 @@ public void setExporter(final MetricType type, final Settings settings) { } } + /** + * Set search query metrics enabled to enable collection of search query categorization metrics + * @param searchQueryMetricsEnabled boolean flag + */ + public void setSearchQueryMetricsEnabled(boolean searchQueryMetricsEnabled) { + this.searchQueryMetricsEnabled = searchQueryMetricsEnabled; + } + + /** + * Is search query metrics feature enabled. + * @return boolean flag + */ + public boolean isSearchQueryMetricsEnabled() { + return this.searchQueryMetricsEnabled; + } + + /** + * Get search query categorizer object + * @return SearchQueryCategorizer object + */ + public SearchQueryCategorizer getSearchQueryCategorizer() { + return this.searchQueryCategorizer; + } + /** * Validate the exporter config for a metricType * diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java new file mode 100644 index 00000000..626e05dc --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java @@ -0,0 +1,102 @@ +/* + * 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.apache.lucene.search.BooleanClause; +import org.opensearch.common.SetOnce; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; + +import java.util.ArrayList; +import java.util.EnumMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +/** + * Class to traverse the QueryBuilder tree and capture the query shape + */ +public final class QueryShapeVisitor implements QueryBuilderVisitor { + private final SetOnce queryType = new SetOnce<>(); + private final Map> childVisitors = new EnumMap<>(BooleanClause.Occur.class); + + @Override + public void accept(QueryBuilder qb) { + queryType.set(qb.getName()); + } + + @Override + public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { + // Should get called once per Occur value + if (childVisitors.containsKey(occur)) { + throw new IllegalStateException("child visitor already called for " + occur); + } + final List childVisitorList = new ArrayList<>(); + QueryBuilderVisitor childVisitorWrapper = new QueryBuilderVisitor() { + QueryShapeVisitor currentChild; + + @Override + public void accept(QueryBuilder qb) { + currentChild = new QueryShapeVisitor(); + childVisitorList.add(currentChild); + currentChild.accept(qb); + } + + @Override + public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { + return currentChild.getChildVisitor(occur); + } + }; + childVisitors.put(occur, childVisitorList); + return childVisitorWrapper; + } + + /** + * Convert query builder tree to json + * @return json query builder tree as a string + */ + public String toJson() { + StringBuilder outputBuilder = new StringBuilder("{\"type\":\"").append(queryType.get()).append("\""); + for (Map.Entry> entry : childVisitors.entrySet()) { + outputBuilder.append(",\"").append(entry.getKey().name().toLowerCase(Locale.ROOT)).append("\"["); + boolean first = true; + for (QueryShapeVisitor child : entry.getValue()) { + if (!first) { + outputBuilder.append(","); + } + outputBuilder.append(child.toJson()); + first = false; + } + outputBuilder.append("]"); + } + outputBuilder.append("}"); + return outputBuilder.toString(); + } + + /** + * Pretty print the query builder tree + * @param indent indent size + * @return Query builder tree as a pretty string + */ + public String prettyPrintTree(String indent) { + StringBuilder outputBuilder = new StringBuilder(indent).append(queryType.get()).append("\n"); + for (Map.Entry> entry : childVisitors.entrySet()) { + outputBuilder.append(indent).append(" ").append(entry.getKey().name().toLowerCase(Locale.ROOT)).append(":\n"); + for (QueryShapeVisitor child : entry.getValue()) { + outputBuilder.append(child.prettyPrintTree(indent + " ")); + } + } + return outputBuilder.toString(); + } + + /** + * Default constructor + */ + public QueryShapeVisitor() {} +} diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryAggregationCategorizer.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryAggregationCategorizer.java new file mode 100644 index 00000000..f252795a --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryAggregationCategorizer.java @@ -0,0 +1,63 @@ +/* + * 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.search.aggregations.AggregationBuilder; +import org.opensearch.search.aggregations.PipelineAggregationBuilder; +import org.opensearch.telemetry.metrics.tags.Tags; + +import java.util.Collection; + +/** + * Increments the counters related to Aggregation Search Queries. + */ +public class SearchQueryAggregationCategorizer { + + private static final String TYPE_TAG = "type"; + private final SearchQueryCounters searchQueryCounters; + + /** + * Constructor for SearchQueryAggregationCategorizer + * @param searchQueryCounters Object containing all query counters + */ + public SearchQueryAggregationCategorizer(SearchQueryCounters searchQueryCounters) { + this.searchQueryCounters = searchQueryCounters; + } + + /** + * Increment aggregation related counters + * @param aggregatorFactories input aggregations + */ + public void incrementSearchQueryAggregationCounters(Collection aggregatorFactories) { + for (AggregationBuilder aggregationBuilder : aggregatorFactories) { + incrementCountersRecursively(aggregationBuilder); + } + } + + private void incrementCountersRecursively(AggregationBuilder aggregationBuilder) { + // Increment counters for the current aggregation + String aggregationType = aggregationBuilder.getType(); + searchQueryCounters.incrementAggCounter(1, Tags.create().addTag(TYPE_TAG, aggregationType)); + + // Recursively process sub-aggregations if any + Collection subAggregations = aggregationBuilder.getSubAggregations(); + if (subAggregations != null && !subAggregations.isEmpty()) { + for (AggregationBuilder subAggregation : subAggregations) { + incrementCountersRecursively(subAggregation); + } + } + + // Process pipeline aggregations + Collection pipelineAggregations = aggregationBuilder.getPipelineAggregations(); + for (PipelineAggregationBuilder pipelineAggregation : pipelineAggregations) { + String pipelineAggregationType = pipelineAggregation.getType(); + searchQueryCounters.incrementAggCounter(1, Tags.create().addTag(TYPE_TAG, pipelineAggregationType)); + } + } +} 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 new file mode 100644 index 00000000..851912a6 --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java @@ -0,0 +1,132 @@ +/* + * 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.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; +import org.opensearch.plugin.insights.rules.model.Attribute; +import org.opensearch.plugin.insights.rules.model.SearchQueryRecord; +import org.opensearch.search.aggregations.AggregatorFactories; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.sort.SortBuilder; +import org.opensearch.telemetry.metrics.MetricsRegistry; +import org.opensearch.telemetry.metrics.tags.Tags; + +import java.util.List; + +/** + * Class to categorize the search queries based on the type and increment the relevant counters. + * Class also logs the query shape. + */ +public final class SearchQueryCategorizer { + + private static final Logger log = LogManager.getLogger(SearchQueryCategorizer.class); + + /** + * Contains all the search query counters + */ + private final SearchQueryCounters searchQueryCounters; + + final SearchQueryAggregationCategorizer searchQueryAggregationCategorizer; + private static SearchQueryCategorizer instance; + + /** + * Constructor for SearchQueryCategorizor + * @param metricsRegistry opentelemetry metrics registry + */ + private SearchQueryCategorizer(MetricsRegistry metricsRegistry) { + searchQueryCounters = new SearchQueryCounters(metricsRegistry); + searchQueryAggregationCategorizer = new SearchQueryAggregationCategorizer(searchQueryCounters); + } + + /** + * Get singleton instance of SearchQueryCategorizer + * @param metricsRegistry metric registry + * @return singleton instance + */ + public static SearchQueryCategorizer getInstance(MetricsRegistry metricsRegistry) { + if (instance == null) { + synchronized (SearchQueryCategorizer.class) { + if (instance == null) { + instance = new SearchQueryCategorizer(metricsRegistry); + } + } + } + return instance; + } + + /** + * Consume records and increment counters for the records + * @param records records to consume + */ + public void consumeRecords(List records) { + for (SearchQueryRecord record : records) { + SearchSourceBuilder source = (SearchSourceBuilder) record.getAttributes().get(Attribute.SOURCE); + categorize(source); + } + } + + /** + * Increment categorizations counters for the given source search query + * @param source search query source + */ + public void categorize(SearchSourceBuilder source) { + QueryBuilder topLevelQueryBuilder = source.query(); + logQueryShape(topLevelQueryBuilder); + incrementQueryTypeCounters(topLevelQueryBuilder); + incrementQueryAggregationCounters(source.aggregations()); + incrementQuerySortCounters(source.sorts()); + } + + private void incrementQuerySortCounters(List> sorts) { + if (sorts != null && sorts.size() > 0) { + for (SortBuilder sortBuilder : sorts) { + String sortOrder = sortBuilder.order().toString(); + searchQueryCounters.incrementSortCounter(1, Tags.create().addTag("sort_order", sortOrder)); + } + } + } + + private void incrementQueryAggregationCounters(AggregatorFactories.Builder aggregations) { + if (aggregations == null) { + return; + } + + searchQueryAggregationCategorizer.incrementSearchQueryAggregationCounters(aggregations.getAggregatorFactories()); + } + + private void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder) { + if (topLevelQueryBuilder == null) { + return; + } + QueryBuilderVisitor searchQueryVisitor = new SearchQueryCategorizingVisitor(searchQueryCounters); + topLevelQueryBuilder.visit(searchQueryVisitor); + } + + private void logQueryShape(QueryBuilder topLevelQueryBuilder) { + if (log.isTraceEnabled()) { + if (topLevelQueryBuilder == null) { + return; + } + QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); + topLevelQueryBuilder.visit(shapeVisitor); + log.trace("Query shape : {}", shapeVisitor.prettyPrintTree(" ")); + } + } + + /** + * Get search query counters + * @return search query counters + */ + public SearchQueryCounters getSearchQueryCounters() { + return this.searchQueryCounters; + } +} diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizingVisitor.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizingVisitor.java new file mode 100644 index 00000000..d4762e69 --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizingVisitor.java @@ -0,0 +1,39 @@ +/* + * 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.apache.lucene.search.BooleanClause; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; + +/** + * Class to visit the query builder tree and also track the level information. + * Increments the counters related to Search Query type. + */ +final class SearchQueryCategorizingVisitor implements QueryBuilderVisitor { + private final int level; + private final SearchQueryCounters searchQueryCounters; + + public SearchQueryCategorizingVisitor(SearchQueryCounters searchQueryCounters) { + this(searchQueryCounters, 0); + } + + private SearchQueryCategorizingVisitor(SearchQueryCounters counters, int level) { + this.searchQueryCounters = counters; + this.level = level; + } + + public void accept(QueryBuilder qb) { + searchQueryCounters.incrementCounter(qb, level); + } + + public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { + return new SearchQueryCategorizingVisitor(searchQueryCounters, level + 1); + } +} diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java new file mode 100644 index 00000000..7a272b61 --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java @@ -0,0 +1,134 @@ +/* + * 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.telemetry.metrics.Counter; +import org.opensearch.telemetry.metrics.MetricsRegistry; +import org.opensearch.telemetry.metrics.tags.Tags; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Class contains all the Counters related to search query types. + */ +public final class SearchQueryCounters { + private static final String LEVEL_TAG = "level"; + private static final String UNIT = "1"; + private final MetricsRegistry metricsRegistry; + /** + * Aggregation counter + */ + private final Counter aggCounter; + /** + * Counter for all other query types (catch all) + */ + private final Counter otherQueryCounter; + /** + * Counter for sort + */ + private final Counter sortCounter; + private final Map, Counter> queryHandlers; + /** + * Counter name to Counter object map + */ + private final ConcurrentHashMap nameToQueryTypeCounters; + + /** + * Constructor + * @param metricsRegistry opentelemetry metrics registry + */ + public SearchQueryCounters(MetricsRegistry metricsRegistry) { + this.metricsRegistry = metricsRegistry; + this.nameToQueryTypeCounters = new ConcurrentHashMap<>(); + this.aggCounter = metricsRegistry.createCounter( + "search.query.type.agg.count", + "Counter for the number of top level agg search queries", + UNIT + ); + this.otherQueryCounter = metricsRegistry.createCounter( + "search.query.type.other.count", + "Counter for the number of top level and nested search queries that do not match any other categories", + UNIT + ); + this.sortCounter = metricsRegistry.createCounter( + "search.query.type.sort.count", + "Counter for the number of top level sort search queries", + UNIT + ); + this.queryHandlers = new HashMap<>(); + + } + + /** + * Increment counter + * @param queryBuilder query builder + * @param level level of query builder, 0 being highest level + */ + public void incrementCounter(QueryBuilder queryBuilder, int level) { + String uniqueQueryCounterName = queryBuilder.getName(); + + Counter counter = nameToQueryTypeCounters.computeIfAbsent(uniqueQueryCounterName, k -> createQueryCounter(k)); + counter.add(1, Tags.create().addTag(LEVEL_TAG, level)); + } + + /** + * Increment aggregate counter + * @param value value to increment + * @param tags tags + */ + public void incrementAggCounter(double value, Tags tags) { + aggCounter.add(value, tags); + } + + /** + * Increment sort counter + * @param value value to increment + * @param tags tags + */ + public void incrementSortCounter(double value, Tags tags) { + sortCounter.add(value, tags); + } + + /** + * Get aggregation counter + * @return aggregation counter + */ + public Counter getAggCounter() { + return aggCounter; + } + + /** + * Get sort counter + * @return sort counter + */ + public Counter getSortCounter() { + return sortCounter; + } + + /** + * Get counter based on the query builder name + * @param queryBuilderName query builder name + * @return counter + */ + public Counter getCounterByQueryBuilderName(String queryBuilderName) { + return nameToQueryTypeCounters.get(queryBuilderName); + } + + private Counter createQueryCounter(String counterName) { + Counter counter = metricsRegistry.createCounter( + "search.query.type." + counterName + ".count", + "Counter for the number of top level and nested " + counterName + " search queries", + UNIT + ); + return counter; + } +} diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/package-info.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/package-info.java new file mode 100644 index 00000000..3bcbc7d4 --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/package-info.java @@ -0,0 +1,12 @@ +/* + * 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. + */ + +/** + * Query Insights search query categorizor to collect metrics related to search queries + */ +package org.opensearch.plugin.insights.core.service.categorizer; diff --git a/src/main/java/org/opensearch/plugin/insights/settings/QueryCategorizationSettings.java b/src/main/java/org/opensearch/plugin/insights/settings/QueryCategorizationSettings.java new file mode 100644 index 00000000..a5a65af9 --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/settings/QueryCategorizationSettings.java @@ -0,0 +1,32 @@ +/* + * 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.settings; + +import org.opensearch.common.settings.Setting; + +/** + * Settings for Query Categorization + */ +public class QueryCategorizationSettings { + /** + * Enabling search query metrics + */ + public static final Setting SEARCH_QUERY_METRICS_ENABLED_SETTING = Setting.boolSetting( + "search.query.metrics.enabled", + false, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Default constructor + */ + public QueryCategorizationSettings() {} + +} diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java index 2efe9085..0bb43da0 100644 --- a/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java +++ b/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java @@ -8,6 +8,7 @@ package org.opensearch.plugin.insights; +import org.junit.Before; import org.opensearch.action.ActionRequest; import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; @@ -18,6 +19,7 @@ import org.opensearch.plugin.insights.core.service.QueryInsightsService; import org.opensearch.plugin.insights.rules.action.top_queries.TopQueriesAction; import org.opensearch.plugin.insights.rules.resthandler.top_queries.RestTopQueriesAction; +import org.opensearch.plugin.insights.settings.QueryCategorizationSettings; import org.opensearch.plugin.insights.settings.QueryInsightsSettings; import org.opensearch.plugins.ActionPlugin; import org.opensearch.rest.RestHandler; @@ -26,7 +28,6 @@ import org.opensearch.threadpool.ExecutorBuilder; import org.opensearch.threadpool.ScalingExecutorBuilder; import org.opensearch.threadpool.ThreadPool; -import org.junit.Before; import java.util.Arrays; import java.util.List; @@ -65,7 +66,8 @@ public void testGetSettings() { QueryInsightsSettings.TOP_N_MEMORY_QUERIES_ENABLED, QueryInsightsSettings.TOP_N_MEMORY_QUERIES_SIZE, QueryInsightsSettings.TOP_N_MEMORY_QUERIES_WINDOW_SIZE, - QueryInsightsSettings.TOP_N_MEMORY_EXPORTER_SETTINGS + QueryInsightsSettings.TOP_N_MEMORY_EXPORTER_SETTINGS, + QueryCategorizationSettings.SEARCH_QUERY_METRICS_ENABLED_SETTING ), queryInsightsPlugin.getSettings() ); @@ -83,6 +85,8 @@ public void testCreateComponent() { null, null, null, + null, + null, null ); assertEquals(2, components.size()); diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java index 9b97e5f3..c1b4e54a 100644 --- a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java +++ b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java @@ -20,7 +20,9 @@ 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.plugin.insights.settings.QueryCategorizationSettings; import org.opensearch.plugin.insights.settings.QueryInsightsSettings; +import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.test.VersionUtils; import java.io.IOException; @@ -35,6 +37,8 @@ import static java.util.Collections.emptyMap; import static java.util.Collections.emptySet; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder; import static org.opensearch.test.OpenSearchTestCase.buildNewFakeTransportAddress; import static org.opensearch.test.OpenSearchTestCase.random; @@ -43,8 +47,6 @@ import static org.opensearch.test.OpenSearchTestCase.randomIntBetween; import static org.opensearch.test.OpenSearchTestCase.randomLong; import static org.opensearch.test.OpenSearchTestCase.randomLongBetween; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNull; final public class QueryInsightsTestUtils { @@ -76,9 +78,13 @@ public static List generateQueryInsightRecords(int lower, int for (int j = 0; j < countOfPhases; ++j) { phaseLatencyMap.put(randomAlphaOfLengthBetween(5, 10), randomLong()); } + + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchSourceBuilder.size(20); // Set the size parameter as needed + Map attributes = new HashMap<>(); attributes.put(Attribute.SEARCH_TYPE, SearchType.QUERY_THEN_FETCH.toString().toLowerCase(Locale.ROOT)); - attributes.put(Attribute.SOURCE, "{\"size\":20}"); + attributes.put(Attribute.SOURCE, searchSourceBuilder); attributes.put(Attribute.TOTAL_SHARDS, randomIntBetween(1, 100)); attributes.put(Attribute.INDICES, randomArray(1, 3, Object[]::new, () -> randomAlphaOfLengthBetween(5, 10))); attributes.put(Attribute.PHASE_LATENCY_MAP, phaseLatencyMap); @@ -222,5 +228,6 @@ public static void registerAllQueryInsightsSettings(ClusterSettings clusterSetti clusterSettings.registerSetting(QueryInsightsSettings.TOP_N_MEMORY_QUERIES_SIZE); clusterSettings.registerSetting(QueryInsightsSettings.TOP_N_MEMORY_QUERIES_WINDOW_SIZE); clusterSettings.registerSetting(QueryInsightsSettings.TOP_N_MEMORY_EXPORTER_SETTINGS); + clusterSettings.registerSetting(QueryCategorizationSettings.SEARCH_QUERY_METRICS_ENABLED_SETTING); } } diff --git a/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java b/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java index 5c2ea577..93876c67 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java @@ -8,6 +8,8 @@ package org.opensearch.plugin.insights.core.listener; +import org.junit.Before; +import org.mockito.ArgumentCaptor; import org.opensearch.action.search.SearchPhaseContext; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchRequestContext; @@ -36,7 +38,6 @@ import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; -import org.junit.Before; import java.util.ArrayList; import java.util.Collections; @@ -48,8 +49,6 @@ import java.util.concurrent.Phaser; import java.util.concurrent.TimeUnit; -import org.mockito.ArgumentCaptor; - import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; @@ -136,7 +135,7 @@ public void testOnRequestEnd() throws InterruptedException { assertEquals(timestamp.longValue(), generatedRecord.getTimestamp()); assertEquals(numberOfShards, generatedRecord.getAttributes().get(Attribute.TOTAL_SHARDS)); assertEquals(searchType.toString().toLowerCase(Locale.ROOT), generatedRecord.getAttributes().get(Attribute.SEARCH_TYPE)); - assertEquals(searchSourceBuilder.toString(), generatedRecord.getAttributes().get(Attribute.SOURCE)); + assertEquals(searchSourceBuilder, generatedRecord.getAttributes().get(Attribute.SOURCE)); Map labels = (Map) generatedRecord.getAttributes().get(Attribute.LABELS); assertEquals("userLabel", labels.get(Task.X_OPAQUE_ID)); } diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java index 75a5768f..60ed6529 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java @@ -8,6 +8,7 @@ package org.opensearch.plugin.insights.core.service; +import org.junit.Before; import org.opensearch.client.Client; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; @@ -15,9 +16,9 @@ import org.opensearch.plugin.insights.rules.model.MetricType; import org.opensearch.plugin.insights.rules.model.SearchQueryRecord; import org.opensearch.plugin.insights.settings.QueryInsightsSettings; +import org.opensearch.telemetry.metrics.noop.NoopMetricsRegistry; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; -import org.junit.Before; import static org.mockito.Mockito.mock; @@ -35,7 +36,7 @@ public void setup() { Settings settings = settingsBuilder.build(); ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); QueryInsightsTestUtils.registerAllQueryInsightsSettings(clusterSettings); - queryInsightsService = new QueryInsightsService(clusterSettings, threadPool, client); + queryInsightsService = new QueryInsightsService(clusterSettings, threadPool, client, NoopMetricsRegistry.INSTANCE); queryInsightsService.enableCollection(MetricType.LATENCY, true); queryInsightsService.enableCollection(MetricType.CPU, true); queryInsightsService.enableCollection(MetricType.MEMORY, true); @@ -62,4 +63,25 @@ public void testClose() { fail("No exception expected when closing query insights service"); } } + + public void testSearchQueryMetricsEnabled() { + // Initially, searchQueryMetricsEnabled should be false and searchQueryCategorizer should be null + assertFalse(queryInsightsService.isSearchQueryMetricsEnabled()); + assertNull(queryInsightsService.getSearchQueryCategorizer()); + + // Enable search query metrics + queryInsightsService.setSearchQueryMetricsEnabled(true); + + // Assert that searchQueryMetricsEnabled is true and searchQueryCategorizer is initialized + assertTrue(queryInsightsService.isSearchQueryMetricsEnabled()); + assertNotNull(queryInsightsService.getSearchQueryCategorizer()); + + // Disable search query metrics + queryInsightsService.setSearchQueryMetricsEnabled(false); + + // Assert that searchQueryMetricsEnabled is false and searchQueryCategorizer is not null + assertFalse(queryInsightsService.isSearchQueryMetricsEnabled()); + assertNotNull(queryInsightsService.getSearchQueryCategorizer()); + + } } diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeVisitorTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeVisitorTests.java new file mode 100644 index 00000000..bffe287e --- /dev/null +++ b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeVisitorTests.java @@ -0,0 +1,38 @@ +/* + * 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.categorizor; + +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.ConstantScoreQueryBuilder; +import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.RangeQueryBuilder; +import org.opensearch.index.query.RegexpQueryBuilder; +import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.index.query.TermsQueryBuilder; +import org.opensearch.plugin.insights.core.service.categorizer.QueryShapeVisitor; +import org.opensearch.test.OpenSearchTestCase; + +public final class QueryShapeVisitorTests extends OpenSearchTestCase { + public void testQueryShapeVisitor() { + QueryBuilder builder = new BoolQueryBuilder().must(new TermQueryBuilder("foo", "bar")) + .filter(new ConstantScoreQueryBuilder(new RangeQueryBuilder("timestamp").from("12345677").to("2345678"))) + .should( + new BoolQueryBuilder().must(new MatchQueryBuilder("text", "this is some text")) + .mustNot(new RegexpQueryBuilder("color", "red.*")) + ) + .must(new TermsQueryBuilder("genre", "action", "drama", "romance")); + QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); + builder.visit(shapeVisitor); + assertEquals( + "{\"type\":\"bool\",\"must\"[{\"type\":\"term\"},{\"type\":\"terms\"}],\"filter\"[{\"type\":\"constant_score\",\"filter\"[{\"type\":\"range\"}]}],\"should\"[{\"type\":\"bool\",\"must\"[{\"type\":\"match\"}],\"must_not\"[{\"type\":\"regexp\"}]}]}", + shapeVisitor.toJson() + ); + } +} diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java new file mode 100644 index 00000000..e511ef3b --- /dev/null +++ b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java @@ -0,0 +1,247 @@ +/* + * 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.categorizor; + +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.BoostingQueryBuilder; +import org.opensearch.index.query.MatchNoneQueryBuilder; +import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.index.query.MultiMatchQueryBuilder; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.index.query.QueryStringQueryBuilder; +import org.opensearch.index.query.RangeQueryBuilder; +import org.opensearch.index.query.RegexpQueryBuilder; +import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.index.query.WildcardQueryBuilder; +import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder; +import org.opensearch.plugin.insights.core.service.categorizer.SearchQueryCategorizer; +import org.opensearch.search.aggregations.bucket.range.RangeAggregationBuilder; +import org.opensearch.search.aggregations.bucket.terms.MultiTermsAggregationBuilder; +import org.opensearch.search.aggregations.support.MultiTermsValuesSourceConfig; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.sort.ScoreSortBuilder; +import org.opensearch.search.sort.SortOrder; +import org.opensearch.telemetry.metrics.Counter; +import org.opensearch.telemetry.metrics.MetricsRegistry; +import org.opensearch.telemetry.metrics.tags.Tags; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.Before; + +import java.util.Arrays; + +import org.mockito.ArgumentCaptor; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public final class SearchQueryCategorizerTests extends OpenSearchTestCase { + + private static final String MULTI_TERMS_AGGREGATION = "multi_terms"; + + private MetricsRegistry metricsRegistry; + + private SearchQueryCategorizer searchQueryCategorizer; + + @Before + public void setup() { + metricsRegistry = mock(MetricsRegistry.class); + when(metricsRegistry.createCounter(any(String.class), any(String.class), any(String.class))).thenAnswer( + invocation -> mock(Counter.class) + ); + searchQueryCategorizer = SearchQueryCategorizer.getInstance(metricsRegistry); + } + + public void testAggregationsQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.aggregation( + new MultiTermsAggregationBuilder("agg1").terms( + Arrays.asList( + new MultiTermsValuesSourceConfig.Builder().setFieldName("username").build(), + new MultiTermsValuesSourceConfig.Builder().setFieldName("rating").build() + ) + ) + ); + sourceBuilder.size(0); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getAggCounter()).add(eq(1.0d), any(Tags.class)); + + ArgumentCaptor valueCaptor = ArgumentCaptor.forClass(Double.class); + ArgumentCaptor tagsCaptor = ArgumentCaptor.forClass(Tags.class); + + verify(searchQueryCategorizer.getSearchQueryCounters().getAggCounter()).add(valueCaptor.capture(), tagsCaptor.capture()); + + double actualValue = valueCaptor.getValue(); + String actualTag = (String) tagsCaptor.getValue().getTagsMap().get("type"); + + assertEquals(1.0d, actualValue, 0.0001); + assertEquals(MULTI_TERMS_AGGREGATION, actualTag); + } + + public void testBoolQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + sourceBuilder.query(new BoolQueryBuilder().must(new MatchQueryBuilder("searchText", "fox"))); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("bool")).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("match")).add(eq(1.0d), any(Tags.class)); + } + + public void testFunctionScoreQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + sourceBuilder.query(new FunctionScoreQueryBuilder(QueryBuilders.prefixQuery("text", "bro"))); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("function_score")).add( + eq(1.0d), + any(Tags.class) + ); + } + + public void testMatchQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + sourceBuilder.query(QueryBuilders.matchQuery("tags", "php")); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("match")).add(eq(1.0d), any(Tags.class)); + } + + public void testMatchPhraseQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + sourceBuilder.query(QueryBuilders.matchPhraseQuery("tags", "php")); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("match_phrase")).add(eq(1.0d), any(Tags.class)); + } + + public void testMultiMatchQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + sourceBuilder.query(new MultiMatchQueryBuilder("foo bar", "myField")); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("multi_match")).add(eq(1.0d), any(Tags.class)); + } + + public void testOtherQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + BoostingQueryBuilder queryBuilder = new BoostingQueryBuilder( + new TermQueryBuilder("unmapped_field", "foo"), + new MatchNoneQueryBuilder() + ); + sourceBuilder.query(queryBuilder); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("boosting")).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("match_none")).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("term")).add(eq(1.0d), any(Tags.class)); + } + + public void testQueryStringQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + QueryStringQueryBuilder queryBuilder = new QueryStringQueryBuilder("foo:*"); + sourceBuilder.query(queryBuilder); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("query_string")).add(eq(1.0d), any(Tags.class)); + } + + public void testRangeQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + RangeQueryBuilder rangeQuery = new RangeQueryBuilder("date"); + rangeQuery.gte("1970-01-01"); + rangeQuery.lt("1982-01-01"); + sourceBuilder.query(rangeQuery); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("range")).add(eq(1.0d), any(Tags.class)); + } + + public void testRegexQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.query(new RegexpQueryBuilder("field", "text")); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("regexp")).add(eq(1.0d), any(Tags.class)); + } + + public void testSortQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.query(QueryBuilders.matchQuery("tags", "ruby")); + sourceBuilder.sort("creationDate", SortOrder.DESC); + sourceBuilder.sort(new ScoreSortBuilder()); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("match")).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.getSearchQueryCounters().getSortCounter(), times(2)).add(eq(1.0d), any(Tags.class)); + } + + public void testTermQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + sourceBuilder.query(QueryBuilders.termQuery("field", "value2")); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("term")).add(eq(1.0d), any(Tags.class)); + } + + public void testWildcardQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + sourceBuilder.query(new WildcardQueryBuilder("field", "text")); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("wildcard")).add(eq(1.0d), any(Tags.class)); + } + + public void testComplexQuery() { + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.size(50); + + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("field", "value2"); + MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("tags", "php"); + RegexpQueryBuilder regexpQueryBuilder = new RegexpQueryBuilder("field", "text"); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder().must(termQueryBuilder) + .filter(matchQueryBuilder) + .should(regexpQueryBuilder); + sourceBuilder.query(boolQueryBuilder); + sourceBuilder.aggregation(new RangeAggregationBuilder("agg1").field("num")); + + searchQueryCategorizer.categorize(sourceBuilder); + + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("term")).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("match")).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("regexp")).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("bool")).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.getSearchQueryCounters().getAggCounter()).add(eq(1.0d), any(Tags.class)); + } +}