--
-- Test explicit subtransactions
--
-- Test table to see if transactions get properly rolled back
CREATE TABLE subtransaction_tbl (
    i integer
);
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":
        raise Exception("Python exception")
$$ LANGUAGE plpython3u;
SELECT subtransaction_ctx_test();
 subtransaction_ctx_test 
-------------------------
 
(1 row)

SELECT * FROM subtransaction_tbl;
 i 
---
 1
 2
(2 rows)

TRUNCATE subtransaction_tbl;
SELECT subtransaction_ctx_test('SPI');
ERROR:  spiexceptions.InvalidTextRepresentation: invalid input syntax for type integer: "oops"
LINE 1: INSERT INTO subtransaction_tbl VALUES ('oops')
                                               ^
QUERY:  INSERT INTO subtransaction_tbl VALUES ('oops')
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_ctx_test", line 6, in <module>
    plpy.execute("INSERT INTO subtransaction_tbl VALUES ('oops')")
PL/Python function "subtransaction_ctx_test"
SELECT * FROM subtransaction_tbl;
 i 
---
(0 rows)

TRUNCATE subtransaction_tbl;
SELECT subtransaction_ctx_test('Python');
ERROR:  Exception: Python exception
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_ctx_test", line 8, in <module>
    raise Exception("Python exception")
PL/Python function "subtransaction_ctx_test"
SELECT * FROM subtransaction_tbl;
 i 
---
(0 rows)

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 as e:
        if not swallow:
            raise
        plpy.notice("Swallowed %s(%r)" % (e.__class__.__name__, e.args[0]))
return "ok"
$$ LANGUAGE plpython3u;
SELECT subtransaction_nested_test();
ERROR:  spiexceptions.SyntaxError: syntax error at or near "error"
LINE 1: error
        ^
QUERY:  error
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_nested_test", line 8, in <module>
    plpy.execute("error")
PL/Python function "subtransaction_nested_test"
SELECT * FROM subtransaction_tbl;
 i 
---
(0 rows)

TRUNCATE subtransaction_tbl;
SELECT subtransaction_nested_test('t');
NOTICE:  Swallowed SyntaxError('syntax error at or near "error"')
 subtransaction_nested_test 
----------------------------
 ok
(1 row)

SELECT * FROM subtransaction_tbl;
 i 
---
 1
 2
(2 rows)

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 plpython3u;
SELECT subtransaction_deeply_nested_test();
NOTICE:  Swallowed SyntaxError('syntax error at or near "error"')
 subtransaction_deeply_nested_test 
-----------------------------------
 ok
(1 row)

SELECT * FROM subtransaction_tbl;
 i 
---
 1
 2
 1
 2
(4 rows)

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 plpython3u;
CREATE FUNCTION subtransaction_enter_without_exit() RETURNS void
AS $$
plpy.subtransaction().__enter__()
$$ LANGUAGE plpython3u;
CREATE FUNCTION subtransaction_exit_twice() RETURNS void
AS $$
plpy.subtransaction().__enter__()
plpy.subtransaction().__exit__(None, None, None)
plpy.subtransaction().__exit__(None, None, None)
$$ LANGUAGE plpython3u;
CREATE FUNCTION subtransaction_enter_twice() RETURNS void
AS $$
plpy.subtransaction().__enter__()
plpy.subtransaction().__enter__()
$$ LANGUAGE plpython3u;
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 plpython3u;
CREATE FUNCTION subtransaction_enter_same_subtransaction_twice() RETURNS void
AS $$
s = plpy.subtransaction()
s.__enter__()
s.__enter__()
s.__exit__(None, None, None)
$$ LANGUAGE plpython3u;
-- 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 plpython3u;
CREATE FUNCTION subtransaction_exit_subtransaction_in_with() RETURNS void
AS $$
try:
    with plpy.subtransaction() as s:
        s.__exit__(None, None, None)
except ValueError as e:
    raise ValueError(e)
$$ LANGUAGE plpython3u;
SELECT subtransaction_exit_without_enter();
ERROR:  ValueError: this subtransaction has not been entered
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_exit_without_enter", line 2, in <module>
    plpy.subtransaction().__exit__(None, None, None)
PL/Python function "subtransaction_exit_without_enter"
SELECT subtransaction_enter_without_exit();
WARNING:  forcibly aborting a subtransaction that has not been exited
 subtransaction_enter_without_exit 
-----------------------------------
 
(1 row)

SELECT subtransaction_exit_twice();
WARNING:  forcibly aborting a subtransaction that has not been exited
ERROR:  ValueError: this subtransaction has not been entered
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_exit_twice", line 3, in <module>
    plpy.subtransaction().__exit__(None, None, None)
PL/Python function "subtransaction_exit_twice"
SELECT subtransaction_enter_twice();
WARNING:  forcibly aborting a subtransaction that has not been exited
WARNING:  forcibly aborting a subtransaction that has not been exited
 subtransaction_enter_twice 
----------------------------
 
(1 row)

SELECT subtransaction_exit_same_subtransaction_twice();
ERROR:  ValueError: this subtransaction has already been exited
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_exit_same_subtransaction_twice", line 5, in <module>
    s.__exit__(None, None, None)
PL/Python function "subtransaction_exit_same_subtransaction_twice"
SELECT subtransaction_enter_same_subtransaction_twice();
WARNING:  forcibly aborting a subtransaction that has not been exited
ERROR:  ValueError: this subtransaction has already been entered
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_enter_same_subtransaction_twice", line 4, in <module>
    s.__enter__()
PL/Python function "subtransaction_enter_same_subtransaction_twice"
SELECT subtransaction_enter_subtransaction_in_with();
ERROR:  ValueError: this subtransaction has already been entered
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_enter_subtransaction_in_with", line 3, in <module>
    s.__enter__()
PL/Python function "subtransaction_enter_subtransaction_in_with"
SELECT subtransaction_exit_subtransaction_in_with();
ERROR:  ValueError: this subtransaction has already been exited
CONTEXT:  Traceback (most recent call last):
  PL/Python function "subtransaction_exit_subtransaction_in_with", line 6, in <module>
    raise ValueError(e)
PL/Python function "subtransaction_exit_subtransaction_in_with"
-- Make sure we don't get a "current transaction is aborted" error
SELECT 1 as test;
 test 
------
    1
(1 row)

-- 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 plpython3u;
SELECT subtransaction_mix_explicit_and_implicit();
WARNING:  Caught a SPI error from an explicit subtransaction
WARNING:  Caught a SPI error
 subtransaction_mix_explicit_and_implicit 
------------------------------------------
 
(1 row)

SELECT * FROM subtransaction_tbl;
 i 
---
 1
 2
(2 rows)

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 plpython3u;
SELECT subtransaction_alternative_names();
 subtransaction_alternative_names 
----------------------------------
 
(1 row)

-- 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 plpython3u;
SELECT try_catch_inside_subtransaction();
NOTICE:  caught
 try_catch_inside_subtransaction 
---------------------------------
 
(1 row)

SELECT * FROM subtransaction_tbl;
 i 
---
 1
(1 row)

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 plpython3u;
SELECT pk_violation_inside_subtransaction();
NOTICE:  caught
 pk_violation_inside_subtransaction 
------------------------------------
 
(1 row)

SELECT * FROM subtransaction_tbl;
 i 
---
 1
(1 row)

DROP TABLE subtransaction_tbl;
-- cursor/subtransactions interactions
CREATE FUNCTION cursor_in_subxact() RETURNS int AS $$
with plpy.subtransaction():
    cur = plpy.cursor("select * from generate_series(1, 20) as gen(i)")
    cur.fetch(10)
fetched = cur.fetch(10);
return int(fetched[5]["i"])
$$ LANGUAGE plpython3u;
CREATE FUNCTION cursor_aborted_subxact() RETURNS int AS $$
try:
    with plpy.subtransaction():
        cur = plpy.cursor("select * from generate_series(1, 20) as gen(i)")
        cur.fetch(10);
        plpy.execute("select no_such_function()")
except plpy.SPIError:
    fetched = cur.fetch(10)
    return int(fetched[5]["i"])
return 0 # not reached
$$ LANGUAGE plpython3u;
CREATE FUNCTION cursor_plan_aborted_subxact() RETURNS int AS $$
try:
    with plpy.subtransaction():
        plpy.execute('create temporary table tmp(i) '
                     'as select generate_series(1, 10)')
        plan = plpy.prepare("select i from tmp")
        cur = plpy.cursor(plan)
        plpy.execute("select no_such_function()")
except plpy.SPIError:
    fetched = cur.fetch(5)
    return fetched[2]["i"]
return 0 # not reached
$$ LANGUAGE plpython3u;
CREATE FUNCTION cursor_close_aborted_subxact() RETURNS boolean AS $$
try:
    with plpy.subtransaction():
        cur = plpy.cursor('select 1')
        plpy.execute("select no_such_function()")
except plpy.SPIError:
    cur.close()
    return True
return False # not reached
$$ LANGUAGE plpython3u;
SELECT cursor_in_subxact();
 cursor_in_subxact 
-------------------
                16
(1 row)

SELECT cursor_aborted_subxact();
ERROR:  ValueError: iterating a cursor in an aborted subtransaction
CONTEXT:  Traceback (most recent call last):
  PL/Python function "cursor_aborted_subxact", line 8, in <module>
    fetched = cur.fetch(10)
PL/Python function "cursor_aborted_subxact"
SELECT cursor_plan_aborted_subxact();
ERROR:  ValueError: iterating a cursor in an aborted subtransaction
CONTEXT:  Traceback (most recent call last):
  PL/Python function "cursor_plan_aborted_subxact", line 10, in <module>
    fetched = cur.fetch(5)
PL/Python function "cursor_plan_aborted_subxact"
SELECT cursor_close_aborted_subxact();
ERROR:  ValueError: closing a cursor in an aborted subtransaction
CONTEXT:  Traceback (most recent call last):
  PL/Python function "cursor_close_aborted_subxact", line 7, in <module>
    cur.close()
PL/Python function "cursor_close_aborted_subxact"