Fix SET and START TRANSACTION in create procedure statements (vitessio#18279)

Signed-off-by: Manan Gupta <[email protected]>

Author: GuptaManan100

go/test/endtoend/vtgate/unsharded/main_test.go

@@ -151,7 +151,10 @@ BEGIN
 	insert into allDefaults(id) values (128);
 	select 128 into val from dual;
 END;
-`}
+`,
+		`CREATE PROCEDURE p1 (in x BIGINT) BEGIN declare y DECIMAL(14,2); set y = 4.2; END`,
+		`CREATE PROCEDURE p2 (in x BIGINT) BEGIN START TRANSACTION; SELECT 128 from dual; COMMIT; END`,
+	}
 )
 
 func TestMain(m *testing.M) {

go/vt/sqlparser/ast.go

@@ -640,8 +640,15 @@ type (
 	// TxAccessMode is an enum for Transaction Access Mode
 	TxAccessMode int8
 
+	// BeginType is an enum for the type of BEGIN statement.
+	BeginType int8
+
 	// Begin represents a Begin statement.
 	Begin struct {
+		// We need to differentiate between BEGIN and START TRANSACTION statements
+		// because inside a stored procedure the former is considered part of a BEGIN...END statement,
+		// while the latter starts a transaction.
+		Type          BeginType
 		TxAccessModes []TxAccessMode
 	}
 

go/vt/sqlparser/ast_equals.go

@@ -1188,7 +1188,7 @@ func (node *Commit) Format(buf *TrackedBuffer) {
 
 // Format formats the node.
 func (node *Begin) Format(buf *TrackedBuffer) {
-	if node.TxAccessModes == nil {
+	if node.Type == BeginStmt {
 		buf.literal("begin")
 		return
 	}

go/vt/sqlparser/ast_format.go

@@ -997,6 +997,12 @@ const (
 	ReadOnly
 )
 
+// BEGIN statement type
+const (
+	BeginStmt BeginType = iota
+	StartTransactionStmt
+)
+
 // Enum Types of WKT functions
 const (
 	GeometryFromText GeomFromWktType = iota

go/vt/sqlparser/ast_format_fast.go

@@ -174,7 +174,7 @@ func (nz *normalizer) determineQueryRewriteStrategy(in Statement) {
 func (nz *normalizer) walkDown(node, _ SQLNode) bool {
 	switch node := node.(type) {
 	case *Begin, *Commit, *Rollback, *Savepoint, *SRollback, *Release, *OtherAdmin, *Analyze,
-		*PrepareStmt, *ExecuteStmt, *FramePoint, *ColName, TableName, *ConvertType:
+		*PrepareStmt, *ExecuteStmt, *FramePoint, *ColName, TableName, *ConvertType, *CreateProcedure:
 		// These statement do not need normalizing
 		return false
 	case *AssignmentExpr:

go/vt/sqlparser/cached_size.go

@@ -441,6 +441,11 @@ func TestNormalize(t *testing.T) {
 			"bv1": sqltypes.Int64BindVariable(1),
 			"bv2": sqltypes.Int64BindVariable(0),
 		},
+	}, {
+		// Verify we don't change anything in the normalization of create procedures.
+		in:      "CREATE PROCEDURE p2 (in x BIGINT) BEGIN declare y DECIMAL(14,2); START TRANSACTION; set y = 4.2; SELECT 128 from dual; COMMIT; END",
+		outstmt: "create procedure p2 (in x BIGINT) begin declare y DECIMAL(14,2); start transaction; set y = 4.2; select 128 from dual; commit; end;",
+		outbv:   map[string]*querypb.BindVariable{},
 	}}
 	parser := NewTestParser()
 	for _, tc := range testcases {

go/vt/sqlparser/constants.go

@@ -2226,6 +2226,8 @@ var (
 	}, {
 		input:  "create procedure ConditionWithSignalAndHandler() begin declare custom_error condition for sqlstate '45000'; declare exit handler for custom_error begin select 'Handled with custom condition and signal'; end; signal sqlstate '45000' set message_text = 'Custom signal triggered'; end;",
 		output: "create procedure ConditionWithSignalAndHandler () begin declare custom_error condition for sqlstate '45000'; declare exit handler for custom_error begin select 'Handled with custom condition and signal' from dual; end; signal sqlstate '45000' set message_text = 'Custom signal triggered'; end;",
+	}, {
+		input: "create procedure t1 (in x BIGINT) begin start transaction; insert into unsharded_a values (1, 'a', 'a'); commit; end;",
 	}, {
 		input: "create /*vt+ strategy=online */ or replace view v as select a, b, c from t",
 	}, {
@@ -2907,8 +2909,7 @@ var (
 		input:  "begin;",
 		output: "begin",
 	}, {
-		input:  "start transaction",
-		output: "begin",
+		input: "start transaction",
 	}, {
 		input: "start transaction with consistent snapshot",
 	}, {