PLUGIN-1823: Rework

Author: sgarg-CS

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

@@ -16,15 +16,12 @@
 
 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;
 import io.cdap.cdap.etl.api.action.Action;
 import io.cdap.cdap.etl.api.action.ActionContext;
 import io.cdap.plugin.common.db.DBErrorDetailsProvider;
 import io.cdap.plugin.util.DBUtils;
-import io.cdap.plugin.util.RetryPolicyUtil;
 
 import java.sql.Driver;
 import java.sql.SQLException;
@@ -36,6 +33,7 @@ public abstract class AbstractDBAction extends Action {
   private static final String JDBC_PLUGIN_ID = "driver";
   private final QueryConfig config;
   private final Boolean enableAutoCommit;
+
   public AbstractDBAction(QueryConfig config, Boolean enableAutoCommit) {
     this.config = config;
     this.enableAutoCommit = enableAutoCommit;
@@ -46,8 +44,7 @@ public void run(ActionContext context) throws Exception {
     Class<? extends Driver> driverClass = context.loadPluginClass(JDBC_PLUGIN_ID);
     DBRun executeQuery = new DBRun(config, driverClass, enableAutoCommit);
     try {
-      Failsafe.with(RetryPolicyUtil.createConnectionRetryPolicy(config.getInitialRetryDuration(),
-        config.getMaxRetryDuration(), config.getMaxRetryCount())).run(()-> executeQuery.run());
+      executeQuery.run();
     } catch (Exception e) {
       if (e instanceof SQLException) {
         DBErrorDetailsProvider dbe = new DBErrorDetailsProvider();

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

@@ -17,6 +17,7 @@
 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;
 import io.cdap.cdap.etl.api.StageConfigurer;
@@ -43,9 +44,12 @@ public class AbstractDBArgumentSetter extends Action {
 
   private static final String JDBC_PLUGIN_ID = "driver";
   private final ArgumentSetterConfig config;
+  private final RetryPolicy<Object> retryPolicy;
 
   public AbstractDBArgumentSetter(ArgumentSetterConfig config) {
     this.config = config;
+    this.retryPolicy = RetryPolicyUtil.getRetryPolicy(config.getInitialRetryDuration(), config.getMaxRetryDuration(),
+      config.getMaxRetryCount());
   }
 
   @Override
@@ -94,10 +98,11 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer)
    */
   private void processArguments(Class<? extends Driver> driverClass,
                                 FailureCollector failureCollector, SettableArguments settableArguments)
-          throws SQLException, IllegalAccessException, InstantiationException {
+    throws SQLException, IllegalAccessException, InstantiationException {
     DriverCleanup driverCleanup;
+
     driverCleanup = DBUtils.ensureJDBCDriverIsAvailable(driverClass, config.getConnectionString(),
-            config.getJdbcPluginName());
+                                                        config.getJdbcPluginName());
     Properties connectionProperties = new Properties();
     connectionProperties.putAll(config.getConnectionArguments());
     try {
@@ -108,46 +113,51 @@ private void processArguments(Class<? extends Driver> driverClass,
   }
 
   private void executeWithRetry(FailureCollector failureCollector, SettableArguments settableArguments,
-                                Properties connectionProperties) {
-    Failsafe.with(RetryPolicyUtil.createConnectionRetryPolicy(config.getInitialRetryDuration(),
-      config.getMaxRetryDuration(), config.getMaxRetryCount())).run(() -> {
-      try (Connection connection = DriverManager
-        .getConnection(config.getConnectionString(), connectionProperties)) {
-        ResultSet resultSet;
-        try (Statement statement = connection.createStatement()) {
-          resultSet = statement.executeQuery(config.getQuery());
-        }
-        boolean hasRecord = resultSet.next();
-        if (!hasRecord) {
-          failureCollector.addFailure("No record found.",
-            "The argument selection conditions must match only one record.");
-          return;
-        }
-        if (settableArguments != null) {
-          setArguments(resultSet, settableArguments);
-        }
-        if (resultSet.next()) {
-          failureCollector
-            .addFailure("More than one records found.",
-              "The argument selection conditions must match only one record.");
-        }
+                                Properties connectionProperties) throws SQLException {
+    try (Connection connection = createConnectionWithRetry(connectionProperties)) {
+      ResultSet resultSet;
+      try (Statement statement = createStatementWithRetry(connection)) {
+        resultSet = Failsafe.with(retryPolicy).get(() -> statement.executeQuery(config.getQuery()));
+      }
+      boolean hasRecord = resultSet.next();
+      if (!hasRecord) {
+        failureCollector.addFailure("No record found.",
+                                    "The argument selection conditions must match only one record.");
+        return;
+      }
+      if (settableArguments != null) {
+        setArguments(resultSet, failureCollector, settableArguments);
       }
-    });
+      if (resultSet.next()) {
+        failureCollector
+          .addFailure("More than one records found.",
+                      "The argument selection conditions must match only one record.");
+      }
+    }
+  }
+
+  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
    *
    * @param resultSet        - result set from db {@link ResultSet}
+   * @param failureCollector - context failure collector @{link FailureCollector}
    * @param arguments        - context argument setter {@link SettableArguments}
+   * @throws SQLException - raises {@link SQLException} when configuration is not valid
    */
-  private void setArguments(ResultSet resultSet, SettableArguments arguments) {
-    Failsafe.with(RetryPolicyUtil.createConnectionRetryPolicy(config.getInitialRetryDuration(),
-      config.getMaxRetryDuration(), config.getMaxRetryCount())).run(() -> {
-      String[] columns = config.getArgumentsColumns().split(",");
-      for (String column : columns) {
-        arguments.set(column, resultSet.getString(column));
-      }
-    });
+  private void setArguments(ResultSet resultSet, FailureCollector failureCollector,
+                            SettableArguments arguments) throws SQLException {
+    String[] columns = config.getArgumentsColumns().split(",");
+    for (String column : columns) {
+      arguments.set(column, resultSet.getString(column));
+    }
   }
 }

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

@@ -16,13 +16,11 @@
 
 package io.cdap.plugin.db.action;
 
-import dev.failsafe.Failsafe;
 import io.cdap.cdap.etl.api.FailureCollector;
 import io.cdap.cdap.etl.api.PipelineConfigurer;
 import io.cdap.cdap.etl.api.batch.BatchActionContext;
 import io.cdap.cdap.etl.api.batch.PostAction;
 import io.cdap.plugin.util.DBUtils;
-import io.cdap.plugin.util.RetryPolicyUtil;
 
 import java.sql.Driver;
 
@@ -53,8 +51,7 @@ public void run(BatchActionContext batchContext) throws Exception {
 
     Class<? extends Driver> driverClass = batchContext.loadPluginClass(JDBC_PLUGIN_ID);
     DBRun executeQuery = new DBRun(config, driverClass, enableAutoCommit);
-    Failsafe.with(RetryPolicyUtil.createConnectionRetryPolicy(config.getInitialRetryDuration(),
-      config.getMaxRetryDuration(), config.getMaxRetryCount())).run(() -> executeQuery.run());
+    executeQuery.run();
   }
 
   @Override

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

@@ -17,6 +17,7 @@
 package io.cdap.plugin.db.action;
 
 import dev.failsafe.Failsafe;
+import dev.failsafe.RetryPolicy;
 import io.cdap.plugin.util.DBUtils;
 import io.cdap.plugin.util.DriverCleanup;
 import io.cdap.plugin.util.RetryPolicyUtil;
@@ -36,60 +37,64 @@ public class DBRun {
   private final QueryConfig config;
   private final Class<? extends Driver> driverClass;
   private boolean enableAutoCommit;
+  private final RetryPolicy<Object> retryPolicy;
 
   public DBRun(QueryConfig config, Class<? extends Driver> driverClass, Boolean enableAutocommit) {
     this.config = config;
     this.driverClass = driverClass;
     if (enableAutocommit != null) {
       this.enableAutoCommit = enableAutocommit;
     }
+    this.retryPolicy = RetryPolicyUtil.getRetryPolicy(config.getInitialRetryDuration(), config.getMaxRetryDuration(),
+      config.getMaxRetryCount());
   }
 
   /**
    * Uses a configured JDBC driver to execute a SQL statement. The configurations of which JDBC driver
    * to use and which connection string to use come from the plugin configuration.
    */
   public void run() throws SQLException, InstantiationException, IllegalAccessException {
-
     DriverCleanup driverCleanup = null;
     try {
       driverCleanup = DBUtils.ensureJDBCDriverIsAvailable(driverClass, config.getConnectionString(),
                                                           config.getJdbcPluginName());
 
       Properties connectionProperties = new Properties();
       connectionProperties.putAll(config.getConnectionArguments());
-      Failsafe.with(RetryPolicyUtil.createConnectionRetryPolicy(config.getInitialRetryDuration(),
-        config.getMaxRetryDuration(), config.getMaxRetryCount())).run(() -> {
-        try (Connection connection = DriverManager.getConnection(config.getConnectionString(), connectionProperties)) {
-          executeInitQueries(connection, config.getInitQueries());
+      try (Connection connection = createConnectionWithRetry(connectionProperties)) {
+        executeInitQueries(connection, config.getInitQueries());
+        if (!enableAutoCommit) {
+          connection.setAutoCommit(false);
+        }
+        try (Statement statement = createStatementWithRetry(connection)) {
+          Failsafe.with(retryPolicy).run(() -> statement.execute(config.query));
+          
           if (!enableAutoCommit) {
-            connection.setAutoCommit(false);
-          }
-          try (Statement statement = connection.createStatement()) {
-            statement.execute(config.query);
-            if (!enableAutoCommit) {
-              connection.commit();
-            }
+            connection.commit();
           }
         }
-      });
+      }
     } finally {
       if (driverCleanup != null) {
         driverCleanup.destroy();
       }
     }
   }
 
-  private void executeInitQueries(Connection connection, List<String> initQueries) {
+  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));
+      }
+    }
+  }
 
-    Failsafe.with(RetryPolicyUtil.createConnectionRetryPolicy(config.getInitialRetryDuration(),
-      config.getMaxRetryDuration(), config.getMaxRetryCount()))
-      .run(() -> {
-        for (String query : initQueries) {
-          try (Statement statement = connection.createStatement()) {
-            statement.execute(query);
-          }
-        }
-      });
+  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

@@ -174,16 +174,15 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
 
   protected Schema loadTableSchema(Connection connection, String query, @Nullable Integer timeoutSec, String sessionID)
     throws SQLException {
-    return Failsafe.with(RetryPolicyUtil.createConnectionRetryPolicy(config.getInitialRetryDuration(),
-      config.getMaxRetryDuration(), config.getMaxRetryCount())).get(() -> {
       Statement statement = connection.createStatement();
       statement.setMaxRows(1);
       if (timeoutSec != null) {
         statement.setQueryTimeout(timeoutSec);
       }
-      ResultSet resultSet = statement.executeQuery(query);
+    ResultSet resultSet = Failsafe.with(RetryPolicyUtil.getRetryPolicy(config.getInitialRetryDuration(),
+      config.getMaxRetryDuration(), config.getMaxRetryCount()))
+      .get(() -> statement.executeQuery(query));
       return Schema.recordOf("outputSchema", getSchemaReader(sessionID).getSchemaFields(resultSet));
-    });
   }
 
   protected void setConnectionProperties(Map<String, String> properties, ConnectorSpecRequest request) {