Fix error handling in SQL's CLI (elastic/x-pack-elasticsearch#2730)

We weren't returning errors correctly from the server
or catching them correctly in the CLI. This fixes that
and adds simple integration tests.

Original commit: elastic/x-pack-elasticsearch@259da0da6f

qa/sql/multinode/src/test/java/org/elasticsearch/xpack/qa/sql/multinode/CliErrorsIT.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.qa.sql.multinode;
+
+import org.elasticsearch.xpack.qa.sql.cli.ErrorsTestCase;
+
+public class CliErrorsIT extends ErrorsTestCase {
+}

qa/sql/multinode/src/test/java/org/elasticsearch/xpack/qa/sql/multinode/CliFetchSizeIT.java

@@ -5,9 +5,7 @@
  */
 package org.elasticsearch.xpack.qa.sql.multinode;
 
-import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.elasticsearch.xpack.qa.sql.cli.FetchSizeTestCase;
 
-@TestLogging("org.elasticsearch.xpack.qa.sql.multinode:DEBUG")
 public class CliFetchSizeIT extends FetchSizeTestCase {
 }

qa/sql/no-security/src/test/java/org/elasticsearch/xpack/qa/sql/nosecurity/CliErrorsIT.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.qa.sql.nosecurity;
+
+import org.elasticsearch.xpack.qa.sql.cli.ErrorsTestCase;
+
+public class CliErrorsIT extends ErrorsTestCase {
+}

qa/sql/security/src/test/java/org/elasticsearch/xpack/qa/sql/security/CliErrorsIT.java

@@ -0,0 +1,19 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.qa.sql.security;
+
+import org.apache.lucene.util.LuceneTestCase.AwaitsFix;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.xpack.qa.sql.cli.ErrorsTestCase;
+
+@AwaitsFix(bugUrl = "https://github.com/elastic/x-pack-elasticsearch/issues/2074")
+public class CliErrorsIT extends ErrorsTestCase {
+    // NOCOMMIT get this working with security....
+    @Override
+    protected Settings restClientSettings() {
+        return RestSqlIT.securitySettings();
+    }
+}

qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/cli/ErrorsTestCase.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.qa.sql.cli;
+
+/**
+ * Tests for error messages.
+ */
+public abstract class ErrorsTestCase extends CliIntegrationTestCase {
+    public void testSelectFromMissingTable() throws Exception {
+        assertEquals("[1;31mBad request [[22;3;33mFound 1 problem(s)", command("SELECT * FROM test"));
+        assertEquals("line 1:15: Unknown index [test][1;23;31m][0m", readLine());
+    }
+}

qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/ErrorsTestCase.java

@@ -9,7 +9,7 @@ import java.sql.Connection;
 import java.sql.SQLException;
 
 /**
- * Tests for error messages.
+ * Tests for exceptions and their messages.
  */
 public class ErrorsTestCase extends JdbcIntegrationTestCase {
     public void testSelectFromMissingTable() throws Exception {

sql/cli-proto/src/main/java/org/elasticsearch/xpack/sql/cli/net/protocol/Proto.java

@@ -61,7 +61,7 @@ public final class Proto extends AbstractProto {
         }
     }
 
-    enum ResponseType implements AbstractProto.ResponseType {
+    public enum ResponseType implements AbstractProto.ResponseType {
         EXCEPTION(ExceptionResponse::new),
         ERROR(ErrorResponse::new),
         INFO(InfoResponse::new),

sql/cli/src/main/java/org/elasticsearch/xpack/sql/cli/Cli.java

@@ -6,10 +6,12 @@
 package org.elasticsearch.xpack.sql.cli;
 
 import org.elasticsearch.xpack.sql.cli.net.protocol.QueryResponse;
+import org.elasticsearch.xpack.sql.cli.net.protocol.Proto.ResponseType;
 import org.elasticsearch.xpack.sql.net.client.SuppressForbidden;
 import org.elasticsearch.xpack.sql.net.client.util.IOUtils;
 import org.elasticsearch.xpack.sql.net.client.util.StringUtils;
 import org.elasticsearch.xpack.sql.protocol.shared.AbstractQueryInitRequest;
+import org.elasticsearch.xpack.sql.protocol.shared.Response;
 import org.jline.reader.EndOfFileException;
 import org.jline.reader.LineReader;
 import org.jline.reader.LineReaderBuilder;
@@ -287,17 +289,21 @@ public class Cli {
     }
 
     private void executeQuery(String line) throws IOException {
-        QueryResponse response = cliClient.queryInit(line, fetchSize);
+        Response response = cliClient.queryInit(line, fetchSize);
         while (true) {
             term.writer().print(ResponseToString.toAnsi(response).toAnsi(term));
             term.writer().flush();
-            if (response.cursor().length == 0) {
+            if (response.responseType() == ResponseType.ERROR || response.responseType() == ResponseType.EXCEPTION) {
+                return;
+            }
+            QueryResponse queryResponse = (QueryResponse) response;
+            if (queryResponse.cursor().length == 0) {
                 return;
             }
             if (false == fetchSeparator.equals("")) {
                 term.writer().println(fetchSeparator);
             }
-            response = cliClient.nextPage(response.cursor());
+            response = cliClient.nextPage(queryResponse.cursor());
         }
     }
 

sql/cli/src/main/java/org/elasticsearch/xpack/sql/cli/CliHttpClient.java

@@ -9,9 +9,7 @@ import org.elasticsearch.xpack.sql.cli.net.protocol.InfoRequest;
 import org.elasticsearch.xpack.sql.cli.net.protocol.InfoResponse;
 import org.elasticsearch.xpack.sql.cli.net.protocol.Proto;
 import org.elasticsearch.xpack.sql.cli.net.protocol.QueryInitRequest;
-import org.elasticsearch.xpack.sql.cli.net.protocol.QueryInitResponse;
 import org.elasticsearch.xpack.sql.cli.net.protocol.QueryPageRequest;
-import org.elasticsearch.xpack.sql.cli.net.protocol.QueryPageResponse;
 import org.elasticsearch.xpack.sql.net.client.util.Bytes;
 import org.elasticsearch.xpack.sql.protocol.shared.Request;
 import org.elasticsearch.xpack.sql.protocol.shared.Response;
@@ -31,18 +29,19 @@ public class CliHttpClient implements AutoCloseable {
 
     public InfoResponse serverInfo() {
         InfoRequest request = new InfoRequest();
+        // TODO think about error handling here
         return (InfoResponse) sendRequest(request);
     }
 
-    public QueryInitResponse queryInit(String query, int fetchSize) {
+    public Response queryInit(String query, int fetchSize) {
         // TODO allow customizing the time zone - this is what session set/reset/get should be about
         QueryInitRequest request = new QueryInitRequest(query, fetchSize, TimeZone.getTimeZone("UTC"), new TimeoutInfo(0, 0, 0));
-        return (QueryInitResponse) sendRequest(request);
+        return sendRequest(request);
     }
 
-    public QueryPageResponse nextPage(byte[] cursor) {
+    public Response nextPage(byte[] cursor) {
         QueryPageRequest request = new QueryPageRequest(cursor, new TimeoutInfo(0, 0, 0));
-        return (QueryPageResponse) sendRequest(request);
+        return sendRequest(request);
     }
 
     private Response sendRequest(Request request) {

sql/cli/src/main/java/org/elasticsearch/xpack/sql/cli/ResponseToString.java

@@ -8,6 +8,7 @@ package org.elasticsearch.xpack.sql.cli;
 import org.elasticsearch.xpack.sql.cli.net.protocol.ErrorResponse;
 import org.elasticsearch.xpack.sql.cli.net.protocol.ExceptionResponse;
 import org.elasticsearch.xpack.sql.cli.net.protocol.InfoResponse;
+import org.elasticsearch.xpack.sql.cli.net.protocol.Proto.ResponseType;
 import org.elasticsearch.xpack.sql.cli.net.protocol.QueryResponse;
 import org.elasticsearch.xpack.sql.net.client.SuppressForbidden;
 import org.elasticsearch.xpack.sql.protocol.shared.Response;
@@ -20,17 +21,19 @@ import java.nio.file.Path;
 
 import static org.jline.utils.AttributedStyle.BOLD;
 import static org.jline.utils.AttributedStyle.BRIGHT;
-import static org.jline.utils.AttributedStyle.CYAN;
 import static org.jline.utils.AttributedStyle.DEFAULT;
 import static org.jline.utils.AttributedStyle.RED;
 import static org.jline.utils.AttributedStyle.WHITE;
+import static org.jline.utils.AttributedStyle.YELLOW;
 
 abstract class ResponseToString {
 
     static AttributedStringBuilder toAnsi(Response response) {
         AttributedStringBuilder sb = new AttributedStringBuilder();
 
-        if (response instanceof QueryResponse) {
+        switch ((ResponseType) response.responseType()) {
+        case QUERY_INIT:
+        case QUERY_PAGE:
             QueryResponse cmd = (QueryResponse) response;
             if (cmd.data != null) {
                 String data = cmd.data.toString();
@@ -41,12 +44,8 @@ abstract class ResponseToString {
                     sb.append(data, DEFAULT.foreground(WHITE));
                 }
             }
-        }
+            return sb;
-        else if (response instanceof ExceptionResponse) {
+        case INFO:
-            ExceptionResponse ex = (ExceptionResponse) response;
-            sb.append(ex.message, BOLD.foreground(CYAN));
-        }
-        else if (response instanceof InfoResponse) {
             InfoResponse info = (InfoResponse) response;
             sb.append("Node:", DEFAULT.foreground(BRIGHT));
             sb.append(info.node, DEFAULT.foreground(WHITE));
@@ -54,17 +53,24 @@ abstract class ResponseToString {
             sb.append(info.cluster, DEFAULT.foreground(WHITE));
             sb.append(" Version:", DEFAULT.foreground(BRIGHT));
             sb.append(info.versionString, DEFAULT.foreground(WHITE));
+            return sb;
+        case ERROR:
+            ErrorResponse err = (ErrorResponse) response;
+            error("Server error", err.message, sb);
+            return sb;
+        case EXCEPTION:
+            ExceptionResponse ex = (ExceptionResponse) response;
+            error("Bad request", ex.message, sb);
+            return sb;
+        default:
+            throw new IllegalArgumentException("Unsupported response: " + response);
         }
-        else if (response instanceof ErrorResponse) {
+    }
-            ErrorResponse error = (ErrorResponse) response;
-            sb.append("Server error:", BOLD.foreground(RED));
-            sb.append(error.message, DEFAULT.italic().foreground(RED));
-        }
-        else {
-            sb.append("Invalid response received from server...", BOLD.foreground(RED));
-        }
 
-        return sb;
+    private static void error(String type, String message, AttributedStringBuilder sb) {
+        sb.append(type + " [", BOLD.foreground(RED));
+        sb.append(message, DEFAULT.boldOff().italic().foreground(YELLOW));
+        sb.append("]", BOLD.underlineOff().foreground(RED));
     }
 
     // NOCOMMIT - is using the default temp folder a problem?

sql/cli/src/test/java/org/elasticsearch/xpack/sql/cli/ResponseToStringTests.java

@@ -33,8 +33,8 @@ public class ResponseToStringTests extends ESTestCase {
     public void testExceptionResponse() {
         AttributedStringBuilder s = ResponseToString.toAnsi(new ExceptionResponse(RequestType.INFO, "test message", "test cause",
                 randomFrom(SqlExceptionType.values())));
-        assertEquals("test message", unstyled(s));
+        assertEquals("Bad request [test message]", unstyled(s));
-        assertEquals("[1;36mtest message[0m", fullyStyled(s));
+        assertEquals("[1;31mBad request [[22;3;33mtest message[1;23;31m][0m", fullyStyled(s));
     }
 
     private String unstyled(AttributedStringBuilder s) {

sql/server/src/main/java/org/elasticsearch/xpack/sql/plugin/RestSqlCliAction.java

@@ -6,65 +6,67 @@
 package org.elasticsearch.xpack.sql.plugin;
 
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.main.MainAction;
 import org.elasticsearch.action.main.MainRequest;
 import org.elasticsearch.client.Client;
-import org.elasticsearch.client.node.NodeClient;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.rest.BaseRestHandler;
-import org.elasticsearch.rest.BytesRestResponse;
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestController;
-import org.elasticsearch.rest.RestRequest;
-import org.elasticsearch.rest.action.RestActionListener;
 import org.elasticsearch.xpack.sql.cli.net.protocol.InfoResponse;
 import org.elasticsearch.xpack.sql.cli.net.protocol.Proto;
 import org.elasticsearch.xpack.sql.cli.net.protocol.Proto.RequestType;
+import org.elasticsearch.xpack.sql.cli.net.protocol.ErrorResponse;
+import org.elasticsearch.xpack.sql.cli.net.protocol.ExceptionResponse;
 import org.elasticsearch.xpack.sql.cli.net.protocol.QueryInitRequest;
 import org.elasticsearch.xpack.sql.cli.net.protocol.QueryInitResponse;
 import org.elasticsearch.xpack.sql.cli.net.protocol.QueryPageRequest;
 import org.elasticsearch.xpack.sql.cli.net.protocol.QueryPageResponse;
 import org.elasticsearch.xpack.sql.plugin.sql.action.SqlAction;
 import org.elasticsearch.xpack.sql.plugin.sql.action.SqlRequest;
+import org.elasticsearch.xpack.sql.protocol.shared.AbstractProto.SqlExceptionType;
 import org.elasticsearch.xpack.sql.protocol.shared.Request;
-import org.elasticsearch.xpack.sql.protocol.shared.Response;
 import org.elasticsearch.xpack.sql.session.Cursor;
 import org.joda.time.DateTimeZone;
 
-import java.io.DataInputStream;
-import java.io.DataOutputStream;
 import java.io.IOException;
 import java.util.function.Consumer;
-import java.util.function.Function;
 
-import static org.elasticsearch.rest.BytesRestResponse.TEXT_CONTENT_TYPE;
 import static org.elasticsearch.rest.RestRequest.Method.POST;
-import static org.elasticsearch.rest.RestStatus.OK;
 
-public class RestSqlCliAction extends BaseRestHandler {
+public class RestSqlCliAction extends AbstractSqlProtocolRestAction {
     private final NamedWriteableRegistry cursorRegistry = new NamedWriteableRegistry(Cursor.getNamedWriteables());
 
     public RestSqlCliAction(Settings settings, RestController controller) {
-        super(settings);
+        super(settings, Proto.INSTANCE);
         controller.registerHandler(POST, "/_sql/cli", this);
     }
 
     @Override
-    protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient client) throws IOException {
+    public String getName() {
-        Request request;
+        return "xpack_sql_cli_action";
-        try (DataInputStream in = new DataInputStream(restRequest.content().streamInput())) {
+    }
-            request = Proto.INSTANCE.readRequest(in);
+
-        }
+    @Override
+    protected RestChannelConsumer innerPrepareRequest(Request request, Client client) throws IOException {
         Consumer<RestChannel> consumer = operation(request, client);
         return consumer::accept;
     }
 
+    @Override
+    protected ErrorResponse buildErrorResponse(Request request, String message, String cause, String stack) {
+        return new ErrorResponse((RequestType) request.requestType(), message, cause, stack);
+    }
+
+    @Override
+    protected ExceptionResponse buildExceptionResponse(Request request, String message, String cause, SqlExceptionType exceptionType) {
+        return new ExceptionResponse((RequestType) request.requestType(), message, cause, exceptionType);
+    }
+
     /**
      * Actual implementation of the operation
      */
@@ -73,7 +75,7 @@ public class RestSqlCliAction extends BaseRestHandler {
         RequestType requestType = (RequestType) request.requestType();
         switch (requestType) {
         case INFO:
-            return channel -> client.execute(MainAction.INSTANCE, new MainRequest(), toActionListener(channel, response ->
+            return channel -> client.execute(MainAction.INSTANCE, new MainRequest(), toActionListener(request, channel, response ->
                     new InfoResponse(response.getNodeName(), response.getClusterName().value(),
                             response.getVersion().major, response.getVersion().minor, response.getVersion().toString(),
                             response.getBuild().shortHash(), response.getBuild().date())));
@@ -86,11 +88,11 @@ public class RestSqlCliAction extends BaseRestHandler {
         }
     }
 
-    private static Consumer<RestChannel> queryInit(Client client, QueryInitRequest request) {
+    private Consumer<RestChannel> queryInit(Client client, QueryInitRequest request) {
         // TODO time zone support for CLI
         SqlRequest sqlRequest = new SqlRequest(request.query, DateTimeZone.forTimeZone(request.timeZone), request.fetchSize, Cursor.EMPTY);
         long start = System.nanoTime();
-        return channel -> client.execute(SqlAction.INSTANCE, sqlRequest, toActionListener(channel, response -> {
+        return channel -> client.execute(SqlAction.INSTANCE, sqlRequest, toActionListener(request, channel, response -> {
             CliFormatter formatter = new CliFormatter(response);
             String data = formatter.formatWithHeader(response);
             return new QueryInitResponse(System.nanoTime() - start, serializeCursor(response.cursor(), formatter), data);
@@ -108,7 +110,7 @@ public class RestSqlCliAction extends BaseRestHandler {
         }
         SqlRequest sqlRequest = new SqlRequest("", SqlRequest.DEFAULT_TIME_ZONE, -1, cursor);
         long start = System.nanoTime();
-        return channel -> client.execute(SqlAction.INSTANCE, sqlRequest, toActionListener(channel, response -> {
+        return channel -> client.execute(SqlAction.INSTANCE, sqlRequest, toActionListener(request, channel, response -> {
             String data = formatter.formatWithoutHeader(response);
             return new QueryPageResponse(System.nanoTime() - start, serializeCursor(response.cursor(), formatter), data);
         }));
@@ -126,26 +128,4 @@ public class RestSqlCliAction extends BaseRestHandler {
             throw new RuntimeException("unexpected trouble building the cursor", e);
         }
     }
-
-    private static <T> ActionListener<T> toActionListener(RestChannel channel, Function<T, Response> responseBuilder) {
-        // NOCOMMIT error response
-        return new RestActionListener<T>(channel) {
-            @Override
-            protected void processResponse(T actionResponse) throws Exception {
-                Response response = responseBuilder.apply(actionResponse);
-                try (BytesStreamOutput bytesStreamOutput = new BytesStreamOutput()) {
-                    try (DataOutputStream dataOutputStream = new DataOutputStream(bytesStreamOutput)) {
-                        // NOCOMMIT use the version from the client
-                        Proto.INSTANCE.writeResponse(response, Proto.CURRENT_VERSION, dataOutputStream);
-                    }
-                    channel.sendResponse(new BytesRestResponse(OK, TEXT_CONTENT_TYPE, bytesStreamOutput.bytes()));
-                }
-            }
-        };
-    }
-
-    @Override
-    public String getName() {
-        return "xpack_sql_cli_action";
-    }
 }

sql/server/src/main/java/org/elasticsearch/xpack/sql/plugin/RestSqlJdbcAction.java

@@ -37,7 +37,6 @@ import org.elasticsearch.xpack.sql.jdbc.net.protocol.QueryPageResponse;
 import org.elasticsearch.xpack.sql.plugin.sql.action.SqlAction;
 import org.elasticsearch.xpack.sql.plugin.sql.action.SqlRequest;
 import org.elasticsearch.xpack.sql.plugin.sql.action.SqlResponse;
-import org.elasticsearch.xpack.sql.protocol.shared.AbstractExceptionResponse;
 import org.elasticsearch.xpack.sql.protocol.shared.AbstractProto.SqlExceptionType;
 import org.elasticsearch.xpack.sql.protocol.shared.Request;
 import org.elasticsearch.xpack.sql.session.Cursor;
@@ -85,8 +84,7 @@ public class RestSqlJdbcAction extends AbstractSqlProtocolRestAction {
     }
 
     @Override
-    protected AbstractExceptionResponse buildExceptionResponse(Request request, String message, String cause,
+    protected ExceptionResponse buildExceptionResponse(Request request, String message, String cause, SqlExceptionType exceptionType) {
-            SqlExceptionType exceptionType) {
         return new ExceptionResponse((RequestType) request.requestType(), message, cause, exceptionType);
     }