SQL: Fix issues with format=txt when paging through result sets and in mixed node environments (#83833)

Resolves #83581
Resolves #83788

SQL REST requests using `format=txt` stand out from the other formats because the cursor needs to carry the formatter state from the initial request to subsequent scroll requests. The state is needed to be able to format the subsequent pages with column widths that match the widths in the first page.

Currently, this is solved by wrapping cursor objects together with the formatter state in `TextFormatterCursor`. Hence, a query might return a `ListCursor` and `TextFormat.PLAIN_TEXT` adds the necessary state when formatting the output. This approach is handy because a `TextFormatterCursor` is a `Cursor` and delegates calls to the wrapped cursor. Unfortunately, it also has some downsides that have been revealed when looking into #83581 and #83788:

- MediaType formatting is a concern of the REST layer of the plugin that does not have access to the serialization logic usually accessible through the `NamedWriteablesRegistry`. Because the formatting required to (de)serialize the cursors for the wrapping, this meant that the `Cursor` deserialization could only use a subset of `NamedWriteable`s specific to Cursors (namely the ones returned by `Cursors.getNamedWriteables`). This subset of writeables is not enough to deserialize `CompositeAggCursor` which lead to the surprising design that `CompositeAggCursor` has a `nextQuery` member consisting of the serialized `SearchSourceBuilder` whose deserialization is suspended until the call to `Cursor.nextPage`.
- `TextFormat.PLAIN_TEXT` had to deserialize the complete cursor to do it's work, not just the information relevant to the formatter. Because of the `min_compatible_version` redirects that occur during rolling upgrades, it can happen though that an upgraded node needs to deserialize a cursor from a redirected SQL query written by an older node. Because we do not offer any form of bwc for cursors, this causes an error and leads to #83581.

In this PR, I propose to change the design such that the REST layer can add state to the cursor strings without having to read or write instances of `Cursor`. This allows to address the issues described above and also made fixing #83788 very straightforward.

Attaching state to the cursor is achieved by extending the format of encoded cursors. Previously, an base64 encoded SQL cursor had the following wire format `<version><zoneId><cursor>`. Now, there are two different cursor types:

* `<version><zoneId><CursorType.NO_STATE><cursor>` for most cases
* `<version><ZoneOffset.UTC><CursorType.WITH_STATE><formatterState><base64encodedCursor>` for txt format cursors. Note, `<ZoneOffset.UTC>` could be any zone id because consumers will use the one encoded in `<base64encodedCursor>`.

This serialization format ensures that every version of ES can read and check the cursor version and produce an according error in case of a mismatch.

Author: Lukas Wegmann

docs/changelog/83833.yaml

@@ -0,0 +1,7 @@
+pr: 83833
+summary: Fix issues with format=txt when paging through result sets and in mixed node environments
+area: SQL
+type: bug
+issues:
+ - 83581
+ - 83788

x-pack/plugin/sql/qa/mixed-node/src/test/java/org/elasticsearch/xpack/sql/qa/mixed_node/SqlCompatIT.java

@@ -197,4 +197,33 @@ private Map<String, Object> performRequestAndReadBodyAsJson(RestClient client, R
         }
     }
 
+    public void testCreateCursorWithFormatTxtOnNewNode() throws IOException {
+        testCreateCursorWithFormatTxt(newNodesClient);
+    }
+
+    public void testCreateCursorWithFormatTxtOnOldNode() throws IOException {
+        testCreateCursorWithFormatTxt(oldNodesClient);
+    }
+
+    /**
+     * Tests covering https://github.com/elastic/elasticsearch/issues/83581
+     */
+    public void testCreateCursorWithFormatTxt(RestClient client) throws IOException {
+        index("{\"foo\":1}", "{\"foo\":2}");
+
+        Request query = new Request("POST", "_sql");
+        XContentBuilder json = XContentFactory.jsonBuilder()
+            .startObject()
+            .field("query", randomFrom("SELECT foo FROM test", "SELECT foo FROM test GROUP BY foo"))
+            .field("fetch_size", 1)
+            .endObject();
+
+        query.setJsonEntity(Strings.toString(json));
+        query.addParameter("format", "txt");
+
+        Response response = client.performRequest(query);
+        assertOK(response);
+        assertFalse(Strings.isNullOrEmpty(response.getHeader("Cursor")));
+    }
+
 }

x-pack/plugin/sql/qa/server/src/main/java/org/elasticsearch/xpack/sql/qa/rest/RestSqlTestCase.java

@@ -77,6 +77,7 @@
 import static org.elasticsearch.xpack.sql.proto.CoreProtocol.WAIT_FOR_COMPLETION_TIMEOUT_NAME;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.lessThan;
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 /**
  * Integration test for the rest sql action. The one that speaks json directly to a
@@ -1175,11 +1176,14 @@ private void executeQueryWithNextPage(String format, String expectedHeader, Stri
         }
         index(docs);
 
-        String request = query("SELECT text, number, number + 5 AS sum FROM " + indexPattern("test") + " ORDER BY number").fetchSize(2)
+        String optionalGroupBy = randomFrom("", " GROUP BY text, number, sum");
+
+        String request = query("SELECT text, number, number + 5 AS sum FROM " + indexPattern("test") + optionalGroupBy + " ORDER BY number")
+            .fetchSize(2)
             .toString();
 
         String cursor = null;
-        for (int i = 0; i <= 20; i += 2) {
+        for (int i = 0; i <= size; i += 2) {
             Tuple<String, String> response;
             if (i == 0) {
                 response = runSqlAsText(StringUtils.EMPTY, new StringEntity(request, ContentType.APPLICATION_JSON), format);
@@ -1194,7 +1198,7 @@ private void executeQueryWithNextPage(String format, String expectedHeader, Stri
             StringBuilder expected = new StringBuilder();
             if (i == 0) {
                 expected.append(expectedHeader);
-                if (format == "text/plain") {
+                if (format.equals("text/plain")) {
                     expected.append("---------------+---------------+---------------\n");
                 }
             }
@@ -1438,6 +1442,127 @@ public void testCompressCursor() throws IOException {
         assertThat(resp.get("cursor").toString().length(), lessThan(5000));
     }
 
+    public void testFetchAllPagesSearchHitCursorTxt() throws IOException {
+        testFetchAllPagesSearchHitCursor("text/plain");
+    }
+
+    public void testFetchAllPagesSearchHitCursorCsv() throws IOException {
+        testFetchAllPagesSearchHitCursor("text/csv");
+    }
+
+    public void testFetchAllPagesSearchHitCursorTsv() throws IOException {
+        testFetchAllPagesSearchHitCursor("text/tab-separated-values");
+    }
+
+    public void testFetchAllPagesSearchHitCursor(String format) throws IOException {
+        int size = randomIntBetween(4, 20);
+        int pageSize = randomIntBetween(1, size + 1);
+
+        List<String> texts = IntStream.range(0, size).mapToObj(i -> String.format(Locale.ROOT, "text%02d", i)).toList();
+        index(texts.stream().map(t -> "{\"field\": \"" + t + "\"}").toArray(String[]::new));
+
+        testFetchAllPages(format, "SELECT field FROM " + indexPattern("test") + " ORDER BY field", texts, pageSize, true);
+    }
+
+    /**
+     * Special case of `testFetchAllPagesSearchHitCursorTxt` covering https://github.com/elastic/elasticsearch/issues/83788.
+     */
+    public void testFetchAllPagesSearchHitCursorWithNonFullLastPageTxt() throws IOException {
+        int size = randomIntBetween(4, 20);
+        int pageSize = size / 2 + 1;
+
+        List<String> texts = IntStream.range(0, size).mapToObj(i -> String.format(Locale.ROOT, "text%02d", i)).toList();
+        index(texts.stream().map(t -> "{\"field\": \"" + t + "\"}").toArray(String[]::new));
+
+        testFetchAllPages("text/plain", "SELECT field FROM " + indexPattern("test") + " ORDER BY field", texts, pageSize, false);
+    }
+
+    public void testFetchAllPagesCompositeAggCursorTxt() throws IOException {
+        testFetchAllPagesCompositeAggCursor("text/plain");
+    }
+
+    public void testFetchAllPagesCompositeAggCursorCsv() throws IOException {
+        testFetchAllPagesCompositeAggCursor("text/csv");
+    }
+
+    public void testFetchAllPagesCompositeAggCursorTsv() throws IOException {
+        testFetchAllPagesCompositeAggCursor("text/tab-separated-values");
+    }
+
+    public void testFetchAllPagesCompositeAggCursor(String format) throws IOException {
+        int size = randomIntBetween(4, 20);
+        int pageSize = randomIntBetween(1, size + 1);
+
+        List<String> texts = IntStream.range(0, size).mapToObj(i -> String.format(Locale.ROOT, "text%02d", i)).toList();
+        index(texts.stream().map(t -> "{\"field\": \"" + t + "\"}").toArray(String[]::new));
+
+        testFetchAllPages(format, "SELECT field FROM " + indexPattern("test") + " GROUP BY field ORDER BY field", texts, pageSize, true);
+    }
+
+    public void testFetchAllPagesListCursorTxt() throws IOException {
+        testFetchAllPagesListCursor("text/plain");
+    }
+
+    public void testFetchAllPagesListCursorCsv() throws IOException {
+        testFetchAllPagesListCursor("text/csv");
+    }
+
+    public void testFetchAllPagesListCursorTsv() throws IOException {
+        testFetchAllPagesListCursor("text/tab-separated-values");
+    }
+
+    public void testFetchAllPagesListCursor(String format) throws IOException {
+        int size = randomIntBetween(4, 20);
+        int pageSize = randomIntBetween(1, size + 1);
+
+        List<String> fields = IntStream.range(0, size).mapToObj(i -> String.format(Locale.ROOT, "text%02d", i)).toList();
+        index(fields.stream().map(f -> "{\"" + f + "\": 1}").toArray(String[]::new));
+
+        testFetchAllPages(format, "SHOW COLUMNS FROM " + indexPattern("test"), fields, pageSize, false);
+    }
+
+    /**
+     * Generic paging test that ensures that
+     * 1. All expected documents are fetched when scrolling through all pages
+     * 2. There are at most `expectedValues.size() / pageSize + 1` pages (the last one might or might not be empty)
+     * 3. Optionally: That the last page is not empty.
+     */
+    private void testFetchAllPages(String format, String query, List<String> expectedValues, int pageSize, boolean allowEmptyLastPage)
+        throws IOException {
+        int remainingPages = expectedValues.size() / pageSize + 1;
+
+        Tuple<String, String> response = runSqlAsText(
+            StringUtils.EMPTY,
+            new StringEntity(query(query).fetchSize(pageSize).toString(), ContentType.APPLICATION_JSON),
+            format
+        );
+        StringBuilder allResults = new StringBuilder(response.v1());
+
+        while (response.v2() != null && remainingPages > 0) {
+            response = runSqlAsText(
+                StringUtils.EMPTY,
+                new StringEntity(cursor(response.v2()).toString(), ContentType.APPLICATION_JSON),
+                format
+            );
+            allResults.append("\n").append(response.v1());
+
+            assertThat(response.v1().split("\n").length, lessThanOrEqualTo(pageSize));
+
+            remainingPages--;
+        }
+
+        if (allowEmptyLastPage == false) {
+            assertFalse(Strings.isNullOrEmpty(response.v1()));
+        }
+
+        assertNull(response.v2());
+
+        String results = allResults.toString();
+        for (String v : expectedValues) {
+            assertThat(results, containsString(v));
+        }
+    }
+
     static Map<String, Object> runSql(RequestObjectBuilder builder, String mode) throws IOException {
         return toMap(runSql(builder.mode(mode)), mode);
     }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/PlanExecutor.java

@@ -120,11 +120,11 @@ public void nextPage(SqlConfiguration cfg, Cursor cursor, ActionListener<Page> l
      * `nextPage` for internal callers (not from the APIs) without metrics reporting.
      */
     public void nextPageInternal(SqlConfiguration cfg, Cursor cursor, ActionListener<Page> listener) {
-        cursor.nextPage(cfg, client, writableRegistry, listener);
+        cursor.nextPage(cfg, client, listener);
     }
 
     public void cleanCursor(Cursor cursor, ActionListener<Boolean> listener) {
-        cursor.clear(client, writableRegistry, listener);
+        cursor.clear(client, listener);
     }
 
     public Client client() {

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/CompositeAggCursor.java

@@ -12,7 +12,6 @@
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.client.internal.Client;
-import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.search.aggregations.Aggregation;
@@ -38,10 +37,8 @@
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
 
-import static org.elasticsearch.xpack.sql.execution.search.Querier.deserializeQuery;
 import static org.elasticsearch.xpack.sql.execution.search.Querier.logSearchResponse;
 import static org.elasticsearch.xpack.sql.execution.search.Querier.prepareRequest;
-import static org.elasticsearch.xpack.sql.execution.search.Querier.serializeQuery;
 
 /**
  * Cursor for composite aggregation (GROUP BY).
@@ -54,15 +51,22 @@ public class CompositeAggCursor implements Cursor {
     public static final String NAME = "c";
 
     private final String[] indices;
-    private final byte[] nextQuery;
+    private final SearchSourceBuilder nextQuery;
     private final List<BucketExtractor> extractors;
     private final BitSet mask;
     private final int limit;
     private final boolean includeFrozen;
 
-    CompositeAggCursor(byte[] next, List<BucketExtractor> exts, BitSet mask, int remainingLimit, boolean includeFrozen, String... indices) {
+    CompositeAggCursor(
+        SearchSourceBuilder nextQuery,
+        List<BucketExtractor> exts,
+        BitSet mask,
+        int remainingLimit,
+        boolean includeFrozen,
+        String... indices
+    ) {
         this.indices = indices;
-        this.nextQuery = next;
+        this.nextQuery = nextQuery;
         this.extractors = exts;
         this.mask = mask;
         this.limit = remainingLimit;
@@ -71,7 +75,7 @@ public class CompositeAggCursor implements Cursor {
 
     public CompositeAggCursor(StreamInput in) throws IOException {
         indices = in.readStringArray();
-        nextQuery = in.readByteArray();
+        nextQuery = new SearchSourceBuilder(in);
         limit = in.readVInt();
 
         extractors = in.readNamedWriteableList(BucketExtractor.class);
@@ -82,7 +86,7 @@ public CompositeAggCursor(StreamInput in) throws IOException {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeStringArray(indices);
-        out.writeByteArray(nextQuery);
+        nextQuery.writeTo(out);
         out.writeVInt(limit);
 
         out.writeNamedWriteableList(extractors);
@@ -99,7 +103,7 @@ String[] indices() {
         return indices;
     }
 
-    byte[] next() {
+    SearchSourceBuilder next() {
         return nextQuery;
     }
 
@@ -120,21 +124,12 @@ boolean includeFrozen() {
     }
 
     @Override
-    public void nextPage(SqlConfiguration cfg, Client client, NamedWriteableRegistry registry, ActionListener<Page> listener) {
-        SearchSourceBuilder q;
-        try {
-            q = deserializeQuery(registry, nextQuery);
-        } catch (Exception ex) {
-            listener.onFailure(ex);
-            return;
-        }
-
-        SearchSourceBuilder query = q;
+    public void nextPage(SqlConfiguration cfg, Client client, ActionListener<Page> listener) {
         if (log.isTraceEnabled()) {
-            log.trace("About to execute composite query {} on {}", StringUtils.toString(query), indices);
+            log.trace("About to execute composite query {} on {}", StringUtils.toString(nextQuery), indices);
         }
 
-        SearchRequest request = prepareRequest(query, cfg.requestTimeout(), includeFrozen, indices);
+        SearchRequest request = prepareRequest(nextQuery, cfg.requestTimeout(), includeFrozen, indices);
 
         client.search(request, new ActionListener.Delegating<>(listener) {
             @Override
@@ -156,15 +151,15 @@ protected Supplier<CompositeAggRowSet> makeRowSet(SearchResponse response) {
         return () -> new CompositeAggRowSet(extractors, mask, response, limit);
     }
 
-    protected BiFunction<byte[], CompositeAggRowSet, CompositeAggCursor> makeCursor() {
+    protected BiFunction<SearchSourceBuilder, CompositeAggRowSet, CompositeAggCursor> makeCursor() {
         return (q, r) -> new CompositeAggCursor(q, r.extractors(), r.mask(), r.remainingData(), includeFrozen, indices);
     }
 
     static void handle(
         SearchResponse response,
         SearchSourceBuilder source,
         Supplier<CompositeAggRowSet> makeRowSet,
-        BiFunction<byte[], CompositeAggRowSet, CompositeAggCursor> makeCursor,
+        BiFunction<SearchSourceBuilder, CompositeAggRowSet, CompositeAggCursor> makeCursor,
         Runnable retry,
         ActionListener<Page> listener,
         Schema schema
@@ -186,13 +181,11 @@ static void handle(
                 CompositeAggRowSet rowSet = makeRowSet.get();
                 Map<String, Object> afterKey = rowSet.afterKey();
 
-                byte[] queryAsBytes = null;
                 if (afterKey != null) {
                     updateSourceAfterKey(afterKey, source);
-                    queryAsBytes = serializeQuery(source);
                 }
 
-                Cursor next = rowSet.remainingData() == 0 ? Cursor.EMPTY : makeCursor.apply(queryAsBytes, rowSet);
+                Cursor next = rowSet.remainingData() == 0 ? Cursor.EMPTY : makeCursor.apply(source, rowSet);
                 listener.onResponse(new Page(rowSet, next));
             } catch (Exception ex) {
                 listener.onFailure(ex);
@@ -247,13 +240,13 @@ private static void updateSourceAfterKey(Map<String, Object> afterKey, SearchSou
     }
 
     @Override
-    public void clear(Client client, NamedWriteableRegistry registry, ActionListener<Boolean> listener) {
+    public void clear(Client client, ActionListener<Boolean> listener) {
         listener.onResponse(true);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(Arrays.hashCode(indices), Arrays.hashCode(nextQuery), extractors, limit, mask, includeFrozen);
+        return Objects.hash(Arrays.hashCode(indices), nextQuery, extractors, limit, mask, includeFrozen);
     }
 
     @Override
@@ -263,7 +256,7 @@ public boolean equals(Object obj) {
         }
         CompositeAggCursor other = (CompositeAggCursor) obj;
         return Arrays.equals(indices, other.indices)
-            && Arrays.equals(nextQuery, other.nextQuery)
+            && Objects.equals(nextQuery, other.nextQuery)
             && Objects.equals(extractors, other.extractors)
             && Objects.equals(limit, other.limit)
             && Objects.equals(includeFrozen, other.includeFrozen);

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/PivotCursor.java

@@ -10,6 +10,7 @@
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.xpack.ql.execution.search.extractor.BucketExtractor;
 import org.elasticsearch.xpack.ql.type.Schema;
 
@@ -29,7 +30,7 @@ public class PivotCursor extends CompositeAggCursor {
 
     PivotCursor(
         Map<String, Object> previousKey,
-        byte[] next,
+        SearchSourceBuilder next,
         List<BucketExtractor> exts,
         BitSet mask,
         int remainingLimit,
@@ -67,7 +68,7 @@ protected Supplier<CompositeAggRowSet> makeRowSet(SearchResponse response) {
     }
 
     @Override
-    protected BiFunction<byte[], CompositeAggRowSet, CompositeAggCursor> makeCursor() {
+    protected BiFunction<SearchSourceBuilder, CompositeAggRowSet, CompositeAggCursor> makeCursor() {
         return (q, r) -> {
             Map<String, Object> lastAfterKey = r instanceof PivotRowSet ? ((PivotRowSet) r).lastAfterKey() : null;
             return new PivotCursor(lastAfterKey, q, r.extractors(), r.mask(), r.remainingData(), includeFrozen(), indices());