SQL: Improve serialization of SQL processors (#45678)

Encapsulate the serialization/deserialization of SQL client classes.
Make configuration specific parameters (such as ZoneId) generic just
like the version and remove the need for consumer classes to manage them
individually.
This is not only consistent but also provides significant savings in the
cursor.

Fix #40216

(cherry picked from commit 5c84479)

Author: costin

x-pack/plugin/sql/sql-action/src/main/java/org/elasticsearch/xpack/sql/action/AbstractSqlQueryRequest.java

@@ -194,7 +194,7 @@ public AbstractSqlQueryRequest(StreamInput in) throws IOException {
         super(in);
         query = in.readString();
         params = in.readList(AbstractSqlQueryRequest::readSqlTypedParamValue);
-        zoneId = ZoneId.of(in.readString());
+        zoneId = in.readZoneId();
         fetchSize = in.readVInt();
         requestTimeout = in.readTimeValue();
         pageTimeout = in.readTimeValue();
@@ -218,7 +218,7 @@ public void writeTo(StreamOutput out) throws IOException {
         for (SqlTypedParamValue param: params) {
             writeSqlTypedParamValue(out, param);
         }
-        out.writeString(zoneId.getId());
+        out.writeZoneId(zoneId);
         out.writeVInt(fetchSize);
         out.writeTimeValue(requestTimeout);
         out.writeTimeValue(pageTimeout);

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/common/io/SqlStreamInput.java

@@ -0,0 +1,54 @@
+/*
+ * 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.sql.common.io;
+
+import org.elasticsearch.Version;
+import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
+
+import java.io.IOException;
+import java.time.ZoneId;
+import java.util.Base64;
+
+/**
+ * SQL-specific stream extension for {@link StreamInput} used for deserializing
+ * SQL components, especially on the client-side.
+ */
+public class SqlStreamInput extends NamedWriteableAwareStreamInput {
+
+    private final ZoneId zoneId;
+
+    public SqlStreamInput(String base64encoded, NamedWriteableRegistry namedWriteableRegistry, Version version) throws IOException {
+        this(Base64.getDecoder().decode(base64encoded), namedWriteableRegistry, version);
+    }
+
+    public SqlStreamInput(byte[] input, NamedWriteableRegistry namedWriteableRegistry, Version version) throws IOException {
+        super(StreamInput.wrap(input), namedWriteableRegistry);
+
+        // version check first
+        Version ver = Version.readVersion(delegate);
+        if (version.compareTo(ver) != 0) {
+            throw new SqlIllegalArgumentException("Unsupported cursor version [{}], expected [{}]", ver, version);
+        }
+        delegate.setVersion(version);
+        // configuration settings
+        zoneId = delegate.readZoneId();
+    }
+
+    public ZoneId zoneId() {
+        return zoneId;
+    }
+
+    public static SqlStreamInput asSqlStream(StreamInput in) {
+        if (in instanceof SqlStreamInput) {
+            return (SqlStreamInput) in;
+        }
+        throw new SqlIllegalArgumentException("Expected SQL cursor stream, received [{}]", in.getClass());
+    }
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/common/io/SqlStreamOutput.java

@@ -0,0 +1,41 @@
+/*
+ * 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.sql.common.io;
+
+import org.elasticsearch.Version;
+import org.elasticsearch.common.io.stream.OutputStreamStreamOutput;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.time.ZoneId;
+import java.util.Base64;
+
+public class SqlStreamOutput extends OutputStreamStreamOutput {
+
+    private final ByteArrayOutputStream bytes;
+
+    public SqlStreamOutput(Version version, ZoneId zoneId) throws IOException {
+        this(new ByteArrayOutputStream(), version, zoneId);
+    }
+
+    private SqlStreamOutput(ByteArrayOutputStream bytes, Version version, ZoneId zoneId) throws IOException {
+        super(Base64.getEncoder().wrap(new OutputStreamStreamOutput(bytes)));
+        this.bytes = bytes;
+
+        Version.writeVersion(version, this);
+        writeZoneId(zoneId);
+    }
+
+    /**
+     * Should be called _after_ closing the stream - there are no guarantees otherwise.
+     */
+    public String streamAsString() {
+        // Base64 uses this encoding instead of UTF-8
+        return new String(bytes.toByteArray(), StandardCharsets.ISO_8859_1);
+    }
+}

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

@@ -9,6 +9,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation.Bucket;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
+import org.elasticsearch.xpack.sql.common.io.SqlStreamInput;
 import org.elasticsearch.xpack.sql.querydsl.container.GroupByRef.Property;
 import org.elasticsearch.xpack.sql.util.DateUtils;
 
@@ -42,15 +43,15 @@ public CompositeKeyExtractor(String key, Property property, ZoneId zoneId, boole
     CompositeKeyExtractor(StreamInput in) throws IOException {
         key = in.readString();
         property = in.readEnum(Property.class);
-        zoneId = ZoneId.of(in.readString());
         isDateTimeBased = in.readBoolean();
+
+        zoneId = SqlStreamInput.asSqlStream(in).zoneId();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(key);
         out.writeEnum(property);
-        out.writeString(zoneId.getId());
         out.writeBoolean(isDateTimeBased);
     }
 

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.index.mapper.IgnoredFieldMapper;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
+import org.elasticsearch.xpack.sql.common.io.SqlStreamInput;
 import org.elasticsearch.xpack.sql.expression.function.scalar.geo.GeoShape;
 import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.util.DateUtils;
@@ -95,11 +96,12 @@ public FieldHitExtractor(String name, String fullFieldName, DataType dataType, Z
         }
         String esType = in.readOptionalString();
         dataType = esType != null ? DataType.fromTypeName(esType) : null;
-        zoneId = ZoneId.of(in.readString());
         useDocValue = in.readBoolean();
         hitName = in.readOptionalString();
         arrayLeniency = in.readBoolean();
         path = sourcePath(fieldName, useDocValue, hitName);
+
+        zoneId = SqlStreamInput.asSqlStream(in).zoneId();
     }
 
     @Override
@@ -114,7 +116,6 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeOptionalString(fullFieldName);
         }
         out.writeOptionalString(dataType == null ? null : dataType.typeName);
-        out.writeString(zoneId.getId());
         out.writeBoolean(useDocValue);
         out.writeOptionalString(hitName);
         out.writeBoolean(arrayLeniency);

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

@@ -21,6 +21,7 @@
 import org.elasticsearch.search.aggregations.metrics.InternalTDigestPercentileRanks;
 import org.elasticsearch.search.aggregations.metrics.InternalTDigestPercentiles;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
+import org.elasticsearch.xpack.sql.common.io.SqlStreamInput;
 import org.elasticsearch.xpack.sql.querydsl.agg.Aggs;
 import org.elasticsearch.xpack.sql.util.DateUtils;
 
@@ -55,7 +56,8 @@ public MetricAggExtractor(String name, String property, String innerKey, ZoneId
         property = in.readString();
         innerKey = in.readOptionalString();
         isDateTimeBased = in.readBoolean();
-        zoneId = ZoneId.of(in.readString());
+
+        zoneId = SqlStreamInput.asSqlStream(in).zoneId();
     }
 
     @Override
@@ -64,7 +66,6 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeString(property);
         out.writeOptionalString(innerKey);
         out.writeBoolean(isDateTimeBased);
-        out.writeString(zoneId.getId());
     }
 
     String name() {

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation.Bucket;
 import org.elasticsearch.search.aggregations.metrics.InternalTopHits;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
+import org.elasticsearch.xpack.sql.common.io.SqlStreamInput;
 import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.util.DateUtils;
 
@@ -34,14 +35,13 @@ public TopHitsAggExtractor(String name, DataType fieldDataType, ZoneId zoneId) {
     TopHitsAggExtractor(StreamInput in) throws IOException {
         name = in.readString();
         fieldDataType = in.readEnum(DataType.class);
-        zoneId = ZoneId.of(in.readString());
+        zoneId = SqlStreamInput.asSqlStream(in).zoneId();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(name);
         out.writeEnum(fieldDataType);
-        out.writeString(zoneId.getId());
     }
 
     String name() {

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/datetime/BaseDateTimeProcessor.java

@@ -7,8 +7,8 @@
 package org.elasticsearch.xpack.sql.expression.function.scalar.datetime;
 
 import org.elasticsearch.common.io.stream.StreamInput;
-import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
+import org.elasticsearch.xpack.sql.common.io.SqlStreamInput;
 import org.elasticsearch.xpack.sql.expression.gen.processor.Processor;
 
 import java.io.IOException;
@@ -24,12 +24,7 @@ public abstract class BaseDateTimeProcessor implements Processor {
     }
 
     BaseDateTimeProcessor(StreamInput in) throws IOException {
-        zoneId = ZoneId.of(in.readString());
-    }
-
-    @Override
-    public void writeTo(StreamOutput out) throws IOException {
-        out.writeString(zoneId.getId());
+        zoneId = SqlStreamInput.asSqlStream(in).zoneId();
     }
 
     ZoneId zoneId() {

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/datetime/DateTimeProcessor.java

@@ -63,7 +63,6 @@ public DateTimeProcessor(StreamInput in) throws IOException {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        super.writeTo(out);
         out.writeEnum(extractor);
     }
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/datetime/NamedDateTimeProcessor.java

@@ -57,7 +57,6 @@ public NamedDateTimeProcessor(StreamInput in) throws IOException {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        super.writeTo(out);
         out.writeEnum(extractor);
     }
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/datetime/NonIsoDateTimeProcessor.java

@@ -72,7 +72,6 @@ public NonIsoDateTimeProcessor(StreamInput in) throws IOException {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        super.writeTo(out);
         out.writeEnum(extractor);
     }
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/datetime/QuarterProcessor.java

@@ -7,6 +7,7 @@
 package org.elasticsearch.xpack.sql.expression.function.scalar.datetime;
 
 import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
 
 import java.io.IOException;
 import java.time.ZoneId;
@@ -16,6 +17,10 @@
 import java.util.Objects;
 
 public class QuarterProcessor extends BaseDateTimeProcessor {
+
+    public static final String NAME = "q";
+    private static final DateTimeFormatter QUARTER_FORMAT = DateTimeFormatter.ofPattern("q", Locale.ROOT);
+
 
     public QuarterProcessor(ZoneId zoneId) {
         super(zoneId);
@@ -25,8 +30,8 @@ public QuarterProcessor(StreamInput in) throws IOException {
         super(in);
     }
 
-    public static final String NAME = "q";
-    private static final DateTimeFormatter QUARTER_FORMAT = DateTimeFormatter.ofPattern("q", Locale.ROOT);
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {}
 
     @Override
     public String getWriteableName() {

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

@@ -110,7 +110,7 @@ public RestResponse buildResponse(SqlQueryResponse response) throws Exception {
                     final String data = textFormat.format(request, response);
 
                     restResponse = new BytesRestResponse(RestStatus.OK, textFormat.contentType(request),
-                            data.getBytes(StandardCharsets.UTF_8));
+                        data.getBytes(StandardCharsets.UTF_8));
 
                     if (response.hasCursor()) {
                         restResponse.addHeader("Cursor", response.cursor());

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

@@ -5,8 +5,8 @@
  */
 package org.elasticsearch.xpack.sql.plugin;
 
-import org.elasticsearch.Version;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
 import org.elasticsearch.xpack.sql.action.BasicFormatter;
@@ -17,6 +17,7 @@
 import org.elasticsearch.xpack.sql.util.DateUtils;
 import org.elasticsearch.xpack.sql.util.StringUtils;
 
+import java.time.ZoneId;
 import java.time.ZonedDateTime;
 import java.util.List;
 import java.util.Locale;
@@ -43,10 +44,13 @@ enum TextFormat {
         String format(RestRequest request, SqlQueryResponse response) {
             BasicFormatter formatter = null;
             Cursor cursor = null;
+            ZoneId zoneId = null;
 
             // check if the cursor is already wrapped first
             if (response.hasCursor()) {
-                cursor = Cursors.decodeFromString(response.cursor());
+                Tuple<Cursor, ZoneId> tuple = Cursors.decodeFromStringWithZone(response.cursor());
+                cursor = tuple.v1();
+                zoneId = tuple.v2();
                 if (cursor instanceof TextFormatterCursor) {
                     formatter = ((TextFormatterCursor) cursor).getFormatter();
                 }
@@ -58,7 +62,7 @@ String format(RestRequest request, SqlQueryResponse response) {
                 formatter = new BasicFormatter(response.columns(), response.rows(), TEXT);
                 // if there's a cursor, wrap the formatter in it
                 if (cursor != null) {
-                    response.cursor(Cursors.encodeToString(Version.CURRENT, new TextFormatterCursor(cursor, formatter)));
+                    response.cursor(Cursors.encodeToString(new TextFormatterCursor(cursor, formatter), zoneId));
                 }
                 // format with header
                 return formatter.formatWithHeader(response.columns(), response.rows());

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

@@ -5,7 +5,6 @@
  */
 package org.elasticsearch.xpack.sql.plugin;
 
-import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
@@ -81,7 +80,7 @@ public static void operation(PlanExecutor planExecutor, SqlQueryRequest request,
                     wrap(p -> listener.onResponse(createResponseWithSchema(request, p)), listener::onFailure));
         } else {
             planExecutor.nextPage(cfg, Cursors.decodeFromString(request.cursor()),
-                    wrap(p -> listener.onResponse(createResponse(request.mode(), request.columnar(), null, p)),
+                    wrap(p -> listener.onResponse(createResponse(request, null, p)),
                             listener::onFailure));
         }
     }
@@ -102,10 +101,10 @@ static SqlQueryResponse createResponseWithSchema(SqlQueryRequest request, Page p
             }
         }
         columns = unmodifiableList(columns);
-        return createResponse(request.mode(), request.columnar(), columns, page);
+        return createResponse(request, columns, page);
     }
 
-    static SqlQueryResponse createResponse(Mode mode, boolean columnar, List<ColumnInfo> header, Page page) {
+    static SqlQueryResponse createResponse(SqlQueryRequest request, List<ColumnInfo> header, Page page) {
         List<List<Object>> rows = new ArrayList<>();
         page.rowSet().forEachRow(rowView -> {
             List<Object> row = new ArrayList<>(rowView.columnCount());
@@ -114,9 +113,9 @@ static SqlQueryResponse createResponse(Mode mode, boolean columnar, List<ColumnI
         });
 
         return new SqlQueryResponse(
-                Cursors.encodeToString(Version.CURRENT, page.next()),
-                mode,
-                columnar,
+                Cursors.encodeToString(page.next(), request.zoneId()),
+                request.mode(),
+                request.columnar(),
                 header,
                 rows);
     }