1 files changed, 99 insertions, 85 deletions

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 3254a140b3a..a287eeee195 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -26,72 +26,90 @@
 #include "mbprint.h"
 
 
-
 static bool ExecQueryUsingCursor(const char *query, double *elapsed_msec);
 static bool command_no_begin(const char *query);
 static bool is_select_command(const char *query);
 
+
 /*
- * setQFout
- * -- handler for -o command line option and \o command
+ * openQueryOutputFile --- attempt to open a query output file
  *
- * Tries to open file fname (or pipe if fname starts with '|')
- * and stores the file handle in pset)
- * Upon failure, sets stdout and returns false.
+ * fname == NULL selects stdout, else an initial '|' selects a pipe,
+ * else plain file.
+ *
+ * Returns output file pointer into *fout, and is-a-pipe flag into *is_pipe.
+ * Caller is responsible for adjusting SIGPIPE state if it's a pipe.
+ *
+ * On error, reports suitable error message and returns FALSE.
  */
 bool
-setQFout(const char *fname)
+openQueryOutputFile(const char *fname, FILE **fout, bool *is_pipe)
 {
-	bool		status = true;
-
-	/* Close old file/pipe */
-	if (pset.queryFout && pset.queryFout != stdout && pset.queryFout != stderr)
-	{
-		if (pset.queryFoutPipe)
-			pclose(pset.queryFout);
-		else
-			fclose(pset.queryFout);
-	}
-
-	/* If no filename, set stdout */
 	if (!fname || fname[0] == '\0')
 	{
-		pset.queryFout = stdout;
-		pset.queryFoutPipe = false;
+		*fout = stdout;
+		*is_pipe = false;
 	}
 	else if (*fname == '|')
 	{
-		pset.queryFout = popen(fname + 1, "w");
-		pset.queryFoutPipe = true;
+		*fout = popen(fname + 1, "w");
+		*is_pipe = true;
 	}
 	else
 	{
-		pset.queryFout = fopen(fname, "w");
-		pset.queryFoutPipe = false;
+		*fout = fopen(fname, "w");
+		*is_pipe = false;
 	}
 
-	if (!(pset.queryFout))
+	if (*fout == NULL)
 	{
 		psql_error("%s: %s\n", fname, strerror(errno));
-		pset.queryFout = stdout;
-		pset.queryFoutPipe = false;
-		status = false;
+		return false;
 	}
 
-	/* Direct signals */
-#ifndef WIN32
-	pqsignal(SIGPIPE, pset.queryFoutPipe ? SIG_IGN : SIG_DFL);
-#endif
-
-	return status;
+	return true;
 }
 
+/*
+ * setQFout
+ * -- handler for -o command line option and \o command
+ *
+ * On success, updates pset with the new output file and returns true.
+ * On failure, returns false without changing pset state.
+ */
+bool
+setQFout(const char *fname)
+{
+	FILE	   *fout;
+	bool		is_pipe;
+
+	/* First make sure we can open the new output file/pipe */
+	if (!openQueryOutputFile(fname, &fout, &is_pipe))
+		return false;
+
+	/* Close old file/pipe */
+	if (pset.queryFout && pset.queryFout != stdout && pset.queryFout != stderr)
+	{
+		if (pset.queryFoutPipe)
+			pclose(pset.queryFout);
+		else
+			fclose(pset.queryFout);
+	}
+
+	pset.queryFout = fout;
+	pset.queryFoutPipe = is_pipe;
+
+	/* Adjust SIGPIPE handling appropriately: ignore signal if is_pipe */
+	set_sigpipe_trap_state(is_pipe);
+	restore_sigpipe_trap();
+
+	return true;
+}
 
 
 /*
  * Error reporting for scripts. Errors should look like
  *	 psql:filename:lineno: message
- *
  */
 void
 psql_error(const char *fmt,...)
@@ -611,27 +629,23 @@ PrintQueryTuples(const PGresult *results)
 	/* write output to \g argument, if any */
 	if (pset.gfname)
 	{
-		/* keep this code in sync with ExecQueryUsingCursor */
-		FILE	   *queryFout_copy = pset.queryFout;
-		bool		queryFoutPipe_copy = pset.queryFoutPipe;
-
-		pset.queryFout = stdout;	/* so it doesn't get closed */
+		FILE	   *fout;
+		bool		is_pipe;
 
-		/* open file/pipe */
-		if (!setQFout(pset.gfname))
-		{
-			pset.queryFout = queryFout_copy;
-			pset.queryFoutPipe = queryFoutPipe_copy;
+		if (!openQueryOutputFile(pset.gfname, &fout, &is_pipe))
 			return false;
-		}
-
-		printQuery(results, &my_popt, pset.queryFout, false, pset.logfile);
+		if (is_pipe)
+			disable_sigpipe_trap();
 
-		/* close file/pipe, restore old setting */
-		setQFout(NULL);
+		printQuery(results, &my_popt, fout, false, pset.logfile);
 
-		pset.queryFout = queryFout_copy;
-		pset.queryFoutPipe = queryFoutPipe_copy;
+		if (is_pipe)
+		{
+			pclose(fout);
+			restore_sigpipe_trap();
+		}
+		else
+			fclose(fout);
 	}
 	else
 		printQuery(results, &my_popt, pset.queryFout, false, pset.logfile);
@@ -1199,10 +1213,10 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 	PGresult   *results;
 	PQExpBufferData buf;
 	printQueryOpt my_popt = pset.popt;
-	FILE	   *queryFout_copy = pset.queryFout;
-	bool		queryFoutPipe_copy = pset.queryFoutPipe;
+	FILE	   *fout;
+	bool		is_pipe;
+	bool		is_pager = false;
 	bool		started_txn = false;
-	bool		did_pager = false;
 	int			ntuples;
 	int			fetch_count;
 	char		fetch_cmd[64];
@@ -1268,21 +1282,22 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 	/* prepare to write output to \g argument, if any */
 	if (pset.gfname)
 	{
-		/* keep this code in sync with PrintQueryTuples */
-		pset.queryFout = stdout;	/* so it doesn't get closed */
-
-		/* open file/pipe */
-		if (!setQFout(pset.gfname))
+		if (!openQueryOutputFile(pset.gfname, &fout, &is_pipe))
 		{
-			pset.queryFout = queryFout_copy;
-			pset.queryFoutPipe = queryFoutPipe_copy;
 			OK = false;
 			goto cleanup;
 		}
+		if (is_pipe)
+			disable_sigpipe_trap();
+	}
+	else
+	{
+		fout = pset.queryFout;
+		is_pipe = false;		/* doesn't matter */
 	}
 
 	/* clear any pre-existing error indication on the output stream */
-	clearerr(pset.queryFout);
+	clearerr(fout);
 
 	for (;;)
 	{
@@ -1302,12 +1317,10 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		if (PQresultStatus(results) != PGRES_TUPLES_OK)
 		{
 			/* shut down pager before printing error message */
-			if (did_pager)
+			if (is_pager)
 			{
-				ClosePager(pset.queryFout);
-				pset.queryFout = queryFout_copy;
-				pset.queryFoutPipe = queryFoutPipe_copy;
-				did_pager = false;
+				ClosePager(fout);
+				is_pager = false;
 			}
 
 			OK = AcceptResult(results);
@@ -1331,17 +1344,17 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 			/* this is the last result set, so allow footer decoration */
 			my_popt.topt.stop_table = true;
 		}
-		else if (pset.queryFout == stdout && !did_pager)
+		else if (fout == stdout && !is_pager)
 		{
 			/*
 			 * If query requires multiple result sets, hack to ensure that
 			 * only one pager instance is used for the whole mess
 			 */
-			pset.queryFout = PageOutput(INT_MAX, &(my_popt.topt));
-			did_pager = true;
+			fout = PageOutput(INT_MAX, &(my_popt.topt));
+			is_pager = true;
 		}
 
-		printQuery(results, &my_popt, pset.queryFout, did_pager, pset.logfile);
+		printQuery(results, &my_popt, fout, is_pager, pset.logfile);
 
 		PQclear(results);
 
@@ -1355,7 +1368,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		 * the pager dies/exits/etc, there's no sense throwing more data at
 		 * it.
 		 */
-		flush_error = fflush(pset.queryFout);
+		flush_error = fflush(fout);
 
 		/*
 		 * Check if we are at the end, if a cancel was pressed, or if there
@@ -1365,24 +1378,25 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		 * stop bothering to pull down more data.
 		 */
 		if (ntuples < fetch_count || cancel_pressed || flush_error ||
-			ferror(pset.queryFout))
+			ferror(fout))
 			break;
 	}
 
-	/* close \g argument file/pipe, restore old setting */
 	if (pset.gfname)
 	{
-		/* keep this code in sync with PrintQueryTuples */
-		setQFout(NULL);
-
-		pset.queryFout = queryFout_copy;
-		pset.queryFoutPipe = queryFoutPipe_copy;
+		/* close \g argument file/pipe */
+		if (is_pipe)
+		{
+			pclose(fout);
+			restore_sigpipe_trap();
+		}
+		else
+			fclose(fout);
 	}
-	else if (did_pager)
+	else if (is_pager)
 	{
-		ClosePager(pset.queryFout);
-		pset.queryFout = queryFout_copy;
-		pset.queryFoutPipe = queryFoutPipe_copy;
+		/* close transient pager */
+		ClosePager(fout);
 	}
 
 cleanup: