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

…-24.3-136738 release-24.3: workload/schemachanger: fix drop columns during PK swaps
cockroachdb · Dec 5, 2024 · d68d1b3 · d68d1b3
2 parents 5cbaf79 + 6e19976
commit d68d1b3
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 39 deletions.
diff --git a/pkg/workload/schemachange/error_screening.go b/pkg/workload/schemachange/error_screening.go
@@ -1417,35 +1417,42 @@ func (og *operationGenerator) tableHasOngoingAlterPKSchemaChanges(
 		ctx,
 		tx,
 		`
-WITH
-	descriptors
-		AS (
-			SELECT
-				crdb_internal.pb_to_json(
-					'cockroach.sql.sqlbase.Descriptor',
-					descriptor
-				)->'table'
-					AS d
-			FROM
-				system.descriptor
-			WHERE
-				id = $1::REGCLASS
-		)
-SELECT
-	EXISTS(
-		SELECT
-			mut
-		FROM
-			(
-				SELECT
-					json_array_elements(d->'mutations')
-						AS mut
-				FROM
-					descriptors
-			)
-		WHERE
-			(mut->'primaryKeySwap') IS NOT NULL
-	);
+WITH descriptors AS (
+                    SELECT crdb_internal.pb_to_json(
+                            'cockroach.sql.sqlbase.Descriptor',
+                            descriptor
+                           )->'table' AS d
+                      FROM system.descriptor
+                     WHERE id = $1::REGCLASS
+                 ),
+     mutations AS (
+                SELECT json_array_elements(
+                        d->'mutations'
+                       ) AS m
+                  FROM descriptors
+               ),
+     primaryindex AS (
+                    SELECT d->'primaryIndex' AS p
+                      FROM descriptors
+                  )
+		 -- Check for legacy primary key swaps which exist as mutations
+     SELECT EXISTS(
+                SELECT mut
+                  FROM (
+                        SELECT json_array_elements(
+                                d->'mutations'
+                               ) AS mut
+                          FROM descriptors
+                       )
+                 WHERE (mut->'primaryKeySwap') IS NOT NULL
+		 -- Check for declarative primary key swaps, which will appear as
+		 -- as new primary indexes with different key columns
+     UNION SELECT p
+             FROM primaryindex AS pk, mutations AS mut
+            WHERE m->'index'->'encodingType' = '1'::JSONB
+                  AND m->'index'->'keyColumnIds'
+                    != p->'keyColumnIds'
+			);
 		`,
 		tableName.String(),
 	)

diff --git a/pkg/workload/schemachange/operation_generator.go b/pkg/workload/schemachange/operation_generator.go
@@ -307,7 +307,7 @@ func (og *operationGenerator) addColumn(ctx context.Context, tx pgx.Tx) (*opStmt
 		{code: pgcode.DuplicateColumn, condition: columnExistsOnTable},
 		{code: pgcode.UndefinedObject, condition: typ == nil},
 		{code: pgcode.NotNullViolation, condition: hasRows && def.Nullable.Nullability == tree.NotNull},
-		{code: pgcode.FeatureNotSupported, condition: hasAlterPKSchemaChange},
+		{code: pgcode.FeatureNotSupported, condition: hasAlterPKSchemaChange && !og.useDeclarativeSchemaChanger},
 		// UNIQUE is only supported for indexable types.
 		{
 			code:      pgcode.FeatureNotSupported,
@@ -399,7 +399,7 @@ func (og *operationGenerator) addUniqueConstraint(ctx context.Context, tx pgx.Tx
 		{code: pgcode.UndefinedColumn, condition: !columnExistsOnTable},
 		{code: pgcode.DuplicateRelation, condition: constraintExists},
 		{code: pgcode.FeatureNotSupported, condition: columnExistsOnTable && !colinfo.ColumnTypeIsIndexable(columnForConstraint.typ)},
-		{pgcode.FeatureNotSupported, hasAlterPKSchemaChange},
+		{pgcode.FeatureNotSupported, hasAlterPKSchemaChange && !og.useDeclarativeSchemaChanger},
 		{code: pgcode.ObjectNotInPrerequisiteState, condition: databaseHasRegionChange && tableIsRegionalByRow},
 	})
 
@@ -1147,7 +1147,7 @@ func (og *operationGenerator) createIndex(ctx context.Context, tx pgx.Tx) (*opSt
 			{code: pgcode.FeatureNotSupported, condition: regionColStored},
 			{code: pgcode.FeatureNotSupported, condition: duplicateRegionColumn},
 			{code: pgcode.FeatureNotSupported, condition: isStoringVirtualComputed},
-			{code: pgcode.FeatureNotSupported, condition: hasAlterPKSchemaChange},
+			{code: pgcode.FeatureNotSupported, condition: hasAlterPKSchemaChange && !og.useDeclarativeSchemaChanger},
 			{code: pgcode.FeatureNotSupported, condition: lastColInvertedIndexIsDescending},
 			{code: pgcode.FeatureNotSupported, condition: pkColUsedInInvertedIndex},
 		})
@@ -1679,13 +1679,16 @@ func (og *operationGenerator) dropColumn(ctx context.Context, tx pgx.Tx) (*opStm
 		{code: pgcode.UndefinedColumn, condition: !columnExists},
 		{code: pgcode.InvalidColumnReference, condition: colIsPrimaryKey || colIsRefByComputed},
 		{code: pgcode.DependentObjectsStillExist, condition: columnIsDependedOn || columnRemovalWillDropFKBackingIndexes},
-		{code: pgcode.FeatureNotSupported, condition: hasAlterPKSchemaChange},
+		{code: pgcode.FeatureNotSupported, condition: hasAlterPKSchemaChange && !og.useDeclarativeSchemaChanger},
+	})
+	stmt.potentialExecErrors.addAll(codesWithConditions{
+		// For legacy schema changer its possible for create index operations to interfere if they
+		// are in progress.
+		{code: pgcode.ObjectNotInPrerequisiteState, condition: !og.useDeclarativeSchemaChanger},
+		// It is possible the column we are dropping is in the new primary key,
+		// so a potential error is an invalid reference in this case.
+		{code: pgcode.InvalidColumnReference, condition: og.useDeclarativeSchemaChanger && hasAlterPKSchemaChange},
 	})
-	// For legacy schema changer its possible for create index operations to interfere if they
-	// are in progress.
-	if !og.useDeclarativeSchemaChanger {
-		stmt.potentialExecErrors.add(pgcode.ObjectNotInPrerequisiteState)
-	}
 	stmt.sql = fmt.Sprintf(`ALTER TABLE %s DROP COLUMN %s`, tableName.String(), columnName.String())
 	return stmt, nil
 }
@@ -1923,7 +1926,7 @@ func (og *operationGenerator) dropIndex(ctx context.Context, tx pgx.Tx) (*opStmt
 	if err != nil {
 		return nil, err
 	}
-	if hasAlterPKSchemaChange {
+	if hasAlterPKSchemaChange && !og.useDeclarativeSchemaChanger {
 		stmt.expectedExecErrors.add(pgcode.FeatureNotSupported)
 	}