sql: ignore shard column during unique constraint searching

pkg/sql/catalog/descpb/index.go

@@ -65,13 +65,27 @@ func (desc *IndexDescriptor) IsValidOriginIndex(originColIDs ColumnIDs) bool {
 	return !desc.IsPartial() && ColumnIDs(desc.KeyColumnIDs).HasPrefix(originColIDs)
 }
 
+// explicitColumnIDsWithoutShardColumn returns explicit column ids of the index
+// excluding the shard column.
+func (desc *IndexDescriptor) explicitColumnIDsWithoutShardColumn() ColumnIDs {
+	explicitColIDs := desc.KeyColumnIDs[desc.ExplicitColumnStartIdx():]
+	explicitColNames := desc.KeyColumnNames[desc.ExplicitColumnStartIdx():]
+	colIDs := make(ColumnIDs, 0, len(explicitColIDs))
+	for i := range explicitColNames {
+		if !desc.IsSharded() || explicitColNames[i] != desc.Sharded.Name {
+			colIDs = append(colIDs, explicitColIDs[i])
+		}
+	}
+	return colIDs
+}
+
 // IsValidReferencedUniqueConstraint  is part of the UniqueConstraint interface.
 // It returns whether the index can serve as a referenced index for a foreign
 // key constraint with the provided set of referencedColumnIDs.
 func (desc *IndexDescriptor) IsValidReferencedUniqueConstraint(referencedColIDs ColumnIDs) bool {
 	return desc.Unique &&
 		!desc.IsPartial() &&
-		ColumnIDs(desc.KeyColumnIDs[desc.Partitioning.NumImplicitColumns:]).PermutationOf(referencedColIDs)
+		desc.explicitColumnIDsWithoutShardColumn().PermutationOf(referencedColIDs)
 }
 
 // GetName is part of the UniqueConstraint interface.

pkg/sql/logictest/testdata/logic_test/hash_sharded_index

@@ -680,3 +680,130 @@ descriptor_id  descriptor_name  index_id  index_name   index_type  is_unique  is
 
 statement ok
 DROP TABLE poor_t
+
+subtest fk_reference_shard_pk
+
+statement ok
+DROP TABLE IF EXISTS child
+
+statement ok
+DROP TABLE IF EXISTS parent
+
+statement ok
+CREATE TABLE parent (
+    id INT PRIMARY KEY USING HASH WITH BUCKET_COUNT = 8
+);
+
+statement ok
+CREATE TABLE child (
+    id INT PRIMARY KEY,
+    pid INT,
+    CONSTRAINT fk_child_pid FOREIGN KEY (pid) REFERENCES parent(id) ON DELETE CASCADE
+);
+
+statement error pq: insert on table "child" violates foreign key constraint "fk_child_pid"
+INSERT INTO child VALUES (1,1)
+
+statement ok
+INSERT INTO parent VALUES (1)
+
+statement ok
+INSERT INTO child VALUES (1,1)
+
+subtest fk_reference_shard_index
+
+statement ok
+DROP TABLE IF EXISTS child
+
+statement ok
+DROP TABLE IF EXISTS parent
+
+statement ok
+CREATE TABLE parent (
+    id INT PRIMARY KEY,
+    oid INT
+);
+
+statement ok
+CREATE UNIQUE INDEX t_idx_oid ON parent(oid) USING HASH WITH BUCKET_COUNT = 8
+
+statement ok
+CREATE TABLE child (
+    id INT PRIMARY KEY,
+    poid INT,
+    CONSTRAINT fk_child_pid FOREIGN KEY (poid) REFERENCES parent(oid) ON DELETE CASCADE
+);
+
+statement error pq: insert on table "child" violates foreign key constraint "fk_child_pid"
+INSERT INTO child VALUES (1,11)
+
+statement ok
+INSERT INTO parent VALUES (1,11)
+
+statement ok
+INSERT INTO child VALUES (1,11)
+
+subtest fk_reference_shard_pk_multi_col
+
+statement ok
+DROP TABLE IF EXISTS child
+
+statement ok
+DROP TABLE IF EXISTS parent
+
+statement ok
+CREATE TABLE parent (
+    a INT NOT NULL,
+    b INT NOT NULL,
+    PRIMARY KEY (a, b) USING HASH WITH BUCKET_COUNT = 8
+);
+
+statement ok
+CREATE TABLE child (
+    ca INT PRIMARY KEY,
+    cb INT,
+    CONSTRAINT fk_child_ca_cb FOREIGN KEY (ca, cb) REFERENCES parent(a, b) ON DELETE CASCADE
+);
+
+statement error pq: insert on table "child" violates foreign key constraint "fk_child_ca_cb"
+INSERT INTO child VALUES (1,1)
+
+statement ok
+INSERT INTO parent VALUES (1,1)
+
+statement ok
+INSERT INTO child VALUES (1,1)
+
+subtest fk_reference_shard_index_multi_col
+
+statement ok
+DROP TABLE IF EXISTS child
+
+statement ok
+DROP TABLE IF EXISTS parent
+
+statement ok
+CREATE TABLE parent (
+    a INT,
+    b INT
+);
+
+statement ok
+CREATE UNIQUE INDEX t_idx_a_b ON parent(a, b) USING HASH WITH BUCKET_COUNT = 8
+
+
+statement ok
+CREATE TABLE child (
+    ca INT PRIMARY KEY,
+    cb INT,
+    CONSTRAINT fk_child_ca_cb FOREIGN KEY (ca, cb) REFERENCES parent(a, b) ON DELETE CASCADE
+);
+
+statement error pq: insert on table "child" violates foreign key constraint "fk_child_ca_cb"
+INSERT INTO child VALUES (1,1)
+
+statement ok
+INSERT INTO parent VALUES (1,1)
+
+statement ok
+INSERT INTO child VALUES (1,1)