Allow psql to print COPY command status in more cases.

Previously, psql would print the "COPY nnn" command status only for COPY
commands executed server-side.  Now it will print that for frontend copies
too (including \copy).  However, we continue to suppress the command status
for COPY TO STDOUT, since in that case the copy data has been routed to the
same place that the command status would go, and there is a risk of the
status line being mistaken for another line of COPY data.  Doing that would
break existing scripts, and it doesn't seem worth the benefit --- this case
seems fairly analogous to SELECT, for which we also suppress the command
status.

Kumar Rajeev Rastogi, with substantial review by Amit Khandekar

3 files changed, 53 insertions, 29 deletions

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 6ca9bbc9d86..6968adfd42c 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -632,7 +632,9 @@ StoreQueryTuple(const PGresult *result)
  * degenerates to an AcceptResult() call.
  *
  * Changes its argument to point to the last PGresult of the command string,
- * or NULL if that result was for a COPY FROM STDIN or COPY TO STDOUT.
+ * or NULL if that result was for a COPY TO STDOUT.  (Returning NULL prevents
+ * the command status from being printed, which we want in that case so that
+ * the status line doesn't get taken as part of the COPY data.)
  *
  * Returns true on complete success, false otherwise.  Possible failure modes
  * include purely client-side problems; check the transaction status for the
@@ -641,14 +643,14 @@ StoreQueryTuple(const PGresult *result)
 static bool
 ProcessResult(PGresult **results)
 {
-	PGresult   *next_result;
 	bool		success = true;
 	bool		first_cycle = true;
 
-	do
+	for (;;)
 	{
 		ExecStatusType result_status;
 		bool		is_copy;
+		PGresult   *next_result;
 
 		if (!AcceptResult(*results))
 		{
@@ -693,6 +695,7 @@ ProcessResult(PGresult **results)
 			 * otherwise use queryFout or cur_cmd_source as appropriate.
 			 */
 			FILE	   *copystream = pset.copyStream;
+			PGresult   *copy_result;
 
 			SetCancelConn();
 			if (result_status == PGRES_COPY_OUT)
@@ -700,7 +703,19 @@ ProcessResult(PGresult **results)
 				if (!copystream)
 					copystream = pset.queryFout;
 				success = handleCopyOut(pset.db,
-										copystream) && success;
+										copystream,
+										&copy_result) && success;
+
+				/*
+				 * Suppress status printing if the report would go to the same
+				 * place as the COPY data just went.  Note this doesn't
+				 * prevent error reporting, since handleCopyOut did that.
+				 */
+				if (copystream == pset.queryFout)
+				{
+					PQclear(copy_result);
+					copy_result = NULL;
+				}
 			}
 			else
 			{
@@ -708,30 +723,37 @@ ProcessResult(PGresult **results)
 					copystream = pset.cur_cmd_source;
 				success = handleCopyIn(pset.db,
 									   copystream,
-									   PQbinaryTuples(*results)) && success;
+									   PQbinaryTuples(*results),
+									   &copy_result) && success;
 			}
 			ResetCancelConn();
 
 			/*
-			 * Call PQgetResult() once more.  In the typical case of a
-			 * single-command string, it will return NULL.	Otherwise, we'll
-			 * have other results to process that may include other COPYs.
+			 * Replace the PGRES_COPY_OUT/IN result with COPY command's exit
+			 * status, or with NULL if we want to suppress printing anything.
 			 */
 			PQclear(*results);
-			*results = next_result = PQgetResult(pset.db);
+			*results = copy_result;
 		}
 		else if (first_cycle)
+		{
 			/* fast path: no COPY commands; PQexec visited all results */
 			break;
-		else if ((next_result = PQgetResult(pset.db)))
-		{
-			/* non-COPY command(s) after a COPY: keep the last one */
-			PQclear(*results);
-			*results = next_result;
 		}
 
+		/*
+		 * Check PQgetResult() again.  In the typical case of a single-command
+		 * string, it will return NULL.  Otherwise, we'll have other results
+		 * to process that may include other COPYs.  We keep the last result.
+		 */
+		next_result = PQgetResult(pset.db);
+		if (!next_result)
+			break;
+
+		PQclear(*results);
+		*results = next_result;
 		first_cycle = false;
-	} while (next_result);
+	}
 
 	/* may need this to recover from conn loss during COPY */
 	if (!first_cycle && !CheckConnection())
diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c
index a058f2ff0d4..d706206cc96 100644
--- a/src/bin/psql/copy.c
+++ b/src/bin/psql/copy.c
@@ -429,16 +429,17 @@ do_copy(const char *args)
  * conn should be a database connection that you just issued COPY TO on
  * and got back a PGRES_COPY_OUT result.
  * copystream is the file stream for the data to go to.
+ * The final status for the COPY is returned into *res (but note
+ * we already reported the error, if it's not a success result).
  *
  * result is true if successful, false if not.
  */
 bool
-handleCopyOut(PGconn *conn, FILE *copystream)
+handleCopyOut(PGconn *conn, FILE *copystream, PGresult **res)
 {
 	bool		OK = true;
 	char	   *buf;
 	int			ret;
-	PGresult   *res;
 
 	for (;;)
 	{
@@ -485,13 +486,12 @@ handleCopyOut(PGconn *conn, FILE *copystream)
 	 * but hasn't exited COPY_OUT state internally.  So we ignore the
 	 * possibility here.
 	 */
-	res = PQgetResult(conn);
-	if (PQresultStatus(res) != PGRES_COMMAND_OK)
+	*res = PQgetResult(conn);
+	if (PQresultStatus(*res) != PGRES_COMMAND_OK)
 	{
 		psql_error("%s", PQerrorMessage(conn));
 		OK = false;
 	}
-	PQclear(res);
 
 	return OK;
 }
@@ -504,6 +504,8 @@ handleCopyOut(PGconn *conn, FILE *copystream)
  * and got back a PGRES_COPY_IN result.
  * copystream is the file stream to read the data from.
  * isbinary can be set from PQbinaryTuples().
+ * The final status for the COPY is returned into *res (but note
+ * we already reported the error, if it's not a success result).
  *
  * result is true if successful, false if not.
  */
@@ -512,12 +514,11 @@ handleCopyOut(PGconn *conn, FILE *copystream)
 #define COPYBUFSIZ 8192
 
 bool
-handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary)
+handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res)
 {
 	bool		OK;
 	const char *prompt;
 	char		buf[COPYBUFSIZ];
-	PGresult   *res;
 
 	/*
 	 * Establish longjmp destination for exiting from wait-for-input. (This is
@@ -679,21 +680,20 @@ copyin_cleanup:
 	 * connection is lost.	But that's fine; it will get us out of COPY_IN
 	 * state, which is what we need.)
 	 */
-	while (res = PQgetResult(conn), PQresultStatus(res) == PGRES_COPY_IN)
+	while (*res = PQgetResult(conn), PQresultStatus(*res) == PGRES_COPY_IN)
 	{
 		OK = false;
-		PQclear(res);
+		PQclear(*res);
 		/* We can't send an error message if we're using protocol version 2 */
 		PQputCopyEnd(conn,
 					 (PQprotocolVersion(conn) < 3) ? NULL :
 					 _("trying to exit copy mode"));
 	}
-	if (PQresultStatus(res) != PGRES_COMMAND_OK)
+	if (PQresultStatus(*res) != PGRES_COMMAND_OK)
 	{
 		psql_error("%s", PQerrorMessage(conn));
 		OK = false;
 	}
-	PQclear(res);
 
 	return OK;
 }
diff --git a/src/bin/psql/copy.h b/src/bin/psql/copy.h
index ec1f0d09a9a..2c71da00603 100644
--- a/src/bin/psql/copy.h
+++ b/src/bin/psql/copy.h
@@ -12,11 +12,13 @@
 
 
 /* handler for \copy */
-bool		do_copy(const char *args);
+extern bool do_copy(const char *args);
 
 /* lower level processors for copy in/out streams */
 
-bool		handleCopyOut(PGconn *conn, FILE *copystream);
-bool		handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary);
+extern bool handleCopyOut(PGconn *conn, FILE *copystream,
+			  PGresult **res);
+extern bool handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary,
+			 PGresult **res);
 
 #endif