Added changes for New committer.

Author: vikasrathee-cs

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

@@ -203,6 +203,7 @@ public void prepareRun(BatchSinkContext context) {
     configAccessor.setInitQueries(dbSinkConfig.getInitQueries());
     configAccessor.getConfiguration().set(DBConfiguration.DRIVER_CLASS_PROPERTY, driverClass.getName());
     configAccessor.getConfiguration().set(DBConfiguration.URL_PROPERTY, connectionString);
+    configAccessor.getConfiguration().set(ETLDBOutputFormat.STAGE_NAME, context.getStageName());
     String fullyQualifiedTableName = dbSchemaName == null ? dbSinkConfig.getEscapedTableName()
             : dbSinkConfig.getEscapedDbSchemaName() + "." + dbSinkConfig.getEscapedTableName();
     configAccessor.getConfiguration().set(DBConfiguration.OUTPUT_TABLE_NAME_PROPERTY, fullyQualifiedTableName);

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

@@ -25,6 +25,8 @@
 import io.cdap.plugin.db.TransactionIsolationLevel;
 import io.cdap.plugin.util.DBUtils;
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.mapreduce.JobContext;
+import org.apache.hadoop.mapreduce.OutputCommitter;
 import org.apache.hadoop.mapreduce.RecordWriter;
 import org.apache.hadoop.mapreduce.TaskAttemptContext;
 import org.apache.hadoop.mapreduce.lib.db.DBConfiguration;
@@ -41,6 +43,7 @@
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
 
@@ -56,15 +59,98 @@
 public class ETLDBOutputFormat<K extends DBWritable, V> extends DBOutputFormat<K, V> {
   // Batch size before submitting a batch to the SQL engine. If set to 0, no batches will be submitted until commit.
   public static final String COMMIT_BATCH_SIZE = "io.cdap.plugin.db.output.commit.batch.size";
+  public static final String STAGE_NAME = "io.cdap.plugin.db.output.stage_name";
   public static final int DEFAULT_COMMIT_BATCH_SIZE = 1000;
   private static final Character ESCAPE_CHAR = '"';
-
+  private static final String CONNECTION_MAP_ID_REGEX = "%s_%s";
   private static final Logger LOG = LoggerFactory.getLogger(ETLDBOutputFormat.class);
 
   private Configuration conf;
   private Driver driver;
   private JDBCDriverShim driverShim;
 
+  private static Map<String, Connection> connectionMap = new HashMap<>();
+
+  @Override
+  public OutputCommitter getOutputCommitter(TaskAttemptContext context)
+    throws IOException, InterruptedException {
+    return new OutputCommitter() {
+      @Override
+      public void setupJob(JobContext jobContext) throws IOException {
+        // do nothing
+      }
+
+      @Override
+      public void setupTask(TaskAttemptContext taskContext) throws IOException {
+        // do nothing
+      }
+
+      @Override
+      public boolean needsTaskCommit(TaskAttemptContext taskContext) throws IOException {
+        return true;
+      }
+
+      @Override
+      public void commitTask(TaskAttemptContext taskContext) throws IOException {
+        conf = context.getConfiguration();
+        String stageName = conf.get(STAGE_NAME);
+        String connectionId = String.format(CONNECTION_MAP_ID_REGEX, context.getTaskAttemptID().toString(), stageName);
+        Connection connection;
+        if ((connection = connectionMap.remove(connectionId)) != null) {
+          try {
+            connection.commit();
+          } catch (SQLException e) {
+            try {
+              connection.rollback();
+            } catch (SQLException ex) {
+              LOG.warn(StringUtils.stringifyException(ex));
+            }
+            throw new IOException(e);
+          } finally {
+            try {
+              connection.close();
+              LOG.debug("Connection Closed after committing the task with taskAttemptId {}", connectionId);
+            } catch (SQLException ex) {
+              LOG.warn(StringUtils.stringifyException(ex));
+            }
+          }
+          try {
+            DriverManager.deregisterDriver(driverShim);
+          } catch (SQLException e) {
+            throw new IOException(e);
+          }
+        }
+      }
+
+      @Override
+      public void abortTask(TaskAttemptContext taskContext) throws IOException {
+        conf = context.getConfiguration();
+        String stageName = conf.get(STAGE_NAME);
+        String connectionId = String.format(CONNECTION_MAP_ID_REGEX, context.getTaskAttemptID().toString(), stageName);
+        Connection connection;
+        if ((connection = connectionMap.remove(connectionId)) != null) {
+          try {
+            connection.rollback();
+          } catch (SQLException e) {
+            throw new IOException(e);
+          } finally {
+            try {
+              connection.close();
+              LOG.debug("Connection Closed after rollback the task with taskAttemptId {}", connectionId);
+            } catch (SQLException ex) {
+              LOG.warn(StringUtils.stringifyException(ex));
+            }
+          }
+          try {
+            DriverManager.deregisterDriver(driverShim);
+          } catch (SQLException e) {
+            throw new IOException(e);
+          }
+        }
+      }
+    };
+  }
+
   @Override
   public RecordWriter<K, V> getRecordWriter(TaskAttemptContext context) throws IOException {
     conf = context.getConfiguration();
@@ -81,6 +167,11 @@ public RecordWriter<K, V> getRecordWriter(TaskAttemptContext context) throws IOE
 
     try {
       Connection connection = getConnection(conf);
+      String stageName = conf.get(STAGE_NAME);
+      // If using multiple sinks, task attemptID can be same in that case, appending stage in the end for uniqueness.
+      String connectionId = String.format(CONNECTION_MAP_ID_REGEX, context.getTaskAttemptID().toString(), stageName);
+      connectionMap.put(connectionId, connection);
+      LOG.debug("Connection Added to the map with connectionId : {}", connectionId);
       PreparedStatement statement = connection.prepareStatement(constructQueryOnOperation(tableName, fieldNames,
         operationName, listKeys));
       return new DBRecordWriter(connection, statement) {
@@ -98,28 +189,15 @@ public void close(TaskAttemptContext context) throws IOException {
             if (!emptyData) {
               getStatement().executeBatch();
             }
-            getConnection().commit();
           } catch (SQLException e) {
-            try {
-              getConnection().rollback();
-            } catch (SQLException ex) {
-              LOG.warn(StringUtils.stringifyException(ex));
-            }
             throw new IOException(e);
           } finally {
             try {
               getStatement().close();
-              getConnection().close();
             } catch (SQLException ex) {
               throw new IOException(ex);
             }
           }
-
-          try {
-            DriverManager.deregisterDriver(driverShim);
-          } catch (SQLException e) {
-            throw new IOException(e);
-          }
         }
 
         @Override