Example #1
0
char *
get_prompt(promptStatus_t status, ConditionalStack cstack)
{
#define MAX_PROMPT_SIZE 256
	static char destination[MAX_PROMPT_SIZE + 1];
	char		buf[MAX_PROMPT_SIZE + 1];
	bool		esc = false;
	const char *p;
	const char *prompt_string = "? ";

	switch (status)
	{
		case PROMPT_READY:
			prompt_string = pset.prompt1;
			break;

		case PROMPT_CONTINUE:
		case PROMPT_SINGLEQUOTE:
		case PROMPT_DOUBLEQUOTE:
		case PROMPT_DOLLARQUOTE:
		case PROMPT_COMMENT:
		case PROMPT_PAREN:
			prompt_string = pset.prompt2;
			break;

		case PROMPT_COPY:
			prompt_string = pset.prompt3;
			break;
	}

	destination[0] = '\0';

	for (p = prompt_string;
		 *p && strlen(destination) < sizeof(destination) - 1;
		 p++)
	{
		memset(buf, 0, sizeof(buf));
		if (esc)
		{
			switch (*p)
			{
					/* Current database */
				case '/':
					if (pset.db)
						strlcpy(buf, PQdb(pset.db), sizeof(buf));
					break;
				case '~':
					if (pset.db)
					{
						const char *var;

						if (strcmp(PQdb(pset.db), PQuser(pset.db)) == 0 ||
							((var = getenv("PGDATABASE")) && strcmp(var, PQdb(pset.db)) == 0))
							strlcpy(buf, "~", sizeof(buf));
						else
							strlcpy(buf, PQdb(pset.db), sizeof(buf));
					}
					break;

					/* DB server hostname (long/short) */
				case 'M':
				case 'm':
					if (pset.db)
					{
						const char *host = PQhost(pset.db);

						/* INET socket */
						if (host && host[0] && !is_absolute_path(host))
						{
							strlcpy(buf, host, sizeof(buf));
							if (*p == 'm')
								buf[strcspn(buf, ".")] = '\0';
						}
#ifdef HAVE_UNIX_SOCKETS
						/* UNIX socket */
						else
						{
							if (!host
								|| strcmp(host, DEFAULT_PGSOCKET_DIR) == 0
								|| *p == 'm')
								strlcpy(buf, "[local]", sizeof(buf));
							else
								snprintf(buf, sizeof(buf), "[local:%s]", host);
						}
#endif
					}
					break;
					/* DB server port number */
				case '>':
					if (pset.db && PQport(pset.db))
						strlcpy(buf, PQport(pset.db), sizeof(buf));
					break;
					/* DB server user name */
				case 'n':
					if (pset.db)
						strlcpy(buf, session_username(), sizeof(buf));
					break;
					/* backend pid */
				case 'p':
					if (pset.db)
					{
						int			pid = PQbackendPID(pset.db);

						if (pid)
							snprintf(buf, sizeof(buf), "%d", pid);
					}
					break;

				case '0':
				case '1':
				case '2':
				case '3':
				case '4':
				case '5':
				case '6':
				case '7':
					*buf = (char) strtol(p, (char **) &p, 8);
					--p;
					break;
				case 'R':
					switch (status)
					{
						case PROMPT_READY:
							if (cstack != NULL && !conditional_active(cstack))
								buf[0] = '@';
							else if (!pset.db)
								buf[0] = '!';
							else if (!pset.singleline)
								buf[0] = '=';
							else
								buf[0] = '^';
							break;
						case PROMPT_CONTINUE:
							buf[0] = '-';
							break;
						case PROMPT_SINGLEQUOTE:
							buf[0] = '\'';
							break;
						case PROMPT_DOUBLEQUOTE:
							buf[0] = '"';
							break;
						case PROMPT_DOLLARQUOTE:
							buf[0] = '$';
							break;
						case PROMPT_COMMENT:
							buf[0] = '*';
							break;
						case PROMPT_PAREN:
							buf[0] = '(';
							break;
						default:
							buf[0] = '\0';
							break;
					}
					break;

				case 'x':
					if (!pset.db)
						buf[0] = '?';
					else
						switch (PQtransactionStatus(pset.db))
						{
							case PQTRANS_IDLE:
								buf[0] = '\0';
								break;
							case PQTRANS_ACTIVE:
							case PQTRANS_INTRANS:
								buf[0] = '*';
								break;
							case PQTRANS_INERROR:
								buf[0] = '!';
								break;
							default:
								buf[0] = '?';
								break;
						}
					break;

				case 'l':
					snprintf(buf, sizeof(buf), UINT64_FORMAT, pset.stmt_lineno);
					break;

				case '?':
					/* not here yet */
					break;

				case '#':
					if (is_superuser())
						buf[0] = '#';
					else
						buf[0] = '>';
					break;

					/* execute command */
				case '`':
					{
						FILE	   *fd;
						char	   *file = pg_strdup(p + 1);
						int			cmdend;

						cmdend = strcspn(file, "`");
						file[cmdend] = '\0';
						fd = popen(file, "r");
						if (fd)
						{
							if (fgets(buf, sizeof(buf), fd) == NULL)
								buf[0] = '\0';
							pclose(fd);
						}
						if (strlen(buf) > 0 && buf[strlen(buf) - 1] == '\n')
							buf[strlen(buf) - 1] = '\0';
						free(file);
						p += cmdend + 1;
						break;
					}

					/* interpolate variable */
				case ':':
					{
						char	   *name;
						const char *val;
						int			nameend;

						name = pg_strdup(p + 1);
						nameend = strcspn(name, ":");
						name[nameend] = '\0';
						val = GetVariable(pset.vars, name);
						if (val)
							strlcpy(buf, val, sizeof(buf));
						free(name);
						p += nameend + 1;
						break;
					}

				case '[':
				case ']':
#if defined(USE_READLINE) && defined(RL_PROMPT_START_IGNORE)

					/*
					 * readline >=4.0 undocumented feature: non-printing
					 * characters in prompt strings must be marked as such, in
					 * order to properly display the line during editing.
					 */
					buf[0] = (*p == '[') ? RL_PROMPT_START_IGNORE : RL_PROMPT_END_IGNORE;
					buf[1] = '\0';
#endif							/* USE_READLINE */
					break;

				default:
					buf[0] = *p;
					buf[1] = '\0';
					break;

			}
			esc = false;
		}
		else if (*p == '%')
			esc = true;
		else
		{
			buf[0] = *p;
			buf[1] = '\0';
			esc = false;
		}

		if (!esc)
			strlcat(destination, buf, sizeof(destination));
	}

	return destination;
}
Example #2
0
/*
 * Main processing loop for reading lines of input
 *	and sending them to the backend.
 *
 * This loop is re-entrant. May be called by \i command
 *	which reads input from a file.
 */
int
MainLoop(FILE *source)
{
	PsqlScanState scan_state;	/* lexer working state */
	ConditionalStack cond_stack;	/* \if status stack */
	volatile PQExpBuffer query_buf; /* buffer for query being accumulated */
	volatile PQExpBuffer previous_buf;	/* if there isn't anything in the new
										 * buffer yet, use this one for \e,
										 * etc. */
	PQExpBuffer history_buf;	/* earlier lines of a multi-line command, not
								 * yet saved to readline history */
	char	   *line;			/* current line of input */
	int			added_nl_pos;
	bool		success;
	bool		line_saved_in_history;
	volatile int successResult = EXIT_SUCCESS;
	volatile backslashResult slashCmdStatus = PSQL_CMD_UNKNOWN;
	volatile promptStatus_t prompt_status = PROMPT_READY;
	volatile int count_eof = 0;
	volatile bool die_on_error = false;
	FILE	   *prev_cmd_source;
	bool		prev_cmd_interactive;
	uint64		prev_lineno;

	/* Save the prior command source */
	prev_cmd_source = pset.cur_cmd_source;
	prev_cmd_interactive = pset.cur_cmd_interactive;
	prev_lineno = pset.lineno;
	/* pset.stmt_lineno does not need to be saved and restored */

	/* Establish new source */
	pset.cur_cmd_source = source;
	pset.cur_cmd_interactive = ((source == stdin) && !pset.notty);
	pset.lineno = 0;
	pset.stmt_lineno = 1;

	/* Create working state */
	scan_state = psql_scan_create(&psqlscan_callbacks);
	cond_stack = conditional_stack_create();
	psql_scan_set_passthrough(scan_state, (void *) cond_stack);

	query_buf = createPQExpBuffer();
	previous_buf = createPQExpBuffer();
	history_buf = createPQExpBuffer();
	if (PQExpBufferBroken(query_buf) ||
		PQExpBufferBroken(previous_buf) ||
		PQExpBufferBroken(history_buf))
	{
		psql_error("out of memory\n");
		exit(EXIT_FAILURE);
	}

	/* main loop to get queries and execute them */
	while (successResult == EXIT_SUCCESS)
	{
		/*
		 * Clean up after a previous Control-C
		 */
		if (cancel_pressed)
		{
			if (!pset.cur_cmd_interactive)
			{
				/*
				 * You get here if you stopped a script with Ctrl-C.
				 */
				successResult = EXIT_USER;
				break;
			}

			cancel_pressed = false;
		}

		/*
		 * Establish longjmp destination for exiting from wait-for-input. We
		 * must re-do this each time through the loop for safety, since the
		 * jmpbuf might get changed during command execution.
		 */
		if (sigsetjmp(sigint_interrupt_jmp, 1) != 0)
		{
			/* got here with longjmp */

			/* reset parsing state */
			psql_scan_finish(scan_state);
			psql_scan_reset(scan_state);
			resetPQExpBuffer(query_buf);
			resetPQExpBuffer(history_buf);
			count_eof = 0;
			slashCmdStatus = PSQL_CMD_UNKNOWN;
			prompt_status = PROMPT_READY;
			pset.stmt_lineno = 1;
			cancel_pressed = false;

			if (pset.cur_cmd_interactive)
			{
				putc('\n', stdout);

				/*
				 * if interactive user is in an \if block, then Ctrl-C will
				 * exit from the innermost \if.
				 */
				if (!conditional_stack_empty(cond_stack))
				{
					psql_error("\\if: escaped\n");
					conditional_stack_pop(cond_stack);
				}
			}
			else
			{
				successResult = EXIT_USER;
				break;
			}
		}

		fflush(stdout);

		/*
		 * get another line
		 */
		if (pset.cur_cmd_interactive)
		{
			/* May need to reset prompt, eg after \r command */
			if (query_buf->len == 0)
				prompt_status = PROMPT_READY;
			line = gets_interactive(get_prompt(prompt_status, cond_stack),
									query_buf);
		}
		else
		{
			line = gets_fromFile(source);
			if (!line && ferror(source))
				successResult = EXIT_FAILURE;
		}

		/*
		 * query_buf holds query already accumulated.  line is the malloc'd
		 * new line of input (note it must be freed before looping around!)
		 */

		/* No more input.  Time to quit, or \i done */
		if (line == NULL)
		{
			if (pset.cur_cmd_interactive)
			{
				/* This tries to mimic bash's IGNOREEOF feature. */
				count_eof++;

				if (count_eof < pset.ignoreeof)
				{
					if (!pset.quiet)
						printf(_("Use \"\\q\" to leave %s.\n"), pset.progname);
					continue;
				}

				puts(pset.quiet ? "" : "\\q");
			}
			break;
		}

		count_eof = 0;

		pset.lineno++;

		/* ignore UTF-8 Unicode byte-order mark */
		if (pset.lineno == 1 && pset.encoding == PG_UTF8 && strncmp(line, "\xef\xbb\xbf", 3) == 0)
			memmove(line, line + 3, strlen(line + 3) + 1);

		/* Detect attempts to run custom-format dumps as SQL scripts */
		if (pset.lineno == 1 && !pset.cur_cmd_interactive &&
			strncmp(line, "PGDMP", 5) == 0)
		{
			free(line);
			puts(_("The input is a PostgreSQL custom-format dump.\n"
				   "Use the pg_restore command-line client to restore this dump to a database.\n"));
			fflush(stdout);
			successResult = EXIT_FAILURE;
			break;
		}

		/* no further processing of empty lines, unless within a literal */
		if (line[0] == '\0' && !psql_scan_in_quote(scan_state))
		{
			free(line);
			continue;
		}

		/* A request for help? Be friendly and give them some guidance */
		if (pset.cur_cmd_interactive && query_buf->len == 0 &&
			pg_strncasecmp(line, "help", 4) == 0 &&
			(line[4] == '\0' || line[4] == ';' || isspace((unsigned char) line[4])))
		{
			free(line);
			puts(_("You are using psql, the command-line interface to PostgreSQL."));
			printf(_("Type:  \\copyright for distribution terms\n"
					 "       \\h for help with SQL commands\n"
					 "       \\? for help with psql commands\n"
					 "       \\g or terminate with semicolon to execute query\n"
					 "       \\q to quit\n"));

			fflush(stdout);
			continue;
		}

		/* echo back if flag is set, unless interactive */
		if (pset.echo == PSQL_ECHO_ALL && !pset.cur_cmd_interactive)
		{
			puts(line);
			fflush(stdout);
		}

		/* insert newlines into query buffer between source lines */
		if (query_buf->len > 0)
		{
			appendPQExpBufferChar(query_buf, '\n');
			added_nl_pos = query_buf->len;
		}
		else
			added_nl_pos = -1;	/* flag we didn't add one */

		/* Setting this will not have effect until next line. */
		die_on_error = pset.on_error_stop;

		/*
		 * Parse line, looking for command separators.
		 */
		psql_scan_setup(scan_state, line, strlen(line),
						pset.encoding, standard_strings());
		success = true;
		line_saved_in_history = false;

		while (success || !die_on_error)
		{
			PsqlScanResult scan_result;
			promptStatus_t prompt_tmp = prompt_status;
			size_t		pos_in_query;
			char	   *tmp_line;

			pos_in_query = query_buf->len;
			scan_result = psql_scan(scan_state, query_buf, &prompt_tmp);
			prompt_status = prompt_tmp;

			if (PQExpBufferBroken(query_buf))
			{
				psql_error("out of memory\n");
				exit(EXIT_FAILURE);
			}

			/*
			 * Increase statement line number counter for each linebreak added
			 * to the query buffer by the last psql_scan() call. There only
			 * will be ones to add when navigating to a statement in
			 * readline's history containing newlines.
			 */
			tmp_line = query_buf->data + pos_in_query;
			while (*tmp_line != '\0')
			{
				if (*(tmp_line++) == '\n')
					pset.stmt_lineno++;
			}

			if (scan_result == PSCAN_EOL)
				pset.stmt_lineno++;

			/*
			 * Send command if semicolon found, or if end of line and we're in
			 * single-line mode.
			 */
			if (scan_result == PSCAN_SEMICOLON ||
				(scan_result == PSCAN_EOL && pset.singleline))
			{
				/*
				 * Save line in history.  We use history_buf to accumulate
				 * multi-line queries into a single history entry.  Note that
				 * history accumulation works on input lines, so it doesn't
				 * matter whether the query will be ignored due to \if.
				 */
				if (pset.cur_cmd_interactive && !line_saved_in_history)
				{
					pg_append_history(line, history_buf);
					pg_send_history(history_buf);
					line_saved_in_history = true;
				}

				/* execute query unless we're in an inactive \if branch */
				if (conditional_active(cond_stack))
				{
					success = SendQuery(query_buf->data);
					slashCmdStatus = success ? PSQL_CMD_SEND : PSQL_CMD_ERROR;
					pset.stmt_lineno = 1;

					/* transfer query to previous_buf by pointer-swapping */
					{
						PQExpBuffer swap_buf = previous_buf;

						previous_buf = query_buf;
						query_buf = swap_buf;
					}
					resetPQExpBuffer(query_buf);

					added_nl_pos = -1;
					/* we need not do psql_scan_reset() here */
				}
				else
				{
					/* if interactive, warn about non-executed query */
					if (pset.cur_cmd_interactive)
						psql_error("query ignored; use \\endif or Ctrl-C to exit current \\if block\n");
					/* fake an OK result for purposes of loop checks */
					success = true;
					slashCmdStatus = PSQL_CMD_SEND;
					pset.stmt_lineno = 1;
					/* note that query_buf doesn't change state */
				}
			}
			else if (scan_result == PSCAN_BACKSLASH)
			{
				/* handle backslash command */

				/*
				 * If we added a newline to query_buf, and nothing else has
				 * been inserted in query_buf by the lexer, then strip off the
				 * newline again.  This avoids any change to query_buf when a
				 * line contains only a backslash command.  Also, in this
				 * situation we force out any previous lines as a separate
				 * history entry; we don't want SQL and backslash commands
				 * intermixed in history if at all possible.
				 */
				if (query_buf->len == added_nl_pos)
				{
					query_buf->data[--query_buf->len] = '\0';
					pg_send_history(history_buf);
				}
				added_nl_pos = -1;

				/* save backslash command in history */
				if (pset.cur_cmd_interactive && !line_saved_in_history)
				{
					pg_append_history(line, history_buf);
					pg_send_history(history_buf);
					line_saved_in_history = true;
				}

				/* execute backslash command */
				slashCmdStatus = HandleSlashCmds(scan_state,
												 cond_stack,
												 query_buf,
												 previous_buf);

				success = slashCmdStatus != PSQL_CMD_ERROR;

				/*
				 * Resetting stmt_lineno after a backslash command isn't
				 * always appropriate, but it's what we've done historically
				 * and there have been few complaints.
				 */
				pset.stmt_lineno = 1;

				if (slashCmdStatus == PSQL_CMD_SEND)
				{
					/* should not see this in inactive branch */
					Assert(conditional_active(cond_stack));

					success = SendQuery(query_buf->data);

					/* transfer query to previous_buf by pointer-swapping */
					{
						PQExpBuffer swap_buf = previous_buf;

						previous_buf = query_buf;
						query_buf = swap_buf;
					}
					resetPQExpBuffer(query_buf);

					/* flush any paren nesting info after forced send */
					psql_scan_reset(scan_state);
				}
				else if (slashCmdStatus == PSQL_CMD_NEWEDIT)
				{
					/* should not see this in inactive branch */
					Assert(conditional_active(cond_stack));
					/* rescan query_buf as new input */
					psql_scan_finish(scan_state);
					free(line);
					line = pg_strdup(query_buf->data);
					resetPQExpBuffer(query_buf);
					/* reset parsing state since we are rescanning whole line */
					psql_scan_reset(scan_state);
					psql_scan_setup(scan_state, line, strlen(line),
									pset.encoding, standard_strings());
					line_saved_in_history = false;
					prompt_status = PROMPT_READY;
				}
				else if (slashCmdStatus == PSQL_CMD_TERMINATE)
					break;
			}

			/* fall out of loop if lexer reached EOL */
			if (scan_result == PSCAN_INCOMPLETE ||
				scan_result == PSCAN_EOL)
				break;
		}

		/* Add line to pending history if we didn't execute anything yet */
		if (pset.cur_cmd_interactive && !line_saved_in_history)
			pg_append_history(line, history_buf);

		psql_scan_finish(scan_state);
		free(line);

		if (slashCmdStatus == PSQL_CMD_TERMINATE)
		{
			successResult = EXIT_SUCCESS;
			break;
		}

		if (!pset.cur_cmd_interactive)
		{
			if (!success && die_on_error)
				successResult = EXIT_USER;
			/* Have we lost the db connection? */
			else if (!pset.db)
				successResult = EXIT_BADCONN;
		}
	}							/* while !endoffile/session */

	/*
	 * If we have a non-semicolon-terminated query at the end of file, we
	 * process it unless the input source is interactive --- in that case it
	 * seems better to go ahead and quit.  Also skip if this is an error exit.
	 */
	if (query_buf->len > 0 && !pset.cur_cmd_interactive &&
		successResult == EXIT_SUCCESS)
	{
		/* save query in history */
		/* currently unneeded since we don't use this block if interactive */
#ifdef NOT_USED
		if (pset.cur_cmd_interactive)
			pg_send_history(history_buf);
#endif

		/* execute query unless we're in an inactive \if branch */
		if (conditional_active(cond_stack))
		{
			success = SendQuery(query_buf->data);
		}
		else
		{
			if (pset.cur_cmd_interactive)
				psql_error("query ignored; use \\endif or Ctrl-C to exit current \\if block\n");
			success = true;
		}

		if (!success && die_on_error)
			successResult = EXIT_USER;
		else if (pset.db == NULL)
			successResult = EXIT_BADCONN;
	}

	/*
	 * Check for unbalanced \if-\endifs unless user explicitly quit, or the
	 * script is erroring out
	 */
	if (slashCmdStatus != PSQL_CMD_TERMINATE &&
		successResult != EXIT_USER &&
		!conditional_stack_empty(cond_stack))
	{
		psql_error("reached EOF without finding closing \\endif(s)\n");
		if (die_on_error && !pset.cur_cmd_interactive)
			successResult = EXIT_USER;
	}

	/*
	 * Let's just make real sure the SIGINT handler won't try to use
	 * sigint_interrupt_jmp after we exit this routine.  If there is an outer
	 * MainLoop instance, it will reset sigint_interrupt_jmp to point to
	 * itself at the top of its loop, before any further interactive input
	 * happens.
	 */
	sigint_interrupt_enabled = false;

	destroyPQExpBuffer(query_buf);
	destroyPQExpBuffer(previous_buf);
	destroyPQExpBuffer(history_buf);

	psql_scan_destroy(scan_state);
	conditional_stack_destroy(cond_stack);

	pset.cur_cmd_source = prev_cmd_source;
	pset.cur_cmd_interactive = prev_cmd_interactive;
	pset.lineno = prev_lineno;

	return successResult;
}								/* MainLoop() */