psql: Refactor ProcessResult()

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
index c9847c8f9a02745d8d5553095d2814b973479f06..f6777bbfc44a9469a3338178a0175f40f2103508 100644 (file)
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -892,6 +892,116 @@ loop_exit:
 }
 
 
+/*
+ * 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 */