SQL: Fix issue with date columns returned always in UTC (#40163)

When selecting columns of ES type `date` (SQL's DATETIME) the
`FieldHitExtractor` was not using the timezone of the client session
but always resorted to UTC. The same behaviour (UTC only) was
encountered also for grouping keys (`CompositeKeyExtractor`) and
for First/Last functions on dates (`TopHitsAggExtractor`).

Fixes: #40152

Author: matriv

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/jdbc/JdbcTestUtils.java

@@ -135,7 +135,7 @@ public static void logLikeCLI(ResultSet rs, Logger logger) throws SQLException {
         logger.info("\n" + formatter.formatWithHeader(cols, data));
     }
 
-    public static String of(long millis) {
-        return StringUtils.toString(ZonedDateTime.ofInstant(Instant.ofEpochMilli(millis), UTC));
+    public static String of(long millis, String zoneId) {
+        return StringUtils.toString(ZonedDateTime.ofInstant(Instant.ofEpochMilli(millis), ZoneId.of(zoneId)));
     }
-}
+}

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/jdbc/ResultSetTestCase.java

@@ -52,6 +52,7 @@
 import org.elasticsearch.xpack.sql.util.StringUtils;
 
 import java.io.IOException;
+import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -64,19 +65,21 @@ public class Querier {
     private final Logger log = LogManager.getLogger(getClass());
 
     private final TimeValue keepAlive, timeout;
+    private final ZoneId zoneId;
     private final int size;
     private final Client client;
     @Nullable
     private final QueryBuilder filter;
 
     public Querier(Client client, Configuration cfg) {
-        this(client, cfg.requestTimeout(), cfg.pageTimeout(), cfg.filter(), cfg.pageSize());
+        this(client, cfg.requestTimeout(), cfg.pageTimeout(), cfg.zoneId(), cfg.filter(), cfg.pageSize());
     }
 
-    public Querier(Client client, TimeValue keepAlive, TimeValue timeout, QueryBuilder filter, int size) {
+    public Querier(Client client, TimeValue keepAlive, TimeValue timeout, ZoneId zoneId, QueryBuilder filter, int size) {
         this.client = client;
         this.keepAlive = keepAlive;
         this.timeout = timeout;
+        this.zoneId = zoneId;
         this.filter = filter;
         this.size = size;
     }
@@ -98,13 +101,13 @@ public void query(Schema schema, QueryContainer query, String index, ActionListe
         ActionListener<SearchResponse> l;
         if (query.isAggsOnly()) {
             if (query.aggs().useImplicitGroupBy()) {
-                l = new ImplicitGroupActionListener(listener, client, timeout, schema, query, search);
+                l = new ImplicitGroupActionListener(listener, client, timeout, zoneId, schema, query, search);
             } else {
-                l = new CompositeActionListener(listener, client, timeout, schema, query, search);
+                l = new CompositeActionListener(listener, client, timeout, zoneId, schema, query, search);
             }
         } else {
             search.scroll(keepAlive);
-            l = new ScrollActionListener(listener, client, timeout, schema, query);
+            l = new ScrollActionListener(listener, client, timeout, zoneId, schema, query);
         }
 
         client.search(search, l);
@@ -149,9 +152,9 @@ public Aggregations getAggregations() {
             }
         });
 
-        ImplicitGroupActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive, Schema schema,
-                QueryContainer query, SearchRequest request) {
-            super(listener, client, keepAlive, schema, query, request);
+        ImplicitGroupActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive, ZoneId zoneId,
+                                    Schema schema, QueryContainer query, SearchRequest request) {
+            super(listener, client, keepAlive, zoneId, schema, query, request);
         }
 
         @Override
@@ -197,12 +200,11 @@ private void handleBuckets(List<? extends Bucket> buckets, SearchResponse respon
      */
     static class CompositeActionListener extends BaseAggActionListener {
 
-        CompositeActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive,
+        CompositeActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive, ZoneId zoneId,
                 Schema schema, QueryContainer query, SearchRequest request) {
-            super(listener, client, keepAlive, schema, query, request);
+            super(listener, client, keepAlive, zoneId, schema, query, request);
         }
 
-
         @Override
         protected void handleResponse(SearchResponse response, ActionListener<SchemaRowSet> listener) {
             // there are some results
@@ -240,9 +242,9 @@ abstract static class BaseAggActionListener extends BaseActionListener {
         final QueryContainer query;
         final SearchRequest request;
 
-        BaseAggActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive, Schema schema,
-                QueryContainer query, SearchRequest request) {
-            super(listener, client, keepAlive, schema);
+        BaseAggActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive, ZoneId zoneId,
+                              Schema schema, QueryContainer query, SearchRequest request) {
+            super(listener, client, keepAlive, zoneId, schema);
 
             this.query = query;
             this.request = request;
@@ -263,7 +265,7 @@ protected List<BucketExtractor> initBucketExtractors(SearchResponse response) {
         private BucketExtractor createExtractor(FieldExtraction ref, BucketExtractor totalCount) {
             if (ref instanceof GroupByRef) {
                 GroupByRef r = (GroupByRef) ref;
-                return new CompositeKeyExtractor(r.key(), r.property(), r.zoneId());
+                return new CompositeKeyExtractor(r.key(), r.property(), zoneId, r.isDateTimeBased());
             }
 
             if (ref instanceof MetricAggRef) {
@@ -297,9 +299,9 @@ private BucketExtractor createExtractor(FieldExtraction ref, BucketExtractor tot
     static class ScrollActionListener extends BaseActionListener {
         private final QueryContainer query;
 
-        ScrollActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive,
+        ScrollActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive, ZoneId zoneId,
                 Schema schema, QueryContainer query) {
-            super(listener, client, keepAlive, schema);
+            super(listener, client, keepAlive, zoneId, schema);
             this.query = query;
         }
 
@@ -344,12 +346,12 @@ protected void handleResponse(SearchResponse response, ActionListener<SchemaRowS
         private HitExtractor createExtractor(FieldExtraction ref) {
             if (ref instanceof SearchHitFieldRef) {
                 SearchHitFieldRef f = (SearchHitFieldRef) ref;
-                return new FieldHitExtractor(f.name(), f.getDataType(), f.useDocValue(), f.hitName());
+                return new FieldHitExtractor(f.name(), f.getDataType(), zoneId, f.useDocValue(), f.hitName());
             }
 
             if (ref instanceof ScriptFieldRef) {
                 ScriptFieldRef f = (ScriptFieldRef) ref;
-                return new FieldHitExtractor(f.name(), null, true);
+                return new FieldHitExtractor(f.name(), null, zoneId, true);
             }
 
             if (ref instanceof ComputedRef) {
@@ -387,13 +389,15 @@ abstract static class BaseActionListener implements ActionListener<SearchRespons
 
         final Client client;
         final TimeValue keepAlive;
+        final ZoneId zoneId;
         final Schema schema;
 
-        BaseActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive, Schema schema) {
+        BaseActionListener(ActionListener<SchemaRowSet> listener, Client client, TimeValue keepAlive, ZoneId zoneId, Schema schema) {
             this.listener = listener;
 
             this.client = client;
             this.keepAlive = keepAlive;
+            this.zoneId = zoneId;
             this.schema = schema;
         }
 

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

@@ -20,44 +20,38 @@
 public class CompositeKeyExtractor implements BucketExtractor {
 
     /**
-     * Key or Komposite extractor.
+     * Key or Composite extractor.
      */
     static final String NAME = "k";
 
     private final String key;
     private final Property property;
     private final ZoneId zoneId;
+    private final boolean isDateTimeBased;
 
     /**
      * Constructs a new <code>CompositeKeyExtractor</code> instance.
-     * The time-zone parameter is used to indicate a date key.
      */
-    public CompositeKeyExtractor(String key, Property property, ZoneId zoneId) {
+    public CompositeKeyExtractor(String key, Property property, ZoneId zoneId, boolean isDateTimeBased) {
         this.key = key;
         this.property = property;
         this.zoneId = zoneId;
+        this.isDateTimeBased = isDateTimeBased;
     }
 
     CompositeKeyExtractor(StreamInput in) throws IOException {
         key = in.readString();
         property = in.readEnum(Property.class);
-        if (in.readBoolean()) {
-            zoneId = ZoneId.of(in.readString());
-        } else {
-            zoneId = null;
-        }
+        zoneId = ZoneId.of(in.readString());
+        isDateTimeBased = in.readBoolean();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(key);
         out.writeEnum(property);
-        if (zoneId == null) {
-            out.writeBoolean(false);
-        } else {
-            out.writeBoolean(true);
-            out.writeString(zoneId.getId());
-        }
+        out.writeString(zoneId.getId());
+        out.writeBoolean(isDateTimeBased);
     }
 
     String key() {
@@ -72,6 +66,10 @@ ZoneId zoneId() {
         return zoneId;
     }
 
+    public boolean isDateTimeBased() {
+        return isDateTimeBased;
+    }
+
     @Override
     public String getWriteableName() {
         return NAME;
@@ -91,7 +89,7 @@ public Object extract(Bucket bucket) {
 
         Object object = ((Map<?, ?>) m).get(key);
 
-        if (zoneId != null) {
+        if (isDateTimeBased) {
             if (object == null) {
                 return object;
             } else if (object instanceof Long) {
@@ -106,7 +104,7 @@ public Object extract(Bucket bucket) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(key, property, zoneId);
+        return Objects.hash(key, property, zoneId, isDateTimeBased);
     }
 
     @Override
@@ -122,7 +120,8 @@ public boolean equals(Object obj) {
         CompositeKeyExtractor other = (CompositeKeyExtractor) obj;
         return Objects.equals(key, other.key)
                 && Objects.equals(property, other.property)
-                && Objects.equals(zoneId, other.zoneId);
+                && Objects.equals(zoneId, other.zoneId)
+                && Objects.equals(isDateTimeBased, other.isDateTimeBased);
     }
 
     @Override

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

@@ -17,6 +17,7 @@
 import org.joda.time.DateTime;
 
 import java.io.IOException;
+import java.time.ZoneId;
 import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.List;
@@ -48,16 +49,18 @@ private static String[] sourcePath(String name, boolean useDocValue, String hitN
 
     private final String fieldName, hitName;
     private final DataType dataType;
+    private final ZoneId zoneId;
     private final boolean useDocValue;
     private final String[] path;
 
-    public FieldHitExtractor(String name, DataType dataType, boolean useDocValue) {
-        this(name, dataType, useDocValue, null);
+    public FieldHitExtractor(String name, DataType dataType, ZoneId zoneId, boolean useDocValue) {
+        this(name, dataType, zoneId, useDocValue, null);
     }
 
-    public FieldHitExtractor(String name, DataType dataType, boolean useDocValue, String hitName) {
+    public FieldHitExtractor(String name, DataType dataType, ZoneId zoneId, boolean useDocValue, String hitName) {
         this.fieldName = name;
         this.dataType = dataType;
+        this.zoneId = zoneId;
         this.useDocValue = useDocValue;
         this.hitName = hitName;
 
@@ -74,6 +77,7 @@ public FieldHitExtractor(String name, DataType dataType, boolean useDocValue, St
         fieldName = in.readString();
         String esType = in.readOptionalString();
         dataType = esType != null ? DataType.fromTypeName(esType) : null;
+        zoneId = ZoneId.of(in.readString());
         useDocValue = in.readBoolean();
         hitName = in.readOptionalString();
         path = sourcePath(fieldName, useDocValue, hitName);
@@ -88,6 +92,7 @@ public String getWriteableName() {
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(fieldName);
         out.writeOptionalString(dataType == null ? null : dataType.esType);
+        out.writeString(zoneId.getId());
         out.writeBoolean(useDocValue);
         out.writeOptionalString(hitName);
     }
@@ -130,7 +135,7 @@ private Object unwrapMultiValue(Object values) {
         }
         if (dataType == DataType.DATE) {
             if (values instanceof String) {
-                return DateUtils.of(Long.parseLong(values.toString()));
+                return DateUtils.of(Long.parseLong(values.toString()), zoneId);
             }
             // returned by nested types...
             if (values instanceof DateTime) {
@@ -214,9 +219,17 @@ public String fieldName() {
         return fieldName;
     }
 
+    public ZoneId zoneId() {
+        return zoneId;
+    }
+
+    DataType dataType() {
+        return dataType;
+    }
+
     @Override
     public String toString() {
-        return fieldName + "@" + hitName;
+        return fieldName + "@" + hitName + "@" + zoneId;
     }
 
     @Override

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

@@ -5,14 +5,14 @@
  */
 package org.elasticsearch.xpack.sql.planner;
 
-import java.util.List;
-import java.util.Map;
-
 import org.elasticsearch.xpack.sql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.sql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.sql.planner.Verifier.Failure;
 import org.elasticsearch.xpack.sql.tree.Node;
 
+import java.util.List;
+import java.util.Map;
+
 import static java.util.stream.Collectors.toMap;
 
 public class Planner {
@@ -64,4 +64,4 @@ public Map<Node<?>, String> verifyExecutingPlanFailures(PhysicalPlan plan) {
         List<Failure> failures = Verifier.verifyExecutingPlan(plan);
         return failures.stream().collect(toMap(Failure::source, Failure::message));
     }
-}
+}