Add QueryRewriter interface

batch_test.go

@@ -239,6 +239,36 @@ func TestConnSendBatchWithPreparedStatement(t *testing.T) {
 	})
 }
 
+func TestConnSendBatchWithQueryRewriter(t *testing.T) {
+	t.Parallel()
+
+	pgxtest.RunWithQueryExecModes(context.Background(), t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		batch := &pgx.Batch{}
+		batch.Queue("something to be replaced", &testQueryRewriter{sql: "select $1::int", args: []any{1}})
+		batch.Queue("something else to be replaced", &testQueryRewriter{sql: "select $1::text", args: []any{"hello"}})
+		batch.Queue("more to be replaced", &testQueryRewriter{sql: "select $1::int", args: []any{3}})
+
+		br := conn.SendBatch(context.Background(), batch)
+
+		var n int32
+		err := br.QueryRow().Scan(&n)
+		require.NoError(t, err)
+		require.EqualValues(t, 1, n)
+
+		var s string
+		err = br.QueryRow().Scan(&s)
+		require.NoError(t, err)
+		require.Equal(t, "hello", s)
+
+		err = br.QueryRow().Scan(&n)
+		require.NoError(t, err)
+		require.EqualValues(t, 3, n)
+
+		err = br.Close()
+		require.NoError(t, err)
+	})
+}
+
 // https://github.com/jackc/pgx/issues/856
 func TestConnSendBatchWithPreparedStatementAndStatementCacheDisabled(t *testing.T) {
 	t.Parallel()