batch insert

Author: todd5167

core/src/main/java/com/dtstack/flink/sql/util/JDBCUtils.java

@@ -0,0 +1,34 @@
+package com.dtstack.flink.sql.util;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.sql.DriverManager;
+
+public class JDBCUtils {
+
+    private static final Logger LOG = LoggerFactory.getLogger(ClassUtil.class);
+
+    public final static String lock_str = "jdbc_lock_str";
+
+    public static void forName(String clazz, ClassLoader classLoader)  {
+        synchronized (lock_str){
+            try {
+                Class.forName(clazz, true, classLoader);
+                DriverManager.setLoginTimeout(10);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+
+    public synchronized static void forName(String clazz) {
+        try {
+            Class<?> driverClass = Class.forName(clazz);
+            driverClass.newInstance();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

rdb/rdb-sink/src/main/java/com/dtstack/flink/sql/sink/rdb/format/RetractJDBCOutputFormat.java

@@ -19,6 +19,7 @@
 package com.dtstack.flink.sql.sink.rdb.format;
 
 import com.dtstack.flink.sql.sink.rdb.RdbSink;
+import com.dtstack.flink.sql.util.JDBCUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.flink.api.java.tuple.Tuple;
 import org.apache.flink.api.java.tuple.Tuple2;
@@ -31,6 +32,8 @@
 import java.sql.*;
 import java.util.*;
 import java.io.IOException;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 
 import com.dtstack.flink.sql.sink.MetricOutputFormat;
 
@@ -53,20 +56,27 @@ public class RetractJDBCOutputFormat extends MetricOutputFormat {
     private String tableName;
     private String dbType;
     private RdbSink dbSink;
+    // trigger preparedStatement execute batch interval
+    private long batchWaitInterval = 10000l;
+
+    // batchNum
     private int batchInterval = 5000;
     private String insertQuery;
     public int[] typesArray;
 
     private Connection dbConn;
     private PreparedStatement upload;
 
-    private int batchCount = 0;
+    /** 存储用于批量写入的数据 */
+    protected List<Row> rows = new ArrayList();
 
     //index field
     private Map<String, List<String>> realIndexes = Maps.newHashMap();
     //full field
     private List<String> fullField = Lists.newArrayList();
 
+    private transient ScheduledThreadPoolExecutor timerService;
+
     public RetractJDBCOutputFormat() {
 
     }
@@ -85,6 +95,7 @@ public void configure(Configuration parameters) {
     @Override
     public void open(int taskNumber, int numTasks) throws IOException {
         try {
+            LOG.info("PreparedStatement execute batch num is {}", batchInterval);
             dbConn = establishConnection();
             initMetric();
             if (dbConn.getMetaData().getTables(null, null, tableName, null).next()) {
@@ -96,6 +107,17 @@ public void open(int taskNumber, int numTasks) throws IOException {
                 throw new SQLException("Table " + tableName + " doesn't exist");
             }
 
+            if (batchWaitInterval > 0 && batchInterval > 1) {
+                LOG.info("open batch wait interval scheduled, interval is {} ms", batchWaitInterval);
+
+                timerService = new ScheduledThreadPoolExecutor(1);
+                timerService.scheduleAtFixedRate(() -> {
+                    submitExecuteBatch();
+                }, 0, batchWaitInterval, TimeUnit.MILLISECONDS);
+
+            }
+
+
         } catch (SQLException sqe) {
             throw new IllegalArgumentException("open() failed.", sqe);
         } catch (ClassNotFoundException cnfe) {
@@ -106,12 +128,13 @@ public void open(int taskNumber, int numTasks) throws IOException {
 
     private Connection establishConnection() throws SQLException, ClassNotFoundException {
         Connection connection ;
-        Class.forName(drivername);
+        JDBCUtils.forName(drivername, getClass().getClassLoader());
         if (username == null) {
             connection = DriverManager.getConnection(dbURL);
         } else {
             connection = DriverManager.getConnection(dbURL, username, password);
         }
+        connection.setAutoCommit(false);
         return connection;
     }
 
@@ -129,7 +152,7 @@ private Connection establishConnection() throws SQLException, ClassNotFoundExcep
      * @see PreparedStatement
      */
     @Override
-    public void writeRecord(Tuple2 tuple2) throws IOException {
+    public void writeRecord(Tuple2 tuple2) {
 
         Tuple2<Boolean, Row> tupleTrans = tuple2;
         Boolean retract = tupleTrans.getField(0);
@@ -139,27 +162,64 @@ public void writeRecord(Tuple2 tuple2) throws IOException {
         if (typesArray != null && typesArray.length > 0 && typesArray.length != row.getArity()) {
             LOG.warn("Column SQL types array doesn't match arity of passed Row! Check the passed array...");
         }
+
+        if (retract) {
+            insertWrite(row);
+            outRecords.inc();
+        } else {
+            //do nothing
+        }
+
+    }
+
+
+    private void insertWrite(Row row) {
+        System.out.println("接受到数据row:" +row );
+        checkConnectionOpen(dbConn);
         try {
-            if (retract) {
-                insertWrite(row);
-                outRecords.inc();
+            if (batchInterval == 1) {
+                writeSingleRecord(row);
             } else {
-                //do nothing
+                updatePreparedStmt(row, upload);
+                rows.add(row);
+                upload.addBatch();
+                if (rows.size() >= batchInterval) {
+                    submitExecuteBatch();
+                }
             }
-        } catch (SQLException | IllegalArgumentException e) {
-            throw new IllegalArgumentException("writeRecord() failed", e);
+        } catch (SQLException e) {
+            LOG.error("", e);
+        }
+    }
+
+    private void writeSingleRecord(Row row) {
+        try {
+            updatePreparedStmt(row, upload);
+            upload.execute();
+            System.out.println("单条插入成功：" + row);
+        } catch (SQLException e) {
+            System.out.println("单条插入失败：" + row);
+            LOG.error("record insert failed ..", row.toString());
+            LOG.error("", e);
         }
     }
 
+    private synchronized void submitExecuteBatch() {
+        try {
+            LOG.info("submitExecuteBatch start......");
+            this.upload.executeBatch();
+            dbConn.commit();
+            rows.forEach(row -> System.out.println("批量插入成功："+ row));
+        } catch (SQLException e) {
+            try {
+                dbConn.rollback();
+            } catch (SQLException e1) {
+                LOG.error("rollback  data error !", e);
+            }
 
-    private void insertWrite(Row row) throws SQLException {
-        checkConnectionOpen(dbConn);
-        updatePreparedStmt(row, upload);
-        upload.addBatch();
-        batchCount++;
-        if (batchCount >= batchInterval) {
-            upload.executeBatch();
-            batchCount = 0;
+            rows.forEach(this::writeSingleRecord);
+        } finally {
+            rows.clear();
         }
     }
 
@@ -282,7 +342,7 @@ public void close() throws IOException {
             LOG.info("Inputformat couldn't be closed - " + se.getMessage());
         } finally {
             upload = null;
-            batchCount = 0;
+
         }
 
         try {