adapt to CCS test with partial results

Author: bpintea

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

@@ -37,6 +37,7 @@
 import static org.elasticsearch.test.MapMatcher.matchesMap;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.instanceOf;
 
 @ThreadLeakFilters(filters = TestClustersThreadFilter.class)
@@ -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/FieldExtractorTestCase.java

@@ -52,7 +52,7 @@
 import static org.elasticsearch.test.ListMatcher.matchesList;
 import static org.elasticsearch.test.MapMatcher.assertMap;
 import static org.elasticsearch.test.MapMatcher.matchesMap;
-import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.entityToMapNoPartialCheck;
+import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.entityToMap;
 import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.runEsqlSync;
 import static org.hamcrest.Matchers.closeTo;
 import static org.hamcrest.Matchers.containsString;
@@ -1738,7 +1738,7 @@ private static void index(String name, String... docs) throws IOException {
         }
         request.setJsonEntity(bulk.toString());
         Response response = client().performRequest(request);
-        Map<String, Object> result = entityToMapNoPartialCheck(response.getEntity(), XContentType.JSON);
+        Map<String, Object> result = entityToMap(response.getEntity(), XContentType.JSON);
         assertMap(result, matchesMap().extraOk().entry("errors", false));
     }
 

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

@@ -29,7 +29,7 @@
 
 import static org.elasticsearch.test.ListMatcher.matchesList;
 import static org.elasticsearch.test.MapMatcher.matchesMap;
-import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.entityToMapNoPartialCheck;
+import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.entityToMap;
 import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.requestObjectBuilder;
 import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.anyOf;
@@ -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);
     }
@@ -279,7 +284,7 @@ protected void indexTimestampDataForClient(RestClient client, int docs, String i
             }""".replace("%differentiator_field_name%", differentiatorFieldName));
         Response response = client.performRequest(createIndex);
         assertThat(
-            entityToMapNoPartialCheck(response.getEntity(), XContentType.JSON),
+            entityToMap(response.getEntity(), XContentType.JSON),
             matchesMap().entry("shards_acknowledged", true).entry("index", indexName).entry("acknowledged", true)
         );
 

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

@@ -1249,13 +1249,22 @@ 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 {
@@ -1511,7 +1520,7 @@ static Map<String, Object> removeAsyncProperties(Map<String, Object> map) {
         return Collections.unmodifiableMap(copy);
     }
 
-    protected static Map<String, Object> entityToMapNoPartialCheck(HttpEntity entity, XContentType expectedContentType) throws IOException {
+    protected static Map<String, Object> entityToMap(HttpEntity entity, XContentType expectedContentType) throws IOException {
         var result = EsqlTestUtils.entityToMap(entity, expectedContentType);
         if (shouldLog()) {
             LOGGER.info("entity={}", result);
@@ -1520,12 +1529,10 @@ protected static Map<String, Object> entityToMapNoPartialCheck(HttpEntity entity
         return result;
     }
 
-    /**
-     * Converts the entity to a map and asserts that this is not a partial response (`is_partial` field is absent or `false`).
-     * Use {@link #entityToMapNoPartialCheck} if you want to skip this check.
-     */
-    protected static Map<String, Object> entityToMap(HttpEntity entity, XContentType expectedContentType) throws IOException {
-        return assertNotPartial(entityToMapNoPartialCheck(entity, expectedContentType));
+    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 {
@@ -1859,7 +1866,7 @@ protected void indexTimestampData(int shards) throws IOException {
             }""".replace("%shards%", Integer.toString(shards)));
         Response response = client().performRequest(createIndex);
         assertThat(
-            entityToMapNoPartialCheck(response.getEntity(), XContentType.JSON),
+            entityToMap(response.getEntity(), XContentType.JSON),
             matchesMap().entry("shards_acknowledged", true).entry("index", testIndexName()).entry("acknowledged", true)
         );