From 21920b17e7ec69098d052df8bd1dcbd55d829d95 Mon Sep 17 00:00:00 2001 From: Vikash Kumar Date: Wed, 9 Oct 2024 13:59:57 +0530 Subject: [PATCH] Use dedicated protocol to render data in local timezone --- .../io/trino/client/ClientCapabilities.java | 3 +- .../java/io/trino/client/ProtocolHeaders.java | 7 +++ .../src/main/java/io/trino/Session.java | 33 +++++++++-- .../java/io/trino/SessionRepresentation.java | 8 ++- .../protocol/JsonArrayResultsIterator.java | 58 +++++++++++++++++++ 5 files changed, 100 insertions(+), 9 deletions(-) diff --git a/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java b/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java index ba55f3872ea4..f9dce89abaee 100644 --- a/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java +++ b/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java @@ -25,5 +25,6 @@ public enum ClientCapabilities // When this capability is not set, the server returns datetime types with precision = 3 PARAMETRIC_DATETIME, // Whether clients support the session authorization set/reset feature - SESSION_AUTHORIZATION; + SESSION_AUTHORIZATION, + USE_SESSION_TIME_ZONE, } diff --git a/client/trino-client/src/main/java/io/trino/client/ProtocolHeaders.java b/client/trino-client/src/main/java/io/trino/client/ProtocolHeaders.java index 2bbcca4ca8b9..902ee429f89b 100644 --- a/client/trino-client/src/main/java/io/trino/client/ProtocolHeaders.java +++ b/client/trino-client/src/main/java/io/trino/client/ProtocolHeaders.java @@ -44,6 +44,7 @@ public final class ProtocolHeaders private final String requestResourceEstimate; private final String requestExtraCredential; private final String requestQueryDataEncoding; + private final String requestRenderTimestampInLocalZone; private final String responseSetCatalog; private final String responseSetSchema; private final String responseSetPath; @@ -91,6 +92,7 @@ private ProtocolHeaders(String name) requestResourceEstimate = prefix + "Resource-Estimate"; requestExtraCredential = prefix + "Extra-Credential"; requestQueryDataEncoding = prefix + "Query-Data-Encoding"; + requestRenderTimestampInLocalZone = prefix + "Timestamp-In-Local-Zone"; responseSetCatalog = prefix + "Set-Catalog"; responseSetSchema = prefix + "Set-Schema"; responseSetPath = prefix + "Set-Path"; @@ -205,6 +207,11 @@ public String requestQueryDataEncoding() return requestQueryDataEncoding; } + public String getRequestRenderTimestampInLocalZone() + { + return requestRenderTimestampInLocalZone; + } + public String responseSetCatalog() { return responseSetCatalog; diff --git a/core/trino-main/src/main/java/io/trino/Session.java b/core/trino-main/src/main/java/io/trino/Session.java index 6aeaab959048..a29c3fabf1b9 100644 --- a/core/trino-main/src/main/java/io/trino/Session.java +++ b/core/trino-main/src/main/java/io/trino/Session.java @@ -92,6 +92,7 @@ public final class Session private final ProtocolHeaders protocolHeaders; private final Optional exchangeEncryptionKey; private final Optional queryDataEncoding; + private final boolean renderTimestampInLocalTimeZone; public Session( QueryId queryId, @@ -120,7 +121,8 @@ public Session( Map preparedStatements, ProtocolHeaders protocolHeaders, Optional exchangeEncryptionKey, - Optional queryDataEncoding) + Optional queryDataEncoding, + boolean renderTimestampInLocalTimeZone) { this.queryId = requireNonNull(queryId, "queryId is null"); this.querySpan = requireNonNull(querySpan, "querySpan is null"); @@ -157,6 +159,7 @@ public Session( this.catalogProperties = catalogPropertiesBuilder.buildOrThrow(); checkArgument(catalog.isPresent() || schema.isEmpty(), "schema is set but catalog is not"); + this.renderTimestampInLocalTimeZone = renderTimestampInLocalTimeZone; } public QueryId getQueryId() @@ -265,6 +268,11 @@ public TransactionId getRequiredTransactionId() return transactionId.orElseThrow(NotInTransactionException::new); } + public boolean isRenderTimestampInLocalTimeZone() + { + return renderTimestampInLocalTimeZone; + } + public boolean isClientTransactionSupport() { return clientTransactionSupport; @@ -393,7 +401,8 @@ public Session beginTransactionId(TransactionId transactionId, TransactionManage preparedStatements, protocolHeaders, exchangeEncryptionKey, - queryDataEncoding); + queryDataEncoding, + renderTimestampInLocalTimeZone); } public Session withDefaultProperties(Map systemPropertyDefaults, Map> catalogPropertyDefaults, AccessControl accessControl) @@ -443,7 +452,8 @@ public Session withDefaultProperties(Map systemPropertyDefaults, preparedStatements, protocolHeaders, exchangeEncryptionKey, - queryDataEncoding); + queryDataEncoding, + renderTimestampInLocalTimeZone); } public Session withExchangeEncryption(Slice encryptionKey) @@ -476,7 +486,8 @@ public Session withExchangeEncryption(Slice encryptionKey) preparedStatements, protocolHeaders, Optional.of(encryptionKey), - queryDataEncoding); + queryDataEncoding, + renderTimestampInLocalTimeZone); } public ConnectorSession toConnectorSession() @@ -530,7 +541,8 @@ public SessionRepresentation toSessionRepresentation() identity.getCatalogRoles(), preparedStatements, protocolHeaders.getProtocolName(), - queryDataEncoding); + queryDataEncoding, + renderTimestampInLocalTimeZone); } @Override @@ -556,6 +568,7 @@ public String toString() .add("clientCapabilities", clientCapabilities) .add("resourceEstimates", resourceEstimates) .add("start", start) + .add("renderTimestampInLocalTimeZone", renderTimestampInLocalTimeZone) .omitNullValues() .toString(); } @@ -663,6 +676,7 @@ public static class SessionBuilder private Optional queryDataEncoding = Optional.empty(); private ResourceEstimates resourceEstimates; private Instant start = Instant.now(); + private boolean renderTimestampInLocalTimeZone; private final Map systemProperties = new HashMap<>(); private final Map> catalogSessionProperties = new HashMap<>(); private final SessionPropertyManager sessionPropertyManager; @@ -950,6 +964,12 @@ public SessionBuilder setQueryDataEncoding(Optional value) return this; } + public SessionBuilder setRenderTimestampInLocalTimeZone(boolean renderTimestampInLocalTimeZone) + { + this.renderTimestampInLocalTimeZone = renderTimestampInLocalTimeZone; + return this; + } + public Session build() { return new Session( @@ -979,7 +999,8 @@ public Session build() preparedStatements, protocolHeaders, Optional.empty(), - queryDataEncoding); + queryDataEncoding, + renderTimestampInLocalTimeZone); } } diff --git a/core/trino-main/src/main/java/io/trino/SessionRepresentation.java b/core/trino-main/src/main/java/io/trino/SessionRepresentation.java index caf538014874..085ac22c5011 100644 --- a/core/trino-main/src/main/java/io/trino/SessionRepresentation.java +++ b/core/trino-main/src/main/java/io/trino/SessionRepresentation.java @@ -72,6 +72,7 @@ public final class SessionRepresentation private final Map preparedStatements; private final String protocolName; private final Optional queryDataEncoding; + private final boolean renderTimestampInLocalTimeZone; @JsonCreator public SessionRepresentation( @@ -104,7 +105,8 @@ public SessionRepresentation( @JsonProperty("catalogRoles") Map catalogRoles, @JsonProperty("preparedStatements") Map preparedStatements, @JsonProperty("protocolName") String protocolName, - @JsonProperty("queryDataEncoding") Optional queryDataEncoding) + @JsonProperty("queryDataEncoding") Optional queryDataEncoding, + @JsonProperty("renderTimestampInLocalTimeZone") boolean renderTimestampInLocalTimeZone) { this.queryId = requireNonNull(queryId, "queryId is null"); this.querySpan = requireNonNull(querySpan, "querySpan is null"); @@ -141,6 +143,7 @@ public SessionRepresentation( catalogPropertiesBuilder.put(entry.getKey(), ImmutableMap.copyOf(entry.getValue())); } this.catalogProperties = catalogPropertiesBuilder.buildOrThrow(); + this.renderTimestampInLocalTimeZone = renderTimestampInLocalTimeZone; } @JsonProperty @@ -388,6 +391,7 @@ public Session toSession(SessionPropertyManager sessionPropertyManager, Map pages; private final List columns; private final boolean supportsParametricDateTime; + private final boolean supportsUseSessionTimeZone; private final Consumer exceptionConsumer; + private final TimeZoneKey sessionTimeZoneKey; private Page currentPage; private int rowPosition = -1; @@ -72,7 +75,9 @@ public JsonArrayResultsIterator(Session session, List pages, List getRowValues() if (!supportsParametricDateTime) { value = getLegacyValue(value, type); } + if (supportsUseSessionTimeZone) { + value = useSessionTimeZone(value, type); + } row.add(value); } catch (Throwable throwable) { @@ -131,6 +139,56 @@ private List getRowValues() return unmodifiableList(row); } + private Object useSessionTimeZone(Object value, Type type) + { + if (value == null) { + return null; + } + + if (type instanceof TimestampWithTimeZoneType) { + SqlTimestampWithTimeZone originalValue = ((SqlTimestampWithTimeZone) value); + return SqlTimestampWithTimeZone.newInstance(originalValue.getPrecision(), originalValue.getEpochMillis(), originalValue.getPicosOfMilli(), sessionTimeZoneKey); + } + + if (type instanceof ArrayType) { + Type elementType = ((ArrayType) type).getElementType(); + + if (!(elementType instanceof TimestampWithTimeZoneType)) { + return value; + } + + List listValue = (List) value; + List updateValues = new ArrayList<>(listValue.size()); + for (Object element : listValue) { + updateValues.add(useSessionTimeZone(element, elementType)); + } + + return unmodifiableList(updateValues); + } + + if (type instanceof MapType) { + Type keyType = ((MapType) type).getKeyType(); + Type valueType = ((MapType) type).getValueType(); + + Map mapValue = (Map) value; + Map result = Maps.newHashMapWithExpectedSize(mapValue.size()); + mapValue.forEach((key, val) -> result.put(useSessionTimeZone(key, keyType), useSessionTimeZone(val, valueType))); + return unmodifiableMap(result); + } + + if (type instanceof RowType) { + List fields = ((RowType) type).getFields(); + List values = (List) value; + + List result = new ArrayList<>(values.size()); + for (int i = 0; i < values.size(); i++) { + result.add(useSessionTimeZone(values.get(i), fields.get(i).getType())); + } + return unmodifiableList(result); + } + return value; + } + private Object getLegacyValue(Object value, Type type) { if (value == null) {