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

doc/src/sgml/ref/copy.sgml

@@ -370,6 +370,16 @@ COPY <replaceable class="parameter">count</replaceable>
    The <replaceable class="parameter">count</replaceable> is the number
    of rows copied.
   </para>
+
+  <note>
+   <para>
+    <application>psql</> will print this command tag only if the command
+    was not <literal>COPY ... TO STDOUT</>, or the
+    equivalent <application>psql</> meta-command
+    <literal>\copy ... to stdout</>.  This is to prevent confusing the
+    command tag with the data that was just printed.
+   </para>
+  </note>
  </refsect1>
 
  <refsect1>

doc/src/sgml/ref/psql-ref.sgml

@@ -863,36 +863,36 @@ testdb=>
         <para>
         When <literal>program</> is specified,
         <replaceable class="parameter">command</replaceable> is
-        executed by <application>psql</application> and the data from
+        executed by <application>psql</application> and the data passed from
         or to <replaceable class="parameter">command</replaceable> is
         routed between the server and the client.
-        This means that the execution privileges are those of
+        Again, the execution privileges are those of
         the local user, not the server, and no SQL superuser
         privileges are required.
         </para>
 
-        <para><literal>\copy ... from stdin | to stdout</literal>
-        reads/writes based on the command input and output respectively.
-        All rows are read from the same source that issued the command,
-        continuing until <literal>\.</literal> is read or the stream
-        reaches <acronym>EOF</>. Output is sent to the same place as
-        command output. To read/write from
-        <application>psql</application>'s standard input or output, use
-        <literal>pstdin</> or <literal>pstdout</>. This option is useful
+        <para>
+        For <literal>\copy ... from stdin</>, data rows are read from the same
+        source that issued the command, continuing until <literal>\.</literal>
+        is read or the stream reaches <acronym>EOF</>. This option is useful
         for populating tables in-line within a SQL script file.
+        For <literal>\copy ... to stdout</>, output is sent to the same place
+        as <application>psql</> command output, and
+        the <literal>COPY <replaceable>count</></literal> command status is
+        not printed (since it might be confused with a data row).
+        To read/write <application>psql</application>'s standard input or
+        output regardless of the current command source or <literal>\o</>
+        option, write <literal>from pstdin</> or <literal>to pstdout</>.
         </para>
 
         <para>
-        The syntax of the command is similar to that of the
+        The syntax of this command is similar to that of the
         <acronym>SQL</acronym> <xref linkend="sql-copy">
-        command, and
-        <replaceable class="parameter">option</replaceable>
-        must indicate one of the options of the
-        <acronym>SQL</acronym> <xref linkend="sql-copy"> command.
-        Note that, because of this,
-        special parsing rules apply to the <command>\copy</command>
-        command. In particular, the variable substitution rules and
-        backslash escapes do not apply.
+        command.  All options other than the data source/destination are
+        as specified for <xref linkend="sql-copy">.
+        Because of this, special parsing rules apply to the <command>\copy</>
+        command. In particular, <application>psql</>'s variable substitution
+        rules and backslash escapes do not apply.
         </para>
 
         <tip>

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())

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;
 }

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