Skip to content

Commit 0814677

Browse files
alvherrealvherre
committed
Have \copy go through SendQuery
This enables a bunch of features, notably ON_ERROR_ROLLBACK. It also makes COPY failure (either in the server or psql) as a whole behave more sanely in psql. Additionally, having more commands in the same command line as COPY works better (though since psql splits lines at semicolons, this doesn't matter much unless you're using -c). Also tighten a couple of switches on PQresultStatus() to add PGRES_COPY_BOTH support and stop assuming that unknown statuses received are errors; have those print diagnostics where warranted. Author: Noah Misch
1 parent 6eb71ac commit 0814677

File tree

5 files changed

+250
-126
lines changed

5 files changed

+250
-126
lines changed

‎src/bin/psql/command.c

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -320,25 +320,10 @@ exec_command(const char *cmd,
320320
/* \copy */
321321
else if (pg_strcasecmp(cmd, "copy") == 0)
322322
{
323-
/* Default fetch-it-all-and-print mode */
324-
instr_time before,
325-
after;
326-
327323
char *opt = psql_scan_slash_option(scan_state,
328324
OT_WHOLE_LINE, NULL, false);
329325

330-
if (pset.timing)
331-
INSTR_TIME_SET_CURRENT(before);
332-
333326
success = do_copy(opt);
334-
335-
if (pset.timing && success)
336-
{
337-
INSTR_TIME_SET_CURRENT(after);
338-
INSTR_TIME_SUBTRACT(after, before);
339-
printf(_("Time: %.3f ms\n"), INSTR_TIME_GET_MILLISEC(after));
340-
}
341-
342327
free(opt);
343328
}
344329

‎src/bin/psql/common.c

Lines changed: 152 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,7 @@ ResetCancelConn(void)
438438
static bool
439439
AcceptResult(const PGresult *result)
440440
{
441-
bool OK = true;
441+
bool OK;
442442

443443
if (!result)
444444
OK = false;
@@ -450,11 +450,21 @@ AcceptResult(const PGresult *result)
450450
case PGRES_EMPTY_QUERY:
451451
case PGRES_COPY_IN:
452452
case PGRES_COPY_OUT:
453+
case PGRES_COPY_BOTH:
453454
/* Fine, do nothing */
455+
OK = true;
456+
break;
457+
458+
case PGRES_BAD_RESPONSE:
459+
case PGRES_NONFATAL_ERROR:
460+
case PGRES_FATAL_ERROR:
461+
OK = false;
454462
break;
455463

456464
default:
457465
OK = false;
466+
psql_error("unexpected PQresultStatus (%d)",
467+
PQresultStatus(result));
458468
break;
459469
}
460470

@@ -620,45 +630,114 @@ PrintQueryTuples(const PGresult *results)
620630

621631

622632
/*
623-
* ProcessCopyResult: if command was a COPY FROM STDIN/TO STDOUT, handle it
633+
* ProcessResult: utility function for use by SendQuery() only
624634
*
625-
* Note: Utility function for use by SendQuery() only.
635+
* When our command string contained a COPY FROM STDIN or COPY TO STDOUT,
636+
* PQexec() has stopped at the PGresult associated with the first such
637+
* command. In that event, we'll marshal data for the COPY and then cycle
638+
* through any subsequent PGresult objects.
626639
*
627-
* Returns true if the query executed successfully, false otherwise.
640+
* When the command string contained no affected COPY command, this function
641+
* degenerates to an AcceptResult() call.
642+
*
643+
* Changes its argument to point to the last PGresult of the command string,
644+
* or NULL if that result was for a COPY FROM STDIN or COPY TO STDOUT.
645+
*
646+
* Returns true on complete success, false otherwise. Possible failure modes
647+
* include purely client-side problems; check the transaction status for the
648+
* server-side opinion.
628649
*/
629650
static bool
630-
ProcessCopyResult(PGresult *results)
651+
ProcessResult(PGresult **results)
631652
{
632-
bool success = false;
653+
PGresult *next_result;
654+
bool success = true;
655+
bool first_cycle = true;
633656

634-
if (!results)
635-
return false;
636-
637-
switch (PQresultStatus(results))
657+
do
638658
{
639-
case PGRES_TUPLES_OK:
640-
case PGRES_COMMAND_OK:
641-
case PGRES_EMPTY_QUERY:
642-
/* nothing to do here */
643-
success = true;
644-
break;
659+
ExecStatusType result_status;
660+
bool is_copy;
645661

646-
case PGRES_COPY_OUT:
647-
SetCancelConn();
648-
success = handleCopyOut(pset.db, pset.queryFout);
649-
ResetCancelConn();
662+
if (!AcceptResult(*results))
663+
{
664+
/*
665+
* Failure at this point is always a server-side failure or a
666+
* failure to submit the command string. Either way, we're
667+
* finished with this command string.
668+
*/
669+
success = false;
650670
break;
671+
}
651672

652-
case PGRES_COPY_IN:
673+
result_status = PQresultStatus(*results);
674+
switch (result_status)
675+
{
676+
case PGRES_COPY_BOTH:
677+
/*
678+
* No now-existing SQL command can yield PGRES_COPY_BOTH, but
679+
* defend against the future. PQexec() can't short-circuit
680+
* it's way out of a PGRES_COPY_BOTH, so the connection will
681+
* be useless at this point. XXX is there a method for
682+
* clearing this status that's likely to work with every
683+
* future command that can initiate it?
684+
*/
685+
psql_error("unexpected PQresultStatus (%d)", result_status);
686+
return false;
687+
688+
case PGRES_COPY_OUT:
689+
case PGRES_COPY_IN:
690+
is_copy = true;
691+
break;
692+
693+
case PGRES_EMPTY_QUERY:
694+
case PGRES_COMMAND_OK:
695+
case PGRES_TUPLES_OK:
696+
is_copy = false;
697+
break;
698+
699+
default:
700+
/* AcceptResult() should have caught anything else. */
701+
is_copy = false;
702+
psql_error("unexpected PQresultStatus (%d)", result_status);
703+
break;
704+
}
705+
706+
if (is_copy)
707+
{
708+
/*
709+
* Marshal the COPY data. Either subroutine will get the
710+
* connection out of its COPY state, then call PQresultStatus()
711+
* once and report any error.
712+
*/
653713
SetCancelConn();
654-
success = handleCopyIn(pset.db, pset.cur_cmd_source,
655-
PQbinaryTuples(results));
714+
if (result_status == PGRES_COPY_OUT)
715+
success = handleCopyOut(pset.db, pset.queryFout) && success;
716+
else
717+
success = handleCopyIn(pset.db, pset.cur_cmd_source,
718+
PQbinaryTuples(*results)) && success;
656719
ResetCancelConn();
657-
break;
658720

659-
default:
721+
/*
722+
* Call PQgetResult() once more. In the typical case of a
723+
* single-command string, it will return NULL. Otherwise, we'll
724+
* have other results to process that may include other COPYs.
725+
*/
726+
PQclear(*results);
727+
*results = next_result = PQgetResult(pset.db);
728+
}
729+
else if (first_cycle)
730+
/* fast path: no COPY commands; PQexec visited all results */
660731
break;
661-
}
732+
else if ((next_result = PQgetResult(pset.db)))
733+
{
734+
/* non-COPY command(s) after a COPY: keep the last one */
735+
PQclear(*results);
736+
*results = next_result;
737+
}
738+
739+
first_cycle = false;
740+
} while (next_result);
662741

663742
/* may need this to recover from conn loss during COPY */
664743
if (!CheckConnection())
@@ -708,7 +787,7 @@ PrintQueryStatus(PGresult *results)
708787
static bool
709788
PrintQueryResults(PGresult *results)
710789
{
711-
bool success = false;
790+
bool success;
712791
const char *cmdstatus;
713792

714793
if (!results)
@@ -738,11 +817,21 @@ PrintQueryResults(PGresult *results)
738817

739818
case PGRES_COPY_OUT:
740819
case PGRES_COPY_IN:
820+
case PGRES_COPY_BOTH:
741821
/* nothing to do here */
742822
success = true;
743823
break;
744824

825+
case PGRES_BAD_RESPONSE:
826+
case PGRES_NONFATAL_ERROR:
827+
case PGRES_FATAL_ERROR:
828+
success = false;
829+
break;
830+
745831
default:
832+
success = false;
833+
psql_error("unexpected PQresultStatus (%d)",
834+
PQresultStatus(results));
746835
break;
747836
}
748837

@@ -867,7 +956,7 @@ SendQuery(const char *query)
867956

868957
/* these operations are included in the timing result: */
869958
ResetCancelConn();
870-
OK = (AcceptResult(results) && ProcessCopyResult(results));
959+
OK = ProcessResult(&results);
871960

872961
if (pset.timing)
873962
{
@@ -877,7 +966,7 @@ SendQuery(const char *query)
877966
}
878967

879968
/* but printing results isn't: */
880-
if (OK)
969+
if (OK && results)
881970
OK = PrintQueryResults(results);
882971
}
883972
else
@@ -891,34 +980,44 @@ SendQuery(const char *query)
891980
/* If we made a temporary savepoint, possibly release/rollback */
892981
if (on_error_rollback_savepoint)
893982
{
894-
const char *svptcmd;
983+
const char *svptcmd = NULL;
895984

896985
transaction_status = PQtransactionStatus(pset.db);
897986

898-
if (transaction_status == PQTRANS_INERROR)
899-
{
900-
/* We always rollback on an error */
901-
svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint";
902-
}
903-
else if (transaction_status != PQTRANS_INTRANS)
987+
switch (transaction_status)
904988
{
905-
/* If they are no longer in a transaction, then do nothing */
906-
svptcmd = NULL;
907-
}
908-
else
909-
{
910-
/*
911-
* Do nothing if they are messing with savepoints themselves: If
912-
* the user did RELEASE or ROLLBACK, our savepoint is gone. If
913-
* they issued a SAVEPOINT, releasing ours would remove theirs.
914-
*/
915-
if (results &&
916-
(strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
917-
strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
918-
strcmp(PQcmdStatus(results), "ROLLBACK") == 0))
919-
svptcmd = NULL;
920-
else
921-
svptcmd = "RELEASE pg_psql_temporary_savepoint";
989+
case PQTRANS_INERROR:
990+
/* We always rollback on an error */
991+
svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint";
992+
break;
993+
994+
case PQTRANS_IDLE:
995+
/* If they are no longer in a transaction, then do nothing */
996+
break;
997+
998+
case PQTRANS_INTRANS:
999+
/*
1000+
* Do nothing if they are messing with savepoints themselves:
1001+
* If the user did RELEASE or ROLLBACK, our savepoint is
1002+
* gone. If they issued a SAVEPOINT, releasing ours would
1003+
* remove theirs.
1004+
*/
1005+
if (results &&
1006+
(strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
1007+
strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
1008+
strcmp(PQcmdStatus(results), "ROLLBACK") == 0))
1009+
svptcmd = NULL;
1010+
else
1011+
svptcmd = "RELEASE pg_psql_temporary_savepoint";
1012+
break;
1013+
1014+
case PQTRANS_ACTIVE:
1015+
case PQTRANS_UNKNOWN:
1016+
default:
1017+
OK = false;
1018+
psql_error("unexpected transaction status (%d)\n",
1019+
transaction_status);
1020+
break;
9221021
}
9231022

9241023
if (svptcmd)

0 commit comments

Comments
 (0)