[BABEL] Fix to prevent SPI_finish crash on partially-initialized SPI connection after OOM (#725) (#732)

SPI_connect_ext() increments _SPI_connected before allocating memory contexts (procCxt, execCxt). If AllocSetContextCreate fails with OOM between the counter increment and context allocation, the SPI stack is left in an inconsistent state as the counter says a connection exists but the memory contexts are NULL.

In vanilla PostgreSQL this is harmless because SPI_finish is never called on error paths — AtEOXact_SPI safely resets the counter during transaction abort. However, Babelfish calls SPI_finish explicitly in error handling and cleanup paths to emulate SQL Server's error handling semantics. This exposes the inconsistent state and causes a crash in MemoryContextDelete(NULL).

SPI_finish crashes on such entries because it unconditionally calls MemoryContextDelete on NULL pointers.

SPI_finish_safe checks whether the connection is fully initialized. If so, it uses SPI_finish. If not, it safely cleans up any partial allocation and pops the stack. 

Extension PR: babelfish-for-postgresql/babelfish_extensions#4629 

Task: BABEL-6349

Authored-by: Rucha Kulkarni ruchask@amazon.com

Author: RuchaSK1

src/backend/executor/spi.c

@@ -3457,3 +3457,52 @@ SPI_get_depth(void)
 {
 	return _SPI_connected;
 }
+
+/*
+ * Modified version of SPI_finish for Babelfish, which safely handles
+ * the case where the current SPI connection was not fully initialized.
+ * For fully initialized connections, delegates to SPI_finish.
+ */
+int
+SPI_finish_safe(void)
+{
+	/* Check if there is an active SPI connection */
+	if (_SPI_current == NULL)
+		return SPI_ERROR_UNCONNECTED;
+
+	/* Fully initialized */
+	if (_SPI_current->procCxt != NULL &&
+		_SPI_current->execCxt != NULL)
+	{
+		return SPI_finish();
+	}
+	
+	if (_SPI_current->internal_xact)
+    	elog(FATAL, "SPI_finish_safe: unexpected state - partially initialized SPI connection should not have internal_xact set");
+
+	/* Clean up any partially allocated context */
+	if (_SPI_current->execCxt)
+	{
+		MemoryContextDelete(_SPI_current->execCxt);
+		_SPI_current->execCxt = NULL;
+	}
+	if (_SPI_current->procCxt)
+	{
+		MemoryContextDelete(_SPI_current->procCxt);
+		_SPI_current->procCxt = NULL;
+	}
+
+	/* Restore outer API variables */
+	SPI_processed = _SPI_current->outer_processed;
+	SPI_tuptable = _SPI_current->outer_tuptable;
+	SPI_result = _SPI_current->outer_result;
+
+	/* Pop the stack */
+	_SPI_connected--;
+	if (_SPI_connected < 0)
+		_SPI_current = NULL;
+	else
+		_SPI_current = &(_SPI_stack[_SPI_connected]);
+
+	return SPI_OK_FINISH;
+}

src/include/executor/spi.h

@@ -113,6 +113,7 @@ extern PGDLLIMPORT int SPI_result;
 extern int	SPI_connect(void);
 extern int	SPI_connect_ext(int options);
 extern int	SPI_finish(void);
+extern int	SPI_finish_safe(void);
 extern int	SPI_execute(const char *src, bool read_only, long tcount);
 extern int	SPI_execute_extended(const char *src,
 								 const SPIExecuteOptions *options);