From 592a2e8a94ef06befe42e50fa42901482f2bf5fc Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Tue, 19 Nov 2019 09:36:06 +1300
Subject: [PATCH v2] Allow invisible PROMPT2 in psql.

Keep track of the visible width of PROMPT1, and provide %w as a way
for PROMPT2 to generate the same number of spaces.

Author: Thomas Munro, using suggestions from others
Reviewed-by: Tom Lane
Discussion: https://postgr.es/m/CA%2BhUKG%2BzGd7RigjWbxwhzGW59gUpf76ydQECeGdEdodH6nd__A%40mail.gmail.com
---
 doc/src/sgml/ref/psql-ref.sgml | 12 ++++++++++
 src/bin/psql/prompt.c          | 40 ++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index 7789fc6177..e9fde65d72 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -4310,6 +4310,18 @@ testdb=&gt; \set PROMPT1 '%[%033[1;33;40m%]%n@%/%R%[%033[0m%]%# '
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+        <term><literal>%w</literal></term>
+        <listitem>
+        <para>
+        Whitespace of the same width as <varname>PROMPT1</varname>.  This can
+        be used as a <varname>PROMPT2</varname> setting, so that multi-line
+        statements are aligned with the first line, but there is no visible
+        secondary prompt.
+        </para>
+        </listitem>
+      </varlistentry>
+
     </variablelist>
 
     To insert a percent sign into your prompt, write
diff --git a/src/bin/psql/prompt.c b/src/bin/psql/prompt.c
index 195192a95d..edb57a4cf7 100644
--- a/src/bin/psql/prompt.c
+++ b/src/bin/psql/prompt.c
@@ -39,6 +39,7 @@
  * %n - database user name
  * %/ - current database
  * %~ - like %/ but "~" when database name equals user name
+ * %w - whitespace of the same width as the most recent output of PROMPT1
  * %# - "#" if superuser, ">" otherwise
  * %R - in prompt1 normally =, or ^ if single line mode,
  *			or a ! if session is not connected to a database;
@@ -74,6 +75,7 @@ get_prompt(promptStatus_t status, ConditionalStack cstack)
 	bool		esc = false;
 	const char *p;
 	const char *prompt_string = "? ";
+	static size_t last_prompt1_width;
 
 	switch (status)
 	{
@@ -124,6 +126,13 @@ get_prompt(promptStatus_t status, ConditionalStack cstack)
 					}
 					break;
 
+					/* Whitespace of the same width as the last PROMPT1 */
+				case 'w':
+					if (pset.db)
+						memset(buf, ' ',
+							   Min(last_prompt1_width, sizeof(buf) - 1));
+					break;
+
 					/* DB server hostname (long/short) */
 				case 'M':
 				case 'm':
@@ -336,5 +345,36 @@ get_prompt(promptStatus_t status, ConditionalStack cstack)
 			strlcat(destination, buf, sizeof(destination));
 	}
 
+	/* Compute the visible width of PROMPT1, for PROMPT2's %w */
+	if (prompt_string == pset.prompt1)
+	{
+		char	   *p = destination;
+		char	   *end = p + strlen(p);
+		bool		visible = true;
+
+		last_prompt1_width = 0;
+		while (p < end && *p)
+		{
+#if defined(USE_READLINE) && defined(RL_PROMPT_START_IGNORE)
+			if (*p == RL_PROMPT_START_IGNORE)
+			{
+				visible = false;
+				++p;
+			}
+			else if (*p == RL_PROMPT_END_IGNORE)
+			{
+				visible = true;
+				++p;
+			}
+			else
+#endif
+			{
+				if (visible)
+					last_prompt1_width += PQdsplen(p, pset.encoding);
+				p += PQmblen(p, pset.encoding);
+			}
+		}
+	}
+
 	return destination;
 }
-- 
2.23.0