Address review comments

Signed-off-by: Siddhant Deshmukh <[email protected]>
deshsidd · Jun 28, 2024 · 093a791 · 093a791
1 parent 9836921
commit 093a791
Show file tree

Hide file tree

Showing 9 changed files with 66 additions and 46 deletions.
diff --git a/...ights/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java b/...ights/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java
@@ -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;

diff --git a/...s/core/categorizer/QueryShapeVisitor.java → ...ervice/categorizer/QueryShapeVisitor.java b/...s/core/categorizer/QueryShapeVisitor.java → ...ervice/categorizer/QueryShapeVisitor.java
@@ -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;

diff --git a/...er/SearchQueryAggregationCategorizer.java → ...er/SearchQueryAggregationCategorizer.java b/...er/SearchQueryAggregationCategorizer.java → ...er/SearchQueryAggregationCategorizer.java
@@ -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;
@@ -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();
@@ -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));
         }
     }
 }
diff --git a/...e/categorizer/SearchQueryCategorizer.java → ...e/categorizer/SearchQueryCategorizer.java b/...e/categorizer/SearchQueryCategorizer.java → ...e/categorizer/SearchQueryCategorizer.java
@@ -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;
@@ -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;
 
@@ -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));
             }
         }
     }
@@ -105,4 +104,7 @@ private void logQueryShape(QueryBuilder topLevelQueryBuilder) {
         log.trace("Query shape : {}", shapeVisitor.prettyPrintTree("  "));
     }
 
+    public SearchQueryCounters getSearchQueryCounters() {
+        return this.searchQueryCounters;
+    }
 }
diff --git a/...rizer/SearchQueryCategorizingVisitor.java → ...rizer/SearchQueryCategorizingVisitor.java b/...rizer/SearchQueryCategorizingVisitor.java → ...rizer/SearchQueryCategorizingVisitor.java
@@ -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;

diff --git a/...core/categorizer/SearchQueryCounters.java → ...vice/categorizer/SearchQueryCounters.java b/...core/categorizer/SearchQueryCounters.java → ...vice/categorizer/SearchQueryCounters.java
@@ -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;
@@ -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
@@ -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",

diff --git a/...sights/core/categorizer/package-info.java → ...ore/service/categorizer/package-info.java b/...sights/core/categorizer/package-info.java → ...ore/service/categorizer/package-info.java
@@ -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;
diff --git a/...e/categorizor/QueryShapeVisitorTests.java → ...e/categorizor/QueryShapeVisitorTests.java b/...e/categorizor/QueryShapeVisitorTests.java → ...e/categorizor/QueryShapeVisitorTests.java
@@ -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;
@@ -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 {

diff --git a/...egorizor/SearchQueryCategorizerTests.java → ...egorizor/SearchQueryCategorizerTests.java b/...egorizor/SearchQueryCategorizerTests.java → ...egorizor/SearchQueryCategorizerTests.java
@@ -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;
@@ -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;
@@ -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");
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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() {
@@ -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));
     }
 }