Merge pull request cockroachdb#32591 from vivekmenezes/backport2.1-32188

backport-2.1: sql: fix panic on UPDATE RETURNING * in middle of schema change

Author: vivekmenezes

pkg/sql/logictest/testdata/logic_test/delete

@@ -187,17 +187,35 @@ SELECT count(*) FROM [DELETE FROM unindexed LIMIT 5 RETURNING v]
 subtest regression_29494
 
 statement ok
-CREATE TABLE t29494(x INT); BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
+CREATE TABLE t29494(x INT PRIMARY KEY); INSERT INTO t29494 VALUES (12)
+
+statement ok
+BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
 
 # Check that the new column is not visible
 query T
 SELECT create_statement FROM [SHOW CREATE t29494]
 ----
 CREATE TABLE t29494 (
-   x INT NULL,
-   FAMILY "primary" (x, rowid)
+   x INT NOT NULL,
+   CONSTRAINT "primary" PRIMARY KEY (x ASC),
+   FAMILY "primary" (x)
 )
 
 # Check that the new column is not usable in RETURNING
 statement error column "y" does not exist
 DELETE FROM t29494 RETURNING y
+
+statement ok
+ROLLBACK
+
+statement ok
+BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
+
+query I
+DELETE FROM t29494 RETURNING *
+----
+12
+
+statement ok
+COMMIT

pkg/sql/logictest/testdata/logic_test/insert

@@ -580,7 +580,10 @@ INSERT INTO derived(y) VALUES('abcd')
 subtest regression_29494
 
 statement ok
-CREATE TABLE t29494(x INT); BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
+CREATE TABLE t29494(x INT); INSERT INTO t29494 VALUES (12)
+
+statement ok
+BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
 
 # Check that the new column is not visible
 query T
@@ -594,3 +597,17 @@ CREATE TABLE t29494 (
 # Check that the new column is not usable in RETURNING
 statement error column "y" does not exist
 INSERT INTO t29494(x) VALUES (123) RETURNING y
+
+statement ok
+ROLLBACK
+
+statement ok
+BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
+
+query I
+INSERT INTO t29494(x) VALUES (123) RETURNING *
+----
+123
+
+statement ok
+COMMIT

pkg/sql/logictest/testdata/logic_test/update

@@ -437,15 +437,19 @@ UPDATE derived SET y = 'abcd'
 subtest regression_29494
 
 statement ok
-CREATE TABLE t29494(x INT); BEGIN; ALTER TABLE t29494 ADD COLUMN y INT
+CREATE TABLE t29494(x INT PRIMARY KEY); INSERT INTO t29494 VALUES (12)
+
+statement ok
+BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
 
 # Check that the new column is not visible
 query T
 SELECT create_statement FROM [SHOW CREATE t29494]
 ----
 CREATE TABLE t29494 (
-   x INT NULL,
-   FAMILY "primary" (x, rowid)
+   x INT NOT NULL,
+   CONSTRAINT "primary" PRIMARY KEY (x ASC),
+   FAMILY "primary" (x)
 )
 
 # Check that the new column is not usable in RETURNING
@@ -457,6 +461,12 @@ UPDATE t29494 SET x = 123 RETURNING y
 statement ok
 ROLLBACK; BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
 
+# Returning * doesn't return y
+query I
+UPDATE t29494 SET x = 124 WHERE x = 12 RETURNING *
+----
+124
+
 statement error column "y" does not exist
 UPDATE t29494 SET y = 123
 

pkg/sql/logictest/testdata/logic_test/upsert

@@ -622,7 +622,10 @@ UPSERT INTO tn2(x, y) VALUES (123, 'abcd')
 subtest regression_29494
 
 statement ok
-CREATE TABLE t29494(x INT); BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
+CREATE TABLE t29494(x INT); INSERT INTO t29494 VALUES (12)
+
+statement ok
+BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
 
 # Check that the new column is not visible
 query T
@@ -647,6 +650,27 @@ INSERT INTO t29494(x) VALUES (123) ON CONFLICT(rowid) DO UPDATE SET x = 400 RETU
 statement ok
 ROLLBACK
 
+statement ok
+BEGIN; ALTER TABLE t29494 ADD COLUMN y INT NOT NULL DEFAULT 123
+
+query I
+UPSERT INTO t29494(x) VALUES (12) RETURNING *
+----
+12
+
+query I
+UPSERT INTO t29494(x) VALUES (123) RETURNING *
+----
+123
+
+query I
+INSERT INTO t29494(x) VALUES (123) ON CONFLICT(rowid) DO UPDATE SET x = 400 RETURNING *
+----
+123
+
+statement ok
+COMMIT
+
 subtest regression_31255
 
 statement ok

pkg/sql/update.go

@@ -275,7 +275,7 @@ func (p *planner) Update(
 		// If rowsNeeded is set, we have requested from the source above
 		// all the columns from the descriptor. However, to ensure that
 		// modified rows include all columns, the construction of the
-		// source has used publicAndNonVivible columns so the source may
+		// source has used publicAndNonPublicColumns so the source may
 		// contain additional columns for every newly added column not yet
 		// visible.
 		// We do not want these to be available for RETURNING below.
@@ -675,7 +675,15 @@ func (u *updateNode) processSourceRow(params runParams, sourceVals tree.Datums)
 
 	// If result rows need to be accumulated, do it.
 	if u.run.rows != nil {
-		if _, err := u.run.rows.AddRow(params.ctx, newValues); err != nil {
+		// The new values can include all columns, the construction of the
+		// values has used publicAndNonPublicColumns so the values may
+		// contain additional columns for every newly added column not yet
+		// visible. We do not want them to be available for RETURNING.
+		//
+		// MakeUpdater guarantees that the first columns of the new values
+		// are those specified u.columns.
+		resultValues := newValues[:len(u.columns)]
+		if _, err := u.run.rows.AddRow(params.ctx, resultValues); err != nil {
 			return err
 		}
 	}