workload/schemachange: only allow PK change in single statement transactions

This will avoid errors of the form: "cannot perform other schema changes
in the same transaction as a primary key change."

Release note: None

Author: rafiss

pkg/workload/schemachange/operation_generator.go

@@ -182,7 +182,7 @@ func (og *operationGenerator) getSupportedDeclarativeOp(
 // stochastic attempts and if verbosity is >= 2 the unsuccessful attempts are
 // recorded in `log` to help with debugging of the workload.
 func (og *operationGenerator) randOp(
-	ctx context.Context, tx pgx.Tx, useDeclarativeSchemaChanger, allowDML bool,
+	ctx context.Context, tx pgx.Tx, useDeclarativeSchemaChanger bool, numOpsInTxn int,
 ) (stmt *opStmt, err error) {
 	for {
 		var op opType
@@ -195,8 +195,11 @@ func (og *operationGenerator) randOp(
 		} else {
 			op = opType(og.params.ops.Int())
 		}
-		if !allowDML && isDMLOpType(op) {
-			continue
+		if numOpsInTxn != 1 {
+			// DML and legacy PK changes are only allowed in single-statement transactions.
+			if isDMLOpType(op) || (op == alterTableAlterPrimaryKey && !useDeclarativeSchemaChanger) {
+				continue
+			}
 		}
 
 		og.resetOpState(useDeclarativeSchemaChanger)

pkg/workload/schemachange/schemachange.go

@@ -456,12 +456,12 @@ func (w *schemaChangeWorker) runInTxn(
 ) error {
 	w.logger.startLog(w.id)
 	w.logger.writeLog("BEGIN")
-	opsNum := 1 + w.opGen.randIntn(w.maxOpsPerWorker)
-	if useDeclarativeSchemaChanger && opsNum > w.workload.declarativeSchemaMaxStmtsPerTxn {
-		opsNum = w.workload.declarativeSchemaMaxStmtsPerTxn
+	numOps := 1 + w.opGen.randIntn(w.maxOpsPerWorker)
+	if useDeclarativeSchemaChanger && numOps > w.workload.declarativeSchemaMaxStmtsPerTxn {
+		numOps = w.workload.declarativeSchemaMaxStmtsPerTxn
 	}
 
-	for i := 0; i < opsNum; i++ {
+	for i := 0; i < numOps; i++ {
 		// Terminating this loop early if there are expected commit errors prevents unexpected commit behavior from being
 		// hidden by subsequent operations. Consider the case where there are expected commit errors.
 		// It is possible that committing the transaction now will fail the workload because the error does not occur
@@ -474,8 +474,7 @@ func (w *schemaChangeWorker) runInTxn(
 			break
 		}
 
-		// Only allow DML for single-statement transactions.
-		op, err := w.opGen.randOp(ctx, tx, useDeclarativeSchemaChanger, opsNum == 1)
+		op, err := w.opGen.randOp(ctx, tx, useDeclarativeSchemaChanger, numOps)
 		if pgErr := new(pgconn.PgError); errors.As(err, &pgErr) &&
 			pgcode.MakeCode(pgErr.Code) == pgcode.SerializationFailure {
 			return errors.Mark(err, errRunInTxnRbkSentinel)