Fix TimestampFuncExpr and SetOp in stored procedures (#3319)

Author: jycor

enginetest/queries/procedure_queries.go

@@ -2352,6 +2352,61 @@ end;
 			},
 		},
 	},
+	{
+		Name: "stored procedure with subquery set operations",
+		SetUpScript: []string{
+			`create procedure sub_intersect() select(select 1 intersect select 2);`,
+			`create procedure sub_except() select(select 1 except select 2);`,
+			`
+CREATE PROCEDURE test() 
+	WITH RECURSIVE 
+		data (str) AS (SELECT "test"), 
+		positions (pos, chr) AS (
+			SELECT 
+				1, 
+				LEFT(data.str,1) 
+			FROM 
+				data 
+			UNION ALL 
+			SELECT 
+				pos + 1, 
+				MID(data.str,pos+1,1) 
+			FROM 
+				positions 
+			CROSS JOIN 
+				data 
+			WHERE 
+				pos < LENGTH(data.str)
+		) 
+	SELECT 
+		pos, chr 
+	FROM 
+		positions;`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "call sub_intersect();",
+				Expected: []sql.Row{
+					{nil},
+				},
+			},
+			{
+				Query: "call sub_except();",
+				Expected: []sql.Row{
+					{1},
+				},
+			},
+			{
+				Query: "call test();",
+				Expected: []sql.Row{
+					{1, "t"},
+					{2, "e"},
+					{3, "s"},
+					{4, "t"},
+				},
+			},
+		},
+	},
 }
 
 var ProcedureCallTests = []ScriptTest{
@@ -2589,6 +2644,31 @@ var ProcedureCallTests = []ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "Procedures with TimestampFuncExpr",
+		SetUpScript: []string{
+			`CREATE PROCEDURE after_2000(IN bday DATE) SELECT TIMESTAMPDIFF(YEAR, '2000-01-01', bday);`,
+			`CREATE PROCEDURE add_month(IN bday DATE) SELECT TIMESTAMPADD(MONTH, 1, bday);`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "CALL after_2000('2001-09-11')",
+				Expected: []sql.Row{{1}},
+			},
+			{
+				Query:    "CALL after_2000('2025-01-01')",
+				Expected: []sql.Row{{25}},
+			},
+			{
+				Query:    "CALL add_month('2001-09-11')",
+				Expected: []sql.Row{{"2001-10-11 00:00:00"}},
+			},
+			{
+				Query:    "CALL add_month('2025-01-01')",
+				Expected: []sql.Row{{"2025-02-01 00:00:00"}},
+			},
+		},
+	},
 	{
 		// https://github.com/dolthub/dolt/pull/7947/
 		Name: "Call a procedure that needs subqueries resolved in an if condition",

sql/procedures/interpreter_logic.go

@@ -146,6 +146,17 @@ func replaceVariablesInExpr(ctx *sql.Context, stack *InterpreterStack, expr ast.
 			return nil, err
 		}
 		e.Subquery = newSubquery.(*ast.Subquery)
+	case *ast.TimestampFuncExpr:
+		newExpr1, err := replaceVariablesInExpr(ctx, stack, e.Expr1, asOf)
+		if err != nil {
+			return nil, err
+		}
+		newExpr2, err := replaceVariablesInExpr(ctx, stack, e.Expr2, asOf)
+		if err != nil {
+			return nil, err
+		}
+		e.Expr1 = newExpr1.(ast.Expr)
+		e.Expr2 = newExpr2.(ast.Expr)
 	case *ast.FuncExpr:
 		for i := range e.Exprs {
 			newExpr, err := replaceVariablesInExpr(ctx, stack, e.Exprs[i], asOf)
@@ -266,7 +277,7 @@ func replaceVariablesInExpr(ctx *sql.Context, stack *InterpreterStack, expr ast.
 		if err != nil {
 			return nil, err
 		}
-		e.Select = newExpr.(*ast.Select)
+		e.Select = newExpr.(ast.SelectStatement)
 	case *ast.SetOp:
 		newLeftExpr, err := replaceVariablesInExpr(ctx, stack, e.Left, asOf)
 		if err != nil {