Add timeout to ESQL

docs/changelog/99526.yaml

@@ -0,0 +1,5 @@
+pr: 99526
+summary: Add timeout to ESQL
+area: ES|QL
+type: enhancement
+issues: []

x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/RestEsqlIT.java

@@ -52,6 +52,50 @@ public void testBasicEsql() throws IOException {
         assertEquals(List.of(List.of(499.5d)), result.get("values"));
     }
 
+    public void testTimeout() throws IOException {
+        StringBuilder b = new StringBuilder();
+        for (int i = 0; i < 10000; i++) {
+            b.append(String.format(Locale.ROOT, """
+                {"create":{"_index":"%s"}}
+                {"test":"value%s"}
+                """, testIndexName(), randomAlphaOfLength(20) + " " + randomAlphaOfLength(20) + " " + randomAlphaOfLength(20)));
+        }
+        Request bulk = new Request("POST", "/_bulk");
+        bulk.addParameter("refresh", "true");
+        bulk.addParameter("filter_path", "errors");
+        bulk.setJsonEntity(b.toString());
+        Response response = client().performRequest(bulk);
+        Assert.assertEquals("{\"errors\":false}", EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8));
+
+        // let's make it a bit expensive, so that the timeout will most likely trigger
+        RequestObjectBuilder builder = new RequestObjectBuilder().query(fromIndex() + """
+            | grok test "%{NOTSPACE:a} %{NOTSPACE:b} %{NOTSPACE:c}"
+            | grok test "%{NOTSPACE:c} %{NOTSPACE:d} %{NOTSPACE:e}"
+            | grok test "%{NOTSPACE:f} %{NOTSPACE:g} %{NOTSPACE:h}"
+            | grok test "%{NOTSPACE:i} %{NOTSPACE:j} %{NOTSPACE:k}"
+            | grok test "%{WORD:l} %{WORD:m} %{WORD:n}"
+            | grok test "%{WORD:o} %{WORD:p} %{WORD:q}"
+            | grok test "%{USERNAME:r} %{USERNAME:s} %{USERNAME:t}"
+            | grok test "%{USERNAME:u} %{USERNAME:v} %{USERNAME:w}"
+            | grok test "%{DATA:x} %{DATA:y} %{DATA:z}"
+            | sort test
+            | limit 9999
+            | sort b desc
+            | limit 9998
+            | sort c
+            | limit 9997
+            | eval s = length(a) + length(b) * length(c)
+            | stats max(s)
+            """);
+        if (Build.current().isSnapshot()) {
+            builder.pragmas(Settings.builder().put("data_partitioning", "shard").build());
+        }
+        builder.timeout(0);
+        builder.build();
+        ResponseException re = expectThrows(ResponseException.class, () -> runEsql(builder));
+        assertThat(EntityUtils.toString(re.getResponse().getEntity()), containsString("ESQL query timed out after 0s"));
+    }
+
     public void testInvalidPragma() throws IOException {
         assumeTrue("pragma only enabled on snapshot builds", Build.current().isSnapshot());
         RequestObjectBuilder builder = new RequestObjectBuilder().query("row a = 1, b = 2");

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

@@ -86,6 +86,11 @@ public RequestObjectBuilder timeZone(ZoneId zoneId) throws IOException {
             return this;
         }
 
+        public RequestObjectBuilder timeout(int timeout) throws IOException {
+            builder.field("timeout", timeout);
+            return this;
+        }
+
         public RequestObjectBuilder pragmas(Settings pragmas) throws IOException {
             builder.startObject("pragma");
             pragmas.toXContent(builder, ToXContent.EMPTY_PARAMS);

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -87,7 +87,8 @@ public byte[] max(String field, DataType dataType) {
         null,
         null,
         new QueryPragmas(Settings.EMPTY),
-        EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(Settings.EMPTY)
+        EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(Settings.EMPTY),
+        null
     );
 
     private EsqlTestUtils() {}

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.tasks.CancellableTask;
@@ -35,6 +36,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
 import static org.elasticsearch.action.ValidateActions.addValidationError;
@@ -65,13 +67,16 @@ public class EsqlQueryRequest extends ActionRequest implements CompositeIndicesR
     private static final ParseField PARAMS_FIELD = new ParseField("params");
     private static final ParseField LOCALE_FIELD = new ParseField("locale");
 
+    private static final ParseField TIMEOUT_FIELD = new ParseField("timeout");
+
     private static final ObjectParser<EsqlQueryRequest, Void> PARSER = objectParser(EsqlQueryRequest::new);
 
     private String query;
     private boolean columnar;
     private ZoneId zoneId;
     private Locale locale;
     private QueryBuilder filter;
+    private TimeValue timeout;
     private QueryPragmas pragmas = new QueryPragmas(Settings.EMPTY);
     private List<TypedParamValue> params = List.of();
 
@@ -149,6 +154,14 @@ public void params(List<TypedParamValue> params) {
         this.params = params;
     }
 
+    public void timeout(TimeValue val) {
+        this.timeout = val;
+    }
+
+    public TimeValue timeout() {
+        return timeout;
+    }
+
     public static EsqlQueryRequest fromXContent(XContentParser parser) {
         return PARSER.apply(parser, null);
     }
@@ -166,6 +179,7 @@ private static ObjectParser<EsqlQueryRequest, Void> objectParser(Supplier<EsqlQu
         );
         parser.declareField(EsqlQueryRequest::params, EsqlQueryRequest::parseParams, PARAMS_FIELD, VALUE_ARRAY);
         parser.declareString((request, localeTag) -> request.locale(Locale.forLanguageTag(localeTag)), LOCALE_FIELD);
+        parser.declareInt((request, value) -> request.timeout(new TimeValue(value, TimeUnit.MILLISECONDS)), TIMEOUT_FIELD);
 
         return parser;
     }

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.plugin;
 
+import org.elasticsearch.ElasticsearchTimeoutException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionListenerResponseHandler;
 import org.elasticsearch.action.search.SearchRequest;
@@ -17,6 +18,7 @@
 import org.elasticsearch.action.support.ChannelActionListener;
 import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.action.support.ListenableActionFuture;
+import org.elasticsearch.action.support.ListenerTimeouts;
 import org.elasticsearch.action.support.RefCountingListener;
 import org.elasticsearch.action.support.RefCountingRunnable;
 import org.elasticsearch.cluster.node.DiscoveryNode;
@@ -139,12 +141,13 @@ public void execute(
         LOGGER.info("Sending data node plan\n{}\n with filter [{}]", dataNodePlan, requestFilter);
 
         String[] originalIndices = PlannerUtils.planOriginalIndices(physicalPlan);
+
         computeTargetNodes(
             rootTask,
             requestFilter,
             concreteIndices,
             originalIndices,
-            listener.delegateFailureAndWrap((delegate, targetNodes) -> {
+            wrapTimeout(configuration, transportService.getThreadPool(), listener.delegateFailureAndWrap((delegate, targetNodes) -> {
                 final ExchangeSourceHandler exchangeSource = exchangeService.createSourceHandler(
                     sessionId,
                     queryPragmas.exchangeBufferSize(),
@@ -172,8 +175,23 @@ public void execute(
                         () -> cancelOnFailure(rootTask, cancelled, requestRefs.acquire()).map(unused -> null)
                     );
                 }
-            })
+            }))
         );
+
+    }
+
+    private ActionListener<List<TargetNode>> wrapTimeout(
+        EsqlConfiguration configuration,
+        ThreadPool pool,
+        ActionListener<List<TargetNode>> l
+    ) {
+        if (configuration.timeout() != null) {
+            return ListenerTimeouts.wrapWithTimeout(pool, configuration.timeout(), esqlExecutor, l, (x) -> {
+                String timeoutMessage = "ESQL query timed out after {}";
+                l.onFailure(new ElasticsearchTimeoutException(timeoutMessage, configuration.timeout()));
+            });
+        }
+        return l;
     }
 
     private void runComputeOnRemoteNodes(

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
@@ -35,9 +36,12 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.concurrent.Executor;
+import java.util.concurrent.TimeUnit;
 
 public class TransportEsqlQueryAction extends HandledTransportAction<EsqlQueryRequest, EsqlQueryResponse> {
 
+    public static TimeValue DEFAULT_TIMEOUT = new TimeValue(4, TimeUnit.MINUTES);
+    public static TimeValue MAX_TIMEOUT = new TimeValue(4, TimeUnit.MINUTES);
     private final PlanExecutor planExecutor;
     private final ComputeService computeService;
     private final ExchangeService exchangeService;
@@ -84,16 +88,20 @@ protected void doExecute(Task task, EsqlQueryRequest request, ActionListener<Esq
     }
 
     private void doExecuteForked(Task task, EsqlQueryRequest request, ActionListener<EsqlQueryResponse> listener) {
+        TimeValue timeout = timeout(request);
+
         EsqlConfiguration configuration = new EsqlConfiguration(
             request.zoneId() != null ? request.zoneId() : ZoneOffset.UTC,
             request.locale() != null ? request.locale() : Locale.US,
             // TODO: plug-in security
             null,
             clusterService.getClusterName().value(),
             request.pragmas(),
-            EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.get(settings)
+            EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.get(settings),
+            timeout
         );
         String sessionId = sessionID(task);
+
         planExecutor.esql(
             request,
             sessionId,
@@ -110,6 +118,18 @@ private void doExecuteForked(Task task, EsqlQueryRequest request, ActionListener
         );
     }
 
+    private TimeValue timeout(EsqlQueryRequest request) {
+        TimeValue rTimeout = request.timeout();
+        if (rTimeout == null) {
+            return DEFAULT_TIMEOUT;
+        }
+        TimeValue timeout = rTimeout.compareTo(MAX_TIMEOUT) > 0 ? MAX_TIMEOUT : rTimeout;
+        if (timeout.duration() < 0) {
+            return MAX_TIMEOUT;
+        }
+        return timeout;
+    }
+
     /**
      * Returns the ID for this compute session. The ID is unique within the cluster, and is used
      * to identify the compute-session across nodes. The ID is just the TaskID of the task that

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 import org.elasticsearch.xpack.ql.session.Configuration;
 
@@ -26,25 +27,30 @@ public class EsqlConfiguration extends Configuration implements Writeable {
 
     private final Locale locale;
 
+    private final TimeValue timeout;
+
     public EsqlConfiguration(
         ZoneId zi,
         Locale locale,
         String username,
         String clusterName,
         QueryPragmas pragmas,
-        int resultTruncationMaxSize
+        int resultTruncationMaxSize,
+        TimeValue timeout
     ) {
         super(zi, username, clusterName);
         this.locale = locale;
         this.pragmas = pragmas;
         this.resultTruncationMaxSize = resultTruncationMaxSize;
+        this.timeout = timeout;
     }
 
     public EsqlConfiguration(StreamInput in) throws IOException {
         super(in.readZoneId(), Instant.ofEpochSecond(in.readVLong(), in.readVInt()), in.readOptionalString(), in.readOptionalString());
-        locale = Locale.forLanguageTag(in.readString());
+        this.locale = Locale.forLanguageTag(in.readString());
         this.pragmas = new QueryPragmas(in);
         this.resultTruncationMaxSize = in.readVInt();
+        this.timeout = in.readOptionalTimeValue();
     }
 
     @Override
@@ -58,6 +64,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeString(locale.toLanguageTag());
         pragmas.writeTo(out);
         out.writeVInt(resultTruncationMaxSize);
+        out.writeOptionalTimeValue(timeout);
     }
 
     public QueryPragmas pragmas() {
@@ -72,6 +79,10 @@ public Locale locale() {
         return locale;
     }
 
+    public TimeValue timeout() {
+        return timeout;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (super.equals(o)) {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java

@@ -151,7 +151,8 @@ public class CsvTests extends ESTestCase {
         null,
         null,
         new QueryPragmas(Settings.builder().put("page_size", randomPageSize()).build()),
-        EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(Settings.EMPTY)
+        EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(Settings.EMPTY),
+        null
     );
     private final FunctionRegistry functionRegistry = new EsqlFunctionRegistry();
     private final EsqlParser parser = new EsqlParser();

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java

@@ -137,7 +137,8 @@ public static List<Object[]> readScriptSpec() {
                     null,
                     null,
                     new QueryPragmas(settings),
-                    EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(settings)
+                    EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(settings),
+                    null
                 ) };
         }).toList();
     }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/EvalMapperTests.java

@@ -61,7 +61,15 @@ public class EvalMapperTests extends ESTestCase {
     private static final FieldAttribute LONG = field("long", DataTypes.LONG);
     private static final FieldAttribute DATE = field("date", DataTypes.DATETIME);
 
-    private static final EsqlConfiguration TEST_CONFIG = new EsqlConfiguration(ZoneOffset.UTC, Locale.US, "test", null, null, 10000000);
+    private static final EsqlConfiguration TEST_CONFIG = new EsqlConfiguration(
+        ZoneOffset.UTC,
+        Locale.US,
+        "test",
+        null,
+        null,
+        10000000,
+        null
+    );
 
     @ParametersFactory(argumentFormatting = "%1$s")
     public static List<Object[]> params() {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java

@@ -133,7 +133,8 @@ private EsqlConfiguration config() {
             "test_user",
             "test_cluser",
             pragmas,
-            EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(null)
+            EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(null),
+            null
         );
     }
 

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestTests.java

@@ -182,7 +182,8 @@ static PhysicalPlan mapAndMaybeOptimize(LogicalPlan logicalPlan) {
             null,
             null,
             new QueryPragmas(Settings.EMPTY),
-            EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(Settings.EMPTY)
+            EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(Settings.EMPTY),
+            null
         );
         var physicalPlanOptimizer = new PhysicalPlanOptimizer(new PhysicalOptimizerContext(configuration));
         FunctionRegistry functionRegistry = new EsqlFunctionRegistry();