get rid of nested switch in apply_foreign_keys

Author: angelamayxie

sql/analyzer/apply_foreign_keys.go

@@ -122,22 +122,14 @@ func applyForeignKeysToNodes(ctx *sql.Context, a *Analyzer, n sql.Node, cache *f
 		if plan.IsEmptyTable(n.Child) {
 			return n, transform.SameTree, nil
 		}
-		updateDest, err := plan.GetUpdatable(n.Child)
-		if err != nil {
-			return nil, transform.SameTree, err
-		}
-		switch updateDest.(type) {
-		case *plan.UpdatableJoinTable:
-			updateTargets := updateDest.(*plan.UpdatableJoinTable).UpdateTargets
+		if n.IsJoin {
+			uj := n.Child.(*plan.UpdateJoin)
+			updateTargets := uj.UpdateTargets
 			fkHandlerMap := make(map[string]sql.Node, len(updateTargets))
 			for tableName, updateTarget := range updateTargets {
 				fkHandlerMap[tableName] = updateTarget
-				updateDest, err := plan.GetUpdatable(updateTarget)
-				if err != nil {
-					return nil, transform.SameTree, err
-				}
 				fkHandler, err :=
-					getForeignKeyHandlerFromUpdateDestination(updateDest, ctx, a, cache, fkChain, updateTarget)
+					getForeignKeyHandlerFromUpdateTarget(updateTarget, ctx, a, cache, fkChain)
 				if err != nil {
 					return nil, transform.SameTree, err
 				}
@@ -147,20 +139,19 @@ func applyForeignKeysToNodes(ctx *sql.Context, a *Analyzer, n sql.Node, cache *f
 					fkHandlerMap[tableName] = fkHandler
 				}
 			}
-			uj := plan.NewUpdateJoin(fkHandlerMap, n.Child.(*plan.UpdateJoin).Child)
+			uj = plan.NewUpdateJoin(fkHandlerMap, uj.Child)
 			nn, err := n.WithChildren(uj)
 			return nn, transform.NewTree, err
-		default:
-			fkHandler, err := getForeignKeyHandlerFromUpdateDestination(updateDest, ctx, a, cache, fkChain, n.Child)
-			if err != nil {
-				return nil, transform.SameTree, err
-			}
-			if fkHandler == nil {
-				return n, transform.SameTree, nil
-			}
-			nn, err := n.WithChildren(fkHandler)
-			return nn, transform.NewTree, err
 		}
+		fkHandler, err := getForeignKeyHandlerFromUpdateTarget(n.Child, ctx, a, cache, fkChain)
+		if err != nil {
+			return nil, transform.SameTree, err
+		}
+		if fkHandler == nil {
+			return n, transform.SameTree, nil
+		}
+		nn, err := n.WithChildren(fkHandler)
+		return nn, transform.NewTree, err
 	case *plan.DeleteFrom:
 		if plan.IsEmptyTable(n.Child) {
 			return n, transform.SameTree, nil
@@ -457,10 +448,14 @@ func getForeignKeyRefActions(ctx *sql.Context, a *Analyzer, tbl sql.ForeignKeyTa
 	return fkEditor, nil
 }
 
-// getForeignKeyHandlerFromUpdateDestination creates a ForeignKeyHandler from a given UpdatableTable. It's used in
-// applying foreign keys to Update nodes
-func getForeignKeyHandlerFromUpdateDestination(updateDest sql.UpdatableTable, ctx *sql.Context, a *Analyzer,
-	cache *foreignKeyCache, fkChain foreignKeyChain, originalNode sql.Node) (*plan.ForeignKeyHandler, error) {
+// getForeignKeyHandlerFromUpdateTarget creates a ForeignKeyHandler from a given update target Node. It is used for
+// applying foreign key constrains to Update nodes
+func getForeignKeyHandlerFromUpdateTarget(updateTarget sql.Node, ctx *sql.Context, a *Analyzer,
+	cache *foreignKeyCache, fkChain foreignKeyChain) (*plan.ForeignKeyHandler, error) {
+	updateDest, err := plan.GetUpdatable(updateTarget)
+	if err != nil {
+		return nil, err
+	}
 	fkTbl, ok := updateDest.(sql.ForeignKeyTable)
 	if !ok {
 		return nil, nil
@@ -477,7 +472,7 @@ func getForeignKeyHandlerFromUpdateDestination(updateDest sql.UpdatableTable, ct
 	return &plan.ForeignKeyHandler{
 		Table:        fkTbl,
 		Sch:          updateDest.Schema(),
-		OriginalNode: originalNode,
+		OriginalNode: updateTarget,
 		Editor:       fkEditor,
 		AllUpdaters:  fkChain.GetUpdaters(),
 	}, nil

sql/plan/update_join.go

@@ -21,14 +21,14 @@ import (
 )
 
 type UpdateJoin struct {
-	updateTargets map[string]sql.Node
+	UpdateTargets map[string]sql.Node
 	UnaryNode
 }
 
 // NewUpdateJoin returns a new *UpdateJoin node.
 func NewUpdateJoin(updateTargets map[string]sql.Node, child sql.Node) *UpdateJoin {
 	return &UpdateJoin{
-		updateTargets: updateTargets,
+		UpdateTargets: updateTargets,
 		UnaryNode:     UnaryNode{Child: child},
 	}
 }
@@ -55,7 +55,7 @@ func (u *UpdateJoin) DebugString() string {
 // GetUpdatable returns an updateJoinTable which implements sql.UpdatableTable.
 func (u *UpdateJoin) GetUpdatable() sql.UpdatableTable {
 	return &UpdatableJoinTable{
-		UpdateTargets: u.updateTargets,
+		UpdateTargets: u.UpdateTargets,
 		joinNode:      u.Child.(*UpdateSource).Child,
 	}
 }
@@ -66,7 +66,7 @@ func (u *UpdateJoin) WithChildren(children ...sql.Node) (sql.Node, error) {
 		return nil, sql.ErrInvalidChildrenNumber.New(u, len(children), 1)
 	}
 
-	return NewUpdateJoin(u.updateTargets, children[0]), nil
+	return NewUpdateJoin(u.UpdateTargets, children[0]), nil
 }
 
 func (u *UpdateJoin) IsReadOnly() bool {
@@ -79,7 +79,7 @@ func (u *UpdateJoin) CollationCoercibility(ctx *sql.Context) (collation sql.Coll
 }
 
 func (u *UpdateJoin) GetUpdaters(ctx *sql.Context) (map[string]sql.RowUpdater, error) {
-	return getUpdaters(u.updateTargets, ctx)
+	return getUpdaters(u.UpdateTargets, ctx)
 }
 
 func getUpdaters(updateTargets map[string]sql.Node, ctx *sql.Context) (map[string]sql.RowUpdater, error) {