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,21 @@ 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(
+        @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

@@ -781,19 +781,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