PL/Python explicit subtransactions

Adds a context manager, obtainable by plpy.subtransaction(), to run a
group of statements in a subtransaction.

Jan Urbański, reviewed by Steve Singer, additional scribbling by me

doc/src/sgml/plpython.sgml

@@ -955,7 +955,7 @@ $$ LANGUAGE plpythonu;
 
   </sect2>
 
-  <sect2>
+  <sect2 id="plpython-trapping">
    <title>Trapping Errors</title>
 
    <para>
@@ -981,6 +981,136 @@ $$ LANGUAGE plpythonu;
   </sect2>
  </sect1>
 
+ <sect1 id="plpython-subtransaction">
+  <title>Explicit Subtransactions</title>
+
+  <para>
+   Recovering from errors caused by database access as described in
+   <xref linkend="plpython-trapping"> can lead to an undesirable
+   situation where some operations succeed before one of them fails,
+   and after recovering from that error the data is left in an
+   inconsistent state.  PL/Python offers a solution to this problem in
+   the form of explicit subtransactions.
+  </para>
+
+  <sect2>
+   <title>Subtransaction Context Managers</title>
+
+   <para>
+    Consider a function that implements a transfer between two
+    accounts:
+<programlisting>
+CREATE FUNCTION transfer_funds() RETURNS void AS $$
+try:
+    plpy.execute("UPDATE accounts SET balance = balance - 100 WHERE account_name = 'joe'")
+    plpy.execute("UPDATE accounts SET balance = balance + 100 WHERE account_name = 'mary'")
+except plpy.SPIError, e:
+    result = "error transferring funds: %s" % e.args
+else:
+    result = "funds transferred correctly"
+plan = plpy.prepare("INSERT INTO operations (result) VALUES ($1)", ["text"])
+plpy.execute(plan, [result])
+$$ LANGUAGE plpythonu;
+</programlisting>
+    If the second <literal>UPDATE</literal> statement results in an
+    exception being raised, this function will report the error, but
+    the result of the first <literal>UPDATE</literal> will
+    nevertheless be committed.  In other words, the funds will be
+    withdrawn from Joe's account, but will not be transferred to
+    Mary's account.
+   </para>
+
+   <para>
+    To avoid such issues, you can wrap your
+    <literal>plpy.execute</literal> calls in an explicit
+    subtransaction.  The <literal>plpy</literal> module provides a
+    helper object to manage explicit subtransactions that gets created
+    with the <literal>plpy.subtransaction()</literal> function.
+    Objects created by this function implement the
+    <ulink url="http://docs.python.org/library/stdtypes.html#context-manager-types">
+    context manager interface</ulink>.  Using explicit subtransactions
+    we can rewrite our function as:
+<programlisting>
+CREATE FUNCTION transfer_funds2() RETURNS void AS $$
+try:
+    with plpy.subtransaction():
+        plpy.execute("UPDATE accounts SET balance = balance - 100 WHERE account_name = 'joe'")
+        plpy.execute("UPDATE accounts SET balance = balance + 100 WHERE account_name = 'mary'")
+except plpy.SPIError, e:
+    result = "error transferring funds: %s" % e.args
+else:
+    result = "funds transferred correctly"
+plan = plpy.prepare("INSERT INTO operations (result) VALUES ($1)", ["text"])
+plpy.execute(plan, [result])
+$$ LANGUAGE plpythonu;
+</programlisting>
+    Note that the use of <literal>try/catch</literal> is still
+    required.  Otherwise the exception would propagate to the top of
+    the Python stack and would cause the whole function to abort with
+    a <productname>PostgreSQL</productname> error, so that the
+    <literal>operations</literal> table would not have any row
+    inserted into it.  The subtransaction context manager does not
+    trap errors, it only assures that all database operations executed
+    inside its scope will be atomically committed or rolled back.  A
+    rollback of the subtransaction block occurrs on any kind of
+    exception exit, not only ones caused by errors originating from
+    database access.  A regular Python exception raised inside an
+    explicit subtransaction block would also cause the subtransaction
+    to be rolled back.
+   </para>
+  </sect2>
+
+  <sect2>
+   <title>Older Python Versions</title>
+
+   <para>
+    Context managers syntax using the <literal>with</literal> keyword
+    is available by default in Python 2.6.  If using PL/Python with an
+    older Python version, it is still possible to use explicit
+    subtransactions, although not as transparently.  You can call the
+    subtransaction manager's <literal>__enter__</literal> and
+    <literal>__exit__</literal> functions using the
+    <literal>enter</literal> and <literal>exit</literal> convenience
+    aliases.  The example function that transfers funds could be
+    written as:
+<programlisting>
+CREATE FUNCTION transfer_funds_old() RETURNS void AS $$
+try:
+    subxact = plpy.subtransaction()
+    subxact.enter()
+    try:
+        plpy.execute("UPDATE accounts SET balance = balance - 100 WHERE account_name = 'joe'")
+        plpy.execute("UPDATE accounts SET balance = balance + 100 WHERE account_name = 'mary'")
+    except:
+        import sys
+        subxact.exit(*sys.exc_info())
+        raise
+    else:
+        subxact.exit(None, None, None)
+except plpy.SPIError, e:
+    result = "error transferring funds: %s" % e.args
+else:
+    result = "funds transferred correctly"
+
+plan = plpy.prepare("INSERT INTO operations (result) VALUES ($1)", ["text"])
+plpy.execute(plan, [result])
+$$ LANGUAGE plpythonu;
+</programlisting>
+   </para>
+
+   <note>
+    <para>
+     Although context managers were implemented in Python 2.5, to use
+     the <literal>with</literal> syntax in that version you need to
+     use a <ulink
+     url="http://docs.python.org/release/2.5/ref/future.html">future
+     statement</ulink>.  Because of implementation details, however,
+     you cannot use future statements in PL/Python functions.
+    </para>
+   </note>
+  </sect2>
+ </sect1>
+
  <sect1 id="plpython-util">
   <title>Utility Functions</title>
   <para>

src/pl/plpython/Makefile

@@ -81,6 +81,7 @@ REGRESS = \
 	plpython_unicode \
 	plpython_quote \
 	plpython_composite \
+	plpython_subtransaction \
 	plpython_drop
 # where to find psql for running the tests
 PSQLDIR = $(bindir)

src/pl/plpython/expected/README

@@ -6,6 +6,8 @@ plpython_unicode.out		server encoding != SQL_ASCII and client encoding == UTF8;
 plpython_unicode_0.out		server encoding != SQL_ASCII and client encoding != UTF8; else ...
 plpython_unicode_3.out		server encoding == SQL_ASCII
 
+plpython_subtransaction_0.out	Python 2.5 and older (without with statement)
+
 plpython_types_3.out		Python 3.x
 
 Note: Building with Python 2.2 is supported, but there are no expected

src/pl/plpython/expected/plpython_test.out

@@ -43,9 +43,9 @@ contents.sort()
 return ", ".join(contents)
 $$ LANGUAGE plpythonu;
 select module_contents();
-                                                            module_contents                                                            
----------------------------------------------------------------------------------------------------------------------------------------
- Error, Fatal, SPIError, debug, error, execute, fatal, info, log, notice, prepare, quote_ident, quote_literal, quote_nullable, warning
+                                                                    module_contents                                                                    
+-------------------------------------------------------------------------------------------------------------------------------------------------------
+ Error, Fatal, SPIError, debug, error, execute, fatal, info, log, notice, prepare, quote_ident, quote_literal, quote_nullable, subtransaction, warning
 (1 row)
 
 CREATE FUNCTION elog_test() RETURNS void

src/pl/plpython/sql/plpython_subtransaction.sql

+--
+-- Test explicit subtransactions
+--
+
+-- Test table to see if transactions get properly rolled back
+
+CREATE TABLE subtransaction_tbl (
+    i integer
+);
+
+-- Explicit case for Python <2.6
+
+CREATE FUNCTION subtransaction_test(what_error text = NULL) RETURNS text
+AS $$
+import sys
+subxact = plpy.subtransaction()
+subxact.__enter__()
+exc = True
+try:
+    try:
+        plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+        plpy.execute("INSERT INTO subtransaction_tbl VALUES (2)")
+        if what_error == "SPI":
+            plpy.execute("INSERT INTO subtransaction_tbl VALUES ('oops')")
+        elif what_error == "Python":
+            plpy.attribute_error
+    except:
+        exc = False
+        subxact.__exit__(*sys.exc_info())
+        raise
+finally:
+    if exc:
+        subxact.__exit__(None, None, None)
+$$ LANGUAGE plpythonu;
+
+SELECT subtransaction_test();
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+SELECT subtransaction_test('SPI');
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+SELECT subtransaction_test('Python');
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+
+-- Context manager case for Python >=2.6
+
+CREATE FUNCTION subtransaction_ctx_test(what_error text = NULL) RETURNS text
+AS $$
+with plpy.subtransaction():
+    plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+    plpy.execute("INSERT INTO subtransaction_tbl VALUES (2)")
+    if what_error == "SPI":
+        plpy.execute("INSERT INTO subtransaction_tbl VALUES ('oops')")
+    elif what_error == "Python":
+        plpy.attribute_error
+$$ LANGUAGE plpythonu;
+
+SELECT subtransaction_ctx_test();
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+SELECT subtransaction_ctx_test('SPI');
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+SELECT subtransaction_ctx_test('Python');
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+
+-- Nested subtransactions
+
+CREATE FUNCTION subtransaction_nested_test(swallow boolean = 'f') RETURNS text
+AS $$
+plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+with plpy.subtransaction():
+    plpy.execute("INSERT INTO subtransaction_tbl VALUES (2)")
+    try:
+        with plpy.subtransaction():
+            plpy.execute("INSERT INTO subtransaction_tbl VALUES (3)")
+            plpy.execute("error")
+    except plpy.SPIError, e:
+        if not swallow:
+            raise
+        plpy.notice("Swallowed %r" % e)
+return "ok"
+$$ LANGUAGE plpythonu;
+
+SELECT subtransaction_nested_test();
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+
+SELECT subtransaction_nested_test('t');
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+
+-- Nested subtransactions that recursively call code dealing with
+-- subtransactions
+
+CREATE FUNCTION subtransaction_deeply_nested_test() RETURNS text
+AS $$
+plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+with plpy.subtransaction():
+    plpy.execute("INSERT INTO subtransaction_tbl VALUES (2)")
+    plpy.execute("SELECT subtransaction_nested_test('t')")
+return "ok"
+$$ LANGUAGE plpythonu;
+
+SELECT subtransaction_deeply_nested_test();
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+
+-- Error conditions from not opening/closing subtransactions
+
+CREATE FUNCTION subtransaction_exit_without_enter() RETURNS void
+AS $$
+plpy.subtransaction().__exit__(None, None, None)
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION subtransaction_enter_without_exit() RETURNS void
+AS $$
+plpy.subtransaction().__enter__()
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION subtransaction_exit_twice() RETURNS void
+AS $$
+plpy.subtransaction().__enter__()
+plpy.subtransaction().__exit__(None, None, None)
+plpy.subtransaction().__exit__(None, None, None)
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION subtransaction_enter_twice() RETURNS void
+AS $$
+plpy.subtransaction().__enter__()
+plpy.subtransaction().__enter__()
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION subtransaction_exit_same_subtransaction_twice() RETURNS void
+AS $$
+s = plpy.subtransaction()
+s.__enter__()
+s.__exit__(None, None, None)
+s.__exit__(None, None, None)
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION subtransaction_enter_same_subtransaction_twice() RETURNS void
+AS $$
+s = plpy.subtransaction()
+s.__enter__()
+s.__enter__()
+s.__exit__(None, None, None)
+$$ LANGUAGE plpythonu;
+
+-- No warnings here, as the subtransaction gets indeed closed
+CREATE FUNCTION subtransaction_enter_subtransaction_in_with() RETURNS void
+AS $$
+with plpy.subtransaction() as s:
+    s.__enter__()
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION subtransaction_exit_subtransaction_in_with() RETURNS void
+AS $$
+with plpy.subtransaction() as s:
+    s.__exit__(None, None, None)
+$$ LANGUAGE plpythonu;
+
+SELECT subtransaction_exit_without_enter();
+SELECT subtransaction_enter_without_exit();
+SELECT subtransaction_exit_twice();
+SELECT subtransaction_enter_twice();
+SELECT subtransaction_exit_same_subtransaction_twice();
+SELECT subtransaction_enter_same_subtransaction_twice();
+SELECT subtransaction_enter_subtransaction_in_with();
+SELECT subtransaction_exit_subtransaction_in_with();
+
+-- Make sure we don't get a "current transaction is aborted" error
+SELECT 1 as test;
+
+-- Mix explicit subtransactions and normal SPI calls
+
+CREATE FUNCTION subtransaction_mix_explicit_and_implicit() RETURNS void
+AS $$
+p = plpy.prepare("INSERT INTO subtransaction_tbl VALUES ($1)", ["integer"])
+try:
+    with plpy.subtransaction():
+        plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+        plpy.execute(p, [2])
+        plpy.execute(p, ["wrong"])
+except plpy.SPIError:
+    plpy.warning("Caught a SPI error from an explicit subtransaction")
+
+try:
+    plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+    plpy.execute(p, [2])
+    plpy.execute(p, ["wrong"])
+except plpy.SPIError:
+    plpy.warning("Caught a SPI error")
+$$ LANGUAGE plpythonu;
+
+SELECT subtransaction_mix_explicit_and_implicit();
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+
+-- Alternative method names for Python <2.6
+
+CREATE FUNCTION subtransaction_alternative_names() RETURNS void
+AS $$
+s = plpy.subtransaction()
+s.enter()
+s.exit(None, None, None)
+$$ LANGUAGE plpythonu;
+
+SELECT subtransaction_alternative_names();
+
+-- try/catch inside a subtransaction block
+
+CREATE FUNCTION try_catch_inside_subtransaction() RETURNS void
+AS $$
+with plpy.subtransaction():
+     plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+     try:
+         plpy.execute("INSERT INTO subtransaction_tbl VALUES ('a')")
+     except plpy.SPIError:
+         plpy.notice("caught")
+$$ LANGUAGE plpythonu;
+
+SELECT try_catch_inside_subtransaction();
+SELECT * FROM subtransaction_tbl;
+TRUNCATE subtransaction_tbl;
+
+ALTER TABLE subtransaction_tbl ADD PRIMARY KEY (i);
+
+CREATE FUNCTION pk_violation_inside_subtransaction() RETURNS void
+AS $$
+with plpy.subtransaction():
+     plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+     try:
+         plpy.execute("INSERT INTO subtransaction_tbl VALUES (1)")
+     except plpy.SPIError:
+         plpy.notice("caught")
+$$ LANGUAGE plpythonu;
+
+SELECT pk_violation_inside_subtransaction();
+SELECT * FROM subtransaction_tbl;
+
+DROP TABLE subtransaction_tbl;