From c99ec28ba40eb30645884c3a1761d92d4dd52068 Mon Sep 17 00:00:00 2001 From: Pizo Date: Thu, 16 Jul 2026 20:28:36 +0800 Subject: [PATCH 1/2] fix(control): ensure native session on first-frame OpCode::Begin OpCode::Begin can arrive before any SQL statement has called ensure_session. SessionStore::begin no-ops when the peer has no entry, so the following write autocommitted instead of entering InBlock. Ensure the session exists at the top of handle_begin (SQL BEGIN already does this in sql.rs before calling here). Guarded by native_first_frame_begin_buffers_following_write. --- nodedb/src/control/server/native/dispatch/transaction.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/nodedb/src/control/server/native/dispatch/transaction.rs b/nodedb/src/control/server/native/dispatch/transaction.rs index ff58d0314..c8455484f 100644 --- a/nodedb/src/control/server/native/dispatch/transaction.rs +++ b/nodedb/src/control/server/native/dispatch/transaction.rs @@ -53,6 +53,7 @@ impl TxnDataPlane for NativeTxnDp<'_> { tenant_id: task.tenant_id, trace_id: TraceId::generate(), database_id: task.database_id, + txn_id: task.txn_id, }; let payloads = gw.execute(&gw_ctx, task.plan).await?; Ok(Response { @@ -93,6 +94,11 @@ impl TxnDataPlane for NativeTxnDp<'_> { } pub(crate) fn handle_begin(ctx: &DispatchCtx<'_>, seq: u64) -> NativeResponse { + // OpCode::Begin can arrive before any SQL statement has created the + // session. `SessionStore::begin` no-ops when the peer has no entry, so + // ensure the session exists first (SQL "BEGIN" already does this in + // `sql.rs` before calling here). + ctx.sessions.ensure_session(*ctx.peer_addr); match lifecycle::run_begin(ctx.sessions, ctx.peer_addr, ctx.state) { Ok(()) => NativeResponse::status_row(seq, "BEGIN"), Err(e) => { From a3b69852c96837c2523c06a736390834766a4688 Mon Sep 17 00:00:00 2001 From: Pizo Date: Thu, 16 Jul 2026 20:28:36 +0800 Subject: [PATCH 2/2] fix(control): thread session txn_id through gateway local dispatch Native SQL and direct-op reads inside an explicit transaction route through Gateway::execute, which previously hard-coded txn_id: None into DispatchRouteParams and used the txn-less local SPSC path. The Data Plane staging overlay was never resolved, so in-transaction SELECTs could not see their own staged writes. - Add QueryContext.txn_id - Pass it into DispatchRouteParams for local and remote routes - dispatch_local uses dispatch_to_data_plane_with_txn - Native sql_gateway/direct_ops/streaming stamp the session task txn_id - Autocommit and system callers keep None (compile-driven field fill) No wire version bump; ExecuteRequest.txn_id already exists on main. Guarded by native_in_txn_sql_select_reads_own_write and native_first_frame_begin_buffers_following_write. --- nodedb-cluster-tests/tests/gateway_execute.rs | 1 + .../tests/http_gateway_migration.rs | 1 + .../tests/ilp_gateway_migration.rs | 1 + .../tests/listeners_gateway_smoke.rs | 1 + .../tests/listeners_typed_not_leader.rs | 1 + .../tests/native_gateway_migration.rs | 1 + .../tests/pgwire_gateway_migration.rs | 1 + .../tests/resp_gateway_migration.rs | 1 + nodedb/src/control/gateway/core.rs | 13 +- nodedb/src/control/gateway/dispatcher.rs | 16 +- nodedb/src/control/planner/calvin/preexec.rs | 1 + nodedb/src/control/scatter_gather.rs | 1 + nodedb/src/control/server/exchange/gather.rs | 1 + .../server/exchange/resolve/exchange.rs | 1 + .../server/http/routes/promql/remote.rs | 1 + .../server/http/routes/query/materialized.rs | 1 + .../server/http/routes/query/ndjson.rs | 1 + .../server/http/routes/query_stream.rs | 1 + .../server/http/routes/ws_rpc/execute_sql.rs | 1 + nodedb/src/control/server/ilp_batch.rs | 1 + .../server/native/dispatch/direct_ops.rs | 1 + .../server/native/dispatch/sql_gateway.rs | 3 + .../server/native/dispatch/streaming.rs | 1 + .../handler/routing/gateway_dispatch.rs | 1 + .../pgwire/handler/routing/streaming.rs | 1 + .../control/server/resp/gateway_dispatch.rs | 2 + nodedb/src/event/cdc/consume.rs | 1 + nodedb/src/event/topic/publish.rs | 1 + nodedb/tests/native_gateway_txn_overlay.rs | 215 ++++++++++++++++++ 29 files changed, 263 insertions(+), 10 deletions(-) create mode 100644 nodedb/tests/native_gateway_txn_overlay.rs diff --git a/nodedb-cluster-tests/tests/gateway_execute.rs b/nodedb-cluster-tests/tests/gateway_execute.rs index af85dfafd..bb257cdf5 100644 --- a/nodedb-cluster-tests/tests/gateway_execute.rs +++ b/nodedb-cluster-tests/tests/gateway_execute.rs @@ -32,6 +32,7 @@ fn test_ctx() -> QueryContext { tenant_id: TenantId::new(0), trace_id: nodedb_types::TraceId::ZERO, database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, } } diff --git a/nodedb-cluster-tests/tests/http_gateway_migration.rs b/nodedb-cluster-tests/tests/http_gateway_migration.rs index 0bcaa8a81..61ad7829e 100644 --- a/nodedb-cluster-tests/tests/http_gateway_migration.rs +++ b/nodedb-cluster-tests/tests/http_gateway_migration.rs @@ -30,6 +30,7 @@ fn test_ctx() -> QueryContext { tenant_id: TenantId::new(0), trace_id: nodedb_types::TraceId::ZERO, database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, } } diff --git a/nodedb-cluster-tests/tests/ilp_gateway_migration.rs b/nodedb-cluster-tests/tests/ilp_gateway_migration.rs index 4f95c42d4..bd7195569 100644 --- a/nodedb-cluster-tests/tests/ilp_gateway_migration.rs +++ b/nodedb-cluster-tests/tests/ilp_gateway_migration.rs @@ -28,6 +28,7 @@ fn test_ctx() -> QueryContext { tenant_id: TenantId::new(1), trace_id: nodedb_types::TraceId::ZERO, database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, } } diff --git a/nodedb-cluster-tests/tests/listeners_gateway_smoke.rs b/nodedb-cluster-tests/tests/listeners_gateway_smoke.rs index 1870daed4..81bc53810 100644 --- a/nodedb-cluster-tests/tests/listeners_gateway_smoke.rs +++ b/nodedb-cluster-tests/tests/listeners_gateway_smoke.rs @@ -36,6 +36,7 @@ fn test_ctx(_trace_id: u64) -> QueryContext { tenant_id: TenantId::new(0), trace_id: nodedb_types::TraceId::ZERO, database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, } } diff --git a/nodedb-cluster-tests/tests/listeners_typed_not_leader.rs b/nodedb-cluster-tests/tests/listeners_typed_not_leader.rs index 44811609b..8cc5fb65f 100644 --- a/nodedb-cluster-tests/tests/listeners_typed_not_leader.rs +++ b/nodedb-cluster-tests/tests/listeners_typed_not_leader.rs @@ -73,6 +73,7 @@ fn test_ctx() -> QueryContext { tenant_id: TenantId::new(0), trace_id: nodedb_types::TraceId::ZERO, database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, } } diff --git a/nodedb-cluster-tests/tests/native_gateway_migration.rs b/nodedb-cluster-tests/tests/native_gateway_migration.rs index 59a3b90fa..3b21ecbb7 100644 --- a/nodedb-cluster-tests/tests/native_gateway_migration.rs +++ b/nodedb-cluster-tests/tests/native_gateway_migration.rs @@ -29,6 +29,7 @@ fn test_ctx() -> QueryContext { tenant_id: TenantId::new(0), trace_id: nodedb_types::TraceId::ZERO, database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, } } diff --git a/nodedb-cluster-tests/tests/pgwire_gateway_migration.rs b/nodedb-cluster-tests/tests/pgwire_gateway_migration.rs index c76deda8f..29707993e 100644 --- a/nodedb-cluster-tests/tests/pgwire_gateway_migration.rs +++ b/nodedb-cluster-tests/tests/pgwire_gateway_migration.rs @@ -33,6 +33,7 @@ fn test_ctx() -> QueryContext { tenant_id: TenantId::new(0), trace_id: nodedb_types::TraceId::ZERO, database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, } } diff --git a/nodedb-cluster-tests/tests/resp_gateway_migration.rs b/nodedb-cluster-tests/tests/resp_gateway_migration.rs index 171b21dbb..eccaed9da 100644 --- a/nodedb-cluster-tests/tests/resp_gateway_migration.rs +++ b/nodedb-cluster-tests/tests/resp_gateway_migration.rs @@ -26,6 +26,7 @@ fn test_ctx() -> QueryContext { tenant_id: TenantId::new(0), trace_id: nodedb_types::TraceId::ZERO, database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, } } diff --git a/nodedb/src/control/gateway/core.rs b/nodedb/src/control/gateway/core.rs index 9ec0bc2a1..13cb2906b 100644 --- a/nodedb/src/control/gateway/core.rs +++ b/nodedb/src/control/gateway/core.rs @@ -27,7 +27,7 @@ use tracing::{Instrument, debug, info_span}; use crate::Error; use crate::control::state::SharedState; use crate::control::trace_export::EmitSpanParams; -use crate::types::{DatabaseId, Lsn, TenantId, TraceId, VShardId}; +use crate::types::{DatabaseId, Lsn, TenantId, TraceId, TxnId, VShardId}; use nodedb_physical::physical_plan::PhysicalPlan; use super::dispatcher::{DispatchRouteParams, default_deadline_ms, dispatch_route}; @@ -48,6 +48,11 @@ pub struct QueryContext { /// catalog lookups. Single-database deployments pass /// [`DatabaseId::DEFAULT`]. pub database_id: DatabaseId, + /// Session-transaction id for resolving the per-transaction staging + /// overlay (read-your-own-writes) on local SPSC dispatch and for + /// forwarding on remote `ExecuteRequest`. `None` for autocommit and + /// non-interactive callers. + pub txn_id: Option, } /// The gateway: routes, dispatches, retries, and caches physical plans. @@ -309,6 +314,7 @@ impl Gateway { let tenant_id = ctx.tenant_id; let database_id = ctx.database_id; let trace_id = ctx.trace_id; + let txn_id = ctx.txn_id; let version_set = version_set_for_route.clone(); async move { let decision = { @@ -351,10 +357,7 @@ impl Gateway { trace_id, deadline_ms, version_set: &version_set, - // The gateway does not yet carry session-transaction - // context across this boundary (known cross-node - // in-transaction read gap), so `None`. - txn_id: None, + txn_id, }) .await } diff --git a/nodedb/src/control/gateway/dispatcher.rs b/nodedb/src/control/gateway/dispatcher.rs index 2f0771da3..c7636d76b 100644 --- a/nodedb/src/control/gateway/dispatcher.rs +++ b/nodedb/src/control/gateway/dispatcher.rs @@ -21,7 +21,7 @@ use std::time::Duration; use nodedb_cluster::rpc_codec::TypedClusterError; use crate::Error; -use crate::control::server::dispatch_utils::dispatch_to_data_plane; +use crate::control::server::dispatch_utils::dispatch_to_data_plane_with_txn; use crate::control::server::result_stream::ResultStream; use crate::control::state::SharedState; use crate::types::{DatabaseId, Lsn, TenantId, TraceId, TxnId, VShardId}; @@ -57,8 +57,9 @@ pub struct DispatchOutcome { /// `trace_id` — distributed trace ID propagated from the client request. /// `deadline_ms` — remaining deadline in milliseconds. /// `version_set` — descriptor versions for the collections touched by the plan. -/// `txn_id` — session-transaction context to forward to the remote executor, or -/// `None` for non-transactional dispatch (the common case). +/// `txn_id` — session-transaction context for local overlay resolution and +/// remote `ExecuteRequest` forwarding, or `None` for non-transactional +/// dispatch (the common case). pub struct DispatchRouteParams<'a> { pub route: TaskRoute, pub shared: &'a Arc, @@ -84,7 +85,7 @@ pub async fn dispatch_route(params: DispatchRouteParams<'_>) -> Result { - dispatch_local(route, shared, tenant_id, database_id, trace_id).await + dispatch_local(route, shared, tenant_id, database_id, trace_id, txn_id).await } RouteDecision::Remote { node_id, vshard_id } => { dispatch_remote(RemoteDispatchArgs { @@ -197,21 +198,26 @@ pub async fn dispatch_route_stream( } /// Local dispatch via SPSC bridge. +/// +/// Carries `txn_id` so the Data Plane can resolve this session transaction's +/// staging overlay (read-your-own-writes) for in-block SQL and direct ops. async fn dispatch_local( route: TaskRoute, shared: &Arc, tenant_id: TenantId, database_id: DatabaseId, trace_id: TraceId, + txn_id: Option, ) -> Result { let vshard_id = VShardId::new(route.vshard_id); - let resp = dispatch_to_data_plane( + let resp = dispatch_to_data_plane_with_txn( shared, tenant_id, database_id, vshard_id, route.plan, trace_id, + txn_id, ) .await?; Ok(DispatchOutcome { diff --git a/nodedb/src/control/planner/calvin/preexec.rs b/nodedb/src/control/planner/calvin/preexec.rs index 159ffbaf4..ed6a1e6dc 100644 --- a/nodedb/src/control/planner/calvin/preexec.rs +++ b/nodedb/src/control/planner/calvin/preexec.rs @@ -126,6 +126,7 @@ pub async fn run_preexec_scan( tenant_id, trace_id: TraceId::ZERO, database_id, + txn_id: None, }; let payloads = gateway diff --git a/nodedb/src/control/scatter_gather.rs b/nodedb/src/control/scatter_gather.rs index a0f3dd928..c0b8d9081 100644 --- a/nodedb/src/control/scatter_gather.rs +++ b/nodedb/src/control/scatter_gather.rs @@ -345,6 +345,7 @@ pub async fn coordinate_cross_shard_hop( tenant_id: crate::types::TenantId::new(tenant_id_u64), trace_id: TraceId::generate(), database_id, + txn_id: None, }; // Build a fresh QueryContext per traversal using cloned inputs diff --git a/nodedb/src/control/server/exchange/gather.rs b/nodedb/src/control/server/exchange/gather.rs index 8836649fc..225b5fa8f 100644 --- a/nodedb/src/control/server/exchange/gather.rs +++ b/nodedb/src/control/server/exchange/gather.rs @@ -385,6 +385,7 @@ pub async fn gather_all_vshards( tenant_id, trace_id, database_id, + txn_id, }; // `Box::pin` breaks an async-fn recursion cycle: the gateway dispatches the diff --git a/nodedb/src/control/server/exchange/resolve/exchange.rs b/nodedb/src/control/server/exchange/resolve/exchange.rs index be8a18ec2..01e3b4b46 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange.rs @@ -169,6 +169,7 @@ async fn resolve_exchange( tenant_id, trace_id, database_id, + txn_id: None, }; // NOTE: cluster mode does not yet thread `txn_id` through // `gateway.execute_stream` — cross-node in-transaction diff --git a/nodedb/src/control/server/http/routes/promql/remote.rs b/nodedb/src/control/server/http/routes/promql/remote.rs index 0e4d500d9..007146491 100644 --- a/nodedb/src/control/server/http/routes/promql/remote.rs +++ b/nodedb/src/control/server/http/routes/promql/remote.rs @@ -90,6 +90,7 @@ pub async fn remote_write( tenant_id, trace_id: TraceId::generate(), database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, }; gw.execute(&gw_ctx, plan).await } diff --git a/nodedb/src/control/server/http/routes/query/materialized.rs b/nodedb/src/control/server/http/routes/query/materialized.rs index aa61ffef0..d7436fda3 100644 --- a/nodedb/src/control/server/http/routes/query/materialized.rs +++ b/nodedb/src/control/server/http/routes/query/materialized.rs @@ -293,6 +293,7 @@ pub async fn query( tenant_id: task.tenant_id, trace_id, database_id, + txn_id: None, }; gw.execute(&gw_ctx, task.plan) .await diff --git a/nodedb/src/control/server/http/routes/query/ndjson.rs b/nodedb/src/control/server/http/routes/query/ndjson.rs index 377ba9061..d354efb15 100644 --- a/nodedb/src/control/server/http/routes/query/ndjson.rs +++ b/nodedb/src/control/server/http/routes/query/ndjson.rs @@ -161,6 +161,7 @@ pub async fn query_ndjson( tenant_id: task.tenant_id, trace_id, database_id, + txn_id: None, }; gw.execute(&gw_ctx, task.plan).await } diff --git a/nodedb/src/control/server/http/routes/query_stream.rs b/nodedb/src/control/server/http/routes/query_stream.rs index 775b82990..d66e6f516 100644 --- a/nodedb/src/control/server/http/routes/query_stream.rs +++ b/nodedb/src/control/server/http/routes/query_stream.rs @@ -61,6 +61,7 @@ pub(super) async fn try_open_stream( tenant_id: task.tenant_id, trace_id, database_id, + txn_id: None, }; gw.execute_stream(&ctx, child_plan).await } else { diff --git a/nodedb/src/control/server/http/routes/ws_rpc/execute_sql.rs b/nodedb/src/control/server/http/routes/ws_rpc/execute_sql.rs index 37c5c40e5..f4e5959ca 100644 --- a/nodedb/src/control/server/http/routes/ws_rpc/execute_sql.rs +++ b/nodedb/src/control/server/http/routes/ws_rpc/execute_sql.rs @@ -228,6 +228,7 @@ pub async fn execute_sql( tenant_id: task.tenant_id, trace_id, database_id: task.database_id, + txn_id: None, }; gw.execute(&gw_ctx, task.plan).await } diff --git a/nodedb/src/control/server/ilp_batch.rs b/nodedb/src/control/server/ilp_batch.rs index 28ab409c0..fd5268348 100644 --- a/nodedb/src/control/server/ilp_batch.rs +++ b/nodedb/src/control/server/ilp_batch.rs @@ -151,6 +151,7 @@ async fn flush_ilp_batch_inner( tenant_id, trace_id: TraceId::generate(), database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, }; gw.execute(&gw_ctx, plan) .await diff --git a/nodedb/src/control/server/native/dispatch/direct_ops.rs b/nodedb/src/control/server/native/dispatch/direct_ops.rs index abd2d3f6c..29a668a13 100644 --- a/nodedb/src/control/server/native/dispatch/direct_ops.rs +++ b/nodedb/src/control/server/native/dispatch/direct_ops.rs @@ -525,6 +525,7 @@ async fn dispatch_single_task_raw( tenant_id, trace_id: TraceId::generate(), database_id: ctx.database_id(), + txn_id, }; match gw.execute(&gw_ctx, plan).await { Ok(payloads) => Ok(gateway_payloads_to_response(payloads)), diff --git a/nodedb/src/control/server/native/dispatch/sql_gateway.rs b/nodedb/src/control/server/native/dispatch/sql_gateway.rs index 61d7c7e02..abc527c26 100644 --- a/nodedb/src/control/server/native/dispatch/sql_gateway.rs +++ b/nodedb/src/control/server/native/dispatch/sql_gateway.rs @@ -39,6 +39,9 @@ pub(super) async fn dispatch_task_via_gateway( tenant_id, trace_id: TraceId::generate(), database_id, + // Propagate the in-block transaction id so gateway local + // dispatch resolves the per-txn staging overlay. + txn_id, }; gw.execute(&gw_ctx, plan) .await diff --git a/nodedb/src/control/server/native/dispatch/streaming.rs b/nodedb/src/control/server/native/dispatch/streaming.rs index 7e12d2e0d..c4a8a4829 100644 --- a/nodedb/src/control/server/native/dispatch/streaming.rs +++ b/nodedb/src/control/server/native/dispatch/streaming.rs @@ -106,6 +106,7 @@ pub(crate) async fn try_open_sql_stream( tenant_id: task.tenant_id, trace_id: crate::types::TraceId::ZERO, database_id, + txn_id: task.txn_id, }; gw.execute_stream(&gw_ctx, child_plan).await? } else { diff --git a/nodedb/src/control/server/pgwire/handler/routing/gateway_dispatch.rs b/nodedb/src/control/server/pgwire/handler/routing/gateway_dispatch.rs index 35197b92f..751209279 100644 --- a/nodedb/src/control/server/pgwire/handler/routing/gateway_dispatch.rs +++ b/nodedb/src/control/server/pgwire/handler/routing/gateway_dispatch.rs @@ -103,6 +103,7 @@ impl NodeDbPgHandler { tenant_id, trace_id: TraceId::generate(), database_id, + txn_id: None, }; let mut responses: Vec = Vec::with_capacity(tasks.len()); diff --git a/nodedb/src/control/server/pgwire/handler/routing/streaming.rs b/nodedb/src/control/server/pgwire/handler/routing/streaming.rs index f1592faee..fe0ed953e 100644 --- a/nodedb/src/control/server/pgwire/handler/routing/streaming.rs +++ b/nodedb/src/control/server/pgwire/handler/routing/streaming.rs @@ -82,6 +82,7 @@ impl NodeDbPgHandler { tenant_id: task.tenant_id, trace_id: crate::types::TraceId::ZERO, database_id: task.database_id, + txn_id: None, }; gw.execute_stream(&ctx, child_plan).await } else { diff --git a/nodedb/src/control/server/resp/gateway_dispatch.rs b/nodedb/src/control/server/resp/gateway_dispatch.rs index 54fa03963..f922d7ed4 100644 --- a/nodedb/src/control/server/resp/gateway_dispatch.rs +++ b/nodedb/src/control/server/resp/gateway_dispatch.rs @@ -37,6 +37,7 @@ pub(super) async fn dispatch_kv( tenant_id: session.tenant_id, trace_id: TraceId::generate(), database_id: DatabaseId::DEFAULT, + txn_id: None, }; gw.execute(&gw_ctx, plan) .await @@ -81,6 +82,7 @@ pub(super) async fn dispatch_kv_write( tenant_id: session.tenant_id, trace_id: TraceId::generate(), database_id: DatabaseId::DEFAULT, + txn_id: None, }; gw.execute(&gw_ctx, plan) .await diff --git a/nodedb/src/event/cdc/consume.rs b/nodedb/src/event/cdc/consume.rs index 3980621c2..5f4226f66 100644 --- a/nodedb/src/event/cdc/consume.rs +++ b/nodedb/src/event/cdc/consume.rs @@ -250,6 +250,7 @@ pub async fn consume_remote( tenant_id: crate::types::TenantId::new(tenant_id), trace_id: nodedb_types::TraceId::generate(), database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, }; let query_ctx = crate::control::planner::context::QueryContext::for_state(state); diff --git a/nodedb/src/event/topic/publish.rs b/nodedb/src/event/topic/publish.rs index a2d582a1b..c7565472f 100644 --- a/nodedb/src/event/topic/publish.rs +++ b/nodedb/src/event/topic/publish.rs @@ -162,6 +162,7 @@ pub async fn publish_remote( tenant_id: crate::types::TenantId::new(tenant_id), trace_id: nodedb_types::TraceId::generate(), database_id: nodedb_types::id::DatabaseId::DEFAULT, + txn_id: None, }; let query_ctx = crate::control::planner::context::QueryContext::for_state(state); diff --git a/nodedb/tests/native_gateway_txn_overlay.rs b/nodedb/tests/native_gateway_txn_overlay.rs new file mode 100644 index 000000000..855ca5159 --- /dev/null +++ b/nodedb/tests/native_gateway_txn_overlay.rs @@ -0,0 +1,215 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Native gateway transaction-overlay regressions: +//! +//! 1. `OpCode::Begin` as the first post-handshake frame must create a session +//! and enter `InBlock`, so the following write is staged (not autocommitted). +//! 2. Planned SQL reads inside an explicit transaction must resolve the +//! session staging overlay when dispatch goes through the gateway (local +//! SPSC with `txn_id`), matching the direct-op path already covered by +//! `native_direct_op_txn_overlay.rs`. + +mod common; + +use common::native_harness::{NativeTestServer, do_handshake, send_request, send_sql}; +use common::pgwire_harness::TestServer; + +use nodedb_types::protocol::opcodes::ResponseStatus; +use nodedb_types::protocol::text_fields::TextFields; +use nodedb_types::protocol::{HelloFrame, OpCode}; +use nodedb_types::value::Value; +use tokio::net::TcpStream; + +/// Open a native session against the server's native listener. +async fn native_session(srv: &TestServer) -> TcpStream { + let addr = format!("127.0.0.1:{}", srv.native_port) + .parse() + .expect("native addr"); + let (stream, _ack) = do_handshake(addr, &HelloFrame::current()) + .await + .expect("native handshake"); + stream +} + +/// `OpCode::Begin` with empty fields — first application frame after handshake. +async fn send_begin_opcode( + stream: &mut TcpStream, + seq: u64, +) -> nodedb_types::protocol::NativeResponse { + send_request(stream, seq, OpCode::Begin, TextFields::default()).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn native_first_frame_begin_buffers_following_write() { + let server = NativeTestServer::start().await; + let (mut stream, _ack) = do_handshake(server.addr, &HelloFrame::current()) + .await + .expect("handshake"); + + // First application frame: binary/protocol BEGIN, not SQL text. + // Pre-fix, `handle_begin` called `run_begin` without `ensure_session`, + // so SessionStore::begin silently no-oped and the following INSERT + // autocommitted outside a transaction block. + let begin = send_begin_opcode(&mut stream, 1).await; + assert_ne!( + begin.status, + ResponseStatus::Error, + "first-frame OpCode::Begin must succeed: {begin:?}" + ); + + let create = send_sql( + &mut stream, + 2, + "CREATE COLLECTION native_ff_begin (id STRING PRIMARY KEY, n INT) \ + WITH (engine='document_schemaless')", + ) + .await; + assert_ne!( + create.status, + ResponseStatus::Error, + "CREATE must succeed: {create:?}" + ); + + let insert = send_sql( + &mut stream, + 3, + "INSERT INTO native_ff_begin (id, n) VALUES ('a', 1)", + ) + .await; + assert_ne!( + insert.status, + ResponseStatus::Error, + "in-block INSERT must succeed: {insert:?}" + ); + + // Mid-transaction SELECT must see the staged row. + let select = send_sql( + &mut stream, + 4, + "SELECT n FROM native_ff_begin WHERE id = 'a'", + ) + .await; + assert_eq!( + select.status, + ResponseStatus::Ok, + "in-txn SELECT must succeed: {select:?}" + ); + let rows = select + .rows + .expect("in-txn SELECT must return rows for the staged write"); + assert_eq!( + rows.first().and_then(|r| r.first()).cloned(), + Some(Value::Integer(1)), + "first-frame BEGIN must leave the session InBlock so the write is staged \ + and visible to the same-connection SELECT" + ); + + let rollback = send_sql(&mut stream, 5, "ROLLBACK").await; + assert_ne!( + rollback.status, + ResponseStatus::Error, + "ROLLBACK must succeed: {rollback:?}" + ); + + // After ROLLBACK the staged row must not be durable. + let after = send_sql( + &mut stream, + 6, + "SELECT n FROM native_ff_begin WHERE id = 'a'", + ) + .await; + server.shutdown().await; + assert_eq!( + after.status, + ResponseStatus::Ok, + "post-rollback SELECT: {after:?}" + ); + let after_rows = after.rows.unwrap_or_default(); + assert!( + after_rows.is_empty() + || after_rows + .iter() + .all(|r| r.iter().all(|v| matches!(v, Value::Null))), + "ROLLBACK must discard the staged write, got: {after_rows:?}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn native_in_txn_sql_select_reads_own_write() { + let server = TestServer::start().await; + + server + .exec( + "CREATE COLLECTION native_gw_ryow (id STRING PRIMARY KEY, n INT) \ + WITH (engine='document_schemaless')", + ) + .await + .unwrap(); + + let mut stream = native_session(&server).await; + let mut seq = 1u64; + + let begin = send_sql(&mut stream, seq, "BEGIN").await; + assert_eq!(begin.status, ResponseStatus::Ok, "BEGIN: {begin:?}"); + + seq += 1; + let insert = send_sql( + &mut stream, + seq, + "INSERT INTO native_gw_ryow (id, n) VALUES ('x', 42)", + ) + .await; + assert_eq!( + insert.status, + ResponseStatus::Ok, + "in-tx INSERT: {insert:?}" + ); + + // Gateway-routed planned SQL SELECT must resolve the staging overlay. + seq += 1; + let select = send_sql( + &mut stream, + seq, + "SELECT n FROM native_gw_ryow WHERE id = 'x'", + ) + .await; + assert_eq!( + select.status, + ResponseStatus::Ok, + "in-tx SELECT: {select:?}" + ); + let rows = select + .rows + .expect("in-tx SELECT must return the staged row"); + assert_eq!( + rows.first().and_then(|r| r.first()).cloned(), + Some(Value::Integer(42)), + "native SQL SELECT inside a transaction must see its own staged write \ + when dispatch goes through the gateway" + ); + + seq += 1; + let rollback = send_sql(&mut stream, seq, "ROLLBACK").await; + assert_eq!( + rollback.status, + ResponseStatus::Ok, + "ROLLBACK: {rollback:?}" + ); + + seq += 1; + let after = send_sql( + &mut stream, + seq, + "SELECT n FROM native_gw_ryow WHERE id = 'x'", + ) + .await; + assert_eq!(after.status, ResponseStatus::Ok); + let after_rows = after.rows.unwrap_or_default(); + assert!( + after_rows.is_empty() + || after_rows + .iter() + .all(|r| r.iter().all(|v| matches!(v, Value::Null))), + "ROLLBACK must hide the staged write, got: {after_rows:?}" + ); +}