reverted most changes, modified UpdateJoin to contain Updatables instead of Updaters

Author: angelamayxie

sql/analyzer/apply_foreign_keys.go

@@ -122,44 +122,33 @@ func applyForeignKeysToNodes(ctx *sql.Context, a *Analyzer, n sql.Node, cache *f
 		if plan.IsEmptyTable(n.Child) {
 			return n, transform.SameTree, nil
 		}
-		targets := n.GetUpdateTargets()
-		foreignKeyHandlers := make([]sql.Node, len(targets))
-		copy(foreignKeyHandlers, targets)
-
-		for i, node := range targets {
-			updateDest, err := plan.GetUpdatable(node)
-			if err != nil {
-				return nil, transform.SameTree, err
-			}
+		// 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
+		}
 
-			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],
-				Editor:       fkEditor,
-				AllUpdaters:  fkChain.GetUpdaters(),
-			}
+		fkEditor, err := getForeignKeyEditor(ctx, a, fkTbl, cache, fkChain, false)
+		if err != nil {
+			return nil, transform.SameTree, err
 		}
-		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
+		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
 	case *plan.DeleteFrom:
 		if plan.IsEmptyTable(n.Child) {
 			return n, transform.SameTree, nil

sql/analyzer/assign_update_join.go

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

sql/plan/update.go

@@ -31,13 +31,11 @@ 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
-	updateJoinTargets []sql.Node
-	joinSchema        sql.Schema
-	HasSingleRel      bool
-	IsProcNested      bool
+	checks       sql.CheckConstraints
+	Ignore       bool
+	IsJoin       bool
+	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.
@@ -232,120 +230,3 @@ func (u *Update) DebugString() string {
 	_ = pr.WriteChildren(sql.DebugString(u.Child))
 	return pr.String()
 }
-
-// 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) GetUpdaterAndSchema(ctx *sql.Context) (sql.RowUpdater, sql.Schema, error) {
-	if u.IsJoin {
-		updaterMap := make(map[string]sql.RowUpdater)
-		for _, target := range u.updateJoinTargets {
-			targetTable, err := GetUpdatable(target)
-			if err != nil {
-				return nil, nil, err
-			}
-			updaterMap[targetTable.Name()] = targetTable.Updater(ctx)
-		}
-		return &joinUpdater{
-			updaterMap: updaterMap,
-			schemaMap:  RecreateTableSchemaFromJoinSchema(u.joinSchema),
-			joinSchema: u.joinSchema,
-		}, u.joinSchema, nil
-	}
-	updatable, err := GetUpdatable(u.Child)
-	if err != nil {
-		return nil, nil, err
-	}
-	return updatable.Updater(ctx), updatable.Schema(), nil
-}
-
-type joinUpdater struct {
-	updaterMap map[string]sql.RowUpdater
-	schemaMap  map[string]sql.Schema
-	joinSchema sql.Schema
-}
-
-var _ sql.RowUpdater = (*joinUpdater)(nil)
-
-// StatementBegins implements the sql.TableEditor interface
-func (u *joinUpdater) StatementBegin(ctx *sql.Context) {
-	for _, updater := range u.updaterMap {
-		updater.StatementBegin(ctx)
-	}
-}
-
-// DiscardChanges implements the sql.TableEditor interface
-func (u *joinUpdater) DiscardChanges(ctx *sql.Context, errorEncountered error) error {
-	for _, updater := range u.updaterMap {
-		err := updater.DiscardChanges(ctx, errorEncountered)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-// StatementComplete implements the sql.TableEditor interface
-func (u *joinUpdater) StatementComplete(ctx *sql.Context) error {
-	for _, updater := range u.updaterMap {
-		err := updater.StatementComplete(ctx)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-func (u *joinUpdater) Update(ctx *sql.Context, old sql.Row, new sql.Row) error {
-	tableToOldRowMap := SplitRowIntoTableRowMap(old, u.joinSchema)
-	tableToNewRowMap := SplitRowIntoTableRowMap(new, u.joinSchema)
-
-	for tableName, updater := range u.updaterMap {
-		oldRow := tableToOldRowMap[tableName]
-		newRow := tableToNewRowMap[tableName]
-		schema := u.schemaMap[tableName]
-
-		eq, err := oldRow.Equals(ctx, newRow, schema)
-		if err != nil {
-			return err
-		}
-
-		if !eq {
-			err = updater.Update(ctx, oldRow, newRow)
-		}
-
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func (u *joinUpdater) Close(ctx *sql.Context) error {
-	for _, updater := range u.updaterMap {
-		err := updater.Close(ctx)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}