Fix for recursive plpython triggers - Mailing list pgsql-hackers

This fixes bug #18456 [1].  Since we're in back-branch release freeze,
I'll just park it for the moment.  But I think we should shove it in
once the freeze lifts so it's in 17beta1.

            regards, tom lane

[1] https://www.postgresql.org/message-id/18456-82d3d70134aefd28%40postgresql.org

diff --git a/src/pl/plpython/expected/plpython_trigger.out b/src/pl/plpython/expected/plpython_trigger.out
index dd1ca32fa4..4cb90cb520 100644
--- a/src/pl/plpython/expected/plpython_trigger.out
+++ b/src/pl/plpython/expected/plpython_trigger.out
@@ -618,3 +618,30 @@ SELECT * FROM trigger_test_generated;
 ---+---
 (0 rows)

+-- recursive call of a trigger mustn't corrupt TD (bug #18456)
+CREATE TABLE recursive_trigger_test (a int, b int);
+CREATE FUNCTION recursive_trigger_func() RETURNS trigger
+LANGUAGE plpython3u
+AS $$
+if TD["event"] == "UPDATE":
+    plpy.execute("INSERT INTO recursive_trigger_test VALUES (1, 2)")
+    plpy.notice("TD[event] => " + str(TD["event"]) + ", expecting UPDATE");
+else:
+    plpy.notice("TD[event] => " + str(TD["event"]) + ", expecting INSERT");
+return None
+$$;
+CREATE TRIGGER recursive_trigger_trig
+  AFTER INSERT OR UPDATE ON recursive_trigger_test
+  FOR EACH ROW EXECUTE PROCEDURE recursive_trigger_func();
+INSERT INTO recursive_trigger_test VALUES (0, 0);
+NOTICE:  TD[event] => INSERT, expecting INSERT
+UPDATE recursive_trigger_test SET a = 11 WHERE b = 0;
+NOTICE:  TD[event] => INSERT, expecting INSERT
+NOTICE:  TD[event] => UPDATE, expecting UPDATE
+SELECT * FROM recursive_trigger_test;
+ a  | b
+----+---
+ 11 | 0
+  1 | 2
+(2 rows)
+
diff --git a/src/pl/plpython/plpy_exec.c b/src/pl/plpython/plpy_exec.c
index 3145c69699..6f7b5e121d 100644
--- a/src/pl/plpython/plpy_exec.c
+++ b/src/pl/plpython/plpy_exec.c
@@ -335,6 +335,13 @@ PLy_exec_trigger(FunctionCallInfo fcinfo, PLyProcedure *proc)
     PLy_output_setup_tuple(&proc->result, rel_descr, proc);
     PLy_input_setup_tuple(&proc->result_in, rel_descr, proc);

+    /*
+     * If the trigger is called recursively, we must push outer-level
+     * arguments into the stack.  This must be immediately before the PG_TRY
+     * to ensure that the corresponding pop happens.
+     */
+    PLy_global_args_push(proc);
+
     PG_TRY();
     {
         int            rc PG_USED_FOR_ASSERTS_ONLY;
@@ -397,6 +404,7 @@ PLy_exec_trigger(FunctionCallInfo fcinfo, PLyProcedure *proc)
     }
     PG_FINALLY();
     {
+        PLy_global_args_pop(proc);
         Py_XDECREF(plargs);
         Py_XDECREF(plrv);
     }
@@ -501,6 +509,13 @@ PLy_function_save_args(PLyProcedure *proc)
     result->args = PyDict_GetItemString(proc->globals, "args");
     Py_XINCREF(result->args);

+    /* If it's a trigger, also save "TD" */
+    if (proc->is_trigger)
+    {
+        result->td = PyDict_GetItemString(proc->globals, "TD");
+        Py_XINCREF(result->td);
+    }
+
     /* Fetch all the named arguments */
     if (proc->argnames)
     {
@@ -550,6 +565,13 @@ PLy_function_restore_args(PLyProcedure *proc, PLySavedArgs *savedargs)
         Py_DECREF(savedargs->args);
     }

+    /* Restore the "TD" object, too */
+    if (savedargs->td)
+    {
+        PyDict_SetItemString(proc->globals, "TD", savedargs->td);
+        Py_DECREF(savedargs->td);
+    }
+
     /* And free the PLySavedArgs struct */
     pfree(savedargs);
 }
@@ -568,8 +590,9 @@ PLy_function_drop_args(PLySavedArgs *savedargs)
         Py_XDECREF(savedargs->namedargs[i]);
     }

-    /* Drop ref to the "args" object, too */
+    /* Drop refs to the "args" and "TD" objects, too */
     Py_XDECREF(savedargs->args);
+    Py_XDECREF(savedargs->td);

     /* And free the PLySavedArgs struct */
     pfree(savedargs);
@@ -578,9 +601,9 @@ PLy_function_drop_args(PLySavedArgs *savedargs)
 /*
  * Save away any existing arguments for the given procedure, so that we can
  * install new values for a recursive call.  This should be invoked before
- * doing PLy_function_build_args().
+ * doing PLy_function_build_args() or PLy_trigger_build_args().
  *
- * NB: caller must ensure that PLy_global_args_pop gets invoked once, and
+ * NB: callers must ensure that PLy_global_args_pop gets invoked once, and
  * only once, per successful completion of PLy_global_args_push.  Otherwise
  * we'll end up out-of-sync between the actual call stack and the contents
  * of proc->argstack.
diff --git a/src/pl/plpython/plpy_procedure.c b/src/pl/plpython/plpy_procedure.c
index 9cbbe7e3c8..ba7786d31c 100644
--- a/src/pl/plpython/plpy_procedure.c
+++ b/src/pl/plpython/plpy_procedure.c
@@ -183,6 +183,7 @@ PLy_procedure_create(HeapTuple procTup, Oid fn_oid, bool is_trigger)
         proc->fn_readonly = (procStruct->provolatile != PROVOLATILE_VOLATILE);
         proc->is_setof = procStruct->proretset;
         proc->is_procedure = (procStruct->prokind == PROKIND_PROCEDURE);
+        proc->is_trigger = is_trigger;
         proc->src = NULL;
         proc->argnames = NULL;
         proc->args = NULL;
diff --git a/src/pl/plpython/plpy_procedure.h b/src/pl/plpython/plpy_procedure.h
index 8968b5c92e..5db854fc8b 100644
--- a/src/pl/plpython/plpy_procedure.h
+++ b/src/pl/plpython/plpy_procedure.h
@@ -16,6 +16,7 @@ typedef struct PLySavedArgs
 {
     struct PLySavedArgs *next;    /* linked-list pointer */
     PyObject   *args;            /* "args" element of globals dict */
+    PyObject   *td;                /* "TD" element of globals dict, if trigger */
     int            nargs;            /* length of namedargs array */
     PyObject   *namedargs[FLEXIBLE_ARRAY_MEMBER];    /* named args */
 } PLySavedArgs;
@@ -32,6 +33,7 @@ typedef struct PLyProcedure
     bool        fn_readonly;
     bool        is_setof;        /* true, if function returns result set */
     bool        is_procedure;
+    bool        is_trigger;        /* called as trigger? */
     PLyObToDatum result;        /* Function result output conversion info */
     PLyDatumToOb result_in;        /* For converting input tuples in a trigger */
     char       *src;            /* textual procedure code, after mangling */
diff --git a/src/pl/plpython/sql/plpython_trigger.sql b/src/pl/plpython/sql/plpython_trigger.sql
index e5504b9ab1..f6c2ef8d6a 100644
--- a/src/pl/plpython/sql/plpython_trigger.sql
+++ b/src/pl/plpython/sql/plpython_trigger.sql
@@ -467,3 +467,27 @@ FOR EACH ROW EXECUTE PROCEDURE generated_test_func1();
 TRUNCATE trigger_test_generated;
 INSERT INTO trigger_test_generated (i) VALUES (1);
 SELECT * FROM trigger_test_generated;
+
+
+-- recursive call of a trigger mustn't corrupt TD (bug #18456)
+
+CREATE TABLE recursive_trigger_test (a int, b int);
+
+CREATE FUNCTION recursive_trigger_func() RETURNS trigger
+LANGUAGE plpython3u
+AS $$
+if TD["event"] == "UPDATE":
+    plpy.execute("INSERT INTO recursive_trigger_test VALUES (1, 2)")
+    plpy.notice("TD[event] => " + str(TD["event"]) + ", expecting UPDATE");
+else:
+    plpy.notice("TD[event] => " + str(TD["event"]) + ", expecting INSERT");
+return None
+$$;
+
+CREATE TRIGGER recursive_trigger_trig
+  AFTER INSERT OR UPDATE ON recursive_trigger_test
+  FOR EACH ROW EXECUTE PROCEDURE recursive_trigger_func();
+
+INSERT INTO recursive_trigger_test VALUES (0, 0);
+UPDATE recursive_trigger_test SET a = 11 WHERE b = 0;
+SELECT * FROM recursive_trigger_test;