diff --git a/src/main/java/org/opensearch/plugin/insights/QueryInsightsPlugin.java b/src/main/java/org/opensearch/plugin/insights/QueryInsightsPlugin.java index bba6764..ba33bbb 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 8f18676..c54f1ad 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 @@ -91,26 +91,24 @@ public QueryInsightsListener(final ClusterService clusterService, final QueryIns * and query insights services. * * @param metricType {@link MetricType} - * @param enabled boolean + * @param isCurrentMetricEnabled boolean */ - public void setEnableTopQueries(final MetricType metricType, final boolean enabled) { - boolean isAllMetricsDisabled = !queryInsightsService.isEnabled(); - this.queryInsightsService.enableCollection(metricType, enabled); - if (!enabled) { - // disable QueryInsightsListener only if all metrics collections are disabled now. - if (!queryInsightsService.isEnabled()) { - super.setEnabled(false); - this.queryInsightsService.stop(); + public void setEnableTopQueries(final MetricType metricType, final boolean isCurrentMetricEnabled) { + boolean isTopNFeaturePreviouslyDisabled = !queryInsightsService.isTopNFeatureEnabled(); + this.queryInsightsService.enableCollection(metricType, isCurrentMetricEnabled); + boolean isTopNFeatureCurrentlyDisabled = !queryInsightsService.isTopNFeatureEnabled(); + + if (isTopNFeatureCurrentlyDisabled) { + super.setEnabled(false); + if (!isTopNFeaturePreviouslyDisabled) { + queryInsightsService.checkAndStopQueryInsights(); } } else { super.setEnabled(true); - // restart QueryInsightsListener only if none of metrics collections is enabled before. - if (isAllMetricsDisabled) { - this.queryInsightsService.stop(); - this.queryInsightsService.start(); + if (isTopNFeaturePreviouslyDisabled) { + queryInsightsService.checkAndRestartQueryInsights(); } } - } @Override @@ -176,7 +174,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 33343aa..6a36c9a 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,11 +196,20 @@ 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 */ - public boolean isEnabled() { + public boolean isAnyFeatureEnabled() { + return isTopNFeatureEnabled() || isSearchQueryMetricsFeatureEnabled(); + } + + /** + * Check if top N enabled for any metric type + * + * @return if top N feature is enabled + */ + public boolean isTopNFeatureEnabled() { for (MetricType t : MetricType.allMetricTypes()) { if (isCollectionEnabled(t)) { return true; @@ -179,6 +218,33 @@ public boolean isEnabled() { return false; } + /** + * Is search query metrics feature enabled. + * @return boolean flag + */ + public boolean isSearchQueryMetricsFeatureEnabled() { + return this.searchQueryMetricsEnabled; + } + + /** + * Stops query insights service if no features enabled + */ + public void checkAndStopQueryInsights() { + if (!isAnyFeatureEnabled()) { + this.stop(); + } + } + + /** + * Restarts query insights service if any feature enabled + */ + public void checkAndRestartQueryInsights() { + if (isAnyFeatureEnabled()) { + this.stop(); + this.start(); + } + } + /** * Validate the window size config for a metricType * @@ -239,6 +305,32 @@ 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) { + boolean oldSearchQueryMetricsEnabled = isSearchQueryMetricsFeatureEnabled(); + this.searchQueryMetricsEnabled = searchQueryMetricsEnabled; + if (searchQueryMetricsEnabled) { + if (!oldSearchQueryMetricsEnabled) { + checkAndRestartQueryInsights(); + } + } else { + if (oldSearchQueryMetricsEnabled) { + checkAndStopQueryInsights(); + } + } + } + + /** + * Get search query categorizer object + * @return SearchQueryCategorizer object + */ + public SearchQueryCategorizer getSearchQueryCategorizer() { + return this.searchQueryCategorizer; + } + /** * Validate the exporter config for a metricType * @@ -253,7 +345,7 @@ public void validateExporterConfig(final MetricType type, final Settings setting @Override protected void doStart() { - if (isEnabled()) { + if (isAnyFeatureEnabled()) { scheduledFuture = threadPool.scheduleWithFixedDelay( this::drainRecords, QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL, 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 0000000..626e05d --- /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 0000000..f252795 --- /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 0000000..b941ed7 --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java @@ -0,0 +1,139 @@ +/* + * 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; + } + + /** + * Reset the search query categorizer and its counters + */ + public void reset() { + instance = null; + } +} 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 0000000..d4762e6 --- /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 0000000..7a272b6 --- /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 0000000..3bcbc7d --- /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/rules/model/Attribute.java b/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java index 9d9f3db..43d7f61 100644 --- a/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java +++ b/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java @@ -13,6 +13,7 @@ import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo; +import org.opensearch.search.builder.SearchSourceBuilder; import java.io.IOException; import java.util.Collections; @@ -90,6 +91,8 @@ static void writeTo(final StreamOutput out, final Attribute attribute) throws IO public static void writeValueTo(StreamOutput out, Object attributeValue) throws IOException { if (attributeValue instanceof List) { out.writeList((List) attributeValue); + } else if (attributeValue instanceof SearchSourceBuilder) { + ((SearchSourceBuilder) attributeValue).writeTo(out); } else { out.writeGenericValue(attributeValue); } @@ -106,6 +109,9 @@ public static void writeValueTo(StreamOutput out, Object attributeValue) throws public static Object readAttributeValue(StreamInput in, Attribute attribute) throws IOException { if (attribute == Attribute.TASK_RESOURCE_USAGES) { return in.readList(TaskResourceInfo::readFromStream); + } else if (attribute == Attribute.SOURCE) { + SearchSourceBuilder builder = new SearchSourceBuilder(in); + return builder; } else { return in.readGenericValue(); } 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 0000000..a5a65af --- /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 2efe908..443ca8b 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,20 +19,24 @@ 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; +import org.opensearch.telemetry.metrics.Counter; +import org.opensearch.telemetry.metrics.MetricsRegistry; import org.opensearch.test.ClusterServiceUtils; import org.opensearch.test.OpenSearchTestCase; 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; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class QueryInsightsPluginTests extends OpenSearchTestCase { @@ -40,6 +45,7 @@ public class QueryInsightsPluginTests extends OpenSearchTestCase { private final Client client = mock(Client.class); private ClusterService clusterService; private final ThreadPool threadPool = mock(ThreadPool.class); + private MetricsRegistry metricsRegistry = mock(MetricsRegistry.class); @Before public void setup() { @@ -49,6 +55,10 @@ public void setup() { ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); QueryInsightsTestUtils.registerAllQueryInsightsSettings(clusterSettings); clusterService = ClusterServiceUtils.createClusterService(settings, clusterSettings, threadPool); + + when(metricsRegistry.createCounter(any(String.class), any(String.class), any(String.class))).thenAnswer( + invocation -> mock(Counter.class) + ); } public void testGetSettings() { @@ -65,7 +75,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,7 +94,9 @@ public void testCreateComponent() { null, null, null, - null + null, + null, + metricsRegistry ); assertEquals(2, components.size()); assertTrue(components.get(0) instanceof QueryInsightsService); diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java index 9b97e5f..c1b4e54 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 5c2ea57..6db7ffe 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,12 +8,17 @@ package org.opensearch.plugin.insights.core.listener; +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.mockito.MockitoAnnotations; +import org.mockito.MockitoSession; import org.opensearch.action.search.SearchPhaseContext; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchRequestContext; import org.opensearch.action.search.SearchTask; import org.opensearch.action.search.SearchType; import org.opensearch.action.support.replication.ClusterStateCreationUtils; +import org.opensearch.client.Client; import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.collect.Tuple; @@ -32,11 +37,12 @@ import org.opensearch.search.aggregations.support.ValueType; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.tasks.Task; +import org.opensearch.telemetry.metrics.MetricsRegistry; +import org.opensearch.telemetry.metrics.tags.Tags; import org.opensearch.test.ClusterServiceUtils; 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,10 +54,11 @@ import java.util.concurrent.Phaser; import java.util.concurrent.TimeUnit; -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.never; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -136,7 +143,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)); } @@ -202,16 +209,62 @@ public void testConcurrentOnRequestEnd() throws InterruptedException { verify(queryInsightsService, times(numRequests)).addRecord(any()); } - public void testSetEnabled() { - when(queryInsightsService.isCollectionEnabled(MetricType.LATENCY)).thenReturn(true); - QueryInsightsListener queryInsightsListener = new QueryInsightsListener(clusterService, queryInsightsService); - queryInsightsListener.setEnableTopQueries(MetricType.LATENCY, true); - assertTrue(queryInsightsListener.isEnabled()); - - when(queryInsightsService.isCollectionEnabled(MetricType.LATENCY)).thenReturn(false); - when(queryInsightsService.isCollectionEnabled(MetricType.CPU)).thenReturn(false); - when(queryInsightsService.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); - queryInsightsListener.setEnableTopQueries(MetricType.LATENCY, false); - assertFalse(queryInsightsListener.isEnabled()); + public void testTopNFeatureEnabledDisabled() { + // Test case 1: Only latency enabled initially, disable latency and verify expected behavior + QueryInsightsService queryInsightsService1 = mock(QueryInsightsService.class); + QueryInsightsListener queryInsightsListener1 = new QueryInsightsListener(clusterService, queryInsightsService1); + + when(queryInsightsService1.isCollectionEnabled(MetricType.LATENCY)).thenReturn(true); + when(queryInsightsService1.isCollectionEnabled(MetricType.CPU)).thenReturn(false); + when(queryInsightsService1.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); + when(queryInsightsService1.isTopNFeatureEnabled()).thenReturn(true).thenReturn(false); + + queryInsightsListener1.setEnableTopQueries(MetricType.LATENCY, false); + assertFalse(queryInsightsListener1.isEnabled()); + verify(queryInsightsService1).checkAndStopQueryInsights(); + verify(queryInsightsService1, never()).checkAndRestartQueryInsights(); + + // Test case 2: All disabled initially, enable latency and verify expected behavior + QueryInsightsService queryInsightsService2 = mock(QueryInsightsService.class); + QueryInsightsListener queryInsightsListener2 = new QueryInsightsListener(clusterService, queryInsightsService2); + + when(queryInsightsService2.isCollectionEnabled(MetricType.LATENCY)).thenReturn(false); + when(queryInsightsService2.isCollectionEnabled(MetricType.CPU)).thenReturn(false); + when(queryInsightsService2.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); + when(queryInsightsService2.isTopNFeatureEnabled()).thenReturn(false).thenReturn(true); + + queryInsightsListener2.setEnableTopQueries(MetricType.LATENCY, true); + assertTrue(queryInsightsListener2.isEnabled()); + verify(queryInsightsService2).checkAndRestartQueryInsights(); + verify(queryInsightsService2, never()).checkAndStopQueryInsights(); + + // Test case 3: latency and CPU enabled initially, disable latency and verify expected behavior + QueryInsightsService queryInsightsService3 = mock(QueryInsightsService.class); + QueryInsightsListener queryInsightsListener3 = new QueryInsightsListener(clusterService, queryInsightsService3); + + when(queryInsightsService3.isCollectionEnabled(MetricType.LATENCY)).thenReturn(true); + when(queryInsightsService3.isCollectionEnabled(MetricType.CPU)).thenReturn(true); + when(queryInsightsService3.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); + when(queryInsightsService3.isTopNFeatureEnabled()).thenReturn(true).thenReturn(true); + + queryInsightsListener3.setEnableTopQueries(MetricType.LATENCY, false); + assertTrue(queryInsightsListener3.isEnabled()); + verify(queryInsightsService3, never()).checkAndStopQueryInsights(); + verify(queryInsightsService3, never()).checkAndRestartQueryInsights(); + + + // Test case 4: Only CPU enabled initially, enable latency and verify expected behavior + QueryInsightsService queryInsightsService4 = mock(QueryInsightsService.class); + QueryInsightsListener queryInsightsListener4 = new QueryInsightsListener(clusterService, queryInsightsService4); + + when(queryInsightsService4.isCollectionEnabled(MetricType.LATENCY)).thenReturn(false); + when(queryInsightsService4.isCollectionEnabled(MetricType.CPU)).thenReturn(true); + when(queryInsightsService4.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); + when(queryInsightsService4.isTopNFeatureEnabled()).thenReturn(true).thenReturn(true); + + queryInsightsListener4.setEnableTopQueries(MetricType.LATENCY, true); + assertTrue(queryInsightsListener4.isEnabled()); + verify(queryInsightsService4, never()).checkAndRestartQueryInsights(); + verify(queryInsightsService4, never()).checkAndStopQueryInsights(); } } 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 75a5768..cd7a7a7 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,11 +16,15 @@ 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; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; /** * Unit Tests for {@link QueryInsightsService}. @@ -28,6 +33,7 @@ public class QueryInsightsServiceTests extends OpenSearchTestCase { private final ThreadPool threadPool = mock(ThreadPool.class); private final Client client = mock(Client.class); private QueryInsightsService queryInsightsService; + private QueryInsightsService queryInsightsServiceSpy; @Before public void setup() { @@ -35,10 +41,11 @@ 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); + queryInsightsServiceSpy = spy(queryInsightsService); } public void testAddRecordToLimitAndDrain() { @@ -62,4 +69,45 @@ 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.isSearchQueryMetricsFeatureEnabled()); + assertNotNull(queryInsightsService.getSearchQueryCategorizer()); + + // Enable search query metrics + queryInsightsService.setSearchQueryMetricsEnabled(true); + + // Assert that searchQueryMetricsEnabled is true and searchQueryCategorizer is initialized + assertTrue(queryInsightsService.isSearchQueryMetricsFeatureEnabled()); + assertNotNull(queryInsightsService.getSearchQueryCategorizer()); + + // Disable search query metrics + queryInsightsService.setSearchQueryMetricsEnabled(false); + + // Assert that searchQueryMetricsEnabled is false and searchQueryCategorizer is not null + assertFalse(queryInsightsService.isSearchQueryMetricsFeatureEnabled()); + assertNotNull(queryInsightsService.getSearchQueryCategorizer()); + + } + + public void testFeaturesEnableDisable() { + // Test case 1: All metric type collection disabled and search query metrics disabled, enable search query metrics + queryInsightsServiceSpy.enableCollection(MetricType.LATENCY, false); + queryInsightsServiceSpy.enableCollection(MetricType.CPU, false); + queryInsightsServiceSpy.enableCollection(MetricType.MEMORY, false); + queryInsightsServiceSpy.setSearchQueryMetricsEnabled(false); + + queryInsightsServiceSpy.setSearchQueryMetricsEnabled(true); + verify(queryInsightsServiceSpy).checkAndRestartQueryInsights(); + + // Test case 2: All metric type collection disabled and search query metrics enabled, disable search query metrics + queryInsightsServiceSpy.enableCollection(MetricType.LATENCY, false); + queryInsightsServiceSpy.enableCollection(MetricType.CPU, false); + queryInsightsServiceSpy.enableCollection(MetricType.MEMORY, false); + queryInsightsServiceSpy.setSearchQueryMetricsEnabled(true); + + queryInsightsServiceSpy.setSearchQueryMetricsEnabled(false); + verify(queryInsightsServiceSpy).checkAndStopQueryInsights(); + } } 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 0000000..bffe287 --- /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 0000000..643b514 --- /dev/null +++ b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java @@ -0,0 +1,254 @@ +/* + * 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.junit.After; +import org.mockito.Mockito; +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); + } + + @After + public void cleanup() { + searchQueryCategorizer.reset(); + } + + 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)); + } +}