Merge pull request #356 from NodeDB-Lab/fix/statement-command-tag

Report one real command tag per DML statement on every engine, protocol, and path

Author: farhan-syah

nodedb-client/src/native/connection/response.rs

@@ -54,6 +54,7 @@ pub(super) fn response_to_query_result(resp: NativeResponse) -> NodeDbResult<Que
         columns: resp.columns.unwrap_or_default(),
         rows: resp.rows.unwrap_or_default(),
         rows_affected: resp.rows_affected.unwrap_or(0),
+        command: resp.command,
     })
 }
 
@@ -69,6 +70,7 @@ mod tests {
                 columns: vec!["x".into()],
                 rows: vec![vec![nodedb_types::Value::Integer(42)]],
                 rows_affected: 0,
+                command: None,
             },
             0,
         );

nodedb-client/src/remote/client/sql_lifecycle.rs

@@ -49,6 +49,7 @@ impl NodeDbRemote {
             columns,
             rows,
             rows_affected: 0,
+            command: None,
         })
     }
 

nodedb-cluster-tests/tests/common_suite/cases/calvin_cluster_pgwire_e2e.rs

@@ -30,10 +30,10 @@ use crate::common;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 
-use nodedb::types::{DatabaseId, VShardId};
 use nodedb_cluster::calvin::SEQUENCER_GROUP_ID;
 use tokio_postgres::SimpleQueryMessage;
 
+use super::vshard_names::distinct_vshard_collections;
 use common::cluster_harness::{TestClusterNode, wait_for};
 
 /// Observed sequencer-group leader id from a node's local Raft status, or `0`
@@ -61,23 +61,6 @@ async fn value_of(client: &tokio_postgres::Client, coll: &str, id: &str) -> Opti
     })
 }
 
-/// Find two collection names whose vShard ids differ.
-fn two_distinct_vshard_collections() -> (String, String) {
-    let mut first: Option<(String, u32)> = None;
-    for i in 0u32..512 {
-        let name = format!("calvin_e2e_{i}");
-        let vshard = VShardId::from_collection_in_database(DatabaseId::DEFAULT, &name).as_u32();
-        if let Some((ref fname, fv)) = first {
-            if fv != vshard {
-                return (fname.clone(), name);
-            }
-        } else {
-            first = Some((name, vshard));
-        }
-    }
-    panic!("could not find two distinct-vshard collections in 512 tries");
-}
-
 /// Calvin multi-shard batch via pgwire `simple_query` COMMITS when sent as an
 /// interactive `BEGIN ... COMMIT` block.
 ///
@@ -124,7 +107,7 @@ async fn calvin_multishard_write_in_explicit_block_commits() {
     )
     .await;
 
-    let (col_a, col_b) = two_distinct_vshard_collections();
+    let (col_a, col_b) = distinct_vshard_collections("calvin_e2e_0", "calvin_e2e");
 
     // Create both collections on this (single) node.
     node.client

nodedb-cluster-tests/tests/common_suite/cases/calvin_multi_shard_bitemporal_best_effort_restart.rs

@@ -11,8 +11,7 @@
 //!
 //! 1. Two `document_schemaless` collections, each created `WITH
 //!    (bitemporal=true)`, are placed on DIFFERENT vShards
-//!    (`distinct_vshard_bitemporal_collections`, same technique as
-//!    `calvin_multi_shard_redo_restart.rs`).
+//!    (`vshard_names::distinct_vshard_collections`).
 //! 2. `BEGIN; INSERT INTO <a>; INSERT INTO <b>; COMMIT` is sent as ONE
 //!    `simple_query` call, so both writes are buffered inside the block and,
 //!    on COMMIT, `classify_dispatch` sees writes on two vShards → MultiShard.
@@ -32,9 +31,9 @@ use crate::common;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 
-use nodedb::types::{DatabaseId, VShardId};
 use tokio_postgres::SimpleQueryMessage;
 
+use super::vshard_names::distinct_vshard_collections;
 use common::cluster_harness::{TestClusterNode, read_once_a_leader_exists, wait_for};
 
 /// Observed sequencer-group leader id from a node's local Raft status, or `0`
@@ -61,26 +60,6 @@ fn admitted_total(node: &TestClusterNode) -> u64 {
         .unwrap_or(0)
 }
 
-/// A `(coll_a, coll_b)` pair of bitemporal collection names whose vShard ids
-/// differ, so a transaction writing to both is genuinely multi-shard.
-/// Deterministic: `VShardId::from_collection_in_database` is a pure function
-/// of the database id + collection name bytes. Same technique as
-/// `calvin_multi_shard_redo_restart.rs::distinct_vshard_collections`.
-fn distinct_vshard_bitemporal_collections() -> (String, String) {
-    let a_name = "bt_a".to_string();
-    let va = VShardId::from_collection_in_database(DatabaseId::DEFAULT, &a_name).as_u32();
-    for i in 0u32..512 {
-        let b_name = format!("bt_b_{i}");
-        if VShardId::from_collection_in_database(DatabaseId::DEFAULT, &b_name).as_u32() != va {
-            return (a_name, b_name);
-        }
-    }
-    panic!(
-        "could not find a second bitemporal collection name on a distinct vShard \
-         from the first in 512 tries"
-    );
-}
-
 /// Current `value` for `id` in a bitemporal document collection, or `None` if
 /// not visible.
 ///
@@ -160,7 +139,7 @@ async fn calvin_multi_shard_bitemporal_best_effort_commit_survives_wal_only_rest
     )
     .await;
 
-    let (coll_a, coll_b) = distinct_vshard_bitemporal_collections();
+    let (coll_a, coll_b) = distinct_vshard_collections("bt_a", "bt_b");
 
     node.client
         .simple_query(&format!(

nodedb-cluster-tests/tests/common_suite/cases/calvin_multi_shard_bitemporal_restart.rs

@@ -9,8 +9,7 @@
 //!
 //! 1. Two `document_schemaless` collections, each created `WITH
 //!    (bitemporal=true)`, are placed on DIFFERENT vShards
-//!    (`distinct_vshard_bitemporal_collections`, same technique as
-//!    `calvin_multi_shard_redo_restart.rs`).
+//!    (`vshard_names::distinct_vshard_collections`).
 //! 2. `BEGIN; INSERT INTO <a>; INSERT INTO <b>; COMMIT` is sent as ONE
 //!    `simple_query` call, so both writes are buffered inside the block and,
 //!    on COMMIT, `classify_dispatch` sees writes on two vShards → MultiShard
@@ -29,9 +28,9 @@ use crate::common;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 
-use nodedb::types::{DatabaseId, VShardId};
 use tokio_postgres::SimpleQueryMessage;
 
+use super::vshard_names::distinct_vshard_collections;
 use common::cluster_harness::{TestClusterNode, read_once_a_leader_exists, wait_for};
 
 /// Observed sequencer-group leader id from a node's local Raft status, or `0`
@@ -58,26 +57,6 @@ fn admitted_total(node: &TestClusterNode) -> u64 {
         .unwrap_or(0)
 }
 
-/// A `(coll_a, coll_b)` pair of bitemporal collection names whose vShard ids
-/// differ, so a transaction writing to both is genuinely multi-shard.
-/// Deterministic: `VShardId::from_collection_in_database` is a pure function
-/// of the database id + collection name bytes. Same technique as
-/// `calvin_multi_shard_redo_restart.rs::distinct_vshard_collections`.
-fn distinct_vshard_bitemporal_collections() -> (String, String) {
-    let a_name = "bt_a".to_string();
-    let va = VShardId::from_collection_in_database(DatabaseId::DEFAULT, &a_name).as_u32();
-    for i in 0u32..512 {
-        let b_name = format!("bt_b_{i}");
-        if VShardId::from_collection_in_database(DatabaseId::DEFAULT, &b_name).as_u32() != va {
-            return (a_name, b_name);
-        }
-    }
-    panic!(
-        "could not find a second bitemporal collection name on a distinct vShard \
-         from the first in 512 tries"
-    );
-}
-
 /// Current `value` for `id` in a bitemporal document collection, or `None` if
 /// not visible.
 ///
@@ -156,7 +135,7 @@ async fn calvin_multi_shard_bitemporal_commit_survives_wal_only_restart() {
     )
     .await;
 
-    let (coll_a, coll_b) = distinct_vshard_bitemporal_collections();
+    let (coll_a, coll_b) = distinct_vshard_collections("bt_a", "bt_b");
 
     node.client
         .simple_query(&format!(

nodedb-cluster-tests/tests/common_suite/cases/calvin_multi_shard_redo_restart.rs

@@ -6,8 +6,8 @@
 //! `vector_index_txn_restart.rs` for the single-shard analogue).
 //!
 //! 1. Two collections — a KV collection and a vector-indexed document
-//!    collection — are created on DIFFERENT vShards (`distinct_vshard_
-//!    collections`, same technique as `calvin_cluster_pgwire_e2e.rs`).
+//!    collection — are created on DIFFERENT vShards
+//!    (`vshard_names::distinct_vshard_collections`).
 //! 2. `BEGIN; INSERT INTO <kv>; INSERT INTO <vecdocs>; COMMIT` is sent as ONE
 //!    `simple_query` call. tokio-postgres ships this as a single wire
 //!    message; the server buffers the two INSERTs during the transaction and,
@@ -25,9 +25,9 @@ use crate::common;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 
-use nodedb::types::{DatabaseId, VShardId};
 use tokio_postgres::SimpleQueryMessage;
 
+use super::vshard_names::distinct_vshard_collections;
 use common::cluster_harness::{TestClusterNode, read_once_a_leader_exists, wait_for};
 
 /// Observed sequencer-group leader id from a node's local Raft status, or `0`
@@ -54,26 +54,6 @@ fn admitted_total(node: &TestClusterNode) -> u64 {
         .unwrap_or(0)
 }
 
-/// A `(kv_name, vec_name)` pair of collection names whose vShard ids differ,
-/// so a transaction writing to both is genuinely multi-shard. Deterministic:
-/// `VShardId::from_collection_in_database` is a pure function of the database
-/// id + collection name bytes. Same technique as
-/// `calvin_cluster_pgwire_e2e.rs::two_distinct_vshard_collections`.
-fn distinct_vshard_collections() -> (String, String) {
-    let kv_name = "cmr_kv".to_string();
-    let vkv = VShardId::from_collection_in_database(DatabaseId::DEFAULT, &kv_name).as_u32();
-    for i in 0u32..512 {
-        let vec_name = format!("cmr_vecdocs_{i}");
-        if VShardId::from_collection_in_database(DatabaseId::DEFAULT, &vec_name).as_u32() != vkv {
-            return (kv_name, vec_name);
-        }
-    }
-    panic!(
-        "could not find a vector-doc collection name on a distinct vShard from \
-         the KV collection in 512 tries"
-    );
-}
-
 /// Single-row `col` value for `id` in a KV/document collection, or `None` if
 /// not visible.
 ///
@@ -140,7 +120,7 @@ async fn calvin_multi_shard_write_in_explicit_block_commits_and_survives_restart
     )
     .await;
 
-    let (kv, vecdocs) = distinct_vshard_collections();
+    let (kv, vecdocs) = distinct_vshard_collections("cmr_kv", "cmr_vecdocs");
 
     node.client
         .simple_query(&format!(