Skip to content

Commit

Permalink
Address review comments
Browse files Browse the repository at this point in the history
Signed-off-by: Siddhant Deshmukh <[email protected]>
  • Loading branch information
deshsidd committed Jun 29, 2024
1 parent 8ff9e21 commit ac7fd90
Show file tree
Hide file tree
Showing 9 changed files with 66 additions and 46 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
import org.opensearch.common.settings.ClusterSettings;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.unit.TimeValue;
import org.opensearch.plugin.insights.core.categorizer.SearchQueryCategorizer;
import org.opensearch.plugin.insights.core.service.categorizer.SearchQueryCategorizer;
import org.opensearch.plugin.insights.core.exporter.QueryInsightsExporterFactory;
import org.opensearch.plugin.insights.rules.model.MetricType;
import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.categorizer;
package org.opensearch.plugin.insights.core.service.categorizer;

import org.apache.lucene.search.BooleanClause;
import org.opensearch.common.SetOnce;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.categorizer;
package org.opensearch.plugin.insights.core.service.categorizer;

import org.opensearch.search.aggregations.AggregationBuilder;
import org.opensearch.search.aggregations.PipelineAggregationBuilder;
Expand Down Expand Up @@ -43,7 +43,7 @@ public void incrementSearchQueryAggregationCounters(Collection<AggregationBuilde
private void incrementCountersRecursively(AggregationBuilder aggregationBuilder) {
// Increment counters for the current aggregation
String aggregationType = aggregationBuilder.getType();
searchQueryCounters.aggCounter.add(1, Tags.create().addTag(TYPE_TAG, aggregationType));
searchQueryCounters.incrementAggCounter(1, Tags.create().addTag(TYPE_TAG, aggregationType));

// Recursively process sub-aggregations if any
Collection<AggregationBuilder> subAggregations = aggregationBuilder.getSubAggregations();
Expand All @@ -57,7 +57,7 @@ private void incrementCountersRecursively(AggregationBuilder aggregationBuilder)
Collection<PipelineAggregationBuilder> pipelineAggregations = aggregationBuilder.getPipelineAggregations();
for (PipelineAggregationBuilder pipelineAggregation : pipelineAggregations) {
String pipelineAggregationType = pipelineAggregation.getType();
searchQueryCounters.aggCounter.add(1, Tags.create().addTag(TYPE_TAG, pipelineAggregationType));
searchQueryCounters.incrementAggCounter(1, Tags.create().addTag(TYPE_TAG, pipelineAggregationType));
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.categorizer;
package org.opensearch.plugin.insights.core.service.categorizer;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
Expand Down Expand Up @@ -34,7 +34,7 @@ public final class SearchQueryCategorizer {
/**
* Contains all the search query counters
*/
public final SearchQueryCounters searchQueryCounters;
private final SearchQueryCounters searchQueryCounters;

final SearchQueryAggregationCategorizer searchQueryAggregationCategorizer;

Expand Down Expand Up @@ -72,10 +72,9 @@ public void categorize(SearchSourceBuilder source) {

private void incrementQuerySortCounters(List<SortBuilder<?>> sorts) {
if (sorts != null && sorts.size() > 0) {
for (ListIterator<SortBuilder<?>> it = sorts.listIterator(); it.hasNext();) {
SortBuilder<?> sortBuilder = it.next();
for (SortBuilder<?> sortBuilder : sorts) {
String sortOrder = sortBuilder.order().toString();
searchQueryCounters.sortCounter.add(1, Tags.create().addTag("sort_order", sortOrder));
searchQueryCounters.incrementSortCounter(1, Tags.create().addTag("sort_order", sortOrder));
}
}
}
Expand Down Expand Up @@ -105,4 +104,7 @@ private void logQueryShape(QueryBuilder topLevelQueryBuilder) {
log.trace("Query shape : {}", shapeVisitor.prettyPrintTree(" "));
}

public SearchQueryCounters getSearchQueryCounters() {
return this.searchQueryCounters;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.categorizer;
package org.opensearch.plugin.insights.core.service.categorizer;

import org.apache.lucene.search.BooleanClause;
import org.opensearch.index.query.QueryBuilder;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.categorizer;
package org.opensearch.plugin.insights.core.service.categorizer;

import org.opensearch.index.query.QueryBuilder;
import org.opensearch.telemetry.metrics.Counter;
Expand All @@ -27,20 +27,20 @@ public final class SearchQueryCounters {
/**
* Aggregation counter
*/
public final Counter aggCounter;
private final Counter aggCounter;
/**
* Counter for all other query types (catch all)
*/
public final Counter otherQueryCounter;
private final Counter otherQueryCounter;
/**
* Counter for sort
*/
public final Counter sortCounter;
private final Counter sortCounter;
private final Map<Class<? extends QueryBuilder>, Counter> queryHandlers;
/**
* Counter name to Counter object map
*/
public final ConcurrentHashMap<String, Counter> nameToQueryTypeCounters;
private final ConcurrentHashMap<String, Counter> nameToQueryTypeCounters;

/**
* Constructor
Expand Down Expand Up @@ -80,6 +80,26 @@ public void incrementCounter(QueryBuilder queryBuilder, int level) {
counter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

public void incrementAggCounter(double value, Tags tags) {
aggCounter.add(value, tags);
}

public void incrementSortCounter(double value, Tags tags) {
sortCounter.add(value, tags);
}

public Counter getAggCounter() {
return aggCounter;
}

public Counter getSortCounter() {
return sortCounter;
}

public Counter getCounterByQueryBuilderName(String queryBuilderName) {
return nameToQueryTypeCounters.get(queryBuilderName);
}

private Counter createQueryCounter(String counterName) {
Counter counter = metricsRegistry.createCounter(
"search.query.type." + counterName + ".count",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,4 +9,4 @@
/**
* Query Insights search query categorizor to collect metrics related to search queries
*/
package org.opensearch.plugin.insights.core.categorizer;
package org.opensearch.plugin.insights.core.service.categorizer;
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.categorizor;
package org.opensearch.plugin.insights.core.service.categorizor;

import org.opensearch.index.query.BoolQueryBuilder;
import org.opensearch.index.query.ConstantScoreQueryBuilder;
Expand All @@ -16,7 +16,7 @@
import org.opensearch.index.query.RegexpQueryBuilder;
import org.opensearch.index.query.TermQueryBuilder;
import org.opensearch.index.query.TermsQueryBuilder;
import org.opensearch.plugin.insights.core.categorizer.QueryShapeVisitor;
import org.opensearch.plugin.insights.core.service.categorizer.QueryShapeVisitor;
import org.opensearch.test.OpenSearchTestCase;

public final class QueryShapeVisitorTests extends OpenSearchTestCase {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.categorizor;
package org.opensearch.plugin.insights.core.service.categorizor;

import org.opensearch.index.query.BoolQueryBuilder;
import org.opensearch.index.query.BoostingQueryBuilder;
Expand All @@ -20,7 +20,7 @@
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.categorizer.SearchQueryCategorizer;
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;
Expand Down Expand Up @@ -75,14 +75,12 @@ public void testAggregationsQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.aggCounter).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getAggCounter()).add(eq(1.0d), any(Tags.class));

// capture the arguments passed to the aggCounter.add method
ArgumentCaptor<Double> valueCaptor = ArgumentCaptor.forClass(Double.class);
ArgumentCaptor<Tags> tagsCaptor = ArgumentCaptor.forClass(Tags.class);

// Verify that aggCounter.add was called with the expected arguments
verify(searchQueryCategorizer.searchQueryCounters.aggCounter).add(valueCaptor.capture(), tagsCaptor.capture());
verify(searchQueryCategorizer.getSearchQueryCounters().getAggCounter()).add(valueCaptor.capture(), tagsCaptor.capture());

double actualValue = valueCaptor.getValue();
String actualTag = (String) tagsCaptor.getValue().getTagsMap().get("type");
Expand All @@ -98,8 +96,8 @@ public void testBoolQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("bool")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("match")).add(eq(1.0d), any(Tags.class));
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() {
Expand All @@ -109,7 +107,7 @@ public void testFunctionScoreQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("function_score")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("function_score")).add(eq(1.0d), any(Tags.class));
}

public void testMatchQuery() {
Expand All @@ -119,7 +117,7 @@ public void testMatchQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("match")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("match")).add(eq(1.0d), any(Tags.class));
}

public void testMatchPhraseQuery() {
Expand All @@ -129,7 +127,7 @@ public void testMatchPhraseQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("match_phrase")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("match_phrase")).add(eq(1.0d), any(Tags.class));
}

public void testMultiMatchQuery() {
Expand All @@ -139,7 +137,7 @@ public void testMultiMatchQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("multi_match")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("multi_match")).add(eq(1.0d), any(Tags.class));
}

public void testOtherQuery() {
Expand All @@ -153,9 +151,9 @@ public void testOtherQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("boosting")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("match_none")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("term")).add(eq(1.0d), any(Tags.class));
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() {
Expand All @@ -166,7 +164,7 @@ public void testQueryStringQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("query_string")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("query_string")).add(eq(1.0d), any(Tags.class));
}

public void testRangeQuery() {
Expand All @@ -178,7 +176,7 @@ public void testRangeQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("range")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("range")).add(eq(1.0d), any(Tags.class));
}

public void testRegexQuery() {
Expand All @@ -187,7 +185,7 @@ public void testRegexQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("regexp")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("regexp")).add(eq(1.0d), any(Tags.class));
}

public void testSortQuery() {
Expand All @@ -198,8 +196,8 @@ public void testSortQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("match")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.searchQueryCounters.sortCounter, times(2)).add(eq(1.0d), any(Tags.class));
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() {
Expand All @@ -209,7 +207,7 @@ public void testTermQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("term")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("term")).add(eq(1.0d), any(Tags.class));
}

public void testWildcardQuery() {
Expand All @@ -219,7 +217,7 @@ public void testWildcardQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("wildcard")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.getSearchQueryCounters().getCounterByQueryBuilderName("wildcard")).add(eq(1.0d), any(Tags.class));
}

public void testComplexQuery() {
Expand All @@ -237,10 +235,10 @@ public void testComplexQuery() {

searchQueryCategorizer.categorize(sourceBuilder);

verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("term")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("match")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("regexp")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.searchQueryCounters.nameToQueryTypeCounters.get("bool")).add(eq(1.0d), any(Tags.class));
verify(searchQueryCategorizer.searchQueryCounters.aggCounter).add(eq(1.0d), any(Tags.class));
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));
}
}

0 comments on commit ac7fd90

Please sign in to comment.