bugfix:throw exception when update pk (apache#3287)

slievrly · Dec 8, 2020 · c2c3488 · c2c3488
1 parent 3530aa0
commit c2c3488
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 14 deletions.
diff --git a/rm-datasource/src/main/java/io/seata/rm/datasource/exec/AbstractDMLBaseExecutor.java b/rm-datasource/src/main/java/io/seata/rm/datasource/exec/AbstractDMLBaseExecutor.java
@@ -26,7 +26,6 @@
 import java.util.stream.Collectors;
 
 import io.seata.common.exception.NotSupportYetException;
-import io.seata.common.exception.ShouldNeverHappenException;
 import io.seata.common.util.CollectionUtils;
 import io.seata.rm.datasource.AbstractConnectionProxy;
 import io.seata.rm.datasource.ConnectionContext;
@@ -35,7 +34,6 @@
 import io.seata.rm.datasource.sql.struct.TableRecords;
 import io.seata.sqlparser.SQLRecognizer;
 import io.seata.sqlparser.util.JdbcConstants;
-import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -202,13 +200,4 @@ public static boolean isLockRetryPolicyBranchRollbackOnConflict() {
             return LOCK_RETRY_POLICY_BRANCH_ROLLBACK_ON_CONFLICT;
         }
     }
-
-    protected void assertContainsPKColumnName(List<String> updateColumns) {
-        for (String columnName : updateColumns) {
-            String standardColumnName = getStandardPkColumnName(columnName);
-            if (StringUtils.isNotEmpty(standardColumnName)) {
-                throw new ShouldNeverHappenException("Sorry, update pk value is not supported!");
-            }
-        }
-    }
 }
diff --git a/rm-datasource/src/main/java/io/seata/rm/datasource/exec/BaseTransactionalExecutor.java b/rm-datasource/src/main/java/io/seata/rm/datasource/exec/BaseTransactionalExecutor.java
@@ -22,6 +22,7 @@
 import java.util.Map;
 import java.util.Objects;
 
+import io.seata.common.exception.ShouldNeverHappenException;
 import io.seata.common.util.CollectionUtils;
 import io.seata.common.util.IOUtil;
 import io.seata.common.util.StringUtils;
@@ -269,7 +270,11 @@ protected void prepareUndoLog(TableRecords beforeImage, TableRecords afterImage)
         if (beforeImage.getRows().isEmpty() && afterImage.getRows().isEmpty()) {
             return;
         }
-
+        if (SQLType.UPDATE == sqlRecognizer.getSQLType()) {
+            if (beforeImage.getRows().size() != afterImage.getRows().size()) {
+                throw new ShouldNeverHappenException("Before image size is not equaled to after image size, probably because you updated the primary keys.");
+            }
+        }
         ConnectionProxy connectionProxy = statementProxy.getConnectionProxy();
 
         TableRecords lockKeyRecords = sqlRecognizer.getSQLType() == SQLType.DELETE ? beforeImage : afterImage;

diff --git a/rm-datasource/src/main/java/io/seata/rm/datasource/exec/MultiExecutor.java b/rm-datasource/src/main/java/io/seata/rm/datasource/exec/MultiExecutor.java
@@ -16,9 +16,11 @@
 package io.seata.rm.datasource.exec;
 
 
+import io.seata.common.exception.ShouldNeverHappenException;
 import io.seata.rm.datasource.StatementProxy;
 import io.seata.rm.datasource.sql.struct.TableRecords;
 import io.seata.sqlparser.SQLRecognizer;
+import io.seata.sqlparser.SQLType;
 
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -116,6 +118,11 @@ protected void prepareUndoLog(TableRecords beforeImage, TableRecords afterImage)
             sqlRecognizer = recognizer = entry.getKey();
             beforeImage = entry.getValue();
             afterImage = afterImagesMap.get(recognizer);
+            if (SQLType.UPDATE == sqlRecognizer.getSQLType()) {
+                if (beforeImage.getRows().size() != afterImage.getRows().size()) {
+                    throw new ShouldNeverHappenException("Before image size is not equaled to after image size, probably because you updated the primary keys.");
+                }
+            }
             super.prepareUndoLog(beforeImage, afterImage);
         }
     }

diff --git a/rm-datasource/src/main/java/io/seata/rm/datasource/exec/MultiUpdateExecutor.java b/rm-datasource/src/main/java/io/seata/rm/datasource/exec/MultiUpdateExecutor.java
@@ -80,7 +80,6 @@ protected TableRecords beforeImage() throws SQLException {
             sqlRecognizer = recognizer;
             SQLUpdateRecognizer sqlUpdateRecognizer = (SQLUpdateRecognizer) recognizer;
             List<String> updateColumns = sqlUpdateRecognizer.getUpdateColumns();
-            assertContainsPKColumnName(updateColumns);
             updateColumnsSet.addAll(updateColumns);
             if (noWhereCondition) {
                 continue;

diff --git a/rm-datasource/src/main/java/io/seata/rm/datasource/exec/UpdateExecutor.java b/rm-datasource/src/main/java/io/seata/rm/datasource/exec/UpdateExecutor.java
@@ -75,7 +75,6 @@ protected TableRecords beforeImage() throws SQLException {
     private String buildBeforeImageSQL(TableMeta tableMeta, ArrayList<List<Object>> paramAppenderList) {
         SQLUpdateRecognizer recognizer = (SQLUpdateRecognizer) sqlRecognizer;
         List<String> updateColumns = recognizer.getUpdateColumns();
-        assertContainsPKColumnName(updateColumns);
         StringBuilder prefix = new StringBuilder("SELECT ");
         StringBuilder suffix = new StringBuilder(" FROM ").append(getFromTableInSQL());
         String whereCondition = buildWhereCondition(recognizer, paramAppenderList);