From 3fe1d1fa9b8b6ecff08f74394be0b0e81ea4ab03 Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@enterprisedb.com>
Date: Mon, 3 Apr 2017 09:39:55 +1200
Subject: [PATCH] Support transition tables in PL/python triggers.

If a trigger is created with transition tables using the new syntax
REFERENCING OLD/NEW TABLE AS ..., allow SQL executed by a PL/python
trigger to see the transient tables.
---
 src/pl/plpython/expected/plpython_trigger.out | 21 ++++++++++++++++++
 src/pl/plpython/plpy_exec.c                   | 31 +++++++++++++++++++++++++++
 src/pl/plpython/sql/plpython_trigger.sql      | 24 +++++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/src/pl/plpython/expected/plpython_trigger.out b/src/pl/plpython/expected/plpython_trigger.out
index 4148963f3a2..d7ab8ac6b8e 100644
--- a/src/pl/plpython/expected/plpython_trigger.out
+++ b/src/pl/plpython/expected/plpython_trigger.out
@@ -503,3 +503,24 @@ SELECT * FROM b;
  1234
 (1 row)
 
+-- check that SQL run in trigger code can see transition tables
+CREATE TABLE transition_table_test (id int, name text);
+INSERT INTO transition_table_test VALUES (1, 'a');
+CREATE FUNCTION transition_table_test_f() RETURNS trigger LANGUAGE plpythonu AS
+$$
+    rv = plpy.execute("SELECT * FROM old_table")
+    assert(rv.nrows() == 1)
+    plpy.info("old: " + str(rv[0]["id"]) + " -> " + rv[0]["name"])
+    rv = plpy.execute("SELECT * FROM new_table")
+    assert(rv.nrows() == 1)
+    plpy.info("new: " + str(rv[0]["id"]) + " -> " + rv[0]["name"])
+    return None
+$$;
+CREATE TRIGGER a_t AFTER UPDATE ON transition_table_test
+  REFERENCING OLD TABLE AS old_table NEW TABLE AS new_table
+  FOR EACH STATEMENT EXECUTE PROCEDURE transition_table_test_f();
+UPDATE transition_table_test SET name = 'b';
+INFO:  old: 1 -> a
+INFO:  new: 1 -> b
+DROP TABLE transition_table_test;
+DROP FUNCTION transition_table_test_f();
diff --git a/src/pl/plpython/plpy_exec.c b/src/pl/plpython/plpy_exec.c
index 697a0e1cc03..7a649e034c6 100644
--- a/src/pl/plpython/plpy_exec.c
+++ b/src/pl/plpython/plpy_exec.c
@@ -345,6 +345,37 @@ PLy_exec_trigger(FunctionCallInfo fcinfo, PLyProcedure *proc)
 
 	PG_TRY();
 	{
+		if (tdata->tg_newtable)
+		{
+			EphemeralNamedRelation enr =
+				palloc(sizeof(EphemeralNamedRelationData));
+			int		rc PG_USED_FOR_ASSERTS_ONLY;
+
+			enr->md.name = tdata->tg_trigger->tgnewtable;
+			enr->md.reliddesc = tdata->tg_relation->rd_id;
+			enr->md.tupdesc = NULL;
+			enr->md.enrtype = ENR_NAMED_TUPLESTORE;
+			enr->md.enrtuples = tuplestore_tuple_count(tdata->tg_oldtable);
+			enr->reldata = tdata->tg_newtable;
+			rc = SPI_register_relation(enr);
+			Assert(rc >= 0);
+		}
+		if (tdata->tg_oldtable)
+		{
+			EphemeralNamedRelation enr =
+				palloc(sizeof(EphemeralNamedRelationData));
+			int		rc PG_USED_FOR_ASSERTS_ONLY;
+
+			enr->md.name = tdata->tg_trigger->tgoldtable;
+			enr->md.reliddesc = tdata->tg_relation->rd_id;
+			enr->md.tupdesc = NULL;
+			enr->md.enrtype = ENR_NAMED_TUPLESTORE;
+			enr->md.enrtuples = tuplestore_tuple_count(tdata->tg_oldtable);
+			enr->reldata = tdata->tg_oldtable;
+			rc = SPI_register_relation(enr);
+			Assert(rc >= 0);
+		}
+
 		plargs = PLy_trigger_build_args(fcinfo, proc, &rv);
 		plrv = PLy_procedure_call(proc, "TD", plargs);
 
diff --git a/src/pl/plpython/sql/plpython_trigger.sql b/src/pl/plpython/sql/plpython_trigger.sql
index a054fe729bc..79c24b714b5 100644
--- a/src/pl/plpython/sql/plpython_trigger.sql
+++ b/src/pl/plpython/sql/plpython_trigger.sql
@@ -406,3 +406,27 @@ SELECT * FROM a;
 DROP TABLE a;
 INSERT INTO b DEFAULT VALUES;
 SELECT * FROM b;
+
+-- check that SQL run in trigger code can see transition tables
+
+CREATE TABLE transition_table_test (id int, name text);
+INSERT INTO transition_table_test VALUES (1, 'a');
+
+CREATE FUNCTION transition_table_test_f() RETURNS trigger LANGUAGE plpythonu AS
+$$
+    rv = plpy.execute("SELECT * FROM old_table")
+    assert(rv.nrows() == 1)
+    plpy.info("old: " + str(rv[0]["id"]) + " -> " + rv[0]["name"])
+    rv = plpy.execute("SELECT * FROM new_table")
+    assert(rv.nrows() == 1)
+    plpy.info("new: " + str(rv[0]["id"]) + " -> " + rv[0]["name"])
+    return None
+$$;
+
+CREATE TRIGGER a_t AFTER UPDATE ON transition_table_test
+  REFERENCING OLD TABLE AS old_table NEW TABLE AS new_table
+  FOR EACH STATEMENT EXECUTE PROCEDURE transition_table_test_f();
+UPDATE transition_table_test SET name = 'b';
+
+DROP TABLE transition_table_test;
+DROP FUNCTION transition_table_test_f();
-- 
2.12.2