moved rowupdateaccumulator

Author: James Cor

sql/analyzer/apply_update_accumulators.go

@@ -26,7 +26,7 @@ import (
 // applyUpdateAccumulators wraps any Insert, Update, or Delete nodes with RowUpdateAccumulators to tally the results
 // for report to the client.
 func applyUpdateAccumulators(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope, sel RuleSelector, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
-	//return n, transform.SameTree, nil
+	return n, transform.SameTree, nil
 	switch n := n.(type) {
 	case *plan.TriggerExecutor, *plan.InsertInto, *plan.DeleteFrom, *plan.Update:
 		accumulatorType, err := getUpdateAccumulatorType(n)

sql/rowexec/dml_iters.go

@@ -516,17 +516,18 @@ type accumulatorIter struct {
 	updateRowHandler accumulatorRowHandler
 }
 
-
-func getRowHandler(ctx *sql.Context, clientFoundRowsToggled bool, iter sql.RowIter) accumulatorRowHandler {
+func getRowHandler(clientFoundRowsToggled bool, iter sql.RowIter) accumulatorRowHandler {
 	switch i := iter.(type) {
 	case *plan.TableEditorIter:
-		return getRowHandler(ctx, clientFoundRowsToggled, i.InnerIter())
+		return getRowHandler(clientFoundRowsToggled, i.InnerIter())
 	case *plan.CheckpointingTableEditorIter:
-		return getRowHandler(ctx, clientFoundRowsToggled, i.InnerIter())
+		return getRowHandler(clientFoundRowsToggled, i.InnerIter())
 	case *ProjectIter:
-		return getRowHandler(ctx, clientFoundRowsToggled, i.childIter)
+		return getRowHandler(clientFoundRowsToggled, i.childIter)
 	case *triggerIter:
-		return getRowHandler(ctx, clientFoundRowsToggled, i.child)
+		return getRowHandler(clientFoundRowsToggled, i.child)
+	case *blockIter:
+		return getRowHandler(clientFoundRowsToggled, i.repIter)
 	case *insertIter:
 		if i.replacer != nil {
 			return &replaceRowHandler{}
@@ -538,14 +539,25 @@ func getRowHandler(ctx *sql.Context, clientFoundRowsToggled bool, iter sql.RowIt
 	case *deleteIter:
 		return &deleteRowHandler{}
 	case *updateIter:
-		if updateJoin, ok := i.childIter.(*updateJoinIter); ok {
-			return &updateJoinRowHandler{
-				joinSchema: updateJoin.joinSchema,
-				tableMap: plan.RecreateTableSchemaFromJoinSchema(updateJoin.joinSchema),
-				updaterMap: updateJoin.updaters,
-			}
+		// it's possible that there's an updateJoinIter that's not the immediate child of updateIter
+		rowHandler := getRowHandler(clientFoundRowsToggled, i.childIter)
+		if rowHandler != nil {
+			return rowHandler
+		}
+		sch := i.schema
+		// special case for foreign keys, plan.ForeignKeyHandler.Schema() returns original schema
+		if fkHandler, isFk := i.updater.(*plan.ForeignKeyHandler); isFk {
+			sch = fkHandler.Sch
 		}
-		return &updateRowHandler{schema: i.schema[:len(i.schema)/2], clientFoundRowsCapability: clientFoundRowsToggled}
+		return &updateRowHandler{schema: sch, clientFoundRowsCapability: clientFoundRowsToggled}
+	case *updateJoinIter:
+		rowHandler := &updateJoinRowHandler{
+			joinSchema: i.joinSchema,
+			tableMap: plan.RecreateTableSchemaFromJoinSchema(i.joinSchema),
+			updaterMap: i.updaters,
+		}
+		i.accumulator = rowHandler
+		return rowHandler
 	default:
 		return nil
 	}
@@ -561,13 +573,9 @@ func AddAccumulatorIter(ctx *sql.Context, node sql.Node, iter sql.RowIter) (sql.
 		childIter, sch, err := AddAccumulatorIter(ctx, node, i.rowIter)
 		i.rowIter = childIter
 		return i, sch, err
-	case *blockIter:
-		childIter, sch, err := AddAccumulatorIter(ctx, node, i.internalIter)
-		i.internalIter = childIter
-		return i, sch, err
 	default:
-		clientFoundRowsToggled := (ctx.Client().Capabilities & mysql.CapabilityClientFoundRows) == mysql.CapabilityClientFoundRows
-		rowHandler := getRowHandler(ctx, clientFoundRowsToggled, iter)
+		clientFoundRowsToggled := (ctx.Client().Capabilities & mysql.CapabilityClientFoundRows) > 0
+		rowHandler := getRowHandler(clientFoundRowsToggled, iter)
 		if rowHandler == nil {
 			return iter, nil, nil
 		}

sql/rowexec/other.go

@@ -148,7 +148,8 @@ func (b *BaseBuilder) buildCachedResults(ctx *sql.Context, n *plan.CachedResults
 func (b *BaseBuilder) buildBlock(ctx *sql.Context, n *plan.Block, row sql.Row) (sql.RowIter, error) {
 	var returnRows []sql.Row
 	var returnNode sql.Node
-	var returnSch sql.Schema
+	var returnIter sql.RowIter
+	var returnSch  sql.Schema
 
 	selectSeen := false
 	for _, s := range n.Children() {
@@ -216,10 +217,12 @@ func (b *BaseBuilder) buildBlock(ctx *sql.Context, n *plan.Block, row sql.Row) (
 			if isSelect = plan.NodeRepresentsSelect(subIterNode); isSelect {
 				selectSeen = true
 				returnNode = subIterNode
-				returnSch = subIterSch
+				returnIter = subIter
+				returnSch  = subIterSch
 			} else if !selectSeen {
 				returnNode = subIterNode
-				returnSch = types.OkResultSchema
+				returnIter = subIter
+				returnSch  = types.OkResultSchema
 				//returnSch = subIterSch
 			}
 
@@ -259,7 +262,8 @@ func (b *BaseBuilder) buildBlock(ctx *sql.Context, n *plan.Block, row sql.Row) (
 	return &blockIter{
 		internalIter: sql.RowsToRowIter(returnRows...),
 		repNode:      returnNode,
-		sch:          returnSch,
+		repIter:      returnIter,
+		repSch:       returnSch,
 	}, nil
 }
 

sql/rowexec/other_iters.go

@@ -115,7 +115,8 @@ func (itr *dropHistogramIter) Close(_ *sql.Context) error {
 type blockIter struct {
 	internalIter sql.RowIter
 	repNode      sql.Node
-	sch          sql.Schema
+	repIter      sql.RowIter
+	repSch       sql.Schema
 }
 
 var _ plan.BlockRowIter = (*blockIter)(nil)
@@ -137,7 +138,7 @@ func (i *blockIter) RepresentingNode() sql.Node {
 
 // Schema implements the sql.BlockRowIter interface.
 func (i *blockIter) Schema() sql.Schema {
-	return i.sch
+	return i.repSch
 }
 
 type prependRowIter struct {

sql/rowexec/proc.go

@@ -232,7 +232,8 @@ func (b *BaseBuilder) buildLoop(ctx *sql.Context, n *plan.Loop, row sql.Row) (sq
 
 	var returnRows []sql.Row
 	var returnNode sql.Node
-	var returnSch sql.Schema
+	var returnIter sql.RowIter
+	var returnSch  sql.Schema
 	selectSeen := false
 
 	// It's technically valid to make an infinite loop, but we don't want to actually allow that
@@ -285,10 +286,12 @@ func (b *BaseBuilder) buildLoop(ctx *sql.Context, n *plan.Loop, row sql.Row) (sq
 				selectSeen = true
 				includeResultSet = true
 				returnNode = subIterNode
+				returnIter = loopBodyIter
 				returnSch = subIterSch
 			} else if !selectSeen {
 				includeResultSet = true
 				returnNode = subIterNode
+				returnIter = loopBodyIter
 				returnSch = subIterSch
 			}
 		}
@@ -336,7 +339,8 @@ func (b *BaseBuilder) buildLoop(ctx *sql.Context, n *plan.Loop, row sql.Row) (sq
 	return &blockIter{
 		internalIter: sql.RowsToRowIter(returnRows...),
 		repNode:      returnNode,
-		sch:          returnSch,
+		repSch:       returnSch,
+		repIter:      returnIter,
 	}, nil
 }