schemachanger: allow dropping column referenced in partial index predicate

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -273,7 +273,9 @@ t7  CREATE TABLE public.t7 (
       INDEX t6i3 (b DESC) WHERE b > 8:::INT8
     ) WITH (schema_locked = true);
 
-# Dropping a column referenced in the predicate drops the index.
+# Dropping a column referenced in the predicate drops the index. The
+# declarative schema changer handles this directly; the legacy schema
+# changer still requires explicitly dropping the index first (see #97813).
 
 statement ok
 CREATE TABLE t8 (
@@ -285,23 +287,28 @@ CREATE TABLE t8 (
     FAMILY (a, b, c)
 )
 
-# TODO(mgartner): Lift this restriction. See #96924.
+onlyif config local-legacy-schema-changer
 statement error cannot drop column "c" because it is referenced by partial index "t8_a_idx1"\nHINT: drop the partial index first, then drop the column
 ALTER TABLE t8 DROP COLUMN c
 
+onlyif config local-legacy-schema-changer
+statement ok
+DROP INDEX t8_a_idx1
+
+statement ok
+ALTER TABLE t8 DROP COLUMN c
+
 onlyif config schema-locked-disabled
 query TT
 SHOW CREATE TABLE t8
 ----
 t8  CREATE TABLE public.t8 (
       a INT8 NULL,
       b INT8 NULL,
-      c STRING NULL,
       rowid INT8 NOT VISIBLE NOT NULL DEFAULT unique_rowid(),
       CONSTRAINT t8_pkey PRIMARY KEY (rowid ASC),
       INDEX t8_a_idx (a ASC) WHERE b > 0:::INT8,
-      INDEX t8_a_idx1 (a ASC) WHERE c = 'foo':::STRING,
-      FAMILY fam_0_a_b_c_rowid (a, b, c, rowid)
+      FAMILY fam_0_a_b_c_rowid (a, b, rowid)
     );
 
 skipif config schema-locked-disabled
@@ -311,12 +318,10 @@ SHOW CREATE TABLE t8
 t8  CREATE TABLE public.t8 (
       a INT8 NULL,
       b INT8 NULL,
-      c STRING NULL,
       rowid INT8 NOT VISIBLE NOT NULL DEFAULT unique_rowid(),
       CONSTRAINT t8_pkey PRIMARY KEY (rowid ASC),
       INDEX t8_a_idx (a ASC) WHERE b > 0:::INT8,
-      INDEX t8_a_idx1 (a ASC) WHERE c = 'foo':::STRING,
-      FAMILY fam_0_a_b_c_rowid (a, b, c, rowid)
+      FAMILY fam_0_a_b_c_rowid (a, b, rowid)
     ) WITH (schema_locked = true);
 
 # CREATE TABLE LIKE ... INCLUDING INDEXES copies partial index predicate
@@ -2368,24 +2373,34 @@ CREATE TABLE t96924 (
 statement ok
 ALTER TABLE t96924 DROP COLUMN a
 
+# Under the declarative schema changer, dropping a column referenced in a
+# partial index predicate drops the dependent index in the same statement.
+# The legacy schema changer still requires the user to drop the index first.
+
+onlyif config local-legacy-schema-changer
 statement error pq: cannot drop column "b" because it is referenced by partial index "t96924_b_key"\nHINT: drop the partial index first, then drop the column
 ALTER TABLE t96924 DROP COLUMN b
 
+onlyif config local-legacy-schema-changer
 statement ok
 DROP INDEX t96924_b_key CASCADE
 
 statement ok
 ALTER TABLE t96924 DROP COLUMN b
 
+onlyif config local-legacy-schema-changer
 statement error pq: cannot drop column "c" because it is referenced by partial index "t96924_c_key"\nHINT: drop the partial index first, then drop the column
 ALTER TABLE t96924 DROP COLUMN c
 
+onlyif config local-legacy-schema-changer
 statement ok
 DROP INDEX t96924_c_key CASCADE
 
 statement ok
 ALTER TABLE t96924 DROP COLUMN c
 
+# Partial unique-without-index predicates still block the column drop in
+# both schema changers (this restriction will be lifted separately).
 statement error pq: cannot drop column "d" because it is referenced by partial unique constraint "unique_d"\nHINT: drop the unique constraint first, then drop the column
 ALTER TABLE t96924 DROP COLUMN d
 
@@ -2395,9 +2410,11 @@ ALTER TABLE t96924 DROP CONSTRAINT unique_d
 statement ok
 ALTER TABLE t96924 DROP COLUMN d
 
+onlyif config local-legacy-schema-changer
 statement error pq: cannot drop column "e" because it is referenced by partial index "t96924_e_f_idx"\nHINT: drop the partial index first, then drop the column
 ALTER TABLE t96924 DROP COLUMN e
 
+onlyif config local-legacy-schema-changer
 statement ok
 DROP INDEX t96924_e_f_idx CASCADE
 
@@ -2409,19 +2426,22 @@ ALTER TABLE t96924 DROP COLUMN e
 statement ok
 ALTER TABLE t96924 DROP COLUMN f
 
-# Column `h` is used in the predicate of another index that does not key on `h`,
-# we will prevent dropping column `h` as well.
+# Column `h` is referenced only in the predicate of another index that does
+# not key on `h`. The declarative schema changer drops both atomically; the
+# legacy schema changer requires dropping the index first.
+onlyif config local-legacy-schema-changer
 statement error pq: cannot drop column "h" because it is referenced by partial index "t96924_g_idx"\nHINT: drop the partial index first, then drop the column
 ALTER TABLE t96924 DROP COLUMN h
 
+onlyif config local-legacy-schema-changer
 statement ok
 DROP INDEX t96924_g_idx
 
 statement ok
 ALTER TABLE t96924 DROP COLUMN h
 
-# Similarly, column `i` cannot be dropped since it's used in the predicate of
-# of a UWI constraint.
+# Column `i` is used in the predicate of a UWI constraint, which still
+# blocks the drop in both schema changers.
 statement error pq: cannot drop column "i" because it is referenced by partial unique constraint "unique_g"\nHINT: drop the unique constraint first, then drop the column
 ALTER TABLE t96924 DROP COLUMN i
 

pkg/sql/schemachanger/dml_injection_test.go

@@ -543,7 +543,6 @@ func TestAlterTableDMLInjection(t *testing.T) {
 				"CREATE INDEX idx ON tbl (val) WHERE i > 1",
 			},
 			schemaChange: "ALTER TABLE tbl DROP COLUMN i",
-			skipIssue:    97813,
 		},
 		{
 			desc:         "create expression index",

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

@@ -89,7 +89,7 @@ func alterTableAlterColumnType(
 		case *scpb.PolicyUsingExpr, *scpb.PolicyWithCheckExpr:
 			panic(sqlerrors.NewAlterDependsOnPolicyExprError(op, objType, t.Column.String()))
 		}
-	}, false /* allowPartialIdxPredicateRef */)
+	}, disallowAllPredicateRefs)
 
 	var err error
 	newColType.Type, err = schemachange.ValidateAlterColumnTypeChecks(
@@ -304,7 +304,7 @@ func handleGeneralColumnConversion(
 		case *scpb.SecondaryIndex:
 			panic(sqlerrors.NewAlterColumnTypeColInIndexNotSupportedErr())
 		}
-	}, false /* allowPartialIdxPredicateRef */)
+	}, disallowAllPredicateRefs)
 
 	// This code path should never be reached for virtual columns, as their values
 	// are always computed dynamically on access and are never stored on disk.
@@ -479,7 +479,7 @@ func maybeWriteNoticeForFKColTypeMismatch(b BuildCtx, col *scpb.Column, colType
 		case *scpb.ForeignKeyConstraintUnvalidated:
 			writeNoticeHelper(e.ColumnIDs, e.ReferencedColumnIDs, e.ReferencedTableID)
 		}
-	}, false /* allowPartialIdxPredicateRef */)
+	}, disallowAllPredicateRefs)
 }
 
 func getComputeExpressionForBackfill(

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

@@ -1376,7 +1376,7 @@ func checkIfColumnCanBeDropped(b BuildCtx, columnToDrop *scpb.Column) bool {
 				canBeDropped = false
 			}
 		}
-	}, false /* allowPartialIdxPredicateRef */)
+	}, disallowAllPredicateRefs)
 	return canBeDropped
 }
 

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

@@ -309,7 +309,7 @@ func dropColumn(
 		default:
 			b.Drop(e)
 		}
-	}, false /* allowPartialIdxPredicateRef */)
+	}, disallowUWIPredicateRefs)
 	// TODO(ajwerner): Track the undropped backrefs to populate a detail
 	// message like postgres does. For example:
 	//  SET serial_normalization = sql_sequence;
@@ -333,25 +333,49 @@ func dropColumn(
 	assertAllColumnElementsAreDropped(colElts)
 }
 
+// predicateRefPolicy controls how walkColumnDependencies handles columns
+// that are referenced in a partial-index predicate or partial
+// unique-without-index predicate.
+type predicateRefPolicy int
+
+const (
+	// disallowAllPredicateRefs panics if the column is referenced in either
+	// a partial-index predicate or a partial UWI predicate. Used by code
+	// paths (ALTER PRIMARY KEY, ALTER COLUMN TYPE) where the dependency
+	// rules don't yet sequence the partial-index drop safely.
+	disallowAllPredicateRefs predicateRefPolicy = iota
+	// disallowUWIPredicateRefs panics only for partial UWI references.
+	// Partial-index references are allowed because the declarative schema
+	// changer's dep rule sequences the partial index into DELETE_ONLY
+	// before the column reaches WRITE_ONLY, allowing the index entries to
+	// be removed via blind DELs without evaluating the predicate.
+	disallowUWIPredicateRefs
+	// allowAllPredicateRefs skips the check entirely. Used by ALTER
+	// COLUMN ... RENAME, which doesn't change column existence.
+	allowAllPredicateRefs
+)
+
 func walkColumnDependencies(
 	b BuildCtx,
 	col *scpb.Column,
 	op string,
 	objType string,
 	fn func(e scpb.Element, op string, objType string),
-	allowPartialIdxPredicateRef bool,
+	predicateRefs predicateRefPolicy,
 ) {
 	var sequenceDeps catalog.DescriptorIDSet
 	var indexDeps catid.IndexSet
 	var columnDeps catalog.TableColSet
 	var constraintDeps catid.ConstraintSet
 	tblElts := b.QueryByID(col.TableID).Filter(orFilter(publicTargetFilter, transientTargetFilter))
 
-	// Panic if `col` is referenced in a predicate of an index or
-	// unique without index constraint.
-	// TODO (xiang): Remove this restriction when #97813 is fixed.
-	if !allowPartialIdxPredicateRef {
+	switch predicateRefs {
+	case disallowAllPredicateRefs:
 		panicIfColReferencedInPredicate(b, col, tblElts, op, objType)
+	case disallowUWIPredicateRefs:
+		panicIfColReferencedInUWIPredicate(b, col, tblElts, op, objType)
+	case allowAllPredicateRefs:
+		// No-op.
 	}
 
 	tblElts.
@@ -525,6 +549,47 @@ func panicIfColReferencedInPredicate(
 	}
 }
 
+// panicIfColReferencedInUWIPredicate disallows dropping a column that is
+// referenced in the predicate of a partial unique-without-index constraint.
+// Unlike partial indexes, UWI predicate evaluation during the column-drop
+// transition is not yet sequenced safely by the declarative schema changer's
+// dependency rules.
+func panicIfColReferencedInUWIPredicate(
+	b BuildCtx, col *scpb.Column, tblElts ElementResultSet, op, objType string,
+) {
+	contains := func(container []catid.ColumnID, target catid.ColumnID) bool {
+		for _, elem := range container {
+			if elem == target {
+				return true
+			}
+		}
+		return false
+	}
+
+	var violatingUWI catid.ConstraintID
+	tblElts.ForEach(func(_ scpb.Status, _ scpb.TargetStatus, e scpb.Element) {
+		if violatingUWI != 0 {
+			return
+		}
+		switch elt := e.(type) {
+		case *scpb.UniqueWithoutIndexConstraint:
+			if elt.Predicate != nil && contains(elt.Predicate.ReferencedColumnIDs, col.ColumnID) {
+				violatingUWI = elt.ConstraintID
+			}
+		case *scpb.UniqueWithoutIndexConstraintUnvalidated:
+			if elt.Predicate != nil && contains(elt.Predicate.ReferencedColumnIDs, col.ColumnID) {
+				violatingUWI = elt.ConstraintID
+			}
+		}
+	})
+	if violatingUWI != 0 {
+		colNameElem := mustRetrieveColumnNameElem(b, col.TableID, col.ColumnID)
+		uwiNameElem := mustRetrieveConstraintWithoutIndexNameElem(b, col.TableID, violatingUWI)
+		panic(sqlerrors.ColumnReferencedByPartialUniqueWithoutIndexConstraint(
+			op, objType, colNameElem.Name, uwiNameElem.Name))
+	}
+}
+
 func handleDropColumnPrimaryIndexes(b BuildCtx, tbl *scpb.Table, col *scpb.Column) {
 	inflatedChain := getInflatedPrimaryIndexChain(b, tbl.TableID)
 

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

@@ -95,7 +95,7 @@ func renameColumnChecks(b BuildCtx, col *scpb.Column, oldName, newName tree.Name
 				tree.ErrString(&oldName), fnName.Name),
 			)
 		}
-	}, true /* allowPartialIdxPredicateRef */)
+	}, allowAllPredicateRefs)
 
 	if col.IsInaccessible {
 		panic(pgerror.Newf(