sql: version gating pkey virtual column validation.

Fixes: cockroachdb#79329

Previously, the schemachange workload on master was failing
because 21.2 does not support virtual columns inside primary
keys. So, the mixed version variant of the schemachange workload
can fail when a primary key with virtual columns exists, specifically
these tables become inaccessible from 21.2 nodes. To address, this
patch introduces a version gate to prevent creating prevent
creating primary indexes with virtual columns and updates the
schemachange workload to detect when this error will be generated
in mixed version workloads.

Release note: None

pkg/sql/alter_primary_key.go

@@ -14,6 +14,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
@@ -134,6 +135,11 @@ func (p *planner) AlterPrimaryKey(
 		if col.IsNullable() {
 			return pgerror.Newf(pgcode.InvalidSchemaDefinition, "cannot use nullable column %q in primary key", col.GetName())
 		}
+		if !p.EvalContext().Settings.Version.IsActive(ctx, clusterversion.Start22_1) {
+			if col.IsVirtual() {
+				return pgerror.Newf(pgcode.FeatureNotSupported, "cannot use virtual column %q in primary key", col.GetName())
+			}
+		}
 	}
 
 	// Validate if the end result is the same as the current

pkg/sql/catalog/schemadesc/synthetic_schema_desc.go

@@ -94,7 +94,8 @@ func (p synthetic) NewBuilder() catalog.DescriptorBuilder {
 func (p synthetic) GetReferencedDescIDs() (catalog.DescriptorIDSet, error) {
 	return catalog.DescriptorIDSet{}, nil
 }
-func (p synthetic) ValidateSelf(_ catalog.ValidationErrorAccumulator) {}
+func (p synthetic) ValidateSelf(_ catalog.ValidationErrorAccumulator) {
+}
 func (p synthetic) ValidateCrossReferences(
 	_ catalog.ValidationErrorAccumulator, _ catalog.ValidationDescGetter,
 ) {

pkg/sql/catalog/tabledesc/validate.go

@@ -13,6 +13,7 @@ package tabledesc
 import (
 	"sort"
 
+	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
@@ -604,7 +605,7 @@ func (desc *wrapper) ValidateSelf(vea catalog.ValidationErrorAccumulator) {
 			desc.validateColumnFamilies(columnIDs),
 			desc.validateCheckConstraints(columnIDs),
 			desc.validateUniqueWithoutIndexConstraints(columnIDs),
-			desc.validateTableIndexes(columnNames),
+			desc.validateTableIndexes(columnNames, vea),
 			desc.validatePartitioning(),
 		}
 		hasErrs := false
@@ -1043,7 +1044,9 @@ func (desc *wrapper) validateUniqueWithoutIndexConstraints(
 // IDs are unique, and the family of the primary key is 0. This does not check
 // if indexes are unique (i.e. same set of columns, direction, and uniqueness)
 // as there are practical uses for them.
-func (desc *wrapper) validateTableIndexes(columnNames map[string]descpb.ColumnID) error {
+func (desc *wrapper) validateTableIndexes(
+	columnNames map[string]descpb.ColumnID, vea catalog.ValidationErrorAccumulator,
+) error {
 	if len(desc.PrimaryIndex.KeyColumnIDs) == 0 {
 		return ErrMissingPrimaryKey
 	}
@@ -1053,6 +1056,15 @@ func (desc *wrapper) validateTableIndexes(columnNames map[string]descpb.ColumnID
 		columnsByID[col.GetID()] = col
 	}
 
+	if !vea.IsActive(clusterversion.Start22_1) {
+		// Verify that the primary index columns are not virtual.
+		for _, pkID := range desc.PrimaryIndex.KeyColumnIDs {
+			if col := columnsByID[pkID]; col != nil && col.IsVirtual() {
+				return errors.Newf("primary index column %q cannot be virtual", col.GetName())
+			}
+		}
+	}
+
 	indexNames := map[string]struct{}{}
 	indexIDs := map[descpb.IndexID]string{}
 	for _, idx := range desc.NonDropIndexes() {
@@ -1204,6 +1216,12 @@ func (desc *wrapper) validateTableIndexes(columnNames map[string]descpb.ColumnID
 			}
 		}
 		for _, colID := range idx.IndexDesc().KeySuffixColumnIDs {
+			if !vea.IsActive(clusterversion.Start22_1) {
+				if col := columnsByID[colID]; col != nil && col.IsVirtual() {
+					return errors.Newf("index %q cannot store virtual column %d", idx.GetName(), colID)
+				}
+			}
+
 			if _, ok := columnsByID[colID]; !ok {
 				return errors.Newf("column %d does not exist in table %s", colID, desc.Name)
 			}

pkg/sql/catalog/tabledesc/validate_test.go

@@ -1901,6 +1901,153 @@ func TestValidateTableDesc(t *testing.T) {
 	}
 }
 
+func TestPrimaryKeyCannotBeVirtualBefore22_1(t *testing.T) {
+	computedExpr := "1 + 1"
+	testData := []struct {
+		err  string
+		desc descpb.TableDescriptor
+	}{
+		{
+			err: `primary index column "c3" cannot be virtual`,
+			desc: descpb.TableDescriptor{
+				ID:            2,
+				ParentID:      1,
+				Name:          "foo",
+				FormatVersion: descpb.InterleavedFormatVersion,
+				Columns: []descpb.ColumnDescriptor{
+					{ID: 1, Name: "c1"},
+					{ID: 2, Name: "c2"},
+					{ID: 3, Name: "c3", ComputeExpr: &computedExpr, Virtual: true},
+				},
+				PrimaryIndex: descpb.IndexDescriptor{
+					ID:                  1,
+					Name:                "primary",
+					Unique:              true,
+					KeyColumnIDs:        []descpb.ColumnID{1, 3},
+					KeyColumnNames:      []string{"c1", "c3"},
+					KeyColumnDirections: []descpb.IndexDescriptor_Direction{descpb.IndexDescriptor_ASC, descpb.IndexDescriptor_ASC},
+					Version:             descpb.LatestIndexDescriptorVersion,
+					EncodingType:        descpb.PrimaryIndexEncoding,
+				},
+				Indexes: []descpb.IndexDescriptor{
+					{ID: 2, Name: "sec", KeyColumnIDs: []descpb.ColumnID{2},
+						KeyColumnNames:      []string{"c2"},
+						KeyColumnDirections: []descpb.IndexDescriptor_Direction{descpb.IndexDescriptor_ASC},
+						KeySuffixColumnIDs:  []descpb.ColumnID{1, 3},
+					},
+				},
+				Families: []descpb.ColumnFamilyDescriptor{
+					{ID: 0, Name: "primary",
+						ColumnIDs:   []descpb.ColumnID{1, 2},
+						ColumnNames: []string{"c1", "c2"},
+					},
+				},
+				Mutations:    []descpb.DescriptorMutation{},
+				NextColumnID: 4,
+				NextFamilyID: 1,
+				NextIndexID:  5,
+				Privileges:   catpb.NewBasePrivilegeDescriptor(security.AdminRoleName()),
+			},
+		},
+		{
+			err: `index "sec" cannot store virtual column 3`,
+			desc: descpb.TableDescriptor{
+				ID:            2,
+				ParentID:      1,
+				Name:          "foo",
+				FormatVersion: descpb.InterleavedFormatVersion,
+				Columns: []descpb.ColumnDescriptor{
+					{ID: 1, Name: "c1"},
+					{ID: 2, Name: "c2"},
+					{ID: 3, Name: "c3", ComputeExpr: &computedExpr, Virtual: true},
+				},
+				PrimaryIndex: descpb.IndexDescriptor{
+					ID:                  1,
+					Name:                "primary",
+					Unique:              true,
+					KeyColumnIDs:        []descpb.ColumnID{1},
+					KeyColumnNames:      []string{"c1"},
+					KeyColumnDirections: []descpb.IndexDescriptor_Direction{descpb.IndexDescriptor_ASC},
+					Version:             descpb.LatestIndexDescriptorVersion,
+					EncodingType:        descpb.PrimaryIndexEncoding,
+				},
+				Indexes: []descpb.IndexDescriptor{
+					{ID: 2, Name: "sec", KeyColumnIDs: []descpb.ColumnID{2},
+						KeyColumnNames:      []string{"c2"},
+						KeyColumnDirections: []descpb.IndexDescriptor_Direction{descpb.IndexDescriptor_ASC},
+						KeySuffixColumnIDs:  []descpb.ColumnID{1, 3},
+					},
+				},
+				Families: []descpb.ColumnFamilyDescriptor{
+					{ID: 0, Name: "primary",
+						ColumnIDs:   []descpb.ColumnID{1, 2},
+						ColumnNames: []string{"c1", "c2"},
+					},
+				},
+				Mutations: []descpb.DescriptorMutation{
+					{
+						Descriptor_: &descpb.DescriptorMutation_Index{
+							Index: &descpb.IndexDescriptor{
+								ID:                  3,
+								Name:                "new_primary_key",
+								Unique:              true,
+								KeyColumnIDs:        []descpb.ColumnID{1, 3},
+								KeyColumnNames:      []string{"c1", "c3"},
+								KeyColumnDirections: []descpb.IndexDescriptor_Direction{descpb.IndexDescriptor_ASC, descpb.IndexDescriptor_ASC},
+								Version:             descpb.LatestIndexDescriptorVersion,
+								EncodingType:        descpb.PrimaryIndexEncoding,
+							},
+						},
+						Direction: descpb.DescriptorMutation_ADD,
+						State:     descpb.DescriptorMutation_DELETE_ONLY,
+					},
+					{
+						Descriptor_: &descpb.DescriptorMutation_Index{
+							Index: &descpb.IndexDescriptor{
+								ID: 4, Name: "new_sec", KeyColumnIDs: []descpb.ColumnID{2},
+								KeyColumnNames:      []string{"c2"},
+								KeyColumnDirections: []descpb.IndexDescriptor_Direction{descpb.IndexDescriptor_ASC},
+								KeySuffixColumnIDs:  []descpb.ColumnID{1, 3},
+							},
+						},
+						Direction: descpb.DescriptorMutation_ADD,
+						State:     descpb.DescriptorMutation_DELETE_ONLY,
+					},
+					{
+						Descriptor_: &descpb.DescriptorMutation_PrimaryKeySwap{
+							PrimaryKeySwap: &descpb.PrimaryKeySwap{
+								OldPrimaryIndexId: 1,
+								NewPrimaryIndexId: 3,
+								NewIndexes:        []descpb.IndexID{4},
+								OldIndexes:        []descpb.IndexID{2},
+							},
+						},
+						Direction: descpb.DescriptorMutation_ADD,
+						State:     descpb.DescriptorMutation_DELETE_ONLY,
+					},
+				},
+				NextColumnID: 4,
+				NextFamilyID: 1,
+				NextIndexID:  5,
+				Privileges:   catpb.NewBasePrivilegeDescriptor(security.AdminRoleName()),
+			},
+		},
+	}
+	for i, d := range testData {
+		t.Run(d.err, func(t *testing.T) {
+			d.desc.Privileges = catpb.NewBasePrivilegeDescriptor(security.RootUserName())
+			desc := NewBuilder(&d.desc).BuildImmutableTable()
+			expectedErr := fmt.Sprintf("%s %q (%d): %s", desc.DescriptorType(), desc.GetName(), desc.GetID(), d.err)
+			err := validate.Self(clusterversion.ClusterVersion{Version: clusterversion.ByKey(clusterversion.V21_2)}, desc)
+			if err == nil {
+				t.Errorf("%d: expected \"%s\", but found success: %+v", i, expectedErr, d.desc)
+			} else if expectedErr != err.Error() {
+				t.Errorf("%d: expected \"%s\", but found \"%+v\"", i, expectedErr, err)
+			}
+		})
+	}
+}
+
 func TestValidateCrossTableReferences(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	ctx := context.Background()

pkg/sql/catalog/typedesc/table_implicit_record_type.go

@@ -205,7 +205,8 @@ func (v TableImplicitRecordType) GetReferencedDescIDs() (catalog.DescriptorIDSet
 }
 
 // ValidateSelf implements the Descriptor interface.
-func (v TableImplicitRecordType) ValidateSelf(_ catalog.ValidationErrorAccumulator) {}
+func (v TableImplicitRecordType) ValidateSelf(_ catalog.ValidationErrorAccumulator) {
+}
 
 // ValidateCrossReferences implements the Descriptor interface.
 func (v TableImplicitRecordType) ValidateCrossReferences(

pkg/sql/create_table.go

@@ -1547,6 +1547,7 @@ func NewTableDesc(
 				if err != nil {
 					return nil, err
 				}
+				primaryIndexColumnSet[shardCol.GetName()] = struct{}{}
 				checkConstraint, err := makeShardCheckConstraintDef(int(buckets), shardCol)
 				if err != nil {
 					return nil, err
@@ -1598,6 +1599,7 @@ func NewTableDesc(
 			if err := desc.AddPrimaryIndex(*implicitColumnDefIdx.idx); err != nil {
 				return nil, err
 			}
+			primaryIndexColumnSet[string(implicitColumnDefIdx.def.Name)] = struct{}{}
 		} else {
 			// If it is a non-primary index that is implicitly created, ensure
 			// partitioning for PARTITION ALL BY.
@@ -1975,6 +1977,14 @@ func NewTableDesc(
 
 	for i := range desc.Columns {
 		if _, ok := primaryIndexColumnSet[desc.Columns[i].Name]; ok {
+			if !st.Version.IsActive(ctx, clusterversion.Start22_1) {
+				if desc.Columns[i].Virtual {
+					return nil, pgerror.Newf(
+						pgcode.FeatureNotSupported,
+						"cannot use virtual column %q in primary key", desc.Columns[i].Name,
+					)
+				}
+			}
 			desc.Columns[i].Nullable = false
 		}
 	}

pkg/sql/schema_changer_test.go

@@ -8229,3 +8229,34 @@ DROP VIEW IF EXISTS v
 	}
 	wg.Wait()
 }
+
+func TestVirtualColumnNotAllowedInPkeyBefore22_1(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+
+	params, _ := tests.CreateTestServerParams()
+	params.Knobs.Server = &server.TestingKnobs{
+		DisableAutomaticVersionUpgrade: make(chan struct{}),
+		BinaryVersionOverride:          clusterversion.ByKey(clusterversion.V21_2),
+	}
+
+	s, sqlDB, _ := serverutils.StartServer(t, params)
+	defer s.Stopper().Stop(ctx)
+
+	_, err := sqlDB.Exec(`CREATE TABLE t (a INT NOT NULL AS (1+1) VIRTUAL, PRIMARY KEY (a))`)
+	require.Error(t, err)
+	require.Equal(t, "pq: cannot use virtual column \"a\" in primary key", err.Error())
+
+	_, err = sqlDB.Exec(`CREATE TABLE t (a INT NOT NULL AS (1+1) VIRTUAL PRIMARY KEY)`)
+	require.Error(t, err)
+	require.Equal(t, "pq: cannot use virtual column \"a\" in primary key", err.Error())
+
+	_, err = sqlDB.Exec(`CREATE TABLE t (a INT PRIMARY KEY, b INT NOT NULL AS (1+1) VIRTUAL)`)
+	require.NoError(t, err)
+
+	_, err = sqlDB.Exec(`ALTER TABLE t ALTER PRIMARY KEY USING COLUMNS (b)`)
+	require.Error(t, err)
+	require.Equal(t, "pq: cannot use virtual column \"b\" in primary key", err.Error())
+}

pkg/workload/schemachange/BUILD.bazel

@@ -15,6 +15,8 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/workload/schemachange",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/clusterversion",
+        "//pkg/roachpb",
         "//pkg/security",
         "//pkg/sql/catalog/catpb",
         "//pkg/sql/catalog/colinfo",

pkg/workload/schemachange/operation_generator.go

@@ -18,6 +18,8 @@ import (
 	"strings"
 	"sync/atomic"
 
+	"github.com/cockroachdb/cockroach/pkg/clusterversion"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
@@ -1116,9 +1118,36 @@ func (og *operationGenerator) createTable(ctx context.Context, tx pgx.Tx) (strin
 	if err != nil {
 		return "", err
 	}
+
+	// Detect if primary indexes contain computed columns, which are disallowed
+	// on older versions of Cockroach.
+	computedColInIndex := false
+	primaryKeyDisallowsComputedCols, err := isClusterVersionLessThan(ctx, tx,
+		clusterversion.TestingBinaryMinSupportedVersion)
+	if err != nil {
+		return "", err
+	}
+	if primaryKeyDisallowsComputedCols {
+		computedCols := make(map[string]struct{})
+		for _, def := range stmt.Defs {
+			if colDef, ok := def.(*tree.ColumnTableDef); ok {
+				if colDef.IsVirtual() {
+					computedCols[colDef.Name.String()] = struct{}{}
+				}
+			} else if indexDef, ok := def.(*tree.IndexTableDef); ok {
+				for _, indexCol := range indexDef.Columns {
+					if _, ok := computedCols[indexCol.Column.String()]; ok {
+						computedColInIndex = true
+					}
+				}
+			}
+		}
+	}
+
 	codesWithConditions{
 		{code: pgcode.DuplicateRelation, condition: tableExists && !stmt.IfNotExists},
 		{code: pgcode.UndefinedSchema, condition: !schemaExists},
+		{code: pgcode.FeatureNotSupported, condition: computedColInIndex},
 	}.add(og.expectedExecErrors)
 
 	return tree.Serialize(stmt), nil
@@ -3030,3 +3059,21 @@ func (og *operationGenerator) typeFromTypeName(
 	}
 	return typ, nil
 }
+
+// Check if the test is running with a mixed version cluster, with a version
+// less than or equal to the target version number. This can be used to detect
+// in mixed version environments if certain errors should be encountered.
+func isClusterVersionLessThan(
+	ctx context.Context, tx pgx.Tx, targetVersion roachpb.Version,
+) (bool, error) {
+	var clusterVersionStr string
+	row := tx.QueryRow(ctx, `SHOW CLUSTER SETTING version`)
+	if err := row.Scan(&clusterVersionStr); err != nil {
+		return false, err
+	}
+	clusterVersion, err := roachpb.ParseVersion(clusterVersionStr)
+	if err != nil {
+		return false, err
+	}
+	return clusterVersion.LessEq(targetVersion), nil
+}