release-21.2: sql: disallow partial index creation referencing new column in predicate

pkg/sql/catalog/descriptor.go

@@ -353,6 +353,12 @@ type TableDescriptor interface {
 	GetRegionalByRowTableRegionColumnName() (tree.Name, error)
 }
 
+// MutableTableDescriptor is both a MutableDescriptor and a TableDescriptor.
+type MutableTableDescriptor interface {
+	TableDescriptor
+	MutableDescriptor
+}
+
 // TypeDescriptor will eventually be called typedesc.Descriptor.
 // It is implemented by (Imm|M)utableTypeDescriptor.
 type TypeDescriptor interface {

pkg/sql/catalog/schemaexpr/partial_index.go

@@ -16,9 +16,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/transform"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/errors"
 )
 
 // ValidatePartialIndexPredicate verifies that an expression is a valid partial
@@ -32,15 +35,17 @@ import (
 //   - It does not include subqueries.
 //   - It does not include non-immutable, aggregate, window, or set returning
 //     functions.
+//   - It does not reference a column which is in the process of being added
+//     or removed.
 //
 func ValidatePartialIndexPredicate(
 	ctx context.Context,
-	desc catalog.TableDescriptor,
+	desc catalog.MutableTableDescriptor,
 	e tree.Expr,
 	tn *tree.TableName,
 	semaCtx *tree.SemaContext,
 ) (string, error) {
-	expr, _, _, err := DequalifyAndValidateExpr(
+	expr, _, cols, err := DequalifyAndValidateExpr(
 		ctx,
 		desc,
 		e,
@@ -53,10 +58,38 @@ func ValidatePartialIndexPredicate(
 	if err != nil {
 		return "", err
 	}
-
+	if !desc.IsNew() {
+		if err := validatePartialIndexExprColsArePublic(desc, cols); err != nil {
+			return "", err
+		}
+	}
 	return expr, nil
 }
 
+func validatePartialIndexExprColsArePublic(
+	desc catalog.TableDescriptor, cols catalog.TableColSet,
+) (err error) {
+	cols.ForEach(func(colID descpb.ColumnID) {
+		if err != nil {
+			return
+		}
+		var col catalog.Column
+		col, err = desc.FindColumnWithID(colID)
+		if err != nil {
+			return
+		}
+		if col.Public() {
+			return
+		}
+		err = pgerror.WithCandidateCode(errors.Errorf(
+			"cannot create partial index on column %q (%d) which is not public",
+			col.GetName(), col.GetID()),
+			pgcode.FeatureNotSupported,
+		)
+	})
+	return err
+}
+
 // MakePartialIndexExprs returns a map of predicate expressions for each
 // partial index in the input list of indexes, or nil if none of the indexes
 // are partial indexes. It also returns a set of all column IDs referenced in

pkg/sql/catalog/schemaexpr/testutils_test.go

@@ -28,7 +28,7 @@ type testCol struct {
 // less verbose way.
 func testTableDesc(
 	name string, columns []testCol, mutationColumns []testCol,
-) catalog.TableDescriptor {
+) catalog.MutableTableDescriptor {
 	cols := make([]descpb.ColumnDescriptor, len(columns))
 	for i := range columns {
 		cols[i] = descpb.ColumnDescriptor{
@@ -57,5 +57,5 @@ func testTableDesc(
 		ID:        1,
 		Columns:   cols,
 		Mutations: muts,
-	}).BuildImmutableTable()
+	}).BuildCreatedMutableTable()
 }

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -526,31 +526,6 @@ SELECT * FROM f@a_b_gt_30_idx WHERE b > 30
 ----
 6  60
 
-# Backfill a partial index with a reference to a new column in the predicate.
-
-statement ok
-CREATE TABLE g (a INT)
-
-statement ok
-INSERT INTO g VALUES (1)
-
-statement ok
-BEGIN
-
-statement ok
-ALTER TABLE g ADD COLUMN b INT
-
-statement ok
-CREATE INDEX a_b_null_idx ON g (a) WHERE b IS NULL
-
-statement ok
-COMMIT
-
-query II rowsort
-SELECT * FROM g@a_b_null_idx WHERE b IS NULL
-----
-1  NULL
-
 # Backfill a partial index with a user defined type.
 
 statement ok
@@ -1832,3 +1807,77 @@ query ITTT
 SELECT * FROM t74385@b_idx
 WHERE b = 'b' AND c IS NULL;
 ----
+
+# Regression test for #79613 to disallow adding a partial index referencing a
+# a column added in the same transaction.
+subtest column_added_in_same_transaction
+
+statement ok
+CREATE TABLE t79613 (i INT PRIMARY KEY);
+
+statement error pgcode 0A000 cannot create partial index on column "k" \(2\) which is not public
+BEGIN;
+   ALTER TABLE t79613 ADD COLUMN k INT DEFAULT 1;
+   CREATE INDEX idx ON t79613(i) WHERE (k > 1);
+
+statement ok
+ROLLBACK
+
+statement error pgcode 0A000 cannot create partial index on column "k" \(2\) which is not public
+BEGIN;
+   ALTER TABLE t79613 ADD COLUMN k INT DEFAULT 1;
+   CREATE UNIQUE INDEX idx ON t79613(i) WHERE (k > 1);
+
+statement ok
+ROLLBACK
+
+statement error pgcode 0A000 cannot create partial index on column "k" \(2\) which is not public
+BEGIN;
+   ALTER TABLE t79613 ADD COLUMN k INT DEFAULT 1;
+   ALTER TABLE t79613 ADD CONSTRAINT c UNIQUE (k) WHERE (k > 1);
+
+statement ok
+ROLLBACK
+
+statement ok
+DROP TABLE t79613
+
+# If the table was created in the same transaction, then it's fine to add
+# these partial indexes because we'll have backfilled the new column
+# synchronously.
+
+statement ok
+BEGIN;
+   CREATE TABLE t79613 (i INT PRIMARY KEY);
+   ALTER TABLE t79613 ADD COLUMN k INT DEFAULT 1;
+   CREATE INDEX idx ON t79613(i) WHERE (k > 1);
+COMMIT;
+SELECT * FROM t79613;
+DROP TABLE t79613;
+
+statement ok
+BEGIN;
+   CREATE TABLE t79613 (i INT PRIMARY KEY);
+   ALTER TABLE t79613 ADD COLUMN k INT DEFAULT 1;
+   CREATE UNIQUE INDEX idx ON t79613(i) WHERE (k > 1);
+COMMIT;
+SELECT * FROM t79613;
+DROP TABLE t79613;
+
+statement ok
+BEGIN;
+   CREATE TABLE t79613 (i INT PRIMARY KEY);
+   ALTER TABLE t79613 ADD COLUMN k INT DEFAULT 1;
+   ALTER TABLE t79613 ADD CONSTRAINT c UNIQUE (k) WHERE (k > 1);
+COMMIT;
+SELECT * FROM t79613;
+DROP TABLE t79613;
+
+statement ok
+BEGIN;
+   CREATE TABLE t79613 (i INT PRIMARY KEY);
+   ALTER TABLE t79613 ADD COLUMN k INT DEFAULT 1,
+       ADD CONSTRAINT c UNIQUE (k) WHERE (k > 1);
+COMMIT;
+SELECT * FROM t79613;
+DROP TABLE t79613;

pkg/sql/logictest/testdata/logic_test/virtual_columns

@@ -899,28 +899,22 @@ DROP INDEX v_idx
 statement ok
 ALTER TABLE sc DROP COLUMN v
 
-# Add a column and a partial index using that column in the predicate in the
-# same transaction.
+# Adding a column and a partial index using that column in the predicate in the
+# same transaction is not allowed.
 statement ok
 BEGIN
 
 statement ok
 ALTER TABLE sc ADD COLUMN v INT AS (a+b) VIRTUAL
 
-statement ok
+statement error pgcode 0A000 cannot create partial index on column "v" \(10\) which is not public
 CREATE INDEX partial_idx ON sc(b) WHERE v > 20
 
 statement ok
 END
 
-query I rowsort
-SELECT a FROM sc@partial_idx WHERE v > 20
-----
-2
-3
-
 statement ok
-DROP INDEX partial_idx
+ALTER TABLE sc ADD COLUMN v INT AS (a+b) VIRTUAL
 
 # Create a partial index on the virtual column and which uses the virtual column in the predicate.
 statement ok