Add a new JDBC driver parameter - timeZoneID

presto-docs/src/main/sphinx/installation/jdbc.rst

@@ -83,6 +83,8 @@ Name                              Description
                                   If neither this property nor ``ApplicationName`` are set, the source
                                   for the query will be ``presto-jdbc``.
 ``accessToken``                   Access token for token based authentication.
+``timeZoneId``                    Timezone to be used for timestamp columns in query output.
+                                  Example: ``timeZoneId=UTC``.
 ``SSL``                           Use HTTPS for connections
 ``SSLKeyStorePath``               The location of the Java KeyStore file that contains the certificate
                                   and private key to use for authentication.

presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java

@@ -54,6 +54,7 @@ final class ConnectionProperties
     public static final ConnectionProperty<File> KERBEROS_KEYTAB_PATH = new KerberosKeytabPath();
     public static final ConnectionProperty<File> KERBEROS_CREDENTIAL_CACHE_PATH = new KerberosCredentialCachePath();
     public static final ConnectionProperty<String> ACCESS_TOKEN = new AccessToken();
+    public static final ConnectionProperty<String> TIMEZONE_ID = new TimeZoneId();
     public static final ConnectionProperty<Map<String, String>> EXTRA_CREDENTIALS = new ExtraCredentials();
     public static final ConnectionProperty<Map<String, String>> CUSTOM_HEADERS = new CustomHeaders();
     public static final ConnectionProperty<Map<String, String>> SESSION_PROPERTIES = new SessionProperties();
@@ -79,6 +80,7 @@ final class ConnectionProperties
             .add(KERBEROS_KEYTAB_PATH)
             .add(KERBEROS_CREDENTIAL_CACHE_PATH)
             .add(ACCESS_TOKEN)
+            .add(TIMEZONE_ID)
             .add(EXTRA_CREDENTIALS)
             .add(CUSTOM_HEADERS)
             .add(SESSION_PROPERTIES)
@@ -301,6 +303,15 @@ public AccessToken()
         }
     }
 
+    private static class TimeZoneId
+            extends AbstractConnectionProperty<String>
+    {
+        public TimeZoneId()
+        {
+            super("timeZoneId", NOT_REQUIRED, ALLOWED, STRING_CONVERTER);
+        }
+    }
+
     private static class ExtraCredentials
             extends AbstractConnectionProperty<Map<String, String>>
     {

presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java

@@ -50,7 +50,6 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
-import java.util.TimeZone;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -118,7 +117,7 @@ public class PrestoConnection
         this.connectionProperties = uri.getProperties();
         this.queryExecutor = requireNonNull(queryExecutor, "queryExecutor is null");
 
-        timeZoneId.set(TimeZone.getDefault().getID());
+        timeZoneId.set(uri.getTimeZoneId());
         locale.set(Locale.getDefault());
 
         this.queryInterceptorInstances = ImmutableList.copyOf(uri.getQueryInterceptors());

presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java

@@ -27,12 +27,14 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.sql.SQLException;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.Properties;
+import java.util.TimeZone;
 
 import static com.facebook.presto.client.GCSOAuthInterceptor.GCS_CREDENTIALS_PATH_KEY;
 import static com.facebook.presto.client.GCSOAuthInterceptor.GCS_OAUTH_SCOPES_KEY;
@@ -66,6 +68,7 @@
 import static com.facebook.presto.jdbc.ConnectionProperties.SSL_KEY_STORE_PATH;
 import static com.facebook.presto.jdbc.ConnectionProperties.SSL_TRUST_STORE_PASSWORD;
 import static com.facebook.presto.jdbc.ConnectionProperties.SSL_TRUST_STORE_PATH;
+import static com.facebook.presto.jdbc.ConnectionProperties.TIMEZONE_ID;
 import static com.facebook.presto.jdbc.ConnectionProperties.USER;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static java.lang.String.format;
@@ -149,6 +152,21 @@ public Properties getProperties()
         return properties;
     }
 
+    public String getTimeZoneId()
+            throws SQLException
+    {
+        Optional<String> timezone = TIMEZONE_ID.getValue(properties);
+
+        if (timezone.isPresent()) {
+            List<String> timeZoneIds = Arrays.asList(TimeZone.getAvailableIDs());
+            if (!timeZoneIds.contains(timezone.get())) {
+                throw new SQLException("Specified timeZoneId is not supported: " + timezone.get());
+            }
+            return timezone.get();
+        }
+        return TimeZone.getDefault().getID();
+    }
+
     public Map<String, String> getExtraCredentials()
             throws SQLException
     {

presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestPrestoDriver.java

@@ -1680,6 +1680,22 @@ public void testEncodeDecodeSessionValue()
         assertTrue(isValidSessionValue);
     }
 
+    @Test
+    public void testTimeZoneIdParameter()
+            throws Exception
+    {
+        String sql = "SELECT current_timezone() zone, TIMESTAMP '2001-02-03 3:04:05' ts";
+
+        try (Connection connection = createConnectionWithParameter("timeZoneId=UTC")) {
+            try (Statement statement = connection.createStatement();
+                    ResultSet rs = statement.executeQuery(sql)) {
+                assertTrue(rs.next());
+                assertEquals(rs.getString("zone"), "UTC");
+                assertEquals(rs.getTimestamp("ts"), new Timestamp(new DateTime(2001, 2, 3, 3, 4, 5, DateTimeZone.UTC).getMillis()));
+            }
+        }
+    }
+
     private QueryState getQueryState(String queryId)
             throws SQLException
     {
@@ -1729,6 +1745,13 @@ private Connection createConnection(String catalog, String schema)
         return DriverManager.getConnection(url, "test", null);
     }
 
+    private Connection createConnectionWithParameter(String parameter)
+            throws SQLException
+    {
+        String url = format("jdbc:presto://%s?%s", server.getAddress(), parameter);
+        return DriverManager.getConnection(url, "test", null);
+    }
+
     private static List<List<Object>> readRows(ResultSet rs)
             throws SQLException
     {

presto-tests/src/test/java/com/facebook/presto/tests/TestBrutalShutdown.java

@@ -42,7 +42,7 @@
 @Test(singleThreaded = true)
 public class TestBrutalShutdown
 {
-    private static final long SHUTDOWN_TIMEOUT_MILLIS = 240_000;
+    private static final long SHUTDOWN_TIMEOUT_MILLIS = 600_000;
     private static final Session TINY_SESSION = testSessionBuilder()
             .setCatalog("tpch")
             .setSchema("tiny")