release-24.3.1-rc: optbuilder: always select tombstone index using the index ordinal

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.