PLUGIN-1823: Rework

Move retry logic into a separate class: RetryUtils and add exception handling

Author: sgarg-CS

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

@@ -109,7 +109,6 @@ public Integer getMaxRetryCount() {
     return maxRetryCount == null ? DEFAULT_MAX_RETRY_COUNT : maxRetryCount;
   }
 
-
   public ConnectionConfig() {
   }
 

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

@@ -16,7 +16,6 @@
 
 package io.cdap.plugin.db.action;
 
-import dev.failsafe.Failsafe;
 import dev.failsafe.RetryPolicy;
 import io.cdap.cdap.etl.api.FailureCollector;
 import io.cdap.cdap.etl.api.PipelineConfigurer;
@@ -28,10 +27,10 @@
 import io.cdap.plugin.util.DBUtils;
 import io.cdap.plugin.util.DriverCleanup;
 import io.cdap.plugin.util.RetryPolicyUtil;
+import io.cdap.plugin.util.RetryUtils;
 
 import java.sql.Connection;
 import java.sql.Driver;
-import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -44,7 +43,7 @@ public class AbstractDBArgumentSetter extends Action {
 
   private static final String JDBC_PLUGIN_ID = "driver";
   private final ArgumentSetterConfig config;
-  private final RetryPolicy<Object> retryPolicy;
+  private final RetryPolicy<?> retryPolicy;
 
   public AbstractDBArgumentSetter(ArgumentSetterConfig config) {
     this.config = config;
@@ -114,10 +113,13 @@ private void processArguments(Class<? extends Driver> driverClass,
 
   private void executeWithRetry(FailureCollector failureCollector, SettableArguments settableArguments,
                                 Properties connectionProperties) throws SQLException {
-    try (Connection connection = createConnectionWithRetry(connectionProperties)) {
+    try (Connection connection = RetryUtils.createConnectionWithRetry((RetryPolicy<Connection>) retryPolicy,
+        config.getConnectionString(), connectionProperties, null)) {
       ResultSet resultSet;
-      try (Statement statement = createStatementWithRetry(connection)) {
-        resultSet = Failsafe.with(retryPolicy).get(() -> statement.executeQuery(config.getQuery()));
+      try (Statement statement = RetryUtils.createStatementWithRetry((RetryPolicy<Statement>) retryPolicy, connection,
+        null)) {
+        resultSet = RetryUtils.executeQueryWithRetry((RetryPolicy<ResultSet>) retryPolicy, statement,
+          config.getQuery(), null);
       }
       boolean hasRecord = resultSet.next();
       if (!hasRecord) {
@@ -136,15 +138,6 @@ private void executeWithRetry(FailureCollector failureCollector, SettableArgumen
     }
   }
 
-  private Connection createConnectionWithRetry(Properties connectionProperties) {
-    return Failsafe.with(retryPolicy).<Connection>get(() -> DriverManager
-      .getConnection(config.getConnectionString(), connectionProperties));
-  }
-
-  private Statement createStatementWithRetry(Connection connection) {
-    return Failsafe.with(retryPolicy).<Statement>get(() -> connection.createStatement());
-  }
-
   /**
    * Converts column from jdbc results set into pipeline arguments
    *

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

@@ -21,6 +21,7 @@
 import io.cdap.plugin.util.DBUtils;
 import io.cdap.plugin.util.DriverCleanup;
 import io.cdap.plugin.util.RetryPolicyUtil;
+import io.cdap.plugin.util.RetryUtils;
 
 import java.sql.Connection;
 import java.sql.Driver;
@@ -37,7 +38,7 @@ public class DBRun {
   private final QueryConfig config;
   private final Class<? extends Driver> driverClass;
   private boolean enableAutoCommit;
-  private final RetryPolicy<Object> retryPolicy;
+  private final RetryPolicy<?> retryPolicy;
 
   public DBRun(QueryConfig config, Class<? extends Driver> driverClass, Boolean enableAutocommit) {
     this.config = config;
@@ -61,14 +62,15 @@ public void run() throws SQLException, InstantiationException, IllegalAccessExce
 
       Properties connectionProperties = new Properties();
       connectionProperties.putAll(config.getConnectionArguments());
-      try (Connection connection = createConnectionWithRetry(connectionProperties)) {
+      try (Connection connection = RetryUtils.createConnectionWithRetry((RetryPolicy<Connection>) retryPolicy,
+        config.getConnectionString(), connectionProperties, null)) {
         executeInitQueries(connection, config.getInitQueries());
         if (!enableAutoCommit) {
           connection.setAutoCommit(false);
         }
-        try (Statement statement = createStatementWithRetry(connection)) {
-          Failsafe.with(retryPolicy).run(() -> statement.execute(config.query));
-          
+        try (Statement statement = RetryUtils.createStatementWithRetry((RetryPolicy<Statement>) retryPolicy, connection,
+          null)) {
+          RetryUtils.executeInitQueryWithRetry(retryPolicy, statement, config.query, null);
           if (!enableAutoCommit) {
             connection.commit();
           }
@@ -83,18 +85,10 @@ public void run() throws SQLException, InstantiationException, IllegalAccessExce
 
   private void executeInitQueries(Connection connection, List<String> initQueries) throws SQLException {
     for (String query : initQueries) {
-      try (Statement statement = createStatementWithRetry(connection)) {
-        Failsafe.with(retryPolicy).run(() -> statement.execute(query));
+      try (Statement statement = RetryUtils.createStatementWithRetry((RetryPolicy<Statement>) retryPolicy, connection,
+        null)) {
+        RetryUtils.executeInitQueryWithRetry(retryPolicy, statement, query, null);
       }
     }
   }
-
-  private Connection createConnectionWithRetry(Properties connectionProperties) {
-    return Failsafe.with(retryPolicy).<Connection>get(() -> DriverManager
-      .getConnection(config.getConnectionString(), connectionProperties));
-  }
-
-  private Statement createStatementWithRetry(Connection connection) {
-    return Failsafe.with(retryPolicy).<Statement>get(() -> connection.createStatement());
-  }
 }

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

@@ -17,7 +17,7 @@
 package io.cdap.plugin.db.connector;
 
 import com.google.common.collect.Maps;
-import dev.failsafe.Failsafe;
+import dev.failsafe.RetryPolicy;
 import io.cdap.cdap.api.data.batch.InputFormatProvider;
 import io.cdap.cdap.api.data.schema.Schema;
 import io.cdap.cdap.etl.api.batch.BatchConnector;
@@ -35,6 +35,7 @@
 import io.cdap.plugin.db.SchemaReader;
 import io.cdap.plugin.db.source.DataDrivenETLDBInputFormat;
 import io.cdap.plugin.util.RetryPolicyUtil;
+import io.cdap.plugin.util.RetryUtils;
 import org.apache.hadoop.io.LongWritable;
 import org.apache.hadoop.mapreduce.MRJobConfig;
 import org.apache.hadoop.mapreduce.lib.db.DBConfiguration;
@@ -58,10 +59,13 @@ public abstract class AbstractDBSpecificConnector<T extends DBWritable> extends
   implements BatchConnector<LongWritable, T> {
 
   private final AbstractDBConnectorConfig config;
+  private final RetryPolicy<?> retryPolicy;
 
   protected AbstractDBSpecificConnector(AbstractDBConnectorConfig config) {
     super(config);
     this.config = config;
+    this.retryPolicy = RetryPolicyUtil.getRetryPolicy(config.getInitialRetryDuration(), config.getMaxRetryDuration(),
+      config.getMaxRetryCount());
   }
 
   public abstract boolean supportSchema();
@@ -174,14 +178,13 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
 
   protected Schema loadTableSchema(Connection connection, String query, @Nullable Integer timeoutSec, String sessionID)
     throws SQLException {
-      Statement statement = connection.createStatement();
+      Statement statement = RetryUtils.createStatementWithRetry((RetryPolicy<Statement>) retryPolicy, connection, null);
       statement.setMaxRows(1);
       if (timeoutSec != null) {
         statement.setQueryTimeout(timeoutSec);
       }
-    ResultSet resultSet = Failsafe.with(RetryPolicyUtil.getRetryPolicy(config.getInitialRetryDuration(),
-      config.getMaxRetryDuration(), config.getMaxRetryCount()))
-      .get(() -> statement.executeQuery(query));
+      ResultSet resultSet = RetryUtils.executeQueryWithRetry((RetryPolicy<ResultSet>) retryPolicy, statement, query,
+        null);
       return Schema.recordOf("outputSchema", getSchemaReader(sessionID).getSchemaFields(resultSet));
   }