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

From b183638bb31c75fab28ecf67112304b52bb12dce Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Tue, 18 Feb 2025 15:32:31 +0100
Subject: [PATCH v10 2/2] 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
---
 src/bin/psql/common.c          | 27 +++++++++++++++++-
 src/bin/psql/prompt.c          | 14 ++++++++++
 src/bin/psql/startup.c         |  5 ++++
 doc/src/sgml/ref/psql-ref.sgml | 50 ++++++++++++++++++++++++++++++++++
 4 files changed, 95 insertions(+), 1 deletion(-)

diff --git src/bin/psql/common.c src/bin/psql/common.c
index bc8c40898f71..ed340a466f99 100644
--- src/bin/psql/common.c
+++ src/bin/psql/common.c
@@ -524,6 +524,26 @@ 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);
+}
+
+
 /*
  * ClearOrSaveResult
  *
@@ -1661,6 +1681,8 @@ ExecQueryAndProcessResults(const char *query,
 
 		CheckConnection();
 
+		SetPipelineVariables();
+
 		return -1;
 	}
 
@@ -1669,8 +1691,10 @@ ExecQueryAndProcessResults(const char *query,
 	{
 		/*
 		 * We are in a pipeline and have not reached the pipeline end, or
-		 * there was no request to read pipeline results, exit.
+		 * there was no request to read pipeline results.  Update the psql
+		 * variables tracking the pipeline activity and exit.
 		 */
+		SetPipelineVariables();
 		return 1;
 	}
 
@@ -2105,6 +2129,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 src/bin/psql/prompt.c src/bin/psql/prompt.c
index 08a14feb3c3f..3aa7d2d06c80 100644
--- src/bin/psql/prompt.c
+++ 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,6 +182,19 @@ get_prompt(promptStatus_t status, ConditionalStack cstack)
 							snprintf(buf, sizeof(buf), "%d", pid);
 					}
 					break;
+					/* pipeline status */
+				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':
diff --git src/bin/psql/startup.c src/bin/psql/startup.c
index 703f3f582c17..5018eedf1e57 100644
--- src/bin/psql/startup.c
+++ 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);
 
 	/*
diff --git doc/src/sgml/ref/psql-ref.sgml doc/src/sgml/ref/psql-ref.sgml
index a8a0b694bbf0..32130242e9c8 100644
--- doc/src/sgml/ref/psql-ref.sgml
+++ doc/src/sgml/ref/psql-ref.sgml
@@ -3727,6 +3727,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>
@@ -4501,6 +4507,39 @@ 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. These results 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 sync messages queued in an ongoing pipeline.
+        </para>
+        </listitem>
+      </varlistentry>
+
       <varlistentry id="app-psql-variables-port">
         <term><varname>PORT</varname></term>
         <listitem>
@@ -4900,6 +4939,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,
+        <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>
-- 
2.47.2