fix: prepare for insert ... subquery

wy1433 · Apr 24, 2024 · a4f8a3a · a4f8a3a
1 parent 91379eb
commit a4f8a3a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/src/logical_plan/prepare_planner.cpp b/src/logical_plan/prepare_planner.cpp
@@ -18,6 +18,7 @@
 #include "insert_planner.h"
 #include "delete_planner.h"
 #include "update_planner.h"
+#include "union_planner.h"
 #include "transaction_planner.h"
 #include "exec_node.h"
 #include "packet_node.h"
@@ -188,6 +189,9 @@ int PreparePlanner::stmt_prepare(const std::string& stmt_name, const std::string
     case parser::NT_DELETE:
         planner.reset(new DeletePlanner(prepare_ctx.get()));
         break;
+    case parser::NT_UNION:
+        planner.reset(new UnionPlanner(prepare_ctx.get()));
+        break;
     default:
         DB_WARNING("un-supported prepare command type: %d", prepare_ctx->stmt_type);
         return -1;
@@ -210,6 +214,9 @@ int PreparePlanner::stmt_prepare(const std::string& stmt_name, const std::string
         return -1;
     }
     prepare_ctx->root->find_place_holder(prepare_ctx->placeholders);
+    for (auto sub_query_ctx : prepare_ctx->sub_query_plans) {
+        sub_query_ctx->root->find_place_holder(prepare_ctx->placeholders);
+    }
     /*
     // 包括类型推导与常量表达式计算
     ret = ExprOptimize().analyze(prepare_ctx.get());
@@ -242,8 +249,8 @@ int PreparePlanner::stmt_execute(const std::string& stmt_name, std::vector<pb::E
     _ctx->copy_query_context(prepare_ctx.get());
 
     auto* p_placeholders = &prepare_ctx->placeholders;
-    if (!prepare_ctx->is_select) {
-        // TODO dml的plan复用
+    // TODO dml的plan复用
+    if (!prepare_ctx->is_select || prepare_ctx->sub_query_plans.size() > 0) {
         // enable_2pc=true or table has global index need generate txn_id
         set_dml_txn_state(prepare_ctx->prepared_table_id);
         _ctx->plan.CopyFrom(prepare_ctx->plan);
@@ -258,6 +265,15 @@ int PreparePlanner::stmt_execute(const std::string& stmt_name, std::vector<pb::E
             return -1;
         }
         _ctx->root->find_place_holder(_ctx->placeholders);
+        for (auto sub_query_ctx : prepare_ctx->sub_query_plans) {
+            int ret = sub_query_ctx->create_plan_tree();
+            if (ret < 0) {
+                DB_WARNING("Failed to pb_plan to execnode");
+                return -1;
+            }
+            _ctx->add_sub_ctx(sub_query_ctx);
+            sub_query_ctx->root->find_place_holder(_ctx->placeholders);
+        }
         p_placeholders = &_ctx->placeholders;
     }
     if (p_placeholders == nullptr) {

diff --git a/src/logical_plan/query_context.cpp b/src/logical_plan/query_context.cpp
@@ -80,6 +80,11 @@ int QueryContext::copy_query_context(QueryContext* p_query_ctx) {
     stat_info.sample_sql << p_query_ctx->stat_info.sample_sql.str();
     need_learner_backup = p_query_ctx->need_learner_backup;
     use_backup = p_query_ctx->use_backup;
+    has_derived_table = p_query_ctx->has_derived_table;
+    derived_table_ctx_mapping.insert(p_query_ctx->derived_table_ctx_mapping.begin(),
+                                     p_query_ctx->derived_table_ctx_mapping.end());
+    slot_column_mapping.insert(p_query_ctx->slot_column_mapping.begin(),
+                               p_query_ctx->slot_column_mapping.end());
 
     // runtime state
     if (p_query_ctx->is_select) {