|
9 | 9 | "bytes"
|
10 | 10 | "context"
|
11 | 11 | "encoding/json"
|
| 12 | + "fmt" |
12 | 13 | "math/rand"
|
13 | 14 | "regexp"
|
14 | 15 | "sort"
|
|
44 | 45 | // indexes.
|
45 | 46 | PartialIndexMutator MultiStatementMutation = partialIndexMutator
|
46 | 47 |
|
| 48 | + // DupPartialIndexMutator create a new random partial index based on an |
| 49 | + // existing index that is selected. |
| 50 | + DupPartialIndexMutator MultiStatementMutation = duplicatePartialIndexMutator |
| 51 | + |
47 | 52 | // PostgresMutator modifies strings such that they execute identically
|
48 | 53 | // in both Postgres and Cockroach (however this mutator does not remove
|
49 | 54 | // features not supported by Postgres; use PostgresCreateTableMutator
|
@@ -1164,6 +1169,84 @@ func partialIndexMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
|
1164 | 1169 | return stmts, changed
|
1165 | 1170 | }
|
1166 | 1171 |
|
| 1172 | +// duplicatePartialIndexMutator is a mutations.MultiStatementMutator that will |
| 1173 | +// copy a random index definition and make it a partial index. Unlike |
| 1174 | +// partialIndexMutator, the existing definition will stay the same. |
| 1175 | +func duplicatePartialIndexMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Statement, bool) { |
| 1176 | + changed := false |
| 1177 | + tables := getTableInfoFromDDLStatements(stmts) |
| 1178 | + var newStmts []tree.Statement |
| 1179 | + for _, stmt := range stmts { |
| 1180 | + switch ast := stmt.(type) { |
| 1181 | + case *tree.CreateIndex: |
| 1182 | + info, ok := tables[ast.Table.ObjectName] |
| 1183 | + if !ok { |
| 1184 | + continue |
| 1185 | + } |
| 1186 | + |
| 1187 | + // If the index is not already a partial index, make it a partial index |
| 1188 | + // with a 50% chance. Do not mutate an index that was created to satisfy a |
| 1189 | + // FK constraint. |
| 1190 | + if ast.Predicate == nil && |
| 1191 | + !hasReferencingConstraint(info, ast.Columns) && |
| 1192 | + rng.Intn(2) == 0 { |
| 1193 | + astString := tree.AsStringWithFlags(ast, tree.FmtParsableNumerics) |
| 1194 | + astCopy, err := parser.ParseOne(astString) |
| 1195 | + if err != nil { |
| 1196 | + panic(err) |
| 1197 | + } |
| 1198 | + createIndexCopy := astCopy.AST.(*tree.CreateIndex) |
| 1199 | + if len(createIndexCopy.Name) > 0 { |
| 1200 | + createIndexCopy.Name = tree.Name(fmt.Sprintf("%s_partial", createIndexCopy.Name)) |
| 1201 | + } |
| 1202 | + tn := tree.MakeUnqualifiedTableName(createIndexCopy.Table.ObjectName) |
| 1203 | + createIndexCopy.Predicate = randPartialIndexPredicateFromCols(rng, info.columnsTableDefs, &tn) |
| 1204 | + changed = true |
| 1205 | + newStmts = append(newStmts, createIndexCopy) |
| 1206 | + } |
| 1207 | + case *tree.CreateTable: |
| 1208 | + info, ok := tables[ast.Table.ObjectName] |
| 1209 | + if !ok { |
| 1210 | + panic("table info could not be found") |
| 1211 | + } |
| 1212 | + var newDefs []tree.TableDef |
| 1213 | + for _, def := range ast.Defs { |
| 1214 | + var idx *tree.IndexTableDef |
| 1215 | + switch defType := def.(type) { |
| 1216 | + case *tree.IndexTableDef: |
| 1217 | + idx = defType |
| 1218 | + case *tree.UniqueConstraintTableDef: |
| 1219 | + if !defType.PrimaryKey && !defType.WithoutIndex { |
| 1220 | + idx = &defType.IndexTableDef |
| 1221 | + } |
| 1222 | + } |
| 1223 | + |
| 1224 | + if idx == nil { |
| 1225 | + continue |
| 1226 | + } |
| 1227 | + |
| 1228 | + // If the index is not already a partial index, make it a partial |
| 1229 | + // index with a 50% chance. |
| 1230 | + if idx.Predicate == nil && |
| 1231 | + !hasReferencingConstraint(info, idx.Columns) && |
| 1232 | + rng.Intn(2) == 0 { |
| 1233 | + idxCopy := *idx |
| 1234 | + if len(idxCopy.Name) > 0 { |
| 1235 | + idxCopy.Name = tree.Name(fmt.Sprintf("%s_partial", idxCopy.Name)) |
| 1236 | + } |
| 1237 | + tn := tree.MakeUnqualifiedTableName(ast.Table.ObjectName) |
| 1238 | + idxCopy.Predicate = randPartialIndexPredicateFromCols(rng, info.columnsTableDefs, &tn) |
| 1239 | + changed = true |
| 1240 | + newDefs = append(newDefs, &idxCopy) |
| 1241 | + } |
| 1242 | + } |
| 1243 | + ast.Defs = append(ast.Defs, newDefs...) |
| 1244 | + } |
| 1245 | + } |
| 1246 | + stmts = append(stmts, newStmts...) |
| 1247 | + return stmts, changed |
| 1248 | +} |
| 1249 | + |
1167 | 1250 | // hasReferencingConstraint returns true if the tableInfo has any referencing
|
1168 | 1251 | // columns that match idxColumns.
|
1169 | 1252 | func hasReferencingConstraint(info tableInfo, idxColumns tree.IndexElemList) bool {
|
|
0 commit comments