Support retrieval of results in chunks with libpq.

This patch generalizes libpq's existing single-row mode to allow
individual partial-result PGresults to contain up to N rows, rather
than always one row.  This reduces malloc overhead compared to plain
single-row mode, and it is very useful for psql's FETCH_COUNT feature,
since otherwise we'd have to add code (and cycles) to either merge
single-row PGresults into a bigger one or teach psql's
results-printing logic to accept arrays of PGresults.

To avoid API breakage, PQsetSingleRowMode() remains the same, and we
add a new function PQsetChunkedRowsMode() to invoke the more general
case.  Also, PGresults obtained the old way continue to carry the
PGRES_SINGLE_TUPLE status code, while if PQsetChunkedRowsMode() is
used then their status code is PGRES_TUPLES_CHUNK.  The underlying
logic is the same either way, though.

Daniel Vérité, reviewed by Laurenz Albe and myself (and whacked
around a bit by me, so any remaining bugs are my fault)

Discussion: https://postgr.es/m/CAKZiRmxsVTkO928CM+-ADvsMyePmU3L9DQCa9NwqjvLPcEe5QA@mail.gmail.com

doc/src/sgml/libpq.sgml

@@ -3588,6 +3588,20 @@ ExecStatusType PQresultStatus(const PGresult *res);
           </listitem>
          </varlistentry>
 
+         <varlistentry id="libpq-pgres-tuples-chunk">
+          <term><literal>PGRES_TUPLES_CHUNK</literal></term>
+          <listitem>
+           <para>
+            The <structname>PGresult</structname> contains several result tuples
+            from the current command.  This status occurs only when
+            chunked mode has been selected for the query
+            (see <xref linkend="libpq-single-row-mode"/>).
+            The number of tuples will not exceed the limit passed to
+            <xref linkend="libpq-PQsetChunkedRowsMode"/>.
+           </para>
+          </listitem>
+         </varlistentry>
+
          <varlistentry id="libpq-pgres-pipeline-sync">
           <term><literal>PGRES_PIPELINE_SYNC</literal></term>
           <listitem>
@@ -3617,8 +3631,9 @@ ExecStatusType PQresultStatus(const PGresult *res);
 
         </variablelist>
 
-        If the result status is <literal>PGRES_TUPLES_OK</literal> or
+        If the result status is <literal>PGRES_TUPLES_OK</literal>,
-        <literal>PGRES_SINGLE_TUPLE</literal>, then
+        <literal>PGRES_SINGLE_TUPLE</literal>, or
+        <literal>PGRES_TUPLES_CHUNK</literal>, then
         the functions described below can be used to retrieve the rows
         returned by the query.  Note that a <command>SELECT</command>
         command that happens to retrieve zero rows still shows
@@ -4030,7 +4045,9 @@ void PQclear(PGresult *res);
     These functions are used to extract information from a
     <structname>PGresult</structname> object that represents a successful
     query result (that is, one that has status
-    <literal>PGRES_TUPLES_OK</literal> or <literal>PGRES_SINGLE_TUPLE</literal>).
+    <literal>PGRES_TUPLES_OK</literal>,
+    <literal>PGRES_SINGLE_TUPLE</literal>, or
+    <literal>PGRES_TUPLES_CHUNK</literal>).
     They can also be used to extract
     information from a successful Describe operation: a Describe's result
     has all the same column information that actual execution of the query
@@ -5235,7 +5252,8 @@ PGresult *PQgetResult(PGconn *conn);
   <para>
    Another frequently-desired feature that can be obtained with
    <xref linkend="libpq-PQsendQuery"/> and <xref linkend="libpq-PQgetResult"/>
-   is retrieving large query results a row at a time.  This is discussed
+   is retrieving large query results a limited number of rows at a time.
+   This is discussed
    in <xref linkend="libpq-single-row-mode"/>.
   </para>
 
@@ -5599,15 +5617,6 @@ int PQflush(PGconn *conn);
      queries in the pipeline; see <xref linkend="libpq-pipeline-interleave"/>.
     </para>
 
-    <para>
-     To enter single-row mode, call <function>PQsetSingleRowMode</function>
-     before retrieving results with <function>PQgetResult</function>.
-     This mode selection is effective only for the query currently
-     being processed. For more information on the use of
-     <function>PQsetSingleRowMode</function>,
-     refer to <xref linkend="libpq-single-row-mode"/>.
-    </para>
-
     <para>
      <function>PQgetResult</function> behaves the same as for normal
      asynchronous processing except that it may contain the new
@@ -5972,36 +5981,49 @@ UPDATE mytable SET x = x + 1 WHERE id = 42;
   </sect2>
  </sect1>
 
+ <!-- keep this not-too-apropos sect1 ID for stability of doc URLs -->
  <sect1 id="libpq-single-row-mode">
-  <title>Retrieving Query Results Row-by-Row</title>
+  <title>Retrieving Query Results in Chunks</title>
 
   <indexterm zone="libpq-single-row-mode">
    <primary>libpq</primary>
    <secondary>single-row mode</secondary>
   </indexterm>
 
+  <indexterm zone="libpq-single-row-mode">
+   <primary>libpq</primary>
+   <secondary>chunked mode</secondary>
+  </indexterm>
+
   <para>
    Ordinarily, <application>libpq</application> collects an SQL command's
    entire result and returns it to the application as a single
    <structname>PGresult</structname>.  This can be unworkable for commands
    that return a large number of rows.  For such cases, applications can use
    <xref linkend="libpq-PQsendQuery"/> and <xref linkend="libpq-PQgetResult"/> in
-   <firstterm>single-row mode</firstterm>.  In this mode, the result row(s) are
+   <firstterm>single-row mode</firstterm> or <firstterm>chunked
-   returned to the application one at a time, as they are received from the
+   mode</firstterm>.  In these modes, result row(s) are returned to the
-   server.
+   application as they are received from the server, one at a time for
+   single-row mode or in groups for chunked mode.
   </para>
 
   <para>
-   To enter single-row mode, call <xref linkend="libpq-PQsetSingleRowMode"/>
+   To enter one of these modes, call <xref linkend="libpq-PQsetSingleRowMode"/>
+    or <xref linkend="libpq-PQsetChunkedRowsMode"/>
    immediately after a successful call of <xref linkend="libpq-PQsendQuery"/>
    (or a sibling function).  This mode selection is effective only for the
    currently executing query.  Then call <xref linkend="libpq-PQgetResult"/>
    repeatedly, until it returns null, as documented in <xref
    linkend="libpq-async"/>.  If the query returns any rows, they are returned
-   as individual <structname>PGresult</structname> objects, which look like
+   as one or more <structname>PGresult</structname> objects, which look like
    normal query results except for having status code
-   <literal>PGRES_SINGLE_TUPLE</literal> instead of
+   <literal>PGRES_SINGLE_TUPLE</literal> for single-row mode or
-   <literal>PGRES_TUPLES_OK</literal>.  After the last row, or immediately if
+   <literal>PGRES_TUPLES_CHUNK</literal> for chunked mode, instead of
+   <literal>PGRES_TUPLES_OK</literal>.  There is exactly one result row in
+   each <literal>PGRES_SINGLE_TUPLE</literal> object, while
+   a <literal>PGRES_TUPLES_CHUNK</literal> object contains at least one
+   row but not more than the specified number of rows per chunk.
+   After the last row, or immediately if
    the query returns zero rows, a zero-row object with status
    <literal>PGRES_TUPLES_OK</literal> is returned; this is the signal that no
    more rows will arrive.  (But note that it is still necessary to continue
@@ -6013,9 +6035,9 @@ UPDATE mytable SET x = x + 1 WHERE id = 42;
   </para>
 
   <para>
-   When using pipeline mode, single-row mode needs to be activated for each
+   When using pipeline mode, single-row or chunked mode needs to be
-   query in the pipeline before retrieving results for that query
+   activated for each query in the pipeline before retrieving results for
-   with <function>PQgetResult</function>.
+   that query with <function>PQgetResult</function>.
    See <xref linkend="libpq-pipeline-mode"/> for more information.
   </para>
 
@@ -6046,6 +6068,36 @@ int PQsetSingleRowMode(PGconn *conn);
       </para>
      </listitem>
     </varlistentry>
+
+    <varlistentry id="libpq-PQsetChunkedRowsMode">
+     <term><function>PQsetChunkedRowsMode</function><indexterm><primary>PQsetChunkedRowsMode</primary></indexterm></term>
+
+     <listitem>
+      <para>
+       Select chunked mode for the currently-executing query.
+
+<synopsis>
+int PQsetChunkedRowsMode(PGconn *conn, int chunkSize);
+</synopsis>
+      </para>
+
+      <para>
+       This function is similar to
+       <xref linkend="libpq-PQsetSingleRowMode"/>, except that it
+       specifies retrieval of up to <replaceable>chunkSize</replaceable> rows
+       per <structname>PGresult</structname>, not necessarily just one row.
+       This function can only be called immediately after
+       <xref linkend="libpq-PQsendQuery"/> or one of its sibling functions,
+       before any other operation on the connection such as
+       <xref linkend="libpq-PQconsumeInput"/> or
+       <xref linkend="libpq-PQgetResult"/>.  If called at the correct time,
+       the function activates chunked mode for the current query and
+       returns 1.  Otherwise the mode stays unchanged and the function
+       returns 0.  In any case, the mode reverts to normal after
+       completion of the current query.
+      </para>
+     </listitem>
+    </varlistentry>
    </variablelist>
   </para>
 
@@ -6054,9 +6106,10 @@ int PQsetSingleRowMode(PGconn *conn);
     While processing a query, the server may return some rows and then
     encounter an error, causing the query to be aborted.  Ordinarily,
     <application>libpq</application> discards any such rows and reports only the
-    error.  But in single-row mode, those rows will have already been
+    error.  But in single-row or chunked mode, some rows may have already
-    returned to the application.  Hence, the application will see some
+    been returned to the application. Hence, the application will see some
-    <literal>PGRES_SINGLE_TUPLE</literal> <structname>PGresult</structname>
+    <literal>PGRES_SINGLE_TUPLE</literal> or <literal>PGRES_TUPLES_CHUNK</literal>
+    <structname>PGresult</structname>
     objects followed by a <literal>PGRES_FATAL_ERROR</literal> object.  For
     proper transactional behavior, the application must be designed to
     discard or undo whatever has been done with the previously-processed

src/backend/replication/libpqwalreceiver/libpqwalreceiver.c

@@ -1248,8 +1248,9 @@ libpqrcv_exec(WalReceiverConn *conn, const char *query,
 
 	switch (PQresultStatus(pgres))
 	{
-		case PGRES_SINGLE_TUPLE:
 		case PGRES_TUPLES_OK:
+		case PGRES_SINGLE_TUPLE:
+		case PGRES_TUPLES_CHUNK:
 			walres->status = WALRCV_OK_TUPLES;
 			libpqrcv_processTuples(pgres, walres, nRetTypes, retTypes);
 			break;

src/bin/pg_amcheck/pg_amcheck.c

@@ -991,6 +991,7 @@ should_processing_continue(PGresult *res)
 		case PGRES_SINGLE_TUPLE:
 		case PGRES_PIPELINE_SYNC:
 		case PGRES_PIPELINE_ABORTED:
+		case PGRES_TUPLES_CHUNK:
 			return false;
 	}
 	return true;

src/interfaces/libpq/exports.txt

@@ -203,3 +203,4 @@ PQcancelErrorMessage      200
 PQcancelReset             201
 PQcancelFinish            202
 PQsocketPoll              203
+PQsetChunkedRowsMode      204

src/interfaces/libpq/fe-exec.c

@@ -41,7 +41,8 @@ char	   *const pgresStatus[] = {
 	"PGRES_COPY_BOTH",
 	"PGRES_SINGLE_TUPLE",
 	"PGRES_PIPELINE_SYNC",
-	"PGRES_PIPELINE_ABORTED"
+	"PGRES_PIPELINE_ABORTED",
+	"PGRES_TUPLES_CHUNK"
 };
 
 /* We return this if we're unable to make a PGresult at all */
@@ -200,6 +201,7 @@ PQmakeEmptyPGresult(PGconn *conn, ExecStatusType status)
 			case PGRES_COPY_IN:
 			case PGRES_COPY_BOTH:
 			case PGRES_SINGLE_TUPLE:
+			case PGRES_TUPLES_CHUNK:
 				/* non-error cases */
 				break;
 			default:
@@ -771,7 +773,7 @@ PQclear(PGresult *res)
 /*
  * Handy subroutine to deallocate any partially constructed async result.
  *
- * Any "next" result gets cleared too.
+ * Any "saved" result gets cleared too.
  */
 void
 pqClearAsyncResult(PGconn *conn)
@@ -779,8 +781,8 @@ pqClearAsyncResult(PGconn *conn)
 	PQclear(conn->result);
 	conn->result = NULL;
 	conn->error_result = false;
-	PQclear(conn->next_result);
+	PQclear(conn->saved_result);
-	conn->next_result = NULL;
+	conn->saved_result = NULL;
 }
 
 /*
@@ -911,14 +913,14 @@ pqPrepareAsyncResult(PGconn *conn)
 	}
 
 	/*
-	 * Replace conn->result with next_result, if any.  In the normal case
+	 * Replace conn->result with saved_result, if any.  In the normal case
-	 * there isn't a next result and we're just dropping ownership of the
+	 * there isn't a saved result and we're just dropping ownership of the
-	 * current result.  In single-row mode this restores the situation to what
+	 * current result.  In partial-result mode this restores the situation to
-	 * it was before we created the current single-row result.
+	 * what it was before we created the current partial result.
 	 */
-	conn->result = conn->next_result;
+	conn->result = conn->saved_result;
-	conn->error_result = false; /* next_result is never an error */
+	conn->error_result = false; /* saved_result is never an error */
-	conn->next_result = NULL;
+	conn->saved_result = NULL;
 
 	return res;
 }
@@ -1199,11 +1201,6 @@ pqSaveParameterStatus(PGconn *conn, const char *name, const char *value)
  * On error, *errmsgp can be set to an error string to be returned.
  * (Such a string should already be translated via libpq_gettext().)
  * If it is left NULL, the error is presumed to be "out of memory".
- *
- * In single-row mode, we create a new result holding just the current row,
- * stashing the previous result in conn->next_result so that it becomes
- * active again after pqPrepareAsyncResult().  This allows the result metadata
- * (column descriptions) to be carried forward to each result row.
  */
 int
 pqRowProcessor(PGconn *conn, const char **errmsgp)
@@ -1215,11 +1212,14 @@ pqRowProcessor(PGconn *conn, const char **errmsgp)
 	int			i;
 
 	/*
-	 * In single-row mode, make a new PGresult that will hold just this one
+	 * In partial-result mode, if we don't already have a partial PGresult
-	 * row; the original conn->result is left unchanged so that it can be used
+	 * then make one by cloning conn->result (which should hold the correct
-	 * again as the template for future rows.
+	 * result metadata by now).  Then the original conn->result is moved over
+	 * to saved_result so that we can re-use it as a reference for future
+	 * partial results.  The saved result will become active again after
+	 * pqPrepareAsyncResult() returns the partial result to the application.
 	 */
-	if (conn->singleRowMode)
+	if (conn->partialResMode && conn->saved_result == NULL)
 	{
 		/* Copy everything that should be in the result at this point */
 		res = PQcopyResult(res,
@@ -1227,6 +1227,11 @@ pqRowProcessor(PGconn *conn, const char **errmsgp)
 						   PG_COPYRES_NOTICEHOOKS);
 		if (!res)
 			return 0;
+		/* Change result status to appropriate special value */
+		res->resultStatus = (conn->singleRowMode ? PGRES_SINGLE_TUPLE : PGRES_TUPLES_CHUNK);
+		/* And stash it as the active result */
+		conn->saved_result = conn->result;
+		conn->result = res;
 	}
 
 	/*
@@ -1241,7 +1246,7 @@ pqRowProcessor(PGconn *conn, const char **errmsgp)
 	tup = (PGresAttValue *)
 		pqResultAlloc(res, nfields * sizeof(PGresAttValue), true);
 	if (tup == NULL)
-		goto fail;
+		return 0;
 
 	for (i = 0; i < nfields; i++)
 	{
@@ -1260,7 +1265,7 @@ pqRowProcessor(PGconn *conn, const char **errmsgp)
 
 			val = (char *) pqResultAlloc(res, clen + 1, isbinary);
 			if (val == NULL)
-				goto fail;
+				return 0;
 
 			/* copy and zero-terminate the data (even if it's binary) */
 			memcpy(val, columns[i].value, clen);
@@ -1273,30 +1278,16 @@ pqRowProcessor(PGconn *conn, const char **errmsgp)
 
 	/* And add the tuple to the PGresult's tuple array */
 	if (!pqAddTuple(res, tup, errmsgp))
-		goto fail;
+		return 0;
 
 	/*
-	 * Success.  In single-row mode, make the result available to the client
+	 * Success.  In partial-result mode, if we have enough rows then make the
-	 * immediately.
+	 * result available to the client immediately.
 	 */
-	if (conn->singleRowMode)
+	if (conn->partialResMode && res->ntups >= conn->maxChunkSize)
-	{
-		/* Change result status to special single-row value */
-		res->resultStatus = PGRES_SINGLE_TUPLE;
-		/* Stash old result for re-use later */
-		conn->next_result = conn->result;
-		conn->result = res;
-		/* And mark the result ready to return */
 		conn->asyncStatus = PGASYNC_READY_MORE;
-	}
 
 	return 1;
-
-fail:
-	/* release locally allocated PGresult, if we made one */
-	if (res != conn->result)
-		PQclear(res);
-	return 0;
 }
 
 
@@ -1745,8 +1736,10 @@ PQsendQueryStart(PGconn *conn, bool newQuery)
 		 */
 		pqClearAsyncResult(conn);
 
-		/* reset single-row processing mode */
+		/* reset partial-result mode */
+		conn->partialResMode = false;
 		conn->singleRowMode = false;
+		conn->maxChunkSize = 0;
 	}
 
 	/* ready to send command message */
@@ -1925,30 +1918,61 @@ sendFailed:
 	return 0;
 }
 
+/*
+ * Is it OK to change partial-result mode now?
+ */
+static bool
+canChangeResultMode(PGconn *conn)
+{
+	/*
+	 * Only allow changing the mode when we have launched a query and not yet
+	 * received any results.
+	 */
+	if (!conn)
+		return false;
+	if (conn->asyncStatus != PGASYNC_BUSY)
+		return false;
+	if (!conn->cmd_queue_head ||
+		(conn->cmd_queue_head->queryclass != PGQUERY_SIMPLE &&
+		 conn->cmd_queue_head->queryclass != PGQUERY_EXTENDED))
+		return false;
+	if (pgHavePendingResult(conn))
+		return false;
+	return true;
+}
+
 /*
  * Select row-by-row processing mode
  */
 int
 PQsetSingleRowMode(PGconn *conn)
 {
-	/*
+	if (canChangeResultMode(conn))
-	 * Only allow setting the flag when we have launched a query and not yet
+	{
-	 * received any results.
+		conn->partialResMode = true;
-	 */
+		conn->singleRowMode = true;
-	if (!conn)
+		conn->maxChunkSize = 1;
-		return 0;
+		return 1;
-	if (conn->asyncStatus != PGASYNC_BUSY)
+	}
-		return 0;
+	else
-	if (!conn->cmd_queue_head ||
-		(conn->cmd_queue_head->queryclass != PGQUERY_SIMPLE &&
-		 conn->cmd_queue_head->queryclass != PGQUERY_EXTENDED))
-		return 0;
-	if (pgHavePendingResult(conn))
 		return 0;
+}
 
-	/* OK, set flag */
+/*
-	conn->singleRowMode = true;
+ * Select chunked results processing mode
-	return 1;
+ */
+int
+PQsetChunkedRowsMode(PGconn *conn, int chunkSize)
+{
+	if (chunkSize > 0 && canChangeResultMode(conn))
+	{
+		conn->partialResMode = true;
+		conn->singleRowMode = false;
+		conn->maxChunkSize = chunkSize;
+		return 1;
+	}
+	else
+		return 0;
 }
 
 /*
@@ -2117,6 +2141,20 @@ PQgetResult(PGconn *conn)
 		case PGASYNC_READY:
 			res = pqPrepareAsyncResult(conn);
 
+			/*
+			 * Normally pqPrepareAsyncResult will have left conn->result
+			 * empty.  Otherwise, "res" must be a not-full PGRES_TUPLES_CHUNK
+			 * result, which we want to return to the caller while staying in
+			 * PGASYNC_READY state.  Then the next call here will return the
+			 * empty PGRES_TUPLES_OK result that was restored from
+			 * saved_result, after which we can proceed.
+			 */
+			if (conn->result)
+			{
+				Assert(res->resultStatus == PGRES_TUPLES_CHUNK);
+				break;
+			}
+
 			/* Advance the queue as appropriate */
 			pqCommandQueueAdvance(conn, false,
 								  res->resultStatus == PGRES_PIPELINE_SYNC);
@@ -3173,10 +3211,12 @@ pqPipelineProcessQueue(PGconn *conn)
 	}
 
 	/*
-	 * Reset single-row processing mode.  (Client has to set it up for each
+	 * Reset partial-result mode.  (Client has to set it up for each query, if
-	 * query, if desired.)
+	 * desired.)
 	 */
+	conn->partialResMode = false;
 	conn->singleRowMode = false;
+	conn->maxChunkSize = 0;
 
 	/*
 	 * If there are no further commands to process in the queue, get us in

src/interfaces/libpq/fe-protocol3.c

@@ -379,7 +379,8 @@ pqParseInput3(PGconn *conn)
 					break;
 				case PqMsg_DataRow:
 					if (conn->result != NULL &&
-						conn->result->resultStatus == PGRES_TUPLES_OK)
+						(conn->result->resultStatus == PGRES_TUPLES_OK ||
+						 conn->result->resultStatus == PGRES_TUPLES_CHUNK))
 					{
 						/* Read another tuple of a normal query response */
 						if (getAnotherTuple(conn, msgLength))

src/interfaces/libpq/libpq-fe.h

@@ -112,8 +112,9 @@ typedef enum
 	PGRES_COPY_BOTH,			/* Copy In/Out data transfer in progress */
 	PGRES_SINGLE_TUPLE,			/* single tuple from larger resultset */
 	PGRES_PIPELINE_SYNC,		/* pipeline synchronization point */
-	PGRES_PIPELINE_ABORTED		/* Command didn't run because of an abort
+	PGRES_PIPELINE_ABORTED,		/* Command didn't run because of an abort
 								 * earlier in a pipeline */
+	PGRES_TUPLES_CHUNK			/* chunk of tuples from larger resultset */
 } ExecStatusType;
 
 typedef enum
@@ -489,6 +490,7 @@ extern int	PQsendQueryPrepared(PGconn *conn,
 								const int *paramFormats,
 								int resultFormat);
 extern int	PQsetSingleRowMode(PGconn *conn);
+extern int	PQsetChunkedRowsMode(PGconn *conn, int chunkSize);
 extern PGresult *PQgetResult(PGconn *conn);
 
 /* Routines for managing an asynchronous query */

src/interfaces/libpq/libpq-int.h

@@ -434,7 +434,10 @@ struct pg_conn
 	bool		nonblocking;	/* whether this connection is using nonblock
 								 * sending semantics */
 	PGpipelineStatus pipelineStatus;	/* status of pipeline mode */
+	bool		partialResMode; /* true if single-row or chunked mode */
 	bool		singleRowMode;	/* return current query result row-by-row? */
+	int			maxChunkSize;	/* return query result in chunks not exceeding
+								 * this number of rows */
 	char		copy_is_binary; /* 1 = copy binary, 0 = copy text */
 	int			copy_already_done;	/* # bytes already returned in COPY OUT */
 	PGnotify   *notifyHead;		/* oldest unreported Notify msg */
@@ -535,12 +538,13 @@ struct pg_conn
 	 * and error_result is true, then we need to return a PGRES_FATAL_ERROR
 	 * result, but haven't yet constructed it; text for the error has been
 	 * appended to conn->errorMessage.  (Delaying construction simplifies
-	 * dealing with out-of-memory cases.)  If next_result isn't NULL, it is a
+	 * dealing with out-of-memory cases.)  If saved_result isn't NULL, it is a
-	 * PGresult that will replace "result" after we return that one.
+	 * PGresult that will replace "result" after we return that one; we use
+	 * that in partial-result mode to remember the query's tuple metadata.
 	 */
 	PGresult   *result;			/* result being constructed */
 	bool		error_result;	/* do we need to make an ERROR result? */
-	PGresult   *next_result;	/* next result (used in single-row mode) */
+	PGresult   *saved_result;	/* original, empty result in partialResMode */
 
 	/* Assorted state for SASL, SSL, GSS, etc */
 	const pg_fe_sasl_mech *sasl;

src/test/modules/libpq_pipeline/libpq_pipeline.c

@@ -1719,6 +1719,46 @@ test_singlerowmode(PGconn *conn)
 	if (PQgetResult(conn) != NULL)
 		pg_fatal("expected NULL result");
 
+	/*
+	 * Try chunked mode as well; make sure that it correctly delivers a
+	 * partial final chunk.
+	 */
+	if (PQsendQueryParams(conn, "SELECT generate_series(1, 5)",
+						  0, NULL, NULL, NULL, NULL, 0) != 1)
+		pg_fatal("failed to send query: %s",
+				 PQerrorMessage(conn));
+	if (PQsendFlushRequest(conn) != 1)
+		pg_fatal("failed to send flush request");
+	if (PQsetChunkedRowsMode(conn, 3) != 1)
+		pg_fatal("PQsetChunkedRowsMode() failed");
+	res = PQgetResult(conn);
+	if (res == NULL)
+		pg_fatal("unexpected NULL");
+	if (PQresultStatus(res) != PGRES_TUPLES_CHUNK)
+		pg_fatal("Expected PGRES_TUPLES_CHUNK, got %s: %s",
+				 PQresStatus(PQresultStatus(res)),
+				 PQerrorMessage(conn));
+	if (PQntuples(res) != 3)
+		pg_fatal("Expected 3 rows, got %d", PQntuples(res));
+	res = PQgetResult(conn);
+	if (res == NULL)
+		pg_fatal("unexpected NULL");
+	if (PQresultStatus(res) != PGRES_TUPLES_CHUNK)
+		pg_fatal("Expected PGRES_TUPLES_CHUNK, got %s",
+				 PQresStatus(PQresultStatus(res)));
+	if (PQntuples(res) != 2)
+		pg_fatal("Expected 2 rows, got %d", PQntuples(res));
+	res = PQgetResult(conn);
+	if (res == NULL)
+		pg_fatal("unexpected NULL");
+	if (PQresultStatus(res) != PGRES_TUPLES_OK)
+		pg_fatal("Expected PGRES_TUPLES_OK, got %s",
+				 PQresStatus(PQresultStatus(res)));
+	if (PQntuples(res) != 0)
+		pg_fatal("Expected 0 rows, got %d", PQntuples(res));
+	if (PQgetResult(conn) != NULL)
+		pg_fatal("expected NULL result");
+
 	if (PQexitPipelineMode(conn) != 1)
 		pg_fatal("failed to end pipeline mode: %s", PQerrorMessage(conn));
 

src/test/modules/libpq_pipeline/traces/singlerow.trace

@@ -56,4 +56,18 @@ B	4	BindComplete
 B	33	RowDescription	 1 "?column?" NNNN 0 NNNN 4 -1 0
 B	11	DataRow	 1 1 '1'
 B	13	CommandComplete	 "SELECT 1"
+F	36	Parse	 "" "SELECT generate_series(1, 5)" 0
+F	14	Bind	 "" "" 0 0 1 0
+F	6	Describe	 P ""
+F	9	Execute	 "" 0
+F	4	Flush
+B	4	ParseComplete
+B	4	BindComplete
+B	40	RowDescription	 1 "generate_series" NNNN 0 NNNN 4 -1 0
+B	11	DataRow	 1 1 '1'
+B	11	DataRow	 1 1 '2'
+B	11	DataRow	 1 1 '3'
+B	11	DataRow	 1 1 '4'
+B	11	DataRow	 1 1 '5'
+B	13	CommandComplete	 "SELECT 5"
 F	4	Terminate