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