feat: session auto-enlistment for 2PC DML routing

Wire Session::execute_statement() to route DML through
TransactionManager::route_dml() when a distributed transaction is active
and sharding is configured. Previously, DML in a session with a
distributed transaction bypassed the 2PC pinned session and went through
a fresh pool connection.

- Add is_distributed() and route_dml() virtual methods to
  TransactionManager with safe defaults
- Override both in DistributedTransactionManager to delegate to
  execute_participant_dml()
- Rewrite Session::execute_statement() DML path to use
  DistributedPlanner::distribute_dml() when sharding is configured,
  routing through route_dml() for distributed txns or
  execute_dml() otherwise
- Add for_each_remote_scan() helper for scatter DML across SET_OP trees
- Add 8 tests covering defaults, overrides, and session-level routing

Author: renecannao

include/sql_engine/distributed_txn.h

@@ -156,13 +156,10 @@ class DistributedTransactionManager : public TransactionManager {
     // backend isn't enlisted and can't be enlisted, the result is an
     // error.
     //
-    // External DML routing (Session::execute_statement) currently does
-    // NOT call this method -- it calls executor_->execute_dml directly,
-    // which means DML inside a distributed transaction currently goes
-    // through a fresh pool connection and is NOT part of the 2PC. The
-    // integration hook is a follow-up. For now, unit tests that want
-    // to simulate in-transaction DML should call this method
-    // explicitly.
+    // External DML routing: Session::execute_statement() now routes DML
+    // through this method when a distributed transaction is active and
+    // sharding is configured. This ensures DML inside a distributed
+    // transaction goes through the pinned session and is part of the 2PC.
     DmlResult execute_participant_dml(const char* backend_name,
                                       sql_parser::StringRef sql) {
         DmlResult r;
@@ -248,6 +245,13 @@ class DistributedTransactionManager : public TransactionManager {
     bool is_auto_commit() const override { return auto_commit_; }
     void set_auto_commit(bool ac) override { auto_commit_ = ac; }
 
+    bool is_distributed() const override { return true; }
+
+    DmlResult route_dml(const char* backend_name,
+                        sql_parser::StringRef sql) override {
+        return execute_participant_dml(backend_name, sql);
+    }
+
     const std::string& txn_id() const { return txn_id_; }
     const std::vector<std::string>& participants() const { return participants_; }
 

include/sql_engine/session.h

@@ -18,6 +18,7 @@
 
 #include <cstring>
 #include <string>
+#include <functional>
 #include <unordered_map>
 #include <list>
 #include <memory>
@@ -201,9 +202,56 @@ class Session {
         bool implicit_txn = txn_mgr_.is_auto_commit() && !txn_mgr_.in_transaction();
         if (implicit_txn) txn_mgr_.begin();
 
-        PlanExecutor<D> executor(functions_, catalog_, parser_.arena());
-        wire_executor(executor);
-        DmlResult result = executor.execute_dml(plan);
+        DmlResult result;
+
+        // If sharding is configured, distribute DML to remote backends.
+        if (shard_map_ && remote_executor_) {
+            DistributedPlanner<D> dp(*shard_map_, catalog_, parser_.arena(),
+                                      remote_executor_, &functions_);
+            PlanNode* dist_plan = dp.distribute_dml(plan);
+
+            if (dist_plan && dist_plan->type == PlanNodeType::REMOTE_SCAN) {
+                // Single-shard DML
+                sql_parser::StringRef sql_ref{dist_plan->remote_scan.remote_sql,
+                                              dist_plan->remote_scan.remote_sql_len};
+                if (txn_mgr_.in_transaction() && txn_mgr_.is_distributed()) {
+                    result = txn_mgr_.route_dml(dist_plan->remote_scan.backend_name, sql_ref);
+                } else {
+                    result = remote_executor_->execute_dml(
+                        dist_plan->remote_scan.backend_name, sql_ref);
+                }
+            } else if (dist_plan && dist_plan->type == PlanNodeType::SET_OP) {
+                // Scatter DML to multiple shards
+                result.success = true;
+                result.affected_rows = 0;
+                for_each_remote_scan(dist_plan, [&](const PlanNode* rs) {
+                    sql_parser::StringRef s{rs->remote_scan.remote_sql,
+                                            rs->remote_scan.remote_sql_len};
+                    DmlResult shard_result;
+                    if (txn_mgr_.in_transaction() && txn_mgr_.is_distributed()) {
+                        shard_result = txn_mgr_.route_dml(rs->remote_scan.backend_name, s);
+                    } else {
+                        shard_result = remote_executor_->execute_dml(
+                            rs->remote_scan.backend_name, s);
+                    }
+                    if (!shard_result.success) {
+                        result.success = false;
+                        result.error_message = shard_result.error_message;
+                    }
+                    result.affected_rows += shard_result.affected_rows;
+                });
+            } else {
+                // Not distributed (table not in shard map) -- local execution
+                PlanExecutor<D> executor(functions_, catalog_, parser_.arena());
+                wire_executor(executor);
+                result = executor.execute_dml(plan);
+            }
+        } else {
+            // No sharding: local execution
+            PlanExecutor<D> executor(functions_, catalog_, parser_.arena());
+            wire_executor(executor);
+            result = executor.execute_dml(plan);
+        }
 
         if (implicit_txn) {
             if (result.success)
@@ -294,6 +342,19 @@ class Session {
         plan_cache_[plan_cache_order_.front().key] = plan_cache_order_.begin();
     }
 
+    static void for_each_remote_scan(const PlanNode* node,
+                                      const std::function<void(const PlanNode*)>& fn) {
+        if (!node) return;
+        if (node->type == PlanNodeType::REMOTE_SCAN) {
+            fn(node);
+            return;
+        }
+        if (node->type == PlanNodeType::SET_OP) {
+            for_each_remote_scan(node->left, fn);
+            for_each_remote_scan(node->right, fn);
+        }
+    }
+
     void wire_executor(PlanExecutor<D>& executor) {
         for (auto& kv : sources_)
             executor.add_data_source(kv.first.c_str(), kv.second);

include/sql_engine/transaction_manager.h

@@ -1,6 +1,9 @@
 #ifndef SQL_ENGINE_TRANSACTION_MANAGER_H
 #define SQL_ENGINE_TRANSACTION_MANAGER_H
 
+#include "sql_engine/dml_result.h"
+#include "sql_parser/common.h"
+
 namespace sql_engine {
 
 // Abstract interface for transaction management.
@@ -24,6 +27,19 @@ class TransactionManager {
     virtual bool in_transaction() const = 0;
     virtual bool is_auto_commit() const = 0;
     virtual void set_auto_commit(bool ac) = 0;
+
+    // Distributed transaction support. Override in DistributedTransactionManager.
+    virtual bool is_distributed() const { return false; }
+
+    // Route DML through the transaction's pinned session for a given backend.
+    // Default: returns DmlResult with success=false.
+    // DistributedTransactionManager overrides to route through execute_participant_dml.
+    virtual DmlResult route_dml(const char* /*backend_name*/,
+                                sql_parser::StringRef /*sql*/) {
+        DmlResult r;
+        r.error_message = "route_dml not supported by this transaction manager";
+        return r;
+    }
 };
 
 } // namespace sql_engine