Add connection properties for Druid connector

Author: yangshangqing95

docs/src/main/sphinx/connector/druid.md

@@ -45,6 +45,27 @@ connection-user=root
 connection-password=secret
 ```
 
+Additionally, the following configuration properties can be set depending on the use-case:
+
+:::{list-table} Druid configuration properties
+:widths: 30, 55, 15
+:header-rows: 1
+
+* - Property name
+  - Description
+  - Default
+* - `druid.execution-timeout`
+  - Query timeout for druid queries, beyond which unfinished queries will be cancelled. 
+    0 timeout means no timeout (up to the server-side maximum query timeout,
+    `druid.server.http.maxQueryTimeout`) 
+    example: `10s`, `1m`, `1000ms`.
+  - None.
+* - `druid.sql-timezone`
+  - Time zone for a connection. For example, "America/Los_Angeles" or an offset 
+    like "-08:00". This parameter affects how time functions and timestamp literals behave.
+  - `UTC`
+:::
+
 Now you can access your Druid database in Trino with the `example` catalog
 name from the properties file.
 

plugin/trino-druid/pom.xml

@@ -26,6 +26,16 @@
             <classifier>classes</classifier>
         </dependency>
 
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>configuration</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>units</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-base-jdbc</artifactId>
@@ -116,14 +126,14 @@
         </dependency>
 
         <dependency>
-            <groupId>io.airlift</groupId>
-            <artifactId>units</artifactId>
-            <scope>runtime</scope>
+            <groupId>com.squareup.okhttp3</groupId>
+            <artifactId>okhttp-jvm</artifactId>
+            <scope>test</scope>
         </dependency>
 
         <dependency>
-            <groupId>com.squareup.okhttp3</groupId>
-            <artifactId>okhttp-jvm</artifactId>
+            <groupId>io.airlift</groupId>
+            <artifactId>configuration-testing</artifactId>
             <scope>test</scope>
         </dependency>
 

plugin/trino-druid/src/main/java/io/trino/plugin/druid/DruidConfig.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.druid;
+
+import io.airlift.configuration.Config;
+import io.airlift.units.Duration;
+
+public class DruidConfig
+{
+    private Duration executionTimeout; // milliseconds
+
+    private String sqlTimeZone;
+
+    @Config("druid.execution-timeout")
+    public DruidConfig setExecutionTimeout(Duration executionTimeout)
+    {
+        this.executionTimeout = executionTimeout;
+        return this;
+    }
+
+    public Duration getExecutionTimeout()
+    {
+        return executionTimeout;
+    }
+
+    @Config("druid.sql-timezone")
+    public DruidConfig setSqlTimeZone(String sqlTimeZone)
+    {
+        this.sqlTimeZone = sqlTimeZone;
+        return this;
+    }
+
+    public String getSqlTimeZone()
+    {
+        return sqlTimeZone;
+    }
+}

plugin/trino-druid/src/main/java/io/trino/plugin/druid/DruidJdbcClientModule.java

@@ -13,6 +13,7 @@
  */
 package io.trino.plugin.druid;
 
+import com.google.common.base.Strings;
 import com.google.inject.Binder;
 import com.google.inject.Module;
 import com.google.inject.Provides;
@@ -29,7 +30,10 @@
 import io.trino.spi.function.table.ConnectorTableFunction;
 import org.apache.calcite.avatica.remote.Driver;
 
+import java.util.Properties;
+
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
+import static io.airlift.configuration.ConfigBinder.configBinder;
 
 public class DruidJdbcClientModule
         implements Module
@@ -38,16 +42,30 @@ public class DruidJdbcClientModule
     public void configure(Binder binder)
     {
         binder.bind(JdbcClient.class).annotatedWith(ForBaseJdbc.class).to(DruidJdbcClient.class).in(Scopes.SINGLETON);
+        configBinder(binder).bindConfig(DruidConfig.class);
         newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(Query.class).in(Scopes.SINGLETON);
     }
 
     @Provides
     @Singleton
     @ForBaseJdbc
-    public static ConnectionFactory createConnectionFactory(BaseJdbcConfig config, CredentialProvider credentialProvider, OpenTelemetry openTelemetry)
+    public static ConnectionFactory createConnectionFactory(BaseJdbcConfig config, CredentialProvider credentialProvider, DruidConfig druidConfig, OpenTelemetry openTelemetry)
     {
         return DriverConnectionFactory.builder(new Driver(), config.getConnectionUrl(), credentialProvider)
+                .setConnectionProperties(getConnectionProperties(druidConfig))
                 .setOpenTelemetry(openTelemetry)
                 .build();
     }
+
+    private static Properties getConnectionProperties(DruidConfig druidConfig)
+    {
+        Properties connectionProperties = new Properties();
+        if (druidConfig.getExecutionTimeout() != null) {
+            connectionProperties.setProperty("timeout", String.valueOf(druidConfig.getExecutionTimeout().toMillis()));
+        }
+        if (!Strings.isNullOrEmpty(druidConfig.getSqlTimeZone())) {
+            connectionProperties.setProperty("sqlTimeZone", druidConfig.getSqlTimeZone());
+        }
+        return connectionProperties;
+    }
 }

plugin/trino-druid/src/test/java/io/trino/plugin/druid/TestDruidConfig.java

@@ -0,0 +1,51 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.druid;
+
+import com.google.common.collect.ImmutableMap;
+import io.airlift.units.Duration;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+class TestDruidConfig
+{
+    @Test
+    void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(DruidConfig.class)
+                .setExecutionTimeout(null)
+                .setSqlTimeZone(null));
+    }
+
+    @Test
+    void testExplicitPropertyMappings()
+    {
+        Map<String, String> properties = ImmutableMap.<String, String>builder()
+                .put("druid.execution-timeout", "10s")
+                .put("druid.sql-timezone", "UTC")
+                .buildOrThrow();
+
+        DruidConfig expected = new DruidConfig()
+                .setExecutionTimeout(new Duration(10, SECONDS))
+                .setSqlTimeZone("UTC");
+
+        assertFullMapping(properties, expected);
+    }
+}