plpython-generator-fix.diff

*** src/pl/plpython/expected/plpython_setof.out
--- src/pl/plpython/expected/plpython_setof.out	2010-11-14 17:38:39.000000000 +0100
***************
*** 31,36 ****
--- 31,44 ----
  		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,109 ****
--- 115,126 ----
   
  (2 rows)
  
+ SELECT test_setof_spi_in_iterator();
+  test_setof_spi_in_iterator 
+ ----------------------------
+  Hello
+  Brave
+  New
+  World
+ (4 rows)
+ 
*** src/pl/plpython/plpython.c
--- src/pl/plpython/plpython.c	2010-11-14 17:38:39.000000000 +0100
***************
*** 1000,1012 ****
  		}
  
  		/*
! 		 * 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");
  
  		if (proc->is_setof)
  		{
--- 1000,1010 ----
  		}
  
  		/*
! 		 * Cannot disconnect from SPI yet, because in case of a SRF returning a
! 		 * Python iterator we will use PyIter_Next() which calls back into
! 		 * Python code, which can contain calls to SPI functions. In that case
! 		 * we will disconnect from SPI after the iterator is exhausted.
  		 */
  
  		if (proc->is_setof)
  		{
***************
*** 1064,1074 ****
--- 1062,1085 ----
  							(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;
*** src/pl/plpython/sql/plpython_setof.sql
--- src/pl/plpython/sql/plpython_setof.sql	2010-11-14 17:38:39.000000000 +0100
***************
*** 35,40 ****
--- 35,50 ----
  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,53 ****
--- 61,65 ----
  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();