diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index 08f6f67a15..56249d688a 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -90,6 +90,7 @@ int
 SPI_connect_ext(int options)
 {
 	int			newdepth;
+	bool		atomic = (options & SPI_OPT_NONATOMIC) == 0;
 
 	/* Enlarge stack if necessary */
 	if (_SPI_stack == NULL)
@@ -98,7 +99,8 @@ SPI_connect_ext(int options)
 			elog(ERROR, "SPI stack corrupted");
 		newdepth = 16;
 		_SPI_stack = (_SPI_connection *)
-			MemoryContextAlloc(TopMemoryContext,
+			MemoryContextAlloc(atomic ?
+							   TopTransactionContext : TopMemoryContext,
 							   newdepth * sizeof(_SPI_connection));
 		_SPI_stack_depth = newdepth;
 	}
@@ -130,7 +132,7 @@ SPI_connect_ext(int options)
 	_SPI_current->execCxt = NULL;
 	_SPI_current->connectSubid = GetCurrentSubTransactionId();
 	_SPI_current->queryEnv = NULL;
-	_SPI_current->atomic = (options & SPI_OPT_NONATOMIC ? false : true);
+	_SPI_current->atomic = atomic;
 	_SPI_current->internal_xact = false;
 
 	/*
@@ -283,6 +285,12 @@ AtEOXact_SPI(bool isCommit)
 				 errmsg("transaction left non-empty SPI stack"),
 				 errhint("Check for missing \"SPI_finish\" calls.")));
 
+	if (_SPI_current != NULL && !_SPI_current->atomic && _SPI_stack != NULL)
+		ereport(WARNING,
+				(errcode(ERRCODE_WARNING),
+				 errmsg("non-atomic transaction left non-empty SPI stack"),
+				 errhint("Check after non-atomic \"SPI_connect_ext\" calls.")));
+
 	_SPI_current = _SPI_stack = NULL;
 	_SPI_stack_depth = 0;
 	_SPI_connected = -1;