Merge pull request #1852 from dolthub/zachmu/procedures

Alter stored procedure execution to deal with statements that commit transactions

Author: zachmu

sql/rowexec/other.go

@@ -250,7 +250,13 @@ func (b *BaseBuilder) buildBlock(ctx *sql.Context, n *plan.Block, row sql.Row) (
 
 	selectSeen := false
 	for _, s := range n.Children() {
-		err := func() error {
+		// TODO: this should happen at iteration time, but this call is where the actual iteration happens
+		err := startTransaction(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		err = func() error {
 			rowCache, disposeFunc := ctx.Memory.NewRowsCache()
 			defer disposeFunc()
 

sql/rowexec/proc.go

@@ -89,6 +89,12 @@ func (b *BaseBuilder) buildIfElseBlock(ctx *sql.Context, n *plan.IfElseBlock, ro
 			continue
 		}
 
+		// TODO: this should happen at iteration time, but this call is where the actual iteration happens
+		err = startTransaction(ctx)
+		if err != nil {
+			return nil, err
+		}
+
 		branchIter, err = b.buildNodeExec(ctx, ifConditional, row)
 		if err != nil {
 			return nil, err
@@ -105,6 +111,12 @@ func (b *BaseBuilder) buildIfElseBlock(ctx *sql.Context, n *plan.IfElseBlock, ro
 		}, nil
 	}
 
+	// TODO: this should happen at iteration time, but this call is where the actual iteration happens
+	err = startTransaction(ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	// All conditions failed so we run the else
 	branchIter, err = b.buildNodeExec(ctx, n.Else, row)
 	if err != nil {

sql/rowexec/proc_iters.go

@@ -39,6 +39,10 @@ var _ plan.BlockRowIter = (*ifElseIter)(nil)
 
 // Next implements the sql.RowIter interface.
 func (i *ifElseIter) Next(ctx *sql.Context) (sql.Row, error) {
+	if err := startTransaction(ctx); err != nil {
+		return nil, err
+	}
+
 	return i.branchIter.Next(ctx)
 }
 
@@ -67,6 +71,10 @@ var _ sql.RowIter = (*beginEndIter)(nil)
 
 // Next implements the interface sql.RowIter.
 func (b *beginEndIter) Next(ctx *sql.Context) (sql.Row, error) {
+	if err := startTransaction(ctx); err != nil {
+		return nil, err
+	}
+
 	row, err := b.rowIter.Next(ctx)
 	if err != nil {
 		if exitErr, ok := err.(expression.ProcedureBlockExitError); ok && b.Pref.CurrentHeight() == int(exitErr) {
@@ -278,6 +286,10 @@ func (l *loopIter) Next(ctx *sql.Context) (sql.Row, error) {
 			}
 		}
 
+		if err := startTransaction(ctx); err != nil {
+			return nil, err
+		}
+
 		nextRow, err := l.blockIter.Next(ctx)
 		if err != nil {
 			restart := false
@@ -395,3 +407,21 @@ func (i *iterateIter) Next(ctx *sql.Context) (sql.Row, error) {
 func (i *iterateIter) Close(ctx *sql.Context) error {
 	return nil
 }
+
+// startTransaction begins a new transaction if necessary, e.g. if a statement in a stored procedure committed the
+// current one
+func startTransaction(ctx *sql.Context) error {
+	if ctx.GetTransaction() == nil {
+		ts, ok := ctx.Session.(sql.TransactionSession)
+		if ok {
+			tx, err := ts.StartTransaction(ctx, sql.ReadWrite)
+			if err != nil {
+				return err
+			}
+
+			ctx.SetTransaction(tx)
+		}
+	}
+
+	return nil
+}