[#22902] YSQL: Add support for dropping a key column

Summary:
Add support for dropping a key column using table rewrite.
Note: dropping any key column also drops the primary key on the table.

Code changes:
 - ATExecDropColumn: mark the table and all of its children to be rewritten in phase 3 (for reason YB_AT_REWRITE_ALTER_PRIMARY_KEY) when the column to be dropped is a key column.
 - YBCPrepareAlterTableCmd: skip yb drops for key columns in the alter table execution flow, as the table will be rewritten in alter table phase 3 anyway.
 - YBCDropColumn: note: this function is called when the column is being dropped outside of an alter table flow (see commit 7c8343d).
        - Fix the condition for missing columns: previously, the function would skip over missing columns only when the IF EXISTS clause was present. Without this clause, it would unnecessarily create a drop column handle for a missing column, as the absence of the column is detected later in ATExecDropColumn. Update the condition to always bypass YB alters for missing columns.
        - Since dropping a key column requires a table rewrite in YB, we must appropriately trigger the table rewrite event. Construct a dummy query for this purpose, since we don't have the original query available.
        - Perform the usual ALTER TABLE execution flow by invoking AlterTableInternal().
        - If we are not dropping a key column, we can use the short path (introduced by commit 7c8343d) of simply executing the YB drop column.
Jira: DB-11810

Test Plan: yb_feature_alter_table, yb_alter_table_rewrite

Reviewers: myang

Reviewed By: myang

Subscribers: yql

Differential Revision: https://phorge.dev.yugabyte.com/D35940

src/postgres/src/backend/commands/tablecmds.c

@@ -432,7 +432,8 @@ static ObjectAddress ATExecSetStorage(Relation rel, const char *colName,
 				 Node *newValue, LOCKMODE lockmode);
 static void ATPrepDropColumn(List **wqueue, Relation rel, bool recurse, bool recursing,
 				 AlterTableCmd *cmd, LOCKMODE lockmode);
-static ObjectAddress ATExecDropColumn(List **wqueue, Relation rel, const char *colName,
+static ObjectAddress ATExecDropColumn(List **wqueue, AlteredTableInfo *yb_tab, Relation rel,
+				 const char *colName,
 				 DropBehavior behavior,
 				 bool recurse, bool recursing,
 				 bool missing_ok, LOCKMODE lockmode);
@@ -4585,12 +4586,12 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab, Relation *mutable_rel,
 			address = ATExecSetStorage(rel, cmd->name, cmd->def, lockmode);
 			break;
 		case AT_DropColumn:		/* DROP COLUMN */
-			address = ATExecDropColumn(wqueue, rel, cmd->name,
+			address = ATExecDropColumn(wqueue, tab, rel, cmd->name,
 									   cmd->behavior, false, false,
 									   cmd->missing_ok, lockmode);
 			break;
 		case AT_DropColumnRecurse:	/* DROP COLUMN with recursion */
-			address = ATExecDropColumn(wqueue, rel, cmd->name,
+			address = ATExecDropColumn(wqueue, tab, rel, cmd->name,
 									   cmd->behavior, true, false,
 									   cmd->missing_ok, lockmode);
 			break;
@@ -7321,7 +7322,8 @@ ATPrepDropColumn(List **wqueue, Relation rel, bool recurse, bool recursing,
  * Return value is the address of the dropped column.
  */
 static ObjectAddress
-ATExecDropColumn(List **wqueue, Relation rel, const char *colName,
+ATExecDropColumn(List **wqueue, AlteredTableInfo *yb_tab, Relation rel,
+				 const char *colName,
 				 DropBehavior behavior,
 				 bool recurse, bool recursing,
 				 bool missing_ok, LOCKMODE lockmode)
@@ -7363,20 +7365,19 @@ ATExecDropColumn(List **wqueue, Relation rel, const char *colName,
 	attnum = targetatt->attnum;
 
 	/*
-	 * In YB, a table cannot drop key columns.
-	 * This check makes sure a consistent state after attempting to drop
-	 * key columns by preventing dropping key columns on postgres side.
+	 * In YB, dropping a key column requires a table rewrite.
 	 */
-	if (IsYBRelation(rel))
+	if (IsYBRelation(rel) && YbIsAttrPrimaryKeyColumn(rel, attnum))
 	{
-		Bitmapset  *pkey   = YBGetTablePrimaryKeyBms(rel);
-
-		/* Can't drop primary-key columns */
-		if (bms_is_member(attnum - YBGetFirstLowInvalidAttributeNumber(rel), pkey))
-			ereport(ERROR,
-					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					 errmsg("cannot drop key column \"%s\"",
-							colName)));
+		/*
+		 * In YB, the ADD/DROP primary key operation involves a table
+		 * rewrite. So if this is partitioned table, we need to add
+		 * its children to the work queue as well.
+		 */
+		if ((rel)->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+			YbATSetPKRewriteChildPartitions(wqueue,
+				yb_tab, false /* skip_copy_split_options */);
+		yb_tab->rewrite |= YB_AT_REWRITE_ALTER_PRIMARY_KEY;
 	}
 
 	/* Can't drop a system attribute, except OID */
@@ -7463,7 +7464,7 @@ ATExecDropColumn(List **wqueue, Relation rel, const char *colName,
 				if (childatt->attinhcount == 1 && !childatt->attislocal)
 				{
 					/* Time to delete this child column, too */
-					ATExecDropColumn(wqueue, childrel, colName,
+					ATExecDropColumn(wqueue, yb_tab, childrel, colName,
 									 behavior, true, true,
 									 false, lockmode);
 				}

src/postgres/src/backend/commands/ybccmds.c

@@ -1297,15 +1297,18 @@ YBCPrepareAlterTableCmd(AlterTableCmd* cmd, Relation rel, List *handles,
 		case AT_DropColumn:
 		case AT_DropColumnRecurse:
 		{
-			/* Skip yb alter for IF EXISTS with non-existent column */
-			if (cmd->missing_ok)
-			{
-				HeapTuple tuple = SearchSysCacheAttName(relationId, cmd->name);
-				if (!HeapTupleIsValid(tuple))
-					break;
-				ReleaseSysCache(tuple);
-			}
-
+			HeapTuple tuple = SearchSysCacheAttName(relationId, cmd->name);
+			/* Skip yb alter for non-existent column */
+			if (!HeapTupleIsValid(tuple))
+				break;
+			AttrNumber attnum = ((Form_pg_attribute) GETSTRUCT(tuple))->attnum;
+			ReleaseSysCache(tuple);
+			/*
+			 * Skip yb alter for primary key columns (the table will be
+			 * rewritten)
+			 */
+			if (YbIsAttrPrimaryKeyColumn(rel, attnum))
+				break;
 			Assert(list_length(handles) == 1);
 			YBCPgStatement drop_col_handle =
 				(YBCPgStatement) lfirst(list_head(handles));
@@ -2037,13 +2040,45 @@ YBCDropColumn(Relation rel, AttrNumber attnum)
 {
 	TupleDesc tupleDesc = RelationGetDescr(rel);
 	Form_pg_attribute attr = TupleDescAttr(tupleDesc, attnum - 1);
-	YBCPgStatement handle = NULL;
-	HandleYBStatus(YBCPgNewAlterTable(
-		YBCGetDatabaseOidByRelid(RelationGetRelid(rel)),
-		YbGetRelfileNodeId(rel),
-		&handle));
-	HandleYBStatus(YBCPgAlterTableDropColumn(
-		handle,
-		attr->attname.data));
-	HandleYBStatus(YBCPgExecAlterTable(handle));
+
+	if (YbIsAttrPrimaryKeyColumn(rel, attnum))
+	{
+		/*
+		 * In YB, dropping a primary key involves a table
+		 * rewrite, so invoke the entire ALTER TABLE logic.
+		 */
+
+		/* Construct a dummy query, as we don't have the original query. */
+		char *query_str = psprintf("ALTER TABLE %s DROP COLUMN %s",
+			quote_qualified_identifier(
+				get_namespace_name(rel->rd_rel->relnamespace),
+				RelationGetRelationName(rel)),
+			attr->attname.data);
+		RawStmt *rawstmt =
+			(RawStmt *) linitial(raw_parser(query_str));
+
+		/* Construct the ALTER TABLE command. */
+		AlterTableCmd *cmd = makeNode(AlterTableCmd);
+		cmd->subtype = AT_DropColumn;
+		cmd->name = attr->attname.data;
+		List *cmds = list_make1(cmd);
+
+		EventTriggerAlterTableStart((Node *) rawstmt->stmt);
+		AlterTableInternal(RelationGetRelid(rel), cmds, true);
+		EventTriggerAlterTableEnd();
+
+		pfree(query_str);
+	}
+	else
+	{
+		YBCPgStatement handle = NULL;
+		HandleYBStatus(YBCPgNewAlterTable(
+			YBCGetDatabaseOidByRelid(RelationGetRelid(rel)),
+			YbGetRelfileNodeId(rel),
+			&handle));
+		HandleYBStatus(YBCPgAlterTableDropColumn(
+			handle,
+			attr->attname.data));
+		HandleYBStatus(YBCPgExecAlterTable(handle));
+	}
 }

src/postgres/src/backend/utils/misc/pg_yb_utils.c

@@ -4974,3 +4974,11 @@ YbReadTimePointHandle YbBuildCurrentReadTimePointHandle()
 bool YbUseFastBackwardScan() {
   return *(YBCGetGFlags()->ysql_use_fast_backward_scan);
 }
+
+/* Used in YB to check if an attribute is a key column. */
+bool YbIsAttrPrimaryKeyColumn(Relation rel, AttrNumber attnum)
+{
+	Bitmapset *pkey = YBGetTablePrimaryKeyBms(rel);
+	return bms_is_member(attnum -
+		YBGetFirstLowInvalidAttributeNumber(rel), pkey);
+}

src/postgres/src/include/pg_yb_utils.h

@@ -1156,4 +1156,6 @@ extern YbReadTimePointHandle YbBuildCurrentReadTimePointHandle();
 
 extern bool YbUseFastBackwardScan();
 
+bool YbIsAttrPrimaryKeyColumn(Relation rel, AttrNumber attnum);
+
 #endif /* PG_YB_UTILS_H */

src/postgres/src/test/regress/expected/yb_alter_table_rewrite.out

@@ -1038,3 +1038,55 @@ SELECT * FROM test_identity;
   3 |   3
   4 |   4
 (4 rows)
+
+--
+-- Test ALTER TABLE ... DROP COLUMN on a primary key column.
+--
+-- basic test
+CREATE TABLE test_drop_pk_column (id int PRIMARY KEY, v int);
+INSERT INTO test_drop_pk_column VALUES (1, 1), (2, 2), (3, 3);
+ALTER TABLE test_drop_pk_column DROP COLUMN id;
+\d test_drop_pk_column;
+        Table "public.test_drop_pk_column"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ v      | integer |           |          |
+Tablespace: "test_tblspc"
+
+INSERT INTO test_drop_pk_column VALUES (1);
+SELECT * FROM test_drop_pk_column ORDER BY v;
+ v
+---
+ 1
+ 1
+ 2
+ 3
+(4 rows)
+
+-- test composite PK
+CREATE TABLE test_drop_pk_column_composite (id int, v int, t int, PRIMARY KEY (id HASH, v ASC));
+INSERT INTO test_drop_pk_column_composite VALUES (1, 1, 1), (2, 2, 2), (3, 3, 3);
+ALTER TABLE test_drop_pk_column_composite DROP COLUMN id;
+\d test_drop_pk_column_composite;
+   Table "public.test_drop_pk_column_composite"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ v      | integer |           | not null |
+ t      | integer |           |          |
+Tablespace: "test_tblspc"
+
+INSERT INTO test_drop_pk_column_composite VALUES (1, 1);
+SELECT * FROM test_drop_pk_column_composite ORDER BY v;
+ v | t
+---+---
+ 1 | 1
+ 1 | 1
+ 2 | 2
+ 3 | 3
+(4 rows)
+
+-- test partitioned table
+CREATE TABLE test_drop_pk_column_part (id int PRIMARY KEY, v int) PARTITION BY RANGE (id);
+CREATE TABLE test_drop_pk_column_part1 PARTITION OF test_drop_pk_column_part FOR VALUES FROM (1) TO (10);
+ALTER TABLE test_drop_pk_column_part DROP COLUMN id; -- should fail.
+ERROR:  cannot drop column named in partition key

src/postgres/src/test/regress/expected/yb_alter_type.out

@@ -132,15 +132,13 @@ DROP TYPE new_name;  -- fail
 ERROR:  cannot drop type new_name because other objects depend on it
 DETAIL:  column happiness_level of table holidays depends on type new_name
 HINT:  Use DROP ... CASCADE to drop the dependent objects too.
-DROP TYPE new_name CASCADE; -- fail (happiness_level is a key column)
+DROP TYPE new_name CASCADE; -- requires table rewrite (happiness_level is a key column)
 NOTICE:  drop cascades to column happiness_level of table holidays
-ERROR:  cannot remove a key column
-ALTER TABLE holidays DROP CONSTRAINT holidays_pkey;
 NOTICE:  table rewrite may lead to inconsistencies
 DETAIL:  Concurrent DMLs may not be reflected in the new table.
 HINT:  See https://github.com/yugabyte/yugabyte-db/issues/19860. Set 'ysql_suppress_unsafe_alter_notice' yb-tserver gflag to true to suppress this notice.
-DROP TYPE new_name CASCADE;
-NOTICE:  drop cascades to column happiness_level of table holidays
+ALTER TABLE holidays DROP CONSTRAINT holidays_pkey; -- fail
+ERROR:  constraint "holidays_pkey" for table "holidays" does not exist
 \d holidays
                Table "public.holidays"
   Column   |  Type   | Collation | Nullable | Default

src/postgres/src/test/regress/expected/yb_feature_alter_table.out

@@ -910,16 +910,18 @@ SELECT * FROM demo;
  1
 (2 rows)
 
--- Test that an attemp to drop primary key column with sequence generator
+-- Test dropping a primary key column with sequence generator
 -- does not delete the associated sequence.
 CREATE TABLE tbl_serial_primary_key (k serial PRIMARY KEY, v text);
 ALTER TABLE tbl_serial_primary_key DROP COLUMN k;
-ERROR:  cannot drop key column "k"
+NOTICE:  table rewrite may lead to inconsistencies
+DETAIL:  Concurrent DMLs may not be reflected in the new table.
+HINT:  See https://github.com/yugabyte/yugabyte-db/issues/19860. Set 'ysql_suppress_unsafe_alter_notice' yb-tserver gflag to true to suppress this notice.
 INSERT INTO tbl_serial_primary_key(v) VALUES ('ABC');
 SELECT * FROM tbl_serial_primary_key;
- k |  v
----+-----
- 1 | ABC
+  v  
+-----
+ ABC
 (1 row)
 
 DROP TABLE tbl_serial_primary_key;

src/postgres/src/test/regress/expected/yb_feature_hash_types.out

@@ -46,10 +46,6 @@ CREATE TABLE ft_h_tab_bool (feature_col BOOLEAN PRIMARY KEY);
 -- Enumerated Types
 CREATE TYPE feature_h_enum AS ENUM('one', 'two', 'three');
 CREATE TABLE ft_h_tab_enum (feature_col feature_h_enum PRIMARY KEY);
--- YB(fizaa): Work-around for the test cleanup code that drops the enum type.
--- Dropping the enum type will cause an error because a key column depends on it (GH #22902).
--- By dropping the table now, the clean up code will succeed.
-DROP TABLE ft_h_tab_enum;
 --
 -- Geometric Types
 CREATE TABLE ft_h_tab_point (feature_col POINT PRIMARY KEY);
@@ -119,10 +115,6 @@ ERROR:  PRIMARY KEY containing column of type 'user_defined_type' not yet suppor
 -- Domain Types
 CREATE DOMAIN feature_h_domain AS INTEGER CHECK (VALUE > 0);
 CREATE TABLE ft_h_tab_domain (feature_col feature_h_domain PRIMARY KEY);
--- YB(fizaa): Work-around for the test cleanup code that drops the domain.
--- Dropping the domain will cause an error because a key column depends on it (GH #22902).
--- By dropping the table now, the clean up code will succeed.
-DROP TABLE ft_h_tab_domain;
 --
 -- Object Identifier Types
 CREATE TABLE ft_h_tab_oid (feature_col OID PRIMARY KEY);

src/postgres/src/test/regress/expected/yb_feature_range_types.out

@@ -68,10 +68,6 @@ CREATE TABLE ft_r_tab_bool (feature_col BOOLEAN, PRIMARY KEY(feature_col ASC))
 -- Enumerated Types
 CREATE TYPE feature_r_enum AS ENUM('one', 'two', 'three');
 CREATE TABLE ft_r_tab_enum (feature_col feature_r_enum, PRIMARY KEY(feature_col ASC));
--- YB(fizaa): Work-around for the test cleanup code that drops the enum type.
--- Dropping the enum type will cause an error because a key column depends on it (GH #22902).
--- By dropping the table now, the clean up code will succeed.
-DROP TABLE ft_r_tab_enum;
 --
 -- Geometric Types
 CREATE TABLE ft_r_tab_point (feature_col POINT, PRIMARY KEY(feature_col ASC));

src/postgres/src/test/regress/sql/yb_alter_table_rewrite.sql

@@ -471,3 +471,25 @@ ALTER TABLE test_identity ADD COLUMN id2 int GENERATED ALWAYS AS IDENTITY;
 INSERT INTO test_identity VALUES (4, 4); -- should fail
 INSERT INTO test_identity OVERRIDING SYSTEM VALUE VALUES (4, 4);
 SELECT * FROM test_identity;
+
+--
+-- Test ALTER TABLE ... DROP COLUMN on a primary key column.
+--
+-- basic test
+CREATE TABLE test_drop_pk_column (id int PRIMARY KEY, v int);
+INSERT INTO test_drop_pk_column VALUES (1, 1), (2, 2), (3, 3);
+ALTER TABLE test_drop_pk_column DROP COLUMN id;
+\d test_drop_pk_column;
+INSERT INTO test_drop_pk_column VALUES (1);
+SELECT * FROM test_drop_pk_column ORDER BY v;
+-- test composite PK
+CREATE TABLE test_drop_pk_column_composite (id int, v int, t int, PRIMARY KEY (id HASH, v ASC));
+INSERT INTO test_drop_pk_column_composite VALUES (1, 1, 1), (2, 2, 2), (3, 3, 3);
+ALTER TABLE test_drop_pk_column_composite DROP COLUMN id;
+\d test_drop_pk_column_composite;
+INSERT INTO test_drop_pk_column_composite VALUES (1, 1);
+SELECT * FROM test_drop_pk_column_composite ORDER BY v;
+-- test partitioned table
+CREATE TABLE test_drop_pk_column_part (id int PRIMARY KEY, v int) PARTITION BY RANGE (id);
+CREATE TABLE test_drop_pk_column_part1 PARTITION OF test_drop_pk_column_part FOR VALUES FROM (1) TO (10);
+ALTER TABLE test_drop_pk_column_part DROP COLUMN id; -- should fail.

src/postgres/src/test/regress/sql/yb_alter_type.sql

@@ -56,9 +56,8 @@ INSERT INTO holidays(num_weeks,happiness_level) VALUES (8, 'ecstatic');
 SELECT * FROM holidays ORDER BY num_weeks;
 
 DROP TYPE new_name;  -- fail
-DROP TYPE new_name CASCADE; -- fail (happiness_level is a key column)
-ALTER TABLE holidays DROP CONSTRAINT holidays_pkey;
-DROP TYPE new_name CASCADE;
+DROP TYPE new_name CASCADE; -- requires table rewrite (happiness_level is a key column)
+ALTER TABLE holidays DROP CONSTRAINT holidays_pkey; -- fail
 
 \d holidays
 SELECT * FROM holidays ORDER BY num_weeks;

src/postgres/src/test/regress/sql/yb_feature_alter_table.sql

@@ -499,7 +499,7 @@ ALTER TABLE demo DROP CONSTRAINT demoi;
 INSERT INTO demo VALUES (1);
 SELECT * FROM demo;
 
--- Test that an attemp to drop primary key column with sequence generator
+-- Test dropping a primary key column with sequence generator
 -- does not delete the associated sequence.
 CREATE TABLE tbl_serial_primary_key (k serial PRIMARY KEY, v text);
 ALTER TABLE tbl_serial_primary_key DROP COLUMN k;

src/postgres/src/test/regress/sql/yb_feature_hash_types.sql

@@ -43,10 +43,6 @@ CREATE TABLE ft_h_tab_bool (feature_col BOOLEAN PRIMARY KEY);
 -- Enumerated Types
 CREATE TYPE feature_h_enum AS ENUM('one', 'two', 'three');
 CREATE TABLE ft_h_tab_enum (feature_col feature_h_enum PRIMARY KEY);
--- YB(fizaa): Work-around for the test cleanup code that drops the enum type.
--- Dropping the enum type will cause an error because a key column depends on it (GH #22902).
--- By dropping the table now, the clean up code will succeed.
-DROP TABLE ft_h_tab_enum;
 --
 -- Geometric Types
 CREATE TABLE ft_h_tab_point (feature_col POINT PRIMARY KEY);
@@ -95,10 +91,6 @@ CREATE TABLE ft_h_tab_range (feature_col feature_h_range PRIMARY KEY);
 -- Domain Types
 CREATE DOMAIN feature_h_domain AS INTEGER CHECK (VALUE > 0);
 CREATE TABLE ft_h_tab_domain (feature_col feature_h_domain PRIMARY KEY);
--- YB(fizaa): Work-around for the test cleanup code that drops the domain.
--- Dropping the domain will cause an error because a key column depends on it (GH #22902).
--- By dropping the table now, the clean up code will succeed.
-DROP TABLE ft_h_tab_domain;
 --
 -- Object Identifier Types
 CREATE TABLE ft_h_tab_oid (feature_col OID PRIMARY KEY);

src/postgres/src/test/regress/sql/yb_feature_range_types.sql

@@ -65,10 +65,6 @@ CREATE TABLE ft_r_tab_bool (feature_col BOOLEAN, PRIMARY KEY(feature_col ASC))
 -- Enumerated Types
 CREATE TYPE feature_r_enum AS ENUM('one', 'two', 'three');
 CREATE TABLE ft_r_tab_enum (feature_col feature_r_enum, PRIMARY KEY(feature_col ASC));
--- YB(fizaa): Work-around for the test cleanup code that drops the enum type.
--- Dropping the enum type will cause an error because a key column depends on it (GH #22902).
--- By dropping the table now, the clean up code will succeed.
-DROP TABLE ft_r_tab_enum;
 --
 -- Geometric Types
 CREATE TABLE ft_r_tab_point (feature_col POINT, PRIMARY KEY(feature_col ASC));