[Improvement] JDBC Connection Configuration with Flexible SSL and Property Injection

Author: Felicity-3786

gateway-ha/pom.xml

@@ -78,6 +78,18 @@
             <classifier>classes</classifier>
         </dependency>
 
+        <dependency>
+            <groupId>com.h2database</groupId>
+            <artifactId>h2</artifactId>
+            <version>1.4.192</version>
+        </dependency>
+
+        <dependency>
+            <groupId>com.mysql</groupId>
+            <artifactId>mysql-connector-j</artifactId>
+            <version>9.3.0</version>
+        </dependency>
+
         <dependency>
             <groupId>com.nimbusds</groupId>
             <artifactId>nimbus-jose-jwt</artifactId>
@@ -91,6 +103,12 @@
             <classifier>jdk11</classifier>
         </dependency>
 
+        <dependency>
+            <groupId>com.oracle.database.jdbc</groupId>
+            <artifactId>ojdbc11</artifactId>
+            <version>23.8.0.25.04</version>
+        </dependency>
+
         <dependency>
             <groupId>com.squareup.okhttp3</groupId>
             <artifactId>okhttp</artifactId>
@@ -268,23 +286,14 @@
         </dependency>
 
         <dependency>
-            <groupId>org.weakref</groupId>
-            <artifactId>jmxutils</artifactId>
-        </dependency>
-
-        <dependency>
-            <groupId>com.mysql</groupId>
-            <artifactId>mysql-connector-j</artifactId>
-            <version>9.3.0</version>
-            <scope>runtime</scope>
+            <groupId>org.postgresql</groupId>
+            <artifactId>postgresql</artifactId>
+            <version>42.7.5</version>
         </dependency>
 
         <dependency>
-            <groupId>com.oracle.database.jdbc</groupId>
-            <artifactId>ojdbc11-production</artifactId>
-            <version>23.8.0.25.04</version>
-            <type>pom</type>
-            <scope>runtime</scope>
+            <groupId>org.weakref</groupId>
+            <artifactId>jmxutils</artifactId>
         </dependency>
 
         <!-- Required for ClusterStatsJdbcMonitor -->
@@ -316,21 +325,7 @@
             <scope>runtime</scope>
         </dependency>
 
-        <dependency>
-            <groupId>org.postgresql</groupId>
-            <artifactId>postgresql</artifactId>
-            <version>42.7.7</version>
-            <scope>runtime</scope>
-        </dependency>
-
         <!-- Test deps -->
-        <dependency>
-            <groupId>com.h2database</groupId>
-            <artifactId>h2</artifactId>
-            <version>1.4.192</version>
-            <scope>test</scope>
-        </dependency>
-
         <dependency>
             <groupId>com.squareup.okhttp3</groupId>
             <artifactId>mockwebserver</artifactId>

gateway-ha/src/main/java/io/trino/gateway/ha/config/DataStoreConfiguration.java

@@ -13,6 +13,13 @@
  */
 package io.trino.gateway.ha.config;
 
+import java.sql.Driver;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+import java.util.Enumeration;
+
+import static java.util.Objects.requireNonNull;
+
 public class DataStoreConfiguration
 {
     private String jdbcUrl;
@@ -21,10 +28,15 @@ public class DataStoreConfiguration
     private String driver;
     private Integer queryHistoryHoursRetention = 4;
     private boolean runMigrationsEnabled = true;
+    private DataStoreType dataStoreType;
+
+    // TODO: Refactor to decouple DataStoreConfiguration from a specific
+    //  database implementation after adopting the Airlift configuration framework (https://github.com/trinodb/trino-gateway/issues/378)
+    private MySqlConfiguration mySqlConfiguration = new MySqlConfiguration();
 
     public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention, boolean runMigrationsEnabled)
     {
-        this.jdbcUrl = jdbcUrl;
+        this.jdbcUrl = requireNonNull(jdbcUrl, "jdbc must be set");
         this.user = user;
         this.password = password;
         this.driver = driver;
@@ -93,4 +105,45 @@ public void setRunMigrationsEnabled(boolean runMigrationsEnabled)
     {
         this.runMigrationsEnabled = runMigrationsEnabled;
     }
+
+    public MySqlConfiguration getMySqlConfiguration()
+    {
+        return mySqlConfiguration;
+    }
+
+    public void setMysqlConfiguration(MySqlConfiguration mySqlConfig)
+    {
+        this.mySqlConfiguration = mySqlConfig;
+    }
+
+    public DataStoreType getDataStoreType()
+    {
+        if (dataStoreType != null) {
+            return dataStoreType;
+        }
+        String jdbcUrl = getJdbcUrl();
+        try {
+            Enumeration<Driver> drivers = DriverManager.getDrivers();
+            while (drivers.hasMoreElements()) {
+                Driver driver = drivers.nextElement();
+                if (driver.acceptsURL(jdbcUrl)) {
+                    for (DataStoreType dataStoreType : DataStoreType.values()) {
+                        if (dataStoreType.getDriverClass().isAssignableFrom(driver.getClass())) {
+                            return dataStoreType;
+                        }
+                    }
+                    break;
+                }
+            }
+        }
+        catch (SQLException e) {
+            throw new RuntimeException("Error enumerating JDBC drivers", e);
+        }
+        throw new IllegalStateException("Unable to infer DataStoreType for URL: " + jdbcUrl);
+    }
+
+    public void setDataStoreType(DataStoreType backendType)
+    {
+        this.dataStoreType = backendType;
+    }
 }

gateway-ha/src/main/java/io/trino/gateway/ha/config/DataStoreType.java

@@ -0,0 +1,38 @@
+/*
+ * 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.gateway.ha.config;
+
+import java.sql.Driver;
+
+public enum DataStoreType {
+    ORACLE(oracle.jdbc.driver.OracleDriver.class),
+    MYSQL(com.mysql.cj.jdbc.Driver.class),
+    POSTGRES(org.postgresql.Driver.class),
+    H2(org.h2.Driver.class);
+
+    private final Class<? extends Driver> driverClass;
+
+    DataStoreType(Class<? extends Driver> driverClass)
+    {
+        this.driverClass = driverClass;
+    }
+
+    /**
+     * Returns the JDBC Driver class associated with this data store type.
+     */
+    public Class<? extends Driver> getDriverClass()
+    {
+        return driverClass;
+    }
+}

gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java

@@ -43,6 +43,7 @@ public class HaGatewayConfiguration
     private List<String> statementPaths = ImmutableList.of(V1_STATEMENT_PATH);
     private boolean includeClusterHostInResponse;
     private ProxyResponseConfiguration proxyResponseConfiguration = new ProxyResponseConfiguration();
+    private boolean oracleBackend;
 
     private RequestAnalyzerConfig requestAnalyzerConfig = new RequestAnalyzerConfig();
 
@@ -297,4 +298,14 @@ public HaGatewayConfigurationException(String message)
             super(message);
         }
     }
+
+    public boolean isOracleBackend()
+    {
+        return oracleBackend;
+    }
+
+    public void setOracleBackend(boolean oracleBackend)
+    {
+        this.oracleBackend = oracleBackend;
+    }
 }

gateway-ha/src/main/java/io/trino/gateway/ha/config/MySqlConfiguration.java

@@ -0,0 +1,95 @@
+/*
+ * 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.gateway.ha.config;
+
+import com.mysql.cj.conf.PropertyDefinitions.SslMode;
+
+/**
+ * Configuration for MySQL SSL (client cert and truststore settings).
+ */
+public class MySqlConfiguration
+{
+    private SslMode sslMode = SslMode.DISABLED;
+    private String clientCertificateKeyStoreUrl;
+    private String clientCertificateKeyStorePassword;
+    private String clientCertificateKeyStoreType;
+    private String trustCertificateKeyStoreUrl;
+    private String trustCertificateKeyStorePassword;
+
+    public SslMode getSslMode()
+    {
+        return sslMode;
+    }
+
+    public MySqlConfiguration setSslMode(SslMode sslMode)
+    {
+        this.sslMode = sslMode;
+        return this;
+    }
+
+    public String getClientCertificateKeyStoreUrl()
+    {
+        return clientCertificateKeyStoreUrl;
+    }
+
+    public MySqlConfiguration setClientCertificateKeyStoreUrl(String url)
+    {
+        this.clientCertificateKeyStoreUrl = url;
+        return this;
+    }
+
+    public String getClientCertificateKeyStorePassword()
+    {
+        return clientCertificateKeyStorePassword;
+    }
+
+    public MySqlConfiguration setClientCertificateKeyStorePassword(String password)
+    {
+        this.clientCertificateKeyStorePassword = password;
+        return this;
+    }
+
+    public String getClientCertificateKeyStoreType()
+    {
+        return clientCertificateKeyStoreType;
+    }
+
+    public MySqlConfiguration setClientCertificateKeyStoreType(String type)
+    {
+        this.clientCertificateKeyStoreType = type;
+        return this;
+    }
+
+    public String getTrustCertificateKeyStoreUrl()
+    {
+        return trustCertificateKeyStoreUrl;
+    }
+
+    public MySqlConfiguration setTrustCertificateKeyStoreUrl(String url)
+    {
+        this.trustCertificateKeyStoreUrl = url;
+        return this;
+    }
+
+    public String getTrustCertificateKeyStorePassword()
+    {
+        return trustCertificateKeyStorePassword;
+    }
+
+    public MySqlConfiguration setTrustCertificateKeyStorePassword(String password)
+    {
+        this.trustCertificateKeyStorePassword = password;
+        return this;
+    }
+}

gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java

@@ -39,6 +39,7 @@
 import io.trino.gateway.ha.config.RoutingRulesConfiguration;
 import io.trino.gateway.ha.config.RulesExternalConfiguration;
 import io.trino.gateway.ha.config.UserConfiguration;
+import io.trino.gateway.ha.persistence.DefaultJdbcPropertiesProvider;
 import io.trino.gateway.ha.persistence.JdbcConnectionManager;
 import io.trino.gateway.ha.router.BackendStateManager;
 import io.trino.gateway.ha.router.ForRouter;
@@ -117,10 +118,10 @@ public HaGatewayProviderModule(HaGatewayConfiguration configuration)
         oAuth2GatewayCookieConfigurationPropertiesProvider.initialize(configuration.getOauth2GatewayCookieConfiguration());
 
         Jdbi jdbi = Jdbi.create(configuration.getDataStore().getJdbcUrl(), configuration.getDataStore().getUser(), configuration.getDataStore().getPassword());
-        JdbcConnectionManager connectionManager = new JdbcConnectionManager(jdbi, configuration.getDataStore());
+        JdbcConnectionManager connectionManager = new JdbcConnectionManager(jdbi, configuration.getDataStore(), new DefaultJdbcPropertiesProvider());
         resourceGroupsManager = new HaResourceGroupsManager(connectionManager);
         gatewayBackendManager = new HaGatewayManager(jdbi, configuration.getRouting());
-        queryHistoryManager = new HaQueryHistoryManager(jdbi, configuration.getDataStore().getJdbcUrl().startsWith("jdbc:oracle"));
+        queryHistoryManager = new HaQueryHistoryManager(jdbi, configuration.getDataStore());
     }
 
     private LbOAuthManager getOAuthManager(HaGatewayConfiguration configuration)

gateway-ha/src/main/java/io/trino/gateway/ha/persistence/DefaultJdbcPropertiesProvider.java

@@ -0,0 +1,47 @@
+/*
+ * 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.gateway.ha.persistence;
+
+import io.trino.gateway.ha.config.DataStoreConfiguration;
+
+import java.util.Properties;
+
+/**
+ * Default JDBC properties provider used as a fallback when no database-specific
+ * {@link JdbcPropertiesProvider} supports the given {@link DataStoreConfiguration}.
+ *
+ * <p>This provider simply sets the basic "user" and "password" properties
+ * and should always be the last provider in the list of available providers.
+ *
+ * <p>If a more specific provider (e.g., for MySQL, Oracle, etc.) supports the configuration,
+ * it should be preferred over this basic fallback.
+ */
+public class DefaultJdbcPropertiesProvider
+        implements JdbcPropertiesProvider
+{
+    @Override
+    public boolean supports(DataStoreConfiguration configuration)
+    {
+        return true;
+    }
+
+    @Override
+    public Properties getProperties(DataStoreConfiguration configuration)
+    {
+        Properties properties = new Properties();
+        properties.setProperty("user", configuration.getUser());
+        properties.setProperty("password", configuration.getPassword());
+        return properties;
+    }
+}