workload/schemachanger: fix drop columns during PK swaps

Previously, if a declarative alter primary key and drop column were executed concurrently this workload did not correctly pick the correct columns. It was possible for the new primary key column to be selected for the drop operation. To address this, this patch allows potential execution errors for invalid column references when a PK swap is in progress. Fixes: cockroachdb#134056 Fixes: cockroachdb#136661 Fixes: cockroachdb#136405 Fixes: cockroachdb#132298 Release note: None
fqazi · Dec 4, 2024 · 8b8b598 · 8b8b598
1 parent 757b9ab
commit 8b8b598
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 34 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 count(*) > 0
+             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
@@ -306,7 +306,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,
@@ -398,7 +398,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},
 	})
 
@@ -1146,7 +1146,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},
 		})
@@ -1678,13 +1678,18 @@ 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},
 	})
 	// 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)
 	}
+	if og.useDeclarativeSchemaChanger && hasAlterPKSchemaChange {
+		// 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.
+		stmt.potentialExecErrors.add(pgcode.InvalidColumnReference)
+	}
 	stmt.sql = fmt.Sprintf(`ALTER TABLE %s DROP COLUMN %s`, tableName.String(), columnName.String())
 	return stmt, nil
 }
@@ -1922,7 +1927,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)
 	}