From 9a5ead904858db5afe2073370d33f474db3c83ad Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sun, 18 Jun 2023 06:43:17 -0500
Subject: [PATCH] Add TxOptions.BeginQuery to allow overriding the default
 BEGIN query

https://github.com/jackc/pgx/issues/1643
---
 tx.go      |  8 ++++++++
 tx_test.go | 20 ++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/tx.go b/tx.go
index e57142a6..575c17a7 100644
--- a/tx.go
+++ b/tx.go
@@ -44,6 +44,10 @@ type TxOptions struct {
 	IsoLevel       TxIsoLevel
 	AccessMode     TxAccessMode
 	DeferrableMode TxDeferrableMode
+
+	// BeginQuery is the SQL query that will be executed to begin the transaction. This allows using non-standard syntax
+	// such as BEGIN PRIORITY HIGH with CockroachDB. If set this will override the other settings.
+	BeginQuery string
 }
 
 var emptyTxOptions TxOptions
@@ -53,6 +57,10 @@ func (txOptions TxOptions) beginSQL() string {
 		return "begin"
 	}
 
+	if txOptions.BeginQuery != "" {
+		return txOptions.BeginQuery
+	}
+
 	var buf strings.Builder
 	buf.Grow(64) // 64 - maximum length of string with available options
 	buf.WriteString("begin")
diff --git a/tx_test.go b/tx_test.go
index 9c1c70d3..786b7eb7 100644
--- a/tx_test.go
+++ b/tx_test.go
@@ -372,6 +372,26 @@ func TestBeginReadOnly(t *testing.T) {
 	}
 }
 
+func TestBeginTxBeginQuery(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	pgxtest.RunWithQueryExecModes(ctx, t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		tx, err := conn.BeginTx(ctx, pgx.TxOptions{BeginQuery: "begin read only"})
+		require.NoError(t, err)
+		defer tx.Rollback(ctx)
+
+		var readOnly bool
+		conn.QueryRow(ctx, "select current_setting('transaction_read_only')::bool").Scan(&readOnly)
+		require.True(t, readOnly)
+
+		err = tx.Rollback(ctx)
+		require.NoError(t, err)
+	})
+}
+
 func TestTxNestedTransactionCommit(t *testing.T) {
 	t.Parallel()