From 8b8b5983e58a030109526b4df3ecd350bd0ebd89 Mon Sep 17 00:00:00 2001 From: Faizan Qazi Date: Wed, 4 Dec 2024 12:36:36 -0500 Subject: [PATCH] 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: #134056 Fixes: #136661 Fixes: #136405 Fixes: #132298 Release note: None --- pkg/workload/schemachange/error_screening.go | 65 ++++++++++--------- .../schemachange/operation_generator.go | 15 +++-- 2 files changed, 46 insertions(+), 34 deletions(-) diff --git a/pkg/workload/schemachange/error_screening.go b/pkg/workload/schemachange/error_screening.go index 617ba34a1921..9912fd7b8d57 100644 --- 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 index a088037334d9..bdc6fbb1211c 100644 --- 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) }