From c1ee511a8ca9d17bc1d876bc55186b9c74dbdc8d Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Wed, 2 Feb 2022 16:58:51 -0300
Subject: [PATCH v12 1/2] Add API of sorts for transaction table handling in
 trigger.c

---
 src/backend/commands/trigger.c | 164 ++++++++++++++++++++++-----------
 1 file changed, 112 insertions(+), 52 deletions(-)

diff --git a/src/backend/commands/trigger.c b/src/backend/commands/trigger.c
index 1a9c1ac290..68ae1e93ce 100644
--- a/src/backend/commands/trigger.c
+++ b/src/backend/commands/trigger.c
@@ -94,6 +94,16 @@ static HeapTuple ExecCallTriggerFunc(TriggerData *trigdata,
 									 FmgrInfo *finfo,
 									 Instrumentation *instr,
 									 MemoryContext per_tuple_context);
+static Tuplestorestate *AfterTriggerGetTransitionTable(int event,
+													   TupleTableSlot *oldslot,
+													   TupleTableSlot *newslot,
+													   TransitionCaptureState *transition_capture);
+static void TransitionTableAddTuple(EState *estate,
+									TransitionCaptureState *transition_capture,
+									ResultRelInfo *relinfo,
+									TupleTableSlot *slot,
+									TupleTableSlot *original_insert_tuple,
+									Tuplestorestate *tuplestore);
 static void AfterTriggerSaveEvent(EState *estate, ResultRelInfo *relinfo,
 								  int event, bool row_trigger,
 								  TupleTableSlot *oldtup, TupleTableSlot *newtup,
@@ -5650,6 +5660,85 @@ AfterTriggerPendingOnRel(Oid relid)
 	return false;
 }
 
+/*
+ * Get the transition table for the given event and depending on whether we are
+ * processing the old or the new tuple.
+ */
+static Tuplestorestate *
+AfterTriggerGetTransitionTable(int event,
+							   TupleTableSlot *oldslot,
+							   TupleTableSlot *newslot,
+							   TransitionCaptureState *transition_capture)
+{
+	Tuplestorestate *tuplestore;
+	bool		delete_old_table = transition_capture->tcs_delete_old_table;
+	bool		update_old_table = transition_capture->tcs_update_old_table;
+	bool		update_new_table = transition_capture->tcs_update_new_table;
+	bool		insert_new_table = transition_capture->tcs_insert_new_table;
+
+	/*
+	 * For INSERT events NEW should be non-NULL, for DELETE events OLD
+	 * should be non-NULL, whereas for UPDATE events normally both OLD and
+	 * NEW are non-NULL.  But for UPDATE events fired for capturing
+	 * transition tuples during UPDATE partition-key row movement, OLD is
+	 * NULL when the event is for a row being inserted, whereas NEW is
+	 * NULL when the event is for a row being deleted.
+	 */
+	Assert(!(event == TRIGGER_EVENT_DELETE && delete_old_table &&
+			 TupIsNull(oldslot)));
+	Assert(!(event == TRIGGER_EVENT_INSERT && insert_new_table &&
+			 TupIsNull(newslot)));
+
+	Assert(TupIsNull(oldslot) ^ TupIsNull(newslot));
+
+	if (!TupIsNull(oldslot) &&
+		((event == TRIGGER_EVENT_DELETE && delete_old_table) ||
+		 (event == TRIGGER_EVENT_UPDATE && update_old_table)))
+		tuplestore = transition_capture->tcs_private->old_tuplestore;
+	else if (!TupIsNull(newslot) &&
+			 ((event == TRIGGER_EVENT_INSERT && insert_new_table) ||
+			  (event == TRIGGER_EVENT_UPDATE && update_new_table)))
+		tuplestore = transition_capture->tcs_private->new_tuplestore;
+	else
+		tuplestore = NULL;
+
+	return tuplestore;
+}
+
+/*
+ * Add the given heap tuple to the given tuplestore, applying the conversion
+ * map if necessary.
+ */
+static void
+TransitionTableAddTuple(EState *estate,
+						TransitionCaptureState *transition_capture,
+						ResultRelInfo *relinfo,
+						TupleTableSlot *slot,
+						TupleTableSlot *original_insert_tuple,
+						Tuplestorestate *tuplestore)
+{
+	TupleConversionMap *map;
+
+	/*
+	 * Nothing needs to be done if we don't have a tuplestore.
+	 */
+	if (tuplestore == NULL)
+		return;
+
+	if (original_insert_tuple)
+		tuplestore_puttupleslot(tuplestore, original_insert_tuple);
+	else if ((map = ExecGetChildToRootMap(relinfo)) != NULL)
+	{
+		AfterTriggersTableData *table = transition_capture->tcs_private;
+		TupleTableSlot *storeslot;
+
+		storeslot = GetAfterTriggersStoreSlot(table, map->outdesc);
+		execute_attr_map_slot(map->attrMap, slot, storeslot);
+		tuplestore_puttupleslot(tuplestore, storeslot);
+	}
+	else
+		tuplestore_puttupleslot(tuplestore, slot);
+}
 
 /* ----------
  * AfterTriggerSaveEvent()
@@ -5709,68 +5798,39 @@ AfterTriggerSaveEvent(EState *estate, ResultRelInfo *relinfo,
 	 */
 	if (row_trigger && transition_capture != NULL)
 	{
-		TupleTableSlot *original_insert_tuple = transition_capture->tcs_original_insert_tuple;
-		TupleConversionMap *map = ExecGetChildToRootMap(relinfo);
-		bool		delete_old_table = transition_capture->tcs_delete_old_table;
-		bool		update_old_table = transition_capture->tcs_update_old_table;
-		bool		update_new_table = transition_capture->tcs_update_new_table;
-		bool		insert_new_table = transition_capture->tcs_insert_new_table;
 
 		/*
-		 * For INSERT events NEW should be non-NULL, for DELETE events OLD
-		 * should be non-NULL, whereas for UPDATE events normally both OLD and
-		 * NEW are non-NULL.  But for UPDATE events fired for capturing
-		 * transition tuples during UPDATE partition-key row movement, OLD is
-		 * NULL when the event is for a row being inserted, whereas NEW is
-		 * NULL when the event is for a row being deleted.
+		 * Capture the old tuple in the appropriate transition table based on
+		 * the event.
 		 */
-		Assert(!(event == TRIGGER_EVENT_DELETE && delete_old_table &&
-				 TupIsNull(oldslot)));
-		Assert(!(event == TRIGGER_EVENT_INSERT && insert_new_table &&
-				 TupIsNull(newslot)));
-
-		if (!TupIsNull(oldslot) &&
-			((event == TRIGGER_EVENT_DELETE && delete_old_table) ||
-			 (event == TRIGGER_EVENT_UPDATE && update_old_table)))
+		if (!TupIsNull(oldslot))
 		{
 			Tuplestorestate *old_tuplestore;
 
-			old_tuplestore = transition_capture->tcs_private->old_tuplestore;
-
-			if (map != NULL)
-			{
-				AfterTriggersTableData *table = transition_capture->tcs_private;
-				TupleTableSlot *storeslot;
-
-				storeslot = GetAfterTriggersStoreSlot(table, map->outdesc);
-				execute_attr_map_slot(map->attrMap, oldslot, storeslot);
-				tuplestore_puttupleslot(old_tuplestore, storeslot);
-			}
-			else
-				tuplestore_puttupleslot(old_tuplestore, oldslot);
+			old_tuplestore = AfterTriggerGetTransitionTable(event,
+															oldslot,
+															NULL,
+															transition_capture);
+			TransitionTableAddTuple(estate, transition_capture, relinfo,
+									oldslot, NULL, old_tuplestore);
 		}
-		if (!TupIsNull(newslot) &&
-			((event == TRIGGER_EVENT_INSERT && insert_new_table) ||
-			 (event == TRIGGER_EVENT_UPDATE && update_new_table)))
+
+		/*
+		 * Capture the new tuple in the appropriate transition table based on
+		 * the event.
+		 */
+		if (!TupIsNull(newslot))
 		{
 			Tuplestorestate *new_tuplestore;
 
-			new_tuplestore = transition_capture->tcs_private->new_tuplestore;
-
-			if (original_insert_tuple != NULL)
-				tuplestore_puttupleslot(new_tuplestore,
-										original_insert_tuple);
-			else if (map != NULL)
-			{
-				AfterTriggersTableData *table = transition_capture->tcs_private;
-				TupleTableSlot *storeslot;
-
-				storeslot = GetAfterTriggersStoreSlot(table, map->outdesc);
-				execute_attr_map_slot(map->attrMap, newslot, storeslot);
-				tuplestore_puttupleslot(new_tuplestore, storeslot);
-			}
-			else
-				tuplestore_puttupleslot(new_tuplestore, newslot);
+			new_tuplestore = AfterTriggerGetTransitionTable(event,
+															NULL,
+															newslot,
+															transition_capture);
+			TransitionTableAddTuple(estate, transition_capture, relinfo,
+									newslot,
+									transition_capture->tcs_original_insert_tuple,
+									new_tuplestore);
 		}
 
 		/*
-- 
2.30.2

