fix prepare statements and user vars in stored procedures

Author: James Cor

engine.go

@@ -593,14 +593,7 @@ func (e *Engine) analyzeNode(ctx *sql.Context, query string, bound sql.Node, qFl
 		// todo(max): improve name resolution so we can cache post name-binding.
 		// this involves expression memoization, which currently screws up aggregation
 		// and order by aliases
-		prepStmt, _, err := e.Parser.ParseOneWithOptions(ctx, query, sqlMode.ParserOptions())
-		if err != nil {
-			return nil, err
-		}
-		prepare, ok := prepStmt.(*sqlparser.Prepare)
-		if !ok {
-			return nil, fmt.Errorf("expected *sqlparser.Prepare, found %T", prepStmt)
-		}
+		prepare := n.PrepStmt
 		cacheStmt, _, err := e.Parser.ParseOneWithOptions(ctx, prepare.Expr, sqlMode.ParserOptions())
 		if err != nil && strings.HasPrefix(prepare.Expr, "@") {
 			val, err := expression.NewUserVar(strings.TrimPrefix(prepare.Expr, "@")).Eval(ctx, nil)

enginetest/queries/procedure_queries.go

@@ -2192,6 +2192,58 @@ END;`,
 			},
 		},
 	},
+	{
+		Name: "user variables are usable within stored procedures",
+		SetUpScript: []string{
+			`
+create procedure proc()
+begin
+  declare v int default 123;
+  set @v = v;
+end;
+`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "call proc();",
+				Expected: []sql.Row{{}},
+			},
+			{
+				Query: "select @v;",
+				Expected: []sql.Row{
+					{123},
+				},
+			},
+		},
+	},
+	{
+		Name: "prepare statement inside of stored procedures",
+		SetUpScript: []string{
+			`
+create procedure create_proc()
+begin
+  set @stmt = 'create table t (i int)';
+  prepare stmt from @stmt;
+  execute stmt;
+  deallocate prepare stmt;
+end;
+`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "call create_proc();",
+				Expected: []sql.Row{
+					{types.NewOkResult(0)},
+				},
+			},
+			{
+				Query: "insert into t values (1), (2), (3);",
+				Expected: []sql.Row{
+					{types.NewOkResult(3)},
+				},
+			},
+		},
+	},
 }
 
 var ProcedureCallTests = []ScriptTest{

sql/plan/prepare.go

@@ -19,20 +19,27 @@ import (
 
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/types"
+
+	"github.com/dolthub/vitess/go/vt/sqlparser"
 )
 
 // PrepareQuery is a node that prepares the query
 type PrepareQuery struct {
-	Name  string
-	Child sql.Node
+	Name     string
+	Child    sql.Node
+	PrepStmt *sqlparser.Prepare
 }
 
 var _ sql.Node = (*PrepareQuery)(nil)
 var _ sql.CollationCoercible = (*PrepareQuery)(nil)
 
 // NewPrepareQuery creates a new PrepareQuery node.
-func NewPrepareQuery(name string, child sql.Node) *PrepareQuery {
-	return &PrepareQuery{Name: name, Child: child}
+func NewPrepareQuery(name string, child sql.Node, prepStmt *sqlparser.Prepare) *PrepareQuery {
+	return &PrepareQuery{
+		Name:     name,
+		Child:    child,
+		PrepStmt: prepStmt,
+	}
 }
 
 // Schema implements the Node interface.

sql/planbuilder/transactions.go

@@ -70,7 +70,7 @@ func (b *Builder) buildPrepare(inScope *scope, n *ast.Prepare) (outScope *scope)
 	// test for query structure; bind variables will be discarded
 	b.bindCtx = &BindvarContext{resolveOnly: true}
 	childScope := b.build(inScope, childStmt, expr)
-	outScope.node = plan.NewPrepareQuery(n.Name, childScope.node)
+	outScope.node = plan.NewPrepareQuery(n.Name, childScope.node, n)
 	return outScope
 }