diff options
| -rw-r--r-- | src/pl/plpython/expected/plpython_setof.out | 17 | ||||
| -rw-r--r-- | src/pl/plpython/plpython.c | 28 | ||||
| -rw-r--r-- | 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(); |