diff --git a/contrib/spi/refint.c b/contrib/spi/refint.c
index adf0490f85..05595c52c3 100644
--- a/contrib/spi/refint.c
+++ b/contrib/spi/refint.c
@@ -12,6 +12,7 @@
 #include "commands/trigger.h"
 #include "executor/spi.h"
 #include "utils/builtins.h"
+#include "utils/memutils.h"
 #include "utils/rel.h"
 
 PG_MODULE_MAGIC;
@@ -186,12 +187,13 @@ check_primary_key(PG_FUNCTION_ARGS)
 
 		/*
 		 * Remember that SPI_prepare places plan in current memory context -
-		 * so, we have to save plan in Top memory context for later use.
+		 * so, we have to save plan in TopMemoryContext for later use.
 		 */
 		if (SPI_keepplan(pplan))
 			/* internal error */
 			elog(ERROR, "check_primary_key: SPI_keepplan failed");
-		plan->splan = (SPIPlanPtr *) malloc(sizeof(SPIPlanPtr));
+		plan->splan = (SPIPlanPtr *) MemoryContextAlloc(TopMemoryContext,
+														sizeof(SPIPlanPtr));
 		*(plan->splan) = pplan;
 		plan->nplans = 1;
 	}
@@ -417,7 +419,8 @@ check_foreign_key(PG_FUNCTION_ARGS)
 		char		sql[8192];
 		char	  **args2 = args;
 
-		plan->splan = (SPIPlanPtr *) malloc(nrefs * sizeof(SPIPlanPtr));
+		plan->splan = (SPIPlanPtr *) MemoryContextAlloc(TopMemoryContext,
+														nrefs * sizeof(SPIPlanPtr));
 
 		for (r = 0; r < nrefs; r++)
 		{
@@ -614,6 +617,13 @@ find_plan(char *ident, EPlan **eplan, int *nplans)
 {
 	EPlan	   *newp;
 	int			i;
+	MemoryContext	oldcontext;
+
+	/*
+	 * All allocations done for the plans need to happen in a session-safe
+	 * context.
+	 */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
 
 	if (*nplans > 0)
 	{
@@ -623,20 +633,24 @@ find_plan(char *ident, EPlan **eplan, int *nplans)
 				break;
 		}
 		if (i != *nplans)
+		{
+			MemoryContextSwitchTo(oldcontext);
 			return (*eplan + i);
-		*eplan = (EPlan *) realloc(*eplan, (i + 1) * sizeof(EPlan));
+		}
+		*eplan = (EPlan *) repalloc(*eplan, (i + 1) * sizeof(EPlan));
 		newp = *eplan + i;
 	}
 	else
 	{
-		newp = *eplan = (EPlan *) malloc(sizeof(EPlan));
+		newp = *eplan = (EPlan *) palloc(sizeof(EPlan));
 		(*nplans) = i = 0;
 	}
 
-	newp->ident = strdup(ident);
+	newp->ident = pstrdup(ident);
 	newp->nplans = 0;
 	newp->splan = NULL;
 	(*nplans)++;
 
+	MemoryContextSwitchTo(oldcontext);
 	return newp;
 }