Add Oracle connection pool

presto-oracle/pom.xml

@@ -14,13 +14,20 @@
 
     <properties>
         <air.main.basedir>${project.parent.basedir}</air.main.basedir>
+        <dep.oracle.version>19.3.0.0</dep.oracle.version>
     </properties>
 
     <dependencies>
         <dependency>
             <groupId>com.oracle.ojdbc</groupId>
             <artifactId>ojdbc8</artifactId>
-            <version>19.3.0.0</version>
+            <version>${dep.oracle.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>com.oracle.ojdbc</groupId>
+            <artifactId>ucp</artifactId>
+            <version>${dep.oracle.version}</version>
         </dependency>
 
         <dependency>

presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleClientModule.java

@@ -53,6 +53,15 @@ public static ConnectionFactory connectionFactory(BaseJdbcConfig config, Credent
         Properties connectionProperties = new Properties();
         connectionProperties.setProperty(OracleConnection.CONNECTION_PROPERTY_INCLUDE_SYNONYMS, String.valueOf(oracleConfig.isSynonymsEnabled()));
 
+        if (oracleConfig.isConnectionPoolEnabled()) {
+            return new OraclePoolConnectorFactory(
+                    config.getConnectionUrl(),
+                    connectionProperties,
+                    credentialProvider,
+                    oracleConfig.getConnectionPoolMinSize(),
+                    oracleConfig.getConnectionPoolMaxSize());
+        }
+
         return new DriverConnectionFactory(
                 new OracleDriver(),
                 config.getConnectionUrl(),

presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleConfig.java

@@ -16,6 +16,7 @@
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
 
+import javax.validation.constraints.AssertTrue;
 import javax.validation.constraints.Max;
 import javax.validation.constraints.Min;
 import javax.validation.constraints.NotNull;
@@ -28,6 +29,9 @@ public class OracleConfig
     private boolean synonymsEnabled;
     private Integer defaultNumberScale;
     private RoundingMode numberRoundingMode = RoundingMode.UNNECESSARY;
+    private boolean connectionPoolEnabled = true;
+    private int connectionPoolMinSize = 1;
+    private int connectionPoolMaxSize = 30;
 
     @NotNull
     public boolean isSynonymsEnabled()
@@ -67,4 +71,49 @@ public OracleConfig setNumberRoundingMode(RoundingMode numberRoundingMode)
         this.numberRoundingMode = numberRoundingMode;
         return this;
     }
+
+    @NotNull
+    public boolean isConnectionPoolEnabled()
+    {
+        return connectionPoolEnabled;
+    }
+
+    @Config("oracle.connection-pool.enabled")
+    public OracleConfig setConnectionPoolEnabled(boolean connectionPoolEnabled)
+    {
+        this.connectionPoolEnabled = connectionPoolEnabled;
+        return this;
+    }
+
+    @Min(0)
+    public int getConnectionPoolMinSize()
+    {
+        return connectionPoolMinSize;
+    }
+
+    @Config("oracle.connection-pool.min-size")
+    public OracleConfig setConnectionPoolMinSize(int connectionPoolMinSize)
+    {
+        this.connectionPoolMinSize = connectionPoolMinSize;
+        return this;
+    }
+
+    @Min(1)
+    public int getConnectionPoolMaxSize()
+    {
+        return connectionPoolMaxSize;
+    }
+
+    @Config("oracle.connection-pool.max-size")
+    public OracleConfig setConnectionPoolMaxSize(int connectionPoolMaxSize)
+    {
+        this.connectionPoolMaxSize = connectionPoolMaxSize;
+        return this;
+    }
+
+    @AssertTrue(message = "Pool min size cannot be larger than max size")
+    public boolean isPoolSizedProperly()
+    {
+        return getConnectionPoolMaxSize() >= getConnectionPoolMinSize();
+    }
 }

presto-oracle/src/main/java/io/prestosql/plugin/oracle/OraclePoolConnectorFactory.java

@@ -0,0 +1,71 @@
+/*
+ * 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.prestosql.plugin.oracle;
+
+import io.prestosql.plugin.jdbc.ConnectionFactory;
+import io.prestosql.plugin.jdbc.JdbcIdentity;
+import io.prestosql.plugin.jdbc.credential.CredentialProvider;
+import oracle.jdbc.pool.OracleDataSource;
+import oracle.ucp.jdbc.PoolDataSource;
+import oracle.ucp.jdbc.PoolDataSourceFactory;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Optional;
+import java.util.Properties;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+public class OraclePoolConnectorFactory
+        implements ConnectionFactory
+{
+    private final PoolDataSource dataSource;
+    private final CredentialProvider credentialProvider;
+
+    public OraclePoolConnectorFactory(
+            String connectionUrl,
+            Properties connectionProperties,
+            CredentialProvider credentialProvider,
+            int connectionPoolMinSize,
+            int connectionPoolMaxSize)
+            throws SQLException
+    {
+        this.credentialProvider = credentialProvider;
+        this.dataSource = PoolDataSourceFactory.getPoolDataSource();
+
+        //Setting connection properties of the data source
+        this.dataSource.setConnectionFactoryClassName(OracleDataSource.class.getName());
+        this.dataSource.setURL(connectionUrl);
+
+        //Setting pool properties
+        this.dataSource.setInitialPoolSize(connectionPoolMinSize);
+        this.dataSource.setMinPoolSize(connectionPoolMinSize);
+        this.dataSource.setMaxPoolSize(connectionPoolMaxSize);
+        this.dataSource.setValidateConnectionOnBorrow(true);
+        this.dataSource.setConnectionProperties(connectionProperties);
+    }
+
+    @Override
+    public Connection openConnection(JdbcIdentity identity)
+            throws SQLException
+    {
+        Optional<String> user = credentialProvider.getConnectionUser(Optional.of(identity));
+        Optional<String> password = credentialProvider.getConnectionPassword(Optional.of(identity));
+
+        checkArgument(user.isPresent(), "Credentials returned null user");
+        checkArgument(password.isPresent(), "Credentials returned null password");
+
+        return dataSource.getConnection(user.get(), password.get());
+    }
+}

presto-oracle/src/test/java/io/prestosql/plugin/oracle/BaseOracleIntegrationSmokeTest.java

@@ -0,0 +1,91 @@
+/*
+ * 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.prestosql.plugin.oracle;
+
+import io.prestosql.testing.AbstractTestIntegrationSmokeTest;
+import io.prestosql.testing.MaterializedResult;
+import io.prestosql.testing.QueryRunner;
+import io.prestosql.tpch.TpchTable;
+import org.testcontainers.shaded.com.google.common.collect.ImmutableList;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.Test;
+
+import static io.prestosql.spi.type.VarcharType.VARCHAR;
+import static io.prestosql.testing.assertions.Assert.assertEquals;
+import static io.prestosql.tpch.TpchTable.CUSTOMER;
+import static io.prestosql.tpch.TpchTable.NATION;
+import static io.prestosql.tpch.TpchTable.ORDERS;
+import static io.prestosql.tpch.TpchTable.REGION;
+import static org.assertj.core.api.Assertions.assertThat;
+
+abstract class BaseOracleIntegrationSmokeTest
+        extends AbstractTestIntegrationSmokeTest
+{
+    private TestingOracleServer oracleServer;
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        oracleServer = new TestingOracleServer();
+        return createOracleQueryRunner(oracleServer, ImmutableList.of(CUSTOMER, NATION, ORDERS, REGION));
+    }
+
+    protected abstract QueryRunner createOracleQueryRunner(TestingOracleServer server, Iterable<TpchTable<?>> tables)
+            throws Exception;
+
+    @AfterClass(alwaysRun = true)
+    public final void destroy()
+    {
+        oracleServer.close();
+    }
+
+    @Test
+    @Override
+    public void testDescribeTable()
+    {
+        MaterializedResult expectedColumns = MaterializedResult.resultBuilder(getQueryRunner().getDefaultSession(), VARCHAR, VARCHAR, VARCHAR, VARCHAR)
+                .row("orderkey", "decimal(19,0)", "", "")
+                .row("custkey", "decimal(19,0)", "", "")
+                .row("orderstatus", "varchar(1)", "", "")
+                .row("totalprice", "double", "", "")
+                .row("orderdate", "timestamp(3)", "", "")
+                .row("orderpriority", "varchar(15)", "", "")
+                .row("clerk", "varchar(15)", "", "")
+                .row("shippriority", "decimal(10,0)", "", "")
+                .row("comment", "varchar(79)", "", "")
+                .build();
+        MaterializedResult actualColumns = computeActual("DESCRIBE orders");
+        assertEquals(actualColumns, expectedColumns);
+    }
+
+    @Test
+    @Override
+    public void testShowCreateTable()
+    {
+        assertThat((String) computeActual("SHOW CREATE TABLE orders").getOnlyValue())
+                // If the connector reports additional column properties, the expected value needs to be adjusted in the test subclass
+                .matches("CREATE TABLE \\w+\\.\\w+\\.orders \\Q(\n" +
+                        "   orderkey decimal(19, 0),\n" +
+                        "   custkey decimal(19, 0),\n" +
+                        "   orderstatus varchar(1),\n" +
+                        "   totalprice double,\n" +
+                        "   orderdate timestamp(3),\n" +
+                        "   orderpriority varchar(15),\n" +
+                        "   clerk varchar(15),\n" +
+                        "   shippriority decimal(10, 0),\n" +
+                        "   comment varchar(79)\n" +
+                        ")");
+    }
+}

presto-oracle/src/test/java/io/prestosql/plugin/oracle/OracleQueryRunner.java

@@ -13,7 +13,6 @@
  */
 package io.prestosql.plugin.oracle;
 
-import com.google.common.collect.ImmutableList;
 import io.airlift.log.Logger;
 import io.airlift.log.Logging;
 import io.prestosql.Session;
@@ -36,19 +35,19 @@ public final class OracleQueryRunner
 {
     private OracleQueryRunner() {}
 
-    public static DistributedQueryRunner createOracleQueryRunner(TestingOracleServer server)
+    public static DistributedQueryRunner createOracleQueryRunner(TestingOracleServer server, Iterable<TpchTable<?>> tables)
             throws Exception
     {
-        return createOracleQueryRunner(server, ImmutableList.of());
+        return createQueryRunner(server, tables, false);
     }
 
-    public static DistributedQueryRunner createOracleQueryRunner(TestingOracleServer server, TpchTable<?>... tables)
+    public static DistributedQueryRunner createOraclePoolQueryRunner(TestingOracleServer server, Iterable<TpchTable<?>> tables)
             throws Exception
     {
-        return createOracleQueryRunner(server, ImmutableList.copyOf(tables));
+        return createQueryRunner(server, tables, true);
     }
 
-    public static DistributedQueryRunner createOracleQueryRunner(TestingOracleServer server, Iterable<TpchTable<?>> tables)
+    private static DistributedQueryRunner createQueryRunner(TestingOracleServer server, Iterable<TpchTable<?>> tables, boolean connectionPoolEnable)
             throws Exception
     {
         DistributedQueryRunner queryRunner = null;
@@ -63,6 +62,7 @@ public static DistributedQueryRunner createOracleQueryRunner(TestingOracleServer
             connectorProperties.putIfAbsent("connection-user", TEST_USER);
             connectorProperties.putIfAbsent("connection-password", TEST_PASS);
             connectorProperties.putIfAbsent("allow-drop-table", "true");
+            connectorProperties.putIfAbsent("oracle.connection-pool.enabled", String.valueOf(connectionPoolEnable));
 
             queryRunner.installPlugin(new OraclePlugin());
             queryRunner.createCatalog("oracle", "oracle", connectorProperties);

presto-oracle/src/test/java/io/prestosql/plugin/oracle/TestOracleConfig.java

@@ -35,7 +35,10 @@ public void testDefaults()
         assertRecordedDefaults(recordDefaults(OracleConfig.class)
                 .setSynonymsEnabled(false)
                 .setDefaultNumberScale(null)
-                .setNumberRoundingMode(RoundingMode.UNNECESSARY));
+                .setNumberRoundingMode(RoundingMode.UNNECESSARY)
+                .setConnectionPoolEnabled(true)
+                .setConnectionPoolMinSize(1)
+                .setConnectionPoolMaxSize(30));
     }
 
     @Test
@@ -45,12 +48,18 @@ public void testExplicitPropertyMappings()
                 .put("oracle.synonyms.enabled", "true")
                 .put("oracle.number.default-scale", "2")
                 .put("oracle.number.rounding-mode", "CEILING")
+                .put("oracle.connection-pool.enabled", "false")
+                .put("oracle.connection-pool.min-size", "10")
+                .put("oracle.connection-pool.max-size", "20")
                 .build();
 
         OracleConfig expected = new OracleConfig()
                 .setSynonymsEnabled(true)
                 .setDefaultNumberScale(2)
-                .setNumberRoundingMode(RoundingMode.CEILING);
+                .setNumberRoundingMode(RoundingMode.CEILING)
+                .setConnectionPoolEnabled(false)
+                .setConnectionPoolMinSize(10)
+                .setConnectionPoolMaxSize(20);
 
         assertFullMapping(properties, expected);
     }
@@ -71,5 +80,19 @@ public void testValidation()
                 "defaultNumberScale",
                 "must be less than or equal to 38",
                 Max.class);
+
+        assertFailsValidation(
+                new OracleConfig()
+                        .setConnectionPoolMinSize(-1),
+                "connectionPoolMinSize",
+                "must be greater than or equal to 0",
+                Min.class);
+
+        assertFailsValidation(
+                new OracleConfig()
+                        .setConnectionPoolMaxSize(0),
+                "connectionPoolMaxSize",
+                "must be greater than or equal to 1",
+                Min.class);
     }
 }