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

GalLalouche · GalLalouche · commit 03839605dd05 · 2025-04-28T17:36:26.000+03:00
Also the individual ID, node, and data node tags. GitHub: Resolves #127187
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractPausableIntegTestCase.java b/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);
         }
     }
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlListQueriesActionIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlListQueriesActionIT.java
@@ -7,24 +7,20 @@
 
 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.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)";
@@ -53,22 +49,17 @@ public void testRunningQueries() throws Exception {
             var listResult = (Map<String, Map<String, Object>>) EsqlTestUtils.singleValue(
                 jsonEntityToMap(listResponse.getEntity()).values()
             );
-            var taskId = new TaskId(EsqlTestUtils.singleValue(listResult.keySet()));
+            String queryId = 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));
+            Response getQueryResponse = getRestClient().performRequest(new Request("GET", "/_query/queries/" + queryId));
             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())
+                basicMatcher.entry("documents_found", IntOrLongMatcher.isIntOrLong()).entry("values_loaded", IntOrLongMatcher.isIntOrLong())
             );
         } finally {
             if (id != null) {
@@ -82,9 +73,42 @@ public void testRunningQueries() throws Exception {
         }
     }
 
+    public void testRunningQueriesSync() throws Exception {
+        var future = sendSyncQueryAsyncly();
+        try {
+            scriptWaits.acquire();
+            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())
+            );
+        } finally {
+            scriptPermits.release(numberOfDocs());
+            future.actionGet(timeValueSeconds(60)).close();
+        }
+    }
+
     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();
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlGetQueryRequest.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlGetQueryRequest.java
@@ -11,29 +11,29 @@
 import org.elasticsearch.action.ActionRequestValidationException;
 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 ActionRequest {
-    private final TaskId id;
+    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
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequest.java b/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;
@@ -22,6 +23,7 @@
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xpack.esql.Column;
 import org.elasticsearch.xpack.esql.parser.QueryParams;
+import org.elasticsearch.xpack.esql.plugin.EsqlDocIdStatus;
 import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 
 import java.io.IOException;
@@ -243,8 +245,14 @@ public EsqlQueryRequest allowPartialResults(boolean allowPartialResults) {
 
     @Override
     public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+        EsqlDocIdStatus status = new EsqlDocIdStatus(UUIDs.randomBase64UUID());
         // Pass the query as the description
-        return new CancellableTask(id, type, action, query, parentTaskId, headers);
+        return new CancellableTask(id, type, action, query, parentTaskId, headers) {
+            @Override
+            public Status getStatus() {
+                return status;
+            }
+        };
     }
 
     // Setter for tests
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/RestEsqlListQueriesAction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/RestEsqlListQueriesAction.java
@@ -15,9 +15,8 @@
 import org.elasticsearch.rest.Scope;
 import org.elasticsearch.rest.ServerlessScope;
 import org.elasticsearch.rest.action.RestToXContentListener;
-import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.xpack.core.async.AsyncExecutionId;
 
-import java.io.IOException;
 import java.util.List;
 
 import static org.elasticsearch.rest.RestRequest.Method.GET;
@@ -37,7 +36,7 @@ public List<Route> routes() {
     }
 
     @Override
-    protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+    protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) {
         return restChannelConsumer(request, client);
     }
 
@@ -46,7 +45,7 @@ private static RestChannelConsumer restChannelConsumer(RestRequest request, Node
 
         String id = request.param("id");
         var action = id != null ? EsqlGetQueryAction.INSTANCE : EsqlListQueriesAction.INSTANCE;
-        var actionRequest = id != null ? new EsqlGetQueryRequest(new TaskId(id)) : new EsqlListQueriesRequest();
+        var actionRequest = id != null ? new EsqlGetQueryRequest(AsyncExecutionId.decode(id)) : new EsqlListQueriesRequest();
 
         return channel -> client.execute(action, actionRequest, new RestToXContentListener<>(channel));
     }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlDocIdStatus.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlDocIdStatus.java
@@ -0,0 +1,52 @@
+/*
+ * 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.esql.plugin;
+
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.core.async.AsyncExecutionId;
+import org.elasticsearch.xpack.esql.action.EsqlQueryAction;
+
+import java.io.IOException;
+
+/**
+ * Keeps track of the doc ID, which is itself used for {@link AsyncExecutionId}.
+ *
+ * The reason this contains just the doc ID and not the entire {@link AsyncExecutionId} is that during the creation of
+ * {@link EsqlQueryAction}, we don't have access to the node ID yet, thus we can't create a {@link TaskId} yet.
+ */
+public record EsqlDocIdStatus(String id) implements Task.Status {
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
+        Task.Status.class,
+        "EsqlDocIdStatus",
+        EsqlDocIdStatus::new
+    );
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
+    private EsqlDocIdStatus(StreamInput stream) throws IOException {
+        this(stream.readString());
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(id);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        return builder.rawValue(id);
+    }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlGetQueryResponse.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlGetQueryResponse.java
@@ -14,7 +14,6 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
-import java.util.List;
 
 public class EsqlGetQueryResponse extends ActionResponse implements ToXContentObject {
     // This is rather limited at the moment, as we don't extract information such as CPU and memory usage, owning user, etc. for the task.
@@ -24,22 +23,16 @@ public record DetailedQuery(
         long runningTimeNanos,
         long documentsFound,
         long valuesLoaded,
-        String query,
-        String coordinatingNode,
-        List<String> dataNodes
+        String query
     ) implements ToXContentObject {
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();
-            builder.field("id", id.getId());
-            builder.field("node", id.getNodeId());
             builder.field("start_time_millis", startTimeMillis);
             builder.field("running_time_nanos", runningTimeNanos);
             builder.field("documents_found", documentsFound);
             builder.field("values_loaded", valuesLoaded);
             builder.field("query", query);
-            builder.field("coordinating_node", coordinatingNode);
-            builder.field("data_nodes", dataNodes);
             builder.endObject();
             return builder;
         }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlListQueriesResponse.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlListQueriesResponse.java
@@ -9,23 +9,21 @@
 
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ToXContentFragment;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.core.async.AsyncExecutionId;
 
 import java.io.IOException;
 import java.util.List;
 
 public class EsqlListQueriesResponse extends ActionResponse implements ToXContentObject {
     private final List<Query> queries;
 
-    public record Query(TaskId taskId, long startTimeMillis, long runningTimeNanos, String query) implements ToXContentFragment {
+    public record Query(AsyncExecutionId id, long startTimeMillis, long runningTimeNanos, String query) implements ToXContentFragment {
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-            builder.startObject(taskId.toString());
-            builder.field("id", taskId.getId());
-            builder.field("node", taskId.getNodeId());
+            builder.startObject(id.getEncoded());
             builder.field("start_time_millis", startTimeMillis);
             builder.field("running_time_nanos", runningTimeNanos);
             builder.field("query", query);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java
@@ -275,6 +275,7 @@ public List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         entries.add(AbstractPageMappingOperator.Status.ENTRY);
         entries.add(AbstractPageMappingToIteratorOperator.Status.ENTRY);
         entries.add(AggregationOperator.Status.ENTRY);
+        entries.add(EsqlDocIdStatus.ENTRY);
         entries.add(ExchangeSinkOperator.Status.ENTRY);
         entries.add(ExchangeSourceOperator.Status.ENTRY);
         entries.add(HashAggregationOperator.Status.ENTRY);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlGetQueryAction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlGetQueryAction.java
@@ -49,7 +49,7 @@ protected void doExecute(Task task, EsqlGetQueryRequest request, ActionListener<
             nodeClient,
             ESQL_ORIGIN,
             TransportGetTaskAction.TYPE,
-            new GetTaskRequest().setTaskId(request.id()),
+            new GetTaskRequest().setTaskId(request.id().getTaskId()),
             new ActionListener<>() {
                 @Override
                 public void onResponse(GetTaskResponse response) {
@@ -64,7 +64,7 @@ public void onResponse(GetTaskResponse response) {
                         TransportListTasksAction.TYPE,
                         new ListTasksRequest().setDetailed(true)
                             .setActions(DriverTaskRunner.ACTION_NAME)
-                            .setTargetParentTaskId(request.id()),
+                            .setTargetParentTaskId(request.id().getTaskId()),
                         new ActionListener<>() {
                             @Override
                             public void onResponse(ListTasksResponse response) {
@@ -91,7 +91,6 @@ public void onFailure(Exception e) {
 
     private static EsqlGetQueryResponse.DetailedQuery toDetailedQuery(TaskInfo main, ListTasksResponse sub) {
         String query = main.description();
-        String coordinatingNode = main.node();
 
         // TODO include completed drivers in documentsFound and valuesLoaded
         long documentsFound = 0;
@@ -110,9 +109,7 @@ private static EsqlGetQueryResponse.DetailedQuery toDetailedQuery(TaskInfo main,
             main.runningTimeNanos(),
             documentsFound,
             valuesLoaded,
-            query,
-            coordinatingNode,
-            sub.getTasks().stream().map(TaskInfo::node).distinct().toList() // Data nodes
+            query
         );
     }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlListQueriesAction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlListQueriesAction.java
@@ -20,6 +20,7 @@
 import org.elasticsearch.tasks.TaskInfo;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.ClientHelper;
+import org.elasticsearch.xpack.core.async.AsyncExecutionId;
 import org.elasticsearch.xpack.esql.action.EsqlListQueriesAction;
 import org.elasticsearch.xpack.esql.action.EsqlListQueriesRequest;
 import org.elasticsearch.xpack.esql.action.EsqlQueryAction;
@@ -72,7 +73,7 @@ public void onFailure(Exception e) {
 
     private static EsqlListQueriesResponse.Query toQuery(TaskInfo taskInfo) {
         return new EsqlListQueriesResponse.Query(
-            taskInfo.taskId(),
+            new AsyncExecutionId(((EsqlDocIdStatus) taskInfo.status()).id(), taskInfo.taskId()),
             taskInfo.startTime(),
             taskInfo.runningTimeNanos(),
             taskInfo.description()
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java
@@ -404,7 +404,12 @@ public EsqlQueryTask createTask(
             originHeaders,
             asyncExecutionId,
             request.keepAlive()
-        );
+        ) {
+            @Override
+            public Status getStatus() {
+                return new EsqlDocIdStatus(asyncExecutionId.getDocId());
+            }
+        };
     }
 
     @Override

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,8 @@`
`30`	`30`	`public abstract class AbstractPausableIntegTestCase extends AbstractEsqlIntegTestCase {`
`31`	`31`
`32`	`32`	`protected static final Semaphore scriptPermits = new Semaphore(0);`
	`33`	`+ // Incremented onWait. Can be used to check if the onWait process has been reached.`
	`34`	`+ protected static final Semaphore scriptWaits = new Semaphore(0);`
`33`	`35`
`34`	`36`	`protected int pageSize = -1;`
`35`	`37`
`@@ -98,6 +100,7 @@ public void setupIndex() throws IOException {`
`98`	`100`	`public static class PausableFieldPlugin extends AbstractPauseFieldPlugin {`
`99`	`101`	`@Override`
`100`	`102`	`protected boolean onWait() throws InterruptedException {`
	`103`	`+ scriptWaits.release();`
`101`	`104`	`return scriptPermits.tryAcquire(1, TimeUnit.MINUTES);`
`102`	`105`	`}`
`103`	`106`	`}`