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