From 79e131c92cadf5c95c9b603cbf37f0f5c5c705ba Mon Sep 17 00:00:00 2001
From: Faizan Qazi <faizan@cockroachlabs.com>
Date: Mon, 21 Mar 2022 12:06:36 -0400
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 | 44 +++++++++++++-------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/pkg/workload/schemachange/error_screening.go b/pkg/workload/schemachange/error_screening.go
index 4b59c6fa30ed..c693266a0f5e 100644
--- a/pkg/workload/schemachange/error_screening.go
+++ b/pkg/workload/schemachange/error_screening.go
@@ -177,16 +177,25 @@ 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 +291,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