modified UpdateJoin to contain target node

Author: angelamayxie

sql/analyzer/apply_foreign_keys.go

@@ -128,27 +128,33 @@ func applyForeignKeysToNodes(ctx *sql.Context, a *Analyzer, n sql.Node, cache *f
 		if err != nil {
 			return nil, transform.SameTree, err
 		}
-		fkTbl, ok := updateDest.(sql.ForeignKeyTable)
-		// If foreign keys aren't supported then we return
-		if !ok {
-			return n, transform.SameTree, nil
-		}
+		switch updateDest.(type) {
+		case *plan.UpdatableJoinTable:
 
-		fkEditor, err := getForeignKeyEditor(ctx, a, fkTbl, cache, fkChain, false)
-		if err != nil {
-			return nil, transform.SameTree, err
-		}
-		if fkEditor == nil {
 			return n, transform.SameTree, nil
+		default:
+			fkTbl, ok := updateDest.(sql.ForeignKeyTable)
+			// If foreign keys aren't supported then we return
+			if !ok {
+				return n, transform.SameTree, nil
+			}
+
+			fkEditor, err := getForeignKeyEditor(ctx, a, fkTbl, cache, fkChain, false)
+			if err != nil {
+				return nil, transform.SameTree, err
+			}
+			if fkEditor == nil {
+				return n, transform.SameTree, nil
+			}
+			nn, err := n.WithChildren(&plan.ForeignKeyHandler{
+				Table:        fkTbl,
+				Sch:          updateDest.Schema(),
+				OriginalNode: n.Child,
+				Editor:       fkEditor,
+				AllUpdaters:  fkChain.GetUpdaters(),
+			})
+			return nn, transform.NewTree, err
 		}
-		nn, err := n.WithChildren(&plan.ForeignKeyHandler{
-			Table:        fkTbl,
-			Sch:          updateDest.Schema(),
-			OriginalNode: n.Child,
-			Editor:       fkEditor,
-			AllUpdaters:  fkChain.GetUpdaters(),
-		})
-		return nn, transform.NewTree, err
 	case *plan.DeleteFrom:
 		if plan.IsEmptyTable(n.Child) {
 			return n, transform.SameTree, nil

sql/analyzer/assign_update_join.go

@@ -34,12 +34,12 @@ func modifyUpdateExprsForJoin(ctx *sql.Context, a *Analyzer, n sql.Node, scope *
 			return n, transform.SameTree, nil
 		}
 
-		updatables, err := updatablesByTable(us, jn)
+		updateTargets, err := updateTargetsByTable(us, jn)
 		if err != nil {
 			return nil, transform.SameTree, err
 		}
 
-		uj := plan.NewUpdateJoin(updatables, us)
+		uj := plan.NewUpdateJoin(updateTargets, us)
 		ret, err := n.WithChildren(uj)
 		if err != nil {
 			return nil, transform.SameTree, err
@@ -52,11 +52,11 @@ func modifyUpdateExprsForJoin(ctx *sql.Context, a *Analyzer, n sql.Node, scope *
 }
 
 // rowUpdatersByTable maps a set of tables to their RowUpdater objects.
-func updatablesByTable(node sql.Node, ij sql.Node) (map[string]sql.UpdatableTable, error) {
+func updateTargetsByTable(node sql.Node, ij sql.Node) (map[string]sql.Node, error) {
 	namesOfTableToBeUpdated := getTablesToBeUpdated(node)
 	resolvedTables := getTablesByName(ij)
 
-	updatables := make(map[string]sql.UpdatableTable)
+	updateTargets := make(map[string]sql.Node)
 	for tableToBeUpdated, _ := range namesOfTableToBeUpdated {
 		resolvedTable, ok := resolvedTables[tableToBeUpdated]
 		if !ok {
@@ -76,10 +76,10 @@ func updatablesByTable(node sql.Node, ij sql.Node) (map[string]sql.UpdatableTabl
 			return nil, sql.ErrUnsupportedFeature.New("error: keyless tables unsupported for UPDATE JOIN")
 		}
 
-		updatables[tableToBeUpdated] = updatable
+		updateTargets[tableToBeUpdated] = resolvedTable
 	}
 
-	return updatables, nil
+	return updateTargets, nil
 }
 
 // getTablesToBeUpdated takes a node and looks for the tables to modified by a SetField.

sql/plan/update_join.go

@@ -21,15 +21,15 @@ import (
 )
 
 type UpdateJoin struct {
-	Updatables map[string]sql.UpdatableTable
+	UpdateTargets map[string]sql.Node
 	UnaryNode
 }
 
 // NewUpdateJoin returns an *UpdateJoin node.
-func NewUpdateJoin(updatablesMap map[string]sql.UpdatableTable, child sql.Node) *UpdateJoin {
+func NewUpdateJoin(updateTargets map[string]sql.Node, child sql.Node) *UpdateJoin {
 	return &UpdateJoin{
-		Updatables: updatablesMap,
-		UnaryNode:  UnaryNode{Child: child},
+		UpdateTargets: updateTargets,
+		UnaryNode:     UnaryNode{Child: child},
 	}
 }
 
@@ -60,8 +60,8 @@ func (u *UpdateJoin) GetUpdatable() sql.UpdatableTable {
 	//       doesn't implement ForeignKeyTable, UpdateJoin statements don't enforce foreign key checks.
 	//       We should revamp this function so that we can communicate multiple tables being updated.
 	return &UpdatableJoinTable{
-		updatables: u.Updatables,
-		joinNode:   u.Child.(*UpdateSource).Child,
+		UpdateTargets: u.UpdateTargets,
+		joinNode:      u.Child.(*UpdateSource).Child,
 	}
 }
 
@@ -71,7 +71,11 @@ func (u *UpdateJoin) WithChildren(children ...sql.Node) (sql.Node, error) {
 		return nil, sql.ErrInvalidChildrenNumber.New(u, len(children), 1)
 	}
 
-	return NewUpdateJoin(u.Updatables, children[0]), nil
+	return NewUpdateJoin(u.UpdateTargets, children[0]), nil
+}
+
+func (u *UpdateJoin) WithUpdateTargets(updateTargets map[string]sql.Node) *UpdateJoin {
+	return NewUpdateJoin(updateTargets, u.Child)
 }
 
 func (u *UpdateJoin) IsReadOnly() bool {
@@ -83,22 +87,26 @@ func (u *UpdateJoin) CollationCoercibility(ctx *sql.Context) (collation sql.Coll
 	return sql.GetCoercibility(ctx, u.Child)
 }
 
-func (u *UpdateJoin) GetUpdaters(ctx *sql.Context) map[string]sql.RowUpdater {
-	return getUpdaters(u.Updatables, ctx)
+func (u *UpdateJoin) GetUpdaters(ctx *sql.Context) (map[string]sql.RowUpdater, error) {
+	return getUpdaters(u.UpdateTargets, ctx)
 }
 
-func getUpdaters(updatables map[string]sql.UpdatableTable, ctx *sql.Context) map[string]sql.RowUpdater {
+func getUpdaters(updateTargets map[string]sql.Node, ctx *sql.Context) (map[string]sql.RowUpdater, error) {
 	updaterMap := make(map[string]sql.RowUpdater)
-	for tableName, updatable := range updatables {
+	for tableName, updateTarget := range updateTargets {
+		updatable, err := GetUpdatable(updateTarget)
+		if err != nil {
+			return nil, err
+		}
 		updaterMap[tableName] = updatable.Updater(ctx)
 	}
-	return updaterMap
+	return updaterMap, nil
 }
 
 // updatableJoinTable manages the update of multiple tables.
 type UpdatableJoinTable struct {
-	updatables map[string]sql.UpdatableTable
-	joinNode   sql.Node
+	UpdateTargets map[string]sql.Node
+	joinNode      sql.Node
 }
 
 var _ sql.UpdatableTable = (*UpdatableJoinTable)(nil)
@@ -135,8 +143,9 @@ func (u *UpdatableJoinTable) Collation() sql.CollationID {
 
 // Updater implements the sql.UpdatableTable interface.
 func (u *UpdatableJoinTable) Updater(ctx *sql.Context) sql.RowUpdater {
+	updaters, _ := getUpdaters(u.UpdateTargets, ctx)
 	return &updatableJoinUpdater{
-		updaterMap: getUpdaters(u.updatables, ctx),
+		updaterMap: updaters,
 		schemaMap:  RecreateTableSchemaFromJoinSchema(u.joinNode.Schema()),
 		joinSchema: u.joinNode.Schema(),
 	}

sql/rowexec/dml.go

@@ -416,10 +416,14 @@ func (b *BaseBuilder) buildUpdateJoin(ctx *sql.Context, n *plan.UpdateJoin, row
 		return nil, err
 	}
 
+	updaters, err := n.GetUpdaters(ctx)
+	if err != nil {
+		return nil, err
+	}
 	return &updateJoinIter{
 		updateSourceIter: ji,
 		joinSchema:       n.Child.(*plan.UpdateSource).Child.Schema(),
-		updaters:         n.GetUpdaters(ctx),
+		updaters:         updaters,
 		caches:           make(map[string]sql.KeyValueCache),
 		disposals:        make(map[string]sql.DisposeFunc),
 		joinNode:         n.Child.(*plan.UpdateSource).Child,