From add0ea88e7722b48d1f3a7c38e3cfd7f1e854674 Mon Sep 17 00:00:00 2001
From: Tom Lane <tgl@sss.pgh.pa.us>
Date: Mon, 15 Nov 2010 14:26:55 -0500
Subject: [PATCH] Fix aboriginal mistake in plpython's set-returning-function
 support.

We must stay in the function's SPI context until done calling the iterator
that returns the set result.  Otherwise, any attempt to invoke SPI features
in the python code called by the iterator will malfunction.  Diagnosis and
patch by Jan Urbanski, per bug report from Jean-Baptiste Quenot.

Back-patch to 8.2; there was no support for SRFs in previous versions of
plpython.
---
 src/pl/plpython/expected/plpython_setof.out | 17 +++++++++++++
 src/pl/plpython/plpython.c                  | 28 +++++++++++++++------
 src/pl/plpython/sql/plpython_setof.sql      | 11 ++++++++
 3 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/src/pl/plpython/expected/plpython_setof.out b/src/pl/plpython/expected/plpython_setof.out
index ebf896df01f..ac9765fc882 100644
--- a/src/pl/plpython/expected/plpython_setof.out
+++ b/src/pl/plpython/expected/plpython_setof.out
@@ -31,6 +31,14 @@ class producer:
 		return self.icontent
 return producer(count, content)
 $$ LANGUAGE plpythonu;
+CREATE FUNCTION test_setof_spi_in_iterator() RETURNS SETOF text AS
+$$
+    for s in ('Hello', 'Brave', 'New', 'World'):
+        plpy.execute('select 1')
+        yield s
+        plpy.execute('select 2')
+$$
+LANGUAGE plpythonu;
 -- Test set returning functions
 SELECT test_setof_as_list(0, 'list');
  test_setof_as_list 
@@ -107,3 +115,12 @@ SELECT test_setof_as_iterator(2, null);
  
 (2 rows)
 
+SELECT test_setof_spi_in_iterator();
+ test_setof_spi_in_iterator 
+----------------------------
+ Hello
+ Brave
+ New
+ World
+(4 rows)
+
diff --git a/src/pl/plpython/plpython.c b/src/pl/plpython/plpython.c
index f3352b3bc6f..1ae12396a12 100644
--- a/src/pl/plpython/plpython.c
+++ b/src/pl/plpython/plpython.c
@@ -985,7 +985,10 @@ PLy_function_handler(FunctionCallInfo fcinfo, PLyProcedure *proc)
 	{
 		if (!proc->is_setof || proc->setof == NULL)
 		{
-			/* Simple type returning function or first time for SETOF function */
+			/*
+			 * Simple type returning function or first time for SETOF function:
+			 * actually execute the function.
+			 */
 			plargs = PLy_function_build_args(fcinfo, proc);
 			plrv = PLy_procedure_call(proc, "args", plargs);
 			if (!proc->is_setof)
@@ -1000,14 +1003,10 @@ PLy_function_handler(FunctionCallInfo fcinfo, PLyProcedure *proc)
 		}
 
 		/*
-		 * Disconnect from SPI manager and then create the return values datum
-		 * (if the input function does a palloc for it this must not be
-		 * allocated in the SPI memory context because SPI_finish would free
-		 * it).
+		 * If it returns a set, call the iterator to get the next return item.
+		 * We stay in the SPI context while doing this, because PyIter_Next()
+		 * calls back into Python code which might contain SPI calls.
 		 */
-		if (SPI_finish() != SPI_OK_FINISH)
-			elog(ERROR, "SPI_finish failed");
-
 		if (proc->is_setof)
 		{
 			bool		has_error = false;
@@ -1064,11 +1063,24 @@ PLy_function_handler(FunctionCallInfo fcinfo, PLyProcedure *proc)
 							(errcode(ERRCODE_DATA_EXCEPTION),
 						  errmsg("error fetching next item from iterator")));
 
+				/* Disconnect from the SPI manager before returning */
+				if (SPI_finish() != SPI_OK_FINISH)
+					elog(ERROR, "SPI_finish failed");
+
 				fcinfo->isnull = true;
 				return (Datum) NULL;
 			}
 		}
 
+		/*
+		 * Disconnect from SPI manager and then create the return values datum
+		 * (if the input function does a palloc for it this must not be
+		 * allocated in the SPI memory context because SPI_finish would free
+		 * it).
+		 */
+		if (SPI_finish() != SPI_OK_FINISH)
+			elog(ERROR, "SPI_finish failed");
+
 		plerrcontext.callback = plpython_return_error_callback;
 		plerrcontext.previous = error_context_stack;
 		error_context_stack = &plerrcontext;
diff --git a/src/pl/plpython/sql/plpython_setof.sql b/src/pl/plpython/sql/plpython_setof.sql
index 53d91a9e7d7..80a8d5b4c1e 100644
--- a/src/pl/plpython/sql/plpython_setof.sql
+++ b/src/pl/plpython/sql/plpython_setof.sql
@@ -35,6 +35,15 @@ class producer:
 return producer(count, content)
 $$ LANGUAGE plpythonu;
 
+CREATE FUNCTION test_setof_spi_in_iterator() RETURNS SETOF text AS
+$$
+    for s in ('Hello', 'Brave', 'New', 'World'):
+        plpy.execute('select 1')
+        yield s
+        plpy.execute('select 2')
+$$
+LANGUAGE plpythonu;
+
 
 -- Test set returning functions
 SELECT test_setof_as_list(0, 'list');
@@ -51,3 +60,5 @@ SELECT test_setof_as_iterator(0, 'list');
 SELECT test_setof_as_iterator(1, 'list');
 SELECT test_setof_as_iterator(2, 'list');
 SELECT test_setof_as_iterator(2, null);
+
+SELECT test_setof_spi_in_iterator();
-- 
GitLab