PL/Python: Refactor subtransaction handling

Lots of repetitive code was moved into new functions
PLy_spi_subtransaction_{begin,commit,abort}.

Jan Urbański

diff --git a/src/pl/plpython/plpython.c b/src/pl/plpython/plpython.c
index 29e0ac7c45466fbf17bb8cbee6327a1c66d964a4..dce8ff247b66b60458dcc2b902dd7c235022ab53 100644 (file)
--- a/src/pl/plpython/plpython.c
+++ b/src/pl/plpython/plpython.c
@@ -2964,6 +2964,12 @@ static int       PLy_result_ass_item(PyObject *, Py_ssize_t, PyObject *);
 static int     PLy_result_ass_slice(PyObject *, Py_ssize_t, Py_ssize_t, PyObject *);
 
 
+/* handling of SPI operations inside subtransactions */
+static void PLy_spi_subtransaction_begin(MemoryContext oldcontext, ResourceOwner oldowner);
+static void PLy_spi_subtransaction_commit(MemoryContext oldcontext, ResourceOwner oldowner);
+static void PLy_spi_subtransaction_abort(MemoryContext oldcontext, ResourceOwner oldowner);
+
+/* SPI operations */
 static PyObject *PLy_spi_prepare(PyObject *, PyObject *);
 static PyObject *PLy_spi_execute(PyObject *, PyObject *);
 static PyObject *PLy_spi_execute_query(char *query, long limit);
@@ -3390,6 +3396,90 @@ PLy_result_ass_slice(PyObject *arg, Py_ssize_t lidx, Py_ssize_t hidx, PyObject *
        return rv;
 }
 
+/*
+ * Utilities for running SPI functions in subtransactions.
+ *
+ * Usage:
+ *
+ *  MemoryContext oldcontext = CurrentMemoryContext;
+ *  ResourceOwner oldowner = CurrentResourceOwner;
+ *
+ *  PLy_spi_subtransaction_begin(oldcontext, oldowner);
+ *  PG_TRY();
+ *  {
+ *      <call SPI functions>
+ *      PLy_spi_subtransaction_commit(oldcontext, oldowner);
+ *  }
+ *  PG_CATCH();
+ *  {
+ *      <do cleanup>
+ *      PLy_spi_subtransaction_abort(oldcontext, oldowner);
+ *      return NULL;
+ *  }
+ *  PG_END_TRY();
+ *
+ * These utilities take care of restoring connection to the SPI manager and
+ * setting a Python exception in case of an abort.
+ */
+static void
+PLy_spi_subtransaction_begin(MemoryContext oldcontext, ResourceOwner oldowner)
+{
+       BeginInternalSubTransaction(NULL);
+       /* Want to run inside function's memory context */
+       MemoryContextSwitchTo(oldcontext);
+}
+
+static void
+PLy_spi_subtransaction_commit(MemoryContext oldcontext, ResourceOwner oldowner)
+{
+       /* Commit the inner transaction, return to outer xact context */
+       ReleaseCurrentSubTransaction();
+       MemoryContextSwitchTo(oldcontext);
+       CurrentResourceOwner = oldowner;
+
+       /*
+        * AtEOSubXact_SPI() should not have popped any SPI context, but just
+        * in case it did, make sure we remain connected.
+        */
+       SPI_restore_connection();
+}
+
+static void
+PLy_spi_subtransaction_abort(MemoryContext oldcontext, ResourceOwner oldowner)
+{
+       ErrorData  *edata;
+       PLyExceptionEntry *entry;
+       PyObject   *exc;
+
+       /* Save error info */
+       MemoryContextSwitchTo(oldcontext);
+       edata = CopyErrorData();
+       FlushErrorState();
+
+       /* Abort the inner transaction */
+       RollbackAndReleaseCurrentSubTransaction();
+       MemoryContextSwitchTo(oldcontext);
+       CurrentResourceOwner = oldowner;
+
+       /*
+        * 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 up the correct exception */
+       entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
+                                               HASH_FIND, NULL);
+       /* We really should find it, but just in case have a fallback */
+       Assert(entry != NULL);
+       exc = entry ? entry->exc : PLy_exc_spi_error;
+       /* Make Python raise the exception */
+       PLy_spi_exception_set(exc, edata);
+       FreeErrorData(edata);
+}
+
+
 /* SPI interface */
 static PyObject *
 PLy_spi_prepare(PyObject *self, PyObject *args)
@@ -3425,8 +3515,7 @@ PLy_spi_prepare(PyObject *self, PyObject *args)
        oldcontext = CurrentMemoryContext;
        oldowner = CurrentResourceOwner;
 
-       BeginInternalSubTransaction(NULL);
-       MemoryContextSwitchTo(oldcontext);
+       PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
        PG_TRY();
        {
@@ -3504,50 +3593,14 @@ PLy_spi_prepare(PyObject *self, PyObject *args)
                if (SPI_keepplan(plan->plan))
                        elog(ERROR, "SPI_keepplan failed");
 
-               /* Commit the inner transaction, return to outer xact context */
-               ReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * AtEOSubXact_SPI() should not have popped any SPI context, but just
-                * in case it did, make sure we remain connected.
-                */
-               SPI_restore_connection();
+               PLy_spi_subtransaction_commit(oldcontext, oldowner);
        }
        PG_CATCH();
        {
-               ErrorData  *edata;
-               PLyExceptionEntry *entry;
-               PyObject   *exc;
-
-               /* Save error info */
-               MemoryContextSwitchTo(oldcontext);
-               edata = CopyErrorData();
-               FlushErrorState();
                Py_DECREF(plan);
                Py_XDECREF(optr);
 
-               /* Abort the inner transaction */
-               RollbackAndReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * 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 up the correct exception */
-               entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
-                                                       HASH_FIND, NULL);
-               /* We really should find it, but just in case have a fallback */
-               Assert(entry != NULL);
-               exc = entry ? entry->exc : PLy_exc_spi_error;
-               /* Make Python raise the exception */
-               PLy_spi_exception_set(exc, edata);
+               PLy_spi_subtransaction_abort(oldcontext, oldowner);
                return NULL;
        }
        PG_END_TRY();
@@ -3626,9 +3679,7 @@ PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
        oldcontext = CurrentMemoryContext;
        oldowner = CurrentResourceOwner;
 
-       BeginInternalSubTransaction(NULL);
-       /* Want to run inside function's memory context */
-       MemoryContextSwitchTo(oldcontext);
+       PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
        PG_TRY();
        {
@@ -3683,28 +3734,11 @@ PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
                if (nargs > 0)
                        pfree(nulls);
 
-               /* Commit the inner transaction, return to outer xact context */
-               ReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * AtEOSubXact_SPI() should not have popped any SPI context, but just
-                * in case it did, make sure we remain connected.
-                */
-               SPI_restore_connection();
+               PLy_spi_subtransaction_commit(oldcontext, oldowner);
        }
        PG_CATCH();
        {
                int                     k;
-               ErrorData  *edata;
-               PLyExceptionEntry *entry;
-               PyObject   *exc;
-
-               /* Save error info */
-               MemoryContextSwitchTo(oldcontext);
-               edata = CopyErrorData();
-               FlushErrorState();
 
                /*
                 * cleanup plan->values array
@@ -3719,26 +3753,7 @@ PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
                        }
                }
 
-               /* Abort the inner transaction */
-               RollbackAndReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * 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 up the correct exception */
-               entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
-                                                       HASH_FIND, NULL);
-               /* We really should find it, but just in case have a fallback */
-               Assert(entry != NULL);
-               exc = entry ? entry->exc : PLy_exc_spi_error;
-               /* Make Python raise the exception */
-               PLy_spi_exception_set(exc, edata);
+               PLy_spi_subtransaction_abort(oldcontext, oldowner);
                return NULL;
        }
        PG_END_TRY();
@@ -3775,9 +3790,7 @@ PLy_spi_execute_query(char *query, long limit)
        oldcontext = CurrentMemoryContext;
        oldowner = CurrentResourceOwner;
 
-       BeginInternalSubTransaction(NULL);
-       /* Want to run inside function's memory context */
-       MemoryContextSwitchTo(oldcontext);
+       PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
        PG_TRY();
        {
@@ -3785,48 +3798,11 @@ PLy_spi_execute_query(char *query, long limit)
                rv = SPI_execute(query, PLy_curr_procedure->fn_readonly, limit);
                ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv);
 
-               /* Commit the inner transaction, return to outer xact context */
-               ReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * AtEOSubXact_SPI() should not have popped any SPI context, but just
-                * in case it did, make sure we remain connected.
-                */
-               SPI_restore_connection();
+               PLy_spi_subtransaction_commit(oldcontext, oldowner);
        }
        PG_CATCH();
        {
-               ErrorData  *edata;
-               PLyExceptionEntry *entry;
-               PyObject   *exc;
-
-               /* Save error info */
-               MemoryContextSwitchTo(oldcontext);
-               edata = CopyErrorData();
-               FlushErrorState();
-
-               /* Abort the inner transaction */
-               RollbackAndReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * 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 up the correct exception */
-               entry = hash_search(PLy_spi_exceptions, &edata->sqlerrcode,
-                                                       HASH_FIND, NULL);
-               /* We really should find it, but just in case have a fallback */
-               Assert(entry != NULL);
-               exc = entry ? entry->exc : PLy_exc_spi_error;
-               /* Make Python raise the exception */
-               PLy_spi_exception_set(exc, edata);
+               PLy_spi_subtransaction_abort(oldcontext, oldowner);
                return NULL;
        }
        PG_END_TRY();
@@ -3944,8 +3920,7 @@ PLy_cursor_query(const char *query)
        oldcontext = CurrentMemoryContext;
        oldowner = CurrentResourceOwner;
 
-       BeginInternalSubTransaction(NULL);
-       MemoryContextSwitchTo(oldcontext);
+       PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
        PG_TRY();
        {
@@ -3969,50 +3944,11 @@ PLy_cursor_query(const char *query)
 
                cursor->portalname = PLy_strdup(portal->name);
 
-               /* Commit the inner transaction, return to outer xact context */
-               ReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * AtEOSubXact_SPI() should not have popped any SPI context, but just
-                * in case it did, make sure we remain connected.
-                */
-               SPI_restore_connection();
+               PLy_spi_subtransaction_commit(oldcontext, oldowner);
        }
        PG_CATCH();
        {
-               ErrorData  *edata;
-               PLyExceptionEntry *entry;
-               PyObject   *exc;
-
-               /* Save error info */
-               MemoryContextSwitchTo(oldcontext);
-               edata = CopyErrorData();
-               FlushErrorState();
-
-               /* Abort the inner transaction */
-               RollbackAndReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               Py_DECREF(cursor);
-
-               /*
-                * 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 up the correct exception */
-               entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
-                                                       HASH_FIND, NULL);
-               /* We really should find it, but just in case have a fallback */
-               Assert(entry != NULL);
-               exc = entry ? entry->exc : PLy_exc_spi_error;
-               /* Make Python raise the exception */
-               PLy_spi_exception_set(exc, edata);
+               PLy_spi_subtransaction_abort(oldcontext, oldowner);
                return NULL;
        }
        PG_END_TRY();
@@ -4072,8 +4008,7 @@ PLy_cursor_plan(PyObject *ob, PyObject *args)
        oldcontext = CurrentMemoryContext;
        oldowner = CurrentResourceOwner;
 
-       BeginInternalSubTransaction(NULL);
-       MemoryContextSwitchTo(oldcontext);
+       PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
        PG_TRY();
        {
@@ -4130,28 +4065,11 @@ PLy_cursor_plan(PyObject *ob, PyObject *args)
 
                cursor->portalname = PLy_strdup(portal->name);
 
-               /* Commit the inner transaction, return to outer xact context */
-               ReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * AtEOSubXact_SPI() should not have popped any SPI context, but just
-                * in case it did, make sure we remain connected.
-                */
-               SPI_restore_connection();
+               PLy_spi_subtransaction_commit(oldcontext, oldowner);
        }
        PG_CATCH();
        {
                int                     k;
-               ErrorData  *edata;
-               PLyExceptionEntry *entry;
-               PyObject   *exc;
-
-               /* Save error info */
-               MemoryContextSwitchTo(oldcontext);
-               edata = CopyErrorData();
-               FlushErrorState();
 
                /* cleanup plan->values array */
                for (k = 0; k < nargs; k++)
@@ -4164,28 +4082,9 @@ PLy_cursor_plan(PyObject *ob, PyObject *args)
                        }
                }
 
-               /* Abort the inner transaction */
-               RollbackAndReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
                Py_DECREF(cursor);
 
-               /*
-                * 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 up the correct exception */
-               entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
-                                                       HASH_FIND, NULL);
-               /* We really should find it, but just in case have a fallback */
-               Assert(entry != NULL);
-               exc = entry ? entry->exc : PLy_exc_spi_error;
-               /* Make Python raise the exception */
-               PLy_spi_exception_set(exc, edata);
+               PLy_spi_subtransaction_abort(oldcontext, oldowner);
                return NULL;
        }
        PG_END_TRY();
@@ -4255,8 +4154,7 @@ PLy_cursor_iternext(PyObject *self)
        oldcontext = CurrentMemoryContext;
        oldowner = CurrentResourceOwner;
 
-       BeginInternalSubTransaction(NULL);
-       MemoryContextSwitchTo(oldcontext);
+       PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
        PG_TRY();
        {
@@ -4277,50 +4175,13 @@ PLy_cursor_iternext(PyObject *self)
 
                SPI_freetuptable(SPI_tuptable);
 
-               /* Commit the inner transaction, return to outer xact context */
-               ReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * AtEOSubXact_SPI() should not have popped any SPI context, but just
-                * in case it did, make sure we remain connected.
-                */
-               SPI_restore_connection();
+               PLy_spi_subtransaction_commit(oldcontext, oldowner);
        }
        PG_CATCH();
        {
-               ErrorData  *edata;
-               PLyExceptionEntry *entry;
-               PyObject   *exc;
-
-               /* Save error info */
-               MemoryContextSwitchTo(oldcontext);
-               edata = CopyErrorData();
-               FlushErrorState();
-
-               /* Abort the inner transaction */
-               RollbackAndReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
                SPI_freetuptable(SPI_tuptable);
 
-               /*
-                * 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 up the correct exception */
-               entry = hash_search(PLy_spi_exceptions, &edata->sqlerrcode,
-                                                       HASH_FIND, NULL);
-               /* We really should find it, but just in case have a fallback */
-               Assert(entry != NULL);
-               exc = entry ? entry->exc : PLy_exc_spi_error;
-               /* Make Python raise the exception */
-               PLy_spi_exception_set(exc, edata);
+               PLy_spi_subtransaction_abort(oldcontext, oldowner);
                return NULL;
        }
        PG_END_TRY();
@@ -4364,8 +4225,7 @@ PLy_cursor_fetch(PyObject *self, PyObject *args)
        oldcontext = CurrentMemoryContext;
        oldowner = CurrentResourceOwner;
 
-       BeginInternalSubTransaction(NULL);
-       MemoryContextSwitchTo(oldcontext);
+       PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
        PG_TRY();
        {
@@ -4398,50 +4258,13 @@ PLy_cursor_fetch(PyObject *self, PyObject *args)
 
                SPI_freetuptable(SPI_tuptable);
 
-               /* Commit the inner transaction, return to outer xact context */
-               ReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
-               /*
-                * AtEOSubXact_SPI() should not have popped any SPI context, but just
-                * in case it did, make sure we remain connected.
-                */
-               SPI_restore_connection();
+               PLy_spi_subtransaction_commit(oldcontext, oldowner);
        }
        PG_CATCH();
        {
-               ErrorData  *edata;
-               PLyExceptionEntry *entry;
-               PyObject   *exc;
-
-               /* Save error info */
-               MemoryContextSwitchTo(oldcontext);
-               edata = CopyErrorData();
-               FlushErrorState();
-
-               /* Abort the inner transaction */
-               RollbackAndReleaseCurrentSubTransaction();
-               MemoryContextSwitchTo(oldcontext);
-               CurrentResourceOwner = oldowner;
-
                SPI_freetuptable(SPI_tuptable);
 
-               /*
-                * 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 up the correct exception */
-               entry = hash_search(PLy_spi_exceptions, &edata->sqlerrcode,
-                                                       HASH_FIND, NULL);
-               /* We really should find it, but just in case have a fallback */
-               Assert(entry != NULL);
-               exc = entry ? entry->exc : PLy_exc_spi_error;
-               /* Make Python raise the exception */
-               PLy_spi_exception_set(exc, edata);
+               PLy_spi_subtransaction_abort(oldcontext, oldowner);
                return NULL;
        }
        PG_END_TRY();