Merge pull request #191 from data-integrations/PLUGIN-994

[PLUGIN-994] Adding setFetchSize to the database plugins

Author: itsankit-google

aurora-mysql-plugin/widgets/AuroraMysql-batchsource.json

@@ -73,6 +73,15 @@
           "widget-attributes": {
             "default": "1"
           }
+        },
+        {
+          "widget-type": "number",
+          "label": "Fetch Size",
+          "name": "fetchSize",
+          "widget-attributes": {
+            "default": "1000",
+            "minimum": "0"
+          }
         }
       ]
     },

aurora-postgresql-plugin/widgets/AuroraPostgres-batchsource.json

@@ -73,6 +73,15 @@
           "widget-attributes": {
             "default": "1"
           }
+        },
+        {
+          "widget-type": "number",
+          "label": "Fetch Size",
+          "name": "fetchSize",
+          "widget-attributes": {
+            "default": "1000",
+            "minimum": "0"
+          }
         }
       ]
     },

cloudsql-mysql-plugin/src/main/java/io/cdap/plugin/cloudsql/mysql/CloudSQLMySQLSource.java

@@ -33,6 +33,7 @@
 import io.cdap.plugin.db.batch.source.AbstractDBSource;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -101,7 +102,14 @@ public static class CloudSQLMySQLSourceConfig extends AbstractDBSpecificSourceCo
 
     @Override
     protected Map<String, String> getDBSpecificArguments() {
-      return Collections.emptyMap();
+      if (getFetchSize() == null || getFetchSize() <= 0) {
+        return Collections.emptyMap();
+      }
+      Map<String, String> arguments = new HashMap<>();
+      // If connected to MySQL > 5.0.2, and setFetchSize() > 0 on a statement,
+      // statement will use cursor-based fetching to retrieve rows
+      arguments.put("useCursorFetch", "true");
+      return arguments;
     }
 
     @Override

cloudsql-mysql-plugin/widgets/CloudSQLMySQL-batchsource.json

@@ -147,6 +147,15 @@
           "widget-attributes": {
             "default": "1"
           }
+        },
+        {
+          "widget-type": "number",
+          "label": "Fetch Size",
+          "name": "fetchSize",
+          "widget-attributes": {
+            "default": "1000",
+            "minimum": "0"
+          }
         }
       ]
     }

cloudsql-postgresql-plugin/widgets/CloudSQLPostgreSQL-batchsource.json

@@ -147,6 +147,15 @@
           "widget-attributes": {
             "default": "1"
           }
+        },
+        {
+          "widget-type": "number",
+          "label": "Fetch Size",
+          "name": "fetchSize",
+          "widget-attributes": {
+            "default": "1000",
+            "minimum": "0"
+          }
         }
       ]
     }

database-commons/src/main/java/io/cdap/plugin/db/ConnectionConfigAccessor.java

@@ -36,6 +36,7 @@ public class ConnectionConfigAccessor {
   private static final String CONNECTION_ARGUMENTS = "io.cdap.plugin.db.connection.arguments";
   private static final String INIT_QUERIES = "io.cdap.plugin.db.init.queries";
   public static final String AUTO_COMMIT_ENABLED = "io.cdap.plugin.db.output.autocommit.enabled";
+  public static final String FETCH_SIZE = "io.cdap.plugin.db.fetch.size";
 
   private static final Gson GSON = new Gson();
   private static final Type STRING_MAP_TYPE = new TypeToken<Map<String, String>>() { }.getType();
@@ -99,6 +100,14 @@ public boolean isAutoCommitEnabled() {
     return configuration.getBoolean(AUTO_COMMIT_ENABLED, false);
   }
 
+  public void setFetchSize(Integer fetchSize) {
+    configuration.setInt(FETCH_SIZE, fetchSize);
+  }
+
+  public Integer getFetchSize() {
+    return configuration.getInt(FETCH_SIZE, 0);
+  }
+
   public Configuration getConfiguration() {
     return configuration;
   }