diff options
| -rw-r--r-- | src/pl/plpython/expected/plpython_function.out | 8 | ||||
| -rw-r--r-- | src/pl/plpython/expected/plpython_test.out | 9 | ||||
| -rw-r--r-- | src/pl/plpython/plpython.c | 28 | ||||
| -rw-r--r-- | src/pl/plpython/sql/plpython_function.sql | 9 | ||||
| -rw-r--r-- | src/pl/plpython/sql/plpython_test.sql | 2 |
5 files changed, 48 insertions, 8 deletions
diff --git a/src/pl/plpython/expected/plpython_function.out b/src/pl/plpython/expected/plpython_function.out index e1ffa7302db..b2bba947040 100644 --- a/src/pl/plpython/expected/plpython_function.out +++ b/src/pl/plpython/expected/plpython_function.out @@ -403,6 +403,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 returning tuples -- diff --git a/src/pl/plpython/expected/plpython_test.out b/src/pl/plpython/expected/plpython_test.out index 170abe7ab6e..1deab6c0d34 100644 --- a/src/pl/plpython/expected/plpython_test.out +++ b/src/pl/plpython/expected/plpython_test.out @@ -298,6 +298,15 @@ 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) + -- Test tuple returning functions SELECT * FROM test_table_record_as('dict', null, null, false); first | second diff --git a/src/pl/plpython/plpython.c b/src/pl/plpython/plpython.c index 63ad15e7afa..edcbc2a7ed8 100644 --- a/src/pl/plpython/plpython.c +++ b/src/pl/plpython/plpython.c @@ -766,7 +766,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) @@ -781,14 +784,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; @@ -845,12 +844,25 @@ 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"); + + /* * If the function is declared to return void, the Python return value * must be None. For void-returning functions, we also treat a None * return value as a special "void datum" rather than NULL (as is the diff --git a/src/pl/plpython/sql/plpython_function.sql b/src/pl/plpython/sql/plpython_function.sql index 224d5196a3c..133704dbbe7 100644 --- a/src/pl/plpython/sql/plpython_function.sql +++ b/src/pl/plpython/sql/plpython_function.sql @@ -444,6 +444,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 returning tuples diff --git a/src/pl/plpython/sql/plpython_test.sql b/src/pl/plpython/sql/plpython_test.sql index dafcae089e9..75d2f06ffc5 100644 --- a/src/pl/plpython/sql/plpython_test.sql +++ b/src/pl/plpython/sql/plpython_test.sql @@ -95,6 +95,8 @@ 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(); + -- Test tuple returning functions SELECT * FROM test_table_record_as('dict', null, null, false); SELECT * FROM test_table_record_as('dict', 'one', null, false); |