Fix tracking of psql script line numbers during \copy from another place.

tglsfdc · tglsfdc · commit b315b767f7a3 · 2014-03-10T15:47:13.000-04:00
Commit 0814677 changed do_copy() to temporarily scribble on pset.cur_cmd_source. That was a mighty ugly bit of code in any case, but in particular it broke handleCopyIn's ability to tell whether it was reading from the current script source file (in which case pset.lineno should be incremented for each line of COPY data), or from someplace else (in which case it shouldn't). The former case still worked, the latter not so much. The visible effect was that line numbers reported for errors in a script file would be wrong if there were an earlier \copy that was reading anything other than inline-in-the-script-file data. To fix, introduce another pset field that holds the file do_copy wants the COPY code to use. This is a little bit ugly, but less so than passing the file down explicitly through several layers that aren't COPY-specific. Extracted from a larger patch by Kumar Rajeev Rastogi; that patch also changes printing of COPY command tags, which is not a bug fix and shouldn't get back-patched. This particular idea was from a suggestion by Amit Khandekar, if I'm reading the thread correctly. Back-patch to 9.2 where the faulty code was introduced.
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
@@ -639,7 +639,7 @@ PrintQueryTuples(const PGresult *results)
  * 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
+ * When the command string contained no such COPY command, this function
  * degenerates to an AcceptResult() call.
  *
  * Changes its argument to point to the last PGresult of the command string,
@@ -699,13 +699,28 @@ ProcessResult(PGresult **results)
 			 * Marshal the COPY data.  Either subroutine will get the
 			 * 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.
 			 */
+			FILE	   *copystream = pset.copyStream;
+
 			SetCancelConn();
 			if (result_status == PGRES_COPY_OUT)
-				success = handleCopyOut(pset.db, pset.queryFout) && success;
+			{
+				if (!copystream)
+					copystream = pset.queryFout;
+				success = handleCopyOut(pset.db,
+										copystream) && success;
+			}
 			else
-				success = handleCopyIn(pset.db, pset.cur_cmd_source,
+			{
+				if (!copystream)
+					copystream = pset.cur_cmd_source;
+				success = handleCopyIn(pset.db,
+									   copystream,
 									   PQbinaryTuples(*results)) && success;
+			}
 			ResetCancelConn();
 
 			/*
diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c
@@ -244,8 +244,6 @@ do_copy(const char *args)
 {
 	PQExpBufferData query;
 	FILE	   *copystream;
-	FILE	   *save_file;
-	FILE	  **override_file;
 	struct copy_options *options;
 	bool		success;
 	struct stat st;
@@ -262,8 +260,6 @@ do_copy(const char *args)
 
 	if (options->from)
 	{
-		override_file = &pset.cur_cmd_source;
-
 		if (options->file)
 			copystream = fopen(options->file, PG_BINARY_R);
 		else if (!options->psql_inout)
@@ -273,8 +269,6 @@ do_copy(const char *args)
 	}
 	else
 	{
-		override_file = &pset.queryFout;
-
 		if (options->file)
 			copystream = fopen(options->file, PG_BINARY_W);
 		else if (!options->psql_inout)
@@ -313,11 +307,10 @@ do_copy(const char *args)
 	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;
+	/* run it like a user command, but with copystream as data source/sink */
+	pset.copyStream = copystream;
 	success = SendQuery(query.data);
-	*override_file = save_file;
+	pset.copyStream = NULL;
 	termPQExpBuffer(&query);
 
 	if (options->file != NULL)
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
@@ -70,6 +70,8 @@ typedef struct _psqlSettings
 	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 */
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
@@ -120,6 +120,7 @@ main(int argc, char *argv[])
 	pset.encoding = PQenv2encoding();
 	pset.queryFout = stdout;
 	pset.queryFoutPipe = false;
+	pset.copyStream = NULL;
 	pset.cur_cmd_source = stdin;
 	pset.cur_cmd_interactive = false;
 
