Merge pull request #109692 from rafiss/backport23.1-109618

release-23.1: tree: handle format of START/BEGIN inside of functions
cockroachdb · Aug 29, 2023 · e6e624c · e6e624c
2 parents ad5ca6e + 75fda63
commit e6e624c
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/pkg/sql/parser/sql.y b/pkg/sql/parser/sql.y
@@ -865,6 +865,9 @@ func (u *sqlSymUnion) showTenantOpts() tree.ShowTenantOptions {
 func (u *sqlSymUnion) showCreateFormatOption() tree.ShowCreateFormatOption {
     return u.val.(tree.ShowCreateFormatOption)
 }
+func (u *sqlSymUnion) beginTransaction() *tree.BeginTransaction {
+    return u.val.(*tree.BeginTransaction)
+}
 %}
 
 // NB: the %token definitions must come before the %type definitions in this
@@ -11482,7 +11485,9 @@ transaction_stmt:
 begin_stmt:
   START TRANSACTION begin_transaction
   {
-    $$.val = $3.stmt()
+    s := $3.beginTransaction()
+    s.FormatWithStart = true
+    $$.val = s
   }
 | START error // SHOW HELP: BEGIN
 

diff --git a/pkg/sql/parser/testdata/begin_commit b/pkg/sql/parser/testdata/begin_commit
@@ -17,10 +17,10 @@ BEGIN TRANSACTION -- identifiers removed
 parse
 START TRANSACTION
 ----
-BEGIN TRANSACTION -- normalized!
-BEGIN TRANSACTION -- fully parenthesized
-BEGIN TRANSACTION -- literals removed
-BEGIN TRANSACTION -- identifiers removed
+START TRANSACTION
+START TRANSACTION -- fully parenthesized
+START TRANSACTION -- literals removed
+START TRANSACTION -- identifiers removed
 
 parse
 BEGIN TRANSACTION READ ONLY

diff --git a/pkg/sql/sem/tree/txn.go b/pkg/sql/sem/tree/txn.go
@@ -221,12 +221,20 @@ func (node *TransactionModes) Merge(other TransactionModes) error {
 
 // BeginTransaction represents a BEGIN statement
 type BeginTransaction struct {
-	Modes TransactionModes
+	// FormatWithStart says whether this statement must be formatted with
+	// "START" rather than "BEGIN". This is needed if this statement is in a
+	// BEGIN ATOMIC block of a procedure or function.
+	FormatWithStart bool
+	Modes           TransactionModes
 }
 
 // Format implements the NodeFormatter interface.
 func (node *BeginTransaction) Format(ctx *FmtCtx) {
-	ctx.WriteString("BEGIN TRANSACTION")
+	if node.FormatWithStart {
+		ctx.WriteString("START TRANSACTION")
+	} else {
+		ctx.WriteString("BEGIN TRANSACTION")
+	}
 	ctx.FormatNode(&node.Modes)
 }