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