psql: Refactor ProcessResult()

petere · petere · commit 5519d5affdfa · 2022-04-01T13:03:33.000+02:00
Separate HandleCopyResult() from ProcessResult() in preparation for a subsequent patch. Author: Fabien COELHO <coelho@cri.ensmp.fr> Discussion: https://www.postgresql.org/message-id/flat/alpine.DEB.2.21.1904132231510.8961@lancre
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
@@ -892,6 +892,116 @@ ExecQueryTuples(const PGresult *result)
 }
 
 
+/*
+ * Marshal the COPY data.  Either subroutine will get the
+ * connection out of its COPY state, then call PQresultStatus()
+ * once and report any error.  Return whether all was ok.
+ *
+ * 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.
+ *
+ * Update result if further processing is necessary, or NULL otherwise.
+ * Return a result when queryFout can safely output a result status: on COPY
+ * IN, or on COPY OUT if written to something other than pset.queryFout.
+ * Returning NULL prevents the command status from being printed, which we
+ * want if the status line doesn't get taken as part of the COPY data.
+ */
+static bool
+HandleCopyResult(PGresult **resultp)
+{
+	bool		success;
+	FILE	   *copystream;
+	PGresult   *copy_result;
+	ExecStatusType result_status = PQresultStatus(*resultp);
+
+	Assert(result_status == PGRES_COPY_OUT ||
+		   result_status == PGRES_COPY_IN);
+
+	SetCancelConn(pset.db);
+
+	if (result_status == PGRES_COPY_OUT)
+	{
+		bool		need_close = false;
+		bool		is_pipe = false;
+
+		if (pset.copyStream)
+		{
+			/* invoked by \copy */
+			copystream = pset.copyStream;
+		}
+		else if (pset.gfname)
+		{
+			/* invoked by \g */
+			if (openQueryOutputFile(pset.gfname,
+									&copystream, &is_pipe))
+			{
+				need_close = true;
+				if (is_pipe)
+					disable_sigpipe_trap();
+			}
+			else
+				copystream = NULL;	/* discard COPY data entirely */
+		}
+		else
+		{
+			/* fall back to the generic query output stream */
+			copystream = pset.queryFout;
+		}
+
+		success = handleCopyOut(pset.db,
+								copystream,
+								&copy_result)
+			&& (copystream != NULL);
+
+		/*
+		 * Suppress status printing if the report would go to the same
+		 * place as the COPY data just went.  Note this doesn't
+		 * prevent error reporting, since handleCopyOut did that.
+		 */
+		if (copystream == pset.queryFout)
+		{
+			PQclear(copy_result);
+			copy_result = NULL;
+		}
+
+		if (need_close)
+		{
+			/* close \g argument file/pipe */
+			if (is_pipe)
+			{
+				pclose(copystream);
+				restore_sigpipe_trap();
+			}
+			else
+			{
+				fclose(copystream);
+			}
+		}
+	}
+	else
+	{
+		/* COPY IN */
+		copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source;
+		success = handleCopyIn(pset.db,
+							   copystream,
+							   PQbinaryTuples(*resultp),
+							   &copy_result);
+	}
+	ResetCancelConn();
+
+	/*
+	 * Replace the PGRES_COPY_OUT/IN result with COPY command's exit
+	 * status, or with NULL if we want to suppress printing anything.
+	 */
+	PQclear(*resultp);
+	*resultp = copy_result;
+
+	return success;
+}
+
+
 /*
  * ProcessResult: utility function for use by SendQuery() only
  *
@@ -957,99 +1067,7 @@ ProcessResult(PGresult **resultp)
 		}
 
 		if (is_copy)
-		{
-			/*
-			 * Marshal the COPY data.  Either subroutine will get the
-			 * connection out of its COPY state, then call PQresultStatus()
-			 * once and report any error.
-			 *
-			 * 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;
-			PGresult   *copy_result;
-
-			SetCancelConn(pset.db);
-			if (result_status == PGRES_COPY_OUT)
-			{
-				bool		need_close = false;
-				bool		is_pipe = false;
-
-				if (pset.copyStream)
-				{
-					/* invoked by \copy */
-					copystream = pset.copyStream;
-				}
-				else if (pset.gfname)
-				{
-					/* invoked by \g */
-					if (openQueryOutputFile(pset.gfname,
-											&copystream, &is_pipe))
-					{
-						need_close = true;
-						if (is_pipe)
-							disable_sigpipe_trap();
-					}
-					else
-						copystream = NULL;	/* discard COPY data entirely */
-				}
-				else
-				{
-					/* fall back to the generic query output stream */
-					copystream = pset.queryFout;
-				}
-
-				success = handleCopyOut(pset.db,
-										copystream,
-										&copy_result)
-					&& success
-					&& (copystream != NULL);
-
-				/*
-				 * Suppress status printing if the report would go to the same
-				 * place as the COPY data just went.  Note this doesn't
-				 * prevent error reporting, since handleCopyOut did that.
-				 */
-				if (copystream == pset.queryFout)
-				{
-					PQclear(copy_result);
-					copy_result = NULL;
-				}
-
-				if (need_close)
-				{
-					/* close \g argument file/pipe */
-					if (is_pipe)
-					{
-						pclose(copystream);
-						restore_sigpipe_trap();
-					}
-					else
-					{
-						fclose(copystream);
-					}
-				}
-			}
-			else
-			{
-				/* COPY IN */
-				copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source;
-				success = handleCopyIn(pset.db,
-									   copystream,
-									   PQbinaryTuples(*resultp),
-									   &copy_result) && success;
-			}
-			ResetCancelConn();
-
-			/*
-			 * Replace the PGRES_COPY_OUT/IN result with COPY command's exit
-			 * status, or with NULL if we want to suppress printing anything.
-			 */
-			PQclear(*resultp);
-			*resultp = copy_result;
-		}
+			success = HandleCopyResult(resultp);
 		else if (first_cycle)
 		{
 			/* fast path: no COPY commands; PQexec visited all results */
