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,157 @@ 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
+
+# Self-reference.
+
+statement ok
+CREATE TABLE self (k int primary key, a int unique, b int references self(a))
+
+statement error pq: upsert on table "self" violates foreign key constraint "fk_b_ref_self"\nDETAIL: Key \(b\)=\(2\) is not present in table "self"
+UPSERT INTO self VALUES (1, 1, 2)
+
+statement ok
+UPSERT INTO self VALUES (1, 1, 1)
+
+statement ok
+UPSERT INTO self VALUES (1, 1, 1)
+
+statement error pq: upsert on table "self" violates foreign key constraint "fk_b_ref_self"\nDETAIL: Key \(b\)=\(2\) is not present in table "self"
+UPSERT INTO self VALUES (1, 1, 2)
+
+statement ok
+UPSERT INTO self VALUES (1, 2, 2)
+
+statement error pq: upsert on table "self" violates foreign key constraint "fk_b_ref_self"\nDETAIL: Key \(b\)=\(2\) is not present in table "self"
+UPSERT INTO self(k,a) VALUES (1, 1)
+
+statement ok
+UPSERT INTO self VALUES (1, 1, 2), (2, 2, 1)
+
+statement ok
+INSERT INTO self VALUES (2, 2, 2) ON CONFLICT (k) DO UPDATE SET b = self.b + 1
+
+statement error pq: insert on table "self" violates foreign key constraint "fk_b_ref_self"\nDETAIL: Key \(b\)=\(3\) is not present in table "self"
+INSERT INTO self VALUES (2, 2, 2) ON CONFLICT (k) DO UPDATE SET b = self.b + 1
+
+statement ok
+DROP TABLE self
+
 # --- 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

@@ -383,6 +383,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,
@@ -393,11 +394,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/norm/custom_funcs.go

@@ -2084,9 +2084,11 @@ func (c *CustomFuncs) WithUses(r opt.Expr) props.WithUsesMap {
 	switch e := r.(type) {
 	case memo.RelExpr:
 		relProps := e.Relational()
+
+		// Lazily calculate and store the WithUses value.
 		if !relProps.IsAvailable(props.WithUses) {
-			relProps.SetAvailable(props.WithUses)
 			relProps.Shared.Rule.WithUses = c.deriveWithUses(r)
+			relProps.SetAvailable(props.WithUses)
 		}
 		return relProps.Shared.Rule.WithUses
 

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