ESQL: Change queries ID to be the same as the async

docs/changelog/127472.yaml

@@ -0,0 +1,6 @@
+pr: 127472
+summary: Change queries ID to be the same as the async
+area: ES|QL
+type: feature
+issues:
+ - 127187

rest-api-spec/src/main/resources/rest-api-spec/api/esql.list_queries.json

@@ -7,10 +7,8 @@
  "stability": "experimental",
  "visibility": "public",
  "headers": {
- "accept": [],
- "content_type": [
- "application/json"
- ]
+ "accept": ["application/json"],
+ "content_type": ["application/json"]
  },
  "url": {
  "paths": [

server/src/main/java/org/elasticsearch/tasks/TaskAwareRequest.java

@@ -9,6 +9,8 @@
 
 package org.elasticsearch.tasks;
 
+import org.elasticsearch.core.Nullable;
+
 import java.util.Map;
 
 /**
@@ -52,6 +54,23 @@ default Task createTask(long id, String type, String action, TaskId parentTaskId
  return new Task(id, type, action, getDescription(), parentTaskId, headers);
  }
 
+ /**
+ * Returns the task object that should be used to keep track of the processing of the request, with an extra local node ID.
+ */
+ // TODO remove the above overload, use only this one.
+ default Task createTask(
+ // TODO this is only nullable in tests, where the MockNode does not guarantee the localNodeId is set before calling this method. We
+ // We should fix the tests, and replace this and id with TaskId instead.
+ @Nullable String localNodeId,
+ long id,
+ String type,
+ String action,
+ TaskId parentTaskId,
+ Map<String, String> headers
+ ) {
+ return createTask(id, type, action, parentTaskId, headers);
+ }
+
  /**
  * Returns optional description of the request to be displayed by the task manager
  */

server/src/main/java/org/elasticsearch/tasks/TaskManager.java

@@ -141,7 +141,14 @@ public Task register(String type, String action, TaskAwareRequest request, boole
  headers.put(key, httpHeader);
  }
  }
- Task task = request.createTask(taskIdGenerator.incrementAndGet(), type, action, request.getParentTask(), headers);
+ Task task = request.createTask(
+ lastDiscoveryNodes.getLocalNodeId(),
+ taskIdGenerator.incrementAndGet(),
+ type,
+ action,
+ request.getParentTask(),
+ headers
+ );
  Objects.requireNonNull(task);
  assert task.getParentTaskId().equals(request.getParentTask()) : "Request [ " + request + "] didn't preserve it parentTaskId";
  if (logger.isTraceEnabled()) {

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

@@ -10,6 +10,8 @@
 import org.elasticsearch.common.io.stream.ByteBufferStreamInput;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.tasks.TaskId;
 
 import java.io.IOException;
@@ -20,7 +22,7 @@
 /**
  * A class that contains all information related to a submitted async execution.
  */
-public final class AsyncExecutionId {
+public final class AsyncExecutionId implements Writeable {
  public static final String ASYNC_EXECUTION_ID_HEADER = "X-Elasticsearch-Async-Id";
  public static final String ASYNC_EXECUTION_IS_RUNNING_HEADER = "X-Elasticsearch-Async-Is-Running";
 
@@ -115,4 +117,13 @@ public static AsyncExecutionId decode(String id) {
  }
  return new AsyncExecutionId(docId, new TaskId(taskId), id);
  }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ out.writeString(getEncoded());
+ }
+
+ public static AsyncExecutionId readFrom(StreamInput input) throws IOException {
+ return decode(input.readString());
+ }
 }

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/async/AsyncExecutionIdWireTests.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.core.async;
+
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.test.AbstractWireSerializingTestCase;
+
+import java.io.IOException;
+
+public class AsyncExecutionIdWireTests extends AbstractWireSerializingTestCase<AsyncExecutionId> {
+ @Override
+ protected Writeable.Reader<AsyncExecutionId> instanceReader() {
+ return AsyncExecutionId::readFrom;
+ }
+
+ @Override
+ protected AsyncExecutionId createTestInstance() {
+ return new AsyncExecutionId(randomAlphaOfLength(15), new TaskId(randomAlphaOfLength(10), randomLong()));
+ }
+
+ @Override
+ protected AsyncExecutionId mutateInstance(AsyncExecutionId instance) throws IOException {
+ return new AsyncExecutionId(
+ instance.getDocId(),
+ new TaskId(instance.getTaskId().getNodeId(), instance.getTaskId().getId() * 12345)
+ );
+ }
+}

x-pack/plugin/esql/qa/security/src/javaRestTest/java/org/elasticsearch/xpack/esql/EsqlSecurityIT.java

@@ -916,19 +916,23 @@ public void testListQueryForbidden() throws Exception {
 
  public void testGetQueryAllowed() throws Exception {
  // This is a bit tricky, since there is no such running query. We just make sure it didn't fail on forbidden privileges.
- Request request = new Request("GET", "_query/queries/foo:1234");
- var resp = expectThrows(ResponseException.class, () -> client().performRequest(request));
- assertThat(resp.getResponse().getStatusLine().getStatusCode(), not(equalTo(404)));
+ setUser(GET_QUERY_REQUEST, "user_with_monitor_privileges");
+ var resp = expectThrows(ResponseException.class, () -> client().performRequest(GET_QUERY_REQUEST));
+ assertThat(resp.getResponse().getStatusLine().getStatusCode(), not(equalTo(403)));
  }
 
  public void testGetQueryForbidden() throws Exception {
- Request request = new Request("GET", "_query/queries/foo:1234");
- setUser(request, "user_without_monitor_privileges");
- var resp = expectThrows(ResponseException.class, () -> client().performRequest(request));
+ setUser(GET_QUERY_REQUEST, "user_without_monitor_privileges");
+ var resp = expectThrows(ResponseException.class, () -> client().performRequest(GET_QUERY_REQUEST));
  assertThat(resp.getResponse().getStatusLine().getStatusCode(), equalTo(403));
  assertThat(resp.getMessage(), containsString("this action is granted by the cluster privileges [monitor_esql,monitor,manage,all]"));
  }
 
+ private static final Request GET_QUERY_REQUEST = new Request(
+ "GET",
+ "_query/queries/FmJKWHpFRi1OU0l5SU1YcnpuWWhoUWcZWDFuYUJBeW1TY0dKM3otWUs2bDJudzo1Mg=="
+ );
+
  private void createEnrichPolicy() throws Exception {
  createIndex("songs", Settings.EMPTY, """
  "properties":{"song_id": {"type": "keyword"}, "title": {"type": "keyword"}, "artist": {"type": "keyword"} }

x-pack/plugin/esql/qa/testFixtures/src/main/resources/query_task.json

@@ -3,6 +3,9 @@
  "id" : 5326,
  "type" : "transport",
  "action" : "indices:data/read/esql",
+ "status" : {
+ "request_id" : "Ks5ApyqMTtWj5LrKigmCjQ"
+ },
  "description" : "FROM test | STATS MAX(d) by a, b", <1>
  "start_time" : "2023-07-31T15:46:32.328Z",
  "start_time_in_millis" : 1690818392328,

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractPausableIntegTestCase.java

@@ -30,6 +30,8 @@
 public abstract class AbstractPausableIntegTestCase extends AbstractEsqlIntegTestCase {
 
  protected static final Semaphore scriptPermits = new Semaphore(0);
+ // Incremented onWait. Can be used to check if the onWait process has been reached.
+ protected static final Semaphore scriptWaits = new Semaphore(0);
 
  protected int pageSize = -1;
 
@@ -98,6 +100,7 @@ public void setupIndex() throws IOException {
  public static class PausableFieldPlugin extends AbstractPauseFieldPlugin {
  @Override
  protected boolean onWait() throws InterruptedException {
+ scriptWaits.release();
  return scriptPermits.tryAcquire(1, TimeUnit.MINUTES);
  }
  }

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlListQueriesActionIT.java

@@ -7,24 +7,21 @@
 
 package org.elasticsearch.xpack.esql.action;
 
+import org.elasticsearch.action.ActionFuture;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.Response;
-import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.test.IntOrLongMatcher;
 import org.elasticsearch.test.MapMatcher;
 import org.elasticsearch.xpack.core.async.GetAsyncResultRequest;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
 
-import java.util.List;
+import java.io.IOException;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.core.TimeValue.timeValueSeconds;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.jsonEntityToMap;
-import static org.hamcrest.Matchers.allOf;
-import static org.hamcrest.Matchers.everyItem;
 import static org.hamcrest.Matchers.is;
-import static org.hamcrest.Matchers.isA;
 
 public class EsqlListQueriesActionIT extends AbstractPausableIntegTestCase {
  private static final String QUERY = "from test | stats sum(pause_me)";
@@ -45,31 +42,10 @@ public void testRunningQueries() throws Exception {
  try (var initialResponse = sendAsyncQuery()) {
  id = initialResponse.asyncExecutionId().get();
 
+ assertRunningQueries();
  var getResultsRequest = new GetAsyncResultRequest(id);
  getResultsRequest.setWaitForCompletionTimeout(timeValueSeconds(1));
  client().execute(EsqlAsyncGetResultAction.INSTANCE, getResultsRequest).get().close();
- Response listResponse = getRestClient().performRequest(new Request("GET", "/_query/queries"));
- @SuppressWarnings("unchecked")
- var listResult = (Map<String, Map<String, Object>>) EsqlTestUtils.singleValue(
- jsonEntityToMap(listResponse.getEntity()).values()
- );
- var taskId = new TaskId(EsqlTestUtils.singleValue(listResult.keySet()));
- MapMatcher basicMatcher = MapMatcher.matchesMap()
- .entry("node", is(taskId.getNodeId()))
- .entry("id", IntOrLongMatcher.matches(taskId.getId()))
- .entry("query", is(QUERY))
- .entry("start_time_millis", IntOrLongMatcher.isIntOrLong())
- .entry("running_time_nanos", IntOrLongMatcher.isIntOrLong());
- MapMatcher.assertMap(EsqlTestUtils.singleValue(listResult.values()), basicMatcher);
-
- Response getQueryResponse = getRestClient().performRequest(new Request("GET", "/_query/queries/" + taskId));
- MapMatcher.assertMap(
- jsonEntityToMap(getQueryResponse.getEntity()),
- basicMatcher.entry("coordinating_node", isA(String.class))
- .entry("data_nodes", allOf(isA(List.class), everyItem(isA(String.class))))
- .entry("documents_found", IntOrLongMatcher.isIntOrLong())
- .entry("values_loaded", IntOrLongMatcher.isIntOrLong())
- );
  } finally {
  if (id != null) {
  // Finish the query.
@@ -82,9 +58,44 @@ public void testRunningQueries() throws Exception {
  }
  }
 
+ public void testRunningQueriesSync() throws Exception {
+ var future = sendSyncQueryAsyncly();
+ try {
+ scriptWaits.acquire();
+ assertRunningQueries();
+ } finally {
+ scriptPermits.release(numberOfDocs());
+ future.actionGet(timeValueSeconds(60)).close();
+ }
+ }
+
+ private static void assertRunningQueries() throws IOException {
+ Response listResponse = getRestClient().performRequest(new Request("GET", "/_query/queries"));
+ @SuppressWarnings("unchecked")
+ var listResult = (Map<String, Map<String, Object>>) EsqlTestUtils.singleValue(jsonEntityToMap(listResponse.getEntity()).values());
+ String queryId = EsqlTestUtils.singleValue(listResult.keySet());
+ MapMatcher basicMatcher = MapMatcher.matchesMap()
+ .entry("query", is(QUERY))
+ .entry("start_time_millis", IntOrLongMatcher.isIntOrLong())
+ .entry("running_time_nanos", IntOrLongMatcher.isIntOrLong());
+ MapMatcher.assertMap(EsqlTestUtils.singleValue(listResult.values()), basicMatcher);
+
+ Response getQueryResponse = getRestClient().performRequest(new Request("GET", "/_query/queries/" + queryId));
+ MapMatcher.assertMap(
+ jsonEntityToMap(getQueryResponse.getEntity()),
+ basicMatcher.entry("documents_found", IntOrLongMatcher.isIntOrLong()).entry("values_loaded", IntOrLongMatcher.isIntOrLong())
+ );
+ }
+
  private EsqlQueryResponse sendAsyncQuery() {
  scriptPermits.drainPermits();
  scriptPermits.release(between(1, 5));
  return EsqlQueryRequestBuilder.newAsyncEsqlQueryRequestBuilder(client()).query(QUERY).execute().actionGet(60, TimeUnit.SECONDS);
  }
+
+ private ActionFuture<EsqlQueryResponse> sendSyncQueryAsyncly() {
+ scriptPermits.drainPermits();
+ scriptPermits.release(between(1, 5));
+ return EsqlQueryRequestBuilder.newSyncEsqlQueryRequestBuilder(client()).query(QUERY).execute();
+ }
 }

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

@@ -7,33 +7,33 @@
 
 package org.elasticsearch.xpack.esql.action;
 
+import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionRequestValidationException;
-import org.elasticsearch.action.LegacyActionRequest;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.xpack.core.async.AsyncExecutionId;
 
 import java.io.IOException;
 
-public class EsqlGetQueryRequest extends LegacyActionRequest {
- private final TaskId id;
+public class EsqlGetQueryRequest extends ActionRequest {
+ private final AsyncExecutionId asyncExecutionId;
 
- public EsqlGetQueryRequest(TaskId id) {
- this.id = id;
+ public EsqlGetQueryRequest(AsyncExecutionId asyncExecutionId) {
+ this.asyncExecutionId = asyncExecutionId;
  }
 
- public TaskId id() {
- return id;
+ public AsyncExecutionId id() {
+ return asyncExecutionId;
  }
 
  public EsqlGetQueryRequest(StreamInput streamInput) throws IOException {
  super(streamInput);
- id = TaskId.readFromStream(streamInput);
+ asyncExecutionId = AsyncExecutionId.decode(streamInput.readString());
  }
 
  @Override
  public void writeTo(StreamOutput out) throws IOException {
- out.writeWriteable(id);
+ out.writeWriteable(asyncExecutionId);
  }
 
  @Override

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.CompositeIndicesRequest;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.UUIDs;
 import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.settings.Settings;
@@ -20,8 +21,10 @@
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.xpack.core.async.AsyncExecutionId;
 import org.elasticsearch.xpack.esql.Column;
 import org.elasticsearch.xpack.esql.parser.QueryParams;
+import org.elasticsearch.xpack.esql.plugin.EsqlQueryStatus;
 import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 
 import java.io.IOException;
@@ -242,9 +245,32 @@ public EsqlQueryRequest allowPartialResults(boolean allowPartialResults) {
  }
 
  @Override
- public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
- // Pass the query as the description
- return new CancellableTask(id, type, action, query, parentTaskId, headers);
+ public Task createTask(String localNodeId, long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+ var status = new EsqlQueryStatus(new AsyncExecutionId(UUIDs.randomBase64UUID(), new TaskId(localNodeId, id)));
+ return new EsqlQueryRequestTask(query, id, type, action, parentTaskId, headers, status);
+ }
+
+ private static class EsqlQueryRequestTask extends CancellableTask {
+ private final Status status;
+
+ EsqlQueryRequestTask(
+ String query,
+ long id,
+ String type,
+ String action,
+ TaskId parentTaskId,
+ Map<String, String> headers,
+ EsqlQueryStatus status
+ ) {
+ // Pass the query as the description
+ super(id, type, action, query, parentTaskId, headers);
+ this.status = status;
+ }
+
+ @Override
+ public Status getStatus() {
+ return status;
+ }
  }
 
  // Setter for tests