Fix dblink and tablefunc to not return with the wrong CurrentMemoryContext.

Per buildfarm results.

2 files changed, 35 insertions, 26 deletions

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index 7a9ddcdb5b..f13b7bc846 100644
--- a/contrib/dblink/dblink.c
+++ b/contrib/dblink/dblink.c
@@ -213,7 +213,6 @@ dblink_connect(PG_FUNCTION_ARGS)
 	char	   *connstr = NULL;
 	char	   *connname = NULL;
 	char	   *msg;
-	MemoryContext oldcontext;
 	PGconn	   *conn = NULL;
 	remoteConn *rconn = NULL;
 
@@ -227,17 +226,14 @@ dblink_connect(PG_FUNCTION_ARGS)
 	else if (PG_NARGS() == 1)
 		connstr = text_to_cstring(PG_GETARG_TEXT_PP(0));
 
-	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
-
 	if (connname)
-		rconn = (remoteConn *) palloc(sizeof(remoteConn));
+		rconn = (remoteConn *) MemoryContextAlloc(TopMemoryContext,
+												  sizeof(remoteConn));
 
 	/* check password in connection string if not superuser */
 	dblink_connstr_check(connstr);
 	conn = PQconnectdb(connstr);
 
-	MemoryContextSwitchTo(oldcontext);
-
 	if (PQstatus(conn) == CONNECTION_BAD)
 	{
 		msg = pstrdup(PQerrorMessage(conn));
@@ -562,10 +558,10 @@ dblink_fetch(PG_FUNCTION_ARGS)
 		funcctx = SRF_FIRSTCALL_INIT();
 
 		/*
-		 * switch to memory context appropriate for multiple function calls
+		 * Try to execute the query.  Note that since libpq uses malloc,
+		 * the PGresult will be long-lived even though we are still in
+		 * a short-lived memory context.
 		 */
-		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
-
 		res = PQexec(conn, buf.data);
 		if (!res ||
 			(PQresultStatus(res) != PGRES_COMMAND_OK &&
@@ -607,9 +603,6 @@ dblink_fetch(PG_FUNCTION_ARGS)
 				break;
 		}
 
-		/* make sure we have a persistent copy of the tupdesc */
-		tupdesc = CreateTupleDescCopy(tupdesc);
-
 		/* check result and tuple descriptor have the same number of columns */
 		if (PQnfields(res) != tupdesc->natts)
 			ereport(ERROR,
@@ -617,14 +610,25 @@ dblink_fetch(PG_FUNCTION_ARGS)
 					 errmsg("remote query result rowtype does not match "
 							"the specified FROM clause rowtype")));
 
-		/* fast track when no results */
+		/*
+		 * fast track when no results.  We could exit earlier, but then
+		 * we'd not report error if the result tuple type is wrong.
+		 */
 		if (funcctx->max_calls < 1)
 		{
-			if (res)
-				PQclear(res);
+			PQclear(res);
 			SRF_RETURN_DONE(funcctx);
 		}
 
+		/*
+		 * switch to memory context appropriate for multiple function calls,
+		 * so we can make long-lived copy of tupdesc etc
+		 */
+		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
+
+		/* make sure we have a persistent copy of the tupdesc */
+		tupdesc = CreateTupleDescCopy(tupdesc);
+
 		/* store needed metadata for subsequent calls */
 		attinmeta = TupleDescGetAttInMetadata(tupdesc);
 		funcctx->attinmeta = attinmeta;
@@ -815,7 +819,10 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async, bool do_get)
 				res = PQgetResult(conn);
 				/* NULL means we're all done with the async results */
 				if (!res)
+				{
+					MemoryContextSwitchTo(oldcontext);
 					SRF_RETURN_DONE(funcctx);
+				}
 			}
 
 			if (!res ||
@@ -825,6 +832,7 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async, bool do_get)
 				dblink_res_error(conname, res, "could not execute query", fail);
 				if (freeconn)
 					PQfinish(conn);
+				MemoryContextSwitchTo(oldcontext);
 				SRF_RETURN_DONE(funcctx);
 			}
 
@@ -894,6 +902,7 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async, bool do_get)
 			{
 				if (res)
 					PQclear(res);
+				MemoryContextSwitchTo(oldcontext);
 				SRF_RETURN_DONE(funcctx);
 			}
 
@@ -1261,8 +1270,11 @@ dblink_get_pkey(PG_FUNCTION_ARGS)
 			funcctx->user_fctx = results;
 		}
 		else
+		{
 			/* fast track when no results */
+			MemoryContextSwitchTo(oldcontext);
 			SRF_RETURN_DONE(funcctx);
+		}
 
 		MemoryContextSwitchTo(oldcontext);
 	}
diff --git a/contrib/tablefunc/tablefunc.c b/contrib/tablefunc/tablefunc.c
index 42e98f7b52..f05a615bf3 100644
--- a/contrib/tablefunc/tablefunc.c
+++ b/contrib/tablefunc/tablefunc.c
@@ -381,11 +381,6 @@ crosstab(PG_FUNCTION_ARGS)
 		/* create a function context for cross-call persistence */
 		funcctx = SRF_FIRSTCALL_INIT();
 
-		/*
-		 * switch to memory context appropriate for multiple function calls
-		 */
-		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
-
 		/* Connect to SPI manager */
 		if ((ret = SPI_connect()) < 0)
 			/* internal error */
@@ -426,9 +421,6 @@ crosstab(PG_FUNCTION_ARGS)
 			SRF_RETURN_DONE(funcctx);
 		}
 
-		/* SPI switches context on us, so reset it */
-		MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
-
 		/* get a tuple descriptor for our result type */
 		switch (get_call_result_type(fcinfo, NULL, &tupdesc))
 		{
@@ -448,9 +440,6 @@ crosstab(PG_FUNCTION_ARGS)
 				break;
 		}
 
-		/* make sure we have a persistent copy of the tupdesc */
-		tupdesc = CreateTupleDescCopy(tupdesc);
-
 		/*
 		 * Check that return tupdesc is compatible with the data we got from
 		 * SPI, at least based on number and type of attributes
@@ -462,6 +451,14 @@ crosstab(PG_FUNCTION_ARGS)
 							"incompatible")));
 
 		/*
+		 * switch to memory context appropriate for multiple function calls
+		 */
+		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
+
+		/* make sure we have a persistent copy of the tupdesc */
+		tupdesc = CreateTupleDescCopy(tupdesc);
+
+		/*
 		 * Generate attribute metadata needed later to produce tuples from raw
 		 * C strings
 		 */