workload/schemachanger: handle foreign key mutations

Previously, the insert operation could fail on foreign key violation
errors, if the foreign key constraint was enforce using a mutation. This
patch detects those scenarios and turns these errors into potential
errors.

Fixes: #148070

Release note: None

Author: fqazi

pkg/workload/schemachange/error_screening.go

@@ -1726,3 +1726,32 @@ func (og *operationGenerator) tableHasUniqueConstraintMutation(
 			AND (m->'index'->>'unique')::BOOL IS TRUE
 		);`, tableName)
 }
+
+// tableHasForeignKeyMutation determines if a table has any foreign key constraint
+// mutation ongoing. This means either being added or dropped.
+func (og *operationGenerator) tableHasForeignKeyMutation(
+	ctx context.Context, tx pgx.Tx, tableName *tree.TableName,
+) (bool, error) {
+	return og.scanBool(ctx, tx, `
+		WITH table_desc AS (
+			SELECT crdb_internal.pb_to_json(
+				'desc',
+				descriptor,
+				false
+			)->'table' as d
+			FROM system.descriptor
+			WHERE id = $1::REGCLASS
+		)
+		SELECT EXISTS (
+			SELECT * FROM (
+			SELECT jsonb_array_elements(
+				CASE WHEN d->'mutations' IS NULL
+				THEN '[]'::JSONB
+				ELSE d->'mutations'
+				END
+			) as m
+			FROM table_desc)
+			WHERE (m->>'direction')::STRING IN ('ADD', 'DROP')
+			AND (m->'constraint'->>'foreign_key') IS NOT NULL
+		);`, tableName)
+}

pkg/workload/schemachange/operation_generator.go

@@ -3082,6 +3082,7 @@ func (og *operationGenerator) insertRow(ctx context.Context, tx pgx.Tx) (stmt *o
 	// we will need to evaluate generated expressions below.
 	hasUniqueConstraints := false
 	hasUniqueConstraintsMutations := false
+	hasForeignKeyMutations := false
 	fkViolation := false
 	if !anyInvalidInserts {
 		// Verify if the new row may violate unique constraints by checking the
@@ -3103,18 +3104,25 @@ func (og *operationGenerator) insertRow(ctx context.Context, tx pgx.Tx) (stmt *o
 		if err != nil {
 			return nil, err
 		}
-
+		// Determine if a foreign key mutation exists.
+		hasForeignKeyMutations, err = og.tableHasForeignKeyMutation(ctx, tx, tableName)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	stmt.potentialExecErrors.addAll(codesWithConditions{
 		{code: pgcode.UniqueViolation, condition: hasUniqueConstraints || hasUniqueConstraintsMutations},
-		{code: pgcode.ForeignKeyViolation, condition: fkViolation},
+		{code: pgcode.ForeignKeyViolation, condition: fkViolation || hasForeignKeyMutations},
 		{code: pgcode.NotNullViolation, condition: true},
 		{code: pgcode.CheckViolation, condition: true},
 		{code: pgcode.InsufficientPrivilege, condition: true}, // For RLS violations
 	})
 	og.expectedCommitErrors.addAll(codesWithConditions{
-		{code: pgcode.ForeignKeyViolation, condition: fkViolation},
+		{code: pgcode.ForeignKeyViolation, condition: fkViolation && !hasForeignKeyMutations},
+	})
+	og.potentialCommitErrors.addAll(codesWithConditions{
+		{code: pgcode.ForeignKeyViolation, condition: hasForeignKeyMutations},
 	})
 
 	var formattedRows []string