psql-copyout-g.patch

text/plain

Filename: psql-copyout-g.patch
Type: text/plain
Part: 0
Message: Re: Alternative to \copy in psql modelled after \g

Patch

Format: unified
File+
src/bin/psql/common.c 35 4
src/bin/psql/copy.c 2 1
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 62c2928..1488bb9 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -1095,14 +1095,31 @@ ProcessResult(PGresult **results)
 			 * If pset.copyStream is set, use that as data source/sink,
 			 * otherwise use queryFout or cur_cmd_source as appropriate.
 			 */
-			FILE	   *copystream = pset.copyStream;
+			FILE	   *copystream = NULL;
 			PGresult   *copy_result;
 
 			SetCancelConn();
 			if (result_status == PGRES_COPY_OUT)
 			{
-				if (!copystream)
+				bool is_pipe;
+				if (pset.gfname)
+				{
+					/*
+					 * COPY TO STDOUT \g [|]file may be used as an alternative
+					 * to \copy
+					 */
+					if (!openQueryOutputFile(pset.gfname, &copystream, &is_pipe))
+					{
+						copystream = NULL; /* will discard the COPY data entirely */
+					}
+					if (is_pipe)
+						disable_sigpipe_trap();
+				}
+				else if (pset.copyStream)
+					copystream = pset.copyStream;
+				else
 					copystream = pset.queryFout;
+
 				success = handleCopyOut(pset.db,
 										copystream,
 										&copy_result) && success;
@@ -1117,11 +1134,25 @@ ProcessResult(PGresult **results)
 					PQclear(copy_result);
 					copy_result = NULL;
 				}
+
+				if (pset.gfname && copystream != NULL)
+				{
+					/* close \g argument file/pipe */
+					if (is_pipe)
+					{
+						pclose(copystream);
+						restore_sigpipe_trap();
+					}
+					else
+					{
+						fclose(copystream);
+					}
+				}
 			}
 			else
 			{
-				if (!copystream)
-					copystream = pset.cur_cmd_source;
+				/* COPY IN */
+				copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source;
 				success = handleCopyIn(pset.db,
 									   copystream,
 									   PQbinaryTuples(*results),
diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c
index 555c633..645970e 100644
--- a/src/bin/psql/copy.c
+++ b/src/bin/psql/copy.c
@@ -426,6 +426,7 @@ 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.
+ * copystream can be NULL to pump the data without writing it anywhere.
  * The final status for the COPY is returned into *res (but note
  * we already reported the error, if it's not a success result).
  *
@@ -434,7 +435,7 @@ do_copy(const char *args)
 bool
 handleCopyOut(PGconn *conn, FILE *copystream, PGresult **res)
 {
-	bool		OK = true;
+	bool		OK = (copystream != NULL);
 	char	   *buf;
 	int			ret;