From 6ec0d650455b8ec09de45c0ac5618cfca217a34c Mon Sep 17 00:00:00 2001
From: Joao Boto <boto@boto.pro>
Date: Mon, 15 Jun 2020 10:17:22 +0200
Subject: [PATCH] Add Oracle fetch size

---
 .../io/prestosql/plugin/oracle/OracleClient.java   |  4 +++-
 .../io/prestosql/plugin/oracle/OracleConfig.java   | 14 ++++++++++++++
 .../prestosql/plugin/oracle/TestOracleConfig.java  | 12 +++++++++++-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleClient.java b/presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleClient.java
index 14f3be659711..39bbeb41a454 100644
--- a/presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleClient.java
+++ b/presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleClient.java
@@ -123,6 +123,7 @@ public class OracleClient
     private static final int PRECISION_OF_UNSPECIFIED_NUMBER = 127;
 
     private final boolean synonymsEnabled;
+    private final int fetchSize;
 
     private static final Map<Type, WriteMapping> WRITE_MAPPINGS = ImmutableMap.<Type, WriteMapping>builder()
             .put(BOOLEAN, oracleBooleanWriteMapping())
@@ -147,6 +148,7 @@ public OracleClient(
 
         requireNonNull(oracleConfig, "oracle config is null");
         this.synonymsEnabled = oracleConfig.isSynonymsEnabled();
+        this.fetchSize = oracleConfig.getFetchSize();
     }
 
     private String[] getTableTypes()
@@ -175,7 +177,7 @@ public PreparedStatement getPreparedStatement(Connection connection, String sql)
             throws SQLException
     {
         PreparedStatement statement = connection.prepareStatement(sql);
-        statement.setFetchSize(1000);
+        statement.setFetchSize(fetchSize);
         return statement;
     }
 
diff --git a/presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleConfig.java b/presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleConfig.java
index 40710e9157f3..3f89078bce04 100644
--- a/presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleConfig.java
+++ b/presto-oracle/src/main/java/io/prestosql/plugin/oracle/OracleConfig.java
@@ -27,6 +27,7 @@ public class OracleConfig
 {
     private boolean synonymsEnabled;
     private Integer defaultNumberScale;
+    private int fetchSize = 1000;
     private RoundingMode numberRoundingMode = RoundingMode.UNNECESSARY;
 
     @NotNull
@@ -67,4 +68,17 @@ public OracleConfig setNumberRoundingMode(RoundingMode numberRoundingMode)
         this.numberRoundingMode = numberRoundingMode;
         return this;
     }
+
+    @Min(1)
+    public int getFetchSize()
+    {
+        return fetchSize;
+    }
+
+    @Config("oracle.fetch-size")
+    public OracleConfig setFetchSize(int fetchSize)
+    {
+        this.fetchSize = fetchSize;
+        return this;
+    }
 }
diff --git a/presto-oracle/src/test/java/io/prestosql/plugin/oracle/TestOracleConfig.java b/presto-oracle/src/test/java/io/prestosql/plugin/oracle/TestOracleConfig.java
index 0dad3433f374..4548b5f909c5 100644
--- a/presto-oracle/src/test/java/io/prestosql/plugin/oracle/TestOracleConfig.java
+++ b/presto-oracle/src/test/java/io/prestosql/plugin/oracle/TestOracleConfig.java
@@ -35,7 +35,8 @@ public void testDefaults()
         assertRecordedDefaults(recordDefaults(OracleConfig.class)
                 .setSynonymsEnabled(false)
                 .setDefaultNumberScale(null)
-                .setNumberRoundingMode(RoundingMode.UNNECESSARY));
+                .setNumberRoundingMode(RoundingMode.UNNECESSARY)
+                .setFetchSize(1000));
     }
 
     @Test
@@ -45,11 +46,13 @@ public void testExplicitPropertyMappings()
                 .put("oracle.synonyms.enabled", "true")
                 .put("oracle.number.default-scale", "2")
                 .put("oracle.number.rounding-mode", "CEILING")
+                .put("oracle.fetch-size", "5000")
                 .build();
 
         OracleConfig expected = new OracleConfig()
                 .setSynonymsEnabled(true)
                 .setDefaultNumberScale(2)
+                .setFetchSize(5000)
                 .setNumberRoundingMode(RoundingMode.CEILING);
 
         assertFullMapping(properties, expected);
@@ -71,5 +74,12 @@ public void testValidation()
                 "defaultNumberScale",
                 "must be less than or equal to 38",
                 Max.class);
+
+        assertFailsValidation(
+                new OracleConfig()
+                        .setFetchSize(0),
+                "fetchSize",
+                "must be greater than or equal to 1",
+                Min.class);
     }
 }