PL/Python: Refactor subtransaction handling

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

Jan Urbański

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);
+	PLy_spi_subtransaction_begin(oldcontext, oldowner);
-	MemoryContextSwitchTo(oldcontext);
 
 	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 */
+		PLy_spi_subtransaction_commit(oldcontext, oldowner);
-		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();
 	}
 	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 */
+		PLy_spi_subtransaction_abort(oldcontext, oldowner);
-		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);
 		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);
+	PLy_spi_subtransaction_begin(oldcontext, oldowner);
-	/* Want to run inside function's memory context */
-	MemoryContextSwitchTo(oldcontext);
 
 	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 */
+		PLy_spi_subtransaction_commit(oldcontext, oldowner);
-		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();
 	}
 	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 */
+		PLy_spi_subtransaction_abort(oldcontext, oldowner);
-		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);
 		return NULL;
 	}
 	PG_END_TRY();
@@ -3775,9 +3790,7 @@ PLy_spi_execute_query(char *query, long limit)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	BeginInternalSubTransaction(NULL);
+	PLy_spi_subtransaction_begin(oldcontext, oldowner);
-	/* Want to run inside function's memory context */
-	MemoryContextSwitchTo(oldcontext);
 
 	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 */
+		PLy_spi_subtransaction_commit(oldcontext, oldowner);
-		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();
 	}
 	PG_CATCH();
 	{
-		ErrorData  *edata;
+		PLy_spi_subtransaction_abort(oldcontext, oldowner);
-		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);
 		return NULL;
 	}
 	PG_END_TRY();
@@ -3944,8 +3920,7 @@ PLy_cursor_query(const char *query)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	BeginInternalSubTransaction(NULL);
+	PLy_spi_subtransaction_begin(oldcontext, oldowner);
-	MemoryContextSwitchTo(oldcontext);
 
 	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 */
+		PLy_spi_subtransaction_commit(oldcontext, oldowner);
-		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();
 	}
 	PG_CATCH();
 	{
-		ErrorData  *edata;
+		PLy_spi_subtransaction_abort(oldcontext, oldowner);
-		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);
 		return NULL;
 	}
 	PG_END_TRY();
@@ -4072,8 +4008,7 @@ PLy_cursor_plan(PyObject *ob, PyObject *args)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	BeginInternalSubTransaction(NULL);
+	PLy_spi_subtransaction_begin(oldcontext, oldowner);
-	MemoryContextSwitchTo(oldcontext);
 
 	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 */
+		PLy_spi_subtransaction_commit(oldcontext, oldowner);
-		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();
 	}
 	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);
 
-		/*
+		PLy_spi_subtransaction_abort(oldcontext, 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);
 		return NULL;
 	}
 	PG_END_TRY();
@@ -4255,8 +4154,7 @@ PLy_cursor_iternext(PyObject *self)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	BeginInternalSubTransaction(NULL);
+	PLy_spi_subtransaction_begin(oldcontext, oldowner);
-	MemoryContextSwitchTo(oldcontext);
 
 	PG_TRY();
 	{
@@ -4277,50 +4175,13 @@ PLy_cursor_iternext(PyObject *self)
 
 		SPI_freetuptable(SPI_tuptable);
 
-		/* Commit the inner transaction, return to outer xact context */
+		PLy_spi_subtransaction_commit(oldcontext, oldowner);
-		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();
 	}
 	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);
 
-		/*
+		PLy_spi_subtransaction_abort(oldcontext, 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);
 		return NULL;
 	}
 	PG_END_TRY();
@@ -4364,8 +4225,7 @@ PLy_cursor_fetch(PyObject *self, PyObject *args)
 	oldcontext = CurrentMemoryContext;
 	oldowner = CurrentResourceOwner;
 
-	BeginInternalSubTransaction(NULL);
+	PLy_spi_subtransaction_begin(oldcontext, oldowner);
-	MemoryContextSwitchTo(oldcontext);
 
 	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 */
+		PLy_spi_subtransaction_commit(oldcontext, oldowner);
-		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();
 	}
 	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);
 
-		/*
+		PLy_spi_subtransaction_abort(oldcontext, 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);
 		return NULL;
 	}
 	PG_END_TRY();