optbuilder: always select tombstone index using the index ordinal

Previously, the index to which tombstones would be written to enforce
uniqueness was chosen using the ordinal of uniqueness constraint, which
happens to be the same as the index ordinal if no non-unique indexes
precede the unique indexes on a relation. However, if a non-unique index
precedes a unique index, we will erroneously apply the uniqueness check
to that index. In most cases, this will cause queries to fail on an
assertion within the row helper code.

However, a mixture of serializable and non-serializable mutations and
some sequences of index additions and non-serializable mutations can cause
uniqueness to be violated. This impacts all regional by row tables with
uniqueness constraints that do not include the region which are mutated
under non-serializable isolations.

This patch plumbs through the index ordinal so that the proper index
will be checked for uniqueness in all cases.

Fixes: #137341
Release note (bug fix): Regional by row tables with uniqueness constraints
where the region is not part of those uniqueness constraints and which also
contain non-unique indices will now have that uniqueness properly enforced
when modified at READ-COMMITTED isolation. This bug was introduced in
release 24.3.0.

pkg/ccl/logictestccl/testdata/logic_test/regional_by_row_read_committed

@@ -272,3 +272,48 @@ CMU                          Scottie Dog  15213
 Central Michigan University  Chippewas    97858
 Evergreen State              geoducks     98505
 Western Oregon               wolves       97361
+
+#
+# Reproduction for #137341
+#
+statement ok
+CREATE TABLE t137341 (
+  region crdb_internal_region NOT NULL,
+  id STRING,
+  u1 STRING,
+  u2 STRING,
+  s STRING,
+  PRIMARY KEY (id),
+  UNIQUE INDEX (u1),
+  INDEX (s),
+  UNIQUE INDEX (u2)
+) LOCALITY REGIONAL BY ROW AS region
+
+statement ok
+INSERT INTO t137341 (region, id, u1, u2, s) VALUES ('ca-central-1', 'a', 'u1-a', 'u2-a', 'common')
+
+statement ok
+CREATE TABLE t137341b (
+  region crdb_internal_region NOT NULL,
+  id STRING,
+  u1 STRING,
+  u2 STRING,
+  s STRING,
+  PRIMARY KEY (id),
+  INDEX (s)
+) LOCALITY REGIONAL BY ROW AS region
+
+statement ok
+INSERT INTO t137341b (region, id, u1, u2, s) VALUES ('ca-central-1', 'a', 'u1-a', 'u2-a', 'common')
+
+statement ok
+INSERT INTO t137341b (region, id, u1, u2, s) VALUES ('us-east-1', 'b', 'u1-b', 'u2-b', 'common')
+
+statement ok
+CREATE UNIQUE INDEX ON t137341b(u1)
+
+statement ok
+CREATE UNIQUE INDEX ON t137341b(u2)
+
+statement error pgcode 23505 pq: duplicate key value violates unique constraint "t137341b_u2_key"
+UPDATE t137341b SET u2 = 'u2-a' WHERE u2 = 'u2-b'

pkg/sql/opt/cat/table.go

@@ -363,9 +363,11 @@ type UniqueConstraint interface {
 	// WithoutIndex is true if this unique constraint is not enforced by an index.
 	WithoutIndex() bool
 
-	// CanUseTombstones is true if this unique constraint can be enforced by
-	// writing tombstones to all partitions.
-	CanUseTombstones() bool
+	// TombstoneIndexOrdinal returns the index ordinal of the index into which to write
+	// tombstones for the enforcement of this uniqueness constraint, or -1 if this
+	// constraint is not enforceable with tombstones. 'ok' is true if this constraint is
+	// enforceable with tombstones, false otherwise.
+	TombstoneIndexOrdinal() (_ IndexOrdinal, ok bool)
 
 	// Validated is true if the constraint is validated (i.e. we know that the
 	// existing data satisfies the constraint). It is possible to set up a unique

pkg/sql/opt/optbuilder/mutation_builder_unique.go

@@ -52,12 +52,12 @@ func (mb *mutationBuilder) buildUniqueChecksForInsert() {
 		// For non-serializable transactions, we guarantee uniqueness by writing tombstones in all
 		// partitions of a unique index with implicit partitioning columns.
 		if mb.b.evalCtx.TxnIsoLevel != isolation.Serializable {
-			if !mb.tab.Unique(i).CanUseTombstones() {
-				panic(unimplemented.NewWithIssue(126592,
-					"unique without index constraint under non-serializable isolation levels"))
+			if indexOrdinal, ok := u.TombstoneIndexOrdinal(); ok {
+				mb.uniqueWithTombstoneIndexes.Add(indexOrdinal)
+				continue
 			}
-			mb.uniqueWithTombstoneIndexes.Add(i)
-			continue
+			panic(unimplemented.NewWithIssue(126592,
+				"unique without index constraint under non-serializable isolation levels"))
 		}
 
 		// If this constraint is an arbiter of an INSERT ... ON CONFLICT ... DO
@@ -111,13 +111,14 @@ func (mb *mutationBuilder) buildUniqueChecksForUpdate() {
 		// For non-serializable transactions, we guarantee uniqueness by writing tombstones in all
 		// partitions of a unique index with implicit partitioning columns.
 		if mb.b.evalCtx.TxnIsoLevel != isolation.Serializable {
-			if !mb.tab.Unique(i).CanUseTombstones() {
-				panic(unimplemented.NewWithIssue(126592,
-					"unique without index constraint under non-serializable isolation levels"))
+			if indexOrdinal, ok := u.TombstoneIndexOrdinal(); ok {
+				mb.uniqueWithTombstoneIndexes.Add(indexOrdinal)
+				continue
 			}
-			mb.uniqueWithTombstoneIndexes.Add(i)
-			continue
+			panic(unimplemented.NewWithIssue(126592,
+				"unique without index constraint under non-serializable isolation levels"))
 		}
+
 		if h.init(mb, i) {
 			// The insertion check works for updates too since it simply checks that
 			// the unique columns in the newly inserted or updated rows do not match
@@ -152,13 +153,14 @@ func (mb *mutationBuilder) buildUniqueChecksForUpsert() {
 		// For non-serializable transactions, we guarantee uniqueness by writing tombstones in all
 		// partitions of a unique index with implicit partitioning columns.
 		if mb.b.evalCtx.TxnIsoLevel != isolation.Serializable {
-			if !mb.tab.Unique(i).CanUseTombstones() {
-				panic(unimplemented.NewWithIssue(126592,
-					"unique without index constraint under non-serializable isolation levels"))
+			if indexOrdinal, ok := u.TombstoneIndexOrdinal(); ok {
+				mb.uniqueWithTombstoneIndexes.Add(indexOrdinal)
+				continue
 			}
-			mb.uniqueWithTombstoneIndexes.Add(i)
-			continue
+			panic(unimplemented.NewWithIssue(126592,
+				"unique without index constraint under non-serializable isolation levels"))
 		}
+
 		// If this constraint is an arbiter of an INSERT ... ON CONFLICT ... DO
 		// UPDATE clause and not updated by the DO UPDATE clause, we don't need to
 		// plan a check (ON CONFLICT ... DO NOTHING does not go through this code

pkg/sql/opt/testutils/testcat/test_catalog.go

@@ -1557,13 +1557,14 @@ func (fk *ForeignKeyConstraint) UpdateReferenceAction() tree.ReferenceAction {
 // UniqueConstraint implements cat.UniqueConstraint. See that interface
 // for more information on the fields.
 type UniqueConstraint struct {
-	name             string
-	tabID            cat.StableID
-	columnOrdinals   []int
-	predicate        string
-	withoutIndex     bool
-	canUseTombstones bool
-	validated        bool
+	name                  string
+	tabID                 cat.StableID
+	columnOrdinals        []int
+	predicate             string
+	withoutIndex          bool
+	canUseTombstones      bool
+	tombstoneIndexOrdinal cat.IndexOrdinal
+	validated             bool
 }
 
 var _ cat.UniqueConstraint = &UniqueConstraint{}
@@ -1604,7 +1605,16 @@ func (u *UniqueConstraint) WithoutIndex() bool {
 	return u.withoutIndex
 }
 
-func (u *UniqueConstraint) CanUseTombstones() bool { return u.canUseTombstones }
+// TombstoneIndexOrdinal is part of the cat.UniqueConstraint interface
+func (u *UniqueConstraint) TombstoneIndexOrdinal() (ordinal cat.IndexOrdinal, ok bool) {
+	ok = u.canUseTombstones
+	if ok {
+		ordinal = u.tombstoneIndexOrdinal
+	} else {
+		ordinal = -1
+	}
+	return ordinal, ok
+}
 
 // Validated is part of the cat.UniqueConstraint interface.
 func (u *UniqueConstraint) Validated() bool {

pkg/sql/opt_catalog.go

@@ -1015,12 +1015,13 @@ func newOptTable(
 				canUseTombstones := idx.ImplicitPartitioningColumnCount() == 1 &&
 					partitionColumn.GetType().Family() == types.EnumFamily
 				ot.uniqueConstraints = append(ot.uniqueConstraints, optUniqueConstraint{
-					name:             idx.GetName(),
-					table:            ot.ID(),
-					columns:          idx.IndexDesc().KeyColumnIDs[idx.IndexDesc().ExplicitColumnStartIdx():],
-					withoutIndex:     true,
-					canUseTombstones: canUseTombstones,
-					predicate:        idx.GetPredicate(),
+					name:                  idx.GetName(),
+					table:                 ot.ID(),
+					columns:               idx.IndexDesc().KeyColumnIDs[idx.IndexDesc().ExplicitColumnStartIdx():],
+					withoutIndex:          true,
+					canUseTombstones:      canUseTombstones,
+					tombstoneIndexOrdinal: idx.Ordinal(),
+					predicate:             idx.GetPredicate(),
 					// TODO(rytaft): will we ever support an unvalidated unique constraint
 					// here?
 					validity: descpb.ConstraintValidity_Validated,
@@ -2007,9 +2008,10 @@ type optUniqueConstraint struct {
 	columns   []descpb.ColumnID
 	predicate string
 
-	withoutIndex     bool
-	canUseTombstones bool
-	validity         descpb.ConstraintValidity
+	withoutIndex          bool
+	canUseTombstones      bool
+	tombstoneIndexOrdinal cat.IndexOrdinal
+	validity              descpb.ConstraintValidity
 
 	uniquenessGuaranteedByAnotherIndex bool
 }
@@ -2054,8 +2056,14 @@ func (u *optUniqueConstraint) WithoutIndex() bool {
 	return u.withoutIndex
 }
 
-func (u *optUniqueConstraint) CanUseTombstones() bool {
-	return u.canUseTombstones
+func (u *optUniqueConstraint) TombstoneIndexOrdinal() (ordinal cat.IndexOrdinal, ok bool) {
+	ok = u.canUseTombstones
+	if ok {
+		ordinal = u.tombstoneIndexOrdinal
+	} else {
+		ordinal = -1
+	}
+	return ordinal, ok
 }
 
 // Validated is part of the cat.UniqueConstraint interface.