tree: handle format of START/BEGIN inside of functions

pkg/sql/parser/sql.y

@@ -880,6 +880,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
@@ -11679,7 +11682,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
 

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

pkg/sql/parser/testdata/create_procedure

@@ -208,3 +208,15 @@ CREATE PROCEDURE _()
 	CALLED ON NULL INPUT
 	LANGUAGE SQL
 	AS $$_$$ -- identifiers removed
+
+parse
+CREATE PROCEDURE FAMILY () BEGIN ATOMIC START TRANSACTION; COMMIT; END;
+----
+CREATE PROCEDURE "family"()
+	BEGIN ATOMIC START TRANSACTION; COMMIT TRANSACTION; END -- normalized!
+CREATE PROCEDURE "family"()
+	BEGIN ATOMIC START TRANSACTION; COMMIT TRANSACTION; END -- fully parenthesized
+CREATE PROCEDURE "family"()
+	BEGIN ATOMIC START TRANSACTION; COMMIT TRANSACTION; END -- literals removed
+CREATE PROCEDURE _()
+	BEGIN ATOMIC START TRANSACTION; COMMIT TRANSACTION; END -- identifiers removed

pkg/sql/sem/tree/txn.go

@@ -258,12 +258,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)
 }