psql COPY vs. ON_ERROR_ROLLBACK, multi-command strings

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 69fac83..c7d024a 100644
*** a/src/bin/psql/command.c
--- b/src/bin/psql/command.c
***************
*** 320,344 **** 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);
  	}
  
--- 320,329 ----
diff --git a/src/bin/psql/comindex 889c157..6eaf794 100644
*** a/src/bin/psql/common.c
--- b/src/bin/psql/common.c
***************
*** 438,444 **** ResetCancelConn(void)
  static bool
  AcceptResult(const PGresult *result)
  {
! 	bool		OK = true;
  
  	if (!result)
  		OK = false;
--- 438,444 ----
  static bool
  AcceptResult(const PGresult *result)
  {
! 	bool		OK;
  
  	if (!result)
  		OK = false;
***************
*** 450,460 **** AcceptResult(const PGresult *result)
--- 450,470 ----
  			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,664 **** PrintQueryTuples(const PGresult *results)
  
  
  /*
!  * ProcessCopyResult: if command was a COPY FROM STDIN/TO STDOUT, handle it
   *
!  * Note: Utility function for use by SendQuery() only.
   *
!  * Returns true if the query executed successfully, false otherwise.
   */
  static bool
! ProcessCopyResult(PGresult *results)
  {
! 	bool		success = false;
  
! 	if (!results)
! 		return false;
! 
! 	switch (PQresultStatus(results))
  	{
! 		case PGRES_TUPLES_OK:
! 		case PGRES_COMMAND_OK:
! 		case PGRES_EMPTY_QUERY:
! 			/* nothing to do here */
! 			success = true;
! 			break;
  
! 		case PGRES_COPY_OUT:
! 			SetCancelConn();
! 			success = handleCopyOut(pset.db, pset.queryFout);
! 			ResetCancelConn();
  			break;
  
! 		case PGRES_COPY_IN:
  			SetCancelConn();
! 			success = handleCopyIn(pset.db, pset.cur_cmd_source,
! 								   PQbinaryTuples(results));
  			ResetCancelConn();
- 			break;
  
! 		default:
  			break;
! 	}
  
  	/* may need this to recover from conn loss during COPY */
  	if (!CheckConnection())
--- 630,743 ----
  
  
  /*
!  * ProcessResult: 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.
   *
!  * 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
! ProcessResult(PGresult **results)
  {
! 	PGresult   *next_result;
! 	bool		success = true;
! 	bool		first_cycle = true;
  
! 	do
  	{
! 		ExecStatusType result_status;
! 		bool		is_copy;
  
! 		if (!AcceptResult(*results))
! 		{
! 			/*
! 			 * 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:
! 			case PGRES_COPY_IN:
! 				is_copy = true;
! 				break;
! 
! 			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)) && 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.
! 			 */
! 			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())
***************
*** 708,714 **** PrintQueryStatus(PGresult *results)
  static bool
  PrintQueryResults(PGresult *results)
  {
! 	bool		success = false;
  	const char *cmdstatus;
  
  	if (!results)
--- 787,793 ----
  static bool
  PrintQueryResults(PGresult *results)
  {
! 	bool		success;
  	const char *cmdstatus;
  
  	if (!results)
***************
*** 738,748 **** PrintQueryResults(PGresult *results)
--- 817,837 ----
  
  		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,873 **** SendQuery(const char *query)
  
  		/* these operations are included in the timing result: */
  		ResetCancelConn();
! 		OK = (AcceptResult(results) && ProcessCopyResult(results));
  
  		if (pset.timing)
  		{
--- 956,962 ----
  
  		/* these operations are included in the timing result: */
  		ResetCancelConn();
! 		OK = ProcessResult(&results);
  
  		if (pset.timing)
  		{
***************
*** 877,883 **** SendQuery(const char *query)
  		}
  
  		/* but printing results isn't: */
! 		if (OK)
  			OK = PrintQueryResults(results);
  	}
  	else
--- 966,972 ----
  		}
  
  		/* but printing results isn't: */
! 		if (OK && results)
  			OK = PrintQueryResults(results);
  	}
  	else
***************
*** 891,924 **** SendQuery(const char *query)
  	/* If we made a temporary savepoint, possibly release/rollback */
  	if (on_error_rollback_savepoint)
  	{
! 		const char *svptcmd;
  
  		transaction_status = PQtransactionStatus(pset.db);
  
! 		if (transaction_status == PQTRANS_INERROR)
! 		{
! 			/* We always rollback on an error */
! 			svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint";
! 		}
! 		else if (transaction_status != PQTRANS_INTRANS)
  		{
! 			/* If they are no longer in a transaction, then do nothing */
! 			svptcmd = NULL;
! 		}
! 		else
! 		{
! 			/*
! 			 * 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 ||
! 				 strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
! 				 strcmp(PQcmdStatus(results), "ROLLBACK") == 0))
! 				svptcmd = NULL;
! 			else
! 				svptcmd = "RELEASE pg_psql_temporary_savepoint";
  		}
  
  		if (svptcmd)
--- 980,1023 ----
  	/* If we made a temporary savepoint, possibly release/rollback */
  	if (on_error_rollback_savepoint)
  	{
! 		const char *svptcmd = NULL;
  
  		transaction_status = PQtransactionStatus(pset.db);
  
! 		switch (transaction_status)
  		{
! 			case PQTRANS_INERROR:
! 				/* We always rollback on an error */
! 				svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint";
! 				break;
! 
! 			case PQTRANS_IDLE:
! 				/* If they are no longer in a transaction, then do nothing */
! 				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.
! 				 */
! 				if (results &&
! 					(strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
! 					 strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
! 					 strcmp(PQcmdStatus(results), "ROLLBACK") == 0))
! 					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)
diff --git a/src/bin/psql/coindex eaf633d..a1dea95 100644
*** a/src/bin/psql/copy.c
--- b/src/bin/psql/copy.c
***************
*** 244,251 **** do_copy(const char *args)
  {
  	PQExpBufferData query;
  	FILE	   *copystream;
  	struct copy_options *options;
- 	PGresult   *result;
  	bool		success;
  	struct stat st;
  
--- 244,252 ----
  {
  	PQExpBufferData query;
  	FILE	   *copystream;
+ 	FILE	   *save_file;
+ 	FILE	  **override_file;
  	struct copy_options *options;
  	bool		success;
  	struct stat st;
  
***************
*** 261,266 **** do_copy(const char *args)
--- 262,269 ----
  
  	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,275 **** do_copy(const char *args)
--- 273,280 ----
  	}
  	else
  	{
+ 		override_file = &pset.queryFout;
+ 
  		if (options->file)
  			copystream = fopen(options->file, PG_BINARY_W);
  		else if (!options->psql_inout)
***************
*** 308,359 **** do_copy(const char *args)
  	if (options->after_tofrom)
  		appendPQExpBufferStr(&query, options->after_tofrom);
  
! 	result = PSQLexec(query.data, true);
  	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)
--- 313,325 ----
  	if (options->after_tofrom)
  		appendPQExpBufferStr(&query, options->after_tofrom);
  
! 	/* 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);
  
  	if (options->file != NULL)
  	{
  		if (fclose(copystream) != 0)
***************
*** 425,432 **** handleCopyOut(PGconn *conn, FILE *copystream)
  		OK = false;
  	}
  
! 	/* Check command status and return to normal libpq state */
! 	res = PQgetResult(conn);
  	if (PQresultStatus(res) != PGRES_COMMAND_OK)
  	{
  		psql_error("%s", PQerrorMessage(conn));
--- 391,420 ----
  		OK = false;
  	}
  
! 	/*
! 	 * 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,483 **** 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;
  	}
  
  	/* Prompt if interactive input */
--- 459,466 ----
  		/* Terminate data transfer */
  		PQputCopyEnd(conn, _("canceled by user"));
  
! 		OK = false;
! 		goto copyin_cleanup;
  	}
  
  	/* Prompt if interactive input */
***************
*** 600,607 **** 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);
  	if (PQresultStatus(res) != PGRES_COMMAND_OK)
  	{
  		psql_error("%s", PQerrorMessage(conn));
--- 583,606 ----
  					 OK ? NULL : _("aborted because of read failure")) <= 0)
  		OK = false;
  
! 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));
diff --git a/src/test/regrindex cbc140c..7c16a61 100644
*** a/src/test/regress/expected/copyselect.out
--- b/src/test/regress/expected/copyselect.out
***************
*** 111,118 **** 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 ...)
  --
--- 111,116 ----
***************
*** 124,126 **** HINT:  Try the COPY (SELECT ...) TO variant.
--- 122,153 ----
  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;
diff --git a/src/test/regress/sql/copyselect.sqindex 621d494..1d98dad 100644
*** a/src/test/regress/sql/copyselect.sql
--- b/src/test/regress/sql/copyselect.sql
***************
*** 80,82 **** copy (select t from test1 where id = 1) to stdout csv header force quote t;
--- 80,96 ----
  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;