From 2d637eabeaa76b7da23b06c1c7f18a898f3a9468 Mon Sep 17 00:00:00 2001
From: Ajin Cherian <ajinc@fast.au.fujitsu.com>
Date: Mon, 12 Dec 2022 04:51:58 -0500
Subject: [PATCH v46 6/6] Support DDL replication of alter type having USING 
 expression on dropped column.

Support DDL replication of alter table having subcommands which include
drop column and alter type which uses dropped column expression like:
CREATE TABLE another (f1 int, f2 text, f3 text);
ALTER TABLE another
  DROP COLUMN f3,
  ALTER f1 TYPE text using f2 || ' and ' || f3 || ' more',
  ALTER f2 TYPE bigint using f1 * 10;

Currently event trigger is called after the execution of the command, i.e.
after the drop column, since the column is dropped we cannot evaluate
the expression using pg_get_expr for a dropped column. Fixed this issue
by adding event trigger before preparation of ALTER TYPE sub command.
---
 src/backend/commands/ddl_deparse.c   |  10 +--
 src/backend/commands/event_trigger.c | 127 +++++++++++++++++++++++++++++++++++
 src/backend/commands/tablecmds.c     |   8 ++-
 src/include/commands/event_trigger.h |   4 ++
 src/include/tcop/deparse_utility.h   |   1 +
 5 files changed, 140 insertions(+), 10 deletions(-)

diff --git a/src/backend/commands/ddl_deparse.c b/src/backend/commands/ddl_deparse.c
index 0c6c7d9..d7141e7 100755
--- a/src/backend/commands/ddl_deparse.c
+++ b/src/backend/commands/ddl_deparse.c
@@ -3858,16 +3858,8 @@ deparse_AlterRelation(CollectedCommand *cmd)
 						tmpobj2 = new_objtree("USING");
 						if (def->raw_default)
 						{
-							Datum		deparsed;
-							char	   *defexpr;
-
 							exprs = lappend(exprs, def->cooked_default);
-							defexpr = nodeToString(def->cooked_default);
-							deparsed = DirectFunctionCall2(pg_get_expr,
-														   CStringGetTextDatum(defexpr),
-														   RelationGetRelid(rel));
-							append_string_object(tmpobj2, "%{expression}s",
-												 TextDatumGetCString(deparsed));
+							append_string_object(tmpobj2, "%{expression}s", sub->usingexpr);
 						}
 						else
 							append_bool_object(tmpobj2, "present", false);
diff --git a/src/backend/commands/event_trigger.c b/src/backend/commands/event_trigger.c
index ab7cd59..98b3029 100644
--- a/src/backend/commands/event_trigger.c
+++ b/src/backend/commands/event_trigger.c
@@ -1783,6 +1783,133 @@ EventTriggerCollectAlterTableSubcmd(Node *subcmd, ObjectAddress address, bool re
 }
 
 /*
+ * EventTriggerAlterTypeStart
+ *		Save data about a single part of an ALTER TYPE.
+ *
+ * ALTER TABLE can have multiple subcommands which might include DROP COLUMN
+ * command and ALTER TYPE referring the drop column in USING expression.
+ * As the dropped column cannot be accessed after the execution of DROP COLUMN,
+ * a special trigger is required to handle this case before the drop column is
+ * executed.
+ */
+void
+EventTriggerAlterTypeStart(AlterTableCmd *subcmd, Relation rel)
+{
+	MemoryContext oldcxt;
+	CollectedATSubcmd *newsub;
+	ColumnDef  *def;
+	Relation	attrelation;
+	HeapTuple	heapTup;
+	Form_pg_attribute attTup;
+	AttrNumber	attnum;
+	ObjectAddress address;
+
+	/* ignore if event trigger context not set, or collection disabled */
+	if (!currentEventTriggerState ||
+		currentEventTriggerState->commandCollectionInhibited)
+		return;
+
+	Assert(IsA(subcmd, AlterTableCmd));
+	Assert(subcmd->subtype == AT_AlterColumnType);
+	Assert(currentEventTriggerState->currentCommand != NULL);
+	Assert(OidIsValid(currentEventTriggerState->currentCommand->d.alterTable.objectId));
+
+	def = (ColumnDef *) subcmd->def;
+	Assert(IsA(def, ColumnDef));
+
+	oldcxt = MemoryContextSwitchTo(currentEventTriggerState->cxt);
+
+	newsub = palloc(sizeof(CollectedATSubcmd));
+	newsub->parsetree = (Node *)copyObject(subcmd);
+
+	attrelation = table_open(AttributeRelationId, RowExclusiveLock);
+
+	/* Look up the target column */
+	heapTup = SearchSysCacheCopyAttName(RelationGetRelid(rel), subcmd->name);
+	if (!HeapTupleIsValid(heapTup)) /* shouldn't happen */
+		ereport(ERROR,
+				errcode(ERRCODE_UNDEFINED_COLUMN),
+				errmsg("column \"%s\" of relation \"%s\" does not exist",
+					   subcmd->name, RelationGetRelationName(rel)));
+	attTup = (Form_pg_attribute) GETSTRUCT(heapTup);
+	attnum = attTup->attnum;
+
+	ObjectAddressSubSet(address, RelationRelationId,
+						RelationGetRelid(rel), attnum);
+	heap_freetuple(heapTup);
+	table_close(attrelation, RowExclusiveLock);
+	newsub->address = address;
+
+	if (def->raw_default)
+	{
+		char	   *defexpr;
+
+		defexpr = nodeToString(def->cooked_default);
+		newsub->usingexpr = TextDatumGetCString(DirectFunctionCall2(pg_get_expr,
+																	CStringGetTextDatum(defexpr),
+																	RelationGetRelid(rel)));
+	}
+	else
+		newsub->usingexpr = NULL;
+
+	currentEventTriggerState->currentCommand->d.alterTable.subcmds =
+		lappend(currentEventTriggerState->currentCommand->d.alterTable.subcmds, newsub);
+
+	MemoryContextSwitchTo(oldcxt);
+}
+
+/*
+ * EventTriggerAlterTypeEnd
+ *		Finish up saving an ALTER TYPE command, and add it to command list.
+ */
+void
+EventTriggerAlterTypeEnd(Node *subcmd, ObjectAddress address, bool rewrite)
+{
+	MemoryContext oldcxt;
+	CollectedATSubcmd *newsub;
+	ListCell   *cell;
+	CollectedCommand *cmd;
+	AlterTableCmd *altsubcmd = (AlterTableCmd *)subcmd;
+
+	/* ignore if event trigger context not set, or collection disabled */
+	if (!currentEventTriggerState ||
+		currentEventTriggerState->commandCollectionInhibited)
+		return;
+
+	cmd = currentEventTriggerState->currentCommand;
+
+	Assert(IsA(subcmd, AlterTableCmd));
+	Assert(cmd != NULL);
+	Assert(OidIsValid(cmd->d.alterTable.objectId));
+
+	foreach(cell, cmd->d.alterTable.subcmds)
+	{
+		CollectedATSubcmd *sub = (CollectedATSubcmd *) lfirst(cell);
+		AlterTableCmd *collcmd = (AlterTableCmd *) sub->parsetree;
+
+		if (collcmd->subtype == altsubcmd->subtype &&
+			address.classId == sub->address.classId &&
+			address.objectId == sub->address.objectId &&
+			address.objectSubId == sub->address.objectSubId)
+		{
+			cmd->d.alterTable.rewrite |= rewrite;
+			return;
+		}
+	}
+
+	oldcxt = MemoryContextSwitchTo(currentEventTriggerState->cxt);
+
+	newsub = palloc(sizeof(CollectedATSubcmd));
+	newsub->address = address;
+	newsub->parsetree = copyObject(subcmd);
+
+	cmd->d.alterTable.rewrite |= rewrite;
+	cmd->d.alterTable.subcmds = lappend(cmd->d.alterTable.subcmds, newsub);
+
+	MemoryContextSwitchTo(oldcxt);
+}
+
+/*
  * EventTriggerAlterTableEnd
  *		Finish up saving an ALTER TABLE command, and add it to command list.
  *
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 191b1ac..a2fffe1 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -4653,6 +4653,9 @@ ATPrepCmd(List **wqueue, Relation rel, AlterTableCmd *cmd,
 			cmd = ATParseTransformCmd(wqueue, tab, rel, cmd, recurse, lockmode,
 									  AT_PASS_UNSET, context);
 			Assert(cmd != NULL);
+
+			EventTriggerAlterTypeStart(cmd, rel);
+
 			/* Performs own recursion */
 			ATPrepAlterColumnType(wqueue, tab, rel, recurse, recursing, cmd,
 								  lockmode, context);
@@ -4924,6 +4927,7 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 {
 	ObjectAddress address = InvalidObjectAddress;
 	Relation	rel = tab->rel;
+	bool		commandCollected = false;
 
 	switch (cmd->subtype)
 	{
@@ -5081,6 +5085,8 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 		case AT_AlterColumnType:	/* ALTER COLUMN TYPE */
 			/* parse transformation was done earlier */
 			address = ATExecAlterColumnType(tab, rel, cmd, lockmode);
+			EventTriggerAlterTypeEnd((Node *) cmd, address, tab->rewrite);
+			commandCollected = true;
 			break;
 		case AT_AlterColumnGenericOptions:	/* ALTER COLUMN OPTIONS */
 			address =
@@ -5253,7 +5259,7 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 	/*
 	 * Report the subcommand to interested event triggers.
 	 */
-	if (cmd)
+	if (cmd && !commandCollected)
 		EventTriggerCollectAlterTableSubcmd((Node *) cmd, address, tab->rewrite);
 
 	/*
diff --git a/src/include/commands/event_trigger.h b/src/include/commands/event_trigger.h
index c241db3..e7aa8f6 100644
--- a/src/include/commands/event_trigger.h
+++ b/src/include/commands/event_trigger.h
@@ -77,6 +77,10 @@ extern void EventTriggerAlterTableRelid(Oid objectId);
 extern void EventTriggerCollectAlterTableSubcmd(Node *subcmd,
 												ObjectAddress address,
 												bool rewrite);
+
+extern void EventTriggerAlterTypeStart(AlterTableCmd *subcmd, Relation rel);
+extern void EventTriggerAlterTypeEnd(Node *subcmd, ObjectAddress address,
+									 bool rewrite);
 extern void EventTriggerAlterTableEnd(void);
 
 extern void EventTriggerCollectGrant(InternalGrant *istmt);
diff --git a/src/include/tcop/deparse_utility.h b/src/include/tcop/deparse_utility.h
index 3d294a0..a1dd248 100644
--- a/src/include/tcop/deparse_utility.h
+++ b/src/include/tcop/deparse_utility.h
@@ -40,6 +40,7 @@ typedef struct CollectedATSubcmd
 {
 	ObjectAddress address;		/* affected column, constraint, index, ... */
 	Node	   *parsetree;
+	char	   *usingexpr;
 } CollectedATSubcmd;
 
 typedef struct CollectedCommand
-- 
1.8.3.1

