From c67cd9ac190532e7255b7936812665999b6e334f Mon Sep 17 00:00:00 2001
From: Kevin Minder <kevin.minder@oracle.com>
Date: Tue, 22 May 2018 21:46:23 -0400
Subject: [PATCH] CALCITE-2322: Add fetch size support to connection url and
 JDBC statement

---
 .../java/org/apache/calcite/avatica/AvaticaStatement.java    | 2 ++
 .../apache/calcite/avatica/BuiltInConnectionProperty.java    | 5 ++++-
 .../java/org/apache/calcite/avatica/ConnectionConfig.java    | 2 ++
 .../org/apache/calcite/avatica/ConnectionConfigImpl.java     | 4 ++++
 core/src/main/java/org/apache/calcite/avatica/MetaImpl.java  | 2 +-
 5 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/avatica/AvaticaStatement.java b/core/src/main/java/org/apache/calcite/avatica/AvaticaStatement.java
index 0676c5d9cf..8b6cf58ee2 100644
--- a/core/src/main/java/org/apache/calcite/avatica/AvaticaStatement.java
+++ b/core/src/main/java/org/apache/calcite/avatica/AvaticaStatement.java
@@ -36,6 +36,7 @@
  */
 public abstract class AvaticaStatement
     implements Statement {
+
   /** The default value for {@link Statement#getFetchSize()}. */
   public static final int DEFAULT_FETCH_SIZE = 100;
 
@@ -108,6 +109,7 @@ protected AvaticaStatement(AvaticaConnection connection,
     this.resultSetType = resultSetType;
     this.resultSetConcurrency = resultSetConcurrency;
     this.resultSetHoldability = resultSetHoldability;
+    this.fetchSize = connection.config().fetchSize(); // Default to connection config fetch size.
     this.signature = signature;
     this.closed = false;
     if (h == null) {
diff --git a/core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java b/core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java
index 1da7025486..20f6050611 100644
--- a/core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java
+++ b/core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java
@@ -77,7 +77,10 @@ public enum BuiltInConnectionProperty implements ConnectionProperty {
   TRUSTSTORE_PASSWORD("truststore_password", Type.STRING, null, false),
 
   HOSTNAME_VERIFICATION("hostname_verification", Type.ENUM, HostnameVerification.STRICT,
-      HostnameVerification.class, false);
+      HostnameVerification.class, false),
+
+  /** Fetch size limit, default is 100 rows. */
+  FETCH_SIZE("fetch_size", Type.NUMBER, AvaticaStatement.DEFAULT_FETCH_SIZE, false);
 
   private final String camelName;
   private final Type type;
diff --git a/core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java b/core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java
index cd18ec5d8f..171a29c7aa 100644
--- a/core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java
+++ b/core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java
@@ -56,6 +56,8 @@ public interface ConnectionConfig {
   String truststorePassword();
   /** @see BuiltInConnectionProperty#HOSTNAME_VERIFICATION */
   HostnameVerification hostnameVerification();
+  /** @see BuiltInConnectionProperty#FETCH_SIZE */
+  int fetchSize();
 }
 
 // End ConnectionConfig.java
diff --git a/core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java b/core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java
index 63fbd8d5b4..ae16a1dad7 100644
--- a/core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java
+++ b/core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java
@@ -111,6 +111,10 @@ public HostnameVerification hostnameVerification() {
         .getEnum(HostnameVerification.class);
   }
 
+  public int fetchSize() {
+    return BuiltInConnectionProperty.FETCH_SIZE.wrap(properties).getInt();
+  }
+
   /** Converts a {@link Properties} object containing (name, value)
    * pairs into a map whose keys are
    * {@link org.apache.calcite.avatica.InternalProperty} objects.
diff --git a/core/src/main/java/org/apache/calcite/avatica/MetaImpl.java b/core/src/main/java/org/apache/calcite/avatica/MetaImpl.java
index f663ebb86a..dd9fb19fc3 100644
--- a/core/src/main/java/org/apache/calcite/avatica/MetaImpl.java
+++ b/core/src/main/java/org/apache/calcite/avatica/MetaImpl.java
@@ -1603,7 +1603,7 @@ private void moveNext() {
         }
         try {
           // currentOffset updated after element is read from `rows` iterator
-          frame = fetch(stmt.handle, currentOffset, AvaticaStatement.DEFAULT_FETCH_SIZE);
+          frame = fetch(stmt.handle, currentOffset, stmt.getFetchSize());
         } catch (NoSuchStatementException e) {
           resetStatement();
           // re-fetch the batch where we left off