Merge #118843

118843: schemachanger: address index validation errors during pk swap r=fqazi a=fqazi

Previously, when mutations were made public for validation queries, the declarative schema changer primary key swaps were not appropriately handled. This could lead to scenarios where queries would be unable to find columns in recreated secondary indexes that should ideally exist (i.e. columns from the old primary key). This would cause validation to fail on these newly created indexes. To address, this patch adds logic for making declarative primary key swaps mutations properly by detecting them.

Fixes: #118626

Release note (bug fix): ALTER PRIMARY KEY could fail with a "non-nullable column <x> with no value! Index scanned .." when validating recreated secondary indexes.

Co-authored-by: Faizan Qazi <faizan@cockroachlabs.com>

pkg/sql/catalog/descriptor.go

@@ -759,6 +759,10 @@ type TableDescriptor interface {
 	// IsSchemaLocked returns true if we don't allow performing schema changes
 	// on this table descriptor.
 	IsSchemaLocked() bool
+	// IsPrimaryKeySwapMutation returns true if the mutation is a primary key
+	// swap mutation or a secondary index used by the declarative schema changer
+	// for a primary index swap.
+	IsPrimaryKeySwapMutation(m *descpb.DescriptorMutation) bool
 }
 
 // MutableTableDescriptor is both a MutableDescriptor and a TableDescriptor.

pkg/sql/catalog/tabledesc/structured.go

@@ -1452,8 +1452,17 @@ func (desc *Mutable) MakeMutationComplete(m descpb.DescriptorMutation) error {
 			desc.AddColumn(t.Column)
 
 		case *descpb.DescriptorMutation_Index:
-			if err := desc.AddSecondaryIndex(*t.Index); err != nil {
-				return err
+			// If a primary index is being made public, then we only need set the
+			// index inside the descriptor directly. Only the declarative schema
+			// changer will use index mutations like this.
+			isPrimaryIndexToPublic := desc.IsPrimaryKeySwapMutation(&m)
+			if isPrimaryIndexToPublic {
+				desc.SetPrimaryIndex(*t.Index)
+			} else {
+				// Otherwise, we need to add this index as a secondary index.
+				if err := desc.AddSecondaryIndex(*t.Index); err != nil {
+					return err
+				}
 			}
 
 		case *descpb.DescriptorMutation_Constraint:
@@ -2065,9 +2074,9 @@ func (desc *wrapper) MakeFirstMutationPublic(
 		}
 		i++
 		switch {
-		case policy.shouldSkip(&mutation):
+		case policy.shouldSkip(desc, &mutation):
 			// Don't add to clone.
-		case policy.shouldRetain(&mutation):
+		case policy.shouldRetain(desc, &mutation):
 			mutation.Direction = descpb.DescriptorMutation_ADD
 			fallthrough
 		default:
@@ -2099,9 +2108,11 @@ func (p mutationPublicationPolicy) includes(f catalog.MutationPublicationFilter)
 	return p.policy.Contains(int(f))
 }
 
-func (p mutationPublicationPolicy) shouldSkip(m *descpb.DescriptorMutation) bool {
+func (p mutationPublicationPolicy) shouldSkip(
+	desc catalog.TableDescriptor, m *descpb.DescriptorMutation,
+) bool {
 	switch {
-	case m.GetPrimaryKeySwap() != nil:
+	case desc.IsPrimaryKeySwapMutation(m):
 		return p.includes(catalog.IgnorePKSwaps)
 	case m.GetConstraint() != nil:
 		return p.includes(catalog.IgnoreConstraints)
@@ -2110,7 +2121,9 @@ func (p mutationPublicationPolicy) shouldSkip(m *descpb.DescriptorMutation) bool
 	}
 }
 
-func (p mutationPublicationPolicy) shouldRetain(m *descpb.DescriptorMutation) bool {
+func (p mutationPublicationPolicy) shouldRetain(
+	desc catalog.TableDescriptor, m *descpb.DescriptorMutation,
+) bool {
 	switch {
 	case m.GetColumn() != nil && m.Direction == descpb.DescriptorMutation_DROP:
 		return p.includes(catalog.RetainDroppingColumns)

pkg/sql/catalog/tabledesc/table_desc.go

@@ -637,3 +637,45 @@ func (desc *wrapper) ForEachUDTDependentForHydration(fn func(t *types.T) error)
 func (desc *wrapper) IsSchemaLocked() bool {
 	return desc.SchemaLocked
 }
+
+// IsPrimaryKeySwapMutation implements the TableDescriptor interface.
+func (desc *wrapper) IsPrimaryKeySwapMutation(m *descpb.DescriptorMutation) bool {
+	switch t := m.Descriptor_.(type) {
+	case *descpb.DescriptorMutation_PrimaryKeySwap:
+		return true
+	case *descpb.DescriptorMutation_Index:
+		// The declarative schema changer handles primary key swaps differently,
+		// and does not use a PrimaryKeySwap mutation for this operation. Instead,
+		// it will create a new index with a primary index encoding as an
+		// index mutation. To detect a primary index swap scenario we are going to
+		// in the declarative case detect the encoding type
+		// and check if a declarative scpb.PrimaryIndex element with a matching ID
+		// exists and is going public. Since a table can only have a single primary
+		// index at a time, this would indicate this index is replacing the existing
+		// primary index.
+		if t.Index.EncodingType != catenumpb.PrimaryIndexEncoding {
+			return false
+		}
+		state := desc.GetDeclarativeSchemaChangerState()
+		if state == nil {
+			return false
+		}
+		// Loop over all targets searching for primary indexes.
+		for _, pk := range state.Targets {
+			// Confirm the target is a primary index going to public.
+			if pk.TargetStatus != scpb.Status_PUBLIC {
+				continue
+			}
+			pk, ok := pk.ElementOneOf.(*scpb.ElementProto_PrimaryIndex)
+			if !ok {
+				continue
+			}
+			// If the primary index going public matches any current mutation
+			// then this is an index swap scenario.
+			if pk.PrimaryIndex.IndexID == t.Index.ID {
+				return true
+			}
+		}
+	}
+	return false
+}

pkg/sql/logictest/testdata/logic_test/alter_table

@@ -3397,7 +3397,7 @@ ALTER TABLE t_99281 ADD COLUMN p INT DEFAULT unique_rowid(), ADD FOREIGN KEY (j)
 # The following statement is not supported using the legacy schema changer.
 skipif config local-legacy-schema-changer
 skipif config local-mixed-23.1
-statement error pq: foreign key violation: "t_99281" row j=0, i=[0-1] has no match in "t_99281_other"
+statement error pq: foreign key violation: "t_99281" row j=0, k=[0-1] has no match in "t_99281_other"
 ALTER TABLE t_99281 ALTER PRIMARY KEY USING COLUMNS (k), ADD FOREIGN KEY (j) REFERENCES t_99281_other;
 
 query TT
@@ -3535,3 +3535,21 @@ CREATE TABLE public.t_118246 (
 )
 
 subtest end
+
+
+# When the new primary key does not have the same columns as the original primary
+# key schema changes recreating secondary indexes could fail incorrectly during validation,
+# if the secondary had some references to the old primary index key (for example
+# partial expressions). This was because an scan of the new secondary index would
+# not return the old primary key, which query execution was expecting.
+subtest 118626
+
+statement ok
+CREATE TABLE public.t_118626(n int primary key, b int NOT NULL, c int NOT NULL);
+INSERT INTO public.t_118626 VALUES(1, 2, 3);
+
+statement ok
+CREATE UNIQUE INDEX ON public.t_118626(c) WHERE n>= 1;
+
+statement ok
+ALTER TABLE public.t_118626 ALTER PRIMARY KEY USING COLUMNS(b);

pkg/sql/schemachanger/testdata/end_to_end/alter_table_alter_primary_key_vanilla/alter_table_alter_primary_key_vanilla.definition

@@ -1,6 +1,9 @@
 setup
-CREATE TABLE t (i INT PRIMARY KEY, j INT NOT NULL);
+CREATE TABLE t (i INT PRIMARY KEY, j INT NOT NULL, k INT DEFAULT 54);
 INSERT INTO t(i, j) VALUES (-4, -4), (-2, -2), (-3, -3);
+CREATE INDEX ON t(i) WHERE i<=0;
+CREATE INDEX ON t(j) WHERE j<= 0;
+CREATE INDEX ON t(k);
 ----
 
 # Note: Unlike other tests, we intentionally avoid