Add a "row processor" API to libpq for better handling of large results.

Traditionally libpq has collected an entire query result before passing
it back to the application.  That provides a simple and transactional API,
but it's pretty inefficient for large result sets.  This patch allows the
application to process each row on-the-fly instead of accumulating the
rows into the PGresult.  Error recovery becomes a bit more complex, but
often that tradeoff is well worth making.

Kyotaro Horiguchi, reviewed by Marko Kreen and Tom Lane

src/interfaces/libpq/exports.txt

@@ -160,3 +160,6 @@ PQconnectStartParams      157
 PQping                    158
 PQpingParams              159
 PQlibVersion              160
+PQsetRowProcessor         161
+PQgetRowProcessor         162
+PQskipResult              163

src/interfaces/libpq/fe-connect.c

@@ -2425,7 +2425,7 @@ keep_going:						/* We will come back to here until there is
 					conn->status = CONNECTION_AUTH_OK;
 
 					/*
-					 * Set asyncStatus so that PQsetResult will think that
+					 * Set asyncStatus so that PQgetResult will think that
 					 * what comes back next is the result of a query.  See
 					 * below.
 					 */
@@ -2686,8 +2686,11 @@ makeEmptyPGconn(void)
 	/* Zero all pointers and booleans */
 	MemSet(conn, 0, sizeof(PGconn));
 
+	/* install default row processor and notice hooks */
+	PQsetRowProcessor(conn, NULL, NULL);
 	conn->noticeHooks.noticeRec = defaultNoticeReceiver;
 	conn->noticeHooks.noticeProc = defaultNoticeProcessor;
+
 	conn->status = CONNECTION_BAD;
 	conn->asyncStatus = PGASYNC_IDLE;
 	conn->xactStatus = PQTRANS_IDLE;
@@ -2721,11 +2724,14 @@ makeEmptyPGconn(void)
 	conn->inBuffer = (char *) malloc(conn->inBufSize);
 	conn->outBufSize = 16 * 1024;
 	conn->outBuffer = (char *) malloc(conn->outBufSize);
+	conn->rowBufLen = 32;
+	conn->rowBuf = (PGdataValue *) malloc(conn->rowBufLen * sizeof(PGdataValue));
 	initPQExpBuffer(&conn->errorMessage);
 	initPQExpBuffer(&conn->workBuffer);
 
 	if (conn->inBuffer == NULL ||
 		conn->outBuffer == NULL ||
+		conn->rowBuf == NULL ||
 		PQExpBufferBroken(&conn->errorMessage) ||
 		PQExpBufferBroken(&conn->workBuffer))
 	{
@@ -2829,6 +2835,8 @@ freePGconn(PGconn *conn)
 		free(conn->inBuffer);
 	if (conn->outBuffer)
 		free(conn->outBuffer);
+	if (conn->rowBuf)
+		free(conn->rowBuf);
 	termPQExpBuffer(&conn->errorMessage);
 	termPQExpBuffer(&conn->workBuffer);
 
@@ -2888,7 +2896,7 @@ closePGconn(PGconn *conn)
 	conn->status = CONNECTION_BAD;		/* Well, not really _bad_ - just
 										 * absent */
 	conn->asyncStatus = PGASYNC_IDLE;
-	pqClearAsyncResult(conn);	/* deallocate result and curTuple */
+	pqClearAsyncResult(conn);	/* deallocate result */
 	pg_freeaddrinfo_all(conn->addrlist_family, conn->addrlist);
 	conn->addrlist = NULL;
 	conn->addr_cur = NULL;

src/interfaces/libpq/fe-exec.c

@@ -50,6 +50,9 @@ static bool static_std_strings = false;
 
 
 static PGEvent *dupEvents(PGEvent *events, int count);
+static bool pqAddTuple(PGresult *res, PGresAttValue *tup);
+static int pqStdRowProcessor(PGresult *res, const PGdataValue *columns,
+				  const char **errmsgp, void *param);
 static bool PQsendQueryStart(PGconn *conn);
 static int PQsendQueryGuts(PGconn *conn,
 				const char *command,
@@ -61,6 +64,8 @@ static int PQsendQueryGuts(PGconn *conn,
 				const int *paramFormats,
 				int resultFormat);
 static void parseInput(PGconn *conn);
+static int dummyRowProcessor(PGresult *res, const PGdataValue *columns,
+				  const char **errmsgp, void *param);
 static bool PQexecStart(PGconn *conn);
 static PGresult *PQexecFinish(PGconn *conn);
 static int PQsendDescribe(PGconn *conn, char desc_type,
@@ -694,14 +699,12 @@ PQclear(PGresult *res)
 /*
  * Handy subroutine to deallocate any partially constructed async result.
  */
-
 void
 pqClearAsyncResult(PGconn *conn)
 {
 	if (conn->result)
 		PQclear(conn->result);
 	conn->result = NULL;
-	conn->curTuple = NULL;
 }
 
 /*
@@ -756,7 +759,6 @@ pqPrepareAsyncResult(PGconn *conn)
 	 */
 	res = conn->result;
 	conn->result = NULL;		/* handing over ownership to caller */
-	conn->curTuple = NULL;		/* just in case */
 	if (!res)
 		res = PQmakeEmptyPGresult(conn, PGRES_FATAL_ERROR);
 	else
@@ -832,7 +834,7 @@ pqInternalNotice(const PGNoticeHooks *hooks, const char *fmt,...)
  *	  add a row pointer to the PGresult structure, growing it if necessary
  *	  Returns TRUE if OK, FALSE if not enough memory to add the row
  */
-int
+static bool
 pqAddTuple(PGresult *res, PGresAttValue *tup)
 {
 	if (res->ntups >= res->tupArrSize)
@@ -978,6 +980,124 @@ pqSaveParameterStatus(PGconn *conn, const char *name, const char *value)
 }
 
 
+/*
+ * PQsetRowProcessor
+ *	  Set function that copies row data out from the network buffer,
+ *	  along with a passthrough parameter for it.
+ */
+void
+PQsetRowProcessor(PGconn *conn, PQrowProcessor func, void *param)
+{
+	if (!conn)
+		return;
+
+	if (func)
+	{
+		/* set custom row processor */
+		conn->rowProcessor = func;
+		conn->rowProcessorParam = param;
+	}
+	else
+	{
+		/* set default row processor */
+		conn->rowProcessor = pqStdRowProcessor;
+		conn->rowProcessorParam = conn;
+	}
+}
+
+/*
+ * PQgetRowProcessor
+ *	  Get current row processor of PGconn.
+ *	  If param is not NULL, also store the passthrough parameter at *param.
+ */
+PQrowProcessor
+PQgetRowProcessor(const PGconn *conn, void **param)
+{
+	if (!conn)
+	{
+		if (param)
+			*param = NULL;
+		return NULL;
+	}
+
+	if (param)
+		*param = conn->rowProcessorParam;
+	return conn->rowProcessor;
+}
+
+/*
+ * pqStdRowProcessor
+ *	  Add the received row to the PGresult structure
+ *	  Returns 1 if OK, -1 if error occurred.
+ *
+ * Note: "param" should point to the PGconn, but we don't actually need that
+ * as of the current coding.
+ */
+static int
+pqStdRowProcessor(PGresult *res, const PGdataValue *columns,
+				  const char **errmsgp, void *param)
+{
+	int			nfields = res->numAttributes;
+	PGresAttValue *tup;
+	int			i;
+
+	if (columns == NULL)
+	{
+		/* New result set ... we have nothing to do in this function. */
+		return 1;
+	}
+
+	/*
+	 * Basically we just allocate space in the PGresult for each field and
+	 * copy the data over.
+	 *
+	 * Note: on malloc failure, we return -1 leaving *errmsgp still NULL,
+	 * which caller will take to mean "out of memory".  This is preferable to
+	 * trying to set up such a message here, because evidently there's not
+	 * enough memory for gettext() to do anything.
+	 */
+	tup = (PGresAttValue *)
+		pqResultAlloc(res, nfields * sizeof(PGresAttValue), TRUE);
+	if (tup == NULL)
+		return -1;
+
+	for (i = 0; i < nfields; i++)
+	{
+		int		clen = columns[i].len;
+
+		if (clen < 0)
+		{
+			/* null field */
+			tup[i].len = NULL_LEN;
+			tup[i].value = res->null_field;
+		}
+		else
+		{
+			bool		isbinary = (res->attDescs[i].format != 0);
+			char	   *val;
+
+			val = (char *) pqResultAlloc(res, clen + 1, isbinary);
+			if (val == NULL)
+				return -1;
+
+			/* copy and zero-terminate the data (even if it's binary) */
+			memcpy(val, columns[i].value, clen);
+			val[clen] = '\0';
+
+			tup[i].len = clen;
+			tup[i].value = val;
+		}
+	}
+
+	/* And add the tuple to the PGresult's tuple array */
+	if (!pqAddTuple(res, tup))
+		return -1;
+
+	/* Success */
+	return 1;
+}
+
+
 /*
  * PQsendQuery
  *	 Submit a query, but don't wait for it to finish
@@ -1223,7 +1343,6 @@ PQsendQueryStart(PGconn *conn)
 
 	/* initialize async result-accumulation state */
 	conn->result = NULL;
-	conn->curTuple = NULL;
 
 	/* ready to send command message */
 	return true;
@@ -1468,6 +1587,9 @@ PQconsumeInput(PGconn *conn)
  * parseInput: if appropriate, parse input data from backend
  * until input is exhausted or a stopping state is reached.
  * Note that this function will NOT attempt to read more data from the backend.
+ *
+ * Note: callers of parseInput must be prepared for a longjmp exit when we are
+ * in PGASYNC_BUSY state, since an external row processor might do that.
  */
 static void
 parseInput(PGconn *conn)
@@ -1615,6 +1737,49 @@ PQgetResult(PGconn *conn)
 	return res;
 }
 
+/*
+ * PQskipResult
+ *	  Get the next PGresult produced by a query, but discard any data rows.
+ *
+ * This is mainly useful for cleaning up after a longjmp out of a row
+ * processor, when resuming processing of the current query result isn't
+ * wanted.  Note that this is of little value in an async-style application,
+ * since any preceding calls to PQisBusy would have already called the regular
+ * row processor.
+ */
+PGresult *
+PQskipResult(PGconn *conn)
+{
+	PGresult   *res;
+	PQrowProcessor savedRowProcessor;
+
+	if (!conn)
+		return NULL;
+
+	/* temporarily install dummy row processor */
+	savedRowProcessor = conn->rowProcessor;
+	conn->rowProcessor = dummyRowProcessor;
+	/* no need to save/change rowProcessorParam */
+
+	/* fetch the next result */
+	res = PQgetResult(conn);
+
+	/* restore previous row processor */
+	conn->rowProcessor = savedRowProcessor;
+
+	return res;
+}
+
+/*
+ * Do-nothing row processor for PQskipResult
+ */
+static int
+dummyRowProcessor(PGresult *res, const PGdataValue *columns,
+				  const char **errmsgp, void *param)
+{
+	return 1;
+}
+
 
 /*
  * PQexec
@@ -1721,7 +1886,7 @@ PQexecStart(PGconn *conn)
 	 * Silently discard any prior query result that application didn't eat.
 	 * This is probably poor design, but it's here for backward compatibility.
 	 */
-	while ((result = PQgetResult(conn)) != NULL)
+	while ((result = PQskipResult(conn)) != NULL)
 	{
 		ExecStatusType resultStatus = result->resultStatus;
 

src/interfaces/libpq/fe-lobj.c

@@ -40,9 +40,7 @@
 #define LO_BUFSIZE		  8192
 
 static int	lo_initialize(PGconn *conn);
-
-static Oid
-			lo_import_internal(PGconn *conn, const char *filename, const Oid oid);
+static Oid	lo_import_internal(PGconn *conn, const char *filename, Oid oid);
 
 /*
  * lo_open
@@ -59,7 +57,7 @@ lo_open(PGconn *conn, Oid lobjId, int mode)
 	PQArgBlock	argv[2];
 	PGresult   *res;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return -1;
@@ -101,7 +99,7 @@ lo_close(PGconn *conn, int fd)
 	int			retval;
 	int			result_len;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return -1;
@@ -139,7 +137,7 @@ lo_truncate(PGconn *conn, int fd, size_t len)
 	int			retval;
 	int			result_len;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return -1;
@@ -192,7 +190,7 @@ lo_read(PGconn *conn, int fd, char *buf, size_t len)
 	PGresult   *res;
 	int			result_len;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return -1;
@@ -234,7 +232,7 @@ lo_write(PGconn *conn, int fd, const char *buf, size_t len)
 	int			result_len;
 	int			retval;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return -1;
@@ -280,7 +278,7 @@ lo_lseek(PGconn *conn, int fd, int offset, int whence)
 	int			retval;
 	int			result_len;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return -1;
@@ -328,7 +326,7 @@ lo_creat(PGconn *conn, int mode)
 	int			retval;
 	int			result_len;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return InvalidOid;
@@ -367,7 +365,7 @@ lo_create(PGconn *conn, Oid lobjId)
 	int			retval;
 	int			result_len;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return InvalidOid;
@@ -413,7 +411,7 @@ lo_tell(PGconn *conn, int fd)
 	PGresult   *res;
 	int			result_len;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return -1;
@@ -451,7 +449,7 @@ lo_unlink(PGconn *conn, Oid lobjId)
 	int			result_len;
 	int			retval;
 
-	if (conn->lobjfuncs == NULL)
+	if (conn == NULL || conn->lobjfuncs == NULL)
 	{
 		if (lo_initialize(conn) < 0)
 			return -1;
@@ -505,7 +503,7 @@ lo_import_with_oid(PGconn *conn, const char *filename, Oid lobjId)
 }
 
 static Oid
-lo_import_internal(PGconn *conn, const char *filename, const Oid oid)
+lo_import_internal(PGconn *conn, const char *filename, Oid oid)
 {
 	int			fd;
 	int			nbytes,
@@ -684,8 +682,13 @@ lo_initialize(PGconn *conn)
 	int			n;
 	const char *query;
 	const char *fname;
+	PQrowProcessor savedRowProcessor;
+	void	   *savedRowProcessorParam;
 	Oid			foid;
 
+	if (!conn)
+		return -1;
+
 	/*
 	 * Allocate the structure to hold the functions OID's
 	 */
@@ -729,7 +732,16 @@ lo_initialize(PGconn *conn)
 			"or proname = 'loread' "
 			"or proname = 'lowrite'";
 
+	/* Ensure the standard row processor is used to collect the result */
+	savedRowProcessor = conn->rowProcessor;
+	savedRowProcessorParam = conn->rowProcessorParam;
+	PQsetRowProcessor(conn, NULL, NULL);
+
 	res = PQexec(conn, query);
+
+	conn->rowProcessor = savedRowProcessor;
+	conn->rowProcessorParam = savedRowProcessorParam;
+
 	if (res == NULL)
 	{
 		free(lobjfuncs);

src/interfaces/libpq/fe-misc.c

@@ -218,6 +218,32 @@ pqGetnchar(char *s, size_t len, PGconn *conn)
 	return 0;
 }
 
+/*
+ * pqSkipnchar:
+ *	skip over len bytes in input buffer.
+ *
+ * Note: this is primarily useful for its debug output, which should
+ * be exactly the same as for pqGetnchar.  We assume the data in question
+ * will actually be used, but just isn't getting copied anywhere as yet.
+ */
+int
+pqSkipnchar(size_t len, PGconn *conn)
+{
+	if (len > (size_t) (conn->inEnd - conn->inCursor))
+		return EOF;
+
+	if (conn->Pfdebug)
+	{
+		fprintf(conn->Pfdebug, "From backend (%lu)> ", (unsigned long) len);
+		fputnbytes(conn->Pfdebug, conn->inBuffer + conn->inCursor, len);
+		fprintf(conn->Pfdebug, "\n");
+	}
+
+	conn->inCursor += len;
+
+	return 0;
+}
+
 /*
  * pqPutnchar:
  *	write exactly len bytes to the current message

src/interfaces/libpq/libpq-fe.h

@@ -38,13 +38,14 @@ extern		"C"
 
 /* Application-visible enum types */
 
-typedef enum
-{
-	/*
-	 * Although it is okay to add to this list, values which become unused
+/*
+ * Although it is okay to add to these lists, values which become unused
  * should never be removed, nor should constants be redefined - that would
  * break compatibility with existing code.
  */
+
+typedef enum
+{
 	CONNECTION_OK,
 	CONNECTION_BAD,
 	/* Non-blocking mode only below here */
@@ -128,6 +129,17 @@ typedef struct pg_conn PGconn;
  */
 typedef struct pg_result PGresult;
 
+/* PGdataValue represents a data field value being passed to a row processor.
+ * It could be either text or binary data; text data is not zero-terminated.
+ * A SQL NULL is represented by len < 0; then value is still valid but there
+ * are no data bytes there.
+ */
+typedef struct pgDataValue
+{
+	int			len;			/* data length in bytes, or <0 if NULL */
+	const char *value;			/* data value, without zero-termination */
+} PGdataValue;
+
 /* PGcancel encapsulates the information needed to cancel a running
  * query on an existing connection.
  * The contents of this struct are not supposed to be known to applications.
@@ -149,6 +161,10 @@ typedef struct pgNotify
 	struct pgNotify *next;		/* list link */
 } PGnotify;
 
+/* Function type for row-processor callback */
+typedef int (*PQrowProcessor) (PGresult *res, const PGdataValue *columns,
+							   const char **errmsgp, void *param);
+
 /* Function types for notice-handling callbacks */
 typedef void (*PQnoticeReceiver) (void *arg, const PGresult *res);
 typedef void (*PQnoticeProcessor) (void *arg, const char *message);
@@ -388,11 +404,16 @@ extern int PQsendQueryPrepared(PGconn *conn,
 					const int *paramFormats,
 					int resultFormat);
 extern PGresult *PQgetResult(PGconn *conn);
+extern PGresult *PQskipResult(PGconn *conn);
 
 /* Routines for managing an asynchronous query */
 extern int	PQisBusy(PGconn *conn);
 extern int	PQconsumeInput(PGconn *conn);
 
+/* Override default per-row processing */
+extern void PQsetRowProcessor(PGconn *conn, PQrowProcessor func, void *param);
+extern PQrowProcessor PQgetRowProcessor(const PGconn *conn, void **param);
+
 /* LISTEN/NOTIFY support */
 extern PGnotify *PQnotifies(PGconn *conn);
 

src/interfaces/libpq/libpq-int.h

@@ -324,6 +324,10 @@ struct pg_conn
 	/* Optional file to write trace info to */
 	FILE	   *Pfdebug;
 
+	/* Callback procedure for per-row processing */
+	PQrowProcessor rowProcessor;	/* function pointer */
+	void	   *rowProcessorParam;	/* passthrough argument */
+
 	/* Callback procedures for notice message processing */
 	PGNoticeHooks noticeHooks;
 
@@ -396,9 +400,14 @@ struct pg_conn
 								 * msg has no length word */
 	int			outMsgEnd;		/* offset to msg end (so far) */
 
+	/* Row processor interface workspace */
+	PGdataValue *rowBuf;		/* array for passing values to rowProcessor */
+	int			rowBufLen;		/* number of entries allocated in rowBuf */
+
 	/* Status for asynchronous result construction */
 	PGresult   *result;			/* result being constructed */
-	PGresAttValue *curTuple;	/* tuple currently being read */
+
+	/* Assorted state for SSL, GSS, etc */
 
 #ifdef USE_SSL
 	bool		allow_ssl_try;	/* Allowed to try SSL negotiation */
@@ -435,7 +444,6 @@ struct pg_conn
 								 * connection */
 #endif
 
-
 	/* Buffer for current error message */
 	PQExpBufferData errorMessage;		/* expansible string */
 
@@ -505,7 +513,6 @@ extern void
 pqInternalNotice(const PGNoticeHooks *hooks, const char *fmt,...)
 /* This lets gcc check the format string for consistency. */
 __attribute__((format(PG_PRINTF_ATTRIBUTE, 2, 3)));
-extern int	pqAddTuple(PGresult *res, PGresAttValue *tup);
 extern void pqSaveMessageField(PGresult *res, char code,
 				   const char *value);
 extern void pqSaveParameterStatus(PGconn *conn, const char *name,
@@ -558,6 +565,7 @@ extern int	pqGets(PQExpBuffer buf, PGconn *conn);
 extern int	pqGets_append(PQExpBuffer buf, PGconn *conn);
 extern int	pqPuts(const char *s, PGconn *conn);
 extern int	pqGetnchar(char *s, size_t len, PGconn *conn);
+extern int	pqSkipnchar(size_t len, PGconn *conn);
 extern int	pqPutnchar(const char *s, size_t len, PGconn *conn);
 extern int	pqGetInt(int *result, size_t bytes, PGconn *conn);
 extern int	pqPutInt(int value, size_t bytes, PGconn *conn);