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 */