v03-0001-psql-Allow-to-add-queries-in-an-ongoing-pipeline.patch

application/octet-stream

Filename: v03-0001-psql-Allow-to-add-queries-in-an-ongoing-pipeline.patch
Type: application/octet-stream
Part: 0
Message: Re: Add Pipelining support in psql
From ae9ce7e44c1e2502429920c5d305e1ffcd30b1a5 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Wed, 5 Mar 2025 14:55:33 +0100
Subject: psql: Allow ';' to add queries in an ongoing pipeline

Currently, the only way to pipe queries in an ongoing pipeline is to
leverage psql meta-commands to create extended queries such as \bind,
\parse or \bind_named. This prevents using psql's pipeline on existing
scripts as it would require to convert all queries to use those
meta-commands.

This patch modifies ';' behaviour within an active pipeline and send all
queries as extended queries, allowing them to be piped in a pipeline.
---
 doc/src/sgml/ref/psql-ref.sgml              |  21 +-
 src/bin/psql/command.c                      |   7 +
 src/bin/psql/common.c                       |  10 +-
 src/test/regress/expected/psql_pipeline.out | 292 ++++++++++++++++++--
 src/test/regress/sql/psql_pipeline.sql      | 139 +++++++++-
 5 files changed, 429 insertions(+), 40 deletions(-)

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index cddf6e07531..2763486e268 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -3698,14 +3698,19 @@ testdb=&gt; <userinput>\setenv LESS -imx4F</userinput>
        </para>
 
        <para>
-        Pipeline mode requires the use of the extended query protocol. All
-        queries need to be sent using the meta-commands
-        <literal>\bind</literal>, <literal>\bind_named</literal>,
-        <literal>\close</literal> or <literal>\parse</literal>. While a
-        pipeline is ongoing, <literal>\sendpipeline</literal> will append the
-        current query buffer to the pipeline. Other meta-commands like
-        <literal>\g</literal>, <literal>\gx</literal> or <literal>\gdesc</literal>
-        are not allowed in pipeline mode.
+        Pipeline mode requires the use of the extended query protocol. Queries
+        can be sent using the meta-commands <literal>\bind</literal>,
+        <literal>\bind_named</literal>, <literal>\close</literal> or
+        <literal>\parse</literal>. While a pipeline is ongoing,
+        <literal>\sendpipeline</literal> will append the current query
+        buffer to the pipeline. Other meta-commands like <literal>\g</literal>,
+        <literal>\gx</literal> or <literal>\gdesc</literal> are not allowed
+        in pipeline mode.
+       </para>
+
+       <para>
+        Queries can also be sent using <literal>;</literal>. While a pipeline is
+        ongoing, they will automatically be sent using extended query protocol.
        </para>
 
        <para>
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index a87ff7e4597..bbe337780ff 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3282,6 +3282,13 @@ exec_command_watch(PsqlScanState scan_state, bool active_branch,
 		int			iter = 0;
 		int			min_rows = 0;
 
+		if (PQpipelineStatus(pset.db) != PQ_PIPELINE_OFF)
+		{
+			pg_log_error("\\watch not allowed in pipeline mode");
+			clean_extended_state();
+			success = false;
+		}
+
 		/*
 		 * Parse arguments.  We allow either an unlabeled interval or
 		 * "name=value", where name is from the set ('i', 'interval', 'c',
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index ed340a466f9..5249336bcf2 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -1668,7 +1668,15 @@ ExecQueryAndProcessResults(const char *query,
 			}
 			break;
 		case PSQL_SEND_QUERY:
-			success = PQsendQuery(pset.db, query);
+			if (PQpipelineStatus(pset.db) != PQ_PIPELINE_OFF)
+			{
+				success = PQsendQueryParams(pset.db, query,
+											0, NULL, NULL, NULL, NULL, 0);
+				if (success)
+					pset.piped_commands++;
+			}
+			else
+				success = PQsendQuery(pset.db, query);
 			break;
 	}
 
diff --git a/src/test/regress/expected/psql_pipeline.out b/src/test/regress/expected/psql_pipeline.out
index 68e3c19ea05..7dddf26a9fd 100644
--- a/src/test/regress/expected/psql_pipeline.out
+++ b/src/test/regress/expected/psql_pipeline.out
@@ -386,6 +386,262 @@ SELECT $1 \bind 3 \sendpipeline
  3
 (1 row)
 
+\endpipeline
+--
+-- Tests pipelining queries with ';'
+--
+-- Single query sent with ';'
+\startpipeline
+SELECT 1;
+\endpipeline
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Multiple queries sent with ';'
+\startpipeline
+SELECT 1;
+SELECT 2;
+SELECT 3;
+\endpipeline
+ ?column? 
+----------
+        1
+(1 row)
+
+ ?column? 
+----------
+        2
+(1 row)
+
+ ?column? 
+----------
+        3
+(1 row)
+
+-- Multiple queries on the same line can be piped with ';'
+\startpipeline
+SELECT 1; SELECT 2; SELECT 3
+;
+\endpipeline
+ ?column? 
+----------
+        1
+(1 row)
+
+ ?column? 
+----------
+        2
+(1 row)
+
+ ?column? 
+----------
+        3
+(1 row)
+
+-- Test \flush with queries piped with ';'
+\startpipeline
+\flush
+SELECT 1;
+\flush
+SELECT 2;
+SELECT 3;
+\endpipeline
+ ?column? 
+----------
+        1
+(1 row)
+
+ ?column? 
+----------
+        2
+(1 row)
+
+ ?column? 
+----------
+        3
+(1 row)
+
+-- Send multiple syncs with queries piped with ';'
+\startpipeline
+\echo :PIPELINE_COMMAND_COUNT
+0
+\echo :PIPELINE_SYNC_COUNT
+0
+\echo :PIPELINE_RESULT_COUNT
+0
+SELECT 1;
+\syncpipeline
+\syncpipeline
+SELECT 2;
+\syncpipeline
+SELECT 3;
+\echo :PIPELINE_COMMAND_COUNT
+1
+\echo :PIPELINE_SYNC_COUNT
+3
+\echo :PIPELINE_RESULT_COUNT
+2
+\endpipeline
+ ?column? 
+----------
+        1
+(1 row)
+
+ ?column? 
+----------
+        2
+(1 row)
+
+ ?column? 
+----------
+        3
+(1 row)
+
+-- Mix queries piped with ';' and \sendpipeline
+\startpipeline
+SELECT 1;
+SELECT $1 \bind 'val1' \sendpipeline
+SELECT $1, $2 \bind 'val2' 'val3' \sendpipeline
+SELECT 2;
+\endpipeline
+ ?column? 
+----------
+        1
+(1 row)
+
+ ?column? 
+----------
+ val1
+(1 row)
+
+ ?column? | ?column? 
+----------+----------
+ val2     | val3
+(1 row)
+
+ ?column? 
+----------
+        2
+(1 row)
+
+-- Piping a query with ';' will replace the unnamed prepared statement
+\startpipeline
+SELECT $1 \parse ''
+SELECT 1;
+\bind_named ''
+\endpipeline
+ ?column? 
+----------
+        1
+(1 row)
+
+-- An extended query can be piped by a ';' after a newline
+\startpipeline
+SELECT $1 \bind 1
+;
+SELECT 2;
+\endpipeline
+ ?column? 
+----------
+ 1
+(1 row)
+
+ ?column? 
+----------
+        2
+(1 row)
+
+-- COPY FROM STDIN, using ';'
+\startpipeline
+SELECT 'val1';
+COPY psql_pipeline FROM STDIN;
+\endpipeline
+ ?column? 
+----------
+ val1
+(1 row)
+
+-- COPY FROM STDIN with \flushrequest + \getresults, using ';'
+\startpipeline
+SELECT 'val1';
+COPY psql_pipeline FROM STDIN;
+\flushrequest
+\getresults
+ ?column? 
+----------
+ val1
+(1 row)
+
+message type 0x5a arrived from server while idle
+\endpipeline
+-- COPY FROM STDIN with \syncpipeline + \getresults, using ';'
+\startpipeline
+SELECT 'val1';
+COPY psql_pipeline FROM STDIN;
+\syncpipeline
+\getresults
+ ?column? 
+----------
+ val1
+(1 row)
+
+\endpipeline
+-- COPY TO STDOUT, using ';'
+\startpipeline
+SELECT 'val1';
+copy psql_pipeline TO STDOUT;
+\endpipeline
+ ?column? 
+----------
+ val1
+(1 row)
+
+1	\N
+2	test2
+3	test3
+4	test4
+20	test2
+30	test3
+40	test4
+-- COPY TO STDOUT with \flushrequest + \getresults, using ';'
+\startpipeline
+SELECT 'val1';
+copy psql_pipeline TO STDOUT;
+\flushrequest
+\getresults
+ ?column? 
+----------
+ val1
+(1 row)
+
+1	\N
+2	test2
+3	test3
+4	test4
+20	test2
+30	test3
+40	test4
+\endpipeline
+-- COPY TO STDOUT with \syncpipeline + \getresults, using ';'
+\startpipeline
+SELECT 'val1';
+copy psql_pipeline TO STDOUT;
+\syncpipeline
+\getresults
+ ?column? 
+----------
+ val1
+(1 row)
+
+1	\N
+2	test2
+3	test3
+4	test4
+20	test2
+30	test3
+40	test4
 \endpipeline
 --
 -- Pipeline errors
@@ -393,18 +649,6 @@ SELECT $1 \bind 3 \sendpipeline
 -- \endpipeline outside of pipeline should fail
 \endpipeline
 cannot send pipeline when not in pipeline mode
--- Query using simple protocol should not be sent and should leave the
--- pipeline usable.
-\startpipeline
-SELECT 1;
-PQsendQuery not allowed in pipeline mode
-SELECT $1 \bind 'val1' \sendpipeline
-\endpipeline
- ?column? 
-----------
- val1
-(1 row)
-
 -- After an aborted pipeline, commands after a \syncpipeline should be
 -- displayed.
 \startpipeline
@@ -425,6 +669,13 @@ SELECT \bind 'val1' \sendpipeline
 SELECT $1 \bind 'val1' \sendpipeline
 \endpipeline
 ERROR:  bind message supplies 1 parameters, but prepared statement "" requires 0
+-- Using ';' with a parameter will trigger an incorrect parameter error and
+-- abort the pipeline
+\startpipeline
+SELECT $1;
+SELECT 1;
+\endpipeline
+ERROR:  bind message supplies 0 parameters, but prepared statement "" requires 1
 -- An explicit transaction with an error needs to be rollbacked after
 -- the pipeline.
 \startpipeline
@@ -435,12 +686,11 @@ ROLLBACK \bind \sendpipeline
 ERROR:  duplicate key value violates unique constraint "psql_pipeline_pkey"
 DETAIL:  Key (a)=(1) already exists.
 ROLLBACK;
--- \watch sends a simple query, something not allowed within a pipeline.
+-- \watch is not allowed within a pipeline.
 \startpipeline
 SELECT \bind \sendpipeline
 \watch 1
-PQsendQuery not allowed in pipeline mode
-
+\watch not allowed in pipeline mode
 \endpipeline
 --
 (1 row)
@@ -530,7 +780,7 @@ SELECT COUNT(*) FROM psql_pipeline \bind \sendpipeline
 
  count 
 -------
-     4
+     7
 (1 row)
 
 -- After an error, pipeline is aborted and requires \syncpipeline to be
@@ -617,11 +867,11 @@ select 1;
 -- Error messages accumulate and are repeated.
 \startpipeline
 SELECT 1 \bind \sendpipeline
-SELECT 1;
-PQsendQuery not allowed in pipeline mode
-SELECT 1;
-PQsendQuery not allowed in pipeline mode
-PQsendQuery not allowed in pipeline mode
+\gdesc
+synchronous command execution functions are not allowed in pipeline mode
+\gdesc
+synchronous command execution functions are not allowed in pipeline mode
+synchronous command execution functions are not allowed in pipeline mode
 \endpipeline
  ?column? 
 ----------
diff --git a/src/test/regress/sql/psql_pipeline.sql b/src/test/regress/sql/psql_pipeline.sql
index e4d7e614af3..6f76adcba82 100644
--- a/src/test/regress/sql/psql_pipeline.sql
+++ b/src/test/regress/sql/psql_pipeline.sql
@@ -210,6 +210,125 @@ SELECT $1 \bind 3 \sendpipeline
 \getresults 0
 \endpipeline
 
+--
+-- Tests pipelining queries with ';'
+--
+
+-- Single query sent with ';'
+\startpipeline
+SELECT 1;
+\endpipeline
+
+-- Multiple queries sent with ';'
+\startpipeline
+SELECT 1;
+SELECT 2;
+SELECT 3;
+\endpipeline
+
+-- Multiple queries on the same line can be piped with ';'
+\startpipeline
+SELECT 1; SELECT 2; SELECT 3
+;
+\endpipeline
+
+-- Test \flush with queries piped with ';'
+\startpipeline
+\flush
+SELECT 1;
+\flush
+SELECT 2;
+SELECT 3;
+\endpipeline
+
+-- Send multiple syncs with queries piped with ';'
+\startpipeline
+\echo :PIPELINE_COMMAND_COUNT
+\echo :PIPELINE_SYNC_COUNT
+\echo :PIPELINE_RESULT_COUNT
+SELECT 1;
+\syncpipeline
+\syncpipeline
+SELECT 2;
+\syncpipeline
+SELECT 3;
+\echo :PIPELINE_COMMAND_COUNT
+\echo :PIPELINE_SYNC_COUNT
+\echo :PIPELINE_RESULT_COUNT
+\endpipeline
+
+-- Mix queries piped with ';' and \sendpipeline
+\startpipeline
+SELECT 1;
+SELECT $1 \bind 'val1' \sendpipeline
+SELECT $1, $2 \bind 'val2' 'val3' \sendpipeline
+SELECT 2;
+\endpipeline
+
+-- Piping a query with ';' will replace the unnamed prepared statement
+\startpipeline
+SELECT $1 \parse ''
+SELECT 1;
+\bind_named ''
+\endpipeline
+
+-- An extended query can be piped by a ';' after a newline
+\startpipeline
+SELECT $1 \bind 1
+;
+SELECT 2;
+\endpipeline
+
+-- COPY FROM STDIN, using ';'
+\startpipeline
+SELECT 'val1';
+COPY psql_pipeline FROM STDIN;
+\endpipeline
+20	test2
+\.
+
+-- COPY FROM STDIN with \flushrequest + \getresults, using ';'
+\startpipeline
+SELECT 'val1';
+COPY psql_pipeline FROM STDIN;
+\flushrequest
+\getresults
+30	test3
+\.
+\endpipeline
+
+-- COPY FROM STDIN with \syncpipeline + \getresults, using ';'
+\startpipeline
+SELECT 'val1';
+COPY psql_pipeline FROM STDIN;
+\syncpipeline
+\getresults
+40	test4
+\.
+\endpipeline
+
+-- COPY TO STDOUT, using ';'
+\startpipeline
+SELECT 'val1';
+copy psql_pipeline TO STDOUT;
+\endpipeline
+
+-- COPY TO STDOUT with \flushrequest + \getresults, using ';'
+\startpipeline
+SELECT 'val1';
+copy psql_pipeline TO STDOUT;
+\flushrequest
+\getresults
+\endpipeline
+
+-- COPY TO STDOUT with \syncpipeline + \getresults, using ';'
+\startpipeline
+SELECT 'val1';
+copy psql_pipeline TO STDOUT;
+\syncpipeline
+\getresults
+\endpipeline
+
 --
 -- Pipeline errors
 --
@@ -217,13 +336,6 @@ SELECT $1 \bind 3 \sendpipeline
 -- \endpipeline outside of pipeline should fail
 \endpipeline
 
--- Query using simple protocol should not be sent and should leave the
--- pipeline usable.
-\startpipeline
-SELECT 1;
-SELECT $1 \bind 'val1' \sendpipeline
-\endpipeline
-
 -- After an aborted pipeline, commands after a \syncpipeline should be
 -- displayed.
 \startpipeline
@@ -239,6 +351,13 @@ SELECT \bind 'val1' \sendpipeline
 SELECT $1 \bind 'val1' \sendpipeline
 \endpipeline
 
+-- Using ';' with a parameter will trigger an incorrect parameter error and
+-- abort the pipeline
+\startpipeline
+SELECT $1;
+SELECT 1;
+\endpipeline
+
 -- An explicit transaction with an error needs to be rollbacked after
 -- the pipeline.
 \startpipeline
@@ -248,7 +367,7 @@ ROLLBACK \bind \sendpipeline
 \endpipeline
 ROLLBACK;
 
--- \watch sends a simple query, something not allowed within a pipeline.
+-- \watch is not allowed within a pipeline.
 \startpipeline
 SELECT \bind \sendpipeline
 \watch 1
@@ -372,8 +491,8 @@ select 1;
 -- Error messages accumulate and are repeated.
 \startpipeline
 SELECT 1 \bind \sendpipeline
-SELECT 1;
-SELECT 1;
+\gdesc
+\gdesc
 \endpipeline
 
 --
-- 
2.39.5 (Apple Git-154)