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

Per buildfarm results.

2 files changed, 36 insertions, 27 deletions

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index b95795a2b1..9cd3a27004 100644
--- a/contrib/dblink/dblink.c
+++ b/contrib/dblink/dblink.c
@@ -8,7 +8,7 @@
  * Darko Prenosil <Darko.Prenosil@finteh.hr>
  * Shridhar Daithankar <shridhar_daithankar@persistent.co.in>
  *
- * $PostgreSQL: pgsql/contrib/dblink/dblink.c,v 1.60.2.2 2008/01/03 21:28:18 tgl Exp $
+ * $PostgreSQL: pgsql/contrib/dblink/dblink.c,v 1.60.2.3 2008/11/30 23:24:08 tgl Exp $
  * Copyright (c) 2001-2006, PostgreSQL Global Development Group
  * ALL RIGHTS RESERVED;
  *
@@ -224,7 +224,6 @@ dblink_connect(PG_FUNCTION_ARGS)
 	char	   *connstr = NULL;
 	char	   *connname = NULL;
 	char	   *msg;
-	MemoryContext oldcontext;
 	PGconn	   *conn = NULL;
 	remoteConn *rconn = NULL;
 
@@ -238,17 +237,14 @@ dblink_connect(PG_FUNCTION_ARGS)
 	else if (PG_NARGS() == 1)
 		connstr = GET_STR(PG_GETARG_TEXT_P(0));
 
-	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
-
 	if (connname)
-		rconn = (remoteConn *) palloc(sizeof(remoteConn));
+		rconn = (remoteConn *) MemoryContextAlloc(TopMemoryContext,
+												  sizeof(remoteConn));
 
 	/* check password used if not superuser */
 	dblink_security_check(conn, rconn, connstr);
 	conn = PQconnectdb(connstr);
 
-	MemoryContextSwitchTo(oldcontext);
-
 	if (PQstatus(conn) == CONNECTION_BAD)
 	{
 		msg = pstrdup(PQerrorMessage(conn));
@@ -581,10 +577,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 &&
@@ -631,9 +627,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,
@@ -641,14 +634,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;
@@ -839,7 +843,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 ||
@@ -853,6 +860,7 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async, bool do_get)
 					DBLINK_RES_ERROR_AS_NOTICE("sql error");
 					if (freeconn)
 						PQfinish(conn);
+					MemoryContextSwitchTo(oldcontext);
 					SRF_RETURN_DONE(funcctx);
 				}
 			}
@@ -923,6 +931,7 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async, bool do_get)
 			{
 				if (res)
 					PQclear(res);
+				MemoryContextSwitchTo(oldcontext);
 				SRF_RETURN_DONE(funcctx);
 			}
 
@@ -1294,8 +1303,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 0be013b908..26886b2af9 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
 		 */