Add transition table support to plpgsql.

Kevin Grittner and Thomas Munro
Reviewed by Heikki Linnakangas, David Fetter, and Thomas Munro
with valuable comments and suggestions from many others

3 files changed, 63 insertions, 11 deletions

diff --git a/src/pl/plpgsql/src/pl_comp.c b/src/pl/plpgsql/src/pl_comp.c
index bed343ea0c0..a6375511f61 100644
--- a/src/pl/plpgsql/src/pl_comp.c
+++ b/src/pl/plpgsql/src/pl_comp.c
@@ -589,11 +589,11 @@ do_compile(FunctionCallInfo fcinfo,
 				  errmsg("trigger functions cannot have declared arguments"),
 						 errhint("The arguments of the trigger can be accessed through TG_NARGS and TG_ARGV instead.")));
 
-			/* Add the record for referencing NEW */
+			/* Add the record for referencing NEW ROW */
 			rec = plpgsql_build_record("new", 0, true);
 			function->new_varno = rec->dno;
 
-			/* Add the record for referencing OLD */
+			/* Add the record for referencing OLD ROW */
 			rec = plpgsql_build_record("old", 0, true);
 			function->old_varno = rec->dno;
 
@@ -2453,15 +2453,16 @@ compute_function_hashkey(FunctionCallInfo fcinfo,
 	hashkey->isTrigger = CALLED_AS_TRIGGER(fcinfo);
 
 	/*
-	 * if trigger, get relation OID.  In validation mode we do not know what
-	 * relation is intended to be used, so we leave trigrelOid zero; the hash
-	 * entry built in this case will never really be used.
+	 * if trigger, get its OID.  In validation mode we do not know what
+	 * relation or transition table names are intended to be used, so we leave
+	 * trigOid zero; the hash entry built in this case will never really be
+	 * used.
 	 */
 	if (hashkey->isTrigger && !forValidator)
 	{
 		TriggerData *trigdata = (TriggerData *) fcinfo->context;
 
-		hashkey->trigrelOid = RelationGetRelid(trigdata->tg_relation);
+		hashkey->trigOid = trigdata->tg_trigger->tgoid;
 	}
 
 	/* get input collation, if known */
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index c27935b51ba..43da986fc0a 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -690,6 +690,47 @@ plpgsql_exec_trigger(PLpgSQL_function *func,
 		elog(ERROR, "unrecognized trigger action: not INSERT, DELETE, or UPDATE");
 
 	/*
+	 * Capture the NEW and OLD transition TABLE tuplestores (if specified for
+	 * this trigger).
+	 */
+	if (trigdata->tg_newtable || trigdata->tg_oldtable)
+	{
+		estate.queryEnv = create_queryEnv();
+		if (trigdata->tg_newtable)
+		{
+			EphemeralNamedRelation enr =
+			  palloc(sizeof(EphemeralNamedRelationData));
+			int rc PG_USED_FOR_ASSERTS_ONLY;
+
+			enr->md.name = trigdata->tg_trigger->tgnewtable;
+			enr->md.reliddesc = RelationGetRelid(trigdata->tg_relation);
+			enr->md.tupdesc = NULL;
+			enr->md.enrtype = ENR_NAMED_TUPLESTORE;
+			enr->md.enrtuples = tuplestore_tuple_count(trigdata->tg_newtable);
+			enr->reldata = trigdata->tg_newtable;
+			register_ENR(estate.queryEnv, enr);
+			rc = SPI_register_relation(enr);
+			Assert(rc >= 0);
+		}
+		if (trigdata->tg_oldtable)
+		{
+			EphemeralNamedRelation enr =
+			  palloc(sizeof(EphemeralNamedRelationData));
+			int rc PG_USED_FOR_ASSERTS_ONLY;
+
+			enr->md.name = trigdata->tg_trigger->tgoldtable;
+			enr->md.reliddesc = RelationGetRelid(trigdata->tg_relation);
+			enr->md.tupdesc = NULL;
+			enr->md.enrtype = ENR_NAMED_TUPLESTORE;
+			enr->md.enrtuples = tuplestore_tuple_count(trigdata->tg_oldtable);
+			enr->reldata = trigdata->tg_oldtable;
+			register_ENR(estate.queryEnv, enr);
+			rc = SPI_register_relation(enr);
+			Assert(rc >= 0);
+		}
+	}
+
+	/*
 	 * Assign the special tg_ variables
 	 */
 
@@ -3442,6 +3483,9 @@ plpgsql_estate_setup(PLpgSQL_execstate *estate,
 	estate->paramLI->paramMask = NULL;
 	estate->params_dirty = false;
 
+	/* default tuplestore cache to "none" */
+	estate->queryEnv = NULL;
+
 	/* set up for use of appropriate simple-expression EState and cast hash */
 	if (simple_eval_estate)
 	{
@@ -7329,6 +7373,9 @@ exec_dynquery_with_params(PLpgSQL_execstate *estate,
 	/* Release transient data */
 	MemoryContextReset(stmt_mcontext);
 
+	/* Make sure the portal knows about any named tuplestores. */
+	portal->queryEnv = estate->queryEnv;
+
 	return portal;
 }
 
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index b7e103b5143..43a62ef34e4 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -20,6 +20,7 @@
 #include "commands/event_trigger.h"
 #include "commands/trigger.h"
 #include "executor/spi.h"
+#include "utils/queryenvironment.h"
 
 /**********************************************************************
  * Definitions
@@ -780,12 +781,12 @@ typedef struct PLpgSQL_func_hashkey
 	/* be careful that pad bytes in this struct get zeroed! */
 
 	/*
-	 * For a trigger function, the OID of the relation triggered on is part of
-	 * the hash key --- we want to compile the trigger separately for each
-	 * relation it is used with, in case the rowtype is different.  Zero if
-	 * not called as a trigger.
+	 * For a trigger function, the OID of the trigger is part of the hash key
+	 * --- we want to compile the trigger function separately for each trigger
+	 * it is used with, in case the rowtype or transition table names are
+	 * different.  Zero if not called as a trigger.
 	 */
-	Oid			trigrelOid;
+	Oid			trigOid;
 
 	/*
 	 * We must include the input collation as part of the hash key too,
@@ -910,6 +911,9 @@ typedef struct PLpgSQL_execstate
 	ParamListInfo paramLI;
 	bool		params_dirty;	/* T if any resettable datum has been passed */
 
+	/* custom environment for parsing/execution of query for this context */
+	QueryEnvironment *queryEnv;
+
 	/* EState to use for "simple" expression evaluation */
 	EState	   *simple_eval_estate;