Implement support for partial search results in SQL CLI (#86982)

Add support for

   allow_partial_search_results = true/false

command in SQL CLI.

If true, returns partial results if there are shard request timeouts or shard failures.
If false, returns an error with no partial results.

docs/changelog/86982.yaml

@@ -0,0 +1,6 @@
+pr: 86982
+summary: Implement support for partial search results in SQL CLI
+area: SQL
+type: enhancement
+issues:
+ - 86082

docs/reference/sql/endpoints/cli.asciidoc

@@ -64,4 +64,138 @@ $ ./java -cp [PATH_TO_CLI_JAR]/elasticsearch-sql-cli-[VERSION].jar org.elasticse
 The jar name will be different for each Elasticsearch version (for example `elasticsearch-sql-cli-7.3.2.jar`),
 thus the generic `VERSION` specified in the example above. Furthermore,
 if not running the command from the folder where the SQL CLI jar resides,
-you'd have to provide the full path, as well.
+you'd have to provide the full path, as well.
+
+
+
+[[cli-commands]]
+[discrete]
+=== CLI commands
+
+Apart from SQL queries, CLI can also execute some specific commands:
+
+
+`allow_partial_search_results = <boolean>` (default `false`)::
+
+If `true`, returns partial results if there are shard request timeouts or
+<<shard-failures,shard failures>>. If `false`, returns an error with
+no partial results.
+
+[source,sqlcli]
+--------------------------------------------------
+sql> allow_partial_search_results = true;
+allow_partial_search_results set to true
+--------------------------------------------------
+
+
+`fetch_size = <number>` (default `1000`)::
+
+Allows to change the size of fetches for query execution.
+Each fetch is delimited by fetch separator (if explicitly set).
+
+[source,sqlcli]
+--------------------------------------------------
+sql> fetch_size = 2000;
+fetch size set to 2000
+--------------------------------------------------
+
+
+
+`fetch_separator = <string>` (empty string by default)::
+
+Allows to change the separator string between fetches.
+
+[source,sqlcli]
+--------------------------------------------------
+sql> fetch_separator = "---------------------";
+fetch separator set to "---------------------"
+--------------------------------------------------
+
+
+`lenient = <boolean>` (default `false`)::
+
+If `false`, {es-sql} returns an error for fields containing <<array,array values>>.
+If `true`, {es-sql} returns the first value from the array with no guarantee of consistent results.
+
+[source,sqlcli]
+--------------------------------------------------
+sql> lenient = true;
+lenient set to true
+--------------------------------------------------
+
+
+
+`info`::
+
+Returns server information.
+
+[source,sqlcli]
+--------------------------------------------------
+sql> info;
+Node:mynode Cluster:elasticsearch Version:8.3
+--------------------------------------------------
+
+
+
+
+`exit`::
+
+Closes the CLI.
+
+[source,sqlcli]
+--------------------------------------------------
+sql> exit;
+Bye!
+--------------------------------------------------
+
+
+
+`cls`::
+
+Clears the screen.
+
+[source,sqlcli]
+--------------------------------------------------
+sql> cls;
+--------------------------------------------------
+
+
+
+`logo`::
+
+Prints Elastic logo.
+
+[source,sqlcli]
+--------------------------------------------------
+sql> logo;
+
+                       asticElasticE
+                     ElasticE  sticEla
+          sticEl  ticEl            Elast
+        lasti Elasti                   tic
+      cEl       ast                     icE
+     icE        as                       cEl
+     icE        as                       cEl
+     icEla     las                        El
+   sticElasticElast                     icElas
+ las           last                    ticElast
+El              asti                 asti    stic
+El              asticEla           Elas        icE
+El            Elas  cElasticE   ticEl           cE
+Ela        ticEl         ticElasti              cE
+ las     astic               last              icE
+   sticElas                   asti           stic
+     icEl                      sticElasticElast
+     icE                       sticE   ticEla
+     icE                       sti       cEla
+     icEl                      sti        Ela
+      cEl                      sti       cEl
+       Ela                    astic    ticE
+         asti               ElasticElasti
+           ticElasti  lasticElas
+              ElasticElast
+
+                       SQL
+                      8.3.0
+
+--------------------------------------------------

x-pack/plugin/sql/qa/server/security/src/test/java/org/elasticsearch/xpack/sql/qa/security/CliPartialResultsIT.java

@@ -0,0 +1,28 @@
+/*
+ * 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.sql.qa.security;
+
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.xpack.sql.qa.cli.EmbeddedCli.SecurityConfig;
+import org.elasticsearch.xpack.sql.qa.cli.PartialResultsTestCase;
+
+public class CliPartialResultsIT extends PartialResultsTestCase {
+    @Override
+    protected Settings restClientSettings() {
+        return RestSqlIT.securitySettings();
+    }
+
+    @Override
+    protected String getProtocol() {
+        return RestSqlIT.SSL_ENABLED ? "https" : "http";
+    }
+
+    @Override
+    protected SecurityConfig securityConfig() {
+        return CliSecurityIT.adminSecurityConfig();
+    }
+}

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

@@ -0,0 +1,13 @@
+/*
+ * 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.sql.qa.single_node;
+
+import org.elasticsearch.xpack.sql.qa.cli.PartialResultsTestCase;
+
+public class CliPartialResultsIT extends PartialResultsTestCase {
+
+}

x-pack/plugin/sql/qa/server/src/main/java/org/elasticsearch/xpack/sql/qa/cli/PartialResultsTestCase.java

@@ -0,0 +1,128 @@
+/*
+ * 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.sql.qa.cli;
+
+import org.elasticsearch.client.Request;
+
+import java.io.IOException;
+import java.util.Locale;
+
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.startsWith;
+
+public abstract class PartialResultsTestCase extends CliIntegrationTestCase {
+
+    private void createTestIndex(int okShards, int badShards) throws IOException {
+        final String mappingTemplate = """
+            {
+              "aliases": {
+                "test": {}
+              },
+              "mappings": {
+                "properties": {
+                  "bool": {
+                    "type": "boolean",
+                      "index": %s,
+                      "doc_values": %s
+                  }
+                }
+              }
+            }""";
+
+        // must match org.elasticsearch.xpack.sql.execution.search.Querier.BaseActionListener.MAX_WARNING_HEADERS
+        final int maxWarningHeaders = 20;
+
+        for (int i = 0; i < maxWarningHeaders - 1 + okShards + badShards; i++) {
+            String indexName = "/test" + i;
+            Request request = new Request("PUT", indexName);
+            boolean indexWithDocVals = i < okShards;
+            request.setJsonEntity(String.format(Locale.ROOT, mappingTemplate, indexWithDocVals, indexWithDocVals));
+            assertOK(client().performRequest(request));
+
+            request = new Request("POST", indexName + "/_doc");
+            request.addParameter("refresh", "true");
+            request.setJsonEntity("{\"bool\": " + (indexWithDocVals || randomBoolean()) + "}");
+            assertOK(client().performRequest(request));
+        }
+    }
+
+    public void testSetAllowPartialResults() throws Exception {
+        createTestIndex(1, 1);
+
+        // default value is false
+        String result = command("SELECT * FROM test WHERE bool = true");
+        assertThat(
+            result,
+            startsWith(
+                "[?1l>[?1000l[?2004l[31;1mServer error [[3;33;22mServer sent bad type [search_phase_execution_exception]. "
+                    + "Original type was [Partial shards failure]. [Failed to execute phase [query], Partial shards failure;"
+            )
+        );
+        consumeStackTrace();
+
+        result = command("allow_partial_search_results = false");
+        assertEquals("[?1l>[?1000l[?2004lallow_partial_search_results set to [90mfalse[0m", result);
+        result = command("SELECT * FROM test WHERE bool = true");
+        assertThat(
+            result,
+            startsWith(
+                "[?1l>[?1000l[?2004l[31;1mServer error [[3;33;22mServer sent bad type [search_phase_execution_exception]. "
+                    + "Original type was [Partial shards failure]. [Failed to execute phase [query], Partial shards failure;"
+            )
+        );
+        consumeStackTrace();
+
+        result = command("allow_partial_search_results = true");
+        assertEquals("[?1l>[?1000l[?2004lallow_partial_search_results set to [90mtrue[0m", result);
+        result = command("SELECT * FROM test WHERE bool = true");
+        assertEquals("[?1l>[?1000l[?2004l     bool      ", result);
+        while (readLine().length() > 0)
+            ;
+
+        result = command("allow_partial_search_results = false");
+        assertEquals("[?1l>[?1000l[?2004lallow_partial_search_results set to [90mfalse[0m", result);
+        result = command("SELECT * FROM test WHERE bool = true");
+        assertThat(
+            result,
+            startsWith(
+                "[?1l>[?1000l[?2004l[31;1mServer error [[3;33;22mServer sent bad type [search_phase_execution_exception]. "
+                    + "Original type was [Partial shards failure]. [Failed to execute phase [query], Partial shards failure;"
+            )
+        );
+        consumeStackTrace();
+    }
+
+    public void testPartialResultHandling() throws Exception {
+        final int okShards = randomIntBetween(1, 5);
+        final int extraBadShards = randomIntBetween(1, 5);
+        createTestIndex(okShards, extraBadShards);
+
+        String query = "SELECT * FROM test WHERE bool=true";
+        String result = command("allow_partial_search_results=true");
+        assertEquals("[?1l>[?1000l[?2004lallow_partial_search_results set to [90mtrue[0m", result);
+
+        result = command(query);
+        assertEquals("[?1l>[?1000l[?2004l     bool      ", result);
+        result = readLine();
+        assertEquals("---------------", result);
+        int rowCount = 0;
+        result = readLine();
+        while (result.isBlank() == false) {
+            result = readLine();
+            rowCount++;
+        }
+        assertThat(rowCount, greaterThanOrEqualTo(okShards));
+    }
+
+    private void consumeStackTrace() throws IOException {
+        String line;
+        do {
+            line = readLine();
+        } while (line.startsWith("][") == false);
+    }
+}

x-pack/plugin/sql/sql-cli/src/main/java/org/elasticsearch/xpack/sql/cli/Cli.java

@@ -14,6 +14,7 @@ import org.elasticsearch.cli.ExitCodes;
 import org.elasticsearch.cli.ProcessInfo;
 import org.elasticsearch.cli.Terminal;
 import org.elasticsearch.cli.UserException;
+import org.elasticsearch.xpack.sql.cli.command.AllowPartialResultsCliCommand;
 import org.elasticsearch.xpack.sql.cli.command.ClearScreenCliCommand;
 import org.elasticsearch.xpack.sql.cli.command.CliCommand;
 import org.elasticsearch.xpack.sql.cli.command.CliCommands;
@@ -131,6 +132,7 @@ public class Cli extends Command {
             new ClearScreenCliCommand(),
             new FetchSizeCliCommand(),
             new LenientCliCommand(),
+            new AllowPartialResultsCliCommand(),
             new FetchSeparatorCliCommand(),
             new ServerInfoCliCommand(),
             new ServerQueryCliCommand()

x-pack/plugin/sql/sql-cli/src/main/java/org/elasticsearch/xpack/sql/cli/command/AllowPartialResultsCliCommand.java

@@ -0,0 +1,32 @@
+/*
+ * 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.sql.cli.command;
+
+import org.elasticsearch.xpack.sql.cli.CliTerminal;
+
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * allows to enable/disable return of partial results.
+ * If true, returns partial results if there are shard request timeouts or shard failures.
+ * If false, returns an error with no partial results.
+ *
+ */
+public class AllowPartialResultsCliCommand extends AbstractCliCommand {
+
+    public AllowPartialResultsCliCommand() {
+        super(Pattern.compile("allow(?: |_)partial(?: |_)search(?: |_)results *= *(.+)", Pattern.CASE_INSENSITIVE));
+    }
+
+    @Override
+    protected boolean doHandle(CliTerminal terminal, CliSession cliSession, Matcher m, String line) {
+        cliSession.cfg().setAllowPartialResults(Boolean.parseBoolean(m.group(1)));
+        terminal.line().text("allow_partial_search_results set to ").em(Boolean.toString(cliSession.cfg().allowPartialResults())).end();
+        return true;
+    }
+}

x-pack/plugin/sql/sql-cli/src/main/java/org/elasticsearch/xpack/sql/cli/command/CliSessionConfiguration.java

@@ -17,10 +17,12 @@ public class CliSessionConfiguration {
     private String fetchSeparator = "";
     private boolean debug;
     private boolean lenient;
+    private boolean allowPartialResults;
 
     public CliSessionConfiguration() {
         this.fetchSize = CoreProtocol.FETCH_SIZE;
         this.lenient = CoreProtocol.FIELD_MULTI_VALUE_LENIENCY;
+        this.allowPartialResults = CoreProtocol.ALLOW_PARTIAL_SEARCH_RESULTS;
     }
 
     public void setFetchSize(int fetchSize) {
@@ -57,4 +59,12 @@ public class CliSessionConfiguration {
     public void setLenient(boolean lenient) {
         this.lenient = lenient;
     }
+
+    public boolean allowPartialResults() {
+        return allowPartialResults;
+    }
+
+    public void setAllowPartialResults(boolean allowPartialResults) {
+        this.allowPartialResults = allowPartialResults;
+    }
 }

x-pack/plugin/sql/sql-cli/src/main/java/org/elasticsearch/xpack/sql/cli/command/ServerQueryCliCommand.java

@@ -26,7 +26,8 @@ public class ServerQueryCliCommand extends AbstractServerCliCommand {
         SimpleFormatter formatter;
         String data;
         try {
-            response = cliClient.basicQuery(line, cliSession.cfg().getFetchSize(), cliSession.cfg().isLenient());
+            CliSessionConfiguration cfg = cliSession.cfg();
+            response = cliClient.basicQuery(line, cfg.getFetchSize(), cfg.isLenient(), cfg.allowPartialResults());
             formatter = new SimpleFormatter(response.columns(), response.rows(), CLI);
             data = formatter.formatWithHeader(response.columns(), response.rows());
             while (true) {

x-pack/plugin/sql/sql-cli/src/test/java/org/elasticsearch/xpack/sql/cli/command/BuiltinCommandTests.java

@@ -115,4 +115,21 @@ public class BuiltinCommandTests extends SqlCliTestCase {
         testTerminal.clear();
     }
 
+    public void testAllowPartialSearchResults() {
+        TestTerminal testTerminal = new TestTerminal();
+        HttpClient httpClient = mock(HttpClient.class);
+        CliSession cliSession = new CliSession(httpClient);
+        AllowPartialResultsCliCommand cliCommand = new AllowPartialResultsCliCommand();
+        assertFalse(cliCommand.handle(testTerminal, cliSession, "allow_partial_search_results"));
+        assertEquals(false, cliSession.cfg().allowPartialResults());
+        assertTrue(cliCommand.handle(testTerminal, cliSession, "allow_partial_search_results = true"));
+        assertEquals(true, cliSession.cfg().allowPartialResults());
+        assertEquals("allow_partial_search_results set to <em>true</em><flush/>", testTerminal.toString());
+        testTerminal.clear();
+        assertTrue(cliCommand.handle(testTerminal, cliSession, "allow_partial_search_results = false"));
+        assertEquals(false, cliSession.cfg().allowPartialResults());
+        assertEquals("allow_partial_search_results set to <em>false</em><flush/>", testTerminal.toString());
+        testTerminal.clear();
+    }
+
 }

x-pack/plugin/sql/sql-cli/src/test/java/org/elasticsearch/xpack/sql/cli/command/ServerQueryCliCommandTests.java

@@ -32,11 +32,11 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
         TestTerminal testTerminal = new TestTerminal();
         HttpClient client = mock(HttpClient.class);
         CliSession cliSession = new CliSession(client);
-        when(client.basicQuery("blah", 1000, false)).thenThrow(new SQLException("test exception"));
+        when(client.basicQuery("blah", 1000, false, false)).thenThrow(new SQLException("test exception"));
         ServerQueryCliCommand cliCommand = new ServerQueryCliCommand();
         assertTrue(cliCommand.handle(testTerminal, cliSession, "blah"));
         assertEquals("<b>Bad request [</b><i>test exception</i><b>]</b>\n", testTerminal.toString());
-        verify(client, times(1)).basicQuery(eq("blah"), eq(1000), eq(false));
+        verify(client, times(1)).basicQuery(eq("blah"), eq(1000), eq(false), eq(false));
         verifyNoMoreInteractions(client);
     }
 
@@ -45,7 +45,7 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
         HttpClient client = mock(HttpClient.class);
         CliSession cliSession = new CliSession(client);
         cliSession.cfg().setFetchSize(10);
-        when(client.basicQuery("test query", 10, false)).thenReturn(fakeResponse("", true, "foo"));
+        when(client.basicQuery("test query", 10, false, false)).thenReturn(fakeResponse("", true, "foo"));
         ServerQueryCliCommand cliCommand = new ServerQueryCliCommand();
         assertTrue(cliCommand.handle(testTerminal, cliSession, "test query"));
         assertEquals("""
@@ -53,7 +53,7 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
             ---------------
             foo           \s
             <flush/>""", testTerminal.toString());
-        verify(client, times(1)).basicQuery(eq("test query"), eq(10), eq(false));
+        verify(client, times(1)).basicQuery(eq("test query"), eq(10), eq(false), eq(false));
         verifyNoMoreInteractions(client);
     }
 
@@ -62,7 +62,7 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
         HttpClient client = mock(HttpClient.class);
         CliSession cliSession = new CliSession(client);
         cliSession.cfg().setFetchSize(10);
-        when(client.basicQuery("test query", 10, false)).thenReturn(fakeResponse("my_cursor1", true, "first"));
+        when(client.basicQuery("test query", 10, false, false)).thenReturn(fakeResponse("my_cursor1", true, "first"));
         when(client.nextPage("my_cursor1")).thenReturn(fakeResponse("my_cursor2", false, "second"));
         when(client.nextPage("my_cursor2")).thenReturn(fakeResponse("", false, "third"));
         ServerQueryCliCommand cliCommand = new ServerQueryCliCommand();
@@ -74,7 +74,7 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
             second        \s
             third         \s
             <flush/>""", testTerminal.toString());
-        verify(client, times(1)).basicQuery(eq("test query"), eq(10), eq(false));
+        verify(client, times(1)).basicQuery(eq("test query"), eq(10), eq(false), eq(false));
         verify(client, times(2)).nextPage(any());
         verifyNoMoreInteractions(client);
     }
@@ -86,7 +86,7 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
         cliSession.cfg().setFetchSize(15);
         // Set a separator
         cliSession.cfg().setFetchSeparator("-----");
-        when(client.basicQuery("test query", 15, false)).thenReturn(fakeResponse("my_cursor1", true, "first"));
+        when(client.basicQuery("test query", 15, false, false)).thenReturn(fakeResponse("my_cursor1", true, "first"));
         when(client.nextPage("my_cursor1")).thenReturn(fakeResponse("", false, "second"));
         ServerQueryCliCommand cliCommand = new ServerQueryCliCommand();
         assertTrue(cliCommand.handle(testTerminal, cliSession, "test query"));
@@ -97,7 +97,7 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
             -----
             second        \s
             <flush/>""", testTerminal.toString());
-        verify(client, times(1)).basicQuery(eq("test query"), eq(15), eq(false));
+        verify(client, times(1)).basicQuery(eq("test query"), eq(15), eq(false), eq(false));
         verify(client, times(1)).nextPage(any());
         verifyNoMoreInteractions(client);
     }
@@ -107,7 +107,7 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
         HttpClient client = mock(HttpClient.class);
         CliSession cliSession = new CliSession(client);
         cliSession.cfg().setFetchSize(15);
-        when(client.basicQuery("test query", 15, false)).thenReturn(fakeResponse("my_cursor1", true, "first"));
+        when(client.basicQuery("test query", 15, false, false)).thenReturn(fakeResponse("my_cursor1", true, "first"));
         when(client.nextPage("my_cursor1")).thenThrow(new SQLException("test exception"));
         when(client.queryClose("my_cursor1", Mode.CLI)).thenReturn(true);
         ServerQueryCliCommand cliCommand = new ServerQueryCliCommand();
@@ -118,7 +118,7 @@ public class ServerQueryCliCommandTests extends SqlCliTestCase {
             first         \s
             <b>Bad request [</b><i>test exception</i><b>]</b>
             """, testTerminal.toString());
-        verify(client, times(1)).basicQuery(eq("test query"), eq(15), eq(false));
+        verify(client, times(1)).basicQuery(eq("test query"), eq(15), eq(false), eq(false));
         verify(client, times(1)).nextPage(any());
         verify(client, times(1)).queryClose(eq("my_cursor1"), eq(Mode.CLI));
         verifyNoMoreInteractions(client);

x-pack/plugin/sql/sql-client/src/main/java/org/elasticsearch/xpack/sql/client/HttpClient.java

@@ -85,6 +85,11 @@ public class HttpClient {
     }
 
     public SqlQueryResponse basicQuery(String query, int fetchSize, boolean fieldMultiValueLeniency) throws SQLException {
+        return basicQuery(query, fetchSize, fieldMultiValueLeniency, cfg.allowPartialSearchResults());
+    }
+
+    public SqlQueryResponse basicQuery(String query, int fetchSize, boolean fieldMultiValueLeniency, boolean allowPartialSearchResults)
+        throws SQLException {
         // TODO allow customizing the time zone - this is what session set/reset/get should be about
         // method called only from CLI
         SqlQueryRequest sqlRequest = new SqlQueryRequest(
@@ -101,7 +106,7 @@ public class HttpClient {
             fieldMultiValueLeniency,
             false,
             cfg.binaryCommunication(),
-            cfg.allowPartialSearchResults()
+            allowPartialSearchResults
         );
         return query(sqlRequest).response();
     }

x-pack/plugin/sql/sql-client/src/test/java/org/elasticsearch/xpack/sql/client/HttpClientRequestTests.java

@@ -106,7 +106,7 @@ public class HttpClientRequestTests extends ESTestCase {
 
         prepareMockResponse();
         try {
-            httpClient.basicQuery(query, fetchSize, randomBoolean());
+            httpClient.basicQuery(query, fetchSize, randomBoolean(), randomBoolean());
         } catch (SQLException e) {
             logger.info("Ignored SQLException", e);
         }