Merge cockroachdb#97372

97372: sql: disallow dropping column referenced in partial index predicate r=Xiang-Gu a=mgartner

Informs cockroachdb#96924

Release note (bug fix/sql change): Columns referenced in partial index
predicates and partial unique constraint predicates can no longer be
dropped. The `ALTER TABLE .. DROP COLUMN` statement now returns an error
with a hint suggesting to drop the indexes and constraints first. This
is a temporary safe-guard to prevent users from hitting cockroachdb#96924. This
restriction will be lifted when that bug is fixed.


Co-authored-by: Xiang Gu <[email protected]>

pkg/sql/alter_table.go

@@ -1653,7 +1653,7 @@ func dropColumnImpl(
 		containsThisColumn := idx.CollectKeyColumnIDs().Contains(colToDrop.GetID()) ||
 			idx.CollectKeySuffixColumnIDs().Contains(colToDrop.GetID()) ||
 			idx.CollectSecondaryStoredColumnIDs().Contains(colToDrop.GetID())
-		if !containsThisColumn && idx.IsPartial() {
+		if idx.IsPartial() {
 			expr, err := parser.ParseExpr(idx.GetPredicate())
 			if err != nil {
 				return nil, err
@@ -1666,6 +1666,7 @@ func dropColumnImpl(
 
 			if colIDs.Contains(colToDrop.GetID()) {
 				containsThisColumn = true
+				return nil, sqlerrors.NewColumnReferencedByPartialIndex(string(colToDrop.ColName()), idx.GetName())
 			}
 		}
 		// Perform the DROP.
@@ -1691,6 +1692,21 @@ func dropColumnImpl(
 
 	// Drop non-index-backed unique constraints which reference the column.
 	for _, uwoi := range tableDesc.EnforcedUniqueConstraintsWithoutIndex() {
+		if uwoi.IsPartial() {
+			expr, err := parser.ParseExpr(uwoi.GetPredicate())
+			if err != nil {
+				return nil, err
+			}
+
+			colIDs, err := schemaexpr.ExtractColumnIDs(tableDesc, expr)
+			if err != nil {
+				return nil, err
+			}
+
+			if colIDs.Contains(colToDrop.GetID()) {
+				return nil, sqlerrors.NewColumnReferencedByPartialUniqueWithoutIndexConstraint(string(colToDrop.ColName()), uwoi.GetName())
+			}
+		}
 		if uwoi.Dropped() || !uwoi.CollectKeyColumnIDs().Contains(colToDrop.GetID()) {
 			continue
 		}

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -186,7 +186,8 @@ CREATE TABLE t8 (
     FAMILY (a, b, c)
 )
 
-statement ok
+# TODO(mgartner): Lift this restriction. See #96924.
+statement error column "c" cannot be dropped because it is referenced by partial index "t8_a_idx1"
 ALTER TABLE t8 DROP COLUMN c
 
 query TT
@@ -195,10 +196,12 @@ 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,
-      FAMILY fam_0_a_b_c_rowid (a, b, rowid)
+      INDEX t8_a_idx1 (a ASC) WHERE c = 'foo':::STRING,
+      FAMILY fam_0_a_b_c_rowid (a, b, c, rowid)
     )
 
 # CREATE TABLE LIKE ... INCLUDING INDEXES copies partial index predicate
@@ -1911,3 +1914,123 @@ BEGIN;
 COMMIT;
 SELECT * FROM t79613;
 DROP TABLE t79613;
+
+# Regression test for #96924 to disallow dropping a column that is referenced in
+# a partial index predicate.
+subtest column_dropped_referenced_in_partial_index
+
+statement ok
+SET experimental_enable_unique_without_index_constraints = true
+
+statement ok
+CREATE TABLE t96924 (
+  a INT,
+  b INT,
+  c INT,
+  d INT,
+  e INT,
+  f JSON,
+  g INT,
+  h INT,
+  i INT,
+  INDEX (a) WHERE (b > 0),
+  UNIQUE INDEX (b) WHERE (b > 0),
+  UNIQUE (c) WHERE (c > 0),
+  UNIQUE WITHOUT INDEX (d) WHERE (d > 0),
+  INVERTED INDEX (e, f) WHERE (e > 0),
+  INDEX (g) WHERE (h > 0),
+  UNIQUE WITHOUT INDEX (g) WHERE (i > 0)
+)
+
+# Column `a` can be dropped because it is not referenced in any partial index
+# predicates.
+statement ok
+ALTER TABLE t96924 DROP COLUMN a
+
+statement error pq: column "b" cannot be dropped because it is referenced by partial index "t96924_b_key"
+ALTER TABLE t96924 DROP COLUMN b
+
+statement ok
+DROP INDEX t96924_b_key CASCADE
+
+statement ok
+ALTER TABLE t96924 DROP COLUMN b
+
+statement error pq: column "c" cannot be dropped because it is referenced by partial index "t96924_c_key"
+ALTER TABLE t96924 DROP COLUMN c
+
+statement ok
+DROP INDEX t96924_c_key CASCADE
+
+statement ok
+ALTER TABLE t96924 DROP COLUMN c
+
+statement error pq: column "d" cannot be dropped because it is referenced by partial unique constraint "unique_d"
+ALTER TABLE t96924 DROP COLUMN d
+
+statement ok
+ALTER TABLE t96924 DROP CONSTRAINT unique_d
+
+statement ok
+ALTER TABLE t96924 DROP COLUMN d
+
+statement error pq: column "e" cannot be dropped because it is referenced by partial index "t96924_e_f_idx"
+ALTER TABLE t96924 DROP COLUMN e
+
+statement ok
+DROP INDEX t96924_e_f_idx CASCADE
+
+statement ok
+ALTER TABLE t96924 DROP COLUMN e
+
+# Column `f` can be dropped because it is not referenced in any partial index
+# predicates.
+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.
+statement error pq: column "h" cannot be dropped because it is referenced by partial index "t96924_g_idx"
+ALTER TABLE t96924 DROP COLUMN h
+
+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.
+statement error pq: column "i" cannot be dropped because it is referenced by partial unique constraint "unique_g"
+ALTER TABLE t96924 DROP COLUMN i
+
+statement ok
+ALTER TABLE t96924 DROP CONSTRAINT unique_g
+
+statement ok
+ALTER TABLE t96924 DROP COLUMN i
+
+statement ok
+ALTER TABLE t96924 DROP COLUMN g
+
+# Another edge case where ALTER PRIMARY KEY implicit drops the `rowid` column
+# yet `rowid` is referenced in a partial index. Note that this behavior is
+# only in the declarative schema changer but not in the legacy schema changer.
+statement ok
+DROP TABLE t96924
+
+statement ok
+CREATE TABLE t96924 (a INT NOT NULL)
+
+statement ok
+CREATE INDEX t96924_idx_1 ON t96924(a) WHERE (rowid > 0);
+
+skipif config local-legacy-schema-changer
+statement error column "rowid" cannot be dropped because it is referenced by partial index "t96924_idx_1"
+ALTER TABLE t96924 ALTER PRIMARY KEY USING COLUMNS (a);
+
+statement ok
+DROP INDEX t96924_idx_1
+
+statement ok
+ALTER TABLE t96924 ALTER PRIMARY KEY USING COLUMNS (a);

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

@@ -508,6 +508,40 @@ func mustRetrieveKeyIndexColumns(
 	return indexColumns
 }
 
+func mustRetrieveIndexNameElem(
+	b BuildCtx, tableID catid.DescID, indexID catid.IndexID,
+) (indexNameElem *scpb.IndexName) {
+	scpb.ForEachIndexName(b.QueryByID(tableID), func(
+		current scpb.Status, target scpb.TargetStatus, e *scpb.IndexName,
+	) {
+		if e.IndexID == indexID {
+			indexNameElem = e
+		}
+	})
+	if indexNameElem == nil {
+		panic(errors.AssertionFailedf("programming error: cannot find an index name element "+
+			"with ID %v from table %v", indexID, tableID))
+	}
+	return indexNameElem
+}
+
+func mustRetrieveConstraintWithoutIndexNameElem(
+	b BuildCtx, tableID catid.DescID, constraintID catid.ConstraintID,
+) (constraintWithoutIndexName *scpb.ConstraintWithoutIndexName) {
+	scpb.ForEachConstraintWithoutIndexName(b.QueryByID(tableID), func(
+		current scpb.Status, target scpb.TargetStatus, e *scpb.ConstraintWithoutIndexName,
+	) {
+		if e.ConstraintID == constraintID {
+			constraintWithoutIndexName = e
+		}
+	})
+	if constraintWithoutIndexName == nil {
+		panic(errors.AssertionFailedf("programming error: cannot find a constraint name "+
+			"element with ID %v from table %v", constraintID, tableID))
+	}
+	return constraintWithoutIndexName
+}
+
 func checkIfConstraintNameAlreadyExists(b BuildCtx, tbl *scpb.Table, t alterPrimaryKeySpec) {
 	if t.Name == "" {
 		return

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

@@ -304,6 +304,10 @@ func walkDropColumnDependencies(b BuildCtx, col *scpb.Column, fn func(e scpb.Ele
 	var indexesToDrop catid.IndexSet
 	var columnsToDrop catalog.TableColSet
 	tblElts := b.QueryByID(col.TableID).Filter(publicTargetFilter)
+	// Panic if `col` is referenced in a predicate of an index or
+	// unique without index constraint.
+	// TODO (xiang): Remove this restriction when #96924 is fixed.
+	panicIfColReferencedInPredicate(b, col, tblElts)
 	tblElts.
 		Filter(referencesColumnIDFilter(col.ColumnID)).
 		ForEachElementStatus(func(_ scpb.Status, _ scpb.TargetStatus, e scpb.Element) {
@@ -323,6 +327,8 @@ func walkDropColumnDependencies(b BuildCtx, col *scpb.Column, fn func(e scpb.Ele
 			case *scpb.SequenceOwner:
 				fn(e)
 				sequencesToDrop.Add(elt.SequenceID)
+			case *scpb.SecondaryIndex:
+				indexesToDrop.Add(elt.IndexID)
 			case *scpb.SecondaryIndexPartial:
 				indexesToDrop.Add(elt.IndexID)
 			case *scpb.IndexColumn:
@@ -350,14 +356,7 @@ func walkDropColumnDependencies(b BuildCtx, col *scpb.Column, fn func(e scpb.Ele
 		case *scpb.SecondaryIndex:
 			if indexesToDrop.Contains(elt.IndexID) {
 				fn(e)
-			} else if elt.EmbeddedExpr != nil {
-				for _, columnID := range elt.EmbeddedExpr.ReferencedColumnIDs {
-					if columnID == col.ColumnID {
-						fn(e)
-					}
-				}
 			}
-
 		}
 	})
 	sequencesToDrop.ForEach(func(id descpb.ID) {
@@ -391,6 +390,56 @@ func walkDropColumnDependencies(b BuildCtx, col *scpb.Column, fn func(e scpb.Ele
 	})
 }
 
+// panicIfColReferencedInPredicate is a temporary fix that disallow dropping a
+// column that is referenced in predicate of a partial index or unique without index.
+// This restriction shall be lifted once #96924 is fixed.
+func panicIfColReferencedInPredicate(b BuildCtx, col *scpb.Column, tblElts ElementResultSet) {
+	contains := func(container []catid.ColumnID, target catid.ColumnID) bool {
+		for _, elem := range container {
+			if elem == target {
+				return true
+			}
+		}
+		return false
+	}
+
+	var violatingIndex catid.IndexID
+	var violatingUWI catid.ConstraintID
+	tblElts.ForEachElementStatus(func(_ scpb.Status, _ scpb.TargetStatus, e scpb.Element) {
+		if violatingIndex != 0 || violatingUWI != 0 {
+			return
+		}
+		switch elt := e.(type) {
+		case *scpb.SecondaryIndex:
+			if elt.EmbeddedExpr != nil && contains(elt.EmbeddedExpr.ReferencedColumnIDs, col.ColumnID) {
+				violatingIndex = elt.IndexID
+			}
+		case *scpb.SecondaryIndexPartial:
+			if contains(elt.ReferencedColumnIDs, col.ColumnID) {
+				violatingIndex = elt.IndexID
+			}
+		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 violatingIndex != 0 {
+		colNameElem := mustRetrieveColumnNameElem(b, col.TableID, col.ColumnID)
+		indexNameElem := mustRetrieveIndexNameElem(b, col.TableID, violatingIndex)
+		panic(sqlerrors.NewColumnReferencedByPartialIndex(colNameElem.Name, indexNameElem.Name))
+	}
+	if violatingUWI != 0 {
+		colNameElem := mustRetrieveColumnNameElem(b, col.TableID, col.ColumnID)
+		uwiNameElem := mustRetrieveConstraintWithoutIndexNameElem(b, col.TableID, violatingUWI)
+		panic(sqlerrors.NewColumnReferencedByPartialUniqueWithoutIndexConstraint(colNameElem.Name, uwiNameElem.Name))
+	}
+}
+
 func handleDropColumnPrimaryIndexes(
 	b BuildCtx, tbl *scpb.Table, n tree.NodeFormatter, col *scpb.Column,
 ) {

pkg/sql/schemachanger/scplan/internal/rules/current/dep_drop_column.go

@@ -158,7 +158,7 @@ func init() {
 		"secondary index partial no longer public before referenced column",
 		scgraph.Precedence,
 		"secondary-partial-index", "column",
-		scpb.Status_DELETE_ONLY, scpb.Status_WRITE_ONLY,
+		scpb.Status_ABSENT, scpb.Status_WRITE_ONLY,
 		func(from, to NodeVars) rel.Clauses {
 			return rel.Clauses{
 				from.Type((*scpb.SecondaryIndexPartial)(nil)),