From 4f357b957a25567d8b494a321977f8ce2a6a2cbe Mon Sep 17 00:00:00 2001 From: James Cor Date: Thu, 1 May 2025 12:51:16 -0700 Subject: [PATCH 1/3] fix prepare statements and user vars in stored procedures --- engine.go | 9 +---- enginetest/queries/procedure_queries.go | 52 +++++++++++++++++++++++++ sql/plan/prepare.go | 15 +++++-- sql/planbuilder/transactions.go | 2 +- 4 files changed, 65 insertions(+), 13 deletions(-) diff --git a/engine.go b/engine.go index bad99e4a8f..fea5d95cec 100644 --- a/engine.go +++ b/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) diff --git a/enginetest/queries/procedure_queries.go b/enginetest/queries/procedure_queries.go index 7ca990a260..7aef0ada88 100644 --- a/enginetest/queries/procedure_queries.go +++ b/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{ diff --git a/sql/plan/prepare.go b/sql/plan/prepare.go index 2bc11aca17..b1c472af8b 100644 --- a/sql/plan/prepare.go +++ b/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. diff --git a/sql/planbuilder/transactions.go b/sql/planbuilder/transactions.go index ab6e9019f7..135519c791 100644 --- a/sql/planbuilder/transactions.go +++ b/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 } From cc7be914a2a5e32950678838512f3168e5232d19 Mon Sep 17 00:00:00 2001 From: James Cor Date: Thu, 1 May 2025 13:10:11 -0700 Subject: [PATCH 2/3] skip results on server engine --- enginetest/queries/procedure_queries.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/enginetest/queries/procedure_queries.go b/enginetest/queries/procedure_queries.go index 7aef0ada88..8d3ae9a71e 100644 --- a/enginetest/queries/procedure_queries.go +++ b/enginetest/queries/procedure_queries.go @@ -2231,7 +2231,8 @@ end; }, Assertions: []ScriptTestAssertion{ { - Query: "call create_proc();", + SkipResultCheckOnServerEngine: true, + Query: "call create_proc();", Expected: []sql.Row{ {types.NewOkResult(0)}, }, From 74c2f1b45cabe957b4aafac1acec56f3dfee721b Mon Sep 17 00:00:00 2001 From: James Cor Date: Fri, 2 May 2025 11:10:51 -0700 Subject: [PATCH 3/3] remove comment --- engine.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/engine.go b/engine.go index fea5d95cec..e85fb1e8dd 100644 --- a/engine.go +++ b/engine.go @@ -587,12 +587,6 @@ func (e *Engine) analyzeNode(ctx *sql.Context, query string, bound sql.Node, qFl switch n := bound.(type) { case *plan.PrepareQuery: sqlMode := sql.LoadSqlMode(ctx) - - // we have to name-resolve to check for structural errors, but we do - // not to cache the name-bound query yet. - // 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 prepare := n.PrepStmt cacheStmt, _, err := e.Parser.ParseOneWithOptions(ctx, prepare.Expr, sqlMode.ParserOptions()) if err != nil && strings.HasPrefix(prepare.Expr, "@") {