Skip to content

Commit 5519d5a

Browse files
peterepetere
committed
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
1 parent d16773c commit 5519d5a

File tree

1 file changed

+111
-93
lines changed

1 file changed

+111
-93
lines changed

‎src/bin/psql/common.c

Lines changed: 111 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -892,6 +892,116 @@ ExecQueryTuples(const PGresult *result)
892892
}
893893

894894

895+
/*
896+
* Marshal the COPY data. Either subroutine will get the
897+
* connection out of its COPY state, then call PQresultStatus()
898+
* once and report any error. Return whether all was ok.
899+
*
900+
* For COPY OUT, direct the output to pset.copyStream if it's set,
901+
* otherwise to pset.gfname if it's set, otherwise to queryFout.
902+
* For COPY IN, use pset.copyStream as data source if it's set,
903+
* otherwise cur_cmd_source.
904+
*
905+
* Update result if further processing is necessary, or NULL otherwise.
906+
* Return a result when queryFout can safely output a result status: on COPY
907+
* IN, or on COPY OUT if written to something other than pset.queryFout.
908+
* Returning NULL prevents the command status from being printed, which we
909+
* want if the status line doesn't get taken as part of the COPY data.
910+
*/
911+
static bool
912+
HandleCopyResult(PGresult **resultp)
913+
{
914+
bool success;
915+
FILE *copystream;
916+
PGresult *copy_result;
917+
ExecStatusType result_status = PQresultStatus(*resultp);
918+
919+
Assert(result_status == PGRES_COPY_OUT ||
920+
result_status == PGRES_COPY_IN);
921+
922+
SetCancelConn(pset.db);
923+
924+
if (result_status == PGRES_COPY_OUT)
925+
{
926+
bool need_close = false;
927+
bool is_pipe = false;
928+
929+
if (pset.copyStream)
930+
{
931+
/* invoked by \copy */
932+
copystream = pset.copyStream;
933+
}
934+
else if (pset.gfname)
935+
{
936+
/* invoked by \g */
937+
if (openQueryOutputFile(pset.gfname,
938+
&copystream, &is_pipe))
939+
{
940+
need_close = true;
941+
if (is_pipe)
942+
disable_sigpipe_trap();
943+
}
944+
else
945+
copystream = NULL; /* discard COPY data entirely */
946+
}
947+
else
948+
{
949+
/* fall back to the generic query output stream */
950+
copystream = pset.queryFout;
951+
}
952+
953+
success = handleCopyOut(pset.db,
954+
copystream,
955+
&copy_result)
956+
&& (copystream != NULL);
957+
958+
/*
959+
* Suppress status printing if the report would go to the same
960+
* place as the COPY data just went. Note this doesn't
961+
* prevent error reporting, since handleCopyOut did that.
962+
*/
963+
if (copystream == pset.queryFout)
964+
{
965+
PQclear(copy_result);
966+
copy_result = NULL;
967+
}
968+
969+
if (need_close)
970+
{
971+
/* close \g argument file/pipe */
972+
if (is_pipe)
973+
{
974+
pclose(copystream);
975+
restore_sigpipe_trap();
976+
}
977+
else
978+
{
979+
fclose(copystream);
980+
}
981+
}
982+
}
983+
else
984+
{
985+
/* COPY IN */
986+
copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source;
987+
success = handleCopyIn(pset.db,
988+
copystream,
989+
PQbinaryTuples(*resultp),
990+
&copy_result);
991+
}
992+
ResetCancelConn();
993+
994+
/*
995+
* Replace the PGRES_COPY_OUT/IN result with COPY command's exit
996+
* status, or with NULL if we want to suppress printing anything.
997+
*/
998+
PQclear(*resultp);
999+
*resultp = copy_result;
1000+
1001+
return success;
1002+
}
1003+
1004+
8951005
/*
8961006
* ProcessResult: utility function for use by SendQuery() only
8971007
*
@@ -957,99 +1067,7 @@ ProcessResult(PGresult **resultp)
9571067
}
9581068

9591069
if (is_copy)
960-
{
961-
/*
962-
* Marshal the COPY data. Either subroutine will get the
963-
* connection out of its COPY state, then call PQresultStatus()
964-
* once and report any error.
965-
*
966-
* For COPY OUT, direct the output to pset.copyStream if it's set,
967-
* otherwise to pset.gfname if it's set, otherwise to queryFout.
968-
* For COPY IN, use pset.copyStream as data source if it's set,
969-
* otherwise cur_cmd_source.
970-
*/
971-
FILE *copystream;
972-
PGresult *copy_result;
973-
974-
SetCancelConn(pset.db);
975-
if (result_status == PGRES_COPY_OUT)
976-
{
977-
bool need_close = false;
978-
bool is_pipe = false;
979-
980-
if (pset.copyStream)
981-
{
982-
/* invoked by \copy */
983-
copystream = pset.copyStream;
984-
}
985-
else if (pset.gfname)
986-
{
987-
/* invoked by \g */
988-
if (openQueryOutputFile(pset.gfname,
989-
&copystream, &is_pipe))
990-
{
991-
need_close = true;
992-
if (is_pipe)
993-
disable_sigpipe_trap();
994-
}
995-
else
996-
copystream = NULL; /* discard COPY data entirely */
997-
}
998-
else
999-
{
1000-
/* fall back to the generic query output stream */
1001-
copystream = pset.queryFout;
1002-
}
1003-
1004-
success = handleCopyOut(pset.db,
1005-
copystream,
1006-
&copy_result)
1007-
&& success
1008-
&& (copystream != NULL);
1009-
1010-
/*
1011-
* Suppress status printing if the report would go to the same
1012-
* place as the COPY data just went. Note this doesn't
1013-
* prevent error reporting, since handleCopyOut did that.
1014-
*/
1015-
if (copystream == pset.queryFout)
1016-
{
1017-
PQclear(copy_result);
1018-
copy_result = NULL;
1019-
}
1020-
1021-
if (need_close)
1022-
{
1023-
/* close \g argument file/pipe */
1024-
if (is_pipe)
1025-
{
1026-
pclose(copystream);
1027-
restore_sigpipe_trap();
1028-
}
1029-
else
1030-
{
1031-
fclose(copystream);
1032-
}
1033-
}
1034-
}
1035-
else
1036-
{
1037-
/* COPY IN */
1038-
copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source;
1039-
success = handleCopyIn(pset.db,
1040-
copystream,
1041-
PQbinaryTuples(*resultp),
1042-
&copy_result) && success;
1043-
}
1044-
ResetCancelConn();
1045-
1046-
/*
1047-
* Replace the PGRES_COPY_OUT/IN result with COPY command's exit
1048-
* status, or with NULL if we want to suppress printing anything.
1049-
*/
1050-
PQclear(*resultp);
1051-
*resultp = copy_result;
1052-
}
1070+
success = HandleCopyResult(resultp);
10531071
else if (first_cycle)
10541072
{
10551073
/* fast path: no COPY commands; PQexec visited all results */

0 commit comments

Comments
 (0)