[PATCH] Add support to psql for edit-and-execute-command

From a314fa15f6bdf5329d3045d736e02b6835107591 Mon Sep 17 00:00:00 2001
From: Joe Wildish <joe@lateraljoin.com>
Date: Sun, 17 May 2020 21:57:10 +0100
Subject: [PATCH] Add support to psql for edit-and-execute-command

Bash has an edit-and-execute-command Readline function that brings the
current line into an editor to be amended rather than having to edit it
in-place at the prompt. This commit adds the same functionality to psql.
We use the same default bindings ("v" in vi mode, "C-x C-e" in emacs
mode) as bash.
---
 doc/src/sgml/ref/psql-ref.sgml |  38 ++++--
 src/bin/psql/command.c         | 234 +-------------------------------
 src/bin/psql/common.c          | 236 +++++++++++++++++++++++++++++++++
 src/bin/psql/common.h          |   3 +
 src/bin/psql/tab-complete.c    |  67 ++++++++++
 5 files changed, 337 insertions(+), 241 deletions(-)

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index 07bf272a20..0cbc6809d4 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -4458,12 +4458,14 @@ testdb=&gt; \set PROMPT1 '%[%033[1;33;40m%]%n@%/%R%[%033[0m%]%# '
     <application>psql</application> supports the <application>Readline</application>
     library for convenient line editing and retrieval. The command
     history is automatically saved when <application>psql</application>
-    exits and is reloaded when
-    <application>psql</application> starts up. Tab-completion is also
-    supported, although the completion logic makes no claim to be an
-    <acronym>SQL</acronym> parser.  The queries generated by tab-completion
-    can also interfere with other SQL commands, e.g. <literal>SET
-    TRANSACTION ISOLATION LEVEL</literal>.
+    exits and is reloaded when <application>psql</application> starts up.
+    </para>
+
+    <para>
+    Tab-completion is also supported, although the completion logic makes no
+    claim to be an <acronym>SQL</acronym> parser.  The queries generated by
+    tab-completion can also interfere with other SQL commands, e.g.
+    <literal>SET TRANSACTION ISOLATION LEVEL</literal>.
     If for some reason you do not like the tab completion, you
     can turn it off by putting this in a file named
     <filename>.inputrc</filename> in your home directory:
@@ -4472,9 +4474,27 @@ $if psql
 set disable-completion on
 $endif
 </programlisting>
-    (This is not a <application>psql</application> but a
-    <application>Readline</application> feature. Read its documentation
-    for further details.)
+    </para>
+
+    <para>
+    In addition to the usual line editing and retrieval commands,
+    <application>psql</application> also supports invoking an editor on the current
+    line to amend the query text, rather than having to edit it in-place at the
+    prompt. This is known as "edit-and-execute-command". The default Readline
+    binding for this is "v" when in vi mode and "C-x C-e" when in emacs mode. However,
+    the binding can be configured in the usual way from the
+    <filename>.inputrc</filename> file:
+<programlisting>
+$if psql
+"Z": edit-and-execute-command
+$endif
+</programlisting>
+    </para>
+
+    <para>
+    (The <filename>.inputrc</filename> configuration file is not a
+    <application>psql</application> but a <application>Readline</application>
+    feature. Read its documentation for further details.)
     </para>
    </refsect3>
   </refsect2>
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index a5160f91de..ecd2605ca9 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -148,8 +148,6 @@ static void discard_query_text(PsqlScanState scan_state, ConditionalStack cstack
 static void copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf);
 static bool do_connect(enum trivalue reuse_previous_specification,
 					   char *dbname, char *user, char *host, char *port);
-static bool do_edit(const char *filename_arg, PQExpBuffer query_buf,
-					int lineno, bool *edited);
 static bool do_shell(const char *command);
 static bool do_watch(PQExpBuffer query_buf, double sleep);
 static bool lookup_object_oid(EditableObjectType obj_type, const char *desc,
@@ -1008,7 +1006,7 @@ exec_command_edit(PsqlScanState scan_state, bool active_branch,
 				/* If query_buf is empty, recall previous query for editing */
 				copy_previous_query(query_buf, previous_buf);
 
-				if (do_edit(fname, query_buf, lineno, NULL))
+				if (do_edit(fname, query_buf, lineno, 0, NULL))
 					status = PSQL_CMD_NEWEDIT;
 				else
 					status = PSQL_CMD_ERROR;
@@ -1131,7 +1129,7 @@ exec_command_ef_ev(PsqlScanState scan_state, bool active_branch,
 		{
 			bool		edited = false;
 
-			if (!do_edit(NULL, query_buf, lineno, &edited))
+			if (!do_edit(NULL, query_buf, lineno, 0, &edited))
 				status = PSQL_CMD_ERROR;
 			else if (!edited)
 				puts(_("No changes"));
@@ -3511,234 +3509,6 @@ UnsyncVariables(void)
 }
 
 
-/*
- * do_edit -- handler for \e
- *
- * If you do not specify a filename, the current query buffer will be copied
- * into a temporary one.
- */
-static bool
-editFile(const char *fname, int lineno)
-{
-	const char *editorName;
-	const char *editor_lineno_arg = NULL;
-	char	   *sys;
-	int			result;
-
-	Assert(fname != NULL);
-
-	/* Find an editor to use */
-	editorName = getenv("PSQL_EDITOR");
-	if (!editorName)
-		editorName = getenv("EDITOR");
-	if (!editorName)
-		editorName = getenv("VISUAL");
-	if (!editorName)
-		editorName = DEFAULT_EDITOR;
-
-	/* Get line number argument, if we need it. */
-	if (lineno > 0)
-	{
-		editor_lineno_arg = getenv("PSQL_EDITOR_LINENUMBER_ARG");
-#ifdef DEFAULT_EDITOR_LINENUMBER_ARG
-		if (!editor_lineno_arg)
-			editor_lineno_arg = DEFAULT_EDITOR_LINENUMBER_ARG;
-#endif
-		if (!editor_lineno_arg)
-		{
-			pg_log_error("environment variable PSQL_EDITOR_LINENUMBER_ARG must be set to specify a line number");
-			return false;
-		}
-	}
-
-	/*
-	 * On Unix the EDITOR value should *not* be quoted, since it might include
-	 * switches, eg, EDITOR="pico -t"; it's up to the user to put quotes in it
-	 * if necessary.  But this policy is not very workable on Windows, due to
-	 * severe brain damage in their command shell plus the fact that standard
-	 * program paths include spaces.
-	 */
-#ifndef WIN32
-	if (lineno > 0)
-		sys = psprintf("exec %s %s%d '%s'",
-					   editorName, editor_lineno_arg, lineno, fname);
-	else
-		sys = psprintf("exec %s '%s'",
-					   editorName, fname);
-#else
-	if (lineno > 0)
-		sys = psprintf("\"%s\" %s%d \"%s\"",
-					   editorName, editor_lineno_arg, lineno, fname);
-	else
-		sys = psprintf("\"%s\" \"%s\"",
-					   editorName, fname);
-#endif
-	result = system(sys);
-	if (result == -1)
-		pg_log_error("could not start editor \"%s\"", editorName);
-	else if (result == 127)
-		pg_log_error("could not start /bin/sh");
-	free(sys);
-
-	return result == 0;
-}
-
-
-/* call this one */
-static bool
-do_edit(const char *filename_arg, PQExpBuffer query_buf,
-		int lineno, bool *edited)
-{
-	char		fnametmp[MAXPGPATH];
-	FILE	   *stream = NULL;
-	const char *fname;
-	bool		error = false;
-	int			fd;
-
-	struct stat before,
-				after;
-
-	if (filename_arg)
-		fname = filename_arg;
-	else
-	{
-		/* make a temp file to edit */
-#ifndef WIN32
-		const char *tmpdir = getenv("TMPDIR");
-
-		if (!tmpdir)
-			tmpdir = "/tmp";
-#else
-		char		tmpdir[MAXPGPATH];
-		int			ret;
-
-		ret = GetTempPath(MAXPGPATH, tmpdir);
-		if (ret == 0 || ret > MAXPGPATH)
-		{
-			pg_log_error("could not locate temporary directory: %s",
-						 !ret ? strerror(errno) : "");
-			return false;
-		}
-
-		/*
-		 * No canonicalize_path() here. EDIT.EXE run from CMD.EXE prepends the
-		 * current directory to the supplied path unless we use only
-		 * backslashes, so we do that.
-		 */
-#endif
-#ifndef WIN32
-		snprintf(fnametmp, sizeof(fnametmp), "%s%spsql.edit.%d.sql", tmpdir,
-				 "/", (int) getpid());
-#else
-		snprintf(fnametmp, sizeof(fnametmp), "%s%spsql.edit.%d.sql", tmpdir,
-				 "" /* trailing separator already present */ , (int) getpid());
-#endif
-
-		fname = (const char *) fnametmp;
-
-		fd = open(fname, O_WRONLY | O_CREAT | O_EXCL, 0600);
-		if (fd != -1)
-			stream = fdopen(fd, "w");
-
-		if (fd == -1 || !stream)
-		{
-			pg_log_error("could not open temporary file \"%s\": %m", fname);
-			error = true;
-		}
-		else
-		{
-			unsigned int ql = query_buf->len;
-
-			/* force newline-termination of what we send to editor */
-			if (ql > 0 && query_buf->data[ql - 1] != '\n')
-			{
-				appendPQExpBufferChar(query_buf, '\n');
-				ql++;
-			}
-
-			if (fwrite(query_buf->data, 1, ql, stream) != ql)
-			{
-				pg_log_error("%s: %m", fname);
-
-				if (fclose(stream) != 0)
-					pg_log_error("%s: %m", fname);
-
-				if (remove(fname) != 0)
-					pg_log_error("%s: %m", fname);
-
-				error = true;
-			}
-			else if (fclose(stream) != 0)
-			{
-				pg_log_error("%s: %m", fname);
-				if (remove(fname) != 0)
-					pg_log_error("%s: %m", fname);
-				error = true;
-			}
-		}
-	}
-
-	if (!error && stat(fname, &before) != 0)
-	{
-		pg_log_error("%s: %m", fname);
-		error = true;
-	}
-
-	/* call editor */
-	if (!error)
-		error = !editFile(fname, lineno);
-
-	if (!error && stat(fname, &after) != 0)
-	{
-		pg_log_error("%s: %m", fname);
-		error = true;
-	}
-
-	if (!error && before.st_mtime != after.st_mtime)
-	{
-		stream = fopen(fname, PG_BINARY_R);
-		if (!stream)
-		{
-			pg_log_error("%s: %m", fname);
-			error = true;
-		}
-		else
-		{
-			/* read file back into query_buf */
-			char		line[1024];
-
-			resetPQExpBuffer(query_buf);
-			while (fgets(line, sizeof(line), stream) != NULL)
-				appendPQExpBufferStr(query_buf, line);
-
-			if (ferror(stream))
-			{
-				pg_log_error("%s: %m", fname);
-				error = true;
-			}
-			else if (edited)
-			{
-				*edited = true;
-			}
-
-			fclose(stream);
-		}
-	}
-
-	/* remove temp file */
-	if (!filename_arg)
-	{
-		if (remove(fname) == -1)
-		{
-			pg_log_error("%s: %m", fname);
-			error = true;
-		}
-	}
-
-	return !error;
-}
-
-
 
 /*
  * process_file
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 06f801764b..4f87f71ed4 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -12,10 +12,15 @@
 #include <math.h>
 #include <signal.h>
 #ifndef WIN32
+#include <sys/stat.h>			/* for stat() */
+#include <fcntl.h>				/* open() flags */
 #include <unistd.h>				/* for write() */
 #else
 #include <io.h>					/* for _write() */
 #include <win32.h>
+#include <fcntl.h>
+#include <direct.h>
+#include <sys/stat.h>			/* for stat() */
 #endif
 
 #include "command.h"
@@ -33,6 +38,7 @@ static bool DescribeQuery(const char *query, double *elapsed_msec);
 static bool ExecQueryUsingCursor(const char *query, double *elapsed_msec);
 static bool command_no_begin(const char *query);
 static bool is_select_command(const char *query);
+static bool editFile(const char *fname, int lineno, int colno);
 
 
 /*
@@ -2304,3 +2310,233 @@ recognized_connection_string(const char *connstr)
 {
 	return uri_prefix_length(connstr) != 0 || strchr(connstr, '=') != NULL;
 }
+
+
+bool
+do_edit(const char *filename_arg, PQExpBuffer query_buf,
+		int lineno, int colno, bool *edited)
+{
+	char		fnametmp[MAXPGPATH];
+	FILE	   *stream = NULL;
+	const char *fname;
+	bool		error = false;
+	int			fd;
+
+	struct stat before,
+		after;
+
+	if (filename_arg)
+		fname = filename_arg;
+	else
+	{
+		/* make a temp file to edit */
+#ifndef WIN32
+		const char *tmpdir = getenv("TMPDIR");
+
+		if (!tmpdir)
+			tmpdir = "/tmp";
+#else
+			char		tmpdir[MAXPGPATH];
+		int			ret;
+
+		ret = GetTempPath(MAXPGPATH, tmpdir);
+		if (ret == 0 || ret > MAXPGPATH)
+		{
+			pg_log_error("could not locate temporary directory: %s",
+						 !ret ? strerror(errno) : "");
+			return false;
+		}
+
+		/*
+		 * No canonicalize_path() here. EDIT.EXE run from CMD.EXE prepends the
+		 * current directory to the supplied path unless we use only
+		 * backslashes, so we do that.
+		 */
+#endif
+#ifndef WIN32
+		snprintf(fnametmp, sizeof(fnametmp), "%s%spsql.edit.%d.sql", tmpdir,
+				 "/", (int) getpid());
+#else
+		snprintf(fnametmp, sizeof(fnametmp), "%s%spsql.edit.%d.sql", tmpdir,
+				 "" /* trailing separator already present */ , (int) getpid());
+#endif
+
+		fname = (const char *) fnametmp;
+
+		fd = open(fname, O_WRONLY | O_CREAT | O_EXCL, 0600);
+		if (fd != -1)
+			stream = fdopen(fd, "w");
+
+		if (fd == -1 || !stream)
+		{
+			pg_log_error("could not open temporary file \"%s\": %m", fname);
+			error = true;
+		}
+		else
+		{
+			unsigned int ql = query_buf->len;
+
+			/* force newline-termination of what we send to editor */
+			if (ql > 0 && query_buf->data[ql - 1] != '\n')
+			{
+				appendPQExpBufferChar(query_buf, '\n');
+				ql++;
+			}
+
+			if (fwrite(query_buf->data, 1, ql, stream) != ql)
+			{
+				pg_log_error("%s: %m", fname);
+
+				if (fclose(stream) != 0)
+					pg_log_error("%s: %m", fname);
+
+				if (remove(fname) != 0)
+					pg_log_error("%s: %m", fname);
+
+				error = true;
+			}
+			else if (fclose(stream) != 0)
+			{
+				pg_log_error("%s: %m", fname);
+				if (remove(fname) != 0)
+					pg_log_error("%s: %m", fname);
+				error = true;
+			}
+		}
+	}
+
+	if (!error && stat(fname, &before) != 0)
+	{
+		pg_log_error("%s: %m", fname);
+		error = true;
+	}
+
+	/* call editor */
+	if (!error)
+		error = !editFile(fname, lineno, colno);
+
+	if (!error && stat(fname, &after) != 0)
+	{
+		pg_log_error("%s: %m", fname);
+		error = true;
+	}
+
+	if (!error && before.st_mtime != after.st_mtime)
+	{
+		stream = fopen(fname, PG_BINARY_R);
+		if (!stream)
+		{
+			pg_log_error("%s: %m", fname);
+			error = true;
+		}
+		else
+		{
+			/* read file back into query_buf */
+			char		line[1024];
+
+			resetPQExpBuffer(query_buf);
+			while (fgets(line, sizeof(line), stream) != NULL)
+				appendPQExpBufferStr(query_buf, line);
+
+			if (ferror(stream))
+			{
+				pg_log_error("%s: %m", fname);
+				error = true;
+			}
+			else if (edited)
+			{
+				*edited = true;
+			}
+
+			fclose(stream);
+		}
+	}
+
+	/* remove temp file */
+	if (!filename_arg)
+	{
+		if (remove(fname) == -1)
+		{
+			pg_log_error("%s: %m", fname);
+			error = true;
+		}
+	}
+
+	return !error;
+}
+
+
+/*
+ * do_edit -- handler for \e
+ *
+ * If you do not specify a filename, the current query buffer will be copied
+ * into a temporary one.
+ */
+static bool
+editFile(const char *fname, int lineno, int colno)
+{
+	const char *editorName;
+	const char *editor_lineno_arg = NULL;
+	char	   *sys;
+	int			result;
+
+	Assert(fname != NULL);
+
+	/* Find an editor to use */
+	editorName = getenv("PSQL_EDITOR");
+	if (!editorName)
+		editorName = getenv("EDITOR");
+	if (!editorName)
+		editorName = getenv("VISUAL");
+	if (!editorName)
+		editorName = DEFAULT_EDITOR;
+
+	/* Get line number argument, if we need it. */
+	if (lineno > 0)
+	{
+		editor_lineno_arg = getenv("PSQL_EDITOR_LINENUMBER_ARG");
+#ifdef DEFAULT_EDITOR_LINENUMBER_ARG
+		if (!editor_lineno_arg)
+			editor_lineno_arg = DEFAULT_EDITOR_LINENUMBER_ARG;
+#endif
+		if (!editor_lineno_arg)
+		{
+			pg_log_error("environment variable PSQL_EDITOR_LINENUMBER_ARG must be set to specify a line number");
+			return false;
+		}
+	}
+
+	/*
+	 * On Unix the EDITOR value should *not* be quoted, since it might include
+	 * switches, eg, EDITOR="pico -t"; it's up to the user to put quotes in it
+	 * if necessary.  But this policy is not very workable on Windows, due to
+	 * severe brain damage in their command shell plus the fact that standard
+	 * program paths include spaces.
+	 */
+#ifndef WIN32
+	if (lineno > 0 && strcmp(editorName, DEFAULT_EDITOR) == 0)
+		sys = psprintf("exec %s -c'call cursor(%d,%d)' '%s'",
+					   editorName, lineno, colno, fname);
+	else if (lineno > 0)
+		sys = psprintf("exec %s %s%d '%s'",
+					   editorName, editor_lineno_arg, lineno, fname);
+	else
+		sys = psprintf("exec %s '%s'",
+					   editorName, fname);
+#else
+	if (lineno > 0)
+		sys = psprintf("\"%s\" %s%d \"%s\"",
+					   editorName, editor_lineno_arg, lineno, fname);
+	else
+		sys = psprintf("\"%s\" \"%s\"",
+					   editorName, fname);
+#endif
+	result = system(sys);
+	if (result == -1)
+		pg_log_error("could not start editor \"%s\"", editorName);
+	else if (result == 127)
+		pg_log_error("could not start /bin/sh");
+	free(sys);
+
+	return result == 0;
+}
\ No newline at end of file
diff --git a/src/bin/psql/common.h b/src/bin/psql/common.h
index ec4e83c9fd..76767732a4 100644
--- a/src/bin/psql/common.h
+++ b/src/bin/psql/common.h
@@ -41,4 +41,7 @@ extern void expand_tilde(char **filename);
 
 extern bool recognized_connection_string(const char *connstr);
 
+extern bool do_edit(const char *filename_arg, PQExpBuffer query_buf,
+					int lineno, int colno, bool *edited);
+
 #endif							/* COMMON_H */
diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c
index eb018854a5..bbdc3bce4f 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -76,6 +76,12 @@
 /* word break characters */
 #define WORD_BREAKS		"\t\n@$><=;|&{() "
 
+#ifdef HAVE_READLINE_READLINE_H
+/* Helpers to define emacs bindings for Control-X */
+#define READLINE_CONTROL_CHARACTER_MASK 0x1f /* 0x20 - 1 */
+#define READLINE_CTRL(c) ((c) & READLINE_CONTROL_CHARACTER_MASK)
+#endif
+
 /*
  * Since readline doesn't let us pass any state through to the tab completion
  * callback, we have to use this global variable to let get_previous_words()
@@ -1138,6 +1144,10 @@ static char **get_previous_words(int point, char **buffer, int *nwords);
 
 static char *get_guctype(const char *varname);
 
+#ifdef HAVE_READLINE_READLINE_H
+static int edit_and_execute_command(int count, int c);
+#endif
+
 #ifdef USE_FILENAME_QUOTING_FUNCTIONS
 static char *quote_file_name(char *fname, int match_type, char *quote_pointer);
 static char *dequote_file_name(char *fname, int quote_char);
@@ -1160,6 +1170,17 @@ initialize_readline(void)
 
 	rl_basic_word_break_characters = WORD_BREAKS;
 
+	/*
+	 * Bind 'v' or 'C-xC-e' to invoke vi or emacs and run result as commands.
+	 */
+#ifdef HAVE_READLINE_READLINE_H
+	{
+		rl_add_defun("edit-and-execute-command", edit_and_execute_command, -1);
+		rl_bind_key_if_unbound_in_map(READLINE_CTRL('E'), edit_and_execute_command, emacs_ctlx_keymap);
+		rl_bind_key_if_unbound_in_map('v', edit_and_execute_command, vi_movement_keymap);
+	}
+#endif
+
 	/*
 	 * We should include '"' in rl_completer_quote_characters too, but that
 	 * will require some upgrades to how we handle quoted identifiers, so
@@ -4814,6 +4835,52 @@ get_guctype(const char *varname)
 	return guctype;
 }
 
+
+/*
+ * Edit the current line in vi (or $EDITOR, $VISUAL).
+ */
+#ifdef HAVE_READLINE_READLINE_H
+static int
+edit_and_execute_command(int count, int c)
+{
+	PQExpBuffer buffer = createPQExpBuffer();
+	bool		edited = false;
+	bool 		result;
+	int			point, lineno, colno;
+
+	appendPQExpBufferStr(buffer, rl_line_buffer);
+
+	/* Discover which line and column we are at. */
+	lineno = 1;
+	colno = 0;
+	for(point = rl_point; point >= 1; point--)
+	{
+		if (rl_line_buffer[point] == '\n')
+			lineno++;
+		if (lineno == 1)
+			colno++;
+	}
+
+	result = do_edit(NULL, buffer, lineno, colno, &edited);
+
+	if (result && edited)
+	{
+		rl_clear_visible_line();
+		rl_extend_line_buffer(buffer->len);
+		rl_replace_line(buffer->data, 1);
+		rl_reset_line_state();
+		rl_redisplay();
+
+		/* accept the line */
+		rl_newline(1, c);
+	}
+
+	destroyPQExpBuffer(buffer);
+
+	return result ? 0 : 1;
+}
+#endif
+
 #ifdef USE_FILENAME_QUOTING_FUNCTIONS
 
 /*
-- 
2.26.2