opt: avoid using virtual columns that are being mutated

Previously, when schema changes involving dropping a virtual
compute column were concurrently run with UPDATE/DELETE, we
would encounter a panic inside the optimizer. Specifically,
because mutated virtual columns would not have information
populated. To address this, this patch similar to SELECT's,
will skip these columns when evaluating DROP's/UPDATE's.

Epic: none

Release note (bug fix): Fixed a bug where concurrent UPDATE/DELETE
while dropping a virtual computed column could fail with a panic.

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -289,9 +289,10 @@ func (mb *mutationBuilder) buildInputForUpdate(
 	mb.fetchScope = mb.b.buildScan(
 		mb.b.addTable(mb.tab, &mb.alias),
 		tableOrdinals(mb.tab, columnKinds{
-			includeMutations: true,
-			includeSystem:    true,
-			includeInverted:  false,
+			includeMutations:                true,
+			includeSystem:                   true,
+			includeInverted:                 false,
+			excludeVirtualComputedMutations: true,
 		}),
 		indexFlags,
 		noRowLocking,
@@ -399,9 +400,10 @@ func (mb *mutationBuilder) buildInputForDelete(
 	mb.fetchScope = mb.b.buildScan(
 		mb.b.addTable(mb.tab, &mb.alias),
 		tableOrdinals(mb.tab, columnKinds{
-			includeMutations: true,
-			includeSystem:    true,
-			includeInverted:  false,
+			includeMutations:                true,
+			includeSystem:                   true,
+			includeInverted:                 false,
+			excludeVirtualComputedMutations: true,
 		}),
 		indexFlags,
 		noRowLocking,

pkg/sql/opt/optbuilder/util.go

@@ -741,6 +741,9 @@ type columnKinds struct {
 
 	// If true, include inverted index columns.
 	includeInverted bool
+
+	// If true, virtual computed column mutations are excluded.
+	excludeVirtualComputedMutations bool
 }
 
 // tableOrdinals returns a slice of ordinals that correspond to table columns of
@@ -758,6 +761,11 @@ func tableOrdinals(tab cat.Table, k columnKinds) []int {
 	for i := 0; i < n; i++ {
 		col := tab.Column(i)
 		if shouldInclude[col.Kind()] {
+			// This behaviour matches the Builder.addComputedColsForTable, which
+			// will exclude any computed columns in mutation.
+			if k.excludeVirtualComputedMutations && col.IsVirtualComputed() && col.IsMutation() {
+				continue
+			}
 			ordinals = append(ordinals, i)
 		}
 	}

pkg/sql/schemachanger/testdata/end_to_end/drop_multiple_columns_separate_statements

@@ -1,5 +1,8 @@
 setup
 CREATE TABLE t (i INT PRIMARY KEY, j INT, k INT DEFAULT 32 ON UPDATE 42, INDEX((j+1), k));
+INSERT INTO t VALUES(-1);
+INSERT INTO t VALUES(-2);
+INSERT INTO t VALUES(-3);
 ----
 ...
 +object {100 101 t} -> 104
@@ -11,37 +14,64 @@ INSERT INTO t (i, j, k) VALUES($stageKey, $stageKey, $stageKey);
 pq: column "j" does not exist
 
 stage-exec phase=PostCommitPhase stage=:
-INSERT INTO t (i) VALUES($stageKey);
+INSERT INTO t VALUES($stageKey);
+INSERT INTO t VALUES($stageKey + 1);
+DELETE FROM t WHERE i=-1;
+DELETE FROM t WHERE i=$stageKey;
+INSERT INTO t VALUES($stageKey);
+INSERT INTO t VALUES(-1);
 ----
 
+stage-exec phase=PostCommitPhase stage=:
+UPDATE t SET k=$stageKey;
+----
+pq: column "k" does not exist
+
 stage-exec phase=PostCommitPhase stage=:
 SELECT j+1, k FROM t
 ----
 pq: column "j" does not exist
 
-stage-exec phase=PostCommitPhase stage=:
-SELECT count(i) FROM t
+# Each insert will be injected twice per stage, plus 3 injected
+# at the start.
+stage-query phase=PostCommitPhase stage=:
+SELECT count(*)=($successfulStageCount*2)+3 FROM t;
 ----
-$successfulStageCount
+true
 
 stage-exec phase=PostCommitNonRevertiblePhase stage=:
 INSERT INTO t (i, j, k) VALUES($stageKey, $stageKey, $stageKey);
 ----
 pq: column "j" does not exist
 
 stage-exec phase=PostCommitNonRevertiblePhase stage=:
-INSERT INTO t (i) VALUES($stageKey);
+INSERT INTO t VALUES($stageKey);
+INSERT INTO t VALUES($stageKey + 1);
+DELETE FROM t WHERE i=-1;
+DELETE FROM t WHERE i=$stageKey;
+INSERT INTO t VALUES($stageKey);
+INSERT INTO t VALUES(-1);
+----
+
+stage-exec phase=PostCommitNonRevertiblePhase stage=:
+UPDATE t SET k=$stageKey;
+----
+pq: column "k" does not exist
+
+stage-exec phase=PostCommitNonRevertiblePhase stage=:
+SELECT j+1, k FROM t
 ----
+pq: column "j" does not exist
 
 stage-exec phase=PostCommitNonRevertiblePhase stage=:
 SELECT j+1, k FROM t
 ----
 pq: column "j" does not exist
 
 stage-exec phase=PostCommitNonRevertiblePhase stage=:
-SELECT count(i) FROM t
+SELECT count(*)=($successfulStageCount*2)+3 FROM t;
 ----
-$successfulStageCount
+true
 
 test
 ALTER TABLE t DROP COLUMN j CASCADE;