From 77ff7d4adb4926eb231fe1cc1390513159f9a880 Mon Sep 17 00:00:00 2001
From: Faizan Qazi <faizan@cockroachlabs.com>
Date: Mon, 21 Mar 2022 16:06:36 +0000
Subject: [PATCH] workload/schemachanger: primary index detection is incorrect

Fixes: #77885

Previously, the schemachanger workload incorrectly detected primary
indexes using their names checking for either the words
'pkey' or 'primary'. This approach was not correct and heuristic
in nature. To address this, this patch uses crdb_internal.table_indexes
to detect primary indexes.

Release note: None
---
 pkg/workload/schemachange/error_screening.go | 41 ++++++++++++--------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/pkg/workload/schemachange/error_screening.go b/pkg/workload/schemachange/error_screening.go
index 4b59c6fa30ed..a76d54c01cb7 100644
--- 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