Merge pull request #78329 from cockroachdb/blathers/backport-release-…

…22.1-78184 release-22.1: workload/schemachanger: primary index detection is incorrect
cockroachdb · Mar 23, 2022 · decb3fa · decb3fa
2 parents 17fb197 + 77ff7d4
commit decb3fa
Showing 1 changed file with 25 additions and 16 deletions.
diff --git a/pkg/workload/schemachange/error_screening.go b/pkg/workload/schemachange/error_screening.go
@@ -177,16 +177,22 @@ func (og *operationGenerator) colIsPrimaryKey(
 	ctx context.Context, tx pgx.Tx, tableName *tree.TableName, columnName string,
 ) (bool, error) {
 	primaryColumns, err := og.scanStringArray(ctx, tx,
-		`SELECT array_agg(column_name)
-		FROM (
-			SELECT DISTINCT column_name
-				FROM information_schema.statistics
-			WHERE (index_name = 'primary' OR index_name LIKE '%pkey%')
-				AND table_schema = $1
-				AND table_name = $2
-				AND storing = 'NO'
-		);
-	`, tableName.Schema(), tableName.Object())
+		`
+SELECT array_agg(column_name)
+  FROM (
+        SELECT DISTINCT column_name
+          FROM information_schema.statistics
+         WHERE index_name
+               IN (
+                  SELECT index_name
+                    FROM crdb_internal.table_indexes
+                   WHERE index_type = 'primary' AND descriptor_id = $3::REGCLASS
+                )
+               AND table_schema = $1
+               AND table_name = $2
+               AND storing = 'NO'
+       );
+	`, tableName.Schema(), tableName.Object(), tableName.String())
 	if err != nil {
 		return false, err
 	}
@@ -282,15 +288,18 @@ func (og *operationGenerator) tableHasPrimaryKeySwapActive(
 	indexName, err := og.scanStringArray(
 		ctx,
 		tx,
-		fmt.Sprintf(`
+		`
 SELECT array_agg(index_name)
   FROM (
-		SELECT index_name
-		  FROM [SHOW INDEXES FROM %s]
-		 WHERE index_name LIKE '%%_pkey%%'
-		 LIMIT 1
+SELECT
+	index_name
+FROM
+	crdb_internal.table_indexes
+WHERE
+	index_type = 'primary'
+	AND descriptor_id = $1::REGCLASS
        );
-	`, tableName.String()),
+	`, tableName.String(),
 	)
 	if err != nil {
 		return err