1 /*
2 * psql - the PostgreSQL interactive terminal
3 *
4 * Copyright (c) 2000-2020, PostgreSQL Global Development Group
5 *
6 * src/bin/psql/common.c
7 */
8 #include "postgres_fe.h"
9
10 #include <ctype.h>
11 #include <limits.h>
12 #include <math.h>
13 #include <signal.h>
14 #ifndef WIN32
15 #include <unistd.h> /* for write() */
16 #else
17 #include <io.h> /* for _write() */
18 #include <win32.h>
19 #endif
20
21 #include "command.h"
22 #include "common.h"
23 #include "common/logging.h"
24 #include "copy.h"
25 #include "crosstabview.h"
26 #include "fe_utils/cancel.h"
27 #include "fe_utils/mbprint.h"
28 #include "fe_utils/string_utils.h"
29 #include "portability/instr_time.h"
30 #include "settings.h"
31
32 #define PQmblenBounded(s, e) strnlen(s, PQmblen(s, e))
33
34 static bool DescribeQuery(const char *query, double *elapsed_msec);
35 static bool ExecQueryUsingCursor(const char *query, double *elapsed_msec);
36 static bool command_no_begin(const char *query);
37 static bool is_select_command(const char *query);
38
39
40 /*
41 * openQueryOutputFile --- attempt to open a query output file
42 *
43 * fname == NULL selects stdout, else an initial '|' selects a pipe,
44 * else plain file.
45 *
46 * Returns output file pointer into *fout, and is-a-pipe flag into *is_pipe.
47 * Caller is responsible for adjusting SIGPIPE state if it's a pipe.
48 *
49 * On error, reports suitable error message and returns false.
50 */
51 bool
openQueryOutputFile(const char * fname,FILE ** fout,bool * is_pipe)52 openQueryOutputFile(const char *fname, FILE **fout, bool *is_pipe)
53 {
54 if (!fname || fname[0] == '\0')
55 {
56 *fout = stdout;
57 *is_pipe = false;
58 }
59 else if (*fname == '|')
60 {
61 *fout = popen(fname + 1, "w");
62 *is_pipe = true;
63 }
64 else
65 {
66 *fout = fopen(fname, "w");
67 *is_pipe = false;
68 }
69
70 if (*fout == NULL)
71 {
72 pg_log_error("%s: %m", fname);
73 return false;
74 }
75
76 return true;
77 }
78
79 /*
80 * setQFout
81 * -- handler for -o command line option and \o command
82 *
83 * On success, updates pset with the new output file and returns true.
84 * On failure, returns false without changing pset state.
85 */
86 bool
setQFout(const char * fname)87 setQFout(const char *fname)
88 {
89 FILE *fout;
90 bool is_pipe;
91
92 /* First make sure we can open the new output file/pipe */
93 if (!openQueryOutputFile(fname, &fout, &is_pipe))
94 return false;
95
96 /* Close old file/pipe */
97 if (pset.queryFout && pset.queryFout != stdout && pset.queryFout != stderr)
98 {
99 if (pset.queryFoutPipe)
100 pclose(pset.queryFout);
101 else
102 fclose(pset.queryFout);
103 }
104
105 pset.queryFout = fout;
106 pset.queryFoutPipe = is_pipe;
107
108 /* Adjust SIGPIPE handling appropriately: ignore signal if is_pipe */
109 set_sigpipe_trap_state(is_pipe);
110 restore_sigpipe_trap();
111
112 return true;
113 }
114
115
116 /*
117 * Variable-fetching callback for flex lexer
118 *
119 * If the specified variable exists, return its value as a string (malloc'd
120 * and expected to be freed by the caller); else return NULL.
121 *
122 * If "quote" isn't PQUOTE_PLAIN, then return the value suitably quoted and
123 * escaped for the specified quoting requirement. (Failure in escaping
124 * should lead to printing an error and returning NULL.)
125 *
126 * "passthrough" is the pointer previously given to psql_scan_set_passthrough.
127 * In psql, passthrough points to a ConditionalStack, which we check to
128 * determine whether variable expansion is allowed.
129 */
130 char *
psql_get_variable(const char * varname,PsqlScanQuoteType quote,void * passthrough)131 psql_get_variable(const char *varname, PsqlScanQuoteType quote,
132 void *passthrough)
133 {
134 char *result = NULL;
135 const char *value;
136
137 /* In an inactive \if branch, suppress all variable substitutions */
138 if (passthrough && !conditional_active((ConditionalStack) passthrough))
139 return NULL;
140
141 value = GetVariable(pset.vars, varname);
142 if (!value)
143 return NULL;
144
145 switch (quote)
146 {
147 case PQUOTE_PLAIN:
148 result = pg_strdup(value);
149 break;
150 case PQUOTE_SQL_LITERAL:
151 case PQUOTE_SQL_IDENT:
152 {
153 /*
154 * For these cases, we use libpq's quoting functions, which
155 * assume the string is in the connection's client encoding.
156 */
157 char *escaped_value;
158
159 if (!pset.db)
160 {
161 pg_log_error("cannot escape without active connection");
162 return NULL;
163 }
164
165 if (quote == PQUOTE_SQL_LITERAL)
166 escaped_value =
167 PQescapeLiteral(pset.db, value, strlen(value));
168 else
169 escaped_value =
170 PQescapeIdentifier(pset.db, value, strlen(value));
171
172 if (escaped_value == NULL)
173 {
174 const char *error = PQerrorMessage(pset.db);
175
176 pg_log_info("%s", error);
177 return NULL;
178 }
179
180 /*
181 * Rather than complicate the lexer's API with a notion of
182 * which free() routine to use, just pay the price of an extra
183 * strdup().
184 */
185 result = pg_strdup(escaped_value);
186 PQfreemem(escaped_value);
187 break;
188 }
189 case PQUOTE_SHELL_ARG:
190 {
191 /*
192 * For this we use appendShellStringNoError, which is
193 * encoding-agnostic, which is fine since the shell probably
194 * is too. In any case, the only special character is "'",
195 * which is not known to appear in valid multibyte characters.
196 */
197 PQExpBufferData buf;
198
199 initPQExpBuffer(&buf);
200 if (!appendShellStringNoError(&buf, value))
201 {
202 pg_log_error("shell command argument contains a newline or carriage return: \"%s\"",
203 value);
204 free(buf.data);
205 return NULL;
206 }
207 result = buf.data;
208 break;
209 }
210
211 /* No default: we want a compiler warning for missing cases */
212 }
213
214 return result;
215 }
216
217
218 /*
219 * for backend Notice messages (INFO, WARNING, etc)
220 */
221 void
NoticeProcessor(void * arg,const char * message)222 NoticeProcessor(void *arg, const char *message)
223 {
224 (void) arg; /* not used */
225 pg_log_info("%s", message);
226 }
227
228
229
230 /*
231 * Code to support query cancellation
232 *
233 * Before we start a query, we enable the SIGINT signal catcher to send a
234 * cancel request to the backend.
235 *
236 * SIGINT is supposed to abort all long-running psql operations, not only
237 * database queries. In most places, this is accomplished by checking
238 * cancel_pressed during long-running loops. However, that won't work when
239 * blocked on user input (in readline() or fgets()). In those places, we
240 * set sigint_interrupt_enabled true while blocked, instructing the signal
241 * catcher to longjmp through sigint_interrupt_jmp. We assume readline and
242 * fgets are coded to handle possible interruption.
243 *
244 * On Windows, currently this does not work, so control-C is less useful
245 * there.
246 */
247 volatile bool sigint_interrupt_enabled = false;
248
249 sigjmp_buf sigint_interrupt_jmp;
250
251 static void
psql_cancel_callback(void)252 psql_cancel_callback(void)
253 {
254 #ifndef WIN32
255 /* if we are waiting for input, longjmp out of it */
256 if (sigint_interrupt_enabled)
257 {
258 sigint_interrupt_enabled = false;
259 siglongjmp(sigint_interrupt_jmp, 1);
260 }
261 #endif
262
263 /* else, set cancel flag to stop any long-running loops */
264 cancel_pressed = true;
265 }
266
267 void
psql_setup_cancel_handler(void)268 psql_setup_cancel_handler(void)
269 {
270 setup_cancel_handler(psql_cancel_callback);
271 }
272
273
274 /* ConnectionUp
275 *
276 * Returns whether our backend connection is still there.
277 */
278 static bool
ConnectionUp(void)279 ConnectionUp(void)
280 {
281 return PQstatus(pset.db) != CONNECTION_BAD;
282 }
283
284
285
286 /* CheckConnection
287 *
288 * Verify that we still have a good connection to the backend, and if not,
289 * see if it can be restored.
290 *
291 * Returns true if either the connection was still there, or it could be
292 * restored successfully; false otherwise. If, however, there was no
293 * connection and the session is non-interactive, this will exit the program
294 * with a code of EXIT_BADCONN.
295 */
296 static bool
CheckConnection(void)297 CheckConnection(void)
298 {
299 bool OK;
300
301 OK = ConnectionUp();
302 if (!OK)
303 {
304 if (!pset.cur_cmd_interactive)
305 {
306 pg_log_fatal("connection to server was lost");
307 exit(EXIT_BADCONN);
308 }
309
310 fprintf(stderr, _("The connection to the server was lost. Attempting reset: "));
311 PQreset(pset.db);
312 OK = ConnectionUp();
313 if (!OK)
314 {
315 fprintf(stderr, _("Failed.\n"));
316
317 /*
318 * Transition to having no connection. Keep this bit in sync with
319 * do_connect().
320 */
321 PQfinish(pset.db);
322 pset.db = NULL;
323 ResetCancelConn();
324 UnsyncVariables();
325 }
326 else
327 {
328 fprintf(stderr, _("Succeeded.\n"));
329
330 /*
331 * Re-sync, just in case anything changed. Keep this in sync with
332 * do_connect().
333 */
334 SyncVariables();
335 connection_warnings(false); /* Must be after SyncVariables */
336 }
337 }
338
339 return OK;
340 }
341
342
343
344
345 /*
346 * AcceptResult
347 *
348 * Checks whether a result is valid, giving an error message if necessary;
349 * and ensures that the connection to the backend is still up.
350 *
351 * Returns true for valid result, false for error state.
352 */
353 static bool
AcceptResult(const PGresult * result)354 AcceptResult(const PGresult *result)
355 {
356 bool OK;
357
358 if (!result)
359 OK = false;
360 else
361 switch (PQresultStatus(result))
362 {
363 case PGRES_COMMAND_OK:
364 case PGRES_TUPLES_OK:
365 case PGRES_EMPTY_QUERY:
366 case PGRES_COPY_IN:
367 case PGRES_COPY_OUT:
368 /* Fine, do nothing */
369 OK = true;
370 break;
371
372 case PGRES_BAD_RESPONSE:
373 case PGRES_NONFATAL_ERROR:
374 case PGRES_FATAL_ERROR:
375 OK = false;
376 break;
377
378 default:
379 OK = false;
380 pg_log_error("unexpected PQresultStatus: %d",
381 PQresultStatus(result));
382 break;
383 }
384
385 if (!OK)
386 {
387 const char *error = PQerrorMessage(pset.db);
388
389 if (strlen(error))
390 pg_log_info("%s", error);
391
392 CheckConnection();
393 }
394
395 return OK;
396 }
397
398
399 /*
400 * Set special variables from a query result
401 * - ERROR: true/false, whether an error occurred on this query
402 * - SQLSTATE: code of error, or "00000" if no error, or "" if unknown
403 * - ROW_COUNT: how many rows were returned or affected, or "0"
404 * - LAST_ERROR_SQLSTATE: same for last error
405 * - LAST_ERROR_MESSAGE: message of last error
406 *
407 * Note: current policy is to apply this only to the results of queries
408 * entered by the user, not queries generated by slash commands.
409 */
410 static void
SetResultVariables(PGresult * results,bool success)411 SetResultVariables(PGresult *results, bool success)
412 {
413 if (success)
414 {
415 const char *ntuples = PQcmdTuples(results);
416
417 SetVariable(pset.vars, "ERROR", "false");
418 SetVariable(pset.vars, "SQLSTATE", "00000");
419 SetVariable(pset.vars, "ROW_COUNT", *ntuples ? ntuples : "0");
420 }
421 else
422 {
423 const char *code = PQresultErrorField(results, PG_DIAG_SQLSTATE);
424 const char *mesg = PQresultErrorField(results, PG_DIAG_MESSAGE_PRIMARY);
425
426 SetVariable(pset.vars, "ERROR", "true");
427
428 /*
429 * If there is no SQLSTATE code, use an empty string. This can happen
430 * for libpq-detected errors (e.g., lost connection, ENOMEM).
431 */
432 if (code == NULL)
433 code = "";
434 SetVariable(pset.vars, "SQLSTATE", code);
435 SetVariable(pset.vars, "ROW_COUNT", "0");
436 SetVariable(pset.vars, "LAST_ERROR_SQLSTATE", code);
437 SetVariable(pset.vars, "LAST_ERROR_MESSAGE", mesg ? mesg : "");
438 }
439 }
440
441
442 /*
443 * ClearOrSaveResult
444 *
445 * If the result represents an error, remember it for possible display by
446 * \errverbose. Otherwise, just PQclear() it.
447 *
448 * Note: current policy is to apply this to the results of all queries,
449 * including "back door" queries, for debugging's sake. It's OK to use
450 * PQclear() directly on results known to not be error results, however.
451 */
452 static void
ClearOrSaveResult(PGresult * result)453 ClearOrSaveResult(PGresult *result)
454 {
455 if (result)
456 {
457 switch (PQresultStatus(result))
458 {
459 case PGRES_NONFATAL_ERROR:
460 case PGRES_FATAL_ERROR:
461 if (pset.last_error_result)
462 PQclear(pset.last_error_result);
463 pset.last_error_result = result;
464 break;
465
466 default:
467 PQclear(result);
468 break;
469 }
470 }
471 }
472
473
474 /*
475 * Print microtiming output. Always print raw milliseconds; if the interval
476 * is >= 1 second, also break it down into days/hours/minutes/seconds.
477 */
478 static void
PrintTiming(double elapsed_msec)479 PrintTiming(double elapsed_msec)
480 {
481 double seconds;
482 double minutes;
483 double hours;
484 double days;
485
486 if (elapsed_msec < 1000.0)
487 {
488 /* This is the traditional (pre-v10) output format */
489 printf(_("Time: %.3f ms\n"), elapsed_msec);
490 return;
491 }
492
493 /*
494 * Note: we could print just seconds, in a format like %06.3f, when the
495 * total is less than 1min. But that's hard to interpret unless we tack
496 * on "s" or otherwise annotate it. Forcing the display to include
497 * minutes seems like a better solution.
498 */
499 seconds = elapsed_msec / 1000.0;
500 minutes = floor(seconds / 60.0);
501 seconds -= 60.0 * minutes;
502 if (minutes < 60.0)
503 {
504 printf(_("Time: %.3f ms (%02d:%06.3f)\n"),
505 elapsed_msec, (int) minutes, seconds);
506 return;
507 }
508
509 hours = floor(minutes / 60.0);
510 minutes -= 60.0 * hours;
511 if (hours < 24.0)
512 {
513 printf(_("Time: %.3f ms (%02d:%02d:%06.3f)\n"),
514 elapsed_msec, (int) hours, (int) minutes, seconds);
515 return;
516 }
517
518 days = floor(hours / 24.0);
519 hours -= 24.0 * days;
520 printf(_("Time: %.3f ms (%.0f d %02d:%02d:%06.3f)\n"),
521 elapsed_msec, days, (int) hours, (int) minutes, seconds);
522 }
523
524
525 /*
526 * PSQLexec
527 *
528 * This is the way to send "backdoor" queries (those not directly entered
529 * by the user). It is subject to -E but not -e.
530 *
531 * Caller is responsible for handling the ensuing processing if a COPY
532 * command is sent.
533 *
534 * Note: we don't bother to check PQclientEncoding; it is assumed that no
535 * caller uses this path to issue "SET CLIENT_ENCODING".
536 */
537 PGresult *
PSQLexec(const char * query)538 PSQLexec(const char *query)
539 {
540 PGresult *res;
541
542 if (!pset.db)
543 {
544 pg_log_error("You are currently not connected to a database.");
545 return NULL;
546 }
547
548 if (pset.echo_hidden != PSQL_ECHO_HIDDEN_OFF)
549 {
550 printf(_("********* QUERY **********\n"
551 "%s\n"
552 "**************************\n\n"), query);
553 fflush(stdout);
554 if (pset.logfile)
555 {
556 fprintf(pset.logfile,
557 _("********* QUERY **********\n"
558 "%s\n"
559 "**************************\n\n"), query);
560 fflush(pset.logfile);
561 }
562
563 if (pset.echo_hidden == PSQL_ECHO_HIDDEN_NOEXEC)
564 return NULL;
565 }
566
567 SetCancelConn(pset.db);
568
569 res = PQexec(pset.db, query);
570
571 ResetCancelConn();
572
573 if (!AcceptResult(res))
574 {
575 ClearOrSaveResult(res);
576 res = NULL;
577 }
578
579 return res;
580 }
581
582
583 /*
584 * PSQLexecWatch
585 *
586 * This function is used for \watch command to send the query to
587 * the server and print out the results.
588 *
589 * Returns 1 if the query executed successfully, 0 if it cannot be repeated,
590 * e.g., because of the interrupt, -1 on error.
591 */
592 int
PSQLexecWatch(const char * query,const printQueryOpt * opt)593 PSQLexecWatch(const char *query, const printQueryOpt *opt)
594 {
595 PGresult *res;
596 double elapsed_msec = 0;
597 instr_time before;
598 instr_time after;
599
600 if (!pset.db)
601 {
602 pg_log_error("You are currently not connected to a database.");
603 return 0;
604 }
605
606 SetCancelConn(pset.db);
607
608 if (pset.timing)
609 INSTR_TIME_SET_CURRENT(before);
610
611 res = PQexec(pset.db, query);
612
613 ResetCancelConn();
614
615 if (!AcceptResult(res))
616 {
617 ClearOrSaveResult(res);
618 return 0;
619 }
620
621 if (pset.timing)
622 {
623 INSTR_TIME_SET_CURRENT(after);
624 INSTR_TIME_SUBTRACT(after, before);
625 elapsed_msec = INSTR_TIME_GET_MILLISEC(after);
626 }
627
628 /*
629 * If SIGINT is sent while the query is processing, the interrupt will be
630 * consumed. The user's intention, though, is to cancel the entire watch
631 * process, so detect a sent cancellation request and exit in this case.
632 */
633 if (cancel_pressed)
634 {
635 PQclear(res);
636 return 0;
637 }
638
639 switch (PQresultStatus(res))
640 {
641 case PGRES_TUPLES_OK:
642 printQuery(res, opt, pset.queryFout, false, pset.logfile);
643 break;
644
645 case PGRES_COMMAND_OK:
646 fprintf(pset.queryFout, "%s\n%s\n\n", opt->title, PQcmdStatus(res));
647 break;
648
649 case PGRES_EMPTY_QUERY:
650 pg_log_error("\\watch cannot be used with an empty query");
651 PQclear(res);
652 return -1;
653
654 case PGRES_COPY_OUT:
655 case PGRES_COPY_IN:
656 case PGRES_COPY_BOTH:
657 pg_log_error("\\watch cannot be used with COPY");
658 PQclear(res);
659 return -1;
660
661 default:
662 pg_log_error("unexpected result status for \\watch");
663 PQclear(res);
664 return -1;
665 }
666
667 PQclear(res);
668
669 fflush(pset.queryFout);
670
671 /* Possible microtiming output */
672 if (pset.timing)
673 PrintTiming(elapsed_msec);
674
675 return 1;
676 }
677
678
679 /*
680 * PrintNotifications: check for asynchronous notifications, and print them out
681 */
682 static void
PrintNotifications(void)683 PrintNotifications(void)
684 {
685 PGnotify *notify;
686
687 PQconsumeInput(pset.db);
688 while ((notify = PQnotifies(pset.db)) != NULL)
689 {
690 /* for backward compatibility, only show payload if nonempty */
691 if (notify->extra[0])
692 fprintf(pset.queryFout, _("Asynchronous notification \"%s\" with payload \"%s\" received from server process with PID %d.\n"),
693 notify->relname, notify->extra, notify->be_pid);
694 else
695 fprintf(pset.queryFout, _("Asynchronous notification \"%s\" received from server process with PID %d.\n"),
696 notify->relname, notify->be_pid);
697 fflush(pset.queryFout);
698 PQfreemem(notify);
699 PQconsumeInput(pset.db);
700 }
701 }
702
703
704 /*
705 * PrintQueryTuples: assuming query result is OK, print its tuples
706 *
707 * Returns true if successful, false otherwise.
708 */
709 static bool
PrintQueryTuples(const PGresult * results)710 PrintQueryTuples(const PGresult *results)
711 {
712 bool result = true;
713
714 /* write output to \g argument, if any */
715 if (pset.gfname)
716 {
717 FILE *fout;
718 bool is_pipe;
719
720 if (!openQueryOutputFile(pset.gfname, &fout, &is_pipe))
721 return false;
722 if (is_pipe)
723 disable_sigpipe_trap();
724
725 printQuery(results, &pset.popt, fout, false, pset.logfile);
726 if (ferror(fout))
727 {
728 pg_log_error("could not print result table: %m");
729 result = false;
730 }
731
732 if (is_pipe)
733 {
734 pclose(fout);
735 restore_sigpipe_trap();
736 }
737 else
738 fclose(fout);
739 }
740 else
741 {
742 printQuery(results, &pset.popt, pset.queryFout, false, pset.logfile);
743 if (ferror(pset.queryFout))
744 {
745 pg_log_error("could not print result table: %m");
746 result = false;
747 }
748 }
749
750 return result;
751 }
752
753
754 /*
755 * StoreQueryTuple: assuming query result is OK, save data into variables
756 *
757 * Returns true if successful, false otherwise.
758 */
759 static bool
StoreQueryTuple(const PGresult * result)760 StoreQueryTuple(const PGresult *result)
761 {
762 bool success = true;
763
764 if (PQntuples(result) < 1)
765 {
766 pg_log_error("no rows returned for \\gset");
767 success = false;
768 }
769 else if (PQntuples(result) > 1)
770 {
771 pg_log_error("more than one row returned for \\gset");
772 success = false;
773 }
774 else
775 {
776 int i;
777
778 for (i = 0; i < PQnfields(result); i++)
779 {
780 char *colname = PQfname(result, i);
781 char *varname;
782 char *value;
783
784 /* concatenate prefix and column name */
785 varname = psprintf("%s%s", pset.gset_prefix, colname);
786
787 if (VariableHasHook(pset.vars, varname))
788 {
789 pg_log_warning("attempt to \\gset into specially treated variable \"%s\" ignored",
790 varname);
791 continue;
792 }
793
794 if (!PQgetisnull(result, 0, i))
795 value = PQgetvalue(result, 0, i);
796 else
797 {
798 /* for NULL value, unset rather than set the variable */
799 value = NULL;
800 }
801
802 if (!SetVariable(pset.vars, varname, value))
803 {
804 free(varname);
805 success = false;
806 break;
807 }
808
809 free(varname);
810 }
811 }
812
813 return success;
814 }
815
816
817 /*
818 * ExecQueryTuples: assuming query result is OK, execute each query
819 * result field as a SQL statement
820 *
821 * Returns true if successful, false otherwise.
822 */
823 static bool
ExecQueryTuples(const PGresult * result)824 ExecQueryTuples(const PGresult *result)
825 {
826 bool success = true;
827 int nrows = PQntuples(result);
828 int ncolumns = PQnfields(result);
829 int r,
830 c;
831
832 /*
833 * We must turn off gexec_flag to avoid infinite recursion. Note that
834 * this allows ExecQueryUsingCursor to be applied to the individual query
835 * results. SendQuery prevents it from being applied when fetching the
836 * queries-to-execute, because it can't handle recursion either.
837 */
838 pset.gexec_flag = false;
839
840 for (r = 0; r < nrows; r++)
841 {
842 for (c = 0; c < ncolumns; c++)
843 {
844 if (!PQgetisnull(result, r, c))
845 {
846 const char *query = PQgetvalue(result, r, c);
847
848 /* Abandon execution if cancel_pressed */
849 if (cancel_pressed)
850 goto loop_exit;
851
852 /*
853 * ECHO_ALL mode should echo these queries, but SendQuery
854 * assumes that MainLoop did that, so we have to do it here.
855 */
856 if (pset.echo == PSQL_ECHO_ALL && !pset.singlestep)
857 {
858 puts(query);
859 fflush(stdout);
860 }
861
862 if (!SendQuery(query))
863 {
864 /* Error - abandon execution if ON_ERROR_STOP */
865 success = false;
866 if (pset.on_error_stop)
867 goto loop_exit;
868 }
869 }
870 }
871 }
872
873 loop_exit:
874
875 /*
876 * Restore state. We know gexec_flag was on, else we'd not be here. (We
877 * also know it'll get turned off at end of command, but that's not ours
878 * to do here.)
879 */
880 pset.gexec_flag = true;
881
882 /* Return true if all queries were successful */
883 return success;
884 }
885
886
887 /*
888 * ProcessResult: utility function for use by SendQuery() only
889 *
890 * When our command string contained a COPY FROM STDIN or COPY TO STDOUT,
891 * PQexec() has stopped at the PGresult associated with the first such
892 * command. In that event, we'll marshal data for the COPY and then cycle
893 * through any subsequent PGresult objects.
894 *
895 * When the command string contained no such COPY command, this function
896 * degenerates to an AcceptResult() call.
897 *
898 * Changes its argument to point to the last PGresult of the command string,
899 * or NULL if that result was for a COPY TO STDOUT. (Returning NULL prevents
900 * the command status from being printed, which we want in that case so that
901 * the status line doesn't get taken as part of the COPY data.)
902 *
903 * Returns true on complete success, false otherwise. Possible failure modes
904 * include purely client-side problems; check the transaction status for the
905 * server-side opinion.
906 */
907 static bool
ProcessResult(PGresult ** results)908 ProcessResult(PGresult **results)
909 {
910 bool success = true;
911 bool first_cycle = true;
912
913 for (;;)
914 {
915 ExecStatusType result_status;
916 bool is_copy;
917 PGresult *next_result;
918
919 if (!AcceptResult(*results))
920 {
921 /*
922 * Failure at this point is always a server-side failure or a
923 * failure to submit the command string. Either way, we're
924 * finished with this command string.
925 */
926 success = false;
927 break;
928 }
929
930 result_status = PQresultStatus(*results);
931 switch (result_status)
932 {
933 case PGRES_EMPTY_QUERY:
934 case PGRES_COMMAND_OK:
935 case PGRES_TUPLES_OK:
936 is_copy = false;
937 break;
938
939 case PGRES_COPY_OUT:
940 case PGRES_COPY_IN:
941 is_copy = true;
942 break;
943
944 default:
945 /* AcceptResult() should have caught anything else. */
946 is_copy = false;
947 pg_log_error("unexpected PQresultStatus: %d", result_status);
948 break;
949 }
950
951 if (is_copy)
952 {
953 /*
954 * Marshal the COPY data. Either subroutine will get the
955 * connection out of its COPY state, then call PQresultStatus()
956 * once and report any error.
957 *
958 * For COPY OUT, direct the output to pset.copyStream if it's set,
959 * otherwise to pset.gfname if it's set, otherwise to queryFout.
960 * For COPY IN, use pset.copyStream as data source if it's set,
961 * otherwise cur_cmd_source.
962 */
963 FILE *copystream;
964 PGresult *copy_result;
965
966 SetCancelConn(pset.db);
967 if (result_status == PGRES_COPY_OUT)
968 {
969 bool need_close = false;
970 bool is_pipe = false;
971
972 if (pset.copyStream)
973 {
974 /* invoked by \copy */
975 copystream = pset.copyStream;
976 }
977 else if (pset.gfname)
978 {
979 /* invoked by \g */
980 if (openQueryOutputFile(pset.gfname,
981 ©stream, &is_pipe))
982 {
983 need_close = true;
984 if (is_pipe)
985 disable_sigpipe_trap();
986 }
987 else
988 copystream = NULL; /* discard COPY data entirely */
989 }
990 else
991 {
992 /* fall back to the generic query output stream */
993 copystream = pset.queryFout;
994 }
995
996 success = handleCopyOut(pset.db,
997 copystream,
998 ©_result)
999 && success
1000 && (copystream != NULL);
1001
1002 /*
1003 * Suppress status printing if the report would go to the same
1004 * place as the COPY data just went. Note this doesn't
1005 * prevent error reporting, since handleCopyOut did that.
1006 */
1007 if (copystream == pset.queryFout)
1008 {
1009 PQclear(copy_result);
1010 copy_result = NULL;
1011 }
1012
1013 if (need_close)
1014 {
1015 /* close \g argument file/pipe */
1016 if (is_pipe)
1017 {
1018 pclose(copystream);
1019 restore_sigpipe_trap();
1020 }
1021 else
1022 {
1023 fclose(copystream);
1024 }
1025 }
1026 }
1027 else
1028 {
1029 /* COPY IN */
1030 copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source;
1031 success = handleCopyIn(pset.db,
1032 copystream,
1033 PQbinaryTuples(*results),
1034 ©_result) && success;
1035 }
1036 ResetCancelConn();
1037
1038 /*
1039 * Replace the PGRES_COPY_OUT/IN result with COPY command's exit
1040 * status, or with NULL if we want to suppress printing anything.
1041 */
1042 PQclear(*results);
1043 *results = copy_result;
1044 }
1045 else if (first_cycle)
1046 {
1047 /* fast path: no COPY commands; PQexec visited all results */
1048 break;
1049 }
1050
1051 /*
1052 * Check PQgetResult() again. In the typical case of a single-command
1053 * string, it will return NULL. Otherwise, we'll have other results
1054 * to process that may include other COPYs. We keep the last result.
1055 */
1056 next_result = PQgetResult(pset.db);
1057 if (!next_result)
1058 break;
1059
1060 PQclear(*results);
1061 *results = next_result;
1062 first_cycle = false;
1063 }
1064
1065 SetResultVariables(*results, success);
1066
1067 /* may need this to recover from conn loss during COPY */
1068 if (!first_cycle && !CheckConnection())
1069 return false;
1070
1071 return success;
1072 }
1073
1074
1075 /*
1076 * PrintQueryStatus: report command status as required
1077 *
1078 * Note: Utility function for use by PrintQueryResults() only.
1079 */
1080 static void
PrintQueryStatus(PGresult * results)1081 PrintQueryStatus(PGresult *results)
1082 {
1083 char buf[16];
1084
1085 if (!pset.quiet)
1086 {
1087 if (pset.popt.topt.format == PRINT_HTML)
1088 {
1089 fputs("<p>", pset.queryFout);
1090 html_escaped_print(PQcmdStatus(results), pset.queryFout);
1091 fputs("</p>\n", pset.queryFout);
1092 }
1093 else
1094 fprintf(pset.queryFout, "%s\n", PQcmdStatus(results));
1095 }
1096
1097 if (pset.logfile)
1098 fprintf(pset.logfile, "%s\n", PQcmdStatus(results));
1099
1100 snprintf(buf, sizeof(buf), "%u", (unsigned int) PQoidValue(results));
1101 SetVariable(pset.vars, "LASTOID", buf);
1102 }
1103
1104
1105 /*
1106 * PrintQueryResults: print out (or store or execute) query results as required
1107 *
1108 * Note: Utility function for use by SendQuery() only.
1109 *
1110 * Returns true if the query executed successfully, false otherwise.
1111 */
1112 static bool
PrintQueryResults(PGresult * results)1113 PrintQueryResults(PGresult *results)
1114 {
1115 bool success;
1116 const char *cmdstatus;
1117
1118 if (!results)
1119 return false;
1120
1121 switch (PQresultStatus(results))
1122 {
1123 case PGRES_TUPLES_OK:
1124 /* store or execute or print the data ... */
1125 if (pset.gset_prefix)
1126 success = StoreQueryTuple(results);
1127 else if (pset.gexec_flag)
1128 success = ExecQueryTuples(results);
1129 else if (pset.crosstab_flag)
1130 success = PrintResultsInCrosstab(results);
1131 else
1132 success = PrintQueryTuples(results);
1133 /* if it's INSERT/UPDATE/DELETE RETURNING, also print status */
1134 cmdstatus = PQcmdStatus(results);
1135 if (strncmp(cmdstatus, "INSERT", 6) == 0 ||
1136 strncmp(cmdstatus, "UPDATE", 6) == 0 ||
1137 strncmp(cmdstatus, "DELETE", 6) == 0)
1138 PrintQueryStatus(results);
1139 break;
1140
1141 case PGRES_COMMAND_OK:
1142 PrintQueryStatus(results);
1143 success = true;
1144 break;
1145
1146 case PGRES_EMPTY_QUERY:
1147 success = true;
1148 break;
1149
1150 case PGRES_COPY_OUT:
1151 case PGRES_COPY_IN:
1152 /* nothing to do here */
1153 success = true;
1154 break;
1155
1156 case PGRES_BAD_RESPONSE:
1157 case PGRES_NONFATAL_ERROR:
1158 case PGRES_FATAL_ERROR:
1159 success = false;
1160 break;
1161
1162 default:
1163 success = false;
1164 pg_log_error("unexpected PQresultStatus: %d",
1165 PQresultStatus(results));
1166 break;
1167 }
1168
1169 fflush(pset.queryFout);
1170
1171 return success;
1172 }
1173
1174
1175 /*
1176 * SendQuery: send the query string to the backend
1177 * (and print out results)
1178 *
1179 * Note: This is the "front door" way to send a query. That is, use it to
1180 * send queries actually entered by the user. These queries will be subject to
1181 * single step mode.
1182 * To send "back door" queries (generated by slash commands, etc.) in a
1183 * controlled way, use PSQLexec().
1184 *
1185 * Returns true if the query executed successfully, false otherwise.
1186 */
1187 bool
SendQuery(const char * query)1188 SendQuery(const char *query)
1189 {
1190 PGresult *results;
1191 PGTransactionStatusType transaction_status;
1192 double elapsed_msec = 0;
1193 bool OK = false;
1194 int i;
1195 bool on_error_rollback_savepoint = false;
1196 static bool on_error_rollback_warning = false;
1197
1198 if (!pset.db)
1199 {
1200 pg_log_error("You are currently not connected to a database.");
1201 goto sendquery_cleanup;
1202 }
1203
1204 if (pset.singlestep)
1205 {
1206 char buf[3];
1207
1208 fflush(stderr);
1209 printf(_("***(Single step mode: verify command)*******************************************\n"
1210 "%s\n"
1211 "***(press return to proceed or enter x and return to cancel)********************\n"),
1212 query);
1213 fflush(stdout);
1214 if (fgets(buf, sizeof(buf), stdin) != NULL)
1215 if (buf[0] == 'x')
1216 goto sendquery_cleanup;
1217 if (cancel_pressed)
1218 goto sendquery_cleanup;
1219 }
1220 else if (pset.echo == PSQL_ECHO_QUERIES)
1221 {
1222 puts(query);
1223 fflush(stdout);
1224 }
1225
1226 if (pset.logfile)
1227 {
1228 fprintf(pset.logfile,
1229 _("********* QUERY **********\n"
1230 "%s\n"
1231 "**************************\n\n"), query);
1232 fflush(pset.logfile);
1233 }
1234
1235 SetCancelConn(pset.db);
1236
1237 transaction_status = PQtransactionStatus(pset.db);
1238
1239 if (transaction_status == PQTRANS_IDLE &&
1240 !pset.autocommit &&
1241 !command_no_begin(query))
1242 {
1243 results = PQexec(pset.db, "BEGIN");
1244 if (PQresultStatus(results) != PGRES_COMMAND_OK)
1245 {
1246 pg_log_info("%s", PQerrorMessage(pset.db));
1247 ClearOrSaveResult(results);
1248 ResetCancelConn();
1249 goto sendquery_cleanup;
1250 }
1251 ClearOrSaveResult(results);
1252 transaction_status = PQtransactionStatus(pset.db);
1253 }
1254
1255 if (transaction_status == PQTRANS_INTRANS &&
1256 pset.on_error_rollback != PSQL_ERROR_ROLLBACK_OFF &&
1257 (pset.cur_cmd_interactive ||
1258 pset.on_error_rollback == PSQL_ERROR_ROLLBACK_ON))
1259 {
1260 if (on_error_rollback_warning == false && pset.sversion < 80000)
1261 {
1262 char sverbuf[32];
1263
1264 pg_log_warning("The server (version %s) does not support savepoints for ON_ERROR_ROLLBACK.",
1265 formatPGVersionNumber(pset.sversion, false,
1266 sverbuf, sizeof(sverbuf)));
1267 on_error_rollback_warning = true;
1268 }
1269 else
1270 {
1271 results = PQexec(pset.db, "SAVEPOINT pg_psql_temporary_savepoint");
1272 if (PQresultStatus(results) != PGRES_COMMAND_OK)
1273 {
1274 pg_log_info("%s", PQerrorMessage(pset.db));
1275 ClearOrSaveResult(results);
1276 ResetCancelConn();
1277 goto sendquery_cleanup;
1278 }
1279 ClearOrSaveResult(results);
1280 on_error_rollback_savepoint = true;
1281 }
1282 }
1283
1284 if (pset.gdesc_flag)
1285 {
1286 /* Describe query's result columns, without executing it */
1287 OK = DescribeQuery(query, &elapsed_msec);
1288 ResetCancelConn();
1289 results = NULL; /* PQclear(NULL) does nothing */
1290 }
1291 else if (pset.fetch_count <= 0 || pset.gexec_flag ||
1292 pset.crosstab_flag || !is_select_command(query))
1293 {
1294 /* Default fetch-it-all-and-print mode */
1295 instr_time before,
1296 after;
1297
1298 if (pset.timing)
1299 INSTR_TIME_SET_CURRENT(before);
1300
1301 results = PQexec(pset.db, query);
1302
1303 /* these operations are included in the timing result: */
1304 ResetCancelConn();
1305 OK = ProcessResult(&results);
1306
1307 if (pset.timing)
1308 {
1309 INSTR_TIME_SET_CURRENT(after);
1310 INSTR_TIME_SUBTRACT(after, before);
1311 elapsed_msec = INSTR_TIME_GET_MILLISEC(after);
1312 }
1313
1314 /* but printing results isn't: */
1315 if (OK && results)
1316 OK = PrintQueryResults(results);
1317 }
1318 else
1319 {
1320 /* Fetch-in-segments mode */
1321 OK = ExecQueryUsingCursor(query, &elapsed_msec);
1322 ResetCancelConn();
1323 results = NULL; /* PQclear(NULL) does nothing */
1324 }
1325
1326 if (!OK && pset.echo == PSQL_ECHO_ERRORS)
1327 pg_log_info("STATEMENT: %s", query);
1328
1329 /* If we made a temporary savepoint, possibly release/rollback */
1330 if (on_error_rollback_savepoint)
1331 {
1332 const char *svptcmd = NULL;
1333
1334 transaction_status = PQtransactionStatus(pset.db);
1335
1336 switch (transaction_status)
1337 {
1338 case PQTRANS_INERROR:
1339 /* We always rollback on an error */
1340 svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint";
1341 break;
1342
1343 case PQTRANS_IDLE:
1344 /* If they are no longer in a transaction, then do nothing */
1345 break;
1346
1347 case PQTRANS_INTRANS:
1348
1349 /*
1350 * Do nothing if they are messing with savepoints themselves:
1351 * If the user did COMMIT AND CHAIN, RELEASE or ROLLBACK, our
1352 * savepoint is gone. If they issued a SAVEPOINT, releasing
1353 * ours would remove theirs.
1354 */
1355 if (results &&
1356 (strcmp(PQcmdStatus(results), "COMMIT") == 0 ||
1357 strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
1358 strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
1359 strcmp(PQcmdStatus(results), "ROLLBACK") == 0))
1360 svptcmd = NULL;
1361 else
1362 svptcmd = "RELEASE pg_psql_temporary_savepoint";
1363 break;
1364
1365 case PQTRANS_ACTIVE:
1366 case PQTRANS_UNKNOWN:
1367 default:
1368 OK = false;
1369 /* PQTRANS_UNKNOWN is expected given a broken connection. */
1370 if (transaction_status != PQTRANS_UNKNOWN || ConnectionUp())
1371 pg_log_error("unexpected transaction status (%d)",
1372 transaction_status);
1373 break;
1374 }
1375
1376 if (svptcmd)
1377 {
1378 PGresult *svptres;
1379
1380 svptres = PQexec(pset.db, svptcmd);
1381 if (PQresultStatus(svptres) != PGRES_COMMAND_OK)
1382 {
1383 pg_log_info("%s", PQerrorMessage(pset.db));
1384 ClearOrSaveResult(svptres);
1385 OK = false;
1386
1387 PQclear(results);
1388 ResetCancelConn();
1389 goto sendquery_cleanup;
1390 }
1391 PQclear(svptres);
1392 }
1393 }
1394
1395 ClearOrSaveResult(results);
1396
1397 /* Possible microtiming output */
1398 if (pset.timing)
1399 PrintTiming(elapsed_msec);
1400
1401 /* check for events that may occur during query execution */
1402
1403 if (pset.encoding != PQclientEncoding(pset.db) &&
1404 PQclientEncoding(pset.db) >= 0)
1405 {
1406 /* track effects of SET CLIENT_ENCODING */
1407 pset.encoding = PQclientEncoding(pset.db);
1408 pset.popt.topt.encoding = pset.encoding;
1409 SetVariable(pset.vars, "ENCODING",
1410 pg_encoding_to_char(pset.encoding));
1411 }
1412
1413 PrintNotifications();
1414
1415 /* perform cleanup that should occur after any attempted query */
1416
1417 sendquery_cleanup:
1418
1419 /* reset \g's output-to-filename trigger */
1420 if (pset.gfname)
1421 {
1422 free(pset.gfname);
1423 pset.gfname = NULL;
1424 }
1425
1426 /* restore print settings if \g changed them */
1427 if (pset.gsavepopt)
1428 {
1429 restorePsetInfo(&pset.popt, pset.gsavepopt);
1430 pset.gsavepopt = NULL;
1431 }
1432
1433 /* reset \gset trigger */
1434 if (pset.gset_prefix)
1435 {
1436 free(pset.gset_prefix);
1437 pset.gset_prefix = NULL;
1438 }
1439
1440 /* reset \gdesc trigger */
1441 pset.gdesc_flag = false;
1442
1443 /* reset \gexec trigger */
1444 pset.gexec_flag = false;
1445
1446 /* reset \crosstabview trigger */
1447 pset.crosstab_flag = false;
1448 for (i = 0; i < lengthof(pset.ctv_args); i++)
1449 {
1450 pg_free(pset.ctv_args[i]);
1451 pset.ctv_args[i] = NULL;
1452 }
1453
1454 return OK;
1455 }
1456
1457
1458 /*
1459 * DescribeQuery: describe the result columns of a query, without executing it
1460 *
1461 * Returns true if the operation executed successfully, false otherwise.
1462 *
1463 * If pset.timing is on, total query time (exclusive of result-printing) is
1464 * stored into *elapsed_msec.
1465 */
1466 static bool
DescribeQuery(const char * query,double * elapsed_msec)1467 DescribeQuery(const char *query, double *elapsed_msec)
1468 {
1469 PGresult *results;
1470 bool OK;
1471 instr_time before,
1472 after;
1473
1474 *elapsed_msec = 0;
1475
1476 if (pset.timing)
1477 INSTR_TIME_SET_CURRENT(before);
1478
1479 /*
1480 * To parse the query but not execute it, we prepare it, using the unnamed
1481 * prepared statement. This is invisible to psql users, since there's no
1482 * way to access the unnamed prepared statement from psql user space. The
1483 * next Parse or Query protocol message would overwrite the statement
1484 * anyway. (So there's no great need to clear it when done, which is a
1485 * good thing because libpq provides no easy way to do that.)
1486 */
1487 results = PQprepare(pset.db, "", query, 0, NULL);
1488 if (PQresultStatus(results) != PGRES_COMMAND_OK)
1489 {
1490 pg_log_info("%s", PQerrorMessage(pset.db));
1491 SetResultVariables(results, false);
1492 ClearOrSaveResult(results);
1493 return false;
1494 }
1495 PQclear(results);
1496
1497 results = PQdescribePrepared(pset.db, "");
1498 OK = AcceptResult(results) &&
1499 (PQresultStatus(results) == PGRES_COMMAND_OK);
1500 if (OK && results)
1501 {
1502 if (PQnfields(results) > 0)
1503 {
1504 PQExpBufferData buf;
1505 int i;
1506
1507 initPQExpBuffer(&buf);
1508
1509 printfPQExpBuffer(&buf,
1510 "SELECT name AS \"%s\", pg_catalog.format_type(tp, tpm) AS \"%s\"\n"
1511 "FROM (VALUES ",
1512 gettext_noop("Column"),
1513 gettext_noop("Type"));
1514
1515 for (i = 0; i < PQnfields(results); i++)
1516 {
1517 const char *name;
1518 char *escname;
1519
1520 if (i > 0)
1521 appendPQExpBufferStr(&buf, ",");
1522
1523 name = PQfname(results, i);
1524 escname = PQescapeLiteral(pset.db, name, strlen(name));
1525
1526 if (escname == NULL)
1527 {
1528 pg_log_info("%s", PQerrorMessage(pset.db));
1529 PQclear(results);
1530 termPQExpBuffer(&buf);
1531 return false;
1532 }
1533
1534 appendPQExpBuffer(&buf, "(%s, '%u'::pg_catalog.oid, %d)",
1535 escname,
1536 PQftype(results, i),
1537 PQfmod(results, i));
1538
1539 PQfreemem(escname);
1540 }
1541
1542 appendPQExpBufferStr(&buf, ") s(name, tp, tpm)");
1543 PQclear(results);
1544
1545 results = PQexec(pset.db, buf.data);
1546 OK = AcceptResult(results);
1547
1548 if (pset.timing)
1549 {
1550 INSTR_TIME_SET_CURRENT(after);
1551 INSTR_TIME_SUBTRACT(after, before);
1552 *elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
1553 }
1554
1555 if (OK && results)
1556 OK = PrintQueryResults(results);
1557
1558 termPQExpBuffer(&buf);
1559 }
1560 else
1561 fprintf(pset.queryFout,
1562 _("The command has no result, or the result has no columns.\n"));
1563 }
1564
1565 SetResultVariables(results, OK);
1566 ClearOrSaveResult(results);
1567
1568 return OK;
1569 }
1570
1571
1572 /*
1573 * ExecQueryUsingCursor: run a SELECT-like query using a cursor
1574 *
1575 * This feature allows result sets larger than RAM to be dealt with.
1576 *
1577 * Returns true if the query executed successfully, false otherwise.
1578 *
1579 * If pset.timing is on, total query time (exclusive of result-printing) is
1580 * stored into *elapsed_msec.
1581 */
1582 static bool
ExecQueryUsingCursor(const char * query,double * elapsed_msec)1583 ExecQueryUsingCursor(const char *query, double *elapsed_msec)
1584 {
1585 bool OK = true;
1586 PGresult *results;
1587 PQExpBufferData buf;
1588 printQueryOpt my_popt = pset.popt;
1589 FILE *fout;
1590 bool is_pipe;
1591 bool is_pager = false;
1592 bool started_txn = false;
1593 int64 total_tuples = 0;
1594 int ntuples;
1595 int fetch_count;
1596 char fetch_cmd[64];
1597 instr_time before,
1598 after;
1599 int flush_error;
1600
1601 *elapsed_msec = 0;
1602
1603 /* initialize print options for partial table output */
1604 my_popt.topt.start_table = true;
1605 my_popt.topt.stop_table = false;
1606 my_popt.topt.prior_records = 0;
1607
1608 if (pset.timing)
1609 INSTR_TIME_SET_CURRENT(before);
1610
1611 /* if we're not in a transaction, start one */
1612 if (PQtransactionStatus(pset.db) == PQTRANS_IDLE)
1613 {
1614 results = PQexec(pset.db, "BEGIN");
1615 OK = AcceptResult(results) &&
1616 (PQresultStatus(results) == PGRES_COMMAND_OK);
1617 ClearOrSaveResult(results);
1618 if (!OK)
1619 return false;
1620 started_txn = true;
1621 }
1622
1623 /* Send DECLARE CURSOR */
1624 initPQExpBuffer(&buf);
1625 appendPQExpBuffer(&buf, "DECLARE _psql_cursor NO SCROLL CURSOR FOR\n%s",
1626 query);
1627
1628 results = PQexec(pset.db, buf.data);
1629 OK = AcceptResult(results) &&
1630 (PQresultStatus(results) == PGRES_COMMAND_OK);
1631 if (!OK)
1632 SetResultVariables(results, OK);
1633 ClearOrSaveResult(results);
1634 termPQExpBuffer(&buf);
1635 if (!OK)
1636 goto cleanup;
1637
1638 if (pset.timing)
1639 {
1640 INSTR_TIME_SET_CURRENT(after);
1641 INSTR_TIME_SUBTRACT(after, before);
1642 *elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
1643 }
1644
1645 /*
1646 * In \gset mode, we force the fetch count to be 2, so that we will throw
1647 * the appropriate error if the query returns more than one row.
1648 */
1649 if (pset.gset_prefix)
1650 fetch_count = 2;
1651 else
1652 fetch_count = pset.fetch_count;
1653
1654 snprintf(fetch_cmd, sizeof(fetch_cmd),
1655 "FETCH FORWARD %d FROM _psql_cursor",
1656 fetch_count);
1657
1658 /* prepare to write output to \g argument, if any */
1659 if (pset.gfname)
1660 {
1661 if (!openQueryOutputFile(pset.gfname, &fout, &is_pipe))
1662 {
1663 OK = false;
1664 goto cleanup;
1665 }
1666 if (is_pipe)
1667 disable_sigpipe_trap();
1668 }
1669 else
1670 {
1671 fout = pset.queryFout;
1672 is_pipe = false; /* doesn't matter */
1673 }
1674
1675 /* clear any pre-existing error indication on the output stream */
1676 clearerr(fout);
1677
1678 for (;;)
1679 {
1680 if (pset.timing)
1681 INSTR_TIME_SET_CURRENT(before);
1682
1683 /* get fetch_count tuples at a time */
1684 results = PQexec(pset.db, fetch_cmd);
1685
1686 if (pset.timing)
1687 {
1688 INSTR_TIME_SET_CURRENT(after);
1689 INSTR_TIME_SUBTRACT(after, before);
1690 *elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
1691 }
1692
1693 if (PQresultStatus(results) != PGRES_TUPLES_OK)
1694 {
1695 /* shut down pager before printing error message */
1696 if (is_pager)
1697 {
1698 ClosePager(fout);
1699 is_pager = false;
1700 }
1701
1702 OK = AcceptResult(results);
1703 Assert(!OK);
1704 SetResultVariables(results, OK);
1705 ClearOrSaveResult(results);
1706 break;
1707 }
1708
1709 if (pset.gset_prefix)
1710 {
1711 /* StoreQueryTuple will complain if not exactly one row */
1712 OK = StoreQueryTuple(results);
1713 ClearOrSaveResult(results);
1714 break;
1715 }
1716
1717 /*
1718 * Note we do not deal with \gdesc, \gexec or \crosstabview modes here
1719 */
1720
1721 ntuples = PQntuples(results);
1722 total_tuples += ntuples;
1723
1724 if (ntuples < fetch_count)
1725 {
1726 /* this is the last result set, so allow footer decoration */
1727 my_popt.topt.stop_table = true;
1728 }
1729 else if (fout == stdout && !is_pager)
1730 {
1731 /*
1732 * If query requires multiple result sets, hack to ensure that
1733 * only one pager instance is used for the whole mess
1734 */
1735 fout = PageOutput(INT_MAX, &(my_popt.topt));
1736 is_pager = true;
1737 }
1738
1739 printQuery(results, &my_popt, fout, is_pager, pset.logfile);
1740
1741 ClearOrSaveResult(results);
1742
1743 /* after the first result set, disallow header decoration */
1744 my_popt.topt.start_table = false;
1745 my_popt.topt.prior_records += ntuples;
1746
1747 /*
1748 * Make sure to flush the output stream, so intermediate results are
1749 * visible to the client immediately. We check the results because if
1750 * the pager dies/exits/etc, there's no sense throwing more data at
1751 * it.
1752 */
1753 flush_error = fflush(fout);
1754
1755 /*
1756 * Check if we are at the end, if a cancel was pressed, or if there
1757 * were any errors either trying to flush out the results, or more
1758 * generally on the output stream at all. If we hit any errors
1759 * writing things to the stream, we presume $PAGER has disappeared and
1760 * stop bothering to pull down more data.
1761 */
1762 if (ntuples < fetch_count || cancel_pressed || flush_error ||
1763 ferror(fout))
1764 break;
1765 }
1766
1767 if (pset.gfname)
1768 {
1769 /* close \g argument file/pipe */
1770 if (is_pipe)
1771 {
1772 pclose(fout);
1773 restore_sigpipe_trap();
1774 }
1775 else
1776 fclose(fout);
1777 }
1778 else if (is_pager)
1779 {
1780 /* close transient pager */
1781 ClosePager(fout);
1782 }
1783
1784 if (OK)
1785 {
1786 /*
1787 * We don't have a PGresult here, and even if we did it wouldn't have
1788 * the right row count, so fake SetResultVariables(). In error cases,
1789 * we already set the result variables above.
1790 */
1791 char buf[32];
1792
1793 SetVariable(pset.vars, "ERROR", "false");
1794 SetVariable(pset.vars, "SQLSTATE", "00000");
1795 snprintf(buf, sizeof(buf), INT64_FORMAT, total_tuples);
1796 SetVariable(pset.vars, "ROW_COUNT", buf);
1797 }
1798
1799 cleanup:
1800 if (pset.timing)
1801 INSTR_TIME_SET_CURRENT(before);
1802
1803 /*
1804 * We try to close the cursor on either success or failure, but on failure
1805 * ignore the result (it's probably just a bleat about being in an aborted
1806 * transaction)
1807 */
1808 results = PQexec(pset.db, "CLOSE _psql_cursor");
1809 if (OK)
1810 {
1811 OK = AcceptResult(results) &&
1812 (PQresultStatus(results) == PGRES_COMMAND_OK);
1813 ClearOrSaveResult(results);
1814 }
1815 else
1816 PQclear(results);
1817
1818 if (started_txn)
1819 {
1820 results = PQexec(pset.db, OK ? "COMMIT" : "ROLLBACK");
1821 OK &= AcceptResult(results) &&
1822 (PQresultStatus(results) == PGRES_COMMAND_OK);
1823 ClearOrSaveResult(results);
1824 }
1825
1826 if (pset.timing)
1827 {
1828 INSTR_TIME_SET_CURRENT(after);
1829 INSTR_TIME_SUBTRACT(after, before);
1830 *elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
1831 }
1832
1833 return OK;
1834 }
1835
1836
1837 /*
1838 * Advance the given char pointer over white space and SQL comments.
1839 */
1840 static const char *
skip_white_space(const char * query)1841 skip_white_space(const char *query)
1842 {
1843 int cnestlevel = 0; /* slash-star comment nest level */
1844
1845 while (*query)
1846 {
1847 int mblen = PQmblenBounded(query, pset.encoding);
1848
1849 /*
1850 * Note: we assume the encoding is a superset of ASCII, so that for
1851 * example "query[0] == '/'" is meaningful. However, we do NOT assume
1852 * that the second and subsequent bytes of a multibyte character
1853 * couldn't look like ASCII characters; so it is critical to advance
1854 * by mblen, not 1, whenever we haven't exactly identified the
1855 * character we are skipping over.
1856 */
1857 if (isspace((unsigned char) *query))
1858 query += mblen;
1859 else if (query[0] == '/' && query[1] == '*')
1860 {
1861 cnestlevel++;
1862 query += 2;
1863 }
1864 else if (cnestlevel > 0 && query[0] == '*' && query[1] == '/')
1865 {
1866 cnestlevel--;
1867 query += 2;
1868 }
1869 else if (cnestlevel == 0 && query[0] == '-' && query[1] == '-')
1870 {
1871 query += 2;
1872
1873 /*
1874 * We have to skip to end of line since any slash-star inside the
1875 * -- comment does NOT start a slash-star comment.
1876 */
1877 while (*query)
1878 {
1879 if (*query == '\n')
1880 {
1881 query++;
1882 break;
1883 }
1884 query += PQmblenBounded(query, pset.encoding);
1885 }
1886 }
1887 else if (cnestlevel > 0)
1888 query += mblen;
1889 else
1890 break; /* found first token */
1891 }
1892
1893 return query;
1894 }
1895
1896
1897 /*
1898 * Check whether a command is one of those for which we should NOT start
1899 * a new transaction block (ie, send a preceding BEGIN).
1900 *
1901 * These include the transaction control statements themselves, plus
1902 * certain statements that the backend disallows inside transaction blocks.
1903 */
1904 static bool
command_no_begin(const char * query)1905 command_no_begin(const char *query)
1906 {
1907 int wordlen;
1908
1909 /*
1910 * First we must advance over any whitespace and comments.
1911 */
1912 query = skip_white_space(query);
1913
1914 /*
1915 * Check word length (since "beginx" is not "begin").
1916 */
1917 wordlen = 0;
1918 while (isalpha((unsigned char) query[wordlen]))
1919 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
1920
1921 /*
1922 * Transaction control commands. These should include every keyword that
1923 * gives rise to a TransactionStmt in the backend grammar, except for the
1924 * savepoint-related commands.
1925 *
1926 * (We assume that START must be START TRANSACTION, since there is
1927 * presently no other "START foo" command.)
1928 */
1929 if (wordlen == 5 && pg_strncasecmp(query, "abort", 5) == 0)
1930 return true;
1931 if (wordlen == 5 && pg_strncasecmp(query, "begin", 5) == 0)
1932 return true;
1933 if (wordlen == 5 && pg_strncasecmp(query, "start", 5) == 0)
1934 return true;
1935 if (wordlen == 6 && pg_strncasecmp(query, "commit", 6) == 0)
1936 return true;
1937 if (wordlen == 3 && pg_strncasecmp(query, "end", 3) == 0)
1938 return true;
1939 if (wordlen == 8 && pg_strncasecmp(query, "rollback", 8) == 0)
1940 return true;
1941 if (wordlen == 7 && pg_strncasecmp(query, "prepare", 7) == 0)
1942 {
1943 /* PREPARE TRANSACTION is a TC command, PREPARE foo is not */
1944 query += wordlen;
1945
1946 query = skip_white_space(query);
1947
1948 wordlen = 0;
1949 while (isalpha((unsigned char) query[wordlen]))
1950 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
1951
1952 if (wordlen == 11 && pg_strncasecmp(query, "transaction", 11) == 0)
1953 return true;
1954 return false;
1955 }
1956
1957 /*
1958 * Commands not allowed within transactions. The statements checked for
1959 * here should be exactly those that call PreventInTransactionBlock() in
1960 * the backend.
1961 */
1962 if (wordlen == 6 && pg_strncasecmp(query, "vacuum", 6) == 0)
1963 return true;
1964 if (wordlen == 7 && pg_strncasecmp(query, "cluster", 7) == 0)
1965 {
1966 /* CLUSTER with any arguments is allowed in transactions */
1967 query += wordlen;
1968
1969 query = skip_white_space(query);
1970
1971 if (isalpha((unsigned char) query[0]))
1972 return false; /* has additional words */
1973 return true; /* it's CLUSTER without arguments */
1974 }
1975
1976 if (wordlen == 6 && pg_strncasecmp(query, "create", 6) == 0)
1977 {
1978 query += wordlen;
1979
1980 query = skip_white_space(query);
1981
1982 wordlen = 0;
1983 while (isalpha((unsigned char) query[wordlen]))
1984 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
1985
1986 if (wordlen == 8 && pg_strncasecmp(query, "database", 8) == 0)
1987 return true;
1988 if (wordlen == 10 && pg_strncasecmp(query, "tablespace", 10) == 0)
1989 return true;
1990
1991 /* CREATE [UNIQUE] INDEX CONCURRENTLY isn't allowed in xacts */
1992 if (wordlen == 6 && pg_strncasecmp(query, "unique", 6) == 0)
1993 {
1994 query += wordlen;
1995
1996 query = skip_white_space(query);
1997
1998 wordlen = 0;
1999 while (isalpha((unsigned char) query[wordlen]))
2000 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
2001 }
2002
2003 if (wordlen == 5 && pg_strncasecmp(query, "index", 5) == 0)
2004 {
2005 query += wordlen;
2006
2007 query = skip_white_space(query);
2008
2009 wordlen = 0;
2010 while (isalpha((unsigned char) query[wordlen]))
2011 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
2012
2013 if (wordlen == 12 && pg_strncasecmp(query, "concurrently", 12) == 0)
2014 return true;
2015 }
2016
2017 return false;
2018 }
2019
2020 if (wordlen == 5 && pg_strncasecmp(query, "alter", 5) == 0)
2021 {
2022 query += wordlen;
2023
2024 query = skip_white_space(query);
2025
2026 wordlen = 0;
2027 while (isalpha((unsigned char) query[wordlen]))
2028 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
2029
2030 /* ALTER SYSTEM isn't allowed in xacts */
2031 if (wordlen == 6 && pg_strncasecmp(query, "system", 6) == 0)
2032 return true;
2033
2034 return false;
2035 }
2036
2037 /*
2038 * Note: these tests will match DROP SYSTEM and REINDEX TABLESPACE, which
2039 * aren't really valid commands so we don't care much. The other four
2040 * possible matches are correct.
2041 */
2042 if ((wordlen == 4 && pg_strncasecmp(query, "drop", 4) == 0) ||
2043 (wordlen == 7 && pg_strncasecmp(query, "reindex", 7) == 0))
2044 {
2045 query += wordlen;
2046
2047 query = skip_white_space(query);
2048
2049 wordlen = 0;
2050 while (isalpha((unsigned char) query[wordlen]))
2051 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
2052
2053 if (wordlen == 8 && pg_strncasecmp(query, "database", 8) == 0)
2054 return true;
2055 if (wordlen == 6 && pg_strncasecmp(query, "system", 6) == 0)
2056 return true;
2057 if (wordlen == 10 && pg_strncasecmp(query, "tablespace", 10) == 0)
2058 return true;
2059 if (wordlen == 5 && (pg_strncasecmp(query, "index", 5) == 0 ||
2060 pg_strncasecmp(query, "table", 5) == 0))
2061 {
2062 query += wordlen;
2063 query = skip_white_space(query);
2064 wordlen = 0;
2065 while (isalpha((unsigned char) query[wordlen]))
2066 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
2067
2068 /*
2069 * REINDEX [ TABLE | INDEX ] CONCURRENTLY are not allowed in
2070 * xacts.
2071 */
2072 if (wordlen == 12 && pg_strncasecmp(query, "concurrently", 12) == 0)
2073 return true;
2074 }
2075
2076 /* DROP INDEX CONCURRENTLY isn't allowed in xacts */
2077 if (wordlen == 5 && pg_strncasecmp(query, "index", 5) == 0)
2078 {
2079 query += wordlen;
2080
2081 query = skip_white_space(query);
2082
2083 wordlen = 0;
2084 while (isalpha((unsigned char) query[wordlen]))
2085 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
2086
2087 if (wordlen == 12 && pg_strncasecmp(query, "concurrently", 12) == 0)
2088 return true;
2089
2090 return false;
2091 }
2092
2093 return false;
2094 }
2095
2096 /* DISCARD ALL isn't allowed in xacts, but other variants are allowed. */
2097 if (wordlen == 7 && pg_strncasecmp(query, "discard", 7) == 0)
2098 {
2099 query += wordlen;
2100
2101 query = skip_white_space(query);
2102
2103 wordlen = 0;
2104 while (isalpha((unsigned char) query[wordlen]))
2105 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
2106
2107 if (wordlen == 3 && pg_strncasecmp(query, "all", 3) == 0)
2108 return true;
2109 return false;
2110 }
2111
2112 return false;
2113 }
2114
2115
2116 /*
2117 * Check whether the specified command is a SELECT (or VALUES).
2118 */
2119 static bool
is_select_command(const char * query)2120 is_select_command(const char *query)
2121 {
2122 int wordlen;
2123
2124 /*
2125 * First advance over any whitespace, comments and left parentheses.
2126 */
2127 for (;;)
2128 {
2129 query = skip_white_space(query);
2130 if (query[0] == '(')
2131 query++;
2132 else
2133 break;
2134 }
2135
2136 /*
2137 * Check word length (since "selectx" is not "select").
2138 */
2139 wordlen = 0;
2140 while (isalpha((unsigned char) query[wordlen]))
2141 wordlen += PQmblenBounded(&query[wordlen], pset.encoding);
2142
2143 if (wordlen == 6 && pg_strncasecmp(query, "select", 6) == 0)
2144 return true;
2145
2146 if (wordlen == 6 && pg_strncasecmp(query, "values", 6) == 0)
2147 return true;
2148
2149 return false;
2150 }
2151
2152
2153 /*
2154 * Test if the current user is a database superuser.
2155 *
2156 * Note: this will correctly detect superuserness only with a protocol-3.0
2157 * or newer backend; otherwise it will always say "false".
2158 */
2159 bool
is_superuser(void)2160 is_superuser(void)
2161 {
2162 const char *val;
2163
2164 if (!pset.db)
2165 return false;
2166
2167 val = PQparameterStatus(pset.db, "is_superuser");
2168
2169 if (val && strcmp(val, "on") == 0)
2170 return true;
2171
2172 return false;
2173 }
2174
2175
2176 /*
2177 * Test if the current session uses standard string literals.
2178 *
2179 * Note: With a pre-protocol-3.0 connection this will always say "false",
2180 * which should be the right answer.
2181 */
2182 bool
standard_strings(void)2183 standard_strings(void)
2184 {
2185 const char *val;
2186
2187 if (!pset.db)
2188 return false;
2189
2190 val = PQparameterStatus(pset.db, "standard_conforming_strings");
2191
2192 if (val && strcmp(val, "on") == 0)
2193 return true;
2194
2195 return false;
2196 }
2197
2198
2199 /*
2200 * Return the session user of the current connection.
2201 *
2202 * Note: this will correctly detect the session user only with a
2203 * protocol-3.0 or newer backend; otherwise it will return the
2204 * connection user.
2205 */
2206 const char *
session_username(void)2207 session_username(void)
2208 {
2209 const char *val;
2210
2211 if (!pset.db)
2212 return NULL;
2213
2214 val = PQparameterStatus(pset.db, "session_authorization");
2215 if (val)
2216 return val;
2217 else
2218 return PQuser(pset.db);
2219 }
2220
2221
2222 /* expand_tilde
2223 *
2224 * substitute '~' with HOME or '~username' with username's home dir
2225 *
2226 */
2227 void
expand_tilde(char ** filename)2228 expand_tilde(char **filename)
2229 {
2230 if (!filename || !(*filename))
2231 return;
2232
2233 /*
2234 * WIN32 doesn't use tilde expansion for file names. Also, it uses tilde
2235 * for short versions of long file names, though the tilde is usually
2236 * toward the end, not at the beginning.
2237 */
2238 #ifndef WIN32
2239
2240 /* try tilde expansion */
2241 if (**filename == '~')
2242 {
2243 char *fn;
2244 char oldp,
2245 *p;
2246 struct passwd *pw;
2247 char home[MAXPGPATH];
2248
2249 fn = *filename;
2250 *home = '\0';
2251
2252 p = fn + 1;
2253 while (*p != '/' && *p != '\0')
2254 p++;
2255
2256 oldp = *p;
2257 *p = '\0';
2258
2259 if (*(fn + 1) == '\0')
2260 get_home_path(home); /* ~ or ~/ only */
2261 else if ((pw = getpwnam(fn + 1)) != NULL)
2262 strlcpy(home, pw->pw_dir, sizeof(home)); /* ~user */
2263
2264 *p = oldp;
2265 if (strlen(home) != 0)
2266 {
2267 char *newfn;
2268
2269 newfn = psprintf("%s%s", home, p);
2270 free(fn);
2271 *filename = newfn;
2272 }
2273 }
2274 #endif
2275 }
2276
2277 /*
2278 * Checks if connection string starts with either of the valid URI prefix
2279 * designators.
2280 *
2281 * Returns the URI prefix length, 0 if the string doesn't contain a URI prefix.
2282 *
2283 * XXX This is a duplicate of the eponymous libpq function.
2284 */
2285 static int
uri_prefix_length(const char * connstr)2286 uri_prefix_length(const char *connstr)
2287 {
2288 /* The connection URI must start with either of the following designators: */
2289 static const char uri_designator[] = "postgresql://";
2290 static const char short_uri_designator[] = "postgres://";
2291
2292 if (strncmp(connstr, uri_designator,
2293 sizeof(uri_designator) - 1) == 0)
2294 return sizeof(uri_designator) - 1;
2295
2296 if (strncmp(connstr, short_uri_designator,
2297 sizeof(short_uri_designator) - 1) == 0)
2298 return sizeof(short_uri_designator) - 1;
2299
2300 return 0;
2301 }
2302
2303 /*
2304 * Recognized connection string either starts with a valid URI prefix or
2305 * contains a "=" in it.
2306 *
2307 * Must be consistent with parse_connection_string: anything for which this
2308 * returns true should at least look like it's parseable by that routine.
2309 *
2310 * XXX This is a duplicate of the eponymous libpq function.
2311 */
2312 bool
recognized_connection_string(const char * connstr)2313 recognized_connection_string(const char *connstr)
2314 {
2315 return uri_prefix_length(connstr) != 0 || strchr(connstr, '=') != NULL;
2316 }
2317