Open JdbcConnectionsPool for extensions

Refactor `JdbcConnectionsPool.connectionFor` splitting into smaller methods.
trinodb · Jun 20, 2023 · f94163c · f94163c
1 parent fb0d9e8
commit f94163c
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tempto-core/src/main/java/io/trino/tempto/query/JdbcConnectionsPool.java b/tempto-core/src/main/java/io/trino/tempto/query/JdbcConnectionsPool.java
@@ -22,13 +22,20 @@
 
 import static com.google.common.collect.Maps.newHashMap;
 import static io.trino.tempto.internal.query.JdbcUtils.dataSource;
+import static java.util.Objects.requireNonNull;
 
 public class JdbcConnectionsPool
 {
     private final Map<JdbcConnectivityParamsState, DataSource> dataSources = newHashMap();
 
     public Connection connectionFor(JdbcConnectivityParamsState jdbcParamsState)
             throws SQLException
+    {
+        return configureConnection(jdbcParamsState, createConnection(jdbcParamsState));
+    }
+
+    protected Connection createConnection(JdbcConnectivityParamsState jdbcParamsState)
+            throws SQLException
     {
         if (!dataSources.containsKey(jdbcParamsState)) {
             dataSources.put(jdbcParamsState, dataSource(jdbcParamsState));
@@ -39,6 +46,13 @@ public Connection connectionFor(JdbcConnectivityParamsState jdbcParamsState)
             // this should never happen, `javax.sql.DataSource#getConnection()` should not return null
             throw new IllegalStateException("No connection was created for: " + jdbcParamsState.getName());
         }
+        return connection;
+    }
+
+    protected static Connection configureConnection(JdbcConnectivityParamsState jdbcParamsState, Connection connection)
+            throws SQLException
+    {
+        requireNonNull(connection, "connection is null");
         if (!jdbcParamsState.prepareStatements.isEmpty()) {
             try (Statement statement = connection.createStatement()) {
                 for (String query : jdbcParamsState.prepareStatements) {