v1-0001-Set-Python-exception-after-failed-commit-or-rollb.patch

From 31bbd62f43ceb0542e0a311dc28704fd702caeb3 Mon Sep 17 00:00:00 2001
From: Peter Eisentraut <peter@eisentraut.org>
Date: Wed, 22 Dec 2021 09:01:24 +0100
Subject: [PATCH v1] Set Python exception after failed commit or rollback

When plpy.commit() or plpy.rollback() failed, the SPI code would
ereport(ERROR) and jump all the way back to the main loop, bypassing
PL/Python.  By contrast, in plpy.execute(), we catch the ereport() and
generate a Python exception.  This has the advantage that you can
catch the exception and you get better backtraces.

An additional problem is apparently that beginning in Python
3.11-to-be, longjmping out of the Python interpreter leaves its
internal state inconsistent, and subsequent invocations of PL/Python
code will crash.  With previous Python versions, this has apparently
not been a problem.

To fix, adapt the logic in PLy_spi_subtransaction_abort() (which is
used by plpy.execute() and others) to catch the PostgreSQL error and
turn it into a Python exception.  Move some code around to avoid
additional cross-file zigzag from this.
---
 .../expected/plpython_transaction.out         |  18 ++-
 src/pl/plpython/plpy_plpymodule.c             |  30 -----
 src/pl/plpython/plpy_spi.c                    | 106 ++++++++++++++++++
 src/pl/plpython/plpy_spi.h                    |   3 +
 4 files changed, 121 insertions(+), 36 deletions(-)

diff --git a/src/pl/plpython/expected/plpython_transaction.out b/src/pl/plpython/expected/plpython_transaction.out
index 14152993c7..9fd9ef500e 100644
--- a/src/pl/plpython/expected/plpython_transaction.out
+++ b/src/pl/plpython/expected/plpython_transaction.out
@@ -55,8 +55,11 @@ for i in range(0, 10):
 return 1
 $$;
 SELECT transaction_test2();
-ERROR:  invalid transaction termination
-CONTEXT:  PL/Python function "transaction_test2"
+ERROR:  spiexceptions.InvalidTransactionTermination: invalid transaction termination
+CONTEXT:  Traceback (most recent call last):
+  PL/Python function "transaction_test2", line 5, in <module>
+    plpy.commit()
+PL/Python function "transaction_test2"
 SELECT * FROM test1;
  a | b 
 ---+---
@@ -70,7 +73,7 @@ plpy.execute("CALL transaction_test1()")
 return 1
 $$;
 SELECT transaction_test3();
-ERROR:  spiexceptions.InvalidTransactionTermination: invalid transaction termination
+ERROR:  spiexceptions.InvalidTransactionTermination: spiexceptions.InvalidTransactionTermination: invalid transaction termination
 CONTEXT:  Traceback (most recent call last):
   PL/Python function "transaction_test3", line 2, in <module>
     plpy.execute("CALL transaction_test1()")
@@ -88,7 +91,7 @@ plpy.execute("DO LANGUAGE plpythonu $x$ plpy.commit() $x$")
 return 1
 $$;
 SELECT transaction_test4();
-ERROR:  spiexceptions.InvalidTransactionTermination: invalid transaction termination
+ERROR:  spiexceptions.InvalidTransactionTermination: spiexceptions.InvalidTransactionTermination: invalid transaction termination
 CONTEXT:  Traceback (most recent call last):
   PL/Python function "transaction_test4", line 2, in <module>
     plpy.execute("DO LANGUAGE plpythonu $x$ plpy.commit() $x$")
@@ -100,8 +103,11 @@ s.enter()
 plpy.commit()
 $$;
 WARNING:  forcibly aborting a subtransaction that has not been exited
-ERROR:  cannot commit while a subtransaction is active
-CONTEXT:  PL/Python anonymous code block
+ERROR:  spiexceptions.InvalidTransactionTermination: cannot commit while a subtransaction is active
+CONTEXT:  Traceback (most recent call last):
+  PL/Python anonymous code block, line 4, in <module>
+    plpy.commit()
+PL/Python anonymous code block
 -- commit inside cursor loop
 CREATE TABLE test2 (x int);
 INSERT INTO test2 VALUES (0), (1), (2), (3), (4);
diff --git a/src/pl/plpython/plpy_plpymodule.c b/src/pl/plpython/plpy_plpymodule.c
index 0365acc95b..907f89d153 100644
--- a/src/pl/plpython/plpy_plpymodule.c
+++ b/src/pl/plpython/plpy_plpymodule.c
@@ -40,8 +40,6 @@ static PyObject *PLy_fatal(PyObject *self, PyObject *args, PyObject *kw);
 static PyObject *PLy_quote_literal(PyObject *self, PyObject *args);
 static PyObject *PLy_quote_nullable(PyObject *self, PyObject *args);
 static PyObject *PLy_quote_ident(PyObject *self, PyObject *args);
-static PyObject *PLy_commit(PyObject *self, PyObject *args);
-static PyObject *PLy_rollback(PyObject *self, PyObject *args);
 
 
 /* A list of all known exceptions, generated from backend/utils/errcodes.txt */
@@ -577,31 +575,3 @@ PLy_output(volatile int level, PyObject *self, PyObject *args, PyObject *kw)
 	 */
 	Py_RETURN_NONE;
 }
-
-static PyObject *
-PLy_commit(PyObject *self, PyObject *args)
-{
-	PLyExecutionContext *exec_ctx = PLy_current_execution_context();
-
-	SPI_commit();
-	SPI_start_transaction();
-
-	/* was cleared at transaction end, reset pointer */
-	exec_ctx->scratch_ctx = NULL;
-
-	Py_RETURN_NONE;
-}
-
-static PyObject *
-PLy_rollback(PyObject *self, PyObject *args)
-{
-	PLyExecutionContext *exec_ctx = PLy_current_execution_context();
-
-	SPI_rollback();
-	SPI_start_transaction();
-
-	/* was cleared at transaction end, reset pointer */
-	exec_ctx->scratch_ctx = NULL;
-
-	Py_RETURN_NONE;
-}
diff --git a/src/pl/plpython/plpy_spi.c b/src/pl/plpython/plpy_spi.c
index 99c1b4f28f..469a9b14ac 100644
--- a/src/pl/plpython/plpy_spi.c
+++ b/src/pl/plpython/plpy_spi.c
@@ -456,6 +456,112 @@ PLy_spi_execute_fetch_result(SPITupleTable *tuptable, uint64 rows, int status)
 	return (PyObject *) result;
 }
 
+PyObject *
+PLy_commit(PyObject *self, PyObject *args)
+{
+	volatile MemoryContext oldcontext;
+	volatile ResourceOwner oldowner;
+
+	oldcontext = CurrentMemoryContext;
+	oldowner = CurrentResourceOwner;
+
+	PG_TRY();
+	{
+		PLyExecutionContext *exec_ctx = PLy_current_execution_context();
+
+		SPI_commit();
+		SPI_start_transaction();
+
+		/* was cleared at transaction end, reset pointer */
+		exec_ctx->scratch_ctx = NULL;
+	}
+	PG_CATCH();
+	{
+		ErrorData  *edata;
+		PLyExceptionEntry *entry;
+		PyObject   *exc;
+
+		/* Save error info */
+		MemoryContextSwitchTo(oldcontext);
+		edata = CopyErrorData();
+		FlushErrorState();
+
+		MemoryContextSwitchTo(oldcontext);
+		CurrentResourceOwner = oldowner;
+
+		/* Look up the correct exception */
+		entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
+							HASH_FIND, NULL);
+
+		/*
+		 * This could be a custom error code, if that's the case fallback to
+		 * SPIError
+		 */
+		exc = entry ? entry->exc : PLy_exc_spi_error;
+		/* Make Python raise the exception */
+		PLy_spi_exception_set(exc, edata);
+		FreeErrorData(edata);
+
+		return NULL;
+	}
+	PG_END_TRY();
+
+	Py_RETURN_NONE;
+}
+
+PyObject *
+PLy_rollback(PyObject *self, PyObject *args)
+{
+	volatile MemoryContext oldcontext;
+	volatile ResourceOwner oldowner;
+
+	oldcontext = CurrentMemoryContext;
+	oldowner = CurrentResourceOwner;
+
+	PG_TRY();
+	{
+		PLyExecutionContext *exec_ctx = PLy_current_execution_context();
+
+		SPI_rollback();
+		SPI_start_transaction();
+
+		/* was cleared at transaction end, reset pointer */
+		exec_ctx->scratch_ctx = NULL;
+	}
+	PG_CATCH();
+	{
+		ErrorData  *edata;
+		PLyExceptionEntry *entry;
+		PyObject   *exc;
+
+		/* Save error info */
+		MemoryContextSwitchTo(oldcontext);
+		edata = CopyErrorData();
+		FlushErrorState();
+
+		MemoryContextSwitchTo(oldcontext);
+		CurrentResourceOwner = oldowner;
+
+		/* Look up the correct exception */
+		entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
+							HASH_FIND, NULL);
+
+		/*
+		 * This could be a custom error code, if that's the case fallback to
+		 * SPIError
+		 */
+		exc = entry ? entry->exc : PLy_exc_spi_error;
+		/* Make Python raise the exception */
+		PLy_spi_exception_set(exc, edata);
+		FreeErrorData(edata);
+
+		return NULL;
+	}
+	PG_END_TRY();
+
+	Py_RETURN_NONE;
+}
+
 /*
  * Utilities for running SPI functions in subtransactions.
  *
diff --git a/src/pl/plpython/plpy_spi.h b/src/pl/plpython/plpy_spi.h
index a5e2e60da7..98ccd21093 100644
--- a/src/pl/plpython/plpy_spi.h
+++ b/src/pl/plpython/plpy_spi.h
@@ -12,6 +12,9 @@ extern PyObject *PLy_spi_prepare(PyObject *self, PyObject *args);
 extern PyObject *PLy_spi_execute(PyObject *self, PyObject *args);
 extern PyObject *PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit);
 
+extern PyObject *PLy_commit(PyObject *self, PyObject *args);
+extern PyObject *PLy_rollback(PyObject *self, PyObject *args);
+
 typedef struct PLyExceptionEntry
 {
 	int			sqlstate;		/* hash key, must be first */
-- 
2.34.1