Minor updates to support UPDATE ... FROM in Doltgres, through the existing UpdateJoin support

Author: fulghum

sql/analyzer/apply_foreign_keys.go

@@ -122,6 +122,8 @@ func applyForeignKeysToNodes(ctx *sql.Context, a *Analyzer, n sql.Node, cache *f
 		if plan.IsEmptyTable(n.Child) {
 			return n, transform.SameTree, nil
 		}
+		// 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

sql/analyzer/assign_update_join.go

@@ -39,7 +39,18 @@ func modifyUpdateExprsForJoin(ctx *sql.Context, a *Analyzer, n sql.Node, scope *
 			return nil, transform.SameTree, err
 		}
 
+		tablesToUpdate, err := getResolvedTablesToUpdate(ctx, us, jn)
+		if err != nil {
+			return n, transform.SameTree, err
+		}
+
+		tableNames := make([]string, len(tablesToUpdate))
+		for i, tableToUpdate := range tablesToUpdate {
+			tableNames[i] = tableToUpdate.Name()
+		}
+
 		uj := plan.NewUpdateJoin(updaters, us)
+		uj.TargetTables = tableNames
 		ret, err := n.WithChildren(uj)
 		if err != nil {
 			return nil, transform.SameTree, err
@@ -99,3 +110,19 @@ func getTablesToBeUpdated(node sql.Node) map[string]struct{} {
 
 	return ret
 }
+
+func getResolvedTablesToUpdate(_ *sql.Context, node sql.Node, ij sql.Node) (resolvedTables []*plan.ResolvedTable, err error) {
+	namesOfTableToBeUpdated := getTablesToBeUpdated(node)
+	resolvedTablesMap := getTablesByName(ij)
+
+	for tableToBeUpdated, _ := range namesOfTableToBeUpdated {
+		resolvedTable, ok := resolvedTablesMap[strings.ToLower(tableToBeUpdated)]
+		if !ok {
+			return nil, plan.ErrUpdateForTableNotSupported.New(tableToBeUpdated)
+		}
+
+		resolvedTables = append(resolvedTables, resolvedTable)
+	}
+
+	return resolvedTables, nil
+}

sql/plan/update_join.go

@@ -21,7 +21,8 @@ import (
 )
 
 type UpdateJoin struct {
-	Updaters map[string]sql.RowUpdater
+	Updaters     map[string]sql.RowUpdater
+	TargetTables []string
 	UnaryNode
 }
 
@@ -54,7 +55,14 @@ func (u *UpdateJoin) DebugString() string {
 
 // GetUpdatable returns an updateJoinTable which implements sql.UpdatableTable.
 func (u *UpdateJoin) GetUpdatable() sql.UpdatableTable {
+	// TODO: UpdateJoin can update multiple tables, but this interface only allows for a single table.
+	//       Additionally, updatableJoinTable doesn't implement interfaces that other parts of the code
+	//       expect, so UpdateJoins don't always work correctly. For example, because updatableJoinTable
+	//       doesn't implement ForeignKeyTable, we UpdateJoin statements don't properly enforce foreign key
+	//       checks. We should revamp this function so that we can support multiple tables being updated,
+	//       but right now we just return the name of the
 	return &updatableJoinTable{
+		name:     u.TargetTables[0],
 		updaters: u.Updaters,
 		joinNode: u.Child.(*UpdateSource).Child,
 	}
@@ -66,7 +74,9 @@ func (u *UpdateJoin) WithChildren(children ...sql.Node) (sql.Node, error) {
 		return nil, sql.ErrInvalidChildrenNumber.New(u, len(children), 1)
 	}
 
-	return NewUpdateJoin(u.Updaters, children[0]), nil
+	newUpdateJoin := NewUpdateJoin(u.Updaters, children[0])
+	newUpdateJoin.TargetTables = u.TargetTables
+	return newUpdateJoin, nil
 }
 
 func (u *UpdateJoin) IsReadOnly() bool {
@@ -80,6 +90,7 @@ func (u *UpdateJoin) CollationCoercibility(ctx *sql.Context) (collation sql.Coll
 
 // updatableJoinTable manages the update of multiple tables.
 type updatableJoinTable struct {
+	name     string
 	updaters map[string]sql.RowUpdater
 	joinNode sql.Node
 }
@@ -98,7 +109,7 @@ func (u *updatableJoinTable) PartitionRows(context *sql.Context, partition sql.P
 
 // Name implements the sql.UpdatableTable interface.
 func (u *updatableJoinTable) Name() string {
-	panic("this method should not be called")
+	return u.name
 }
 
 // String implements the sql.UpdatableTable interface.

sql/planbuilder/dml.go

@@ -534,20 +534,29 @@ func (b *Builder) buildUpdate(inScope *scope, u *ast.Update) (outScope *scope) {
 	update.IsProcNested = b.ProcCtx().DbName != ""
 
 	var checks []*sql.CheckConstraint
-	if join, ok := outScope.node.(*plan.JoinNode); ok {
-		// TODO this doesn't work, a lot of the time the top node
-		// is a filter. This would have to go before we build the
-		// filter/accessory nodes. But that errors for a lot of queries.
+	if hasJoinNode(outScope.node) {
 		source := plan.NewUpdateSource(
-			join,
+			outScope.node,
 			ignore,
 			updateExprs,
 		)
-		updaters, err := rowUpdatersByTable(b.ctx, source, join)
+		updaters, err := rowUpdatersByTable(b.ctx, source, outScope.node)
 		if err != nil {
 			b.handleErr(err)
 		}
+
+		tablesToUpdate, err := getResolvedTablesToUpdate(b.ctx, source, outScope.node)
+		if err != nil {
+			b.handleErr(err)
+		}
+
+		tableNames := make([]string, len(tablesToUpdate))
+		for i, tableToUpdate := range tablesToUpdate {
+			tableNames[i] = tableToUpdate.Name()
+		}
+
 		updateJoin := plan.NewUpdateJoin(updaters, source)
+		updateJoin.TargetTables = tableNames
 		update.Child = updateJoin
 		transform.Inspect(update, func(n sql.Node) bool {
 			// todo maybe this should be later stage
@@ -594,6 +603,34 @@ func (b *Builder) buildUpdate(inScope *scope, u *ast.Update) (outScope *scope) {
 	return
 }
 
+// hasJoinNode returns true if |node| or any child is a JoinNode.
+func hasJoinNode(node sql.Node) bool {
+	updateJoinFound := false
+	transform.Inspect(node, func(n sql.Node) bool {
+		if _, ok := n.(*plan.JoinNode); ok {
+			updateJoinFound = true
+		}
+		return !updateJoinFound
+	})
+	return updateJoinFound
+}
+
+func getResolvedTablesToUpdate(_ *sql.Context, node sql.Node, ij sql.Node) (resolvedTables []*plan.ResolvedTable, err error) {
+	namesOfTableToBeUpdated := getTablesToBeUpdated(node)
+	resolvedTablesMap := getTablesByName(ij)
+
+	for tableToBeUpdated, _ := range namesOfTableToBeUpdated {
+		resolvedTable, ok := resolvedTablesMap[strings.ToLower(tableToBeUpdated)]
+		if !ok {
+			return nil, plan.ErrUpdateForTableNotSupported.New(tableToBeUpdated)
+		}
+
+		resolvedTables = append(resolvedTables, resolvedTable)
+	}
+
+	return resolvedTables, 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)

sql/rowexec/update.go

@@ -258,8 +258,12 @@ func (u *updateJoinIter) Next(ctx *sql.Context) (sql.Row, error) {
 			if errors.Is(err, sql.ErrKeyNotFound) {
 				cache.Put(hash, struct{}{})
 
-				// updateJoin counts matched rows from join output
-				u.accumulator.handleRowMatched()
+				// updateJoin counts matched rows from join output, unless a RETURNING clause
+				// is in use, in which case there will not be an accumulator assigned, since we
+				// don't need to return the count of updated rows, just the RETURNING expressions.
+				if u.accumulator != nil {
+					u.accumulator.handleRowMatched()
+				}
 
 				continue
 			} else if err != nil {