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 											&copystream, &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 										&copy_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 									   &copy_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