psql-copyout-g-v2.patch

text/plain

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

Patch

Format: unified
Series: patch v2
File+
doc/src/sgml/ref/psql-ref.sgml 11 0
src/bin/psql/common.c 50 7
src/bin/psql/copy.c 3 2
diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index 6c76cf2..8d17708 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -1044,6 +1044,17 @@ testdb=>
         </para>
         </tip>
 
+        <tip>
+        <para>
+        As an alternative to <literal>\copy ... to 'filename' | program 'command'</literal>,
+        an equivalent <acronym>SQL</acronym> <literal>COPY ... TO STDOUT</literal> command
+        terminated by <literal>\g filename</literal> or <literal>\g |program</literal>
+        may be used to the same effect. The latter form has the advantages that it
+        can span multiple lines and allows for variable interpolation, both in the query
+        and in the argument following <literal>\g</literal>.
+        </para>
+        </tip>
+
         </listitem>
       </varlistentry>
 
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index b11d7ac6..8d2fa59 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -1092,20 +1092,49 @@ ProcessResult(PGresult **results)
 			 * connection out of its COPY state, then call PQresultStatus()
 			 * once and report any error.
 			 *
-			 * If pset.copyStream is set, use that as data source/sink,
-			 * otherwise use queryFout or cur_cmd_source as appropriate.
+			 * For COPY OUT, direct the output to pset.copyStream if it's set,
+			 * otherwise to pset.gfname if it's set, otherwise to queryFout.
+			 * For COPY IN, use pset.copyStream as data source if it's set,
+			 * otherwise cur_cmd_source.
 			 */
-			FILE	   *copystream = pset.copyStream;
+			FILE	   *copystream;
 			PGresult   *copy_result;
 
 			SetCancelConn();
 			if (result_status == PGRES_COPY_OUT)
 			{
-				if (!copystream)
+				bool is_pipe;
+
+				if (pset.copyStream)
+				{
+					/* instantiated by \copy */
+					copystream = pset.copyStream;
+				}
+				else 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 */
+						is_pipe = false;
+					}
+					if (is_pipe)
+						disable_sigpipe_trap();
+				}
+				else
+				{
+					/* fall back to the generic query output stream */
 					copystream = pset.queryFout;
+				}
+
 				success = handleCopyOut(pset.db,
 										copystream,
-										&copy_result) && success;
+										&copy_result)
+					&& success
+					&& (copystream != NULL);
 
 				/*
 				 * Suppress status printing if the report would go to the same
@@ -1117,11 +1146,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 02c8511..0f3ea33 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).
  *
@@ -447,7 +448,7 @@ handleCopyOut(PGconn *conn, FILE *copystream, PGresult **res)
 
 		if (buf)
 		{
-			if (OK && fwrite(buf, 1, ret, copystream) != ret)
+			if (OK && copystream && fwrite(buf, 1, ret, copystream) != ret)
 			{
 				psql_error("could not write COPY data: %s\n",
 						   strerror(errno));
@@ -458,7 +459,7 @@ handleCopyOut(PGconn *conn, FILE *copystream, PGresult **res)
 		}
 	}
 
-	if (OK && fflush(copystream))
+	if (OK && copystream && fflush(copystream))
 	{
 		psql_error("could not write COPY data: %s\n",
 				   strerror(errno));