sql: scan only primary column for FK checks

This change improves the legacy FK path to only scan the primary
column family. The new path already does this (as long as the best
plan is a lookup join). Adding a test for both paths.

Fixes #30852.

Release note: None

pkg/sql/logictest/testdata/logic_test/fk

@@ -165,8 +165,8 @@ query T rowsort
 SELECT message FROM [SHOW KV TRACE FOR SESSION]
 WHERE message LIKE 'FKScan%'
 ----
-FKScan /Table/54/1/"fake"{-/#}
-FKScan /Table/53/1/2{-/#}
+FKScan /Table/54/1/"fake"/{0-1}
+FKScan /Table/53/1/2/{0-1}
 
 statement error pgcode 23503 foreign key violation: values \['780'\] in columns \[sku\] referenced in table "orders"
 DELETE FROM products
@@ -715,8 +715,8 @@ query T rowsort
 SELECT message FROM [SHOW KV TRACE FOR SESSION]
 WHERE message LIKE 'FKScan%'
 ----
-FKScan /Table/93/1/3{-/#}
-FKScan /Table/94/1/2{-/#}
+FKScan /Table/93/1/3/{0-1}
+FKScan /Table/94/1/2/{0-1}
 
 statement ok
 CREATE TABLE tx (
@@ -2679,3 +2679,49 @@ CREATE TABLE c_persistent(c INT PRIMARY KEY)
 
 statement error pq: constraints on temporary tables may reference only temporary tables
 CREATE TEMP TABLE c_temp(c INT, FOREIGN KEY (c) REFERENCES c_persistent(c))
+
+# Test that when the foreign key is a primary index we only look up the primary
+# family.
+subtest families
+
+statement ok
+CREATE TABLE fam_parent (
+  k INT PRIMARY KEY,
+  a INT,
+  b INT,
+  FAMILY (k, a),
+  FAMILY (b)
+)
+
+statement ok
+CREATE TABLE fam_child (
+  k INT PRIMARY KEY,
+  fk INT REFERENCES fam_parent(k)
+)
+
+statement ok
+INSERT INTO fam_parent VALUES (1, 1, 1)
+
+statement ok
+GRANT ALL ON fam_parent TO testuser;
+GRANT ALL ON fam_child TO testuser;
+
+# Open a transaction that modifies b.
+statement ok
+BEGIN
+
+statement count 1
+UPDATE fam_parent SET b = b+1 WHERE k = 1
+
+user testuser
+
+# Run an INSERT which needs to check existence of the row. If we try to scan
+# the entire row, this blocks on the other transaction. We should only be
+# scanning the primary column family.
+statement ok
+INSERT INTO fam_child VALUES (1, 1)
+
+user root
+
+statement ok
+COMMIT

pkg/sql/logictest/testdata/logic_test/fk_opt

@@ -2881,3 +2881,49 @@ INSERT INTO match_simple VALUES (1, 1)
 
 statement ok
 UPDATE match_simple SET a = NULL
+
+# Test that when the foreign key is a primary index we only look up the primary
+# family.
+subtest families
+
+statement ok
+CREATE TABLE fam_parent (
+  k INT PRIMARY KEY,
+  a INT,
+  b INT,
+  FAMILY (k, a),
+  FAMILY (b)
+)
+
+statement ok
+CREATE TABLE fam_child (
+  k INT PRIMARY KEY,
+  fk INT REFERENCES fam_parent(k)
+)
+
+statement ok
+INSERT INTO fam_parent VALUES (1, 1, 1)
+
+statement ok
+GRANT ALL ON fam_parent TO testuser;
+GRANT ALL ON fam_child TO testuser;
+
+# Open a transaction that modifies b.
+statement ok
+BEGIN
+
+statement count 1
+UPDATE fam_parent SET b = b+1 WHERE k = 1
+
+user testuser
+
+# Run an INSERT which needs to check existence of the row. If we try to scan
+# the entire row, this blocks on the other transaction. We should only be
+# scanning the primary column family.
+statement ok
+INSERT INTO fam_child VALUES (1, 1)
+
+user root
+
+statement ok
+COMMIT

pkg/sql/row/fk_spans.go

@@ -11,6 +11,7 @@
 package row
 
 import (
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
@@ -19,12 +20,29 @@ import (
 // spanForValues produce access spans for a single FK constraint and a
 // tuple of columns.
 func (f fkExistenceCheckBaseHelper) spanForValues(values tree.Datums) (roachpb.Span, error) {
-	var key roachpb.Key
-	if values != nil {
-		span, _, err := sqlbase.EncodePartialIndexSpan(
-			f.searchTable.TableDesc(), f.searchIdx, f.prefixLen, f.ids, values, f.searchPrefix)
-		return span, err
+	if values == nil {
+		key := roachpb.Key(f.searchPrefix)
+		return roachpb.Span{Key: key, EndKey: key.PrefixEnd()}, nil
 	}
-	key = roachpb.Key(f.searchPrefix)
-	return roachpb.Span{Key: key, EndKey: key.PrefixEnd()}, nil
+
+	// If we are scanning the (entire) primary key, only scan family 0 (which is
+	// always present).
+	// TODO(radu): this logic will need to be improved when secondary indexes also
+	// conform to families.
+	if f.searchIdx.ID == f.searchTable.PrimaryIndex.ID && f.prefixLen == len(f.searchIdx.ColumnIDs) {
+		// This code is equivalent to calling EncodePartialIndexSpan followed by
+		// MakeFamilyKey but saves an unnecessary allocation.
+		key, _, err := sqlbase.EncodePartialIndexKey(
+			f.searchTable.TableDesc(), f.searchIdx, f.prefixLen, f.ids, values, f.searchPrefix,
+		)
+		if err != nil {
+			return roachpb.Span{}, err
+		}
+		key = keys.MakeFamilyKey(key, 0)
+		return roachpb.Span{Key: key, EndKey: roachpb.Key(key).PrefixEnd()}, nil
+	}
+
+	span, _, err := sqlbase.EncodePartialIndexSpan(
+		f.searchTable.TableDesc(), f.searchIdx, f.prefixLen, f.ids, values, f.searchPrefix)
+	return span, err
 }