opt: add inbound FK checks for upsert

This change adds inbound FK checks for upsert and switches execution over to the
new style FK checks for upsert.

Similar to UPDATE, the inbound FK checks run on the set difference between "old"
values for the FK columns and "new" values.

Release note (performance improvement): improved execution plans of foreign key
checks for UPSERT and INSERT ON CONFLICT in some cases (in particular
multi-region).

pkg/sql/logictest/testdata/logic_test/fk_opt

@@ -133,6 +133,122 @@ DROP TABLE child2
 statement ok
 DROP TABLE parent2
 
+# Upsert
+# ------
+
+statement ok
+CREATE TABLE parent (p INT PRIMARY KEY, other INT)
+
+statement ok
+CREATE TABLE child (c INT PRIMARY KEY, p INT NOT NULL REFERENCES parent(p))
+
+statement ok
+INSERT INTO parent VALUES (1), (2)
+
+# Insert case.
+statement ok
+INSERT INTO child VALUES (1, 1) ON CONFLICT (c) DO UPDATE SET p = 2
+
+statement error foreign key
+INSERT INTO child VALUES (2, 10) ON CONFLICT (c) DO UPDATE SET p = 2
+
+# Update case.
+statement ok
+INSERT INTO child VALUES (1, 1) ON CONFLICT (c) DO UPDATE SET p = 1
+
+statement ok
+INSERT INTO child VALUES (1, 10) ON CONFLICT (c) DO UPDATE SET p = 1
+
+statement error foreign key
+INSERT INTO child VALUES (1, 10) ON CONFLICT (c) DO UPDATE SET p = 10
+
+statement ok
+TRUNCATE child
+
+statement ok
+INSERT INTO child VALUES (1, 1)
+
+# Both insert and update case.
+
+# Both insert and update are invalid.
+
+statement error foreign key
+INSERT INTO child VALUES (1, 1), (2, 3) ON CONFLICT (c) DO UPDATE SET p = 3
+
+# Insert is invalid, update is valid.
+
+statement error foreign key
+INSERT INTO child VALUES (1, 2), (2, 3) ON CONFLICT (c) DO UPDATE SET p = 1
+
+# Insert is valid, update is invalid.
+
+statement error foreign key
+INSERT INTO child VALUES (1, 2), (2, 1) ON CONFLICT (c) DO UPDATE SET p = 3
+
+# Both insert and update are valid.
+
+statement ok
+INSERT INTO child VALUES (1, 2), (2, 1) ON CONFLICT (c) DO UPDATE SET p = 2
+
+statement ok
+DROP TABLE child
+
+statement ok
+DROP TABLE parent
+
+# Pseudo-deletions
+
+statement ok
+CREATE TABLE parent (a INT PRIMARY KEY, b INT, UNIQUE (b))
+
+statement ok
+CREATE TABLE child (a INT PRIMARY KEY, b INT REFERENCES parent (b))
+
+statement ok
+INSERT INTO parent VALUES (1, 2)
+
+statement ok
+INSERT INTO child VALUES (10, 2)
+
+statement error pq: upsert on table "parent" violates foreign key constraint "fk_b_ref_parent" on table "child"\nDETAIL: Key \(b\)=\(2\) is still referenced from table "child"\.
+UPSERT INTO parent VALUES (1, 3)
+
+statement ok
+INSERT INTO parent VALUES (1, 3), (2, 2) ON CONFLICT (a) DO UPDATE SET b = 3
+
+query II
+SELECT * FROM child
+----
+10  2
+
+query II rowsort
+SELECT * FROM parent
+----
+1  3
+2  2
+
+# child references the second '2' column in parent. This mutation removes that
+# row via an update, and is disallowed.
+statement error pq: insert on table "parent" violates foreign key constraint "fk_b_ref_parent" on table "child"\nDETAIL: Key \(b\)=\(2\) is still referenced from table "child"\.
+INSERT INTO parent VALUES (2, 2) ON CONFLICT (a) DO UPDATE SET b = parent.b - 1
+
+# This mutation *also* removes that row, but also via an update, introduces a
+# new one, making it acceptable.
+statement ok
+INSERT INTO parent VALUES (2, 2), (1, 3) ON CONFLICT (a) DO UPDATE SET b = parent.b - 1
+
+query II rowsort
+SELECT * FROM parent
+----
+1  2
+2  1
+
+statement ok
+DROP TABLE child
+
+statement ok
+DROP TABLE parent
+
 # --- Tests that follow are copied from the fk tests and adjusted as needed.
 
 statement ok

pkg/sql/opt/bench/stub_factory.go

@@ -279,6 +279,7 @@ func (f *stubFactory) ConstructUpsert(
 	returnCols exec.ColumnOrdinalSet,
 	checks exec.CheckOrdinalSet,
 	allowAutoCommit bool,
+	skipFKChecks bool,
 ) (exec.Node, error) {
 	return struct{}{}, nil
 }

pkg/sql/opt/exec/execbuilder/mutation.go

@@ -371,6 +371,7 @@ func (b *Builder) buildUpsert(ups *memo.UpsertExpr) (execPlan, error) {
 	updateColOrds := ordinalSetFromColList(ups.UpdateCols)
 	returnColOrds := ordinalSetFromColList(ups.ReturnCols)
 	checkOrds := ordinalSetFromColList(ups.CheckCols)
+	disableExecFKs := !ups.FKFallback
 	node, err := b.factory.ConstructUpsert(
 		input.root,
 		tab,
@@ -381,11 +382,16 @@ func (b *Builder) buildUpsert(ups *memo.UpsertExpr) (execPlan, error) {
 		returnColOrds,
 		checkOrds,
 		b.allowAutoCommit && len(ups.Checks) == 0,
+		disableExecFKs,
 	)
 	if err != nil {
 		return execPlan{}, err
 	}
 
+	if err := b.buildFKChecks(ups.Checks); err != nil {
+		return execPlan{}, err
+	}
+
 	// If UPSERT returns rows, they contain all non-mutation columns from the
 	// table, in the same order they're defined in the table. Each output column
 	// value is taken from an insert, fetch, or update column, depending on the

pkg/sql/opt/exec/factory.go

@@ -425,6 +425,11 @@ type Factory interface {
 	// transaction (if appropriate, i.e. if it is in an implicit transaction).
 	// This is false if there are multiple mutations in a statement, or the output
 	// of the mutation is processed through side-effecting expressions.
+	//
+	// If skipFKChecks is set, foreign keys are not checked as part of the
+	// execution of the upsert for the insert half. This is used when the FK
+	// checks are planned by the optimizer and are run separately as plan
+	// postqueries.
 	ConstructUpsert(
 		input Node,
 		table cat.Table,
@@ -435,6 +440,7 @@ type Factory interface {
 		returnCols ColumnOrdinalSet,
 		checks CheckOrdinalSet,
 		allowAutoCommit bool,
+		skipFKChecks bool,
 	) (Node, error)
 
 	// ConstructDelete creates a node that implements a DELETE statement. The
@@ -451,8 +457,8 @@ type Factory interface {
 	// of the mutation is processed through side-effecting expressions.
 	//
 	// If skipFKChecks is set, foreign keys are not checked as part of the
-	// execution of the insertion. This is used when the FK checks are planned by
-	// the optimizer and are run separately as plan postqueries.
+	// execution of the delete. This is used when the FK checks are planned
+	// by the optimizer and are run separately as plan postqueries.
 	ConstructDelete(
 		input Node,
 		table cat.Table,

pkg/sql/opt/optbuilder/insert.go

@@ -353,6 +353,9 @@ func (b *Builder) buildInsert(ins *tree.Insert, inScope *scope) (outScope *scope
 //      values specified for them.
 //   3. Each update value is the same as the corresponding insert value.
 //
+// TODO(radu): once FKs no longer require indexes, this function will have to
+// take FKs into account explicitly.
+//
 // TODO(andyk): The fast path is currently only enabled when the UPSERT alias
 // is explicitly selected by the user. It's possible to fast path some queries
 // of the form INSERT ... ON CONFLICT, but the utility is low and there are lots

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -876,6 +876,15 @@ func (mb *mutationBuilder) parseDefaultOrComputedExpr(colID opt.ColumnID) tree.E
 
 // buildFKChecks* methods populate mb.checks with queries that check the
 // integrity of foreign key relations that involve modified rows.
+//
+// The foreign key checks are queries that run after the statement (including
+// the relevant mutation) completes; any row that is returned by these
+// FK check queries indicates a foreign key violation.
+//
+// In the case of insert, each FK check query is an anti-join with the left side
+// being a WithScan of the mutation input and the right side being the
+// referenced table.
+//
 func (mb *mutationBuilder) buildFKChecksForInsert() {
 	if mb.tab.OutboundForeignKeyCount() == 0 {
 		// No relevant FKs.
@@ -899,6 +908,15 @@ func (mb *mutationBuilder) buildFKChecksForInsert() {
 
 // buildFKChecks* methods populate mb.checks with queries that check the
 // integrity of foreign key relations that involve modified rows.
+//
+// The foreign key checks are queries that run after the statement (including
+// the relevant mutation) completes; any row that is returned by these
+// FK check queries indicates a foreign key violation.
+//
+// In the case of delete, each FK check query is a semi-join with the left side
+// being a WithScan of the mutation input and the right side being the
+// referencing table.
+//
 func (mb *mutationBuilder) buildFKChecksForDelete() {
 	if mb.tab.InboundForeignKeyCount() == 0 {
 		// No relevant FKs.
@@ -931,6 +949,23 @@ func (mb *mutationBuilder) buildFKChecksForDelete() {
 
 // buildFKChecks* methods populate mb.checks with queries that check the
 // integrity of foreign key relations that involve modified rows.
+//
+// The foreign key checks are queries that run after the statement (including
+// the relevant mutation) completes; any row that is returned by these
+// FK check queries indicates a foreign key violation.
+//
+// In the case of update, there are two types of FK check queries:
+//  - insertion-side checks are very similar to the checks we issue for insert;
+//    they are an anti-join with the left side being a WithScan of the "new"
+//    values for each row.
+//  - deletion-side checks are similar to the checks we issue for delete; they
+//    are a semi-join but the left side input is more complicated: it is an
+//    Except between a WithScan of the "old" values and a WithScan of the "new"
+//    values for each row (this is the set of values that are effectively
+//    removed from the table).
+//
+// Only FK relations that involve updated columns result in FK checks.
+//
 func (mb *mutationBuilder) buildFKChecksForUpdate() {
 	if mb.tab.OutboundForeignKeyCount() == 0 && mb.tab.InboundForeignKeyCount() == 0 {
 		return
@@ -992,12 +1027,15 @@ func (mb *mutationBuilder) buildFKChecksForUpdate() {
 			return
 		}
 
+		// Construct an Except expression for the set difference between "old"
+		// FK values and "new" FK values.
+
 		oldRows, colsForOldRow, _ := h.makeFKInputScan(fkInputScanFetchedVals)
 		newRows, colsForNewRow, _ := h.makeFKInputScan(fkInputScanNewVals)
 
 		// The rows that no longer exist are the ones that were "deleted" by virtue
 		// of being updated _from_, minus the ones that were "added" by virtue of
-		// being updated _to_. Note that this could equivalently be ExceptAll.
+		// being updated _to_.
 		deletedRows := mb.b.factory.ConstructExcept(
 			oldRows,
 			newRows,
@@ -1012,13 +1050,41 @@ func (mb *mutationBuilder) buildFKChecksForUpdate() {
 	}
 }
 
+// buildFKChecks* methods populate mb.checks with queries that check the
+// integrity of foreign key relations that involve modified rows.
+//
+// The foreign key checks are queries that run after the statement (including
+// the relevant mutation) completes; any row that is returned by these
+// FK check queries indicates a foreign key violation.
+//
+// The case of upsert is similar to update; there are two types of FK check
+// queries:
+//  - insertion-side checks are very similar to the checks we issue for insert;
+//    they are an anti-join with the left side being a WithScan of the "new"
+//    values for each row. In some cases, the "new" value can be the result of
+//    an expression of the form:
+//      CASE WHEN canary IS NULL THEN inserted-value ELSE updated-value END
+//    These expressions are already projected as part of the mutation input and
+//    are accessible through WithScan.
+//
+//  - deletion-side checks are similar to the checks we issue for delete; they
+//    are a semi-join but the left side input is more complicated: it is an
+//    Except between a WithScan of the "old" values and a WithScan of the "new"
+//    values for each row (this is the set of values that are effectively
+//    removed from the table).
+//
+// Only FK relations that involve updated columns result in deletion-side FK
+// check. The insertion-side FK checks are always needed (similar to insert)
+// because any of the rows might result in an insert rather than an update.
+//
 func (mb *mutationBuilder) buildFKChecksForUpsert() {
 	numOutbound := mb.tab.OutboundForeignKeyCount()
 	numInbound := mb.tab.InboundForeignKeyCount()
 
 	if numOutbound == 0 && numInbound == 0 {
 		return
 	}
+
 	if !mb.b.evalCtx.SessionData.OptimizerFKs {
 		mb.fkFallback = true
 		return
@@ -1029,6 +1095,50 @@ func (mb *mutationBuilder) buildFKChecksForUpsert() {
 	for i := 0; i < numOutbound; i++ {
 		mb.addInsertionCheck(i)
 	}
+
+	for i := 0; i < numInbound; i++ {
+		// Verify that at least one FK column is updated by the Upsert; columns that
+		// are not updated can get new values (through the insert path) but existing
+		// values are never removed.
+		if !mb.inboundFKColsUpdated(i) {
+			continue
+		}
+
+		h := &mb.fkCheckHelper
+		if !h.initWithInboundFK(mb, i) {
+			continue
+		}
+
+		if a := h.fk.UpdateReferenceAction(); a != tree.Restrict && a != tree.NoAction {
+			// Bail, so that exec FK checks pick up on FK checks and perform them.
+			mb.checks = nil
+			mb.fkFallback = true
+			return
+		}
+
+		// Construct an Except expression for the set difference between "old" FK
+		// values and "new" FK values. Note that technically, to get the old rows we
+		// should be selecting only the rows that are being updated (using a
+		// "canaryCol IS NOT NULL" condition); but these rows are harmless because
+		// they have all-null fetched values and thus will never match in the semi
+		// join.
+		oldRows, colsForOldRow, _ := h.makeFKInputScan(fkInputScanFetchedVals)
+		newRows, colsForNewRow, _ := h.makeFKInputScan(fkInputScanNewVals)
+
+		// The rows that no longer exist are the ones that were "deleted" by virtue
+		// of being updated _from_, minus the ones that were "added" by virtue of
+		// being updated _to_.
+		deletedRows := mb.b.factory.ConstructExcept(
+			oldRows,
+			newRows,
+			&memo.SetPrivate{
+				LeftCols:  colsForOldRow,
+				RightCols: colsForNewRow,
+				OutCols:   colsForOldRow,
+			},
+		)
+		mb.addDeletionCheck(h, deletedRows, colsForOldRow)
+	}
 }
 
 // addInsertionCheck adds a FK check for rows which are added to a table.