ESQL: Log profile on test error

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

@@ -38,6 +38,9 @@
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.Rule;
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
 
 import java.io.IOException;
 import java.math.BigDecimal;
@@ -302,12 +305,10 @@ protected final void doTest(String query) throws Throwable {
         }
     }
 
-    static Map<String, Object> assertNotPartial(Map<String, Object> answer) {
+    static void 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 {
@@ -336,12 +337,37 @@ protected boolean deduplicateExactWarnings() {
         return false;
     }
 
+    public class ProfileLogger extends TestWatcher {
+        private Object profile;
+
+        void setProfile(RestEsqlTestCase.EsqlResponse response) {
+            profile = response.json().get("profile");
+        }
+
+        @Override
+        protected void failed(Throwable e, Description description) {
+            LOGGER.warn("Profile: {}", profile);
+        }
+    }
+
+    @Rule(order = Integer.MIN_VALUE)
+    public ProfileLogger profileLogger = new ProfileLogger();
+
     private Map<String, Object> runEsql(RequestObjectBuilder requestObject, AssertWarnings assertWarnings) throws IOException {
-        if (mode == Mode.ASYNC) {
-            return RestEsqlTestCase.runEsqlAsync(requestObject, assertWarnings);
-        } else {
-            return RestEsqlTestCase.runEsqlSync(requestObject, assertWarnings);
+        requestObject.profile(true);
+
+        RestEsqlTestCase.EsqlResponse response = mode == Mode.ASYNC
+            ? RestEsqlTestCase.runEsqlAsyncNoWarningsChecks(requestObject, randomBoolean())
+            : RestEsqlTestCase.runEsqlSyncNoWarningsChecks(requestObject);
+
+        profileLogger.setProfile(response);
+
+        RestEsqlTestCase.assertWarnings(response.response(), assertWarnings);
+        if (response.asyncInitialResponse() != response.response()) {
+            RestEsqlTestCase.assertWarnings(response.response(), assertWarnings);
         }
+
+        return response.json();
     }
 
     protected void assertResults(

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

@@ -396,7 +396,9 @@ public void testCSVNoHeaderMode() throws IOException {
         options.addHeader("Content-Type", mediaType);
         options.addHeader("Accept", "text/csv; header=absent");
         request.setOptions(options);
-        HttpEntity entity = performRequest(request, new AssertWarnings.NoWarnings());
+        Response response = performRequest(request);
+        assertWarnings(response, new AssertWarnings.NoWarnings());
+        HttpEntity entity = response.getEntity();
         String actual = Streams.copyToString(new InputStreamReader(entity.getContent(), StandardCharsets.UTF_8));
         assertEquals("keyword0,0\r\n", actual);
     }
@@ -1258,19 +1260,43 @@ public static Map<String, Object> runEsql(
         var results = mode == ASYNC
             ? runEsqlAsync(requestObject, randomBoolean(), assertWarnings)
             : runEsqlSync(requestObject, assertWarnings);
-        return checkPartialResults ? assertNotPartial(results) : results;
+        if (checkPartialResults) {
+            assertNotPartial(results);
+        }
+        return results;
     }
 
     public static Map<String, Object> runEsql(RequestObjectBuilder requestObject, AssertWarnings assertWarnings, Mode mode)
         throws IOException {
         return runEsql(requestObject, assertWarnings, mode, true);
     }
 
+    /**
+     * A response from an ESQL query.
+     * @param response the HTTP response that contains the JSON response body
+     * @param json the parsed JSON response body
+     * @param asyncInitialResponse for async only. The response that was returned from the initial async request.
+     *                             May be the same as {@code response} if the async request completed immediately.
+     */
+    public record EsqlResponse(Response response, Map<String, Object> json, Response asyncInitialResponse) {
+        public EsqlResponse(Response response, Map<String, Object> json) {
+            this(response, json, null);
+        }
+    }
+
     public static Map<String, Object> runEsqlSync(RequestObjectBuilder requestObject, AssertWarnings assertWarnings) throws IOException {
-        Request request = prepareRequestWithOptions(requestObject, SYNC);
+        EsqlResponse response = runEsqlSyncNoWarningsChecks(requestObject);
+        assertWarnings(response.response, assertWarnings);
+        return response.json;
+    }
 
-        HttpEntity entity = performRequest(request, assertWarnings);
-        return entityToMap(entity, requestObject.contentType());
+    public static EsqlResponse runEsqlSyncNoWarningsChecks(RequestObjectBuilder requestObject) throws IOException {
+        Boolean profileEnabled = requestObject.profile;
+        requestObject.profile(true);
+        Request request = prepareRequestWithOptions(requestObject, SYNC);
+        Response response = performRequest(request);
+        Map<String, Object> json = entityToMap(response.getEntity(), requestObject.contentType());
+        return new EsqlResponse(response, json);
     }
 
     public static Map<String, Object> runEsqlAsync(RequestObjectBuilder requestObject, AssertWarnings assertWarnings) throws IOException {
@@ -1282,6 +1308,18 @@ public static Map<String, Object> runEsqlAsync(
         boolean keepOnCompletion,
         AssertWarnings assertWarnings
     ) throws IOException {
+        EsqlResponse response = runEsqlAsyncNoWarningsChecks(requestObject, keepOnCompletion);
+
+        assertWarnings(response.response(), assertWarnings);
+        if (response.asyncInitialResponse() != response.response()) {
+            assertWarnings(response.asyncInitialResponse(), assertWarnings);
+        }
+
+        return response.json;
+    }
+
+    public static EsqlResponse runEsqlAsyncNoWarningsChecks(RequestObjectBuilder requestObject, boolean keepOnCompletion)
+        throws IOException {
         addAsyncParameters(requestObject, keepOnCompletion);
         Request request = prepareRequestWithOptions(requestObject, ASYNC);
 
@@ -1290,6 +1328,7 @@ public static Map<String, Object> runEsqlAsync(
         }
 
         Response response = performRequest(request);
+        Response initialResponse = response;
         HttpEntity entity = response.getEntity();
 
         Object initialColumns = null;
@@ -1309,17 +1348,15 @@ public static Map<String, Object> runEsqlAsync(
                 assertThat(response.getHeader("X-Elasticsearch-Async-Id"), nullValue());
                 assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), is("?0"));
             }
-            assertWarnings(response, assertWarnings);
             json.remove("is_running"); // remove this to not mess up later map assertions
-            return Collections.unmodifiableMap(json);
+            return new EsqlResponse(response, Collections.unmodifiableMap(json), initialResponse);
         } else {
             // async may not return results immediately, so may need an async get
             assertThat(id, is(not(emptyOrNullString())));
             boolean isRunning = (boolean) json.get("is_running");
             if (isRunning == false) {
                 // must have completed immediately so keep_on_completion must be true
                 assertThat(requestObject.keepOnCompletion(), is(true));
-                assertWarnings(response, assertWarnings);
                 // we already have the results, but let's remember them so that we can compare to async get
                 initialColumns = json.get("columns");
                 initialValues = json.get("values");
@@ -1356,9 +1393,8 @@ public static Map<String, Object> runEsqlAsync(
             assertEquals(initialValues, result.get("values"));
         }
 
-        assertWarnings(response, assertWarnings);
         assertDeletable(id);
-        return removeAsyncProperties(result);
+        return new EsqlResponse(response, removeAsyncProperties(result), initialResponse);
     }
 
     private static Object removeOriginalTypesAndSuggestedCast(Object response) {
@@ -1589,8 +1625,9 @@ static String runEsqlAsTextWithFormat(RequestObjectBuilder builder, String forma
         }
 
         Response response = performRequest(request);
-        HttpEntity entity = assertWarnings(response, new AssertWarnings.NoWarnings());
+        assertWarnings(response, new AssertWarnings.NoWarnings());
 
+        HttpEntity entity = response.getEntity();
         // get the content, it could be empty because the request might have not completed
         String initialValue = Streams.copyToString(new InputStreamReader(entity.getContent(), StandardCharsets.UTF_8));
         String id = response.getHeader("X-Elasticsearch-Async-Id");
@@ -1642,7 +1679,8 @@ static String runEsqlAsTextWithFormat(RequestObjectBuilder builder, String forma
             // if `addParam` is false, `options` will already have an `Accept` header
             getRequest.setOptions(options);
             response = performRequest(getRequest);
-            entity = assertWarnings(response, new AssertWarnings.NoWarnings());
+            assertWarnings(response, new AssertWarnings.NoWarnings());
+            entity = response.getEntity();
         }
         String newValue = Streams.copyToString(new InputStreamReader(entity.getContent(), StandardCharsets.UTF_8));
 
@@ -1681,10 +1719,6 @@ private static String attachBody(RequestObjectBuilder requestObject, Request req
         return mediaType;
     }
 
-    private static HttpEntity performRequest(Request request, AssertWarnings assertWarnings) throws IOException {
-        return assertWarnings(performRequest(request), assertWarnings);
-    }
-
     protected static Response performRequest(Request request) throws IOException {
         Response response = client().performRequest(request);
         if (shouldLog()) {
@@ -1695,14 +1729,13 @@ protected static Response performRequest(Request request) throws IOException {
         return response;
     }
 
-    private static HttpEntity assertWarnings(Response response, AssertWarnings assertWarnings) {
+    public static void assertWarnings(Response response, AssertWarnings assertWarnings) {
         List<String> warnings = new ArrayList<>(response.getWarnings());
         warnings.removeAll(mutedWarnings());
         if (shouldLog()) {
             LOGGER.info("RESPONSE warnings (after muted)={}", warnings);
         }
         assertWarnings.assertWarnings(warnings);
-        return response.getEntity();
     }
 
     private static Set<String> mutedWarnings() {