support DECLARE in BEGIN...END BLOCK in TRIGGER

enginetest/queries/trigger_queries.go

@@ -2206,18 +2206,244 @@ INSERT INTO t0 (v1, v2) VALUES (i, s); END;`,
 	},
 
 	{
-		Name: "delete me",
+		Name: "triggers with nested begin-end blocks",
 		SetUpScript: []string{
 			"create table t (i int primary key);",
+			`
+    create trigger trig
+    before insert on t
+    for each row
+    begin
+      declare x int;
+      set x = new.i * 10;
+      begin
+        declare y int;
+        set y = new.i + 10;
+        set new.i = x + y;
+      end;
+    end;
+    `,
 		},
 		Assertions: []ScriptTestAssertion{
 			{
-				Query:       "create trigger trig1 before insert on t for each row begin declare x int; select new.i + 10 into x; set new.i = x; end;",
-				ExpectedErr: sql.ErrUnsupportedFeature,
+				Query: "insert into t values (1), (2), (3);",
+				Expected: []sql.Row{
+					{types.NewOkResult(3)},
+				},
 			},
 			{
-				Query:       "create trigger trig2 before insert on t for each row begin declare x int; set x = new.i * 10; set new.i = x; end;",
-				ExpectedErr: sql.ErrUnsupportedFeature,
+				Query: "select * from t;",
+				Expected: []sql.Row{
+					{21},
+					{32},
+					{43},
+				},
+			},
+		},
+	},
+	{
+		Name: "triggers with declare statements and select into",
+		SetUpScript: []string{
+			"create table t (i int primary key);",
+			"create trigger trig before insert on t for each row begin declare x int; select new.i + 10 into x; set new.i = x; end;",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "insert into t values (1), (2), (3);",
+				Expected: []sql.Row{
+					{types.NewOkResult(3)},
+				},
+			},
+			{
+				Query: "select * from t;",
+				Expected: []sql.Row{
+					{11},
+					{12},
+					{13},
+				},
+			},
+		},
+	},
+	{
+		Name: "triggers with declare statements and set",
+		SetUpScript: []string{
+			"create table t (i int primary key);",
+			"create trigger trig before insert on t for each row begin declare x int; set x = new.i + 10; set new.i = x; end;",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "insert into t values (1), (2), (3);",
+				Expected: []sql.Row{
+					{types.NewOkResult(3)},
+				},
+			},
+			{
+				Query: "select * from t;",
+				Expected: []sql.Row{
+					{11},
+					{12},
+					{13},
+				},
+			},
+		},
+	},
+	{
+		Name: "triggers with declare statements and insert",
+		SetUpScript: []string{
+			"create table t (i int primary key);",
+			"create table t2 (i int primary key);",
+			`
+create trigger trig before
+insert on t for each row begin
+	declare x int;
+	set x = new.i * 10;
+	insert into t2 values (x);
+end;
+`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "insert into t values (1), (2), (3);",
+				Expected: []sql.Row{
+					{types.NewOkResult(3)},
+				},
+			},
+			{
+				Query: "select * from t;",
+				Expected: []sql.Row{
+					{1},
+					{2},
+					{3},
+				},
+			},
+			{
+				Query: "select * from t2;",
+				Expected: []sql.Row{
+					{10},
+					{20},
+					{30},
+				},
+			},
+		},
+	},
+	{
+		Name: "triggers with declare statements and update",
+		SetUpScript: []string{
+			"create table t (i int primary key);",
+			"create table t2 (i int primary key);",
+			"insert into t2 values (1), (2), (3);",
+			`
+create trigger trig before
+insert on t for each row begin
+	declare x int;
+	set x = new.i * 10;
+	update t2 set i = x where i = new.i;
+end;
+`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "insert into t values (1), (2), (3);",
+				Expected: []sql.Row{
+					{types.NewOkResult(3)},
+				},
+			},
+			{
+				Query: "select * from t;",
+				Expected: []sql.Row{
+					{1},
+					{2},
+					{3},
+				},
+			},
+			{
+				Query: "select * from t2;",
+				Expected: []sql.Row{
+					{10},
+					{20},
+					{30},
+				},
+			},
+		},
+	},
+	{
+		Name: "triggers with declare statements and delete",
+		SetUpScript: []string{
+			"create table t (i int primary key);",
+			"create table t2 (i int primary key);",
+			"insert into t2 values (1), (2), (3);",
+			`
+create trigger trig before
+insert on t for each row begin
+	declare x int;
+	set x = new.i;
+	delete from t2 where i = x;
+end;
+`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "insert into t values (1), (2), (3);",
+				Expected: []sql.Row{
+					{types.NewOkResult(3)},
+				},
+			},
+			{
+				Query: "select * from t;",
+				Expected: []sql.Row{
+					{1},
+					{2},
+					{3},
+				},
+			},
+			{
+				Query:    "select * from t2;",
+				Expected: []sql.Row{},
+			},
+		},
+	},
+	{
+		Name: "triggers with declare statements and stored procedure",
+		SetUpScript: []string{
+			"create table t (i int primary key);",
+			"create table t2 (i int primary key);",
+			`
+create procedure proc(in i int)
+begin
+	insert into t2 values (i);
+end;
+`,
+			`
+create trigger trig before
+insert on t for each row begin
+	declare x int;
+	set x = new.i + 10;
+	call proc(x);
+end;
+`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "insert into t values (1), (2), (3);",
+				Expected: []sql.Row{
+					{types.NewOkResult(3)},
+				},
+			},
+			{
+				Query: "select * from t;",
+				Expected: []sql.Row{
+					{1},
+					{2},
+					{3},
+				},
+			},
+			{
+				Query: "select * from t2;",
+				Expected: []sql.Row{
+					{11},
+					{12},
+					{13},
+				},
 			},
 		},
 	},

sql/analyzer/stored_procedures.go

@@ -445,6 +445,7 @@ func applyProceduresCall(ctx *sql.Context, a *Analyzer, call *plan.Call, scope *
 			}
 			return n, transform.SameTree, nil
 		case expression.ProcedureReferencable:
+			// BeginEndBlocks need to reference the same ParameterReference as the Call
 			return n.WithParamReference(pRef), transform.NewTree, nil
 		default:
 			return transform.NodeExprsWithOpaque(n, procParamTransformFunc)

sql/analyzer/triggers.go

@@ -249,6 +249,60 @@ func applyTrigger(ctx *sql.Context, a *Analyzer, originalNode, n sql.Node, scope
 		return nil, transform.SameTree, err
 	}
 
+	if _, ok := triggerLogic.(*plan.TriggerBeginEndBlock); ok {
+		pRef := expression.NewProcedureReference()
+		// assignProcParam transforms any ProcedureParams to reference the ProcedureReference
+		assignProcParam := func(expr sql.Expression) (sql.Expression, transform.TreeIdentity, error) {
+			switch e := expr.(type) {
+			case *expression.ProcedureParam:
+				return e.WithParamReference(pRef), transform.NewTree, nil
+			default:
+				return expr, transform.SameTree, nil
+			}
+		}
+		// assignProcRef calls assignProcParam on all nodes that sql.Expressioner
+		assignProcRef := func(node sql.Node) (sql.Node, transform.TreeIdentity, error) {
+			switch n := node.(type) {
+			case sql.Expressioner:
+				newExprs, same, err := transform.Exprs(n.Expressions(), assignProcParam)
+				if err != nil {
+					return nil, transform.SameTree, err
+				}
+				if same {
+					return node, transform.SameTree, nil
+				}
+				newNode, err := n.WithExpressions(newExprs...)
+				if err != nil {
+					return nil, transform.SameTree, err
+				}
+				return newNode, transform.NewTree, nil
+			default:
+				return node, transform.SameTree, nil
+			}
+		}
+		assignProcs := func(node sql.Node) (sql.Node, transform.TreeIdentity, error) {
+			switch n := node.(type) {
+			case *plan.InsertInto:
+				newSource, same, err := transform.NodeWithOpaque(n.Source, assignProcRef)
+				if err != nil {
+					return nil, transform.SameTree, err
+				}
+				if same {
+					return node, transform.SameTree, nil
+				}
+				return n.WithSource(newSource), transform.NewTree, nil
+			case expression.ProcedureReferencable:
+				return n.WithParamReference(pRef), transform.NewTree, nil
+			default:
+				return assignProcRef(node)
+			}
+		}
+		triggerLogic, _, err = transform.NodeWithOpaque(triggerLogic, assignProcs)
+		if err != nil {
+			return nil, transform.SameTree, err
+		}
+	}
+
 	return transform.NodeWithCtx(n, nil, func(c transform.Context) (sql.Node, transform.TreeIdentity, error) {
 		// Don't double-apply trigger executors to the bodies of triggers. To avoid this, don't apply the trigger if the
 		// parent is a trigger body.