compose: deflake TestComposeCompare

Previously, the compose test had a mutation that added partial indexes
to random indexes. This could cause unique indexes to become unusable
for foreign keys, an issue that was being incorrectly ignored in the
schema changer. A recent change fixed this bug, which now causes these
scenarios to fail correctly. To address this, this patch modifies the
mutator inside compose to only add partial indexes on new, duplicate
indexes so that original unique constraints are not lost.

Fixes: #150817
Fixes: #149565

Release note: None

Author: fqazi

pkg/compose/compare/compare/compare_test.go

@@ -93,7 +93,7 @@ func TestCompare(t *testing.T) {
 						randgen.ForeignKeyMutator,
 						randgen.ColumnFamilyMutator,
 						randgen.IndexStoringMutator,
-						randgen.PartialIndexMutator,
+						randgen.DupPartialIndexMutator,
 					},
 				},
 			},

pkg/sql/randgen/mutator.go

@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"math/rand"
 	"regexp"
 	"sort"
@@ -44,6 +45,10 @@ var (
 	// indexes.
 	PartialIndexMutator MultiStatementMutation = partialIndexMutator
 
+	// DupPartialIndexMutator create a new random partial index based on an
+	// existing index that is selected.
+	DupPartialIndexMutator MultiStatementMutation = duplicatePartialIndexMutator
+
 	// PostgresMutator modifies strings such that they execute identically
 	// in both Postgres and Cockroach (however this mutator does not remove
 	// features not supported by Postgres; use PostgresCreateTableMutator
@@ -1140,6 +1145,84 @@ func partialIndexMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
 	return stmts, changed
 }
 
+// duplicatePartialIndexMutator is a mutations.MultiStatementMutator that will
+// copy a random index definition and make it a partial index. Unlike
+// partialIndexMutator, the existing definition will stay the same.
+func duplicatePartialIndexMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Statement, bool) {
+	changed := false
+	tables := getTableInfoFromDDLStatements(stmts)
+	var newStmts []tree.Statement
+	for _, stmt := range stmts {
+		switch ast := stmt.(type) {
+		case *tree.CreateIndex:
+			info, ok := tables[ast.Table.ObjectName]
+			if !ok {
+				continue
+			}
+
+			// If the index is not already a partial index, make it a partial index
+			// with a 50% chance. Do not mutate an index that was created to satisfy a
+			// FK constraint.
+			if ast.Predicate == nil &&
+				!hasReferencingConstraint(info, ast.Columns) &&
+				rng.Intn(2) == 0 {
+				astString := tree.AsStringWithFlags(ast, tree.FmtParsableNumerics)
+				astCopy, err := parser.ParseOne(astString)
+				if err != nil {
+					panic(err)
+				}
+				createIndexCopy := astCopy.AST.(*tree.CreateIndex)
+				if len(createIndexCopy.Name) > 0 {
+					createIndexCopy.Name = tree.Name(fmt.Sprintf("%s_partial", createIndexCopy.Name))
+				}
+				tn := tree.MakeUnqualifiedTableName(createIndexCopy.Table.ObjectName)
+				createIndexCopy.Predicate = randPartialIndexPredicateFromCols(rng, info.columnsTableDefs, &tn)
+				changed = true
+				newStmts = append(newStmts, createIndexCopy)
+			}
+		case *tree.CreateTable:
+			info, ok := tables[ast.Table.ObjectName]
+			if !ok {
+				panic("table info could not be found")
+			}
+			var newDefs []tree.TableDef
+			for _, def := range ast.Defs {
+				var idx *tree.IndexTableDef
+				switch defType := def.(type) {
+				case *tree.IndexTableDef:
+					idx = defType
+				case *tree.UniqueConstraintTableDef:
+					if !defType.PrimaryKey && !defType.WithoutIndex {
+						idx = &defType.IndexTableDef
+					}
+				}
+
+				if idx == nil {
+					continue
+				}
+
+				// If the index is not already a partial index, make it a partial
+				// index with a 50% chance.
+				if idx.Predicate == nil &&
+					!hasReferencingConstraint(info, idx.Columns) &&
+					rng.Intn(2) == 0 {
+					idxCopy := *idx
+					if len(idxCopy.Name) > 0 {
+						idxCopy.Name = tree.Name(fmt.Sprintf("%s_partial", idxCopy.Name))
+					}
+					tn := tree.MakeUnqualifiedTableName(ast.Table.ObjectName)
+					idxCopy.Predicate = randPartialIndexPredicateFromCols(rng, info.columnsTableDefs, &tn)
+					changed = true
+					newDefs = append(newDefs, &idxCopy)
+				}
+			}
+			ast.Defs = append(ast.Defs, newDefs...)
+		}
+	}
+	stmts = append(stmts, newStmts...)
+	return stmts, changed
+}
+
 // hasReferencingConstraint returns true if the tableInfo has any referencing
 // columns that match idxColumns.
 func hasReferencingConstraint(info tableInfo, idxColumns tree.IndexElemList) bool {