Introduce RemoteComputeException that forwards the cause's status code

Author: pawankartik-elastic

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

@@ -284,6 +284,6 @@ public void testCancelSkipUnavailable() throws Exception {
         }
 
         Exception error = expectThrows(Exception.class, requestFuture::actionGet);
-        assertThat(error, instanceOf(TaskCancelledException.class));
+        assertThat(error.getCause(), instanceOf(TaskCancelledException.class));
     }
 }

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

@@ -29,6 +29,7 @@
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
+import org.elasticsearch.xpack.esql.RemoteComputeException;
 import org.elasticsearch.xpack.esql.VerificationException;
 import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 
@@ -813,10 +814,13 @@ public void testRemoteFailureSkipUnavailableTrue() throws IOException {
         Map<String, Object> testClusterInfo = setupFailClusters();
         String localIndex = (String) testClusterInfo.get("local.index");
         String remote1Index = (String) testClusterInfo.get("remote.index");
-        int localNumShards = (Integer) testClusterInfo.get("local.num_shards");
         String q = Strings.format("FROM %s,cluster-a:%s*", localIndex, remote1Index);
-        IllegalStateException e = expectThrows(IllegalStateException.class, () -> runQuery(q, false));
-        assertThat(e.getMessage(), containsString("Accessing failing field"));
+
+        RemoteComputeException rce = expectThrows(RemoteComputeException.class, () -> runQuery(q, false));
+        Throwable t = rce.getCause();
+
+        assertThat(t, instanceOf(IllegalStateException.class));
+        assertThat(t.getMessage(), containsString("Accessing failing field"));
     }
 
     private static void assertClusterMetadataInResponse(EsqlQueryResponse resp, boolean responseExpectMeta) {

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

@@ -24,6 +24,7 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
+import org.elasticsearch.xpack.esql.RemoteComputeException;
 import org.elasticsearch.xpack.esql.plugin.ComputeService;
 
 import java.io.IOException;
@@ -40,6 +41,7 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.in;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
@@ -73,8 +75,11 @@ public void testPartialResults() throws Exception {
         request.includeCCSMetadata(randomBoolean());
         {
             request.allowPartialResults(false);
-            IllegalStateException error = expectThrows(IllegalStateException.class, () -> runQuery(request).close());
-            assertThat(error.getMessage(), containsString("Accessing failing field"));
+            RemoteComputeException rce = expectThrows(RemoteComputeException.class, () -> runQuery(request).close());
+            Throwable t = rce.getCause();
+
+            assertThat(t, instanceOf(IllegalStateException.class));
+            assertThat(t.getMessage(), containsString("Accessing failing field"));
         }
         request.allowPartialResults(true);
         try (var resp = runQuery(request)) {

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

@@ -0,0 +1,38 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.xpack.esql;
+
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.rest.RestStatus;
+
+import java.util.Objects;
+
+/**
+ * Represents an error that occurred when starting compute on a remote node.
+ * It allows capturing some context such as the cluster alias that encountered the error.
+ */
+public class RemoteComputeException extends ElasticsearchException {
+
+    /**
+     * @param clusterAlias Name of the cluster.
+     * @param cause Error that was encountered.
+     */
+    public RemoteComputeException(String clusterAlias, Throwable cause) {
+        super("Remote [" + clusterAlias + "] encountered an error", cause);
+        Objects.requireNonNull(cause);
+    }
+
+    @Override
+    public RestStatus status() {
+        // This is similar to what we do in SearchPhaseExecutionException.
+        return ExceptionsHelper.status(getCause());
+    }
+}

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

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.plugin;
 
-import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.OriginalIndices;
 import org.elasticsearch.action.search.SearchRequest;
@@ -41,6 +40,7 @@
 import org.elasticsearch.transport.RemoteClusterAware;
 import org.elasticsearch.transport.TransportRequest;
 import org.elasticsearch.transport.TransportService;
+import org.elasticsearch.xpack.esql.RemoteComputeException;
 import org.elasticsearch.xpack.esql.action.EsqlExecutionInfo;
 import org.elasticsearch.xpack.esql.action.EsqlQueryAction;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
@@ -301,9 +301,7 @@ public void execute(
                         cancelQueryOnFailure,
                         execInfo,
                         computeListener.acquireCompute()
-                            .delegateResponse(
-                                (l, ex) -> l.onFailure(new ElasticsearchException(cluster.clusterAlias() + " encountered an error", ex))
-                            )
+                            .delegateResponse((l, ex) -> l.onFailure(new RemoteComputeException(cluster.clusterAlias(), ex)))
                     );
                 }
             }