Add exception handling

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
deshsidd · Oct 16, 2023 · b148f88 · b148f88
1 parent 8632615
commit b148f88
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java
@@ -55,12 +55,20 @@ private void incrementQueryAggregationCounters(AggregatorFactories.Builder aggre
         }
     }
 
-    private static void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder) {
-        QueryBuilderVisitor searhQueryVisitor = new SearchQueryCategorizingVisitor(searchQueryCounters);
-        topLevelQueryBuilder.visit(searhQueryVisitor);
+    private void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder) {
+        if (topLevelQueryBuilder == null) {
+            searchQueryCounters.skippedCounter.add(1);
+            return;
+        }
+        QueryBuilderVisitor searchQueryVisitor = new SearchQueryCategorizingVisitor(searchQueryCounters);
+        topLevelQueryBuilder.visit(searchQueryVisitor);
     }
 
-    private static void logQueryShape(QueryBuilder topLevelQueryBuilder) {
+    private void logQueryShape(QueryBuilder topLevelQueryBuilder) {
+        if (topLevelQueryBuilder == null) {
+            searchQueryCounters.skippedCounter.add(1);
+            return;
+        }
         QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
         topLevelQueryBuilder.visit(shapeVisitor);
         String queryShapeJson = shapeVisitor.prettyPrintTree("  ");

diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java
@@ -30,6 +30,7 @@ public class SearchQueryCounters {
     public final Counter regexCounter;
 
     public final Counter sortCounter;
+    public final Counter skippedCounter;
     public final Counter termCounter;
     public final Counter totalCounter;
     public final Counter wildcardCounter;
@@ -86,6 +87,11 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) {
             "Counter for the number of top level and nested regex search queries",
             "0"
         );
+        this.skippedCounter = metricsRegistry.createCounter(
+            "skippedQueryCounter",
+            "Counter for the number queries skipped due to error",
+            "0"
+        );
         this.sortCounter = metricsRegistry.createCounter(
             "sortSearchQueryCounter",
             "Counter for the number of top level sort search queries",

diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -459,7 +459,11 @@ private void executeRequest(
         }
 
         if (isSearchQueryCategorizationEnabled) {
-            searchQueryCategorizor.categorize(searchRequest.source());
+            try {
+                searchQueryCategorizor.categorize(searchRequest.source());
+            } catch (Exception e) {
+                logger.error("Error while trying to categorize the query.");
+            }
         }
 
         ActionListener<SearchSourceBuilder> rewriteListener = ActionListener.wrap(source -> {