Have \copy go through SendQuery

This enables a bunch of features, notably ON_ERROR_ROLLBACK.  It also
makes COPY failure (either in the server or psql) as a whole behave more
sanely in psql.

Additionally, having more commands in the same command line as COPY
works better (though since psql splits lines at semicolons, this doesn't
matter much unless you're using -c).

Also tighten a couple of switches on PQresultStatus() to add
PGRES_COPY_BOTH support and stop assuming that unknown statuses received
are errors; have those print diagnostics where warranted.

Author: Noah Misch

src/bin/psql/command.c

@@ -320,25 +320,10 @@ exec_command(const char *cmd,
 	/* \copy */
 	else if (pg_strcasecmp(cmd, "copy") == 0)
 	{
-		/* Default fetch-it-all-and-print mode */
-		instr_time	before,
-					after;
-
 		char	   *opt = psql_scan_slash_option(scan_state,
 												 OT_WHOLE_LINE, NULL, false);
 
-		if (pset.timing)
-			INSTR_TIME_SET_CURRENT(before);
-
 		success = do_copy(opt);
-
-		if (pset.timing && success)
-		{
-			INSTR_TIME_SET_CURRENT(after);
-			INSTR_TIME_SUBTRACT(after, before);
-			printf(_("Time: %.3f ms\n"), INSTR_TIME_GET_MILLISEC(after));
-		}
-
 		free(opt);
 	}
 

src/bin/psql/common.c

@@ -438,7 +438,7 @@ ResetCancelConn(void)
 static bool
 AcceptResult(const PGresult *result)
 {
-	bool		OK = true;
+	bool		OK;
 
 	if (!result)
 		OK = false;
@@ -450,11 +450,21 @@ AcceptResult(const PGresult *result)
 			case PGRES_EMPTY_QUERY:
 			case PGRES_COPY_IN:
 			case PGRES_COPY_OUT:
+			case PGRES_COPY_BOTH:
 				/* Fine, do nothing */
+				OK = true;
+				break;
+
+			case PGRES_BAD_RESPONSE:
+			case PGRES_NONFATAL_ERROR:
+			case PGRES_FATAL_ERROR:
+				OK = false;
 				break;
 
 			default:
 				OK = false;
+				psql_error("unexpected PQresultStatus (%d)",
+						   PQresultStatus(result));
 				break;
 		}
 
@@ -620,46 +630,115 @@ PrintQueryTuples(const PGresult *results)
 
 
 /*
- * ProcessCopyResult: if command was a COPY FROM STDIN/TO STDOUT, handle it
+ * ProcessResult: utility function for use by SendQuery() only
  *
- * Note: Utility function for use by SendQuery() only.
+ * When our command string contained a COPY FROM STDIN or COPY TO STDOUT,
+ * PQexec() has stopped at the PGresult associated with the first such
+ * command.  In that event, we'll marshal data for the COPY and then cycle
+ * through any subsequent PGresult objects.
  *
- * Returns true if the query executed successfully, false otherwise.
+ * When the command string contained no affected COPY command, this function
+ * 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.
+ *
+ * Returns true on complete success, false otherwise.  Possible failure modes
+ * include purely client-side problems; check the transaction status for the
+ * server-side opinion.
  */
 static bool
-ProcessCopyResult(PGresult *results)
+ProcessResult(PGresult **results)
 {
-	bool		success = false;
+	PGresult   *next_result;
+	bool		success = true;
+	bool		first_cycle = true;
 
-	if (!results)
-		return false;
+	do
+	{
+		ExecStatusType result_status;
+		bool		is_copy;
 
-	switch (PQresultStatus(results))
+		if (!AcceptResult(*results))
 		{
-		case PGRES_TUPLES_OK:
-		case PGRES_COMMAND_OK:
-		case PGRES_EMPTY_QUERY:
-			/* nothing to do here */
-			success = true;
+			/*
+			 * Failure at this point is always a server-side failure or a
+			 * failure to submit the command string.  Either way, we're
+			 * finished with this command string.
+			 */
+			success = false;
 			break;
+		}
+
+		result_status = PQresultStatus(*results);
+		switch (result_status)
+		{
+			case PGRES_COPY_BOTH:
+				/*
+				 * No now-existing SQL command can yield PGRES_COPY_BOTH, but
+				 * defend against the future.  PQexec() can't short-circuit
+				 * it's way out of a PGRES_COPY_BOTH, so the connection will
+				 * be useless at this point.  XXX is there a method for
+				 * clearing this status that's likely to work with every
+				 * future command that can initiate it?
+				 */
+				psql_error("unexpected PQresultStatus (%d)", result_status);
+				return false;
 
 			case PGRES_COPY_OUT:
-			SetCancelConn();
-			success = handleCopyOut(pset.db, pset.queryFout);
-			ResetCancelConn();
+			case PGRES_COPY_IN:
+				is_copy = true;
 				break;
 
-		case PGRES_COPY_IN:
+			case PGRES_EMPTY_QUERY:
+			case PGRES_COMMAND_OK:
+			case PGRES_TUPLES_OK:
+				is_copy = false;
+				break;
+
+			default:
+				/* AcceptResult() should have caught anything else. */
+				is_copy = false;
+				psql_error("unexpected PQresultStatus (%d)", result_status);
+				break;
+		}
+
+		if (is_copy)
+		{
+			/*
+			 * Marshal the COPY data.  Either subroutine will get the
+			 * connection out of its COPY state, then call PQresultStatus()
+			 * once and report any error.
+			 */
 			SetCancelConn();
+			if (result_status == PGRES_COPY_OUT)
+				success = handleCopyOut(pset.db, pset.queryFout) && success;
+			else
 				success = handleCopyIn(pset.db, pset.cur_cmd_source,
-								   PQbinaryTuples(results));
+									   PQbinaryTuples(*results)) && success;
 			ResetCancelConn();
-			break;
 
-		default:
+			/*
+			 * 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.
+			 */
+			PQclear(*results);
+			*results = next_result = PQgetResult(pset.db);
+		}
+		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;
 		}
 
+		first_cycle = false;
+	} while (next_result);
+
 	/* may need this to recover from conn loss during COPY */
 	if (!CheckConnection())
 		return false;
@@ -708,7 +787,7 @@ PrintQueryStatus(PGresult *results)
 static bool
 PrintQueryResults(PGresult *results)
 {
-	bool		success = false;
+	bool		success;
 	const char *cmdstatus;
 
 	if (!results)
@@ -738,11 +817,21 @@ PrintQueryResults(PGresult *results)
 
 		case PGRES_COPY_OUT:
 		case PGRES_COPY_IN:
+		case PGRES_COPY_BOTH:
 			/* nothing to do here */
 			success = true;
 			break;
 
+		case PGRES_BAD_RESPONSE:
+		case PGRES_NONFATAL_ERROR:
+		case PGRES_FATAL_ERROR:
+			success = false;
+			break;
+
 		default:
+			success = false;
+			psql_error("unexpected PQresultStatus (%d)",
+					   PQresultStatus(results));
 			break;
 	}
 
@@ -867,7 +956,7 @@ SendQuery(const char *query)
 
 		/* these operations are included in the timing result: */
 		ResetCancelConn();
-		OK = (AcceptResult(results) && ProcessCopyResult(results));
+		OK = ProcessResult(&results);
 
 		if (pset.timing)
 		{
@@ -877,7 +966,7 @@ SendQuery(const char *query)
 		}
 
 		/* but printing results isn't: */
-		if (OK)
+		if (OK && results)
 			OK = PrintQueryResults(results);
 	}
 	else
@@ -891,26 +980,27 @@ SendQuery(const char *query)
 	/* If we made a temporary savepoint, possibly release/rollback */
 	if (on_error_rollback_savepoint)
 	{
-		const char *svptcmd;
+		const char *svptcmd = NULL;
 
 		transaction_status = PQtransactionStatus(pset.db);
 
-		if (transaction_status == PQTRANS_INERROR)
+		switch (transaction_status)
 		{
+			case PQTRANS_INERROR:
 				/* We always rollback on an error */
 				svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint";
-		}
-		else if (transaction_status != PQTRANS_INTRANS)
-		{
+				break;
+
+			case PQTRANS_IDLE:
 				/* If they are no longer in a transaction, then do nothing */
-			svptcmd = NULL;
-		}
-		else
-		{
+				break;
+
+			case PQTRANS_INTRANS:
 				/*
-			 * Do nothing if they are messing with savepoints themselves: If
-			 * the user did RELEASE or ROLLBACK, our savepoint is gone. If
-			 * they issued a SAVEPOINT, releasing ours would remove theirs.
+				 * Do nothing if they are messing with savepoints themselves:
+				 * If the user did RELEASE or ROLLBACK, our savepoint is
+				 * gone. If they issued a SAVEPOINT, releasing ours would
+				 * remove theirs.
 				 */
 				if (results &&
 					(strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
@@ -919,6 +1009,15 @@ SendQuery(const char *query)
 					svptcmd = NULL;
 				else
 					svptcmd = "RELEASE pg_psql_temporary_savepoint";
+				break;
+
+			case PQTRANS_ACTIVE:
+			case PQTRANS_UNKNOWN:
+			default:
+				OK = false;
+				psql_error("unexpected transaction status (%d)\n",
+						   transaction_status);
+				break;
 		}
 
 		if (svptcmd)

src/bin/psql/copy.c

@@ -244,8 +244,9 @@ do_copy(const char *args)
 {
 	PQExpBufferData query;
 	FILE	   *copystream;
+	FILE	   *save_file;
+	FILE	  **override_file;
 	struct copy_options *options;
-	PGresult   *result;
 	bool		success;
 	struct stat st;
 
@@ -261,6 +262,8 @@ do_copy(const char *args)
 
 	if (options->from)
 	{
+		override_file = &pset.cur_cmd_source;
+
 		if (options->file)
 			copystream = fopen(options->file, PG_BINARY_R);
 		else if (!options->psql_inout)
@@ -270,6 +273,8 @@ do_copy(const char *args)
 	}
 	else
 	{
+		override_file = &pset.queryFout;
+
 		if (options->file)
 			copystream = fopen(options->file, PG_BINARY_W);
 		else if (!options->psql_inout)
@@ -308,52 +313,13 @@ do_copy(const char *args)
 	if (options->after_tofrom)
 		appendPQExpBufferStr(&query, options->after_tofrom);
 
-	result = PSQLexec(query.data, true);
+	/* Run it like a user command, interposing the data source or sink. */
+	save_file = *override_file;
+	*override_file = copystream;
+	success = SendQuery(query.data);
+	*override_file = save_file;
 	termPQExpBuffer(&query);
 
-	switch (PQresultStatus(result))
-	{
-		case PGRES_COPY_OUT:
-			SetCancelConn();
-			success = handleCopyOut(pset.db, copystream);
-			ResetCancelConn();
-			break;
-		case PGRES_COPY_IN:
-			SetCancelConn();
-			success = handleCopyIn(pset.db, copystream,
-								   PQbinaryTuples(result));
-			ResetCancelConn();
-			break;
-		case PGRES_NONFATAL_ERROR:
-		case PGRES_FATAL_ERROR:
-		case PGRES_BAD_RESPONSE:
-			success = false;
-			psql_error("\\copy: %s", PQerrorMessage(pset.db));
-			break;
-		default:
-			success = false;
-			psql_error("\\copy: unexpected response (%d)\n",
-					   PQresultStatus(result));
-			break;
-	}
-
-	PQclear(result);
-
-	/*
-	 * Make sure we have pumped libpq dry of results; else it may still be in
-	 * ASYNC_BUSY state, leading to false readings in, eg, get_prompt().
-	 */
-	while ((result = PQgetResult(pset.db)) != NULL)
-	{
-		success = false;
-		psql_error("\\copy: unexpected response (%d)\n",
-				   PQresultStatus(result));
-		/* if still in COPY IN state, try to get out of it */
-		if (PQresultStatus(result) == PGRES_COPY_IN)
-			PQputCopyEnd(pset.db, _("trying to exit copy mode"));
-		PQclear(result);
-	}
-
 	if (options->file != NULL)
 	{
 		if (fclose(copystream) != 0)
@@ -425,8 +391,30 @@ handleCopyOut(PGconn *conn, FILE *copystream)
 		OK = false;
 	}
 
-	/* Check command status and return to normal libpq state */
-	res = PQgetResult(conn);
+	/*
+	 * Check command status and return to normal libpq state.  After a
+	 * client-side error, the server will remain ready to deliver data.  The
+	 * cleanest thing is to fully drain and discard that data.  If the
+	 * client-side error happened early in a large file, this takes a long
+	 * time.  Instead, take advantage of the fact that PQexec() will silently
+	 * end any ongoing PGRES_COPY_OUT state.  This does cause us to lose the
+	 * results of any commands following the COPY in a single command string.
+	 * It also only works for protocol version 3.  XXX should we clean up
+	 * using the slow way when the connection is using protocol version 2?
+	 *
+	 * We must not ever return with the status still PGRES_COPY_OUT.  Our
+	 * caller is unable to distinguish that situation from reaching the next
+	 * COPY in a command string that happened to contain two consecutive COPY
+	 * TO STDOUT commands.  We trust that no condition can make PQexec() fail
+	 * indefinitely while retaining status PGRES_COPY_OUT.
+	 */
+	while (res = PQgetResult(conn), PQresultStatus(res) == PGRES_COPY_OUT)
+	{
+		OK = false;
+		PQclear(res);
+
+		PQexec(conn, "-- clear PGRES_COPY_OUT state");
+	}
 	if (PQresultStatus(res) != PGRES_COMMAND_OK)
 	{
 		psql_error("%s", PQerrorMessage(conn));
@@ -471,13 +459,8 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary)
 		/* Terminate data transfer */
 		PQputCopyEnd(conn, _("canceled by user"));
 
-		/* Check command status and return to normal libpq state */
-		res = PQgetResult(conn);
-		if (PQresultStatus(res) != PGRES_COMMAND_OK)
-			psql_error("%s", PQerrorMessage(conn));
-		PQclear(res);
-
-		return false;
+		OK = false;
+		goto copyin_cleanup;
 	}
 
 	/* Prompt if interactive input */
@@ -600,8 +583,24 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary)
 					 OK ? NULL : _("aborted because of read failure")) <= 0)
 		OK = false;
 
-	/* Check command status and return to normal libpq state */
-	res = PQgetResult(conn);
+copyin_cleanup:
+	/*
+	 * Check command status and return to normal libpq state
+	 *
+	 * We must not ever return with the status still PGRES_COPY_IN.  Our
+	 * caller is unable to distinguish that situation from reaching the next
+	 * COPY in a command string that happened to contain two consecutive COPY
+	 * FROM STDIN commands.  XXX if something makes PQputCopyEnd() fail
+	 * indefinitely while retaining status PGRES_COPY_IN, we get an infinite
+	 * loop.  This is more realistic than handleCopyOut()'s counterpart risk.
+	 */
+	while (res = PQgetResult(conn), PQresultStatus(res) == PGRES_COPY_IN)
+	{
+		OK = false;
+		PQclear(res);
+
+		PQputCopyEnd(pset.db, _("trying to exit copy mode"));
+	}
 	if (PQresultStatus(res) != PGRES_COMMAND_OK)
 	{
 		psql_error("%s", PQerrorMessage(conn));

src/test/regress/expected/copyselect.out

@@ -111,8 +111,6 @@ t
 \copy v_test1 to stdout
 ERROR:  cannot copy from view "v_test1"
 HINT:  Try the COPY (SELECT ...) TO variant.
-\copy: ERROR:  cannot copy from view "v_test1"
-HINT:  Try the COPY (SELECT ...) TO variant.
 --
 -- Test \copy (select ...)
 --
@@ -124,3 +122,32 @@ HINT:  Try the COPY (SELECT ...) TO variant.
 drop table test2;
 drop view v_test1;
 drop table test1;
+-- psql handling of COPY in multi-command strings
+copy (select 1) to stdout\; select 1/0;	-- row, then error
+1
+ERROR:  division by zero
+select 1/0\; copy (select 1) to stdout; -- error only
+ERROR:  division by zero
+copy (select 1) to stdout\; copy (select 2) to stdout\; select 0\; select 3; -- 1 2 3
+1
+2
+ ?column? 
+----------
+        3
+(1 row)
+
+create table test3 (c int);
+select 0\; copy test3 from stdin\; copy test3 from stdin\; select 1; -- 1
+ ?column? 
+----------
+        1
+(1 row)
+
+select * from test3;
+ c 
+---
+ 1
+ 2
+(2 rows)
+
+drop table test3;

src/test/regress/sql/copyselect.sql

@@ -80,3 +80,17 @@ copy (select t from test1 where id = 1) to stdout csv header force quote t;
 drop table test2;
 drop view v_test1;
 drop table test1;
+
+-- psql handling of COPY in multi-command strings
+copy (select 1) to stdout\; select 1/0;	-- row, then error
+select 1/0\; copy (select 1) to stdout; -- error only
+copy (select 1) to stdout\; copy (select 2) to stdout\; select 0\; select 3; -- 1 2 3
+
+create table test3 (c int);
+select 0\; copy test3 from stdin\; copy test3 from stdin\; select 1; -- 1
+1
+\.
+2
+\.
+select * from test3;
+drop table test3;