From fbf1492f0f51c9050b8a5be06185659ef3d5482e Mon Sep 17 00:00:00 2001
From: Rafi Shamim <rafi@cockroachlabs.com>
Date: Tue, 14 Jan 2025 18:32:17 -0500
Subject: [PATCH] sql: relax validation for IDENTITY columns

The validations added in 4f99f9cb8bda4faee94f5d16ff1f1739acb62efa were
too strict. This patch makes it so we can continue having IDENTITY
columns that are not backed by a sequence. This can happen if a SERIAL
column is created with IDENTITY, and serial_normalization=rowid.

Release note (bug fix): Fixed a bug where columns created with GENERATED
... BY IDENTITY with the SERIAL type could incorrectly fail internal
validations.
---
 .../catalog/tabledesc/table_desc_builder.go   | 14 +++---
 pkg/sql/catalog/tabledesc/validate.go         |  7 +--
 pkg/sql/catalog/tabledesc/validate_test.go    |  4 +-
 .../logictest/testdata/logic_test/alter_table | 48 +++++++++++++++++++
 4 files changed, 61 insertions(+), 12 deletions(-)

diff --git a/pkg/sql/catalog/tabledesc/table_desc_builder.go b/pkg/sql/catalog/tabledesc/table_desc_builder.go
index f608a5e0ae41..0bfbf7b80273 100644
--- a/pkg/sql/catalog/tabledesc/table_desc_builder.go
+++ b/pkg/sql/catalog/tabledesc/table_desc_builder.go
@@ -1313,15 +1313,15 @@ func maybeFixUsesSequencesIDForIdentityColumns(
 		if len(column.UsesSequenceIds) == 1 {
 			continue
 		}
-		// The first ownership will point to the owner from CREATE TABLE.
-		tableDesc = builder.getOrInitModifiedDesc()
-		column = &tableDesc.Columns[idx]
 		if len(column.OwnsSequenceIds) == 0 {
-			return false, errors.AssertionFailedf("identity column %s (%d) on table %s must own a sequence",
-				column.Name,
-				column.ID,
-				tableDesc.Name)
+			// With serial_normalization=rowid, a table definition like
+			// `CREATE TABLE t (a SERIAL GENERATED ALWAYS AS IDENTITY)`
+			// creates an identity column without any backing sequence. If that's the
+			// case, there's no need to add a sequence reference.
+			continue
 		}
+		tableDesc = builder.getOrInitModifiedDesc()
+		column = &tableDesc.Columns[idx]
 		column.UsesSequenceIds = append(column.UsesSequenceIds, column.OwnsSequenceIds[0])
 		wasRepaired = true
 	}
diff --git a/pkg/sql/catalog/tabledesc/validate.go b/pkg/sql/catalog/tabledesc/validate.go
index 7a3cf93b8f8c..b6c72a8822e6 100644
--- a/pkg/sql/catalog/tabledesc/validate.go
+++ b/pkg/sql/catalog/tabledesc/validate.go
@@ -1225,9 +1225,10 @@ func (desc *wrapper) validateColumns() error {
 				}
 			}
 
-			// For generated as identity columns ensure that the column uses sequences.
-			if column.NumUsesSequences() != 1 {
-				return errors.Newf("column is identity without sequence references %q", column.GetName())
+			// For generated as identity columns ensure that the column uses sequences
+			// if it is backed by a sequence.
+			if column.NumOwnsSequences() == 1 && column.NumUsesSequences() != 1 {
+				return errors.Newf("column %q is GENERATED BY IDENTITY without sequence references", column.GetName())
 			}
 		}
 
diff --git a/pkg/sql/catalog/tabledesc/validate_test.go b/pkg/sql/catalog/tabledesc/validate_test.go
index dda790034dbb..3d13969bab6e 100644
--- a/pkg/sql/catalog/tabledesc/validate_test.go
+++ b/pkg/sql/catalog/tabledesc/validate_test.go
@@ -3069,14 +3069,14 @@ func TestValidateTableDesc(t *testing.T) {
 					},
 				}
 			})},
-		{err: `column is identity without sequence references "bar"`,
+		{err: `column "bar" is GENERATED BY IDENTITY without sequence references`,
 			desc: descpb.TableDescriptor{
 				ID:            2,
 				ParentID:      1,
 				Name:          "foo",
 				FormatVersion: descpb.InterleavedFormatVersion,
 				Columns: []descpb.ColumnDescriptor{
-					{ID: 1, Name: "bar", GeneratedAsIdentityType: catpb.GeneratedAsIdentityType_GENERATED_ALWAYS},
+					{ID: 1, Name: "bar", OwnsSequenceIds: []descpb.ID{5}, GeneratedAsIdentityType: catpb.GeneratedAsIdentityType_GENERATED_ALWAYS},
 				},
 				NextColumnID: 2,
 			}},
diff --git a/pkg/sql/logictest/testdata/logic_test/alter_table b/pkg/sql/logictest/testdata/logic_test/alter_table
index f64dc81937eb..fd49fc5d0adc 100644
--- a/pkg/sql/logictest/testdata/logic_test/alter_table
+++ b/pkg/sql/logictest/testdata/logic_test/alter_table
@@ -3842,6 +3842,54 @@ CREATE TABLE t_identity_drop_dependency(id INT PRIMARY KEY DEFAULT nextval('publ
 statement error pgcode 2BP01 pq: cannot drop table t_identity_drop because other objects depend on it
 ALTER TABLE t_identity_drop ALTER COLUMN b DROP IDENTITY;
 
+let $old_serial_normalization
+SHOW serial_normalization
+
+statement ok
+SET serial_normalization = rowid
+
+# Allow creating an IDENTITY column that isn't backed by a sequence.
+statement ok
+CREATE TABLE t_serial_identity_no_sequence(a SERIAL PRIMARY KEY GENERATED ALWAYS AS IDENTITY)
+
+query T
+SELECT crdb_internal.pb_to_json('desc', descriptor)->'table'->'columns'->0->>'defaultExpr'
+FROM system.descriptor
+WHERE id = 't_serial_identity_no_sequence'::regclass::oid
+----
+unique_rowid()
+
+query T
+SELECT crdb_internal.pb_to_json('desc', descriptor)->'table'->'columns'->0->>'generatedAsIdentityType'
+FROM system.descriptor
+WHERE id = 't_serial_identity_no_sequence'::regclass::oid
+----
+GENERATED_ALWAYS
+
+statement ok
+SET serial_normalization = sql_sequence
+
+# Allow creating an IDENTITY column that is backed by a sequence.
+statement ok
+CREATE TABLE t_serial_identity_with_sequence(a SERIAL PRIMARY KEY GENERATED ALWAYS AS IDENTITY)
+
+query T
+SELECT substring(crdb_internal.pb_to_json('desc', descriptor)->'table'->'columns'->0->>'defaultExpr' FOR 7)
+FROM system.descriptor
+WHERE id = 't_serial_identity_with_sequence'::regclass::oid
+----
+nextval
+
+query T
+SELECT crdb_internal.pb_to_json('desc', descriptor)->'table'->'columns'->0->>'generatedAsIdentityType'
+FROM system.descriptor
+WHERE id = 't_serial_identity_with_sequence'::regclass::oid
+----
+GENERATED_ALWAYS
+
+statement ok
+SET serial_normalization = $old_serial_normalization
+
 subtest comma_syntax
 
 statement ok