v08-0002-Add-prompt-interpolation-and-variables-for-psql-.patch

From 06962deb46be6af15dfe694e6e6342af02964638 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Tue, 18 Feb 2025 15:32:31 +0100
Subject: Add prompt interpolation and variables for psql pipeline

Add %P prompt interpolation that reports the status pipeline: on, off or
abort. Additionally, 3 new special variables are added to report the
status of an ongoing pipeline:
- PIPELINE_SYNC_COUNT: reports the number of piped sync
- PIPELINE_COMMAND_COUNT: reports the number of piped commands, a
  command being either \bind, \bind_named, \close or \parse
- PIPELINE_RESULT_COUNT: reports the results available to read with
  \getresults
---
 doc/src/sgml/ref/psql-ref.sgml | 49 ++++++++++++++++++++++++++++++++++
 src/bin/psql/common.c          | 27 ++++++++++++++++++-
 src/bin/psql/prompt.c          | 12 +++++++++
 src/bin/psql/startup.c         |  5 ++++
 4 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index fdae1ed0479..091cbaefc15 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -3723,6 +3723,12 @@ testdb=&gt; <userinput>\setenv LESS -imx4F</userinput>
           generate one result to get.
        </para>
 
+        <para>
+          When pipeline mode is active, a dedicated prompt variable is
+          available to report the pipeline status. See
+          <xref linkend="app-psql-prompting-p-uc"/> for more details
+       </para>
+
        <para>
         Example:
 <programlisting>
@@ -4504,6 +4510,38 @@ bar
         </listitem>
       </varlistentry>
 
+      <varlistentry id="app-psql-variables-pipeline-command-count">
+        <term><varname>PIPELINE_COMMAND_COUNT</varname></term>
+        <listitem>
+        <para>
+        The number of commands generated by <literal>\bind</literal>,
+        <literal>\bind_named</literal>, <literal>\close</literal> or
+        <literal>\parse</literal> queued in an ongoing pipeline.
+        </para>
+        </listitem>
+      </varlistentry>
+
+      <varlistentry id="app-psql-variables-pipeline-result-count">
+        <term><varname>PIPELINE_RESULT_COUNT</varname></term>
+        <listitem>
+        <para>
+        The number of commands of an ongoing pipeline that were followed
+        by either a <command>\flushrequest</command> or a
+        <command>\syncpipeline</command>, forcing the server to send the
+        results and can be retrieved with <command>\getresults</command>.
+        </para>
+        </listitem>
+      </varlistentry>
+
+      <varlistentry id="app-psql-variables-pipeline-sync-count">
+        <term><varname>PIPELINE_SYNC_COUNT</varname></term>
+        <listitem>
+        <para>
+        The number of syncs queued in an ongoing pipeline.
+        </para>
+        </listitem>
+      </varlistentry>
+
       <varlistentry id="app-psql-variables-port">
         <term><varname>PORT</varname></term>
         <listitem>
@@ -4903,6 +4941,17 @@ testdb=&gt; <userinput>INSERT INTO my_table VALUES (:'content');</userinput>
         </listitem>
       </varlistentry>
 
+      <varlistentry id="app-psql-prompting-p-uc">
+        <term><literal>%P</literal></term>
+        <listitem>
+        <para>
+        Pipeline status: <literal>off</literal> when not in a pipeline,
+        or <literal>on</literal> when in an ongoing pipeline, or
+        <literal>abort</literal> when in an aborted pipeline.
+        </para>
+        </listitem>
+      </varlistentry>
+
       <varlistentry id="app-psql-prompting-r">
         <term><literal>%R</literal></term>
         <listitem>
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index b7fe555d56c..70200e8594d 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -524,6 +524,27 @@ SetShellResultVariables(int wait_result)
 }
 
 
+/*
+ * Set special pipeline variables
+ * - PIPELINE_SYNC_COUNT: The number of piped syncs
+ * - PIPELINE_COMMAND_COUNT: The number of piped commands
+ * - PIPELINE_RESULT_COUNT: The number of results available to read
+ */
+static void
+SetPipelineVariables(void)
+{
+	char		buf[32];
+
+	snprintf(buf, sizeof(buf), "%d", pset.piped_syncs);
+	SetVariable(pset.vars, "PIPELINE_SYNC_COUNT", buf);
+	snprintf(buf, sizeof(buf), "%d", pset.piped_commands);
+	SetVariable(pset.vars, "PIPELINE_COMMAND_COUNT", buf);
+	snprintf(buf, sizeof(buf), "%d", pset.available_results);
+	SetVariable(pset.vars, "PIPELINE_RESULT_COUNT", buf);
+	Assert(pset.requested_results == 0);
+}
+
+
 /*
  * ClearOrSaveResult
  *
@@ -1655,6 +1676,8 @@ ExecQueryAndProcessResults(const char *query,
 
 		CheckConnection();
 
+		SetPipelineVariables();
+
 		return -1;
 	}
 
@@ -1663,8 +1686,9 @@ ExecQueryAndProcessResults(const char *query,
 	{
 		/*
 		 * We're in a pipeline and haven't reached the pipeline end or there
-		 * was no request to read pipeline results, exit.
+		 * was no request to read pipeline results, update psql variables and exit.
 		 */
+		SetPipelineVariables();
 		return 1;
 	}
 
@@ -2097,6 +2121,7 @@ ExecQueryAndProcessResults(const char *query,
 		Assert(pset.available_results == 0);
 	}
 	Assert(pset.requested_results == 0);
+	SetPipelineVariables();
 
 	/* may need this to recover from conn loss during COPY */
 	if (!CheckConnection())
diff --git a/src/bin/psql/prompt.c b/src/bin/psql/prompt.c
index 08a14feb3c3..78505222e01 100644
--- a/src/bin/psql/prompt.c
+++ b/src/bin/psql/prompt.c
@@ -31,6 +31,7 @@
  *		sockets, "[local:/dir/name]" if not default
  * %m - like %M, but hostname only (before first dot), or always "[local]"
  * %p - backend pid
+ * %P - pipeline status: on, off or abort
  * %> - database server port number
  * %n - database user name
  * %s - service
@@ -181,7 +182,18 @@ get_prompt(promptStatus_t status, ConditionalStack cstack)
 							snprintf(buf, sizeof(buf), "%d", pid);
 					}
 					break;
+				case 'P':
+					{
+						PGpipelineStatus status = PQpipelineStatus(pset.db);
 
+						if (status == PQ_PIPELINE_ON)
+							strlcpy(buf, "on", sizeof(buf));
+						else if (status == PQ_PIPELINE_ABORTED)
+							strlcpy(buf, "abort", sizeof(buf));
+						else
+							strlcpy(buf, "off", sizeof(buf));
+						break;
+					}
 				case '0':
 				case '1':
 				case '2':
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
index 703f3f582c1..5018eedf1e5 100644
--- a/src/bin/psql/startup.c
+++ b/src/bin/psql/startup.c
@@ -205,6 +205,11 @@ main(int argc, char *argv[])
 	SetVariable(pset.vars, "PROMPT3", DEFAULT_PROMPT3);
 	SetVariableBool(pset.vars, "SHOW_ALL_RESULTS");
 
+	/* Initialize pipeline variables */
+	SetVariable(pset.vars, "PIPELINE_SYNC_COUNT", "0");
+	SetVariable(pset.vars, "PIPELINE_COMMAND_COUNT", "0");
+	SetVariable(pset.vars, "PIPELINE_RESULT_COUNT", "0");
+
 	parse_psql_options(argc, argv, &options);
 
 	/*
-- 
2.39.5 (Apple Git-154)