sql: add ColumnGeneratedAsIdentity as a new schema change element

Previously, GeneratedAsIdentity was part of the Column element.
Creating a new schema change element will allow the declarative schema changer
to add/remove/update such columns without dropping the Column element.

Epic: CRDB-31283

Part of #142918

Release note: None

Author: shghasemi

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_column.go

@@ -116,11 +116,10 @@ func alterTableAddColumn(
 	spec := addColumnSpec{
 		tbl: tbl,
 		col: &scpb.Column{
-			TableID:                 tbl.TableID,
-			ColumnID:                desc.ID,
-			IsHidden:                desc.Hidden,
-			IsInaccessible:          desc.Inaccessible,
-			GeneratedAsIdentityType: desc.GeneratedAsIdentityType,
+			TableID:        tbl.TableID,
+			ColumnID:       desc.ID,
+			IsHidden:       desc.Hidden,
+			IsInaccessible: desc.Inaccessible,
 		},
 		unique:  d.Unique.IsUnique,
 		notNull: !desc.Nullable,
@@ -139,9 +138,6 @@ func alterTableAddColumn(
 	if pgAttNum := desc.GetPGAttributeNum(); pgAttNum != catid.PGAttributeNum(desc.ID) {
 		spec.col.PgAttributeNum = pgAttNum
 	}
-	if ptr := desc.GeneratedAsIdentitySequenceOption; ptr != nil {
-		spec.col.GeneratedAsIdentitySequenceOption = *ptr
-	}
 	spec.name = &scpb.ColumnName{
 		TableID:  tbl.TableID,
 		ColumnID: spec.col.ColumnID,
@@ -290,6 +286,26 @@ func alterTableAddColumn(
 		}
 		b.IncrementSchemaChangeAddColumnQualificationCounter("on_update")
 	}
+	if d.GeneratedIdentity.IsGeneratedAsIdentity {
+		generatedAsIdentityType := desc.GeneratedAsIdentityType
+		seqOptions := ""
+		if ptr := desc.GeneratedAsIdentitySequenceOption; ptr != nil {
+			seqOptions = *ptr
+		}
+		// Versions from 26.1 GeneratedAsIdentity will have a separate element for
+		// GeneratedAsIdentity. Older versios store it in the column element.
+		if spec.colType.ElementCreationMetadata.In_26_1OrLater {
+			spec.generatedAsID = &scpb.ColumnGeneratedAsIdentity{
+				TableID:        tbl.TableID,
+				ColumnID:       spec.col.ColumnID,
+				Type:           generatedAsIdentityType,
+				SequenceOption: seqOptions,
+			}
+		} else {
+			spec.col.GeneratedAsIdentityType = generatedAsIdentityType
+			spec.col.GeneratedAsIdentitySequenceOption = seqOptions
+		}
+	}
 	// Add secondary indexes for this column.
 	backing := addColumn(b, spec, t)
 	if idx != nil {
@@ -458,6 +474,7 @@ type addColumnSpec struct {
 	compute          *scpb.ColumnComputeExpression
 	transientCompute *scpb.ColumnComputeExpression
 	comment          *scpb.ColumnComment
+	generatedAsID    *scpb.ColumnGeneratedAsIdentity
 	unique           bool
 	notNull          bool
 }
@@ -494,6 +511,9 @@ func addColumn(b BuildCtx, spec addColumnSpec, n tree.NodeFormatter) (backing *s
 		if spec.comment != nil {
 			b.Add(spec.comment)
 		}
+		if spec.generatedAsID != nil {
+			b.Add(spec.generatedAsID)
+		}
 		// Don't need to modify primary indexes for virtual columns.
 		if spec.colType.IsVirtual {
 			return getLatestPrimaryIndex(b, spec.tbl.TableID)

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_alter_column_drop_not_null.go

@@ -6,7 +6,6 @@
 package scbuildstmt
 
 import (
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scpb"
@@ -37,8 +36,7 @@ func alterTableDropNotNull(
 			`column "%s" is in a primary index`, colName.Name))
 	}
 	// Ensure that we are not dropping not-null on a generated column.
-	colEl := mustRetrieveColumnElem(b, tbl.TableID, columnID)
-	if colEl.GeneratedAsIdentityType != catpb.GeneratedAsIdentityType_NOT_IDENTITY_COLUMN {
+	if isColumnGeneratedAsIdentity(b, tbl.TableID, columnID) {
 		colName := mustRetrieveColumnName(b, tbl.TableID, columnID)
 		panic(pgerror.Newf(pgcode.Syntax,
 			`column "%s" of relation "%s" is an identity column`, colName.Name, tn.ObjectName))

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_alter_column_set_default.go

@@ -8,7 +8,6 @@ package scbuildstmt
 import (
 	"context"
 
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemaexpr"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/seqexpr"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
@@ -75,7 +74,7 @@ func panicIfInvalidNonComputedColumnExpr(
 	newExpr tree.Expr,
 	schemaChange tree.SchemaExprContext,
 ) tree.TypedExpr {
-	if col.GeneratedAsIdentityType != catpb.GeneratedAsIdentityType_NOT_IDENTITY_COLUMN {
+	if isColumnGeneratedAsIdentity(b, tbl.TableID, col.ColumnID) {
 		panic(sqlerrors.NewSyntaxErrorf("column %q is an identity column", colName))
 	}
 

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_alter_column_type.go

@@ -10,7 +10,6 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/build"
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemaexpr"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
@@ -98,7 +97,7 @@ func alterTableAlterColumnType(
 	var err error
 	newColType.Type, err = schemachange.ValidateAlterColumnTypeChecks(
 		b, t, b.ClusterSettings(), newColType.Type,
-		col.GeneratedAsIdentityType != catpb.GeneratedAsIdentityType_NOT_IDENTITY_COLUMN,
+		isColumnGeneratedAsIdentity(b, tbl.TableID, col.ColumnID),
 		newColType.IsVirtual)
 	if err != nil {
 		panic(err)

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_drop_column.go

@@ -326,7 +326,7 @@ func walkColumnDependencies(
 				*scpb.UniqueWithoutIndexConstraint, *scpb.CheckConstraint,
 				*scpb.UniqueWithoutIndexConstraintUnvalidated, *scpb.CheckConstraintUnvalidated,
 				*scpb.RowLevelTTL, *scpb.PolicyUsingExpr, *scpb.PolicyWithCheckExpr,
-				*scpb.TriggerDeps:
+				*scpb.TriggerDeps, *scpb.ColumnGeneratedAsIdentity:
 				fn(e, op, objType)
 			case *scpb.ColumnType:
 				if elt.ColumnID == col.ColumnID {

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/helpers.go

@@ -2218,3 +2218,30 @@ func isShardColumn(b BuildCtx, col *scpb.Column) bool {
 
 	return found
 }
+
+// retrieveColumnGeneratedAsIdentityType returns the GeneratedAsIdentityType.
+// Handles versions before 26.1 that store it in Column.GeneratedAsIdentityType.
+func retrieveColumnGeneratedAsIdentityType(
+	b BuildCtx, tableID catid.DescID, columnID catid.ColumnID,
+) (generatedAsIdentityType catpb.GeneratedAsIdentityType) {
+	// First try to retrieve it from ColumnGeneratedAsIdentity. This may be nil
+	// if the column is not a generated identity column or the cluster version is older
+	// than 26.1. In that case, will return Column.GeneratedAsIdentityType.
+	colGeneratedAsID := b.QueryByID(tableID).FilterColumnGeneratedAsIdentity().Filter(
+		func(_ scpb.Status, _ scpb.TargetStatus, e *scpb.ColumnGeneratedAsIdentity) bool {
+			return e.ColumnID == columnID
+		}).MustGetZeroOrOneElement()
+	if colGeneratedAsID != nil {
+		return colGeneratedAsID.Type
+	}
+	// Check the ColumnType in case this is an older version.
+	col := mustRetrieveColumnElem(b, tableID, columnID)
+	return col.GeneratedAsIdentityType
+}
+
+func isColumnGeneratedAsIdentity(
+	b BuildCtx, tableID catid.DescID, columnID catid.ColumnID,
+) (isGenerated bool) {
+	generatedAsIdentityType := retrieveColumnGeneratedAsIdentityType(b, tableID, columnID)
+	return generatedAsIdentityType != catpb.GeneratedAsIdentityType_NOT_IDENTITY_COLUMN
+}

pkg/sql/schemachanger/scbuild/testdata/alter_table_add_column

@@ -354,13 +354,15 @@ ALTER TABLE defaultdb.foo ADD COLUMN j INT GENERATED ALWAYS AS IDENTITY
 - [[SequenceOwner:{DescID: 104, ColumnID: 2, ReferencedDescID: 107}, PUBLIC], ABSENT]
   {columnId: 2, sequenceId: 107, tableId: 104}
 - [[Column:{DescID: 104, ColumnID: 2}, PUBLIC], ABSENT]
-  {columnId: 2, generatedAsIdentityType: 1, tableId: 104}
+  {columnId: 2, tableId: 104}
 - [[ColumnName:{DescID: 104, Name: j, ColumnID: 2}, PUBLIC], ABSENT]
   {columnId: 2, name: j, tableId: 104}
 - [[ColumnType:{DescID: 104, ColumnFamilyID: 0, ColumnID: 2, TypeName: INT8}, PUBLIC], ABSENT]
   {columnId: 2, elementCreationMetadata: {in243OrLater: true, in261OrLater: true}, tableId: 104, type: {family: IntFamily, oid: 20, width: 64}, typeName: INT8}
 - [[ColumnDefaultExpression:{DescID: 104, ColumnID: 2, ReferencedSequenceIDs: [107], Expr: nextval(107:::REGCLASS)}, PUBLIC], ABSENT]
   {columnId: 2, expr: 'nextval(107:::REGCLASS)', tableId: 104, usesSequenceIds: [107]}
+- [[ColumnGeneratedAsIdentity:{DescID: 104, ColumnID: 2}, PUBLIC], ABSENT]
+  {columnId: 2, tableId: 104, type: 1}
 - [[PrimaryIndex:{DescID: 104, IndexID: 2, ConstraintID: 2, TemporaryIndexID: 3, SourceIndexID: 1}, PUBLIC], ABSENT]
   {constraintId: 2, indexId: 2, isUnique: true, sourceIndexId: 1, tableId: 104, temporaryIndexId: 3}
 - [[IndexName:{DescID: 104, Name: foo_pkey, IndexID: 2}, PUBLIC], ABSENT]
@@ -396,13 +398,15 @@ ALTER TABLE defaultdb.foo ADD COLUMN j SERIAL GENERATED ALWAYS AS IDENTITY
 - [[TableData:{DescID: 104, ReferencedDescID: 100}, PUBLIC], PUBLIC]
   {databaseId: 100, tableId: 104}
 - [[Column:{DescID: 104, ColumnID: 2}, PUBLIC], ABSENT]
-  {columnId: 2, generatedAsIdentityType: 1, tableId: 104}
+  {columnId: 2, tableId: 104}
 - [[ColumnName:{DescID: 104, Name: j, ColumnID: 2}, PUBLIC], ABSENT]
   {columnId: 2, name: j, tableId: 104}
 - [[ColumnType:{DescID: 104, ColumnFamilyID: 0, ColumnID: 2, TypeName: INT8}, PUBLIC], ABSENT]
   {columnId: 2, elementCreationMetadata: {in243OrLater: true, in261OrLater: true}, tableId: 104, type: {family: IntFamily, oid: 20, width: 64}, typeName: INT8}
 - [[ColumnDefaultExpression:{DescID: 104, ColumnID: 2, Expr: unique_rowid()}, PUBLIC], ABSENT]
   {columnId: 2, expr: unique_rowid(), tableId: 104}
+- [[ColumnGeneratedAsIdentity:{DescID: 104, ColumnID: 2}, PUBLIC], ABSENT]
+  {columnId: 2, tableId: 104, type: 1}
 - [[PrimaryIndex:{DescID: 104, IndexID: 2, ConstraintID: 2, TemporaryIndexID: 3, SourceIndexID: 1}, PUBLIC], ABSENT]
   {constraintId: 2, indexId: 2, isUnique: true, sourceIndexId: 1, tableId: 104, temporaryIndexId: 3}
 - [[IndexName:{DescID: 104, Name: foo_pkey, IndexID: 2}, PUBLIC], ABSENT]

pkg/sql/schemachanger/scdecomp/decomp.go

@@ -541,13 +541,11 @@ func (w *walkCtx) walkColumn(tbl catalog.TableDescriptor, col catalog.Column) {
 			col.GetName(), tbl.GetName(), tbl.GetID()))
 	}
 	column := &scpb.Column{
-		TableID:                           tbl.GetID(),
-		ColumnID:                          col.GetID(),
-		IsHidden:                          col.IsHidden(),
-		IsInaccessible:                    col.IsInaccessible(),
-		GeneratedAsIdentityType:           col.GetGeneratedAsIdentityType(),
-		GeneratedAsIdentitySequenceOption: col.GetGeneratedAsIdentitySequenceOptionStr(),
-		IsSystemColumn:                    col.IsSystemColumn(),
+		TableID:        tbl.GetID(),
+		ColumnID:       col.GetID(),
+		IsHidden:       col.IsHidden(),
+		IsInaccessible: col.IsInaccessible(),
+		IsSystemColumn: col.IsSystemColumn(),
 	}
 	// Only set PgAttributeNum if it differs from ColumnID.
 	if pgAttNum := col.GetPGAttributeNum(); pgAttNum != catid.PGAttributeNum(col.GetID()) {
@@ -590,6 +588,23 @@ func (w *walkCtx) walkColumn(tbl catalog.TableDescriptor, col catalog.Column) {
 				columnType.ComputeExpr = expr
 			}
 		}
+
+		if columnType.ElementCreationMetadata.In_26_1OrLater {
+			if col.IsGeneratedAsIdentity() {
+				w.ev(scpb.Status_PUBLIC,
+					&scpb.ColumnGeneratedAsIdentity{
+						TableID:        tbl.GetID(),
+						ColumnID:       col.GetID(),
+						Type:           col.GetGeneratedAsIdentityType(),
+						SequenceOption: col.GetGeneratedAsIdentitySequenceOptionStr(),
+					})
+				column.GeneratedAsIdentityType = catpb.GeneratedAsIdentityType_NOT_IDENTITY_COLUMN
+				column.GeneratedAsIdentitySequenceOption = ""
+			}
+		} else {
+			column.GeneratedAsIdentityType = col.GetGeneratedAsIdentityType()
+			column.GeneratedAsIdentitySequenceOption = col.GetGeneratedAsIdentitySequenceOptionStr()
+		}
 		w.ev(scpb.Status_PUBLIC, columnType)
 	}
 	if !col.IsNullable() {

pkg/sql/schemachanger/scexec/scmutationexec/column.go

@@ -162,6 +162,52 @@ func (i *immediateVisitor) updateColumnComputeExpression(
 	return updateColumnExprFunctionsUsage(col)
 }
 
+// AddColumnGeneratedAsIdentity will set generated as identity for a column.
+func (i *immediateVisitor) AddColumnGeneratedAsIdentity(
+	ctx context.Context, op scop.AddColumnGeneratedAsIdentity,
+) error {
+	return i.updateColumnGeneratedAsIdentity(
+		ctx, op.GeneratedAsIdentity.TableID, op.GeneratedAsIdentity.ColumnID,
+		op.GeneratedAsIdentity.Type, &op.GeneratedAsIdentity.SequenceOption)
+}
+
+// RemoveColumnGeneratedAsIdentity will drop generated as identity from a column.
+func (i *immediateVisitor) RemoveColumnGeneratedAsIdentity(
+	ctx context.Context, op scop.RemoveColumnGeneratedAsIdentity,
+) error {
+	return i.updateColumnGeneratedAsIdentity(
+		ctx, op.TableID, op.ColumnID, catpb.GeneratedAsIdentityType_NOT_IDENTITY_COLUMN, nil)
+}
+
+// updateColumnGeneratedAsIdentity will handle add or removal of generated as identity
+func (i *immediateVisitor) updateColumnGeneratedAsIdentity(
+	ctx context.Context,
+	tableID descpb.ID,
+	columnID descpb.ColumnID,
+	genType catpb.GeneratedAsIdentityType,
+	sequenceOption *string,
+) error {
+	tbl, err := i.checkOutTable(ctx, tableID)
+	if err != nil {
+		return err
+	}
+
+	catCol, err := catalog.MustFindColumnByID(tbl, columnID)
+	if err != nil {
+		return err
+	}
+
+	col := catCol.ColumnDesc()
+	col.GeneratedAsIdentityType = genType
+	if sequenceOption != nil && *sequenceOption != "" {
+		so := *sequenceOption
+		col.GeneratedAsIdentitySequenceOption = &so
+	} else {
+		col.GeneratedAsIdentitySequenceOption = nil
+	}
+	return nil
+}
+
 func (i *immediateVisitor) MakeDeleteOnlyColumnWriteOnly(
 	ctx context.Context, op scop.MakeDeleteOnlyColumnWriteOnly,
 ) error {

pkg/sql/schemachanger/scop/immediate_mutation.go

@@ -247,6 +247,19 @@ type RemoveColumnComputeExpression struct {
 	ColumnID descpb.ColumnID
 }
 
+// AddColumnGeneratedAsIdentity will add a new generated as identity to a column.
+type AddColumnGeneratedAsIdentity struct {
+	immediateMutationOp
+	GeneratedAsIdentity scpb.ColumnGeneratedAsIdentity
+}
+
+// RemoveColumnGeneratedAsIdentity will remove the generated as identity from a column.
+type RemoveColumnGeneratedAsIdentity struct {
+	immediateMutationOp
+	TableID  descpb.ID
+	ColumnID descpb.ColumnID
+}
+
 // MakeWriteOnlyColumnPublic moves a new column from its mutation to public.
 type MakeWriteOnlyColumnPublic struct {
 	immediateMutationOp