fix: Harden the code

Author: as51340

src/batch_import.cpp

@@ -166,9 +166,9 @@ struct BatchExecutionContext {
       // to `workers_number` (e.g. 32) ROUTE calls at startup. The TTL expiry of
       // the last route determines when sessions are refreshed (see Run).
       if (config.routed_connection) {
-        sessions[thread_i] = utils::bolt::MakeRoutedBoltSession(config, &expiry);
+        sessions.push_back(utils::bolt::MakeRoutedBoltSession(config, &expiry));
       } else {
-        sessions[thread_i] = MakeBoltSession(config);
+        sessions.push_back(MakeBoltSession(config));
       }
       if (!sessions[thread_i].get()) {
         MG_FAIL("a session uninitialized");
@@ -311,7 +311,8 @@ int Run(const utils::bolt::Config &bolt_config, int batch_size, int workers_numb
     // refresh out of the parallel rounds where it would be race-prone.
     if (bolt_config.routed_connection && std::chrono::steady_clock::now() >= execution_context.expiry) {
       for (uint64_t thread_i = 0; thread_i < execution_context.max_concurrent_executions; ++thread_i) {
-        execution_context.sessions[thread_i] = utils::bolt::MakeRoutedBoltSession(bolt_config, &execution_context.expiry);
+        execution_context.sessions[thread_i] =
+            utils::bolt::MakeRoutedBoltSession(bolt_config, &execution_context.expiry);
         if (!execution_context.sessions[thread_i].get()) {
           MG_FAIL("failed to re-route a worker session after TTL expiry");
         }

src/interactive.cpp

@@ -188,7 +188,12 @@ int Run(utils::bolt::Config &bolt_config, const std::string &history, bool no_hi
         }
         continue;
       }
-      auto ret = query::ExecuteQuery(session.Get(), query->query, param_store.AsMap().get(), bolt_config.db);
+      // Resolve the session (may proactively re-route) before observing this query: the re-route decision must
+      // use the transaction state as it was *before* this query, so a COMMIT/ROLLBACK still runs on the session
+      // that holds the open transaction.
+      auto *session_ptr = session.Get();
+      session.ObserveQuery(query->query);
+      auto ret = query::ExecuteQuery(session_ptr, query->query, param_store.AsMap().get(), bolt_config.db);
       if (ret.records.size() > 0) {
         Output(ret.header, ret.records, output_opts, csv_opts);
       }

src/serial_import.cpp

@@ -36,7 +36,11 @@ int Run(const utils::bolt::Config &bolt_config, const format::CsvOptions &csv_op
     }
 
     try {
-      auto ret = query::ExecuteQuery(session.Get(), query->query, nullptr, bolt_config.db);
+      // Resolve the session (may proactively re-route) before observing this query so transaction-control
+      // statements in a dump stream are not interrupted by a re-route mid-transaction.
+      auto *session_ptr = session.Get();
+      session.ObserveQuery(query->query);
+      auto ret = query::ExecuteQuery(session_ptr, query->query, nullptr, bolt_config.db);
       if (ret.records.size() > 0) {
         Output(ret.header, ret.records, output_opts, csv_opts);
       }

src/utils/bolt.cpp

@@ -139,7 +139,9 @@ mg_memory::MgSessionPtr MakeRoutedBoltSession(const Config &config, std::chrono:
     return mg_memory::MakeCustomUnique<mg_session>(nullptr);
   }
   mg_map_insert(routing.get(), "address", mg_value_make_string(coord_address.c_str()));
-  mg_map_insert(extra.get(), "db", mg_value_make_string(config.db.c_str()));
+  if (!config.db.empty()) {
+    mg_map_insert(extra.get(), "db", mg_value_make_string(config.db.c_str()));
+  }
 
   mg_map *rt_raw = nullptr;
   const int status = mg_session_route(coord.get(), routing.get(), nullptr, extra.get(), &rt_raw);
@@ -152,7 +154,7 @@ mg_memory::MgSessionPtr MakeRoutedBoltSession(const Config &config, std::chrono:
 
   // 3. Parse the routing table: TTL (seconds) and the WRITE (main) instance.
   const mg_value *ttl_val = mg_map_at(rt.get(), "ttl");
-  if (ttl_val != nullptr && expiry_out != nullptr) {
+  if (ttl_val != nullptr && mg_value_get_type(ttl_val) == MG_VALUE_TYPE_INTEGER && expiry_out != nullptr) {
     *expiry_out = std::chrono::steady_clock::now() + std::chrono::seconds(mg_value_integer(ttl_val));
   }
 
@@ -224,10 +226,18 @@ void RoutedSession::Rebuild() {
 }
 
 mg_session *RoutedSession::Get() {
-  if (routed_ && std::chrono::steady_clock::now() >= expiry_) {
-    // Proactive TTL re-route: the previous routing table has expired, fetch a
-    // fresh one (and possibly a new main) before handing back the session.
-    Rebuild();
+  // Proactive TTL re-route: once the previous routing table has expired, fetch a fresh one (and possibly a new
+  // main) before handing back the session. Suppressed while an explicit transaction is open (re-routing would
+  // silently drop it) and made non-destructive: a failed re-route keeps the existing working session rather than
+  // discarding a healthy connection over a transient coordinator hiccup.
+  if (routed_ && !in_transaction_ && std::chrono::steady_clock::now() >= expiry_) {
+    auto refreshed = MakeRoutedBoltSession(config_, &expiry_);
+    if (refreshed.get() != nullptr || session_.get() == nullptr) {
+      session_ = std::move(refreshed);
+    } else {
+      // Keep the still-usable session; back off so we retry periodically instead of on every call.
+      expiry_ = std::chrono::steady_clock::now() + kRerouteRetryBackoffSec;
+    }
   }
   return session_.get();
 }
@@ -236,4 +246,15 @@ void RoutedSession::Reconnect() { Rebuild(); }
 
 bool RoutedSession::Connected() const { return session_.get() != nullptr; }
 
+void RoutedSession::ObserveQuery(const std::string &query) {
+  // Track explicit-transaction state from the transaction-control keyword so proactive re-routing can avoid
+  // tearing down an open transaction. Best-effort: matches the leading keyword of the trimmed query.
+  const auto upper = utils::ToUpperCase(utils::Trim(query));
+  if (upper.rfind("BEGIN", 0) == 0) {
+    in_transaction_ = true;
+  } else if (upper.rfind("COMMIT", 0) == 0 || upper.rfind("ROLLBACK", 0) == 0) {
+    in_transaction_ = false;
+  }
+}
+
 }  // namespace utils::bolt

src/utils/bolt.hpp

@@ -51,8 +51,10 @@ class RoutedSession {
   explicit RoutedSession(Config config);
 
   // Returns the underlying session. In routed mode, if the routing table TTL
-  // has expired, transparently re-routes (rebuilds the session) first. May
-  // return nullptr if a (re-)connection attempt failed.
+  // has expired (and no explicit transaction is open), transparently re-routes
+  // first; a failed re-route keeps the existing session. May return nullptr if
+  // a (re-)connection attempt failed and there is no prior session to fall back
+  // on.
   mg_session *Get();
 
   // Forces a rebuild of the session (direct or routed per config). Used by the
@@ -62,13 +64,24 @@ class RoutedSession {
   // True if a session is currently established.
   bool Connected() const;
 
+  // Observes a query about to be executed so the session can track explicit
+  // transaction boundaries (BEGIN/COMMIT/ROLLBACK) and suppress proactive
+  // re-routing while a transaction is open. Safe to call in direct mode (no-op
+  // effect on routing).
+  void ObserveQuery(const std::string &query);
+
  private:
+  // Backoff applied after a failed proactive re-route so we don't retry on
+  // every Get() while the coordinator is briefly unreachable.
+  static constexpr std::chrono::seconds kRerouteRetryBackoffSec{2};
+
   void Rebuild();
 
   Config config_;
   mg_memory::MgSessionPtr session_;
   std::chrono::steady_clock::time_point expiry_{};
   bool routed_;
+  bool in_transaction_{false};
 };
 
 }  // namespace utils::bolt