Merge cockroachdb#139135

139135: sql/schemachanger: support generated columns in ADD COLUMN r=fqazi a=fqazi Previously, the add column statement in the declarative schema changer only supported serial columns. This patch adds support for adding generated columns inside the declarative schema changer. Fixes: cockroachdb#128087 Release note: None Co-authored-by: Faizan Qazi <[email protected]>
mgartner · Jan 15, 2025 · 62d79ab · 62d79ab
2 parents cf07561 + ac9a722
commit 62d79ab
Show file tree

Hide file tree

Showing 28 changed files with 7,347 additions and 14 deletions.
diff --git a/pkg/ccl/schemachangerccl/backup_base_generated_test.go b/pkg/ccl/schemachangerccl/backup_base_generated_test.go
diff --git a/pkg/sql/logictest/testdata/logic_test/alter_table b/pkg/sql/logictest/testdata/logic_test/alter_table
@@ -2044,9 +2044,58 @@ t  CREATE TABLE public.t (
 statement ok
 INSERT INTO t (c) VALUES (2)
 
+statement ok
+ALTER TABLE t ADD COLUMN d SERIAL GENERATED BY DEFAULT AS IDENTITY
+
+query TT
+SHOW CREATE TABLE t
+----
+t  CREATE TABLE public.t (
+     a INT8 NULL,
+     rowid INT8 NOT VISIBLE NOT NULL DEFAULT unique_rowid(),
+     b INT8 NOT NULL GENERATED ALWAYS AS IDENTITY,
+     c INT8 NOT NULL GENERATED BY DEFAULT AS IDENTITY,
+     d INT8 NOT NULL GENERATED BY DEFAULT AS IDENTITY,
+     CONSTRAINT t_pkey PRIMARY KEY (rowid ASC),
+     UNIQUE INDEX t_a_key (a ASC)
+   )
+
+statement ok
+INSERT INTO t (c) VALUES (2)
+
 statement ok
 DROP TABLE t
 
+subtest generated_with_serial_unique_row_id
+
+statement ok
+SET serial_normalization = rowid
+
+statement ok
+CREATE TABLE t(a INT PRIMARY KEY)
+
+statement ok
+ALTER TABLE t ADD COLUMN b SERIAL GENERATED ALWAYS AS IDENTITY
+
+query T
+SELECT crdb_internal.pb_to_json('desc', descriptor)->'table'->'columns'->1->>'defaultExpr'
+FROM system.descriptor
+WHERE id = 't'::regclass::oid
+----
+unique_rowid()
+
+query T
+SELECT crdb_internal.pb_to_json('desc', descriptor)->'table'->'columns'->1->>'generatedAsIdentityType'
+FROM system.descriptor
+WHERE id = 't'::regclass::oid
+----
+GENERATED_ALWAYS
+
+statement ok
+DROP TABLE t;
+
+subtest end
+
 subtest generated_as_identity_with_seq_option
 statement ok
 DROP TABLE IF EXISTS t;
@@ -2450,8 +2499,8 @@ FROM (
 LEFT JOIN pg_catalog.pg_depend r ON l.table_id = r.objid;
 ----
 table_id  name               state   refobjid
-201       test_serial_b_seq  PUBLIC  200
-200       test_serial        PUBLIC  NULL
+202       test_serial_b_seq  PUBLIC  201
+201       test_serial        PUBLIC  NULL
 
 statement ok
 DROP TABLE test_serial;
@@ -2485,8 +2534,8 @@ FROM (
 LEFT JOIN pg_catalog.pg_depend r ON l.table_id = r.objid;
 ----
 table_id  name               state   refobjid
-203       test_serial_b_seq  PUBLIC  202
-202       test_serial        PUBLIC  NULL
+204       test_serial_b_seq  PUBLIC  203
+203       test_serial        PUBLIC  NULL
 
 statement ok
 ALTER TABLE test_serial DROP COLUMN b;
@@ -2501,7 +2550,7 @@ FROM (
 LEFT JOIN pg_catalog.pg_depend r ON l.table_id = r.objid;
 ----
 table_id  name         state   refobjid
-202       test_serial  PUBLIC  NULL
+203       test_serial  PUBLIC  NULL
 
 statement ok
 DROP TABLE test_serial;

diff --git a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_column.go b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_column.go
@@ -68,11 +68,8 @@ func alterTableAddColumn(
 		panic(sqlerrors.NewColumnAlreadyExistsInRelationError(string(d.Name), tn.Object()))
 	}
 	var colSerialDefaultExpression *scpb.Expression
-	if d.IsSerial {
-		d, colSerialDefaultExpression = alterTableAddColumnSerial(b, d, tn)
-	}
-	if d.GeneratedIdentity.IsGeneratedAsIdentity {
-		panic(scerrors.NotImplementedErrorf(d, "contains generated identity type"))
+	if d.IsSerial || d.GeneratedIdentity.IsGeneratedAsIdentity {
+		d, colSerialDefaultExpression = alterTableAddColumnSerialOrGeneratedIdentity(b, d, tn)
 	}
 	// Unique without an index is unsupported.
 	if d.Unique.WithoutIndex {
@@ -308,22 +305,30 @@ func alterTableAddColumn(
 	}
 }
 
-func alterTableAddColumnSerial(
+func alterTableAddColumnSerialOrGeneratedIdentity(
 	b BuildCtx, d *tree.ColumnTableDef, tn *tree.TableName,
 ) (newDef *tree.ColumnTableDef, colDefaultExpression *scpb.Expression) {
 	if err := catalog.AssertValidSerialColumnDef(d, tn); err != nil {
 		panic(err)
 	}
+	// A generated column can also be serial at the same time.
+	isGeneratedColumn := !d.IsSerial && d.GeneratedIdentity.IsGeneratedAsIdentity
 
 	defType, err := tree.ResolveType(b, d.Type, b.SemaCtx().GetTypeResolver())
 	if err != nil {
 		panic(err)
 	}
-
-	telemetry.Inc(sqltelemetry.SerialColumnNormalizationCounter(
-		defType.Name(), b.SessionData().SerialNormalizationMode.String()))
+	if d.IsSerial {
+		telemetry.Inc(sqltelemetry.SerialColumnNormalizationCounter(
+			defType.Name(), b.SessionData().SerialNormalizationMode.String()))
+	}
 
 	serialNormalizationMode := b.SessionData().SerialNormalizationMode
+	// Generate identity is always a SQL sequence based column.
+	if isGeneratedColumn {
+		serialNormalizationMode = sessiondatapb.SerialUsesSQLSequences
+	}
+
 	switch serialNormalizationMode {
 	// The type will be upgraded when the columns are setup or before a
 	// sequence is created.
@@ -379,6 +384,10 @@ func alterTableAddColumnSerial(
 	if err != nil {
 		panic(err)
 	}
+	// For generated identities inherit the sequence options from the AST.
+	if d.GeneratedIdentity.IsGeneratedAsIdentity {
+		seqOptions = d.GeneratedIdentity.SeqOptions
+	}
 
 	// Create the sequence and fetch the element after. The full descriptor
 	// will be created after the build phase, so we cannot fully resolve it.

diff --git a/pkg/sql/schemachanger/scbuild/testdata/alter_table_add_column b/pkg/sql/schemachanger/scbuild/testdata/alter_table_add_column
@@ -309,3 +309,117 @@ ALTER TABLE t ALTER PRIMARY KEY USING COLUMNS (j)
   {columnId: 2, indexId: 3, kind: KEY_SUFFIX, tableId: 106}
 - [[IndexName:{DescID: 106, Name: t_i_key, IndexID: 2}, PUBLIC], ABSENT]
   {indexId: 2, name: t_i_key, tableId: 106}
+
+build
+ALTER TABLE defaultdb.foo ADD COLUMN j INT GENERATED ALWAYS AS IDENTITY
+----
+- [[IndexColumn:{DescID: 104, ColumnID: 1, IndexID: 1}, ABSENT], PUBLIC]
+  {columnId: 1, indexId: 1, tableId: 104}
+- [[PrimaryIndex:{DescID: 104, IndexID: 1, ConstraintID: 1}, ABSENT], PUBLIC]
+  {constraintId: 1, indexId: 1, isUnique: true, tableId: 104}
+- [[IndexName:{DescID: 104, Name: foo_pkey, IndexID: 1}, ABSENT], PUBLIC]
+  {indexId: 1, name: foo_pkey, tableId: 104}
+- [[IndexData:{DescID: 104, IndexID: 1}, ABSENT], PUBLIC]
+  {indexId: 1, tableId: 104}
+- [[TableData:{DescID: 104, ReferencedDescID: 100}, PUBLIC], PUBLIC]
+  {databaseId: 100, tableId: 104}
+- [[Sequence:{DescID: 107}, PUBLIC], ABSENT]
+  {sequenceId: 107}
+- [[Namespace:{DescID: 107, Name: foo_j_seq, ReferencedDescID: 100}, PUBLIC], ABSENT]
+  {databaseId: 100, descriptorId: 107, name: foo_j_seq, schemaId: 101}
+- [[SchemaChild:{DescID: 107, ReferencedDescID: 101}, PUBLIC], ABSENT]
+  {childObjectId: 107, schemaId: 101}
+- [[TableData:{DescID: 107, ReferencedDescID: 100}, PUBLIC], ABSENT]
+  {databaseId: 100, tableId: 107}
+- [[Column:{DescID: 107, ColumnID: 1}, PUBLIC], ABSENT]
+  {columnId: 1, tableId: 107}
+- [[ColumnType:{DescID: 107, ColumnFamilyID: 0, ColumnID: 1, TypeName: INT8}, PUBLIC], ABSENT]
+  {columnId: 1, elementCreationMetadata: {in231OrLater: true, in243OrLater: true}, tableId: 107, type: {family: IntFamily, oid: 20, width: 64}, typeName: INT8}
+- [[ColumnNotNull:{DescID: 107, ColumnID: 1, IndexID: 0}, PUBLIC], ABSENT]
+  {columnId: 1, tableId: 107}
+- [[ColumnName:{DescID: 107, Name: value, ColumnID: 1}, PUBLIC], ABSENT]
+  {columnId: 1, name: value, tableId: 107}
+- [[PrimaryIndex:{DescID: 107, IndexID: 1, ConstraintID: 0}, PUBLIC], ABSENT]
+  {indexId: 1, isUnique: true, tableId: 107}
+- [[IndexName:{DescID: 107, Name: primary, IndexID: 1}, PUBLIC], ABSENT]
+  {indexId: 1, name: primary, tableId: 107}
+- [[IndexColumn:{DescID: 107, ColumnID: 1, IndexID: 1}, PUBLIC], ABSENT]
+  {columnId: 1, indexId: 1, tableId: 107}
+- [[Owner:{DescID: 107}, PUBLIC], ABSENT]
+  {descriptorId: 107, owner: root}
+- [[UserPrivileges:{DescID: 107, Name: admin}, PUBLIC], ABSENT]
+  {descriptorId: 107, privileges: "2", userName: admin, withGrantOption: "2"}
+- [[UserPrivileges:{DescID: 107, Name: root}, PUBLIC], ABSENT]
+  {descriptorId: 107, privileges: "2", userName: root, withGrantOption: "2"}
+- [[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}
+- [[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: {in231OrLater: true, in243OrLater: 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]}
+- [[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]
+  {indexId: 2, name: foo_pkey, tableId: 104}
+- [[IndexColumn:{DescID: 104, ColumnID: 1, IndexID: 2}, PUBLIC], ABSENT]
+  {columnId: 1, indexId: 2, tableId: 104}
+- [[IndexData:{DescID: 104, IndexID: 2}, PUBLIC], ABSENT]
+  {indexId: 2, tableId: 104}
+- [[TemporaryIndex:{DescID: 104, IndexID: 3, ConstraintID: 3, SourceIndexID: 1}, TRANSIENT_ABSENT], ABSENT]
+  {constraintId: 3, indexId: 3, isUnique: true, sourceIndexId: 1, tableId: 104}
+- [[IndexColumn:{DescID: 104, ColumnID: 1, IndexID: 3}, TRANSIENT_ABSENT], ABSENT]
+  {columnId: 1, indexId: 3, tableId: 104}
+- [[IndexData:{DescID: 104, IndexID: 3}, TRANSIENT_ABSENT], ABSENT]
+  {indexId: 3, tableId: 104}
+- [[IndexColumn:{DescID: 104, ColumnID: 2, IndexID: 2}, PUBLIC], ABSENT]
+  {columnId: 2, indexId: 2, kind: STORED, tableId: 104}
+- [[IndexColumn:{DescID: 104, ColumnID: 2, IndexID: 3}, TRANSIENT_ABSENT], ABSENT]
+  {columnId: 2, indexId: 3, kind: STORED, tableId: 104}
+- [[ColumnNotNull:{DescID: 104, ColumnID: 2, IndexID: 2}, PUBLIC], ABSENT]
+  {columnId: 2, indexIdForValidation: 2, tableId: 104}
+
+build
+ALTER TABLE defaultdb.foo ADD COLUMN j SERIAL GENERATED ALWAYS AS IDENTITY
+----
+- [[IndexColumn:{DescID: 104, ColumnID: 1, IndexID: 1}, ABSENT], PUBLIC]
+  {columnId: 1, indexId: 1, tableId: 104}
+- [[PrimaryIndex:{DescID: 104, IndexID: 1, ConstraintID: 1}, ABSENT], PUBLIC]
+  {constraintId: 1, indexId: 1, isUnique: true, tableId: 104}
+- [[IndexName:{DescID: 104, Name: foo_pkey, IndexID: 1}, ABSENT], PUBLIC]
+  {indexId: 1, name: foo_pkey, tableId: 104}
+- [[IndexData:{DescID: 104, IndexID: 1}, ABSENT], PUBLIC]
+  {indexId: 1, tableId: 104}
+- [[TableData:{DescID: 104, ReferencedDescID: 100}, PUBLIC], PUBLIC]
+  {databaseId: 100, tableId: 104}
+- [[Column:{DescID: 104, ColumnID: 2}, PUBLIC], ABSENT]
+  {columnId: 2, generatedAsIdentityType: 1, 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: {in231OrLater: true, in243OrLater: 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}
+- [[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]
+  {indexId: 2, name: foo_pkey, tableId: 104}
+- [[IndexColumn:{DescID: 104, ColumnID: 1, IndexID: 2}, PUBLIC], ABSENT]
+  {columnId: 1, indexId: 2, tableId: 104}
+- [[IndexData:{DescID: 104, IndexID: 2}, PUBLIC], ABSENT]
+  {indexId: 2, tableId: 104}
+- [[TemporaryIndex:{DescID: 104, IndexID: 3, ConstraintID: 3, SourceIndexID: 1}, TRANSIENT_ABSENT], ABSENT]
+  {constraintId: 3, indexId: 3, isUnique: true, sourceIndexId: 1, tableId: 104}
+- [[IndexColumn:{DescID: 104, ColumnID: 1, IndexID: 3}, TRANSIENT_ABSENT], ABSENT]
+  {columnId: 1, indexId: 3, tableId: 104}
+- [[IndexData:{DescID: 104, IndexID: 3}, TRANSIENT_ABSENT], ABSENT]
+  {indexId: 3, tableId: 104}
+- [[IndexColumn:{DescID: 104, ColumnID: 2, IndexID: 2}, PUBLIC], ABSENT]
+  {columnId: 2, indexId: 2, kind: STORED, tableId: 104}
+- [[IndexColumn:{DescID: 104, ColumnID: 2, IndexID: 3}, TRANSIENT_ABSENT], ABSENT]
+  {columnId: 2, indexId: 3, kind: STORED, tableId: 104}
+- [[ColumnNotNull:{DescID: 104, ColumnID: 2, IndexID: 2}, PUBLIC], ABSENT]
+  {columnId: 2, indexIdForValidation: 2, tableId: 104}