feat: 解决update模式下查找主键失败的问题

Author: lvyanquan

flinkx-kingbase/flinkx-kingbase-core/src/main/java/com/dtstack/flinkx/kingbase/constants/KingbaseCons.java

@@ -28,7 +28,13 @@
 public class KingbaseCons {
 
     public static final String DRIVER = "com.kingbase8.Driver";
+    /**
+     * kingbase 主键索引名后缀
+     */
     public static final String KEY_UPDATE_KEY = "key";
+
+    public static final String KEY_PRIMARY_SUFFIX = "_PKEY";
+
     public static final String INSERT_SQL_MODE_TYPE = "copy";
 
     /**

flinkx-kingbase/flinkx-kingbase-core/src/main/java/com/dtstack/flinkx/kingbase/util/KingBaseDatabaseMeta.java

@@ -23,13 +23,15 @@
 import org.apache.commons.lang3.StringUtils;
 
 import java.util.ArrayList;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 
 import static com.dtstack.flinkx.constants.ConstantValue.COMMA_SYMBOL;
 import static com.dtstack.flinkx.constants.ConstantValue.LEFT_PARENTHESIS_SYMBOL;
 import static com.dtstack.flinkx.constants.ConstantValue.RIGHT_PARENTHESIS_SYMBOL;
 import static com.dtstack.flinkx.kingbase.constants.KingbaseCons.DRIVER;
+import static com.dtstack.flinkx.kingbase.constants.KingbaseCons.KEY_PRIMARY_SUFFIX;
 import static com.dtstack.flinkx.kingbase.constants.KingbaseCons.KEY_UPDATE_KEY;
 
 /**
@@ -75,12 +77,26 @@ public String getSqlQueryColumnFields(List<String> column, String table) {
         return "SELECT " + quoteColumns(column) + " FROM " + quoteTable(table) + " LIMIT 0";
     }
 
+    /**
+     * Kingbase 的主键索引名为TABLE_PKEY格式
+     * @param column column名
+     * @param table 表名
+     * @param updateKey 索引
+     * @return updateSql
+     */
     @Override
     public String getUpsertStatement(List<String> column, String table, Map<String,List<String>> updateKey) {
+        List<String> columnList = new LinkedList<>();
+        updateKey.forEach((key, value) -> {
+            // 兼顾查询主键索引名或者填入key map的情况
+            if (StringUtils.endsWith(key, KEY_PRIMARY_SUFFIX) || StringUtils.equals(key, KEY_UPDATE_KEY)) {
+                columnList.addAll(value);
+            }
+        });
         return "INSERT INTO " + quoteTable(table)
                 + " (" + quoteColumns(column) + ") VALUES "
                 + makeValues(column.size())
-                + " ON CONFLICT " +makeValues(updateKey.get(KEY_UPDATE_KEY)) + " DO UPDATE SET "
+                + " ON CONFLICT " +makeValues(columnList) + " DO UPDATE SET "
                 + makeUpdatePart(column);
     }
 

flinkx-kingbase/flinkx-kingbase-writer/src/main/java/com/dtstack/flinkx/kingbase/format/KingbaseOutputFormat.java

@@ -18,11 +18,13 @@
 
 package com.dtstack.flinkx.kingbase.format;
 
+import com.dtstack.flinkx.constants.ConstantValue;
 import com.dtstack.flinkx.enums.EWriteMode;
 import com.dtstack.flinkx.exception.WriteRecordException;
 import com.dtstack.flinkx.rdb.outputformat.JdbcOutputFormat;
+import com.dtstack.flinkx.rdb.util.DbUtil;
+import com.dtstack.flinkx.util.ClassUtil;
 import com.dtstack.flinkx.util.ExceptionUtil;
-import com.dtstack.flinkx.util.StringUtil;
 import com.kingbase8.copy.CopyManager;
 import com.kingbase8.core.BaseConnection;
 import org.apache.commons.collections.CollectionUtils;
@@ -54,6 +56,61 @@ public class KingbaseOutputFormat extends JdbcOutputFormat {
 
     private CopyManager copyManager;
 
+    /**
+     * schema名
+     */
+    public String schema;
+
+    @Override
+    protected void openInternal(int taskNumber, int numTasks){
+        try {
+            ClassUtil.forName(driverName, getClass().getClassLoader());
+            dbConn = DbUtil.getConnection(dbUrl, username, password);
+
+            if (restoreConfig.isRestore()){
+                dbConn.setAutoCommit(false);
+            }
+            // 查询主键时用table格式
+            if(CollectionUtils.isEmpty(fullColumn)) {
+                fullColumn = probeFullColumns(table, dbConn);
+            }
+
+            if (!EWriteMode.INSERT.name().equalsIgnoreCase(mode)){
+                if(updateKey == null || updateKey.size() == 0) {
+                    updateKey = probePrimaryKeys(table, dbConn);
+                }
+            }
+            // 其他情况，使用schema.table作为表名
+            table = schema + ConstantValue.POINT_SYMBOL + table;
+            if(fullColumnType == null) {
+                fullColumnType = analyzeTable();
+            }
+
+            for(String col : column) {
+                for (int i = 0; i < fullColumn.size(); i++) {
+                    if (col.equalsIgnoreCase(fullColumn.get(i))){
+                        columnType.add(fullColumnType.get(i));
+                        break;
+                    }
+                }
+            }
+
+            preparedStatement = prepareTemplates();
+            readyCheckpoint = false;
+
+            LOG.info("subTask[{}}] wait finished", taskNumber);
+        } catch (SQLException sqe) {
+            throw new IllegalArgumentException("open() failed.", sqe);
+        }
+        try {
+            if (batchInterval > 1) {
+                dbConn.setAutoCommit(false);
+            }
+        } catch (Exception e) {
+            LOG.warn(ExceptionUtil.getErrorMessage(e));
+        }
+    }
+
 
     @Override
     protected PreparedStatement prepareTemplates() throws SQLException {
@@ -71,18 +128,6 @@ protected PreparedStatement prepareTemplates() throws SQLException {
         return super.prepareTemplates();
     }
 
-    @Override
-    protected void openInternal(int taskNumber, int numTasks){
-        super.openInternal(taskNumber, numTasks);
-        try {
-            if (batchInterval > 1) {
-                dbConn.setAutoCommit(false);
-            }
-        } catch (Exception e) {
-            LOG.warn(ExceptionUtil.getErrorMessage(e));
-        }
-    }
-
     @Override
     protected void writeSingleRecordInternal(Row row) throws WriteRecordException {
         if(!checkIsCopyMode(insertSqlMode)){
@@ -195,4 +240,8 @@ private boolean checkIsCopyMode(String insertMode){
         return true;
     }
 
+    public void setSchema(String schema){
+        this.schema = schema;
+    }
+
 }

flinkx-kingbase/flinkx-kingbase-writer/src/main/java/com/dtstack/flinkx/kingbase/writer/KingbaseWriter.java

@@ -41,17 +41,19 @@
 
 public class KingbaseWriter extends JdbcDataWriter {
 
+    public String schema;
+
     public KingbaseWriter(DataTransferConfig config) {
         super(config);
-        String schema = config.getJob().getContent().get(0).getWriter().getParameter().getConnection().get(0).getSchema();
-        table = schema + ConstantValue.POINT_SYMBOL + table;
+        schema = config.getJob().getContent().get(0).getWriter().getParameter().getConnection().get(0).getSchema();
         setDatabaseInterface(new KingBaseDatabaseMeta());
         setTypeConverterInterface(new KingBaseTypeConverter());
     }
 
     @Override
     public DataStreamSink<?> writeData(DataStream<Row> dataSet) {
         KingbaseOutputFormat kingBaseOutputFormat = new KingbaseOutputFormat();
+        kingBaseOutputFormat.setSchema(schema);
         JdbcOutputFormatBuilder builder = new JdbcOutputFormatBuilder(kingBaseOutputFormat);
         builder.setDriverName(databaseInterface.getDriverClass());
         builder.setDbUrl(dbUrl);