Merge pull request #159252 from spilchen/backport24.3-159188

release-24.3: workload: remove FK dependency check in schemachanger drop column

Author: spilchen

pkg/workload/schemachange/error_screening.go

@@ -131,150 +131,6 @@ func (og *operationGenerator) tableHasDependencies(
 	return og.scanBool(ctx, tx, q, tableName.Object(), tableName.Schema())
 }
 
-// columnRemovalWillDropFKBackingIndexes determines if dropping this column
-// will lead to no indexes backing a foreign key.
-func (og *operationGenerator) columnRemovalWillDropFKBackingIndexes(
-	ctx context.Context, tx pgx.Tx, tableName *tree.TableName, columName tree.Name,
-) (bool, error) {
-	return og.scanBool(ctx, tx, fmt.Sprintf(`
-WITH
-	fk
-		AS (
-			SELECT
-				oid,
-				(
-					SELECT
-						r.relname
-					FROM
-						pg_class AS r
-					WHERE
-						r.oid = c.confrelid
-				)
-					AS base_table,
-				a.attname AS base_col,
-				array_position(c.confkey, a.attnum)
-					AS base_ordinal,
-				(
-					SELECT
-						r.relname
-					FROM
-						pg_class AS r
-					WHERE
-						r.oid = c.conrelid
-				)
-					AS referencing_table,
-				unnest(
-					(
-						SELECT
-							array_agg(attname)
-						FROM
-							pg_attribute
-						WHERE
-							attrelid = c.conrelid
-							AND ARRAY[attnum] <@ c.conkey
-							AND array_position(
-									c.confkey,
-									a.attnum
-								)
-								= array_position(
-										c.conkey,
-										attnum
-									)
-					)
-				)
-					AS referencing_col
-			FROM
-				pg_constraint AS c
-				JOIN pg_attribute AS a ON
-						c.confrelid = a.attrelid
-						AND ARRAY[attnum] <@ c.confkey
-			WHERE
-				c.confrelid = $1::REGCLASS::OID
-		),
-	valid_indexes
-		AS (
-			SELECT
-				*
-			FROM
-				[SHOW INDEXES FROM %s]
-			WHERE
-				index_name
-				NOT IN (
-						SELECT
-							DISTINCT index_name
-						FROM
-							[SHOW INDEXES FROM %s]
-						WHERE
-							column_name = $2
-							AND index_name
-								NOT LIKE '%%_pkey' -- renames would keep the old table name
-					)
-		),
-	matching_indexes
-		AS (
-			SELECT
-				oid,
-				index_name,
-				count(base_ordinal) AS count_base_ordinal,
-				count(seq_in_index) AS count_seq_in_index
-			FROM
-				fk, valid_indexes
-			WHERE
-				storing = 'f'
-				AND non_unique = 'f'
-				AND base_col = column_name
-			GROUP BY
-				(oid, index_name)
-		),
-	valid_index_attrib_count
-		AS (
-			SELECT
-				index_name,
-				max(seq_in_index) AS max_seq_in_index
-			FROM
-				valid_indexes
-			WHERE
-				storing = 'f'
-			GROUP BY
-				index_name
-		),
-	valid_fk_count
-		AS (
-			SELECT
-				oid, max(base_ordinal) AS max_base_ordinal
-			FROM
-				fk
-			GROUP BY
-				fk
-		),
-	matching_fks
-		AS (
-			SELECT
-				DISTINCT f.oid
-			FROM
-				valid_index_attrib_count AS i,
-				valid_fk_count AS f,
-				matching_indexes AS m
-			WHERE
-				f.oid = m.oid
-				AND i.index_name = m.index_name
-				AND i.max_seq_in_index
-					= m.count_seq_in_index
-				AND f.max_base_ordinal
-					= m.count_base_ordinal
-		)
-SELECT
-	EXISTS(
-		SELECT
-			*
-		FROM
-			fk
-		WHERE
-			oid NOT IN (SELECT oid FROM matching_fks)
-	);
-`, tableName.String(), tableName.String()), tableName.String(), columName)
-}
-
 func (og *operationGenerator) columnIsDependedOn(
 	ctx context.Context, tx pgx.Tx, tableName *tree.TableName, columnName tree.Name,
 ) (bool, error) {

pkg/workload/schemachange/operation_generator.go

@@ -1652,6 +1652,14 @@ func (og *operationGenerator) dropColumn(ctx context.Context, tx pgx.Tx) (*opStm
 		return nil, err
 	}
 
+	// Check if the table has any foreign keys.
+	tableHasFK := false
+	if tableExists {
+		if tableHasFK, err = og.tableHasFK(ctx, tx, tableName); err != nil {
+			return nil, err
+		}
+	}
+
 	columnName, err := og.randColumn(ctx, tx, *tableName, og.pctExisting(true))
 	if err != nil {
 		return nil, err
@@ -1668,10 +1676,6 @@ func (og *operationGenerator) dropColumn(ctx context.Context, tx pgx.Tx) (*opStm
 	if err != nil {
 		return nil, err
 	}
-	columnRemovalWillDropFKBackingIndexes, err := og.columnRemovalWillDropFKBackingIndexes(ctx, tx, tableName, columnName)
-	if err != nil {
-		return nil, err
-	}
 	columnIsInAddingOrDroppingIndex, err := og.columnIsInAddingOrDroppingIndex(ctx, tx, tableName, columnName)
 	if err != nil {
 		return nil, err
@@ -1690,7 +1694,7 @@ func (og *operationGenerator) dropColumn(ctx context.Context, tx pgx.Tx) (*opStm
 		{code: pgcode.ObjectNotInPrerequisiteState, condition: columnIsInAddingOrDroppingIndex},
 		{code: pgcode.UndefinedColumn, condition: !columnExists},
 		{code: pgcode.InvalidColumnReference, condition: colIsPrimaryKey || colIsRefByComputed},
-		{code: pgcode.DependentObjectsStillExist, condition: columnIsDependedOn || columnRemovalWillDropFKBackingIndexes},
+		{code: pgcode.DependentObjectsStillExist, condition: columnIsDependedOn},
 		{code: pgcode.FeatureNotSupported, condition: hasAlterPKSchemaChange && !og.useDeclarativeSchemaChanger},
 	})
 	stmt.potentialExecErrors.addAll(codesWithConditions{
@@ -1700,11 +1704,27 @@ func (og *operationGenerator) dropColumn(ctx context.Context, tx pgx.Tx) (*opStm
 		// 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},
+		// It is possible that we cannot drop column because
+		// it is depended on by a foreign key.
+		{code: pgcode.DependentObjectsStillExist, condition: tableHasFK},
 	})
 	stmt.sql = fmt.Sprintf(`ALTER TABLE %s DROP COLUMN %s`, tableName.String(), columnName.String())
 	return stmt, nil
 }
 
+// tableHasFK checks if a table participates in any foreign key constraints.
+func (og *operationGenerator) tableHasFK(
+	ctx context.Context, tx pgx.Tx, tableName *tree.TableName,
+) (bool, error) {
+	return og.scanBool(ctx, tx, `
+SELECT EXISTS (
+	SELECT 1
+	  FROM pg_constraint
+	 WHERE contype = 'f'
+	   AND (conrelid = $1::REGCLASS OR confrelid = $1::REGCLASS)
+)`, tableName.String())
+}
+
 func (og *operationGenerator) dropColumnDefault(ctx context.Context, tx pgx.Tx) (*opStmt, error) {
 	tableName, err := og.randTable(ctx, tx, og.pctExisting(true), "")
 	if err != nil {