Optionally ignore session zone from client

Backport of trinodb/trino#1164 PR: prestodb#13140 Co-authored-by: Ariel Weisberg <[email protected]>
cjjxfli · Aug 5, 2019 · 6b3c211 · 6b3c211
1 parent bafa2a1
commit 6b3c211
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 5 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/server/QuerySessionSupplier.java b/presto-main/src/main/java/com/facebook/presto/server/QuerySessionSupplier.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.security.AccessControl;
 import com.facebook.presto.spi.QueryId;
 import com.facebook.presto.spi.security.Identity;
+import com.facebook.presto.spi.type.TimeZoneKey;
 import com.facebook.presto.sql.SqlEnvironmentConfig;
 import com.facebook.presto.sql.SqlPath;
 import com.facebook.presto.transaction.TransactionManager;
@@ -42,6 +43,7 @@ public class QuerySessionSupplier
     private final AccessControl accessControl;
     private final SessionPropertyManager sessionPropertyManager;
     private final Optional<String> path;
+    private final Optional<TimeZoneKey> forcedSessionTimeZone;
 
     @Inject
     public QuerySessionSupplier(
@@ -53,7 +55,9 @@ public QuerySessionSupplier(
         this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
         this.accessControl = requireNonNull(accessControl, "accessControl is null");
         this.sessionPropertyManager = requireNonNull(sessionPropertyManager, "sessionPropertyManager is null");
+        requireNonNull(config, "config is null");
         this.path = requireNonNull(config.getPath(), "path is null");
+        this.forcedSessionTimeZone = requireNonNull(config.getForcedSessionTimeZone(), "forcedSessionTimeZone is null");
     }
 
     @Override
@@ -81,7 +85,10 @@ public Session createSession(QueryId queryId, SessionContext context)
             sessionBuilder.setPath(new SqlPath(Optional.of(context.getPath())));
         }
 
-        if (context.getTimeZoneId() != null) {
+        if (forcedSessionTimeZone.isPresent()) {
+            sessionBuilder.setTimeZoneKey(forcedSessionTimeZone.get());
+        }
+        else if (context.getTimeZoneId() != null) {
             sessionBuilder.setTimeZoneKey(getTimeZoneKey(context.getTimeZoneId()));
         }
 

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/SqlEnvironmentConfig.java b/presto-main/src/main/java/com/facebook/presto/sql/SqlEnvironmentConfig.java
@@ -13,13 +13,25 @@
  */
 package com.facebook.presto.sql;
 
+import com.facebook.presto.spi.type.TimeZoneKey;
 import io.airlift.configuration.Config;
+import io.airlift.configuration.ConfigDescription;
+
+import javax.annotation.Nullable;
+import javax.validation.constraints.NotNull;
 
 import java.util.Optional;
 
 public class SqlEnvironmentConfig
 {
     private Optional<String> path = Optional.empty();
+    private Optional<TimeZoneKey> forcedSessionTimeZone = Optional.empty();
+
+    @NotNull
+    public Optional<String> getPath()
+    {
+        return path;
+    }
 
     @Config("sql.path")
     public SqlEnvironmentConfig setPath(String path)
@@ -28,8 +40,18 @@ public SqlEnvironmentConfig setPath(String path)
         return this;
     }
 
-    public Optional<String> getPath()
+    @NotNull
+    public Optional<TimeZoneKey> getForcedSessionTimeZone()
     {
-        return path;
+        return forcedSessionTimeZone;
+    }
+
+    @Config("sql.forced-session-time-zone")
+    @ConfigDescription("User session time zone overriding value sent by client")
+    public SqlEnvironmentConfig setForcedSessionTimeZone(@Nullable String timeZoneId)
+    {
+        this.forcedSessionTimeZone = Optional.ofNullable(timeZoneId)
+                .map(TimeZoneKey::getTimeZoneKey);
+        return this;
     }
 }
diff --git a/presto-main/src/test/java/com/facebook/presto/sql/TestSqlEnvironmentConfig.java b/presto-main/src/test/java/com/facebook/presto/sql/TestSqlEnvironmentConfig.java
@@ -29,18 +29,21 @@ public class TestSqlEnvironmentConfig
     public void testDefaults()
     {
         assertRecordedDefaults(recordDefaults(SqlEnvironmentConfig.class)
-                .setPath(null));
+                .setPath(null)
+                .setForcedSessionTimeZone(null));
     }
 
     @Test
     public void testExplicitPropertyMappings()
     {
         Map<String, String> properties = new ImmutableMap.Builder<String, String>()
                 .put("sql.path", "a.b, c.d")
+                .put("sql.forced-session-time-zone", "UTC")
                 .build();
 
         SqlEnvironmentConfig expected = new SqlEnvironmentConfig()
-                .setPath("a.b, c.d");
+                .setPath("a.b, c.d")
+                .setForcedSessionTimeZone("UTC");
 
         assertFullMapping(properties, expected);
     }