Rethink plpgsql's way of handling SPI execution during an exception block.

We don't really want to start a new SPI connection, just keep using the old
one; otherwise we have memory management problems as illustrated by
John Kennedy's bug report of today.  This requires a bit of a hack to
ensure the SPI stack state is properly restored, but then again what we
were doing before was a hack too, strictly speaking.  Add a regression
test to cover this case.

5 files changed, 72 insertions, 20 deletions

diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index b533470e1d5..4549ba00080 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/executor/spi.c,v 1.131 2004/10/13 01:25:10 neilc Exp $
+ *	  $PostgreSQL: pgsql/src/backend/executor/spi.c,v 1.132 2004/11/16 18:10:13 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -270,6 +270,14 @@ SPI_pop(void)
 	_SPI_curid--;
 }
 
+/* Restore state of SPI stack after aborting a subtransaction */
+void
+SPI_restore_connection(void)
+{
+	Assert(_SPI_connected >= 0);
+	_SPI_curid = _SPI_connected - 1;
+}
+
 /* Parse, plan, and execute a querystring */
 int
 SPI_execute(const char *src, bool read_only, int tcount)
diff --git a/src/include/executor/spi.h b/src/include/executor/spi.h
index 3014dbce840..d5ba89fa3ea 100644
--- a/src/include/executor/spi.h
+++ b/src/include/executor/spi.h
@@ -2,7 +2,7 @@
  *
  * spi.h
  *
- * $PostgreSQL: pgsql/src/include/executor/spi.h,v 1.49 2004/09/16 16:58:40 tgl Exp $
+ * $PostgreSQL: pgsql/src/include/executor/spi.h,v 1.50 2004/11/16 18:10:13 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -81,6 +81,7 @@ extern int	SPI_connect(void);
 extern int	SPI_finish(void);
 extern void SPI_push(void);
 extern void SPI_pop(void);
+extern void SPI_restore_connection(void);
 extern int	SPI_execute(const char *src, bool read_only, int tcount);
 extern int	SPI_execute_plan(void *plan, Datum *Values, const char *Nulls,
 							 bool read_only, int tcount);
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index 12ca6c0888a..fddbfefac47 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -3,7 +3,7 @@
  *			  procedural language
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.120 2004/09/16 16:58:44 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.121 2004/11/16 18:10:14 tgl Exp $
  *
  *	  This software is copyrighted by Jan Wieck - Hamburg.
  *
@@ -899,20 +899,11 @@ exec_stmt_block(PLpgSQL_execstate *estate, PLpgSQL_stmt_block *block)
 		MemoryContext oldcontext = CurrentMemoryContext;
 		ResourceOwner oldowner = CurrentResourceOwner;
 		volatile bool caught = false;
-		int			xrc;
 
-		/*
-		 * Start a subtransaction, and re-connect to SPI within it
-		 */
-		SPI_push();
 		BeginInternalSubTransaction(NULL);
 		/* Want to run statements inside function's memory context */
 		MemoryContextSwitchTo(oldcontext);
 
-		if ((xrc = SPI_connect()) != SPI_OK_CONNECT)
-			elog(ERROR, "SPI_connect failed: %s",
-				 SPI_result_code_string(xrc));
-
 		PG_TRY();
 		{
 			rc = exec_stmts(estate, block->body);
@@ -928,12 +919,17 @@ exec_stmt_block(PLpgSQL_execstate *estate, PLpgSQL_stmt_block *block)
 			edata = CopyErrorData();
 			FlushErrorState();
 
-			/* Abort the inner transaction (and inner SPI connection) */
+			/* Abort the inner transaction */
 			RollbackAndReleaseCurrentSubTransaction();
 			MemoryContextSwitchTo(oldcontext);
 			CurrentResourceOwner = oldowner;
 
-			SPI_pop();
+			/*
+			 * If AtEOSubXact_SPI() popped any SPI context of the subxact,
+			 * it will have left us in a disconnected state.  We need this
+			 * hack to return to connected state.
+			 */
+			SPI_restore_connection();
 
 			/* Look for a matching exception handler */
 			exceptions = block->exceptions;
@@ -960,15 +956,14 @@ exec_stmt_block(PLpgSQL_execstate *estate, PLpgSQL_stmt_block *block)
 		/* Commit the inner transaction, return to outer xact context */
 		if (!caught)
 		{
-			if ((xrc = SPI_finish()) != SPI_OK_FINISH)
-				elog(ERROR, "SPI_finish failed: %s",
-					 SPI_result_code_string(xrc));
-
 			ReleaseCurrentSubTransaction();
 			MemoryContextSwitchTo(oldcontext);
 			CurrentResourceOwner = oldowner;
-
-			SPI_pop();
+			/*
+			 * AtEOSubXact_SPI() should not have popped any SPI context,
+			 * but just in case it did, make sure we remain connected.
+			 */
+			SPI_restore_connection();
 		}
 	}
 	else
diff --git a/src/test/regress/expected/plpgsql.out b/src/test/regress/expected/plpgsql.out
index e1b57f9716b..8674fa9531d 100644
--- a/src/test/regress/expected/plpgsql.out
+++ b/src/test/regress/expected/plpgsql.out
@@ -1931,6 +1931,36 @@ select * from foo;
  20
 (2 rows)
 
+-- Test for pass-by-ref values being stored in proper context
+create function test_variable_storage() returns text as $$
+declare x text;
+begin
+  x := '1234';
+  begin
+    x := x || '5678';
+    -- force error inside subtransaction SPI context
+    perform trap_zero_divide(-100);
+  exception
+    when others then
+      x := x || '9012';
+  end;
+  return x;
+end$$ language plpgsql;
+select test_variable_storage();
+NOTICE:  should see this
+CONTEXT:  SQL statement "SELECT  trap_zero_divide(-100)"
+PL/pgSQL function "test_variable_storage" line 7 at perform
+NOTICE:  should see this only if -100 <> 0
+CONTEXT:  SQL statement "SELECT  trap_zero_divide(-100)"
+PL/pgSQL function "test_variable_storage" line 7 at perform
+NOTICE:  should see this only if -100 fits in smallint
+CONTEXT:  SQL statement "SELECT  trap_zero_divide(-100)"
+PL/pgSQL function "test_variable_storage" line 7 at perform
+ test_variable_storage 
+-----------------------
+ 123456789012
+(1 row)
+
 --
 -- test foreign key error trapping
 --
diff --git a/src/test/regress/sql/plpgsql.sql b/src/test/regress/sql/plpgsql.sql
index 367a73986e1..f9f307b1ac0 100644
--- a/src/test/regress/sql/plpgsql.sql
+++ b/src/test/regress/sql/plpgsql.sql
@@ -1696,6 +1696,24 @@ reset statement_timeout;
 
 select * from foo;
 
+-- Test for pass-by-ref values being stored in proper context
+create function test_variable_storage() returns text as $$
+declare x text;
+begin
+  x := '1234';
+  begin
+    x := x || '5678';
+    -- force error inside subtransaction SPI context
+    perform trap_zero_divide(-100);
+  exception
+    when others then
+      x := x || '9012';
+  end;
+  return x;
+end$$ language plpgsql;
+
+select test_variable_storage();
+
 --
 -- test foreign key error trapping
 --