Fix transition tables for wCTEs.

The original coding didn't handle this case properly; each separate
DML substatement needs its own set of transitions.

Patch by Thomas Munro

Discussion: https://postgr.es/m/CAL9smLCDQ%3D2o024rBgtD4WihzX8B3C6u_oSQ2K3%2BR5grJrV0bg%40mail.gmail.com

src/backend/commands/copy.c

@@ -1416,6 +1416,12 @@ BeginCopy(ParseState *pstate,
 					 errmsg("table \"%s\" does not have OIDs",
 							RelationGetRelationName(cstate->rel))));
 
+		/*
+		 * If there are any triggers with transition tables on the named
+		 * relation, we need to be prepared to capture transition tuples.
+		 */
+		cstate->transition_capture = MakeTransitionCaptureState(rel->trigdesc);
+
 		/* Initialize state for CopyFrom tuple routing. */
 		if (is_from && rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
 		{
@@ -1439,14 +1445,6 @@ BeginCopy(ParseState *pstate,
 			cstate->partition_tupconv_maps = partition_tupconv_maps;
 			cstate->partition_tuple_slot = partition_tuple_slot;
 
-			/*
-			 * If there are any triggers with transition tables on the named
-			 * relation, we need to be prepared to capture transition tuples
-			 * from child relations too.
-			 */
-			cstate->transition_capture =
-				MakeTransitionCaptureState(rel->trigdesc);
-
 			/*
 			 * If we are capturing transition tuples, they may need to be
 			 * converted from partition format back to partitioned table
@@ -2807,7 +2805,7 @@ CopyFrom(CopyState cstate)
 		pq_endmsgread();
 
 	/* Execute AFTER STATEMENT insertion triggers */
-	ExecASInsertTriggers(estate, resultRelInfo);
+	ExecASInsertTriggers(estate, resultRelInfo, cstate->transition_capture);
 
 	/* Handle queued AFTER triggers */
 	AfterTriggerEndQuery(estate);
@@ -2935,7 +2933,7 @@ CopyFromInsertBatch(CopyState cstate, EState *estate, CommandId mycid,
 			cstate->cur_lineno = firstBufferedLineNo + i;
 			ExecARInsertTriggers(estate, resultRelInfo,
 								 bufferedTuples[i],
-								 NIL, NULL);
+								 NIL, cstate->transition_capture);
 		}
 	}
 

src/backend/executor/execReplication.c

@@ -419,6 +419,12 @@ ExecSimpleRelationInsert(EState *estate, TupleTableSlot *slot)
 		ExecARInsertTriggers(estate, resultRelInfo, tuple,
 							 recheckIndexes, NULL);
 
+		/*
+		 * XXX we should in theory pass a TransitionCaptureState object to the
+		 * above to capture transition tuples, but after statement triggers
+		 * don't actually get fired by replication yet anyway
+		 */
+
 		list_free(recheckIndexes);
 	}
 }

src/backend/executor/nodeModifyTable.c

@@ -1442,14 +1442,18 @@ fireASTriggers(ModifyTableState *node)
 		case CMD_INSERT:
 			if (node->mt_onconflict == ONCONFLICT_UPDATE)
 				ExecASUpdateTriggers(node->ps.state,
-									 resultRelInfo);
-			ExecASInsertTriggers(node->ps.state, resultRelInfo);
+									 resultRelInfo,
+									 node->mt_transition_capture);
+			ExecASInsertTriggers(node->ps.state, resultRelInfo,
+								 node->mt_transition_capture);
 			break;
 		case CMD_UPDATE:
-			ExecASUpdateTriggers(node->ps.state, resultRelInfo);
+			ExecASUpdateTriggers(node->ps.state, resultRelInfo,
+								 node->mt_transition_capture);
 			break;
 		case CMD_DELETE:
-			ExecASDeleteTriggers(node->ps.state, resultRelInfo);
+			ExecASDeleteTriggers(node->ps.state, resultRelInfo,
+								 node->mt_transition_capture);
 			break;
 		default:
 			elog(ERROR, "unknown operation");
@@ -2304,6 +2308,10 @@ ExecEndModifyTable(ModifyTableState *node)
 {
 	int			i;
 
+	/* Free transition tables */
+	if (node->mt_transition_capture != NULL)
+		DestroyTransitionCaptureState(node->mt_transition_capture);
+
 	/*
 	 * Allow any FDWs to shut down
 	 */

src/include/commands/trigger.h

@@ -42,8 +42,8 @@ typedef struct TriggerData
 } TriggerData;
 
 /*
- * Meta-data to control the capture of old and new tuples into transition
- * tables from child tables.
+ * The state for capturing old and new tuples into transition tables for a
+ * single ModifyTable node.
  */
 typedef struct TransitionCaptureState
 {
@@ -72,6 +72,10 @@ typedef struct TransitionCaptureState
 	 * the original tuple directly.
 	 */
 	HeapTuple	tcs_original_insert_tuple;
+
+	/* The tuplestores backing the transition tables. */
+	Tuplestorestate *tcs_old_tuplestore;
+	Tuplestorestate *tcs_new_tuplestore;
 } TransitionCaptureState;
 
 /*
@@ -162,13 +166,15 @@ extern TriggerDesc *CopyTriggerDesc(TriggerDesc *trigdesc);
 
 extern const char *FindTriggerIncompatibleWithInheritance(TriggerDesc *trigdesc);
 extern TransitionCaptureState *MakeTransitionCaptureState(TriggerDesc *trigdesc);
+extern void DestroyTransitionCaptureState(TransitionCaptureState *tcs);
 
 extern void FreeTriggerDesc(TriggerDesc *trigdesc);
 
 extern void ExecBSInsertTriggers(EState *estate,
 					 ResultRelInfo *relinfo);
 extern void ExecASInsertTriggers(EState *estate,
-					 ResultRelInfo *relinfo);
+					 ResultRelInfo *relinfo,
+					 TransitionCaptureState *transition_capture);
 extern TupleTableSlot *ExecBRInsertTriggers(EState *estate,
 					 ResultRelInfo *relinfo,
 					 TupleTableSlot *slot);
@@ -183,7 +189,8 @@ extern TupleTableSlot *ExecIRInsertTriggers(EState *estate,
 extern void ExecBSDeleteTriggers(EState *estate,
 					 ResultRelInfo *relinfo);
 extern void ExecASDeleteTriggers(EState *estate,
-					 ResultRelInfo *relinfo);
+					 ResultRelInfo *relinfo,
+					 TransitionCaptureState *transition_capture);
 extern bool ExecBRDeleteTriggers(EState *estate,
 					 EPQState *epqstate,
 					 ResultRelInfo *relinfo,
@@ -200,7 +207,8 @@ extern bool ExecIRDeleteTriggers(EState *estate,
 extern void ExecBSUpdateTriggers(EState *estate,
 					 ResultRelInfo *relinfo);
 extern void ExecASUpdateTriggers(EState *estate,
-					 ResultRelInfo *relinfo);
+					 ResultRelInfo *relinfo,
+					 TransitionCaptureState *transition_capture);
 extern TupleTableSlot *ExecBRUpdateTriggers(EState *estate,
 					 EPQState *epqstate,
 					 ResultRelInfo *relinfo,

src/test/regress/expected/triggers.out

@@ -2194,6 +2194,26 @@ DETAIL:  ROW triggers with transition tables are not supported in inheritance hi
 drop trigger child_row_trig on child;
 alter table child inherit parent;
 drop table child, parent;
+--
+-- Verify behavior of queries with wCTEs, where multiple transition
+-- tuplestores can be active at the same time because there are
+-- multiple DML statements that might fire triggers with transition
+-- tables
+--
+create table table1 (a int);
+create table table2 (a text);
+create trigger table1_trig
+  after insert on table1 referencing new table as new_table
+  for each statement execute procedure dump_insert();
+create trigger table2_trig
+  after insert on table2 referencing new table as new_table
+  for each statement execute procedure dump_insert();
+with wcte as (insert into table1 values (42))
+  insert into table2 values ('hello world');
+NOTICE:  trigger = table2_trig, new table = ("hello world")
+NOTICE:  trigger = table1_trig, new table = (42)
+drop table table1;
+drop table table2;
 -- cleanup
 drop function dump_insert();
 drop function dump_update();

src/test/regress/sql/triggers.sql

@@ -1704,6 +1704,27 @@ alter table child inherit parent;
 
 drop table child, parent;
 
+--
+-- Verify behavior of queries with wCTEs, where multiple transition
+-- tuplestores can be active at the same time because there are
+-- multiple DML statements that might fire triggers with transition
+-- tables
+--
+create table table1 (a int);
+create table table2 (a text);
+create trigger table1_trig
+  after insert on table1 referencing new table as new_table
+  for each statement execute procedure dump_insert();
+create trigger table2_trig
+  after insert on table2 referencing new table as new_table
+  for each statement execute procedure dump_insert();
+
+with wcte as (insert into table1 values (42))
+  insert into table2 values ('hello world');
+
+drop table table1;
+drop table table2;
+
 -- cleanup
 drop function dump_insert();
 drop function dump_update();