[improve](txn insert) txn insert support update stmt (apache#33034)

feiniaofeiafei · Mar 29, 2024 · 347b048 · 347b048
1 parent 3f883ce
commit 347b048
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 3 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java
@@ -140,6 +140,7 @@
 import org.apache.doris.nereids.trees.plans.commands.CreateTableCommand;
 import org.apache.doris.nereids.trees.plans.commands.Forward;
 import org.apache.doris.nereids.trees.plans.commands.NotAllowFallback;
+import org.apache.doris.nereids.trees.plans.commands.UpdateCommand;
 import org.apache.doris.nereids.trees.plans.commands.insert.BatchInsertIntoTableCommand;
 import org.apache.doris.nereids.trees.plans.commands.insert.InsertIntoTableCommand;
 import org.apache.doris.nereids.trees.plans.commands.insert.InsertOverwriteTableCommand;
@@ -649,9 +650,9 @@ private void executeByNereids(TUniqueId queryId) throws Exception {
         LogicalPlan logicalPlan = ((LogicalPlanAdapter) parsedStmt).getLogicalPlan();
         // when we in transaction mode, we only support insert into command and transaction command
         if (context.isTxnModel()) {
-            if (!(logicalPlan instanceof BatchInsertIntoTableCommand
-                    || logicalPlan instanceof InsertIntoTableCommand)) {
-                String errMsg = "This is in a transaction, only insert, commit, rollback is acceptable.";
+            if (!(logicalPlan instanceof BatchInsertIntoTableCommand || logicalPlan instanceof InsertIntoTableCommand
+                    || logicalPlan instanceof UpdateCommand)) {
+                String errMsg = "This is in a transaction, only insert, update, commit, rollback is acceptable.";
                 throw new NereidsException(errMsg, new AnalysisException(errMsg));
             }
         }

diff --git a/regression-test/data/insert_p0/txn_insert.out b/regression-test/data/insert_p0/txn_insert.out
@@ -298,3 +298,9 @@
 2	3.3	xyz	[1]	[1, 0]
 2	3.3	xyz	[1]	[1, 0]
 
+-- !select25 --
+1	a	101
+
+-- !select26 --
+1	a	100
+
diff --git a/regression-test/suites/insert_p0/txn_insert.groovy b/regression-test/suites/insert_p0/txn_insert.groovy
@@ -213,5 +213,34 @@ suite("txn_insert") {
             order_qt_select23 """select * from ${table}_1"""
             order_qt_select24 """select * from ${table}_2"""
         }
+
+        // 7. update stmt
+        if (use_nereids_planner) {
+            def ut_table = "txn_insert_ut"
+            for (def i in 1..2) {
+                def tableName = ut_table + "_" + i
+                sql """ DROP TABLE IF EXISTS ${tableName} """
+                sql """
+                    CREATE TABLE ${tableName} (
+                        `ID` int(11) NOT NULL,
+                        `NAME` varchar(100) NULL,
+                        `score` int(11) NULL
+                    ) ENGINE=OLAP
+                    unique KEY(`id`)
+                    COMMENT 'OLAP'
+                    DISTRIBUTED BY HASH(`id`) BUCKETS 1
+                    PROPERTIES (
+                        "replication_num" = "1"
+                    );
+                """
+            }
+            sql """ insert into ${ut_table}_1 values(1, "a", 100); """
+            sql """ begin; """
+            sql """ insert into ${ut_table}_2 select * from ${ut_table}_1; """
+            sql """ update ${ut_table}_1 set score = 101 where id = 1; """
+            sql """ commit; """
+            order_qt_select25 """select * from ${ut_table}_1 """
+            order_qt_select26 """select * from ${ut_table}_2 """
+        }
     }
 }