From 72b76d289ed31f5130f86f92df4b85a7e43d0fba Mon Sep 17 00:00:00 2001
From: Faizan Qazi <faizan@cockroachlabs.com>
Date: Tue, 14 Mar 2023 09:17:14 -0400
Subject: [PATCH] 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    | 14 +++---
 pkg/sql/opt/optbuilder/util.go                |  8 ++++
 .../drop_multiple_columns_separate_statements | 44 ++++++++++++++++---
 3 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/pkg/sql/opt/optbuilder/mutation_builder.go b/pkg/sql/opt/optbuilder/mutation_builder.go
index 101164006ae2..8855c14e7a71 100644
--- a/pkg/sql/opt/optbuilder/mutation_builder.go
+++ b/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,
diff --git a/pkg/sql/opt/optbuilder/util.go b/pkg/sql/opt/optbuilder/util.go
index b9755b7a66e9..83b2aa262e01 100644
--- a/pkg/sql/opt/optbuilder/util.go
+++ b/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)
 		}
 	}
diff --git a/pkg/sql/schemachanger/testdata/end_to_end/drop_multiple_columns_separate_statements b/pkg/sql/schemachanger/testdata/end_to_end/drop_multiple_columns_separate_statements
index 4b483137250e..b756b3c5005b 100644
--- a/pkg/sql/schemachanger/testdata/end_to_end/drop_multiple_columns_separate_statements
+++ b/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,18 +14,30 @@ 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);
@@ -30,8 +45,23 @@ 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
@@ -39,9 +69,9 @@ 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;