schemachanger: address index validation errors during pk swap

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