dblink_rowproc_2012-03-07v2.diff

*** a/contrib/dblink/dblink.c
--- b/contrib/dblink/dblink.c
***************
*** 63,73 **** typedef struct remoteConn
  	bool		newXactForCursor;		/* Opened a transaction for a cursor */
  } remoteConn;
  
  /*
   * Internal declarations
   */
  static Datum dblink_record_internal(FunctionCallInfo fcinfo, bool is_async);
- static void materializeResult(FunctionCallInfo fcinfo, PGresult *res);
  static remoteConn *getConnectionByName(const char *name);
  static HTAB *createConnHash(void);
  static void createNewConnection(const char *name, remoteConn *rconn);
--- 63,85 ----
  	bool		newXactForCursor;		/* Opened a transaction for a cursor */
  } remoteConn;
  
+ typedef struct storeInfo
+ {
+ 	Tuplestorestate *tuplestore;
+ 	int nattrs;
+ 	MemoryContext oldcontext;
+ 	AttInMetadata *attinmeta;
+ 	char** valbuf;
+ 	int *valbuflen;
+ 	char **cstrs;
+ 	bool error_occurred;
+ 	bool nummismatch;
+ } storeInfo;
+ 
  /*
   * Internal declarations
   */
  static Datum dblink_record_internal(FunctionCallInfo fcinfo, bool is_async);
  static remoteConn *getConnectionByName(const char *name);
  static HTAB *createConnHash(void);
  static void createNewConnection(const char *name, remoteConn *rconn);
***************
*** 90,95 **** static char *escape_param_str(const char *from);
--- 102,111 ----
  static void validate_pkattnums(Relation rel,
  				   int2vector *pkattnums_arg, int32 pknumatts_arg,
  				   int **pkattnums, int *pknumatts);
+ static void initStoreInfo(storeInfo *sinfo, FunctionCallInfo fcinfo);
+ static void finishStoreInfo(storeInfo *sinfo);
+ static int storeHandler(PGresult *res, PGrowValue *columns, void *param);
+ 
  
  /* Global */
  static remoteConn *pconn = NULL;
***************
*** 503,508 **** dblink_fetch(PG_FUNCTION_ARGS)
--- 519,525 ----
  	char	   *curname = NULL;
  	int			howmany = 0;
  	bool		fail = true;	/* default to backward compatible */
+ 	storeInfo   storeinfo;
  
  	DBLINK_INIT;
  
***************
*** 559,573 **** dblink_fetch(PG_FUNCTION_ARGS)
  	appendStringInfo(&buf, "FETCH %d FROM %s", howmany, curname);
  
  	/*
  	 * 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.
  	 */
! 	res = PQexec(conn, buf.data);
  	if (!res ||
  		(PQresultStatus(res) != PGRES_COMMAND_OK &&
  		 PQresultStatus(res) != PGRES_TUPLES_OK))
  	{
  		dblink_res_error(conname, res, "could not fetch from cursor", fail);
  		return (Datum) 0;
  	}
--- 576,626 ----
  	appendStringInfo(&buf, "FETCH %d FROM %s", howmany, curname);
  
  	/*
+ 	 * Result is stored into storeinfo.tuplestore instead of
+ 	 * res->result retuned by PQexec below
+ 	 */
+ 	initStoreInfo(&storeinfo, fcinfo);
+ 	PQsetRowProcessor(conn, storeHandler, &storeinfo);
+ 
+ 	/*
  	 * 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.
  	 */
! 	PG_TRY();
! 	{
! 		res = PQexec(conn, buf.data);
! 	}
! 	PG_CATCH();
! 	{
! 		ErrorData *edata;
! 
! 		finishStoreInfo(&storeinfo);
! 		edata = CopyErrorData();
! 		FlushErrorState();
! 
! 		/* Skip remaining results when storeHandler raises exception. */
! 		PQskipResult(conn, TRUE);
! 		ReThrowError(edata);
! 	}
! 	PG_END_TRY();
! 
! 	finishStoreInfo(&storeinfo);
! 
  	if (!res ||
  		(PQresultStatus(res) != PGRES_COMMAND_OK &&
  		 PQresultStatus(res) != PGRES_TUPLES_OK))
  	{
+ 		/* finishStoreInfo saves the fields referred to below. */
+ 		if (storeinfo.nummismatch)
+ 		{
+ 			/* This is only for backward compatibility */
+ 			ereport(ERROR,
+ 					(errcode(ERRCODE_DATATYPE_MISMATCH),
+ 					 errmsg("remote query result rowtype does not match "
+ 							"the specified FROM clause rowtype")));
+ 		}
+ 
  		dblink_res_error(conname, res, "could not fetch from cursor", fail);
  		return (Datum) 0;
  	}
***************
*** 579,586 **** dblink_fetch(PG_FUNCTION_ARGS)
  				(errcode(ERRCODE_INVALID_CURSOR_NAME),
  				 errmsg("cursor \"%s\" does not exist", curname)));
  	}
  
- 	materializeResult(fcinfo, res);
  	return (Datum) 0;
  }
  
--- 632,639 ----
  				(errcode(ERRCODE_INVALID_CURSOR_NAME),
  				 errmsg("cursor \"%s\" does not exist", curname)));
  	}
+ 	PQclear(res);
  
  	return (Datum) 0;
  }
  
***************
*** 640,645 **** dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
--- 693,699 ----
  	remoteConn *rconn = NULL;
  	bool		fail = true;	/* default to backward compatible */
  	bool		freeconn = false;
+ 	storeInfo   storeinfo;
  
  	/* check to see if caller supports us returning a tuplestore */
  	if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
***************
*** 660,665 **** dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
--- 714,720 ----
  		{
  			/* text,text,bool */
  			DBLINK_GET_CONN;
+ 			conname = text_to_cstring(PG_GETARG_TEXT_PP(0));
  			sql = text_to_cstring(PG_GETARG_TEXT_PP(1));
  			fail = PG_GETARG_BOOL(2);
  		}
***************
*** 675,680 **** dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
--- 730,736 ----
  			else
  			{
  				DBLINK_GET_CONN;
+ 				conname = text_to_cstring(PG_GETARG_TEXT_PP(0));
  				sql = text_to_cstring(PG_GETARG_TEXT_PP(1));
  			}
  		}
***************
*** 705,710 **** dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
--- 761,768 ----
  		else
  			/* shouldn't happen */
  			elog(ERROR, "wrong number of arguments");
+ 
+ 		conname = text_to_cstring(PG_GETARG_TEXT_PP(0));
  	}
  
  	if (!conn)
***************
*** 715,878 **** dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
  	rsinfo->setResult = NULL;
  	rsinfo->setDesc = NULL;
  
  	/* synchronous query, or async result retrieval */
! 	if (!is_async)
! 		res = PQexec(conn, sql);
! 	else
  	{
! 		res = PQgetResult(conn);
! 		/* NULL means we're all done with the async results */
! 		if (!res)
! 			return (Datum) 0;
  	}
  
! 	/* if needed, close the connection to the database and cleanup */
! 	if (freeconn)
! 		PQfinish(conn);
  
! 	if (!res ||
! 		(PQresultStatus(res) != PGRES_COMMAND_OK &&
! 		 PQresultStatus(res) != PGRES_TUPLES_OK))
  	{
! 		dblink_res_error(conname, res, "could not execute query", fail);
! 		return (Datum) 0;
  	}
  
- 	materializeResult(fcinfo, res);
  	return (Datum) 0;
  }
  
- /*
-  * Materialize the PGresult to return them as the function result.
-  * The res will be released in this function.
-  */
  static void
! materializeResult(FunctionCallInfo fcinfo, PGresult *res)
  {
  	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
  
! 	Assert(rsinfo->returnMode == SFRM_Materialize);
! 
! 	PG_TRY();
  	{
! 		TupleDesc	tupdesc;
! 		bool		is_sql_cmd = false;
! 		int			ntuples;
! 		int			nfields;
  
! 		if (PQresultStatus(res) == PGRES_COMMAND_OK)
! 		{
! 			is_sql_cmd = true;
  
! 			/*
! 			 * need a tuple descriptor representing one TEXT column to return
! 			 * the command status string as our result tuple
! 			 */
! 			tupdesc = CreateTemplateTupleDesc(1, false);
! 			TupleDescInitEntry(tupdesc, (AttrNumber) 1, "status",
! 							   TEXTOID, -1, 0);
! 			ntuples = 1;
! 			nfields = 1;
! 		}
! 		else
! 		{
! 			Assert(PQresultStatus(res) == PGRES_TUPLES_OK);
  
! 			is_sql_cmd = false;
  
! 			/* get a tuple descriptor for our result type */
! 			switch (get_call_result_type(fcinfo, NULL, &tupdesc))
! 			{
! 				case TYPEFUNC_COMPOSITE:
! 					/* success */
! 					break;
! 				case TYPEFUNC_RECORD:
! 					/* failed to determine actual type of RECORD */
! 					ereport(ERROR,
! 							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
! 						errmsg("function returning record called in context "
! 							   "that cannot accept type record")));
! 					break;
! 				default:
! 					/* result type isn't composite */
! 					elog(ERROR, "return type must be a row type");
! 					break;
! 			}
  
! 			/* make sure we have a persistent copy of the tupdesc */
! 			tupdesc = CreateTupleDescCopy(tupdesc);
! 			ntuples = PQntuples(res);
! 			nfields = PQnfields(res);
! 		}
  
! 		/*
! 		 * check result and tuple descriptor have the same number of columns
! 		 */
! 		if (nfields != tupdesc->natts)
! 			ereport(ERROR,
! 					(errcode(ERRCODE_DATATYPE_MISMATCH),
! 					 errmsg("remote query result rowtype does not match "
! 							"the specified FROM clause rowtype")));
  
! 		if (ntuples > 0)
  		{
! 			AttInMetadata *attinmeta;
! 			Tuplestorestate *tupstore;
! 			MemoryContext oldcontext;
! 			int			row;
! 			char	  **values;
! 
! 			attinmeta = TupleDescGetAttInMetadata(tupdesc);
! 
! 			oldcontext = MemoryContextSwitchTo(
! 									rsinfo->econtext->ecxt_per_query_memory);
! 			tupstore = tuplestore_begin_heap(true, false, work_mem);
! 			rsinfo->setResult = tupstore;
! 			rsinfo->setDesc = tupdesc;
! 			MemoryContextSwitchTo(oldcontext);
  
! 			values = (char **) palloc(nfields * sizeof(char *));
  
! 			/* put all tuples into the tuplestore */
! 			for (row = 0; row < ntuples; row++)
! 			{
! 				HeapTuple	tuple;
  
! 				if (!is_sql_cmd)
! 				{
! 					int			i;
  
! 					for (i = 0; i < nfields; i++)
! 					{
! 						if (PQgetisnull(res, row, i))
! 							values[i] = NULL;
! 						else
! 							values[i] = PQgetvalue(res, row, i);
! 					}
! 				}
! 				else
! 				{
! 					values[0] = PQcmdStatus(res);
! 				}
  
! 				/* build the tuple and put it into the tuplestore. */
! 				tuple = BuildTupleFromCStrings(attinmeta, values);
! 				tuplestore_puttuple(tupstore, tuple);
! 			}
  
! 			/* clean up and return the tuplestore */
! 			tuplestore_donestoring(tupstore);
! 		}
  
! 		PQclear(res);
  	}
! 	PG_CATCH();
  	{
! 		/* be sure to release the libpq result */
! 		PQclear(res);
! 		PG_RE_THROW();
  	}
! 	PG_END_TRY();
  }
  
  /*
--- 773,1023 ----
  	rsinfo->setResult = NULL;
  	rsinfo->setDesc = NULL;
  
+ 
+ 	/*
+ 	 * Result is stored into storeinfo.tuplestore instead of
+ 	 * res->result retuned by PQexec/PQgetResult below
+ 	 */
+ 	initStoreInfo(&storeinfo, fcinfo);
+ 	PQsetRowProcessor(conn, storeHandler, &storeinfo);
+ 
  	/* synchronous query, or async result retrieval */
! 	PG_TRY();
  	{
! 		if (!is_async)
! 			res = PQexec(conn, sql);
! 		else
! 			res = PQgetResult(conn);
  	}
+ 	PG_CATCH();
+ 	{
+ 		ErrorData *edata;
  
! 		finishStoreInfo(&storeinfo);
! 		edata = CopyErrorData();
! 		FlushErrorState();
  
! 		/* Skip remaining results when storeHandler raises exception. */
! 		PQskipResult(conn, TRUE);
! 		ReThrowError(edata);
! 	}
! 	PG_END_TRY();
! 
! 	finishStoreInfo(&storeinfo);
! 
! 	/* NULL res from async get means we're all done with the results */
! 	if (res || !is_async)
  	{
! 		if (freeconn)
! 			PQfinish(conn);
! 
! 		if (!res ||
! 			(PQresultStatus(res) != PGRES_COMMAND_OK &&
! 			 PQresultStatus(res) != PGRES_TUPLES_OK))
! 		{
! 			/* finishStoreInfo saves the fields referred to below. */
! 			if (storeinfo.nummismatch)
! 			{
! 				/* This is only for backward compatibility */
! 				ereport(ERROR,
! 						(errcode(ERRCODE_DATATYPE_MISMATCH),
! 						 errmsg("remote query result rowtype does not match "
! 								"the specified FROM clause rowtype")));
! 			}
! 
! 			dblink_res_error(conname, res, "could not execute query", fail);
! 			return (Datum) 0;
! 		}
  	}
+ 	PQclear(res);
  
  	return (Datum) 0;
  }
  
  static void
! initStoreInfo(storeInfo *sinfo, FunctionCallInfo fcinfo)
  {
  	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+ 	TupleDesc	tupdesc;
+ 	int i;
  
! 	switch (get_call_result_type(fcinfo, NULL, &tupdesc))
  	{
! 		case TYPEFUNC_COMPOSITE:
! 			/* success */
! 			break;
! 		case TYPEFUNC_RECORD:
! 			/* failed to determine actual type of RECORD */
! 			ereport(ERROR,
! 					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
! 					 errmsg("function returning record called in context "
! 							"that cannot accept type record")));
! 			break;
! 		default:
! 			/* result type isn't composite */
! 			elog(ERROR, "return type must be a row type");
! 			break;
! 	}
  
! 	sinfo->oldcontext = MemoryContextSwitchTo(
! 		rsinfo->econtext->ecxt_per_query_memory);
  
! 	/* make sure we have a persistent copy of the tupdesc */
! 	tupdesc = CreateTupleDescCopy(tupdesc);
  
! 	sinfo->error_occurred = FALSE;
! 	sinfo->nummismatch = FALSE;
! 	sinfo->attinmeta = TupleDescGetAttInMetadata(tupdesc);
! 	sinfo->nattrs = tupdesc->natts;
! 	sinfo->tuplestore = tuplestore_begin_heap(true, false, work_mem);
! 	sinfo->valbuf = NULL;
! 	sinfo->valbuflen = NULL;
  
! 	/* Preallocate memory of same size with c string array for values. */
! 	sinfo->valbuf = (char **)malloc(sinfo->nattrs * sizeof(char*));
! 	if (sinfo->valbuf)
! 		sinfo->valbuflen = (int *)malloc(sinfo->nattrs * sizeof(int));
! 	if (sinfo->valbuflen)
! 		sinfo->cstrs = (char **)malloc(sinfo->nattrs * sizeof(char*));
  
! 	if (sinfo->cstrs == NULL)
! 	{
! 		if (sinfo->valbuf)
! 			free(sinfo->valbuf);
! 		if (sinfo->valbuflen)
! 			free(sinfo->valbuflen);
  
! 		ereport(ERROR,
! 				(errcode(ERRCODE_OUT_OF_MEMORY),
! 				 errmsg("out of memory")));
! 	}
! 
! 	for (i = 0 ; i < sinfo->nattrs ; i++)
! 	{
! 		sinfo->valbuf[i] = NULL;
! 		sinfo->valbuflen[i] = -1;
! 	}
  
! 	rsinfo->setResult = sinfo->tuplestore;
! 	rsinfo->setDesc = tupdesc;
! }
! 
! static void
! finishStoreInfo(storeInfo *sinfo)
! {
! 	int i;
! 
! 	if (sinfo->valbuf)
! 	{
! 		for (i = 0 ; i < sinfo->nattrs ; i++)
  		{
! 			if (sinfo->valbuf[i])
! 				free(sinfo->valbuf[i]);
! 		}
! 		free(sinfo->valbuf);
! 		sinfo->valbuf = NULL;
! 	}
  
! 	if (sinfo->valbuflen)
! 	{
! 		free(sinfo->valbuflen);
! 		sinfo->valbuflen = NULL;
! 	}
  
! 	if (sinfo->cstrs)
! 	{
! 		free(sinfo->cstrs);
! 		sinfo->cstrs = NULL;
! 	}
  
! 	MemoryContextSwitchTo(sinfo->oldcontext);
! }
  
! /* Prototype of this function is PQrowProcessor */
! static int
! storeHandler(PGresult *res, PGrowValue *columns, void *param)
! {
! 	storeInfo *sinfo = (storeInfo *)param;
! 	HeapTuple  tuple;
! 	int        fields = PQnfields(res);
! 	int        i;
! 	char      **cstrs = sinfo->cstrs;
  
! 	if (sinfo->error_occurred)
! 		return -1;
  
! 	if (sinfo->nattrs != fields)
! 	{
! 		sinfo->error_occurred = TRUE;
! 		sinfo->nummismatch = TRUE;
! 		finishStoreInfo(sinfo);
  
! 		/* This error will be processed in dblink_record_internal() */
! 		return -1;
  	}
! 
! 	/*
! 	 * value input functions assumes that the input string is
! 	 * terminated by zero. We should make the values to be so.
! 	 */
! 	for(i = 0 ; i < fields ; i++)
  	{
! 		int len = columns[i].len;
! 		if (len < 0)
! 			cstrs[i] = NULL;
! 		else
! 		{
! 			char *tmp = sinfo->valbuf[i];
! 			int tmplen = sinfo->valbuflen[i];
! 
! 			/*
! 			 * Divide calls to malloc and realloc so that things will
! 			 * go fine even on the systems of which realloc() does not
! 			 * accept NULL as old memory block.
! 			 *
! 			 * Also try to (re)allocate in bigger steps to
! 			 * avoid flood of allocations on weird data.
! 			 */
! 			if (tmp == NULL)
! 			{
! 				tmplen = len + 1;
! 				if (tmplen < 64)
! 					tmplen = 64;
! 				tmp = (char *)malloc(tmplen);
! 			}
! 			else if (tmplen < len + 1)
! 			{
! 				if (len + 1 > tmplen * 2)
! 					tmplen = len + 1;
! 				else
! 					tmplen = tmplen * 2;
! 				tmp = (char *)realloc(tmp, tmplen);
! 			}
! 
! 			/*
! 			 * sinfo->valbuf[n] will be freed in finishStoreInfo()
! 			 * when realloc returns NULL.
! 			 */
! 			if (tmp == NULL)
! 				return -1;  /* Inform out of memory to the caller */
! 
! 			sinfo->valbuf[i] = tmp;
! 			sinfo->valbuflen[i] = tmplen;
! 
! 			cstrs[i] = sinfo->valbuf[i];
! 			memcpy(cstrs[i], columns[i].value, len);
! 			cstrs[i][len] = '\0';
! 		}
  	}
! 
! 	/*
! 	 * These functions may throw exception. It will be caught in
! 	 * dblink_record_internal()
! 	 */
! 	tuple = BuildTupleFromCStrings(sinfo->attinmeta, cstrs);
! 	tuplestore_puttuple(sinfo->tuplestore, tuple);
! 
! 	return 1;
  }
  
  /*
*** a/contrib/dblink/expected/dblink.out
--- b/contrib/dblink/expected/dblink.out
***************
*** 371,377 **** SELECT *
  FROM dblink('myconn','SELECT * FROM foobar',false) AS t(a int, b text, c text[])
  WHERE t.a > 7;
  NOTICE:  relation "foobar" does not exist
! CONTEXT:  Error occurred on dblink connection named "unnamed": could not execute query.
   a | b | c 
  ---+---+---
  (0 rows)
--- 371,377 ----
  FROM dblink('myconn','SELECT * FROM foobar',false) AS t(a int, b text, c text[])
  WHERE t.a > 7;
  NOTICE:  relation "foobar" does not exist
! CONTEXT:  Error occurred on dblink connection named "myconn": could not execute query.
   a | b | c 
  ---+---+---
  (0 rows)