Merge pull request #97663 from Xiang-Gu/backport22.1-97372

pkg/sql/alter_table.go

@@ -61,8 +61,9 @@ type alterTableNode struct {
 
 // AlterTable applies a schema change on a table.
 // Privileges: CREATE on table.
-//   notes: postgres requires CREATE on the table.
-//          mysql requires ALTER, CREATE, INSERT on the table.
+//
+//	notes: postgres requires CREATE on the table.
+//	       mysql requires ALTER, CREATE, INSERT on the table.
 func (p *planner) AlterTable(ctx context.Context, n *tree.AlterTable) (planNode, error) {
 	if err := checkSchemaChangeEnabled(
 		ctx,
@@ -1708,7 +1709,7 @@ func dropColumnImpl(
 
 		// If the column being dropped is referenced in the partial
 		// index predicate, then the index should be dropped.
-		if !containsThisColumn && idx.IsPartial() {
+		if idx.IsPartial() {
 			expr, err := parser.ParseExpr(idx.GetPredicate())
 			if err != nil {
 				return nil, err
@@ -1721,6 +1722,7 @@ func dropColumnImpl(
 
 			if colIDs.Contains(colToDrop.GetID()) {
 				containsThisColumn = true
+				return nil, sqlerrors.NewColumnReferencedByPartialIndex(colToDrop.GetName(), idx.GetName())
 			}
 		}
 
@@ -1751,6 +1753,25 @@ func dropColumnImpl(
 		constraint := &tableDesc.UniqueWithoutIndexConstraints[i]
 		tableDesc.UniqueWithoutIndexConstraints[sliceIdx] = *constraint
 		sliceIdx++
+
+		// Temporarily disallow dropping a column that is referenced in the
+		// predicate of a UWI constraint. Lift when #96924 is fixed.
+		if constraint.IsPartial() {
+			expr, err := parser.ParseExpr(constraint.Predicate)
+			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()), constraint.GetName())
+			}
+		}
+
 		if descpb.ColumnIDs(constraint.ColumnIDs).Contains(colToDrop.GetID()) {
 			sliceIdx--
 

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
@@ -1906,3 +1909,101 @@ 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

pkg/sql/sqlerrors/errors.go

@@ -250,6 +250,35 @@ func NewAggInAggError() error {
 	return pgerror.New(pgcode.Grouping, "aggregate function calls cannot be nested")
 }
 
+// NewColumnReferencedByPartialIndex is returned when we drop a column that is
+// referenced in a partial index's predicate.
+func NewColumnReferencedByPartialIndex(droppingColumn, partialIndex string) error {
+	return errors.WithIssueLink(errors.WithHint(
+		pgerror.Newf(
+			pgcode.InvalidColumnReference,
+			"column %q cannot be dropped because it is referenced by partial index %q",
+			droppingColumn, partialIndex,
+		),
+		"drop the partial index first, then drop the column",
+	), errors.IssueLink{IssueURL: "https://github.com/cockroachdb/cockroach/pull/97372"})
+}
+
+// NewColumnReferencedByPartialUniqueWithoutIndexConstraint is almost the same as
+// NewColumnReferencedByPartialIndex except it's used when dropping column that is
+// referenced in a partial unique without index constraint's predicate.
+func NewColumnReferencedByPartialUniqueWithoutIndexConstraint(
+	droppingColumn, partialUWIConstraint string,
+) error {
+	return errors.WithIssueLink(errors.WithHint(
+		pgerror.Newf(
+			pgcode.InvalidColumnReference,
+			"column %q cannot be dropped because it is referenced by partial unique constraint %q",
+			droppingColumn, partialUWIConstraint,
+		),
+		"drop the unique constraint first, then drop the column",
+	), errors.IssueLink{IssueURL: "https://github.com/cockroachdb/cockroach/pull/97372"})
+}
+
 // QueryTimeoutError is an error representing a query timeout.
 var QueryTimeoutError = pgerror.New(
 	pgcode.QueryCanceled, "query execution canceled due to statement timeout")