serial not accepted as datatype in ALTER TABLE ... ALTER COLUMN

From 4485a9af33c9c7d493e23c2804bde4c15df0b79a Mon Sep 17 00:00:00 2001
From: "yizhi.fzh" <yizhi.fzh@alibaba-inc.com>
Date: Sun, 18 Feb 2024 16:14:29 +0800
Subject: [PATCH v0 1/1] POC: Support ALTER TABLE tableName ALERT COLUMN col
 type serial.

---
 src/backend/commands/tablecmds.c          | 60 +++++++++++++++++--
 src/backend/parser/gram.y                 |  1 +
 src/backend/parser/parse_utilcmd.c        | 41 ++++++++++++-
 src/include/nodes/parsenodes.h            |  1 +
 src/include/parser/parse_utilcmd.h        |  4 +-
 src/test/regress/expected/alter_table.out | 70 +++++++++++++++++++++--
 src/test/regress/sql/alter_table.sql      | 20 +++++--
 7 files changed, 180 insertions(+), 17 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 46ece07338..7e13b04efa 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -231,6 +231,8 @@ typedef struct NewColumnValue
 	Expr	   *expr;			/* expression to compute */
 	ExprState  *exprstate;		/* execution state */
 	bool		is_generated;	/* is it a GENERATED expression? */
+	bool		new_on_null;	/* set the new value only when the old value
+								 * is NULL. */
 } NewColumnValue;
 
 /*
@@ -5594,7 +5596,8 @@ ATParseTransformCmd(List **wqueue, AlteredTableInfo *tab, Relation rel,
 									 atstmt,
 									 context->queryString,
 									 &beforeStmts,
-									 &afterStmts);
+									 &afterStmts,
+									 context);
 
 	/* Execute any statements that should happen before these subcommand(s) */
 	foreach(lc, beforeStmts)
@@ -6230,10 +6233,11 @@ ATRewriteTable(AlteredTableInfo *tab, Oid OIDNewHeap, LOCKMODE lockmode)
 					if (ex->is_generated)
 						continue;
 
-					newslot->tts_values[ex->attnum - 1]
-						= ExecEvalExpr(ex->exprstate,
-									   econtext,
-									   &newslot->tts_isnull[ex->attnum - 1]);
+					if (!ex->new_on_null || oldslot->tts_isnull[ex->attnum - 1])
+						newslot->tts_values[ex->attnum - 1]
+							= ExecEvalExpr(ex->exprstate,
+										   econtext,
+										   &newslot->tts_isnull[ex->attnum - 1]);
 				}
 
 				ExecStoreVirtualTuple(newslot);
@@ -13284,6 +13288,7 @@ ATPrepAlterColumnType(List **wqueue,
 		newval->attnum = attnum;
 		newval->expr = (Expr *) transform;
 		newval->is_generated = false;
+		newval->new_on_null = def->from_serial;
 
 		tab->newvals = lappend(tab->newvals, newval);
 		if (ATColumnChangeRequiresRewrite(transform, attnum))
@@ -13495,6 +13500,48 @@ ATExecAlterColumnType(AlteredTableInfo *tab, Relation rel,
 	HeapTuple	depTup;
 	ObjectAddress address;
 
+	if (def->from_serial)
+	{
+		/*
+		 * Store the DEFAULT for the from_serial
+		 */
+		/* XXX: copy from ATExecAddColumn */
+		RawColumnDefault *rawEnt;
+
+		heapTup = SearchSysCacheCopyAttName(RelationGetRelid(rel), colName);
+		attTup = (Form_pg_attribute) GETSTRUCT(heapTup);
+
+		rawEnt = (RawColumnDefault *) palloc(sizeof(RawColumnDefault));
+		rawEnt->attnum = attTup->attnum;
+		rawEnt->raw_default = copyObject(def->raw_default);
+
+		/*
+		 * Attempt to skip a complete table rewrite by storing the specified
+		 * DEFAULT value outside of the heap.  This may be disabled inside
+		 * AddRelationNewConstraints if the optimization cannot be applied.
+		 */
+		rawEnt->missingMode = (!def->generated);
+
+		rawEnt->generated = def->generated;
+
+		/*
+		 * This function is intended for CREATE TABLE, so it processes a
+		 * _list_ of defaults, but we just do one.
+		 */
+		AddRelationNewConstraints(rel, list_make1(rawEnt), NIL,
+								  false, true, false, NULL);
+
+		/* Make the additional catalog changes visible */
+		CommandCounterIncrement();
+
+		/*
+		 * Did the request for a missing value work? If not we'll have to do a
+		 * rewrite
+		 */
+		if (!rawEnt->missingMode)
+			tab->rewrite |= AT_REWRITE_DEFAULT_VAL;
+	}
+
 	/*
 	 * Clear all the missing values if we're rewriting the table, since this
 	 * renders them pointless.
@@ -14368,7 +14415,8 @@ ATPostAlterTypeParse(Oid oldId, Oid oldRelId, Oid refRelId, char *cmd,
 													(AlterTableStmt *) stmt,
 													cmd,
 													&beforeStmts,
-													&afterStmts);
+													&afterStmts,
+													NULL);
 			querytree_list = list_concat(querytree_list, beforeStmts);
 			querytree_list = lappend(querytree_list, stmt);
 			querytree_list = list_concat(querytree_list, afterStmts);
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 60b31d9f85..0aea1673c1 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -2599,6 +2599,7 @@ alter_table_cmd:
 					def->collClause = (CollateClause *) $7;
 					def->raw_default = $8;
 					def->location = @3;
+					def->colname = $3;
 					$$ = (Node *) n;
 				}
 			/* ALTER FOREIGN TABLE <name> ALTER [COLUMN] <colname> OPTIONS */
diff --git a/src/backend/parser/parse_utilcmd.c b/src/backend/parser/parse_utilcmd.c
index 8dcf794ca2..622e7f629a 100644
--- a/src/backend/parser/parse_utilcmd.c
+++ b/src/backend/parser/parse_utilcmd.c
@@ -28,6 +28,7 @@
 #include "access/reloptions.h"
 #include "access/table.h"
 #include "access/toast_compression.h"
+#include "access/xact.h"
 #include "catalog/dependency.h"
 #include "catalog/heap.h"
 #include "catalog/index.h"
@@ -59,6 +60,7 @@
 #include "parser/parse_utilcmd.h"
 #include "parser/parser.h"
 #include "rewrite/rewriteManip.h"
+#include "tcop/utility.h"
 #include "utils/acl.h"
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
@@ -385,6 +387,8 @@ generateSerialExtraStmts(CreateStmtContext *cxt, ColumnDef *column,
 	/* Make a copy of this as we may end up modifying it in the code below */
 	seqoptions = list_copy(seqoptions);
 
+	column->from_serial = true;
+
 	/*
 	 * Determine namespace and name to use for the sequence.
 	 *
@@ -3423,11 +3427,15 @@ transformRuleStmt(RuleStmt *stmt, const char *queryString,
  * To avoid race conditions, it's important that this function rely only on
  * the passed-in relid (and not on stmt->relation) to determine the target
  * relation.
+ *
+ * Note: context is used for creating the beforeStmt earlier, maybe NULL
+ * sometimes.
  */
 AlterTableStmt *
 transformAlterTableStmt(Oid relid, AlterTableStmt *stmt,
 						const char *queryString,
-						List **beforeStmts, List **afterStmts)
+						List **beforeStmts, List **afterStmts,
+						struct AlterTableUtilityContext *context)
 {
 	Relation	rel;
 	TupleDesc	tupdesc;
@@ -3440,6 +3448,7 @@ transformAlterTableStmt(Oid relid, AlterTableStmt *stmt,
 	bool		skipValidation = true;
 	AlterTableCmd *newcmd;
 	ParseNamespaceItem *nsitem;
+	bool		beforeStmtExecuted = false;
 
 	/* Caller is responsible for locking the relation */
 	rel = relation_open(relid, NoLock);
@@ -3540,6 +3549,30 @@ transformAlterTableStmt(Oid relid, AlterTableStmt *stmt,
 					ColumnDef  *def = castNode(ColumnDef, cmd->def);
 					AttrNumber	attnum;
 
+					/*
+					 * transform the serial data type,  including 1. generate
+					 * the seq cmd as before stmt. 2. change the serial type
+					 * into intX 3. add not null constraint.
+					 */
+					transformColumnDefinition(&cxt, def);
+					if (def->from_serial)
+					{
+						/*
+						 * XXX: we have to execute this sooner for the later
+						 * transformExpr(.., raw_default) works.
+						 */
+						ListCell   *lc;
+
+						foreach(lc, cxt.blist)
+						{
+							Node	   *stmt = (Node *) lfirst(lc);
+
+							ProcessUtilityForAlterTable(stmt, context);
+							CommandCounterIncrement();
+						}
+						beforeStmtExecuted = true;
+					}
+
 					/*
 					 * For ALTER COLUMN TYPE, transform the USING clause if
 					 * one was specified.
@@ -3776,7 +3809,11 @@ transformAlterTableStmt(Oid relid, AlterTableStmt *stmt,
 	 */
 	stmt->cmds = newcmds;
 
-	*beforeStmts = cxt.blist;
+	if (!beforeStmtExecuted)
+		*beforeStmts = cxt.blist;
+	else
+		*beforeStmts = NIL;
+
 	*afterStmts = list_concat(cxt.alist, save_alist);
 
 	return stmt;
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index d5b08ded44..7ff2abfc53 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -729,6 +729,7 @@ typedef struct ColumnDef
 	List	   *constraints;	/* other constraints on column */
 	List	   *fdwoptions;		/* per-column FDW options */
 	int			location;		/* parse location, or -1 if none/unknown */
+	bool		from_serial;	/* transform from type serial. */
 } ColumnDef;
 
 /*
diff --git a/src/include/parser/parse_utilcmd.h b/src/include/parser/parse_utilcmd.h
index 1406589477..7b70c785fb 100644
--- a/src/include/parser/parse_utilcmd.h
+++ b/src/include/parser/parse_utilcmd.h
@@ -18,12 +18,14 @@
 
 struct AttrMap;					/* avoid including attmap.h here */
 
+struct AlterTableUtilityContext;
 
 extern List *transformCreateStmt(CreateStmt *stmt, const char *queryString);
 extern AlterTableStmt *transformAlterTableStmt(Oid relid, AlterTableStmt *stmt,
 											   const char *queryString,
 											   List **beforeStmts,
-											   List **afterStmts);
+											   List **afterStmts,
+											   struct AlterTableUtilityContext *context);
 extern IndexStmt *transformIndexStmt(Oid relid, IndexStmt *stmt,
 									 const char *queryString);
 extern CreateStatsStmt *transformStatsStmt(Oid relid, CreateStatsStmt *stmt,
diff --git a/src/test/regress/expected/alter_table.out b/src/test/regress/expected/alter_table.out
index 7666c76238..0daf296e13 100644
--- a/src/test/regress/expected/alter_table.out
+++ b/src/test/regress/expected/alter_table.out
@@ -930,10 +930,25 @@ ERROR:  duplicate key value violates unique constraint "atacc_test1"
 DETAIL:  Key (test)=(2) already exists.
 -- should succeed
 insert into atacc1 (test) values (4);
+-- insert a NULL.
+insert into atacc1 select;
 -- try to create duplicates via alter table using - should fail
 alter table atacc1 alter column test type integer using 0;
 ERROR:  could not create unique index "atacc_test1"
 DETAIL:  Key (test)=(0) is duplicated.
+alter table atacc1 alter column test type serial;
+select * from atacc1 order by 1;
+ test 
+------
+    1
+    2
+    4
+(3 rows)
+
+-- fail since the nextval of seq is 2 and 2 exist already.
+insert into atacc1 select;
+ERROR:  duplicate key value violates unique constraint "atacc_test1"
+DETAIL:  Key (test)=(2) already exists.
 drop table atacc1;
 -- let's do one where the unique constraint fails when added
 create table atacc1 ( test int );
@@ -2086,11 +2101,13 @@ alter table at_tab1 alter column b type varchar; -- fails
 ERROR:  cannot alter table "at_tab1" because column "at_tab2.y" uses its row type
 drop table at_tab1, at_tab2;
 -- Alter column type that's part of a partitioned index
-create table at_partitioned (a int, b text) partition by range (a);
+create table at_partitioned (a int, b text, c int) partition by range (a);
 create table at_part_1 partition of at_partitioned for values from (0) to (1000);
-insert into at_partitioned values (512, '0.123');
-create table at_part_2 (b text, a int);
-insert into at_part_2 values ('1.234', 1024);
+insert into at_partitioned values (512, '0.123', 1);
+insert into at_partitioned values (512, '0.123', NULL);
+create table at_part_2 (b text, c int, a int);
+insert into at_part_2 values ('1.234', 2, 1024);
+insert into at_part_2 values ('1.234', NULL, 1024);
 create index on at_partitioned (b);
 create index on at_partitioned (a);
 \d at_part_1
@@ -2099,6 +2116,7 @@ create index on at_partitioned (a);
 --------+---------+-----------+----------+---------
  a      | integer |           |          | 
  b      | text    |           |          | 
+ c      | integer |           |          | 
 Partition of: at_partitioned FOR VALUES FROM (0) TO (1000)
 Indexes:
     "at_part_1_a_idx" btree (a)
@@ -2109,6 +2127,7 @@ Indexes:
  Column |  Type   | Collation | Nullable | Default 
 --------+---------+-----------+----------+---------
  b      | text    |           |          | 
+ c      | integer |           |          | 
  a      | integer |           |          | 
 
 alter table at_partitioned attach partition at_part_2 for values from (1000) to (2000);
@@ -2117,6 +2136,7 @@ alter table at_partitioned attach partition at_part_2 for values from (1000) to
  Column |  Type   | Collation | Nullable | Default 
 --------+---------+-----------+----------+---------
  b      | text    |           |          | 
+ c      | integer |           |          | 
  a      | integer |           |          | 
 Partition of: at_partitioned FOR VALUES FROM (1000) TO (2000)
 Indexes:
@@ -2130,6 +2150,7 @@ alter table at_partitioned alter column b type numeric using b::numeric;
 --------+---------+-----------+----------+---------
  a      | integer |           |          | 
  b      | numeric |           |          | 
+ c      | integer |           |          | 
 Partition of: at_partitioned FOR VALUES FROM (0) TO (1000)
 Indexes:
     "at_part_1_a_idx" btree (a)
@@ -2140,12 +2161,53 @@ Indexes:
  Column |  Type   | Collation | Nullable | Default 
 --------+---------+-----------+----------+---------
  b      | numeric |           |          | 
+ c      | integer |           |          | 
  a      | integer |           |          | 
 Partition of: at_partitioned FOR VALUES FROM (1000) TO (2000)
 Indexes:
     "at_part_2_a_idx" btree (a)
     "at_part_2_b_idx" btree (b)
 
+alter table at_partitioned alter column c type serial;
+select * from at_partitioned order by 1;
+  a   |   b   | c 
+------+-------+---
+  512 | 0.123 | 1
+  512 | 0.123 | 1
+ 1024 | 1.234 | 2
+ 1024 | 1.234 | 2
+(4 rows)
+
+\d+ at_part_1
+                                                  Table "public.at_part_1"
+ Column |  Type   | Collation | Nullable |                  Default                  | Storage | Stats target | Description 
+--------+---------+-----------+----------+-------------------------------------------+---------+--------------+-------------
+ a      | integer |           |          |                                           | plain   |              | 
+ b      | numeric |           |          |                                           | main    |              | 
+ c      | integer |           | not null | nextval('at_partitioned_c_seq'::regclass) | plain   |              | 
+Partition of: at_partitioned FOR VALUES FROM (0) TO (1000)
+Partition constraint: ((a IS NOT NULL) AND (a >= 0) AND (a < 1000))
+Indexes:
+    "at_part_1_a_idx" btree (a)
+    "at_part_1_b_idx" btree (b)
+Not-null constraints:
+    "at_partitioned_c_not_null" NOT NULL "c" (inherited)
+
+\d+ at_part_2
+                                                  Table "public.at_part_2"
+ Column |  Type   | Collation | Nullable |                  Default                  | Storage | Stats target | Description 
+--------+---------+-----------+----------+-------------------------------------------+---------+--------------+-------------
+ b      | numeric |           |          |                                           | main    |              | 
+ c      | integer |           | not null | nextval('at_partitioned_c_seq'::regclass) | plain   |              | 
+ a      | integer |           |          |                                           | plain   |              | 
+Partition of: at_partitioned FOR VALUES FROM (1000) TO (2000)
+Partition constraint: ((a IS NOT NULL) AND (a >= 1000) AND (a < 2000))
+Indexes:
+    "at_part_2_a_idx" btree (a)
+    "at_part_2_b_idx" btree (b)
+Not-null constraints:
+    "at_partitioned_c_not_null" NOT NULL "c" (inherited)
+
 drop table at_partitioned;
 -- Alter column type when no table rewrite is required
 -- Also check that comments are preserved
diff --git a/src/test/regress/sql/alter_table.sql b/src/test/regress/sql/alter_table.sql
index 9df5a63bdf..876e6b1fa4 100644
--- a/src/test/regress/sql/alter_table.sql
+++ b/src/test/regress/sql/alter_table.sql
@@ -680,8 +680,14 @@ insert into atacc1 (test) values (2);
 insert into atacc1 (test) values (2);
 -- should succeed
 insert into atacc1 (test) values (4);
+-- insert a NULL.
+insert into atacc1 select;
 -- try to create duplicates via alter table using - should fail
 alter table atacc1 alter column test type integer using 0;
+alter table atacc1 alter column test type serial;
+select * from atacc1 order by 1;
+-- fail since the nextval of seq is 2 and 2 exist already.
+insert into atacc1 select;
 drop table atacc1;
 
 -- let's do one where the unique constraint fails when added
@@ -1431,11 +1437,13 @@ alter table at_tab1 alter column b type varchar; -- fails
 drop table at_tab1, at_tab2;
 
 -- Alter column type that's part of a partitioned index
-create table at_partitioned (a int, b text) partition by range (a);
+create table at_partitioned (a int, b text, c int) partition by range (a);
 create table at_part_1 partition of at_partitioned for values from (0) to (1000);
-insert into at_partitioned values (512, '0.123');
-create table at_part_2 (b text, a int);
-insert into at_part_2 values ('1.234', 1024);
+insert into at_partitioned values (512, '0.123', 1);
+insert into at_partitioned values (512, '0.123', NULL);
+create table at_part_2 (b text, c int, a int);
+insert into at_part_2 values ('1.234', 2, 1024);
+insert into at_part_2 values ('1.234', NULL, 1024);
 create index on at_partitioned (b);
 create index on at_partitioned (a);
 \d at_part_1
@@ -1445,6 +1453,10 @@ alter table at_partitioned attach partition at_part_2 for values from (1000) to
 alter table at_partitioned alter column b type numeric using b::numeric;
 \d at_part_1
 \d at_part_2
+alter table at_partitioned alter column c type serial;
+select * from at_partitioned order by 1;
+\d+ at_part_1
+\d+ at_part_2
 drop table at_partitioned;
 
 -- Alter column type when no table rewrite is required
-- 
2.34.1