0003-Enhance-the-event-trigger-to-support-DDL--2023_06_22.patch

From cead449ae708edb3648aa5751520b1584fead902 Mon Sep 17 00:00:00 2001
From: Shveta Malik <shveta.malik@gmail.com>
Date: Wed, 21 Jun 2023 09:18:49 +0530
Subject: [PATCH 3/5] Enhance the event trigger to support DDL deparsing

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 added to handle this case before the drop column is executed.
---
 src/backend/commands/ddldeparse.c    |  11 +-
 src/backend/commands/event_trigger.c | 188 +++++++++++++++++++++------
 src/backend/commands/tablecmds.c     |  10 +-
 src/include/commands/event_trigger.h |  46 ++++++-
 src/include/tcop/deparse_utility.h   |   2 +
 5 files changed, 204 insertions(+), 53 deletions(-)

diff --git a/src/backend/commands/ddldeparse.c b/src/backend/commands/ddldeparse.c
index faaf0de3f9..7f8e416d6a 100644
--- a/src/backend/commands/ddldeparse.c
+++ b/src/backend/commands/ddldeparse.c
@@ -2552,22 +2552,13 @@ deparse_AlterTableStmt(CollectedCommand *cmd)
 					 */
 					if (def->raw_default)
 					{
-						Datum		deparsed;
-						char	   *defexpr;
-						List	   *exprs = NIL;
-
-						exprs = lappend(exprs, def->cooked_default);
-						defexpr = nodeToString(def->cooked_default);
-						deparsed = DirectFunctionCall2(pg_get_expr,
-													   CStringGetTextDatum(defexpr),
-													   RelationGetRelid(rel));
 						appendStringInfoString(&fmtSub, " %{using}s");
 						insert_jsonb_key(state, "using");
 						pushJsonbValue(&state, WJB_BEGIN_OBJECT, NULL);
 						new_jsonb_VA(state, 2,
 									 "fmt", jbvString, "USING %{expression}s",
 									 "expression", jbvString,
-									 TextDatumGetCString(deparsed));
+									 sub->usingexpr);
 						pushJsonbValue(&state, WJB_END_OBJECT, NULL);
 					}
 
diff --git a/src/backend/commands/event_trigger.c b/src/backend/commands/event_trigger.c
index 4d48e490ed..8c2a494dcb 100644
--- a/src/backend/commands/event_trigger.c
+++ b/src/backend/commands/event_trigger.c
@@ -36,6 +36,7 @@
 #include "lib/ilist.h"
 #include "miscadmin.h"
 #include "parser/parse_func.h"
+#include "parser/parser.h"
 #include "pgstat.h"
 #include "tcop/ddldeparse.h"
 #include "tcop/deparse_utility.h"
@@ -49,45 +50,7 @@
 #include "utils/rel.h"
 #include "utils/syscache.h"
 
-typedef struct EventTriggerQueryState
-{
-	/* memory context for this state's objects */
-	MemoryContext cxt;
-
-	/* sql_drop */
-	slist_head	SQLDropList;
-	bool		in_sql_drop;
-
-	/* table_rewrite */
-	Oid			table_rewrite_oid;	/* InvalidOid, or set for table_rewrite
-									 * event */
-	int			table_rewrite_reason;	/* AT_REWRITE reason */
-
-	/* Support for command collection */
-	bool		commandCollectionInhibited;
-	CollectedCommand *currentCommand;
-	List	   *commandList;	/* list of CollectedCommand; see
-								 * deparse_utility.h */
-	struct EventTriggerQueryState *previous;
-} EventTriggerQueryState;
-
-static EventTriggerQueryState *currentEventTriggerState = NULL;
-
-/* Support for dropped objects */
-typedef struct SQLDropObject
-{
-	ObjectAddress address;
-	const char *schemaname;
-	const char *objname;
-	const char *objidentity;
-	const char *objecttype;
-	List	   *addrnames;
-	List	   *addrargs;
-	bool		original;
-	bool		normal;
-	bool		istemp;
-	slist_node	next;
-} SQLDropObject;
+EventTriggerQueryState *currentEventTriggerState = NULL;
 
 static void AlterEventTriggerOwner_internal(Relation rel,
 											HeapTuple tup,
@@ -1538,6 +1501,7 @@ EventTriggerAlterTableStart(Node *parsetree)
 
 	command->d.alterTable.classId = RelationRelationId;
 	command->d.alterTable.objectId = InvalidOid;
+	command->d.alterTable.rewrite = false;
 	command->d.alterTable.subcmds = NIL;
 	command->parsetree = copyObject(parsetree);
 
@@ -1571,7 +1535,7 @@ EventTriggerAlterTableRelid(Oid objectId)
  * internally, so that's all that this code needs to handle at the moment.
  */
 void
-EventTriggerCollectAlterTableSubcmd(Node *subcmd, ObjectAddress address)
+EventTriggerCollectAlterTableSubcmd(Node *subcmd, ObjectAddress address, bool rewrite)
 {
 	MemoryContext oldcxt;
 	CollectedATSubcmd *newsub;
@@ -1591,12 +1555,156 @@ EventTriggerCollectAlterTableSubcmd(Node *subcmd, ObjectAddress address)
 	newsub->address = address;
 	newsub->parsetree = copyObject(subcmd);
 
+	currentEventTriggerState->currentCommand->d.alterTable.rewrite |= rewrite;
 	currentEventTriggerState->currentCommand->d.alterTable.subcmds =
 		lappend(currentEventTriggerState->currentCommand->d.alterTable.subcmds, newsub);
 
 	MemoryContextSwitchTo(oldcxt);
 }
 
+/*
+ * 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)
+	{
+		OverrideSearchPath *overridePath;
+		char	   *defexpr;
+
+		/*
+		 * We want all object names to be qualified when deparsing the
+		 * expression, so that results are "portable" to environments with
+		 * different search_path settings. Rather than inject what would be
+		 * repetitive calls to override search path all over the place, we do
+		 * it centrally here.
+		 */
+		overridePath = GetOverrideSearchPath(CurrentMemoryContext);
+		overridePath->schemas = NIL;
+		overridePath->addCatalog = false;
+		overridePath->addTemp = true;
+		PushOverrideSearchPath(overridePath);
+
+		defexpr = nodeToString(def->cooked_default);
+		newsub->usingexpr = TextDatumGetCString(DirectFunctionCall2(pg_get_expr,
+																	CStringGetTextDatum(defexpr),
+																	RelationGetRelid(rel)));
+
+		PopOverrideSearchPath();
+	}
+	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 c15e943f1d..d0e649baaf 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -4742,6 +4742,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);
@@ -5013,6 +5016,7 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 {
 	ObjectAddress address = InvalidObjectAddress;
 	Relation	rel = tab->rel;
+	bool		commandCollected = false;
 
 	switch (cmd->subtype)
 	{
@@ -5136,6 +5140,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 =
@@ -5308,8 +5314,8 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 	/*
 	 * Report the subcommand to interested event triggers.
 	 */
-	if (cmd)
-		EventTriggerCollectAlterTableSubcmd((Node *) cmd, address);
+	if (cmd && !commandCollected)
+		EventTriggerCollectAlterTableSubcmd((Node *) cmd, address, tab->rewrite);
 
 	/*
 	 * Bump the command counter to ensure the next subcommand in the sequence
diff --git a/src/include/commands/event_trigger.h b/src/include/commands/event_trigger.h
index 5ed6ece555..28b3486b9e 100644
--- a/src/include/commands/event_trigger.h
+++ b/src/include/commands/event_trigger.h
@@ -16,6 +16,7 @@
 #include "catalog/dependency.h"
 #include "catalog/objectaddress.h"
 #include "catalog/pg_event_trigger.h"
+#include "lib/ilist.h"
 #include "nodes/parsenodes.h"
 #include "tcop/cmdtag.h"
 #include "tcop/deparse_utility.h"
@@ -29,6 +30,44 @@ typedef struct EventTriggerData
 	CommandTag	tag;
 } EventTriggerData;
 
+typedef struct EventTriggerQueryState
+{
+	/* memory context for this state's objects */
+	MemoryContext cxt;
+
+	/* sql_drop */
+	slist_head	SQLDropList;
+	bool		in_sql_drop;
+
+	/* table_rewrite */
+	Oid			table_rewrite_oid;	/* InvalidOid, or set for table_rewrite
+									 * event */
+	int			table_rewrite_reason;	/* AT_REWRITE reason */
+
+	/* Support for command collection */
+	bool		commandCollectionInhibited;
+	CollectedCommand *currentCommand;
+	List	   *commandList;	/* list of CollectedCommand; see
+								 * deparse_utility.h */
+	struct EventTriggerQueryState *previous;
+} EventTriggerQueryState;
+
+/* Support for dropped objects */
+typedef struct SQLDropObject
+{
+	ObjectAddress address;
+	const char *schemaname;
+	const char *objname;
+	const char *objidentity;
+	const char *objecttype;
+	List	   *addrnames;
+	List	   *addrargs;
+	bool		original;
+	bool		normal;
+	bool		istemp;
+	slist_node	next;
+} SQLDropObject;
+
 #define AT_REWRITE_ALTER_PERSISTENCE	0x01
 #define AT_REWRITE_DEFAULT_VAL			0x02
 #define AT_REWRITE_COLUMN_REWRITE		0x04
@@ -71,7 +110,12 @@ extern void EventTriggerCollectSimpleCommand(ObjectAddress address,
 extern void EventTriggerAlterTableStart(Node *parsetree);
 extern void EventTriggerAlterTableRelid(Oid objectId);
 extern void EventTriggerCollectAlterTableSubcmd(Node *subcmd,
-												ObjectAddress address);
+												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 b585810b9a..1831ec9aae 100644
--- a/src/include/tcop/deparse_utility.h
+++ b/src/include/tcop/deparse_utility.h
@@ -39,6 +39,7 @@ typedef struct CollectedATSubcmd
 {
 	ObjectAddress address;		/* affected column, constraint, index, ... */
 	Node	   *parsetree;
+	char	   *usingexpr;
 } CollectedATSubcmd;
 
 typedef struct CollectedCommand
@@ -62,6 +63,7 @@ typedef struct CollectedCommand
 		{
 			Oid			objectId;
 			Oid			classId;
+			bool		rewrite;
 			List	   *subcmds;
 		}			alterTable;
 
-- 
2.34.1