ESQL: Enable async get to support formatting

docs/changelog/111104.yaml

@@ -0,0 +1,6 @@
+pr: 111104
+summary: "ESQL: Enable async get to support formatting"
+area: ES|QL
+type: feature
+issues:
+ - 110926

docs/reference/esql/esql-async-query-get-api.asciidoc

@@ -39,6 +39,10 @@ parameter is `true`.
 [[esql-async-query-get-api-query-params]]
 ==== {api-query-parms-title}
 
+The API accepts the same parameters as the synchronous
+<<esql-query-api-query-params,query API>>, along with the following
+parameters:
+
 `wait_for_completion_timeout`::
 (Optional, <<time-units,time value>>)
 Timeout duration to wait for the request to finish. Defaults to no timeout,

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/async/AsyncTaskManagementService.java

@@ -208,7 +208,7 @@ private ActionListener<Response> wrapStoringListener(
  ActionListener<Response> listener
  ) {
  AtomicReference<ActionListener<Response>> exclusiveListener = new AtomicReference<>(listener);
- // This is will performed in case of timeout
+ // This will be performed in case of timeout
  Scheduler.ScheduledCancellable timeoutHandler = threadPool.schedule(() -> {
  ActionListener<Response> acquiredListener = exclusiveListener.getAndSet(null);
  if (acquiredListener != null) {

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

@@ -52,6 +52,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.function.IntFunction;
 
@@ -350,21 +351,22 @@ public void testTextMode() throws IOException {
  int count = randomIntBetween(0, 100);
  bulkLoadTestData(count);
  var builder = requestObjectBuilder().query(fromIndex() + " | keep keyword, integer | sort integer asc | limit 100");
- assertEquals(expectedTextBody("txt", count, null), runEsqlAsTextWithFormat(builder, "txt", null));
+ assertEquals(expectedTextBody("txt", count, null), runEsqlAsTextWithFormat(builder, "txt", null, mode));
+
  }
 
  public void testCSVMode() throws IOException {
  int count = randomIntBetween(0, 100);
  bulkLoadTestData(count);
  var builder = requestObjectBuilder().query(fromIndex() + " | keep keyword, integer | sort integer asc | limit 100");
- assertEquals(expectedTextBody("csv", count, '|'), runEsqlAsTextWithFormat(builder, "csv", '|'));
+ assertEquals(expectedTextBody("csv", count, '|'), runEsqlAsTextWithFormat(builder, "csv", '|', mode));
  }
 
  public void testTSVMode() throws IOException {
  int count = randomIntBetween(0, 100);
  bulkLoadTestData(count);
  var builder = requestObjectBuilder().query(fromIndex() + " | keep keyword, integer | sort integer asc | limit 100");
- assertEquals(expectedTextBody("tsv", count, null), runEsqlAsTextWithFormat(builder, "tsv", null));
+ assertEquals(expectedTextBody("tsv", count, null), runEsqlAsTextWithFormat(builder, "tsv", null, mode));
  }
 
  public void testCSVNoHeaderMode() throws IOException {
@@ -1167,14 +1169,19 @@ static void deleteNonExistent(Request request) throws IOException {
  assertEquals(404, response.getStatusLine().getStatusCode());
  }
 
- static String runEsqlAsTextWithFormat(RequestObjectBuilder builder, String format, @Nullable Character delimiter) throws IOException {
- Request request = prepareRequest(SYNC);
+ static String runEsqlAsTextWithFormat(RequestObjectBuilder builder, String format, @Nullable Character delimiter, Mode mode)
+ throws IOException {
+ Request request = prepareRequest(mode);
+ if (mode == ASYNC) {
+ addAsyncParameters(builder);
+ }
  String mediaType = attachBody(builder.build(), request);
 
  RequestOptions.Builder options = request.getOptions().toBuilder();
  options.addHeader("Content-Type", mediaType);
 
- if (randomBoolean()) {
+ boolean addParam = randomBoolean();
+ if (addParam) {
  request.addParameter("format", format);
  } else {
  switch (format) {
@@ -1188,8 +1195,71 @@ static String runEsqlAsTextWithFormat(RequestObjectBuilder builder, String forma
  }
  request.setOptions(options);
 
- HttpEntity entity = performRequest(request, new AssertWarnings.NoWarnings());
- return Streams.copyToString(new InputStreamReader(entity.getContent(), StandardCharsets.UTF_8));
+ Response response = performRequest(request);
+ HttpEntity entity = assertWarnings(response, new AssertWarnings.NoWarnings());
+
+ // 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");
+
+ if (mode == SYNC) {
+ assertThat(id, is(emptyOrNullString()));
+ return initialValue;
+ }
+
+ if (id == null) {
+ // no id returned from an async call, must have completed immediately and without keep_on_completion
+ assertThat(builder.keepOnCompletion(), either(nullValue()).or(is(false)));
+ assertNull(response.getHeader("is_running"));
+ // the content cant be empty
+ assertThat(initialValue, not(emptyOrNullString()));
+ return initialValue;
+ } else {
+ // async may not return results immediately, so may need an async get
+ assertThat(id, is(not(emptyOrNullString())));
+ String isRunning = response.getHeader("X-Elasticsearch-Async-Is-Running");
+ if (Objects.equals(isRunning, "?0")) {
+ // must have completed immediately so keep_on_completion must be true
+ assertThat(builder.keepOnCompletion(), is(true));
+ } else {
+ // did not return results immediately, so we will need an async get
+ // Also, different format modes return different results.
+ switch (format) {
+ case "txt" -> assertThat(initialValue, emptyOrNullString());
+ case "csv" -> {
+ assertEquals(initialValue, "\r\n");
+ initialValue = "";
+ }
+ case "tsv" -> {
+ assertEquals(initialValue, "\n");
+ initialValue = "";
+ }
+ }
+ }
+ // issue a second request to "async get" the results
+ Request getRequest = prepareAsyncGetRequest(id);
+ if (delimiter != null) {
+ getRequest.addParameter("delimiter", String.valueOf(delimiter));
+ }
+ // If the `format` parameter is not added, the GET request will return a response
+ // with the `Content-Type` type due to the lack of an `Accept` header.
+ if (addParam) {
+ getRequest.addParameter("format", format);
+ }
+ // if `addParam` is false, `options` will already have an `Accept` header
+ getRequest.setOptions(options);
+ response = performRequest(getRequest);
+ entity = assertWarnings(response, new AssertWarnings.NoWarnings());
+ }
+ String newValue = Streams.copyToString(new InputStreamReader(entity.getContent(), StandardCharsets.UTF_8));
+
+ // assert initial contents, if any, are the same as async get contents
+ if (initialValue != null && initialValue.isEmpty() == false) {
+ assertEquals(initialValue, newValue);
+ }
+
+ assertDeletable(id);
+ return newValue;
  }
 
  private static Request prepareRequest(Mode mode) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlResponseListener.java

@@ -87,7 +87,7 @@ public TimeValue stop() {
  /**
  * Keep the initial query for logging purposes.
  */
- private final String esqlQuery;
+ private final String esqlQueryOrId;
  /**
  * Stop the time it took to build a response to later log it. Use something thread-safe here because stopping time requires state and
  * {@link EsqlResponseListener} might be used from different threads.
@@ -98,29 +98,23 @@ public TimeValue stop() {
  * To correctly time the execution of a request, a {@link EsqlResponseListener} must be constructed immediately before execution begins.
  */
  public EsqlResponseListener(RestChannel channel, RestRequest restRequest, EsqlQueryRequest esqlRequest) {
- super(channel);
+ this(channel, restRequest, esqlRequest.query(), EsqlMediaTypeParser.getResponseMediaType(restRequest, esqlRequest));
+ }
 
+ /**
+ * Async query get API does not have an EsqlQueryRequest, store the async ID as an alternative
+ */
+ public EsqlResponseListener(RestChannel channel, RestRequest restRequest) {
+ this(channel, restRequest, restRequest.param("id"), EsqlMediaTypeParser.getResponseMediaType(restRequest));
+ }
+
+ private EsqlResponseListener(RestChannel channel, RestRequest restRequest, String esqlQueryOrId, MediaType mediaType) {
+ super(channel);
  this.channel = channel;
  this.restRequest = restRequest;
- this.esqlQuery = esqlRequest.query();
- mediaType = EsqlMediaTypeParser.getResponseMediaType(restRequest, esqlRequest);
-
- /*
- * Special handling for the "delimiter" parameter which should only be
- * checked for being present or not in the case of CSV format. We cannot
- * override {@link BaseRestHandler#responseParams()} because this
- * parameter should only be checked for CSV, not other formats.
- */
- if (mediaType != CSV && restRequest.hasParam(URL_PARAM_DELIMITER)) {
- String message = String.format(
- Locale.ROOT,
- "parameter: [%s] can only be used with the format [%s] for request [%s]",
- URL_PARAM_DELIMITER,
- CSV.queryParameter(),
- restRequest.path()
- );
- throw new IllegalArgumentException(message);
- }
+ this.esqlQueryOrId = esqlQueryOrId;
+ this.mediaType = mediaType;
+ checkDelimiter();
  }
 
  @Override
@@ -197,14 +191,18 @@ public ActionListener<EsqlQueryResponse> wrapWithLogging() {
  listener.onResponse(r);
  // At this point, the StopWatch should already have been stopped, so we log a consistent time.
  LOGGER.debug(
- "Finished execution of ESQL query.\nQuery string: [{}]\nExecution time: [{}]ms",
- esqlQuery,
+ "Finished execution of ESQL query.\nQuery string or async ID: [{}]\nExecution time: [{}]ms",
+ esqlQueryOrId,
  getTook(r, TimeUnit.MILLISECONDS)
  );
  }, ex -> {
  // In case of failure, stop the time manually before sending out the response.
  long timeMillis = getTook(null, TimeUnit.MILLISECONDS);
- LOGGER.debug("Failed execution of ESQL query.\nQuery string: [{}]\nExecution time: [{}]ms", esqlQuery, timeMillis);
+ LOGGER.debug(
+ "Failed execution of ESQL query.\nQuery string or async ID: [{}]\nExecution time: [{}]ms",
+ esqlQueryOrId,
+ timeMillis
+ );
  listener.onFailure(ex);
  });
  }
@@ -213,4 +211,23 @@ static void logOnFailure(Throwable throwable) {
  RestStatus status = ExceptionsHelper.status(throwable);
  LOGGER.log(status.getStatus() >= 500 ? Level.WARN : Level.DEBUG, () -> "Request failed with status [" + status + "]: ", throwable);
  }
+
+ /*
+ * Special handling for the "delimiter" parameter which should only be
+ * checked for being present or not in the case of CSV format. We cannot
+ * override {@link BaseRestHandler#responseParams()} because this
+ * parameter should only be checked for CSV, not other formats.
+ */
+ private void checkDelimiter() {
+ if (mediaType != CSV && restRequest.hasParam(URL_PARAM_DELIMITER)) {
+ String message = String.format(
+ Locale.ROOT,
+ "parameter: [%s] can only be used with the format [%s] for request [%s]",
+ URL_PARAM_DELIMITER,
+ CSV.queryParameter(),
+ restRequest.path()
+ );
+ throw new IllegalArgumentException(message);
+ }
+ }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/RestEsqlGetAsyncResultAction.java

@@ -12,7 +12,6 @@
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.Scope;
 import org.elasticsearch.rest.ServerlessScope;
-import org.elasticsearch.rest.action.RestRefCountedChunkedToXContentListener;
 import org.elasticsearch.xpack.core.async.GetAsyncResultRequest;
 
 import java.util.List;
@@ -43,7 +42,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
  if (request.hasParam("keep_alive")) {
  get.setKeepAlive(request.paramAsTime("keep_alive", get.getKeepAlive()));
  }
- return channel -> client.execute(EsqlAsyncGetResultAction.INSTANCE, get, new RestRefCountedChunkedToXContentListener<>(channel));
+ return channel -> client.execute(EsqlAsyncGetResultAction.INSTANCE, get, new EsqlResponseListener(channel, request));
  }
 
  @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlMediaTypeParser.java

@@ -42,12 +42,20 @@ public class EsqlMediaTypeParser {
  * combinations are detected.
  */
  public static MediaType getResponseMediaType(RestRequest request, EsqlQueryRequest esqlRequest) {
- var mediaType = request.hasParam(URL_PARAM_FORMAT) ? mediaTypeFromParams(request) : mediaTypeFromHeaders(request);
+ var mediaType = getResponseMediaType(request);
  validateColumnarRequest(esqlRequest.columnar(), mediaType);
  validateIncludeCCSMetadata(esqlRequest.includeCCSMetadata(), mediaType);
  return checkNonNullMediaType(mediaType, request);
  }
 
+ /*
+ * Retrieve the mediaType of a request without validating the EsqlQueryRequest.
+ */
+ public static MediaType getResponseMediaType(RestRequest request) {
+ var mediaType = request.hasParam(URL_PARAM_FORMAT) ? mediaTypeFromParams(request) : mediaTypeFromHeaders(request);
+ return checkNonNullMediaType(mediaType, request);
+ }
+
  private static MediaType mediaTypeFromHeaders(RestRequest request) {
  ParsedMediaType acceptType = request.getParsedAccept();
  MediaType mediaType = acceptType != null ? acceptType.toMediaType(MEDIA_TYPE_REGISTRY) : request.getXContentType();