Merge pull request #59 from ProxySQL/fix/shard-having-join-2pc

fix: close remaining sharding correctness holes

Author: renecannao

include/sql_engine/distributed_planner.h

@@ -999,12 +999,72 @@ class DistributedPlanner {
         return result;
     }
 
-    // Case 5: Cross-backend join
+    bool join_on_shard_keys(const sql_parser::AstNode* cond,
+                            sql_parser::StringRef left_key,
+                            sql_parser::StringRef right_key) const {
+        if (!cond || cond->type != sql_parser::NodeType::NODE_BINARY_OP) return false;
+        sql_parser::StringRef op = cond->value();
+        if (op.len != 1 || op.ptr[0] != '=') return false;
+        const sql_parser::AstNode* l = cond->first_child;
+        const sql_parser::AstNode* r = l ? l->next_sibling : nullptr;
+        if (!l || !r) return false;
+        return (is_shard_key_ref(l, left_key) && is_shard_key_ref(r, right_key)) ||
+               (is_shard_key_ref(l, right_key) && is_shard_key_ref(r, left_key));
+    }
+
+    PlanNode* distribute_colocated_join(PlanNode* join_node,
+                                        const TableInfo* left_table,
+                                        const TableInfo* right_table) {
+        ScanContext lctx = extract_scan_context(join_node->left);
+        ScanContext rctx = extract_scan_context(join_node->right);
+        const sql_parser::AstNode* where_expr = nullptr;
+        if (lctx.where_expr && rctx.where_expr) {
+            sql_parser::AstNode* and_node = sql_parser::make_node(
+                arena_, sql_parser::NodeType::NODE_BINARY_OP,
+                sql_parser::StringRef{"AND", 3});
+            and_node->add_child(const_cast<sql_parser::AstNode*>(lctx.where_expr));
+            and_node->add_child(const_cast<sql_parser::AstNode*>(rctx.where_expr));
+            where_expr = and_node;
+        } else if (lctx.where_expr) {
+            where_expr = lctx.where_expr;
+        } else {
+            where_expr = rctx.where_expr;
+        }
+
+        const auto& shard_list = shards_.get_shards(left_table->table_name);
+        PlanNode* current = nullptr;
+        for (const auto& shard : shard_list) {
+            sql_parser::StringRef sql = qb_.build_select_join(
+                left_table, right_table, join_node->join.condition, where_expr);
+            PlanNode* rs = make_remote_scan(shard.backend_name.c_str(), sql, left_table);
+            if (!current) {
+                current = rs;
+            } else {
+                PlanNode* union_node = make_plan_node(arena_, PlanNodeType::SET_OP);
+                union_node->set_op.op = SET_OP_UNION;
+                union_node->set_op.all = true;
+                union_node->left = current;
+                union_node->right = rs;
+                current = union_node;
+            }
+        }
+        return current ? current : join_node;
+    }
+
     PlanNode* distribute_join(PlanNode* join_node) {
-        // Get tables from each side
         const TableInfo* left_table = find_table(join_node->left);
         const TableInfo* right_table = find_table(join_node->right);
 
+        if (left_table && right_table &&
+            shards_.is_sharded(left_table->table_name) &&
+            shards_.is_sharded(right_table->table_name) &&
+            shards_.same_routing(left_table->table_name, right_table->table_name) &&
+            join_on_shard_keys(join_node->join.condition,
+                               shards_.get_shard_key(left_table->table_name),
+                               shards_.get_shard_key(right_table->table_name))) {
+            return distribute_colocated_join(join_node, left_table, right_table);
+        }
+
         PlanNode* left_dist = nullptr;
         PlanNode* right_dist = nullptr;
 
@@ -1172,18 +1232,13 @@ class DistributedPlanner {
     PlanNode* distribute_update(PlanNode* plan) {
         const auto& up = plan->update_plan;
         const TableInfo* table = up.table;
-        if (!table || !shards_.has_table(table->table_name)) return plan;
 
-        // Multi-table UPDATE: emit full SQL from AST, route to primary table's backend
         if (up.original_ast) {
-            sql_parser::StringRef sql = qb_.build_update_from_ast(up.original_ast);
-            if (!shards_.is_sharded(table->table_name)) {
-                return make_remote_scan(shards_.get_backend(table->table_name), sql, table);
-            }
-            const auto& shard_list = shards_.get_shards(table->table_name);
-            return scatter_dml_to_shards(table, shard_list, [&]() { return sql; });
+            return distribute_multi_table_dml(up.original_ast, table, true);
         }
 
+        if (!table || !shards_.has_table(table->table_name)) return plan;
+
         // Check for cross-shard subqueries in WHERE and rewrite
         const sql_parser::AstNode* where_expr = up.where_expr;
         if (where_expr && has_subquery(where_expr) && remote_executor_) {
@@ -1220,18 +1275,13 @@ class DistributedPlanner {
     PlanNode* distribute_delete(PlanNode* plan) {
         const auto& dp = plan->delete_plan;
         const TableInfo* table = dp.table;
-        if (!table || !shards_.has_table(table->table_name)) return plan;
 
-        // Multi-table DELETE: emit full SQL from AST, route to primary table's backend
         if (dp.original_ast) {
-            sql_parser::StringRef sql = qb_.build_delete_from_ast(dp.original_ast);
-            if (!shards_.is_sharded(table->table_name)) {
-                return make_remote_scan(shards_.get_backend(table->table_name), sql, table);
-            }
-            const auto& shard_list = shards_.get_shards(table->table_name);
-            return scatter_dml_to_shards(table, shard_list, [&]() { return sql; });
+            return distribute_multi_table_dml(dp.original_ast, table, false);
         }
 
+        if (!table || !shards_.has_table(table->table_name)) return plan;
+
         // Check for cross-shard subqueries in WHERE and rewrite
         const sql_parser::AstNode* where_expr = dp.where_expr;
         if (where_expr && has_subquery(where_expr) && remote_executor_) {
@@ -1318,7 +1368,68 @@ class DistributedPlanner {
         return false;
     }
 
-    // Scatter DML SQL to all shards, combining results via UNION ALL
+    void collect_ast_table_names(const sql_parser::AstNode* n,
+                                 std::vector<sql_parser::StringRef>& out) const {
+        if (!n) return;
+        if (n->type == sql_parser::NodeType::NODE_TABLE_REF && n->first_child) {
+            const sql_parser::AstNode* name = n->first_child;
+            if (name->type == sql_parser::NodeType::NODE_IDENTIFIER) {
+                out.push_back(name->value());
+            } else if (name->type == sql_parser::NodeType::NODE_QUALIFIED_NAME) {
+                const sql_parser::AstNode* schema = name->first_child;
+                const sql_parser::AstNode* table = schema ? schema->next_sibling : nullptr;
+                if (table) out.push_back(table->value());
+                else if (schema) out.push_back(schema->value());
+            }
+        }
+        for (const sql_parser::AstNode* c = n->first_child; c; c = c->next_sibling) {
+            collect_ast_table_names(c, out);
+        }
+    }
+
+    PlanNode* distribute_multi_table_dml(const sql_parser::AstNode* ast,
+                                         const TableInfo* primary,
+                                         bool is_update) {
+        std::vector<sql_parser::StringRef> names;
+        collect_ast_table_names(ast, names);
+        const char* backend = nullptr;
+        bool saw_mapped = false;
+        for (sql_parser::StringRef name : names) {
+            if (!shards_.has_table(name)) continue;
+            saw_mapped = true;
+            if (shards_.is_sharded(name)) {
+                return fail_dml(is_update
+                    ? "multi-table UPDATE is not supported on sharded tables"
+                    : "multi-table DELETE is not supported on sharded tables");
+            }
+            const char* b = shards_.get_backend(name);
+            if (backend && b && std::strcmp(backend, b) != 0) {
+                return fail_dml(is_update
+                    ? "multi-table UPDATE spans multiple backends"
+                    : "multi-table DELETE spans multiple backends");
+            }
+            if (b) backend = b;
+        }
+        if (!backend && primary && shards_.has_table(primary->table_name)) {
+            if (shards_.is_sharded(primary->table_name)) {
+                return fail_dml(is_update
+                    ? "multi-table UPDATE is not supported on sharded tables"
+                    : "multi-table DELETE is not supported on sharded tables");
+            }
+            backend = shards_.get_backend(primary->table_name);
+            saw_mapped = true;
+        }
+        if (!backend || !saw_mapped) {
+            return fail_dml(is_update
+                ? "multi-table UPDATE is not supported on sharded tables"
+                : "multi-table DELETE is not supported on sharded tables");
+        }
+        sql_parser::StringRef sql = is_update
+            ? qb_.build_update_from_ast(ast)
+            : qb_.build_delete_from_ast(ast);
+        return make_remote_scan(backend, sql, primary);
+    }
+
     PlanNode* scatter_dml_to_shards(const TableInfo* table,
                                      const std::vector<ShardInfo>& shard_list,
                                      std::function<sql_parser::StringRef()> build_sql) {

include/sql_engine/distributed_txn.h

@@ -178,6 +178,18 @@ class DistributedTransactionManager : public TransactionManager {
         return executor_.execute_dml(backend_name, sql);
     }
 
+    ResultSet route_query(const char* backend_name,
+                          sql_parser::StringRef sql) override {
+        if (!active_) return executor_.execute(backend_name, sql);
+        auto it = sessions_.find(backend_name);
+        if (it != sessions_.end() && it->second) {
+            return it->second->execute(sql);
+        }
+        return executor_.execute(backend_name, sql);
+    }
+
+    bool route_query_supported() const override { return true; }
+
     bool commit() override {
         if (!active_) return false;
         if (participants_.empty()) {

include/sql_engine/plan_executor.h

@@ -57,6 +57,7 @@
 #include <vector>
 #include <memory>
 #include <cstdlib>
+#include <cstring>
 
 namespace sql_engine {
 
@@ -847,15 +848,128 @@ class PlanExecutor {
         return ptr;
     }
 
+    static bool same_agg_call(const sql_parser::AstNode* a, const sql_parser::AstNode* b) {
+        if (!a || !b) return false;
+        if (a->type != sql_parser::NodeType::NODE_FUNCTION_CALL ||
+            b->type != sql_parser::NodeType::NODE_FUNCTION_CALL) return false;
+        return a->value().equals_ci(b->value().ptr, b->value().len);
+    }
+
+    const sql_parser::AstNode* rewrite_having_expr(const sql_parser::AstNode* expr,
+                                                    PlanNode* agg_node) {
+        if (!expr || !agg_node) return expr;
+        uint16_t group_count = 0;
+        uint16_t agg_count = 0;
+        const sql_parser::AstNode** agg_exprs = nullptr;
+        const sql_parser::AstNode** group_by = nullptr;
+        if (agg_node->type == PlanNodeType::AGGREGATE) {
+            group_count = agg_node->aggregate.group_count;
+            agg_count = agg_node->aggregate.agg_count;
+            agg_exprs = agg_node->aggregate.agg_exprs;
+            group_by = agg_node->aggregate.group_by;
+        } else if (agg_node->type == PlanNodeType::MERGE_AGGREGATE) {
+            group_count = agg_node->merge_aggregate.group_key_count;
+            if (agg_node->merge_aggregate.output_exprs &&
+                agg_node->merge_aggregate.output_expr_count > group_count) {
+                agg_exprs = agg_node->merge_aggregate.output_exprs + group_count;
+                agg_count = static_cast<uint16_t>(
+                    agg_node->merge_aggregate.output_expr_count - group_count);
+                group_by = agg_node->merge_aggregate.output_exprs;
+            }
+        } else {
+            return expr;
+        }
+
+        if (expr->type == sql_parser::NodeType::NODE_FUNCTION_CALL && agg_exprs) {
+            for (uint16_t i = 0; i < agg_count; ++i) {
+                if (same_agg_call(expr, agg_exprs[i])) {
+                    sql_parser::StringRef name = expr->value();
+                    return sql_parser::make_node(
+                        arena_, sql_parser::NodeType::NODE_IDENTIFIER, name);
+                }
+            }
+        }
+
+        bool changed = false;
+        sql_parser::AstNode* clone = sql_parser::make_node(
+            arena_, expr->type, expr->value(), expr->flags);
+        for (const sql_parser::AstNode* c = expr->first_child; c; c = c->next_sibling) {
+            const sql_parser::AstNode* rw = rewrite_having_expr(c, agg_node);
+            if (rw != c) changed = true;
+            if (rw) clone->add_child(const_cast<sql_parser::AstNode*>(rw));
+        }
+        (void)group_count;
+        (void)group_by;
+        return changed ? clone : expr;
+    }
+
+    const TableInfo* make_agg_output_table(PlanNode* agg_node) {
+        if (!agg_node) return nullptr;
+        uint16_t group_count = 0;
+        uint16_t agg_count = 0;
+        const sql_parser::AstNode** group_by = nullptr;
+        const sql_parser::AstNode** agg_exprs = nullptr;
+        if (agg_node->type == PlanNodeType::AGGREGATE) {
+            group_count = agg_node->aggregate.group_count;
+            agg_count = agg_node->aggregate.agg_count;
+            group_by = agg_node->aggregate.group_by;
+            agg_exprs = agg_node->aggregate.agg_exprs;
+        } else if (agg_node->type == PlanNodeType::MERGE_AGGREGATE &&
+                   agg_node->merge_aggregate.output_exprs) {
+            group_count = agg_node->merge_aggregate.group_key_count;
+            group_by = agg_node->merge_aggregate.output_exprs;
+            if (agg_node->merge_aggregate.output_expr_count > group_count) {
+                agg_exprs = agg_node->merge_aggregate.output_exprs + group_count;
+                agg_count = static_cast<uint16_t>(
+                    agg_node->merge_aggregate.output_expr_count - group_count);
+            }
+        } else {
+            return nullptr;
+        }
+
+        uint16_t n = static_cast<uint16_t>(group_count + agg_count);
+        if (n == 0) return nullptr;
+        auto* cols = static_cast<ColumnInfo*>(arena_.allocate(sizeof(ColumnInfo) * n));
+        if (!cols) return nullptr;
+        for (uint16_t i = 0; i < group_count; ++i) {
+            cols[i].ordinal = i;
+            cols[i].nullable = true;
+            cols[i].type = SqlType::make_int();
+            cols[i].name = (group_by && group_by[i]) ? group_by[i]->value()
+                                                     : sql_parser::StringRef{};
+        }
+        for (uint16_t i = 0; i < agg_count; ++i) {
+            cols[group_count + i].ordinal = static_cast<uint16_t>(group_count + i);
+            cols[group_count + i].nullable = true;
+            cols[group_count + i].type = SqlType::make_int();
+            cols[group_count + i].name = (agg_exprs && agg_exprs[i])
+                ? agg_exprs[i]->value() : sql_parser::StringRef{};
+        }
+        auto* ti = static_cast<TableInfo*>(arena_.allocate(sizeof(TableInfo)));
+        if (!ti) return nullptr;
+        std::memset(ti, 0, sizeof(TableInfo));
+        ti->columns = cols;
+        ti->column_count = n;
+        return ti;
+    }
+
     Operator* build_filter(PlanNode* node) {
         Operator* child = build_operator(node->left);
         if (!child && node->left) return nullptr;
 
         std::vector<const TableInfo*> tables;
-        collect_tables(node->left, tables);
+        const sql_parser::AstNode* expr = node->filter.expr;
+        if (node->left && (node->left->type == PlanNodeType::AGGREGATE ||
+                           node->left->type == PlanNodeType::MERGE_AGGREGATE)) {
+            expr = rewrite_having_expr(expr, node->left);
+            const TableInfo* synth = make_agg_output_table(node->left);
+            if (synth) tables.push_back(synth);
+        } else {
+            collect_tables(node->left, tables);
+        }
 
         auto op = std::make_unique<FilterOperator<D>>(
-            child, node->filter.expr, catalog_, tables, functions_, arena_,
+            child, expr, catalog_, tables, functions_, arena_,
             &subquery_exec_, outer_resolver_);
         Operator* ptr = op.get();
         operators_.push_back(std::move(op));

include/sql_engine/remote_query_builder.h

@@ -94,6 +94,28 @@ class RemoteQueryBuilder {
         return sb.finish();
     }
 
+    sql_parser::StringRef build_select_join(
+            const TableInfo* left,
+            const TableInfo* right,
+            const sql_parser::AstNode* on_expr,
+            const sql_parser::AstNode* where_expr)
+    {
+        sql_parser::StringBuilder sb(arena_, 512);
+        sb.append("SELECT * FROM ");
+        if (left) sb.append(left->table_name.ptr, left->table_name.len);
+        sb.append(" JOIN ");
+        if (right) sb.append(right->table_name.ptr, right->table_name.len);
+        if (on_expr) {
+            sb.append(" ON ");
+            emit_expr(on_expr, sb);
+        }
+        if (where_expr) {
+            sb.append(" WHERE ");
+            emit_expr(where_expr, sb);
+        }
+        return sb.finish();
+    }
+
     // Build an INSERT statement string.
     sql_parser::StringRef build_insert(
             const TableInfo* table,