ESQL: Extend the check on partial results to all REST ITs

x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/RequestIndexFilteringIT.java

@@ -38,6 +38,7 @@
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
 
 @ThreadLeakFilters(filters = TestClustersThreadFilter.class)
 public class RequestIndexFilteringIT extends RequestIndexFilteringTestCase {
@@ -91,14 +92,20 @@ protected String from(String... indexName) {
 
     @Override
     public Map<String, Object> runEsql(RestEsqlTestCase.RequestObjectBuilder requestObject) throws IOException {
+        return runEsql(requestObject, true);
+    }
+
+    @Override
+    public Map<String, Object> runEsql(RestEsqlTestCase.RequestObjectBuilder requestObject, boolean checkPartialResults)
+        throws IOException {
         if (requestObject.allowPartialResults() != null) {
             assumeTrue(
                 "require allow_partial_results on local cluster",
                 clusterHasCapability("POST", "/_query", List.of(), List.of("support_partial_results")).orElse(false)
             );
         }
         requestObject.includeCCSMetadata(true);
-        return super.runEsql(requestObject);
+        return super.runEsql(requestObject, checkPartialResults);
     }
 
     @After
@@ -154,8 +161,13 @@ public void testIndicesDontExistRemote() throws IOException {
         indexTimestampData(docsTest1, "test1", "2024-11-26", "id1");
 
         Map<String, Object> result = runEsql(
-            timestampFilter("gte", "2020-01-01").query("FROM *:foo,*:test1 METADATA _index | SORT id1 | KEEP _index, id*")
+            timestampFilter("gte", "2020-01-01").query("FROM *:foo,*:test1 METADATA _index | SORT id1 | KEEP _index, id*"),
+            false
         );
+
+        // `foo` index doesn't exist, so the request will currently be successful, but with partial results
+        var isPartial = result.get("is_partial");
+        assertThat(isPartial, is(true));
         @SuppressWarnings("unchecked")
         var columns = (List<List<Object>>) result.get("columns");
         assertThat(

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/EsqlSpecTestCase.java

@@ -276,9 +276,7 @@ protected final void doTest(String query) throws Throwable {
         Map<?, ?> prevTooks = supportsTook() ? tooks() : null;
         Map<String, Object> answer = runEsql(builder.query(query), testCase.assertWarnings(deduplicateExactWarnings()));
 
-        var clusters = answer.get("_clusters");
-        var reason = "unexpected partial results" + (clusters != null ? ": _clusters=" + clusters : "");
-        assertThat(reason, answer.get("is_partial"), anyOf(nullValue(), is(false)));
+        assertNotPartial(answer);
 
         var expectedColumnsWithValues = loadCsvSpecValues(testCase.expectedResults);
 
@@ -303,6 +301,14 @@ protected final void doTest(String query) throws Throwable {
         }
     }
 
+    static Map<String, Object> assertNotPartial(Map<String, Object> answer) {
+        var clusters = answer.get("_clusters");
+        var reason = "unexpected partial results" + (clusters != null ? ": _clusters=" + clusters : "");
+        assertThat(reason, answer.get("is_partial"), anyOf(nullValue(), is(false)));
+
+        return answer;
+    }
+
     private Map<?, ?> tooks() throws IOException {
         Request request = new Request("GET", "/_xpack/usage");
         HttpEntity entity = client().performRequest(request).getEntity();

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RequestIndexFilteringTestCase.java

@@ -249,6 +249,11 @@ public Map<String, Object> runEsql(RestEsqlTestCase.RequestObjectBuilder request
         return RestEsqlTestCase.runEsql(requestObject, new AssertWarnings.NoWarnings(), RestEsqlTestCase.Mode.SYNC);
     }
 
+    public Map<String, Object> runEsql(RestEsqlTestCase.RequestObjectBuilder requestObject, boolean checkPartialResults)
+        throws IOException {
+        return RestEsqlTestCase.runEsql(requestObject, new AssertWarnings.NoWarnings(), RestEsqlTestCase.Mode.SYNC, checkPartialResults);
+    }
+
     protected void indexTimestampData(int docs, String indexName, String date, String differentiatorFieldName) throws IOException {
         indexTimestampDataForClient(client(), docs, indexName, date, differentiatorFieldName);
     }

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java

@@ -60,6 +60,7 @@
 import static org.elasticsearch.test.MapMatcher.assertMap;
 import static org.elasticsearch.test.MapMatcher.matchesMap;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
+import static org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase.assertNotPartial;
 import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode.ASYNC;
 import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode.SYNC;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateTimeToString;
@@ -1248,13 +1249,21 @@ public static Map<String, Object> runEsqlAsync(RequestObjectBuilder requestObjec
         return runEsqlAsync(requestObject, randomBoolean(), new AssertWarnings.NoWarnings());
     }
 
+    public static Map<String, Object> runEsql(
+        RequestObjectBuilder requestObject,
+        AssertWarnings assertWarnings,
+        Mode mode,
+        boolean checkPartialResults
+    ) throws IOException {
+        var results = mode == ASYNC
+            ? runEsqlAsync(requestObject, randomBoolean(), assertWarnings)
+            : runEsqlSync(requestObject, assertWarnings);
+        return checkPartialResults ? assertNotPartial(results) : results;
+    }
+
     public static Map<String, Object> runEsql(RequestObjectBuilder requestObject, AssertWarnings assertWarnings, Mode mode)
         throws IOException {
-        if (mode == ASYNC) {
-            return runEsqlAsync(requestObject, randomBoolean(), assertWarnings);
-        } else {
-            return runEsqlSync(requestObject, assertWarnings);
-        }
+        return runEsql(requestObject, assertWarnings, mode, true);
     }
 
     public static Map<String, Object> runEsqlSync(RequestObjectBuilder requestObject, AssertWarnings assertWarnings) throws IOException {
@@ -1515,9 +1524,16 @@ protected static Map<String, Object> entityToMap(HttpEntity entity, XContentType
         if (shouldLog()) {
             LOGGER.info("entity={}", result);
         }
+
         return result;
     }
 
+    protected static Map<String, Object> entityToMap(HttpEntity entity, XContentType expectedContentType, boolean allowPartialResults)
+        throws IOException {
+        var response = EsqlTestUtils.entityToMap(entity, expectedContentType);
+        return allowPartialResults ? response : assertNotPartial(response);
+    }
+
     static void addAsyncParameters(RequestObjectBuilder requestObject, boolean keepOnCompletion) throws IOException {
         // deliberately short in order to frequently trigger return without results
         requestObject.waitForCompletion(TimeValue.timeValueNanos(randomIntBetween(1, 100)));