Added error handling and set a completion listener to the action

Author: ivancea

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.rest.action.RestCancellableNodeClient;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.esql.action.stream.EsqlQueryResponseStream;
+import org.elasticsearch.xpack.esql.action.stream.NonStreamingEsqlQueryResponseStream;
 
 import java.io.IOException;
 import java.util.List;
@@ -53,21 +54,19 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
 
     protected static RestChannelConsumer restChannelConsumer(EsqlQueryRequest esqlRequest, RestRequest request, NodeClient client) {
         final Boolean partialResults = request.paramAsBoolean("allow_partial_results", null);
+        final boolean stream = request.paramAsBoolean("stream", false);
         if (partialResults != null) {
             esqlRequest.allowPartialResults(partialResults);
         }
         LOGGER.debug("Beginning execution of ESQL query.\nQuery string: [{}]", esqlRequest.query());
         // TODO: Create responseStream here, and add to the request object (?). See RestRepositoryVerifyIntegrityAction
         return channel -> {
-            final var responseStream = EsqlQueryResponseStream.forMediaType(channel, request);
+            final var responseStream = stream
+                ? EsqlQueryResponseStream.forMediaType(channel, request)
+                : new NonStreamingEsqlQueryResponseStream(channel, request, esqlRequest);
             esqlRequest.responseStream(responseStream);
             RestCancellableNodeClient cancellableClient = new RestCancellableNodeClient(client, request.getHttpChannel());
-            cancellableClient.execute(
-                EsqlQueryAction.INSTANCE,
-                esqlRequest,
-                // TODO: Replace this with a call to the responseStream
-                new EsqlResponseListener(channel, request, esqlRequest).wrapWithLogging()
-            );
+            cancellableClient.execute(EsqlQueryAction.INSTANCE, esqlRequest, responseStream.completionListener());
         };
     }
 

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

@@ -73,6 +73,11 @@ protected void doSendEverything(EsqlQueryResponse response) {
         doFinishResponse(response);
     }
 
+    @Override
+    protected void doHandleException(Exception e) {
+        // TODO: Implement this
+    }
+
     private void sendStartResponseDroppingNullColumns(List<ColumnInfoImpl> columns, List<Page> pages) {
         assert params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false) : "this method should only be called when dropping null columns";
 

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

@@ -7,10 +7,13 @@
 
 package org.elasticsearch.xpack.esql.action.stream;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.StreamingXContentResponse;
@@ -35,6 +38,7 @@
  * TODO: Took header wouldn't be available on streaming
  */
 public abstract class EsqlQueryResponseStream implements Releasable {
+    private static final Logger LOGGER = LogManager.getLogger(EsqlQueryResponseStream.class);
 
     public static EsqlQueryResponseStream forMediaType(RestChannel restChannel, RestRequest request) {
         MediaType mediaType = EsqlMediaTypeParser.getResponseMediaType(request, XContentType.JSON);
@@ -53,7 +57,7 @@ public static EsqlQueryResponseStream forMediaType(RestChannel restChannel, Rest
     private final RestChannel restChannel;
     protected final ToXContent.Params params;
     /**
-     * Initialized on the first call to {@link #startResponse()} and used to write the response chunks.
+     * Initialized on the first call to {@link #startResponse} and used to write the response chunks.
      */
     @Nullable
     private StreamingXContentResponse streamingXContentResponse;
@@ -79,7 +83,7 @@ protected EsqlQueryResponseStream(RestChannel restChannel, ToXContent.Params par
      */
     public final void startResponse(List<ColumnInfoImpl> columns) throws IOException {
         assert streamingXContentResponse == null : "startResponse() called more than once";
-        assert finished == false : "sendPages() called after finishResponse()";
+        assert finished == false : "sendPages() called on a finished stream";
 
         streamingXContentResponse = new StreamingXContentResponse(restChannel, restChannel.request(), () -> {});
 
@@ -93,7 +97,7 @@ public final void startResponse(List<ColumnInfoImpl> columns) throws IOException
 
     public final void sendPages(Iterable<Page> pages) {
         assert streamingXContentResponse != null : "sendPages() called before startResponse()";
-        assert finished == false : "sendPages() called after finishResponse()";
+        assert finished == false : "sendPages() called on a finished stream";
 
         if (initialStreamChunkSent) {
             doSendPages(pages);
@@ -111,6 +115,36 @@ public final void finishResponse(EsqlQueryResponse response) {
         finished = true;
     }
 
+    public final void handleException(Exception e) {
+        assert finished == false : "handleException() called on a finished stream";
+
+        // TODO: To be overridden by subclasses. This should append the error to the stream, if possible
+        LOGGER.error("Error while streaming response", e);
+
+        finished = true;
+    }
+
+    // TODO: For error handling, check RestActionListener error listener
+    // TODO: Also ensure that we check if the channel is closed at some points (Also see RestActionListener)
+
+    public final ActionListener<EsqlQueryResponse> completionListener() {
+        return new ActionListener<>() {
+            @Override
+            public void onResponse(EsqlQueryResponse esqlResponse) {
+                assert finished == false : "completionListener() called on a finished stream";
+
+                finishResponse(esqlResponse);
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                assert finished == false : "onFailure() called on a finished stream";
+
+                handleException(e);
+            }
+        };
+    }
+
     /**
      * Returns true if the response can be streamed, false otherwise.
      * <p>
@@ -126,6 +160,8 @@ public final void finishResponse(EsqlQueryResponse response) {
 
     protected abstract void doFinishResponse(EsqlQueryResponse response);
 
+    protected abstract void doHandleException(Exception e);
+
     protected void doSendEverything(EsqlQueryResponse response) {
         // TODO: Is this safe? Should this be abstract to ensure proper implementation? Add tests for both cases
         doStartResponse(response.columns());

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

@@ -31,7 +31,7 @@ public class NonStreamingEsqlQueryResponseStream extends EsqlQueryResponseStream
 
     private final ActionListener<EsqlQueryResponse> listener;
 
-    NonStreamingEsqlQueryResponseStream(RestChannel restChannel, RestRequest restRequest, EsqlQueryRequest esqlRequest) {
+    public NonStreamingEsqlQueryResponseStream(RestChannel restChannel, RestRequest restRequest, EsqlQueryRequest esqlRequest) {
         super(restChannel, restRequest);
 
         this.listener = new EsqlResponseListener(restChannel, restRequest, esqlRequest).wrapWithLogging();
@@ -57,6 +57,11 @@ protected void doFinishResponse(EsqlQueryResponse response) {
         throw new UnsupportedOperationException("This class does not support streaming");
     }
 
+    @Override
+    protected void doHandleException(Exception e) {
+        listener.onFailure(e);
+    }
+
     @Override
     protected void doSendEverything(EsqlQueryResponse response) {
         // The base class will close the response, and the listener will do so too.