PL/pgSQL: Nested CALL with transactions

So far, a nested CALL or DO in PL/pgSQL would not establish a context
where transaction control statements were allowed.  This fixes that by
handling CALL and DO specially in PL/pgSQL, passing the atomic/nonatomic
execution context through and doing the required management around
transaction boundaries.

Reviewed-by: Tomas Vondra <tomas.vondra@2ndquadrant.com>

Author: petere

doc/src/sgml/plpgsql.sgml

@@ -3463,9 +3463,9 @@ END LOOP <optional> <replaceable>label</replaceable> </optional>;
    <title>Transaction Management</title>
 
    <para>
-    In procedures invoked by the <command>CALL</command> command from the top
-    level as well as in anonymous code blocks (<command>DO</command> command)
-    called from the top level, it is possible to end transactions using the
+    In procedures invoked by the <command>CALL</command> command
+    as well as in anonymous code blocks (<command>DO</command> command),
+    it is possible to end transactions using the
     commands <command>COMMIT</command> and <command>ROLLBACK</command>.  A new
     transaction is started automatically after a transaction is ended using
     these commands, so there is no separate <command>START
@@ -3495,6 +3495,20 @@ CALL transaction_test1();
 </programlisting>
    </para>
 
+   <para>
+    Transaction control is only possible in <command>CALL</command> or
+    <command>DO</command> invocations from the top level or nested
+    <command>CALL</command> or <command>DO</command> invocations without any
+    other intervening command.  For example, if the call stack is
+    <command>CALL proc1()</command> &rarr; <command>CALL proc2()</command>
+    &rarr; <command>CALL proc3()</command>, then the second and third
+    procedures can perform transaction control actions.  But if the call stack
+    is <command>CALL proc1()</command> &rarr; <command>SELECT
+    func2()</command> &rarr; <command>CALL proc3()</command>, then the last
+    procedure cannot do transaction control, because of the
+    <command>SELECT</command> in between.
+   </para>
+
    <para>
     A transaction cannot be ended inside a loop over a query result, nor
     inside a block with exception handlers.

src/backend/executor/spi.c

@@ -2041,8 +2041,11 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 	 *
 	 * In the first two cases, we can just push the snap onto the stack once
 	 * for the whole plan list.
+	 *
+	 * But if the plan has no_snapshots set to true, then don't manage
+	 * snapshots at all.  The caller should then take care of that.
 	 */
-	if (snapshot != InvalidSnapshot)
+	if (snapshot != InvalidSnapshot && !plan->no_snapshots)
 	{
 		if (read_only)
 		{
@@ -2121,7 +2124,7 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 		 * In the default non-read-only case, get a new snapshot, replacing
 		 * any that we pushed in a previous cycle.
 		 */
-		if (snapshot == InvalidSnapshot && !read_only)
+		if (snapshot == InvalidSnapshot && !read_only && !plan->no_snapshots)
 		{
 			if (pushed_active_snap)
 				PopActiveSnapshot();
@@ -2172,7 +2175,7 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 			 * If not read-only mode, advance the command counter before each
 			 * command and update the snapshot.
 			 */
-			if (!read_only)
+			if (!read_only && !plan->no_snapshots)
 			{
 				CommandCounterIncrement();
 				UpdateActiveSnapshotCommandId();
@@ -2203,10 +2206,23 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 			else
 			{
 				char		completionTag[COMPLETION_TAG_BUFSIZE];
+				ProcessUtilityContext context;
+
+				/*
+				 * If the SPI context is atomic, or we are asked to manage
+				 * snapshots, then we are in an atomic execution context.
+				 * Conversely, to propagate a nonatomic execution context, the
+				 * caller must be in a nonatomic SPI context and manage
+				 * snapshots itself.
+				 */
+				if (_SPI_current->atomic || !plan->no_snapshots)
+					context = PROCESS_UTILITY_QUERY;
+				else
+					context = PROCESS_UTILITY_QUERY_NONATOMIC;
 
 				ProcessUtility(stmt,
 							   plansource->query_string,
-							   PROCESS_UTILITY_QUERY,
+							   context,
 							   paramLI,
 							   _SPI_current->queryEnv,
 							   dest,
@@ -2638,11 +2654,8 @@ _SPI_make_plan_non_temp(SPIPlanPtr plan)
 	oldcxt = MemoryContextSwitchTo(plancxt);
 
 	/* Copy the SPI_plan struct and subsidiary data into the new context */
-	newplan = (SPIPlanPtr) palloc(sizeof(_SPI_plan));
+	newplan = (SPIPlanPtr) palloc0(sizeof(_SPI_plan));
 	newplan->magic = _SPI_PLAN_MAGIC;
-	newplan->saved = false;
-	newplan->oneshot = false;
-	newplan->plancache_list = NIL;
 	newplan->plancxt = plancxt;
 	newplan->cursor_options = plan->cursor_options;
 	newplan->nargs = plan->nargs;
@@ -2705,11 +2718,8 @@ _SPI_save_plan(SPIPlanPtr plan)
 	oldcxt = MemoryContextSwitchTo(plancxt);
 
 	/* Copy the SPI plan into its own context */
-	newplan = (SPIPlanPtr) palloc(sizeof(_SPI_plan));
+	newplan = (SPIPlanPtr) palloc0(sizeof(_SPI_plan));
 	newplan->magic = _SPI_PLAN_MAGIC;
-	newplan->saved = false;
-	newplan->oneshot = false;
-	newplan->plancache_list = NIL;
 	newplan->plancxt = plancxt;
 	newplan->cursor_options = plan->cursor_options;
 	newplan->nargs = plan->nargs;

src/backend/tcop/utility.c

@@ -382,7 +382,7 @@ standard_ProcessUtility(PlannedStmt *pstmt,
 {
 	Node	   *parsetree = pstmt->utilityStmt;
 	bool		isTopLevel = (context == PROCESS_UTILITY_TOPLEVEL);
-	bool		isAtomicContext = (context != PROCESS_UTILITY_TOPLEVEL || IsTransactionBlock());
+	bool		isAtomicContext = (!(context == PROCESS_UTILITY_TOPLEVEL || context == PROCESS_UTILITY_QUERY_NONATOMIC) || IsTransactionBlock());
 	ParseState *pstate;
 
 	check_xact_readonly(parsetree);

src/include/executor/spi_priv.h

@@ -86,6 +86,7 @@ typedef struct _SPI_plan
 	int			magic;			/* should equal _SPI_PLAN_MAGIC */
 	bool		saved;			/* saved or unsaved plan? */
 	bool		oneshot;		/* one-shot plan? */
+	bool		no_snapshots;	/* let the caller handle the snapshots */
 	List	   *plancache_list; /* one CachedPlanSource per parsetree */
 	MemoryContext plancxt;		/* Context containing _SPI_plan and data */
 	int			cursor_options; /* Cursor options used for planning */

src/include/tcop/utility.h

@@ -20,6 +20,7 @@ typedef enum
 {
 	PROCESS_UTILITY_TOPLEVEL,	/* toplevel interactive command */
 	PROCESS_UTILITY_QUERY,		/* a complete query, but not toplevel */
+	PROCESS_UTILITY_QUERY_NONATOMIC, /* a complete query, nonatomic execution context */
 	PROCESS_UTILITY_SUBCOMMAND	/* a portion of a query */
 } ProcessUtilityContext;
 

src/pl/plpgsql/src/expected/plpgsql_transaction.out

@@ -1,10 +1,10 @@
 CREATE TABLE test1 (a int, b text);
-CREATE PROCEDURE transaction_test1()
+CREATE PROCEDURE transaction_test1(x int, y text)
 LANGUAGE plpgsql
 AS $$
 BEGIN
-    FOR i IN 0..9 LOOP
-        INSERT INTO test1 (a) VALUES (i);
+    FOR i IN 0..x LOOP
+        INSERT INTO test1 (a, b) VALUES (i, y);
         IF i % 2 = 0 THEN
             COMMIT;
         ELSE
@@ -13,15 +13,15 @@ BEGIN
     END LOOP;
 END
 $$;
-CALL transaction_test1();
+CALL transaction_test1(9, 'foo');
 SELECT * FROM test1;
- a | b 
----+---
- 0 | 
- 2 | 
- 4 | 
- 6 | 
- 8 | 
+ a |  b  
+---+-----
+ 0 | foo
+ 2 | foo
+ 4 | foo
+ 6 | foo
+ 8 | foo
 (5 rows)
 
 TRUNCATE test1;
@@ -51,9 +51,9 @@ SELECT * FROM test1;
 
 -- transaction commands not allowed when called in transaction block
 START TRANSACTION;
-CALL transaction_test1();
+CALL transaction_test1(9, 'error');
 ERROR:  invalid transaction termination
-CONTEXT:  PL/pgSQL function transaction_test1() line 6 at COMMIT
+CONTEXT:  PL/pgSQL function transaction_test1(integer,text) line 6 at COMMIT
 COMMIT;
 START TRANSACTION;
 DO LANGUAGE plpgsql $$ BEGIN COMMIT; END $$;
@@ -90,14 +90,14 @@ CREATE FUNCTION transaction_test3() RETURNS int
 LANGUAGE plpgsql
 AS $$
 BEGIN
-    CALL transaction_test1();
+    CALL transaction_test1(9, 'error');
     RETURN 1;
 END;
 $$;
 SELECT transaction_test3();
 ERROR:  invalid transaction termination
-CONTEXT:  PL/pgSQL function transaction_test1() line 6 at COMMIT
-SQL statement "CALL transaction_test1()"
+CONTEXT:  PL/pgSQL function transaction_test1(integer,text) line 6 at COMMIT
+SQL statement "CALL transaction_test1(9, 'error')"
 PL/pgSQL function transaction_test3() line 3 at CALL
 SELECT * FROM test1;
  a | b 
@@ -130,6 +130,57 @@ $$;
 CALL transaction_test5();
 ERROR:  invalid transaction termination
 CONTEXT:  PL/pgSQL function transaction_test5() line 3 at COMMIT
+TRUNCATE test1;
+-- nested procedure calls
+CREATE PROCEDURE transaction_test6(c text)
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    CALL transaction_test1(9, c);
+END;
+$$;
+CALL transaction_test6('bar');
+SELECT * FROM test1;
+ a |  b  
+---+-----
+ 0 | bar
+ 2 | bar
+ 4 | bar
+ 6 | bar
+ 8 | bar
+(5 rows)
+
+TRUNCATE test1;
+CREATE PROCEDURE transaction_test7()
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    DO 'BEGIN CALL transaction_test1(9, $x$baz$x$); END;';
+END;
+$$;
+CALL transaction_test7();
+SELECT * FROM test1;
+ a |  b  
+---+-----
+ 0 | baz
+ 2 | baz
+ 4 | baz
+ 6 | baz
+ 8 | baz
+(5 rows)
+
+CREATE PROCEDURE transaction_test8()
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    EXECUTE 'CALL transaction_test1(10, $x$baz$x$)';
+END;
+$$;
+CALL transaction_test8();
+ERROR:  invalid transaction termination
+CONTEXT:  PL/pgSQL function transaction_test1(integer,text) line 6 at COMMIT
+SQL statement "CALL transaction_test1(10, $x$baz$x$)"
+PL/pgSQL function transaction_test8() line 3 at EXECUTE
 -- commit inside cursor loop
 CREATE TABLE test2 (x int);
 INSERT INTO test2 VALUES (0), (1), (2), (3), (4);