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 1446f0d4558..ea47067a855 100644
--- a/src/pl/plpython/plpython.c
+++ b/src/pl/plpython/plpython.c
@@ -984,7 +984,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)
@@ -999,14 +1002,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;
@@ -1063,11 +1062,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();