Refactor crosstab() to build and return a tuplestore instead of using

value-per-call mode.  This should be more efficient in normal usage,
but the real problem with the prior coding was that it returned with
a SPI call still active.  That could cause problems if execution was
interleaved with anything else that might use SPI.

contrib/tablefunc/tablefunc.c

 /*
- * $PostgreSQL: pgsql/contrib/tablefunc/tablefunc.c,v 1.56 2008/11/30 23:23:52 tgl Exp $ 
+ * $PostgreSQL: pgsql/contrib/tablefunc/tablefunc.c,v 1.57 2008/12/01 01:30:18 tgl Exp $ 
  *
  *
  * tablefunc
@@ -94,12 +94,6 @@ typedef struct
 	bool		use_carry;		/* use second generated value */
 }	normal_rand_fctx;
 
-typedef struct
-{
-	SPITupleTable *spi_tuptable;	/* sql results from user query */
-	char	   *lastrowid;		/* rowid of the last tuple sent */
-}	crosstab_fctx;
-
 #define xpfree(var_) \
 	do { \
 		if (var_ != NULL) \
@@ -356,30 +350,36 @@ PG_FUNCTION_INFO_V1(crosstab);
 Datum
 crosstab(PG_FUNCTION_ARGS)
 {
-	FuncCallContext *funcctx;
-	TupleDesc	ret_tupdesc;
+	char	   *sql = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	Tuplestorestate *tupstore;
+	TupleDesc	tupdesc;
 	int			call_cntr;
 	int			max_calls;
 	AttInMetadata *attinmeta;
-	SPITupleTable *spi_tuptable = NULL;
+	SPITupleTable *spi_tuptable;
 	TupleDesc	spi_tupdesc;
-	char	   *lastrowid = NULL;
-	crosstab_fctx *fctx;
+	bool		firstpass;
+	char	   *lastrowid;
 	int			i;
 	int			num_categories;
-	bool		firstpass = false;
+	MemoryContext per_query_ctx;
 	MemoryContext oldcontext;
-
-	/* stuff done only on the first call of the function */
-	if (SRF_IS_FIRSTCALL())
-	{
-		char	   *sql = text_to_cstring(PG_GETARG_TEXT_PP(0));
-		TupleDesc	tupdesc;
 	int			ret;
 	int			proc;
 
-		/* create a function context for cross-call persistence */
-		funcctx = SRF_FIRSTCALL_INIT();
+	/* check to see if caller supports us returning a tuplestore */
+	if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("set-valued function called in context that cannot accept a set")));
+	if (!(rsinfo->allowedModes & SFRM_Materialize))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("materialize mode required, but it is not " \
+						"allowed in this context")));
+
+	per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
 
 	/* Connect to SPI manager */
 	if ((ret = SPI_connect()) < 0)
@@ -390,9 +390,14 @@ crosstab(PG_FUNCTION_ARGS)
 	ret = SPI_execute(sql, true, 0);
 	proc = SPI_processed;
 
-		/* Check for qualifying tuples */
-		if ((ret == SPI_OK_SELECT) && (proc > 0))
+	/* If no qualifying tuples, fall out early */
+	if (ret != SPI_OK_SELECT || proc <= 0)
 	{
+		SPI_finish();
+		rsinfo->isDone = ExprEndResult;
+		PG_RETURN_NULL();
+	}
+
 	spi_tuptable = SPI_tuptable;
 	spi_tupdesc = spi_tuptable->tupdesc;
 
@@ -413,13 +418,6 @@ crosstab(PG_FUNCTION_ARGS)
 				 errmsg("invalid source data SQL statement"),
 				 errdetail("The provided SQL must return 3 "
 						   "columns: rowid, category, and values.")));
-		}
-		else
-		{
-			/* no qualifying tuples */
-			SPI_finish();
-			SRF_RETURN_DONE(funcctx);
-		}
 
 	/* get a tuple descriptor for our result type */
 	switch (get_call_result_type(fcinfo, NULL, &tupdesc))
@@ -451,75 +449,42 @@ crosstab(PG_FUNCTION_ARGS)
 						"incompatible")));
 
 	/*
-		 * switch to memory context appropriate for multiple function calls
+	 * switch to long-lived memory context
 	 */
-		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
+	oldcontext = MemoryContextSwitchTo(per_query_ctx);
 
-		/* make sure we have a persistent copy of the tupdesc */
+	/* make sure we have a persistent copy of the result tupdesc */
 	tupdesc = CreateTupleDescCopy(tupdesc);
 
+	/* initialize our tuplestore in long-lived context */
+	tupstore =
+		tuplestore_begin_heap(rsinfo->allowedModes & SFRM_Materialize_Random,
+							  false, work_mem);
+
+	MemoryContextSwitchTo(oldcontext);
+
 	/*
 	 * Generate attribute metadata needed later to produce tuples from raw
 	 * C strings
 	 */
 	attinmeta = TupleDescGetAttInMetadata(tupdesc);
-		funcctx->attinmeta = attinmeta;
 
-		/* allocate memory for user context */
-		fctx = (crosstab_fctx *) palloc(sizeof(crosstab_fctx));
+	/* total number of tuples to be examined */
+	max_calls = proc;
 
-		/*
-		 * Save spi data for use across calls
-		 */
-		fctx->spi_tuptable = spi_tuptable;
-		fctx->lastrowid = NULL;
-		funcctx->user_fctx = fctx;
-
-		/* total number of tuples to be returned */
-		funcctx->max_calls = proc;
+	/* the return tuple always must have 1 rowid + num_categories columns */
+	num_categories = tupdesc->natts - 1;
 
-		MemoryContextSwitchTo(oldcontext);
 	firstpass = true;
-	}
-
-	/* stuff done on every call of the function */
-	funcctx = SRF_PERCALL_SETUP();
-
-	/*
-	 * initialize per-call variables
-	 */
-	call_cntr = funcctx->call_cntr;
-	max_calls = funcctx->max_calls;
+	lastrowid = NULL;
 
-	/* user context info */
-	fctx = (crosstab_fctx *) funcctx->user_fctx;
-	lastrowid = fctx->lastrowid;
-	spi_tuptable = fctx->spi_tuptable;
-
-	/* the sql tuple */
-	spi_tupdesc = spi_tuptable->tupdesc;
-
-	/* attribute return type and return tuple description */
-	attinmeta = funcctx->attinmeta;
-	ret_tupdesc = attinmeta->tupdesc;
-
-	/* the return tuple always must have 1 rowid + num_categories columns */
-	num_categories = ret_tupdesc->natts - 1;
-
-	if (call_cntr < max_calls)	/* do when there is more left to send */
+	for (call_cntr = 0; call_cntr < max_calls; call_cntr++)
 	{
-		HeapTuple	tuple;
-		Datum		result;
-		char	  **values;
 		bool		skip_tuple = false;
+		char	  **values;
 
-		while (true)
-		{
-			/* allocate space */
-			values = (char **) palloc((1 + num_categories) * sizeof(char *));
-
-			/* and make sure it's clear */
-			memset(values, '\0', (1 + num_categories) * sizeof(char *));
+		/* allocate and zero space */
+		values = (char **) palloc0((1 + num_categories) * sizeof(char *));
 
 		/*
 		 * now loop through the sql results and assign each value in
@@ -528,7 +493,7 @@ crosstab(PG_FUNCTION_ARGS)
 		for (i = 0; i < num_categories; i++)
 		{
 			HeapTuple	spi_tuple;
-				char	   *rowid = NULL;
+			char	   *rowid;
 
 			/* see if we've gone too far already */
 			if (call_cntr >= max_calls)
@@ -554,13 +519,14 @@ crosstab(PG_FUNCTION_ARGS)
 				 */
 				if (!firstpass && xstreq(lastrowid, rowid))
 				{
+					xpfree(rowid);
 					skip_tuple = true;
 					break;
 				}
 			}
 
 			/*
-				 * If rowid hasn't changed on us, continue building the ouput
+			 * If rowid hasn't changed on us, continue building the output
 			 * tuple.
 			 */
 			if (xstreq(rowid, values[0]))
@@ -576,11 +542,12 @@ crosstab(PG_FUNCTION_ARGS)
 
 				/*
 				 * increment the counter since we consume a row for each
-					 * category, but not for last pass because the API will do
-					 * that for us
+				 * category, but not for last pass because the outer loop
+				 * will do that for us
 				 */
 				if (i < (num_categories - 1))
-						call_cntr = ++funcctx->call_cntr;
+					call_cntr++;
+				xpfree(rowid);
 			}
 			else
 			{
@@ -589,71 +556,48 @@ crosstab(PG_FUNCTION_ARGS)
 				 * to decrement the counter since this sql result row
 				 * doesn't belong to the current output tuple.
 				 */
-					call_cntr = --funcctx->call_cntr;
+				call_cntr--;
+				xpfree(rowid);
 				break;
 			}
-				xpfree(rowid);
 		}
 
-			/*
-			 * switch to memory context appropriate for multiple function
-			 * calls
-			 */
-			oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
-
-			xpfree(fctx->lastrowid);
-			xpstrdup(fctx->lastrowid, values[0]);
-			lastrowid = fctx->lastrowid;
-
-			MemoryContextSwitchTo(oldcontext);
-
 		if (!skip_tuple)
 		{
+			HeapTuple	tuple;
+
 			/* build the tuple */
 			tuple = BuildTupleFromCStrings(attinmeta, values);
 
-				/* make the tuple into a datum */
-				result = HeapTupleGetDatum(tuple);
+			/* switch to appropriate context while storing the tuple */
+			oldcontext = MemoryContextSwitchTo(per_query_ctx);
+			tuplestore_puttuple(tupstore, tuple);
+			MemoryContextSwitchTo(oldcontext);
+
+			heap_freetuple(tuple);
+		}
+
+		/* Remember current rowid */
+		xpfree(lastrowid);
+		xpstrdup(lastrowid, values[0]);
+		firstpass = false;
 
 		/* Clean up */
 		for (i = 0; i < num_categories + 1; i++)
 			if (values[i] != NULL)
-						xpfree(values[i]);
-				xpfree(values);
-
-				SRF_RETURN_NEXT(funcctx, result);
+				pfree(values[i]);
+		pfree(values);
 	}
-			else
-			{
-				/*
-				 * Skipping this tuple entirely, but we need to advance the
-				 * counter like the API would if we had returned one.
-				 */
-				call_cntr = ++funcctx->call_cntr;
 
-				/* we'll start over at the top */
-				xpfree(values);
+	/* let the caller know we're sending back a tuplestore */
+	rsinfo->returnMode = SFRM_Materialize;
+	rsinfo->setResult = tupstore;
+	rsinfo->setDesc = tupdesc;
 
-				/* see if we've gone too far already */
-				if (call_cntr >= max_calls)
-				{
-					/* release SPI related resources */
+	/* release SPI related resources (and return to caller's context) */
 	SPI_finish();
-					SRF_RETURN_DONE(funcctx);
-				}
 
-				/* need to reset this before the next tuple is started */
-				skip_tuple = false;
-			}
-		}
-	}
-	else
-		/* do when there is no more left */
-	{
-		/* release SPI related resources */
-		SPI_finish();
-		SRF_RETURN_DONE(funcctx);
-	}
+	return (Datum) 0;
 }
 
 /*
@@ -1613,6 +1557,10 @@ compatCrosstabTupleDescs(TupleDesc ret_tupdesc, TupleDesc sql_tupdesc)
 	Form_pg_attribute sql_attr;
 	Oid			sql_atttypid;
 
+	if (ret_tupdesc->natts < 2 ||
+		sql_tupdesc->natts < 3)
+		return false;
+
 	/* check the rowid types match */
 	ret_atttypid = ret_tupdesc->attrs[0]->atttypid;
 	sql_atttypid = sql_tupdesc->attrs[0]->atttypid;