review: psql command copy count tag

*** ./src/bin/psql/common.c.orig	2014-01-29 21:09:07.108862537 +0100
--- ./src/bin/psql/common.c	2014-01-29 21:09:42.810920907 +0100
***************
*** 631,638 ****
   * 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
--- 631,637 ----
   * 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.
   *
   * Returns true on complete success, false otherwise.  Possible failure modes
   * include purely client-side problems; check the transaction status for the
***************
*** 689,709 ****
  			 * 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 */
--- 688,723 ----
  			 * connection out of its COPY state, then call PQresultStatus()
  			 * once and report any error.
  			 */
+ 
+ 			/*
+ 			 * If this is second copy; then it will be definately not \copy,
+ 			 * and also it can not be from any user given file. So pass the
+ 			 * value of copystream as NULL, so that read/wrie happens on
+ 			 * already set cur_cmd_source/queryFout.
+ 			 */
  			SetCancelConn();
  			if (result_status == PGRES_COPY_OUT)
! 				success = handleCopyOut(pset.db,
! 								 (first_cycle ? pset.copyStream : NULL),
! 											 results) && success;
  			else
! 				success = handleCopyIn(pset.db,
! 								 (first_cycle ? pset.copyStream : NULL),
! 									   PQbinaryTuples(*results),
! 											 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.
+ 			 * It keeps last result.
  			 */
! 			if ((next_result = PQgetResult(pset.db)))
! 			{
! 				PQclear(*results);
! 				*results = next_result;
! 			}
  		}
  		else if (first_cycle)
  			/* fast path: no COPY commands; PQexec visited all results */
*** ./src/bin/psql/copy.c.orig	2014-01-29 21:09:15.131875660 +0100
--- ./src/bin/psql/copy.c	2014-01-29 21:09:42.811920909 +0100
***************
*** 269,276 ****
  {
  	PQExpBufferData query;
  	FILE	   *copystream;
- 	FILE	   *save_file;
- 	FILE	  **override_file;
  	struct copy_options *options;
  	bool		success;
  	struct stat st;
--- 269,274 ----
***************
*** 287,294 ****
  
  	if (options->from)
  	{
- 		override_file = &pset.cur_cmd_source;
- 
  		if (options->file)
  		{
  			if (options->program)
--- 285,290 ----
***************
*** 308,315 ****
  	}
  	else
  	{
- 		override_file = &pset.queryFout;
- 
  		if (options->file)
  		{
  			if (options->program)
--- 304,309 ----
***************
*** 369,378 ****
  		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)
--- 363,371 ----
  		appendPQExpBufferStr(&query, options->after_tofrom);
  
  	/* Run it like a user command, interposing the data source or sink. */
! 	pset.copyStream = copystream;
  	success = SendQuery(query.data);
! 	pset.copyStream = NULL;
  	termPQExpBuffer(&query);
  
  	if (options->file != NULL)
***************
*** 433,444 ****
   * result is true if successful, false if not.
   */
  bool
! handleCopyOut(PGconn *conn, FILE *copystream)
  {
  	bool		OK = true;
  	char	   *buf;
  	int			ret;
! 	PGresult   *res;
  
  	for (;;)
  	{
--- 426,439 ----
   * result is true if successful, false if not.
   */
  bool
! handleCopyOut(PGconn *conn, FILE *copystream, PGresult **res)
  {
  	bool		OK = true;
  	char	   *buf;
  	int			ret;
! 
! 	if (!copystream)
! 		copystream = pset.queryFout;
  
  	for (;;)
  	{
***************
*** 490,508 ****
  	 * 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));
  		OK = false;
  	}
- 	PQclear(res);
  
  	return OK;
  }
--- 485,502 ----
  	 * 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));
  		OK = false;
  	}
  
  	return OK;
  }
***************
*** 523,534 ****
  #define COPYBUFSIZ 8192
  
  bool
! handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary)
  {
  	bool		OK;
  	const char *prompt;
  	char		buf[COPYBUFSIZ];
! 	PGresult   *res;
  
  	/*
  	 * Establish longjmp destination for exiting from wait-for-input. (This is
--- 517,530 ----
  #define COPYBUFSIZ 8192
  
  bool
! handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res)
  {
  	bool		OK;
  	const char *prompt;
  	char		buf[COPYBUFSIZ];
! 
! 	if (!copystream)
! 		copystream = pset.cur_cmd_source;
  
  	/*
  	 * Establish longjmp destination for exiting from wait-for-input. (This is
***************
*** 656,663 ****
  				}
  			}
  
! 			if (copystream == pset.cur_cmd_source)
! 				pset.lineno++;
  		}
  	}
  
--- 652,658 ----
  				}
  			}
  
! 			pset.lineno++;
  		}
  	}
  
***************
*** 682,700 ****
  	 * 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));
  		OK = false;
  	}
- 	PQclear(res);
  
  	return OK;
  }
--- 677,694 ----
  	 * 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));
  		OK = false;
  	}
  
  	return OK;
  }
*** ./src/bin/psql/copy.h.orig	2014-01-29 21:09:21.264885699 +0100
--- ./src/bin/psql/copy.h	2014-01-29 21:09:42.811920909 +0100
***************
*** 16,22 ****
  
  /* lower level processors for copy in/out streams */
  
! bool		handleCopyOut(PGconn *conn, FILE *copystream);
! bool		handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary);
  
  #endif
--- 16,22 ----
  
  /* lower level processors for copy in/out streams */
  
! bool		handleCopyOut(PGconn *conn, FILE *copystream, PGresult **res);
! bool		handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res);
  
  #endif
*** ./src/bin/psql/settings.h.orig	2014-01-29 21:09:28.060896814 +0100
--- ./src/bin/psql/settings.h	2014-01-29 21:09:42.811920909 +0100
***************
*** 70,75 ****
--- 70,77 ----
  	FILE	   *queryFout;		/* where to send the query results */
  	bool		queryFoutPipe;	/* queryFout is from a popen() */
  
+ 	FILE	   *copyStream;		/* Stream to read/write for copy command */
+ 
  	printQueryOpt popt;
  
  	char	   *gfname;			/* one-shot file output argument for \g */