Merge pull request #20241 from vespa-engine/olaa/add-query-item-limit

Add query item limit

2 files changed, 24 insertions, 2 deletions

diff --git a/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java b/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
index 8c4c5c84a28..a93dd1b9de4 100644
--- a/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
+++ b/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.prelude.query;
 
+import com.yahoo.processing.request.CompoundName;
 import com.yahoo.search.Query;
 import com.yahoo.search.query.QueryTree;
 
@@ -19,6 +20,8 @@ public class QueryCanonicalizer {
     /** The name of the operation performed by this, for use in search chain ordering */
     public static final String queryCanonicalization = "queryCanonicalization";
 
+    private static final CompoundName MAX_QUERY_ITEMS = new CompoundName("maxQueryItems");
+
     /**
      * Validates this query and carries out possible operations on this query
      * which simplifies it without changing its semantics.
@@ -26,7 +29,17 @@ public class QueryCanonicalizer {
      * @return null if the query is valid, an error message if it is invalid
      */
     public static String canonicalize(Query query) {
-        return canonicalize(query.getModel().getQueryTree());
+        Integer maxQueryItems = query.properties().getInteger(MAX_QUERY_ITEMS, Integer.MAX_VALUE);
+        return canonicalize(query.getModel().getQueryTree(), maxQueryItems);
+    }
+
+    /**
+     * Canonicalizes this query, allowing any query tree size
+     *
+     * @return null if the query is valid, an error message if it is invalid
+     */
+    public static String canonicalize(QueryTree queryTree) {
+        return canonicalize(queryTree, Integer.MAX_VALUE);
     }
 
     /**
@@ -34,10 +47,12 @@ public class QueryCanonicalizer {
      * 
      * @return null if the query is valid, an error message if it is invalid
      */
-    public static String canonicalize(QueryTree query) {
+    private static String canonicalize(QueryTree query, Integer maxQueryItems) {
         ListIterator<Item> rootItemIterator = query.getItemIterator();
         CanonicalizationResult result = recursivelyCanonicalize(rootItemIterator.next(), rootItemIterator);
         if (query.isEmpty() && ! result.isError()) result = CanonicalizationResult.error("No query");
+        int itemCount = query.treeSize();
+        if (itemCount > maxQueryItems) result = CanonicalizationResult.error(String.format("Query tree exceeds allowed item count. Configured limit: %d - Item count: %d", maxQueryItems, itemCount));
         return result.error().orElse(null); // preserve old API, unfortunately
     }
 
diff --git a/container-search/src/main/java/com/yahoo/prelude/statistics/StatisticsSearcher.java b/container-search/src/main/java/com/yahoo/prelude/statistics/StatisticsSearcher.java
index 1d1c446f6e1..6fae5c97cd2 100644
--- a/container-search/src/main/java/com/yahoo/prelude/statistics/StatisticsSearcher.java
+++ b/container-search/src/main/java/com/yahoo/prelude/statistics/StatisticsSearcher.java
@@ -72,6 +72,7 @@ public class StatisticsSearcher extends Searcher {
     private static final String RELEVANCE_AT_1_METRIC = "relevance.at_1";
     private static final String RELEVANCE_AT_3_METRIC = "relevance.at_3";
     private static final String RELEVANCE_AT_10_METRIC = "relevance.at_10";
+    private static final String QUERY_ITEM_COUNT = "query_item_count";
 
     @SuppressWarnings("unused") // all the work is done by the callback
     private final PeakQpsReporter peakQpsReporter;
@@ -264,6 +265,8 @@ public class StatisticsSearcher extends Searcher {
 
         addRelevanceMetrics(query, execution, result);
 
+        addItemCountMetric(query, metricContext);
+
         return result;
     }
 
@@ -396,6 +399,10 @@ public class StatisticsSearcher extends Searcher {
         }
     }
 
+    private void addItemCountMetric(Query query, Metric.Context context) {
+        metric.set(QUERY_ITEM_COUNT, query.getModel().getQueryTree().treeSize(), context);
+    }
+
     /**
      * Returns the relative start time from request was received by jdisc
      */