modified Update to be more similar to Delete, reduces need to UpdateJoin node, need to implement joinUpdater

Author: angelamayxie

enginetest/queries/update_queries.go

@@ -485,7 +485,7 @@ var UpdateScriptTests = []ScriptTest{
 				// TODO: Foreign key constraints are not honored for UDPATE ... JOIN statements
 				Skip:        true,
 				Query:       "UPDATE orders o JOIN customers c ON o.customer_id = c.id SET o.customer_id = 123 where o.customer_id != 1;",
-				ExpectedErr: sql.ErrCheckConstraintViolated,
+				ExpectedErr: sql.ErrForeignKeyChildViolation,
 			},
 			{
 				Query: "SELECT * FROM orders;",

sql/analyzer/apply_foreign_keys.go

@@ -124,31 +124,43 @@ func applyForeignKeysToNodes(ctx *sql.Context, a *Analyzer, n sql.Node, cache *f
 		}
 		// TODO: UPDATE JOIN can update multiple tables. Because updatableJoinTable does not implement
 		//       sql.ForeignKeyTable, we do not currenly support FK checks for UPDATE JOIN statements.
-		updateDest, err := plan.GetUpdatable(n.Child)
-		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
-		}
+		targets := n.GetUpdateTargets()
+		foreignKeyHandlers := make([]sql.Node, len(targets))
+		copy(foreignKeyHandlers, targets)
 
-		fkEditor, err := getForeignKeyEditor(ctx, a, fkTbl, cache, fkChain, false)
-		if err != nil {
-			return nil, transform.SameTree, err
+		for i, node := range targets {
+			updateDest, err := plan.GetUpdatable(node)
+			if err != nil {
+				return nil, transform.SameTree, err
+			}
+
+			tbl, ok := updateDest.(sql.ForeignKeyTable)
+			if !ok {
+				continue
+			}
+			fkEditor, err := getForeignKeyEditor(ctx, a, tbl, cache, fkChain, false)
+			if err != nil {
+				return nil, transform.SameTree, err
+			}
+			if fkEditor == nil {
+				continue
+			}
+			foreignKeyHandlers[i] = &plan.ForeignKeyHandler{
+				Table:        tbl,
+				Sch:          updateDest.Schema(),
+				OriginalNode: targets[i],
+				AllUpdaters:  fkChain.GetUpdaters(),
+			}
 		}
-		if fkEditor == nil {
-			return n, transform.SameTree, nil
+		if n.IsJoin {
+			return n.WithUpdateJoinTargets(foreignKeyHandlers), transform.NewTree, nil
+		} else {
+			newNode, err := n.WithChildren(foreignKeyHandlers...)
+			if err != nil {
+				return nil, transform.SameTree, err
+			}
+			return newNode, transform.NewTree, 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
 	case *plan.DeleteFrom:
 		if plan.IsEmptyTable(n.Child) {
 			return n, transform.SameTree, nil

sql/analyzer/assign_update_join.go

@@ -34,63 +34,55 @@ func modifyUpdateExprsForJoin(ctx *sql.Context, a *Analyzer, n sql.Node, scope *
 			return n, transform.SameTree, nil
 		}
 
-		updaters, err := rowUpdatersByTable(ctx, us, jn)
+		updateJoinTargets, err := getTablesToBeUpdated(us, jn)
 		if err != nil {
 			return nil, transform.SameTree, err
 		}
-
-		uj := plan.NewUpdateJoin(updaters, us)
-		ret, err := n.WithChildren(uj)
-		if err != nil {
-			return nil, transform.SameTree, err
-		}
-
+		ret := n.WithUpdateJoinTargets(updateJoinTargets)
+		ret = ret.WithJoinSchema(jn.Schema())
 		return ret, transform.NewTree, nil
 	}
 
 	return n, transform.SameTree, nil
 }
 
-// rowUpdatersByTable maps a set of tables to their RowUpdater objects.
-func rowUpdatersByTable(ctx *sql.Context, node sql.Node, ij sql.Node) (map[string]sql.RowUpdater, error) {
-	namesOfTableToBeUpdated := getTablesToBeUpdated(node)
-	resolvedTables := getTablesByName(ij)
+func getTablesToBeUpdated(us sql.Node, jn sql.Node) ([]sql.Node, error) {
+	namesOfTablesToBeUpdated := getNamesOfTablesToBeUpdated(us)
+	resolvedTables := getTablesByName(jn)
+	tablesToBeUpdated := make([]sql.Node, len(namesOfTablesToBeUpdated))
 
-	rowUpdatersByTable := make(map[string]sql.RowUpdater)
-	for tableToBeUpdated, _ := range namesOfTableToBeUpdated {
-		resolvedTable, ok := resolvedTables[tableToBeUpdated]
+	for i, tableName := range namesOfTablesToBeUpdated {
+		resolvedTable, ok := resolvedTables[tableName]
 		if !ok {
-			return nil, plan.ErrUpdateForTableNotSupported.New(tableToBeUpdated)
+			return nil, plan.ErrUpdateForTableNotSupported.New(tableName)
 		}
 
 		var table = resolvedTable.UnderlyingTable()
 
-		// If there is no UpdatableTable for a table being updated, error out
 		updatable, ok := table.(sql.UpdatableTable)
 		if !ok && updatable == nil {
-			return nil, plan.ErrUpdateForTableNotSupported.New(tableToBeUpdated)
+			return nil, plan.ErrUpdateForTableNotSupported.New(tableName)
 		}
 
 		keyless := sql.IsKeyless(updatable.Schema())
 		if keyless {
 			return nil, sql.ErrUnsupportedFeature.New("error: keyless tables unsupported for UPDATE JOIN")
 		}
-
-		rowUpdatersByTable[tableToBeUpdated] = updatable.Updater(ctx)
+		tablesToBeUpdated[i] = resolvedTable
 	}
 
-	return rowUpdatersByTable, nil
+	return tablesToBeUpdated, nil
 }
 
-// getTablesToBeUpdated takes a node and looks for the tables to modified by a SetField.
-func getTablesToBeUpdated(node sql.Node) map[string]struct{} {
-	ret := make(map[string]struct{})
+// getNamesOfTablesToBeUpdated takes a node and looks for the tables to modified by a SetField.
+func getNamesOfTablesToBeUpdated(node sql.Node) []string {
+	ret := make([]string, 0)
 
 	transform.InspectExpressions(node, func(e sql.Expression) bool {
 		switch e := e.(type) {
 		case *expression.SetField:
 			gf := e.LeftChild.(*expression.GetField)
-			ret[strings.ToLower(gf.Table())] = struct{}{}
+			ret = append(ret, strings.ToLower(gf.Table()))
 			return false
 		}
 

sql/plan/update.go

@@ -31,11 +31,13 @@ var ErrUpdateUnexpectedSetResult = errors.NewKind("attempted to set field but ex
 // Update is a node for updating rows on tables.
 type Update struct {
 	UnaryNode
-	checks       sql.CheckConstraints
-	Ignore       bool
-	IsJoin       bool
-	HasSingleRel bool
-	IsProcNested bool
+	checks            sql.CheckConstraints
+	Ignore            bool
+	IsJoin            bool
+	updateJoinTargets []sql.Node
+	joinSchema        sql.Schema
+	HasSingleRel      bool
+	IsProcNested      bool
 
 	// Returning is a list of expressions to return after the update operation. This feature is not
 	// supported in MySQL's syntax, but is exposed through PostgreSQL's syntax.
@@ -168,8 +170,17 @@ func (u *Update) Expressions() []sql.Expression {
 	return exprs
 }
 
+func (u *Update) updateJoinTargetsResolved() bool {
+	for _, target := range u.updateJoinTargets {
+		if target.Resolved() == false {
+			return false
+		}
+	}
+	return true
+}
+
 func (u *Update) Resolved() bool {
-	return u.Child.Resolved() &&
+	return u.Child.Resolved() && u.updateJoinTargetsResolved() &&
 		expression.ExpressionsResolved(u.checks.ToExpressions()...) &&
 		expression.ExpressionsResolved(u.Returning...)
 
@@ -192,6 +203,57 @@ func (u Update) WithExpressions(newExprs ...sql.Expression) (sql.Node, error) {
 	return &u, nil
 }
 
+// WithUpdateJoinTargets returns a new Update node instance with the specified |targets| set as the update join targets
+// of the update operation
+func (u *Update) WithUpdateJoinTargets(targets []sql.Node) *Update {
+	ret := *u
+	ret.updateJoinTargets = targets
+	return &ret
+}
+
+// GetUpdateTargets returns the sql.Nodes representing the tables from which rows should be updated
+func (u *Update) GetUpdateTargets() []sql.Node {
+	if u.IsJoin {
+		return u.updateJoinTargets
+	}
+	return []sql.Node{u.Child}
+}
+
+func (u *Update) WithJoinSchema(schema sql.Schema) *Update {
+	ret := *u
+	ret.joinSchema = schema
+	return &ret
+}
+
+func (u *Update) JoinUpdater() sql.RowUpdater {
+	updaters := make([]sql.RowUpdater, len(u.updateJoinTargets))
+	return &joinUpdater{
+		updaters:   updaters,
+		joinSchema: u.joinSchema,
+	}
+}
+
+type joinUpdater struct {
+	updaters   []sql.RowUpdater
+	joinSchema sql.Schema
+}
+
+var _ sql.RowUpdater = (*joinUpdater)(nil)
+
+func (u *joinUpdater) StatementBegin(ctx *sql.Context) {}
+func (u *joinUpdater) DiscardChanges(ctx *sql.Context, errorEncountered error) error {
+	return nil
+}
+func (u *joinUpdater) StatementComplete(ctx *sql.Context) error {
+	return nil
+}
+func (u *joinUpdater) Update(ctx *sql.Context, old sql.Row, new sql.Row) error {
+	return nil
+}
+func (u *joinUpdater) Close(ctx *sql.Context) error {
+	return nil
+}
+
 // UpdateInfo is the Info for OKResults returned by Update nodes.
 type UpdateInfo struct {
 	Matched, Updated, Warnings int

sql/rowexec/dml.go

@@ -157,18 +157,26 @@ func (b *BaseBuilder) buildForeignKeyHandler(ctx *sql.Context, n *plan.ForeignKe
 }
 
 func (b *BaseBuilder) buildUpdate(ctx *sql.Context, n *plan.Update, row sql.Row) (sql.RowIter, error) {
-	updatable, err := plan.GetUpdatable(n.Child)
-	if err != nil {
-		return nil, err
+	var updater sql.RowUpdater
+	var schema sql.Schema
+	if n.IsJoin {
+		updater = n.JoinUpdater()
+		schema = n.Schema()
+	} else {
+		updatable, err := plan.GetUpdatable(n.Child)
+		if err != nil {
+			return nil, err
+		}
+		updater = updatable.Updater(ctx)
+		schema = updatable.Schema()
 	}
-	updater := updatable.Updater(ctx)
 
 	iter, err := b.buildNodeExec(ctx, n.Child, row)
 	if err != nil {
 		return nil, err
 	}
 
-	return newUpdateIter(iter, updatable.Schema(), updater, n.Checks(), n.Ignore, n.Returning, n.Schema()), nil
+	return newUpdateIter(iter, schema, updater, n.Checks(), n.Ignore, n.Returning, n.Schema()), nil
 }
 
 func (b *BaseBuilder) buildDropForeignKey(ctx *sql.Context, n *plan.DropForeignKey, row sql.Row) (sql.RowIter, error) {