Ensure we handle pipelined cross-shard queries (#1161)

When mulitple queries are received before a Sync message, we need to
treat them as inside an implicit transaction. Before this commit, we'd
see that the first statement routes to a single shard, open a
direct-to-shard connection to that shard, and forward the message. This
would result in the implicit transaction on the shard's end, purely
because it's receiving the messages in the same order we do. But we were
completely unaware of the transaction-like behavior on our end, unless
the client was explicitly starting a transaction.

With this change, we now correctly recognize that we're in a transaction
any time we receive multiple executable messages prior to a sync, and
behave the same as we would if the first query we received in the
pipeline were `BEGIN`

Co-authored-by: Lev Kokotov <lev.kokotov@gmail.com>

Author: sgrif

integration/go/go_pgx/batch_sharded_test.go

@@ -0,0 +1,86 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func countOnShardByComment(t *testing.T, conn *pgx.Conn, shard int, id int64) int {
+	t.Helper()
+
+	var count int
+	err := conn.QueryRow(context.Background(),
+		fmt.Sprintf("/* pgdog_shard: %d */ SELECT COUNT(*) FROM sharded WHERE id = $1", shard), id).Scan(&count)
+	require.NoError(t, err)
+	return count
+}
+
+func TestBatchRoutesShardedTable(t *testing.T) {
+	ctx := context.Background()
+
+	proxy, err := pgx.Connect(ctx, testConnStr)
+	require.NoError(t, err)
+	defer proxy.Close(ctx)
+
+	_, err = proxy.Exec(ctx, "TRUNCATE TABLE sharded")
+	require.NoError(t, err)
+
+	const shard0ID int64 = 1
+	const shard1ID int64 = 11
+
+	batch := &pgx.Batch{}
+	batch.Queue("INSERT INTO sharded (id, value) VALUES ($1, $2)", shard0ID, "shard_0")
+	batch.Queue("INSERT INTO sharded (id, value) VALUES ($1, $2)", shard1ID, "shard_1")
+
+	results := proxy.SendBatch(ctx, batch)
+	for i := range 2 {
+		tag, err := results.Exec()
+		require.NoError(t, err, "batch insert %d", i)
+		assert.EqualValues(t, 1, tag.RowsAffected(), "batch insert %d rows affected", i)
+	}
+	require.NoError(t, results.Close())
+
+	assert.Equal(t, 1, countOnShardByComment(t, proxy, 0, shard0ID), "id=%d must be on shard 0", shard0ID)
+	assert.Equal(t, 0, countOnShardByComment(t, proxy, 1, shard0ID), "id=%d must not be on shard 1", shard0ID)
+	assert.Equal(t, 0, countOnShardByComment(t, proxy, 0, shard1ID), "id=%d must not be on shard 0", shard1ID)
+	assert.Equal(t, 1, countOnShardByComment(t, proxy, 1, shard1ID), "id=%d must be on shard 1", shard1ID)
+}
+
+func TestBatchRollsBackOnError(t *testing.T) {
+	ctx := context.Background()
+
+	proxy, err := pgx.Connect(ctx, testConnStr)
+	require.NoError(t, err)
+	defer proxy.Close(ctx)
+
+	_, err = proxy.Exec(ctx, "TRUNCATE TABLE sharded")
+	require.NoError(t, err)
+
+	const shard0ID int64 = 1
+	const shard1ID int64 = 11
+
+	batch := &pgx.Batch{}
+	batch.Queue("INSERT INTO sharded (id, value) VALUES ($1, $2)", shard0ID, "shard_0")
+	batch.Queue("INSERT INTO sharded (id, value) VALUES ($1, $2)", shard1ID, "shard_1")
+	batch.Queue("INSERT INTO sharded (id, value) VALUES ($1, $2)", shard0ID, "shard_0")
+
+	results := proxy.SendBatch(ctx, batch)
+	tag, err := results.Exec()
+	require.NoError(t, err, "batch insert 1")
+	assert.EqualValues(t, 1, tag.RowsAffected(), "batch insert 1 rows affected")
+	tag, err = results.Exec()
+	require.NoError(t, err, "batch insert 2")
+	assert.EqualValues(t, 1, tag.RowsAffected(), "batch insert 2 rows affected")
+
+	_, err = results.Exec()
+	require.ErrorContains(t, err, "duplicate key value violates unique constraint")
+	require.Error(t, results.Close())
+
+	assert.Equal(t, 0, countOnShardByComment(t, proxy, 0, shard0ID), "transaction must have rolled back")
+	assert.Equal(t, 0, countOnShardByComment(t, proxy, 1, shard0ID), "transaction must have rolled back")
+}

pgdog/src/frontend/client/mod.rs

@@ -551,6 +551,7 @@ impl Client {
         } else {
             let total = spliced.len();
             let mut reqs = spliced.into_iter().enumerate();
+            self.transaction.get_or_insert(TransactionType::Implicit);
             while let Some((num, mut req)) = reqs.next() {
                 debug!("processing spliced request {}/{}", num + 1, total);
                 let mut context = QueryEngineContext::new(self).spliced(&mut req, reqs.len());

pgdog/src/frontend/client/query_engine/query.rs

@@ -172,7 +172,9 @@ impl QueryEngine {
                 TransactionState::Error => {
                     let error_state = match context.transaction {
                         Some(TransactionType::ReadOnly) => Some(TransactionType::ErrorReadOnly),
-                        Some(TransactionType::ReadWrite) => Some(TransactionType::ErrorReadWrite),
+                        Some(TransactionType::ReadWrite | TransactionType::Implicit) => {
+                            Some(TransactionType::ErrorReadWrite)
+                        }
                         _ => None,
                     };
                     context.transaction = error_state;

pgdog/src/frontend/client/transaction_type.rs

@@ -3,6 +3,7 @@ pub enum TransactionType {
     ReadOnly,
     #[default]
     ReadWrite,
+    Implicit,
     ErrorReadWrite,
     ErrorReadOnly,
 }

pgdog/src/frontend/router/parser/context.rs

@@ -81,7 +81,7 @@ impl<'a> QueryParserContext<'a> {
         let role = self.router_context.parameter_hints.compute_role();
         let txn_write = matches!(
             self.router_context.transaction(),
-            Some(TransactionType::ReadWrite)
+            Some(TransactionType::ReadWrite | TransactionType::Implicit)
         ) && self.rw_conservative();
         // prefer_primary defaults reads to the primary; an explicit replica hint opts out.
         txn_write