refactor: 优化upsert去重逻辑

Author: zhou-hao

hsweb-easy-orm-rdb/src/main/java/org/hswebframework/ezorm/rdb/operator/builder/fragments/insert/BatchInsertSqlBuilder.java

@@ -1,6 +1,7 @@
 package org.hswebframework.ezorm.rdb.operator.builder.fragments.insert;
 
 import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
 import lombok.AllArgsConstructor;
 import org.apache.commons.lang3.StringUtils;
 import org.hswebframework.ezorm.core.RuntimeDefaultValue;
@@ -32,6 +33,10 @@ protected int computeSqlSize(int columnSize, int valueSize) {
         return (columnSize * valueSize) * 2 + valueSize * 2 + columnSize * 3 + 2;
     }
 
+    protected boolean isPrimaryKey(RDBColumnMetadata col) {
+        return col.isPrimaryKey();
+    }
+
     @Override
     public SqlRequest build(InsertOperatorParameter parameter) {
 //        PrepareSqlFragments fragments = beforeBuild(parameter, PrepareSqlFragments.of()).addSql("(");
@@ -52,7 +57,7 @@ public SqlRequest build(InsertOperatorParameter parameter) {
         LinkedHashMap<Integer, SqlFragments> functionValues = Maps.newLinkedHashMapWithExpectedSize(columns.size());
 
         int index = 0;
-        int primaryIndex = -1;
+        List<Integer> primaryIndex = new ArrayList<>(1);
 
         //如果只有一条数据则忽略null的列
         boolean ignoreNullColumn = parameter.getValues().size() == 1;
@@ -62,8 +67,8 @@ public SqlRequest build(InsertOperatorParameter parameter) {
                 .flatMap(table::getColumn)
                 .orElse(null);
             if (columnMetadata != null && columnMetadata.isInsertable()) {
-                if (columnMetadata.isPrimaryKey()) {
-                    primaryIndex = index;
+                if (isPrimaryKey(columnMetadata)) {
+                    primaryIndex.add(index);
                 }
                 //忽略null的列
                 if (ignoreNullColumn) {
@@ -100,14 +105,31 @@ public SqlRequest build(InsertOperatorParameter parameter) {
         // ) values
         fragments.add(VALUES);
         index = 0;
-        Set<Object> duplicatePrimary = new HashSet<>(32);
+        Set<Object> duplicatePrimary = new HashSet<>(8);
         for (List<Object> values : valueList) {
-            if (primaryIndex >= 0) {
-                //重复的id 则不进行处理
-                if (values.size() > primaryIndex && !duplicatePrimary.add(values.get(primaryIndex))) {
+            int indexSize = primaryIndex.size();
+            int vSize = values.size();
+            // id
+            if (indexSize == 1) {
+                int idx = primaryIndex.get(0);
+                if (vSize > idx && !duplicatePrimary.add(values.get(idx))) {
+                    continue;
+                }
+            }
+            // 唯一索引?
+            else if (indexSize >= 1) {
+                Set<Object> dis = Sets.newHashSetWithExpectedSize(indexSize);
+                for (Integer i : primaryIndex) {
+                    if (vSize > i) {
+                        dis.add(values.get(i));
+                    }
+                }
+                // 存在重复数据 ?
+                if(!duplicatePrimary.add(dis)){
                     continue;
                 }
             }
+
             if (index++ != 0) {
                 fragments.add(SqlFragments.COMMA);
             }

hsweb-easy-orm-rdb/src/main/java/org/hswebframework/ezorm/rdb/supports/postgres/PostgresqlBatchUpsertOperator.java

@@ -18,10 +18,7 @@
 import org.hswebframework.ezorm.rdb.utils.ExceptionUtils;
 import reactor.core.publisher.Mono;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Objects;
-import java.util.Set;
+import java.util.*;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
@@ -36,6 +33,8 @@ public class PostgresqlBatchUpsertOperator implements SaveOrUpdateOperator {
 
     private SaveOrUpdateOperator fallback;
 
+    private Set<String> primaryColumns;
+
     public PostgresqlBatchUpsertOperator(RDBTableMetadata table) {
         this.table = table;
         this.fallback = new DefaultSaveOrUpdateOperator(table);
@@ -58,14 +57,18 @@ SqlFragments getOrCreateOnConflict() {
     }
 
     SqlFragments createOnConflict() {
+        if (primaryColumns == null) {
+            primaryColumns = new HashSet<>();
+        }
         RDBColumnMetadata idColumn = table
             .getColumns()
             .stream()
             .filter(RDBColumnMetadata::isPrimaryKey)
             .findFirst()
             .orElse(null);
         if (idColumn != null) {
-            return SqlFragments.of("on conflict (", idColumn.getName(), ") do ");
+            primaryColumns.add(idColumn.getName());
+            return SqlFragments.of("on conflict (", idColumn.getQuoteName(), ") do ");
         }
         RDBIndexMetadata indexMetadata = table
             .getIndexes()
@@ -80,7 +83,10 @@ SqlFragments createOnConflict() {
                 .stream()
                 .map(c -> table.getColumn(c.getColumn()).orElse(null))
                 .filter(Objects::nonNull)
-                .map(RDBColumnMetadata::getQuoteName)
+                .map(c -> {
+                    primaryColumns.add(c.getName());
+                    return c.getQuoteName();
+                })
                 .collect(Collectors.joining(","));
 
             return SqlFragments.of("on conflict( ", columns, ") do ");
@@ -136,6 +142,15 @@ public PostgresqlUpsertBatchInsertSqlBuilder(RDBTableMetadata table) {
             super(table);
         }
 
+        @Override
+        protected boolean isPrimaryKey(RDBColumnMetadata col) {
+            getOrCreateOnConflict();
+            if (primaryColumns != null && primaryColumns.contains(col.getName())) {
+                return true;
+            }
+            return super.isPrimaryKey(col);
+        }
+
         @Override
         protected int computeSqlSize(int columnSize, int valueSize) {
             return super.computeSqlSize(columnSize, valueSize) + columnSize * 3 + 2;

hsweb-easy-orm-rdb/src/test/java/org/hswebframework/ezorm/rdb/supports/BasicReactiveTests.java

@@ -138,15 +138,15 @@ public void init() {
                 .commit()
                 .reactive()
                 .block();
-        table.<EntityColumnMapping>getFeatureNow(MappingFeatureType.columnPropertyMapping.createFeatureId(BasicTestEntity.class))
+        table
+            .<EntityColumnMapping>getFeatureNow(MappingFeatureType.columnPropertyMapping.createFeatureId(BasicTestEntity.class))
             .reload();
 
         EntityResultWrapper<BasicTestEntity> wrapper = new EntityResultWrapper<>(BasicTestEntity::new);
         wrapper.setMapping(table
                                .<EntityColumnMapping>getFeature(MappingFeatureType.columnPropertyMapping.createFeatureId(BasicTestEntity.class))
                                .orElseThrow(NullPointerException::new));
 
-
         repository = new DefaultReactiveRepository<>(operator, table, BasicTestEntity.class, wrapper);
         addressRepository = operator.dml().createReactiveRepository("test_address");
 
@@ -286,7 +286,7 @@ public void testInsertMerge() {
         repository
             .createQuery()
             .where(BasicTestEntity::getId, first.getId())
-            .select("id", "name","ext_name")
+            .select("id", "name", "ext_name")
             .fetch()
             .doOnNext(System.out::println)
             .map(BasicTestEntity::getName)
@@ -310,8 +310,9 @@ public void testInsertDuplicate() {
                 .stateEnum(StateEnum.enabled)
                 .enabled(true)
                 .build())
-            .collectList()
-            .as(repository::insertBatch)
+//            .collectList()
+            .as(repository::save)
+            .map(r -> r.getTotal())
             .as(StepVerifier::create)
             .expectNext(3)
             .verifyComplete();