fix flaky integration tests

Signed-off-by: Anthony Leong <[email protected]>

Author: ajleong623

server/src/main/java/org/opensearch/index/query/IntervalBuilder.java

@@ -144,6 +144,20 @@ protected IntervalsSource analyzeTerm(TokenStream ts) throws IOException {
         return Intervals.term(BytesRef.deepCopyOf(bytesAtt.getBytesRef()));
     }
 
+    static boolean canCombineSources(List<IntervalsSource> sources) {
+        int sourceIndex = 0;
+        long disjunctionCount = 1;
+
+        while (sourceIndex < sources.size()) {
+            disjunctionCount = disjunctionCount * sources.get(sourceIndex).pullUpDisjunctions().size();
+            if (disjunctionCount > IndexSearcher.getMaxClauseCount()) {
+                return false;
+            }
+            sourceIndex += 1;
+        }
+        return true;
+    }
+
     protected static IntervalsSource combineSources(List<IntervalsSource> sources, int maxGaps, IntervalMode mode) {
         if (sources.size() == 0) {
             return NO_INTERVALS;

server/src/main/java/org/opensearch/index/query/IntervalsSourceProvider.java

@@ -443,6 +443,9 @@ public IntervalsSource getSource(QueryShardContext ctx, MappedFieldType fieldTyp
             for (IntervalsSourceProvider provider : subSources) {
                 ss.add(provider.getSource(ctx, fieldType));
             }
+            if (maxGaps == 0 && mode == IntervalMode.ORDERED && IntervalBuilder.canCombineSources(ss) == false) {
+                throw new IllegalArgumentException("Too many disjunctions to expand");
+            }
             IntervalsSource source = IntervalBuilder.combineSources(ss, maxGaps, mode);
             if (filter != null) {
                 return filter.filter(source, ctx, fieldType);

test/framework/src/main/java/org/opensearch/test/AbstractQueryTestCase.java

@@ -448,7 +448,13 @@ public void testToQuery() throws IOException {
              * We do it this way in SearchService where
              * we first rewrite the query with a private context, then reset the context and then build the actual lucene query*/
             QueryBuilder rewritten = rewriteQuery(firstQuery, new QueryShardContext(context));
-            Query firstLuceneQuery = rewritten.toQuery(context);
+            Query firstLuceneQuery;
+            try {
+                firstLuceneQuery = rewritten.toQuery(context);
+            } catch (IllegalArgumentException e) {
+                assertEquals("Too many disjunctions to expand", e.getMessage());
+                continue;
+            }
             assertNotNull("toQuery should not return null", firstLuceneQuery);
             assertLuceneQuery(firstQuery, firstLuceneQuery, context);
             // remove after assertLuceneQuery since the assertLuceneQuery impl might access the context as well