From 5475f855141251776e6d1a08902fde895fb99782 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Mon, 22 Jan 2024 09:19:20 +0100
Subject: [PATCH v3 1/2] Abort pgbench if script end is reached with an opened
 pipeline

When a pipeline is opened with \startpipeline and not closed, pgbench
will either error on the next transaction with a "already in pipeline
mode" error or successfully ends if this was the last transaction
despite not sending anything that was piped in the pipeline.

This change adds an error if the end of script is reached while there's
an opened pipeline.
---
 src/bin/pgbench/pgbench.c                    |  8 +++++-
 src/bin/pgbench/t/001_pgbench_with_server.pl | 28 ++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 2574454839..03ad74a738 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -3766,7 +3766,13 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
 				/* Transition to script end processing if done */
 				if (command == NULL)
 				{
-					st->state = CSTATE_END_TX;
+					if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF)
+					{
+						pg_log_error("client %d aborted: end of script reached without closing the ongoing pipeline",
+									 st->id);
+						st->state = CSTATE_ABORTED;
+					} else
+						st->state = CSTATE_END_TX;
 					break;
 				}
 
diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl
index fc57facf9e..532e4a5e98 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -876,6 +876,34 @@ select 1 \gset f
 }
 	});
 
+# Try \startpipeline without \endpipeline in a single transaction
+$node->pgbench(
+	'-t 1 -n -M extended',
+	2,
+	[],
+	[qr{end of script reached without closing the ongoing pipeline}],
+	'error: call \startpipeline without \endpipeline in a single transaction',
+	{
+		'001_pgbench_pipeline_5' => q{
+-- startpipeline only with single transaction
+\startpipeline
+}
+	});
+
+# Try \startpipeline without \endpipeline
+$node->pgbench(
+	'-t 2 -n -M extended',
+	2,
+	[],
+	[qr{end of script reached without closing the ongoing pipeline}],
+	'error: call \startpipeline without \endpipeline',
+	{
+		'001_pgbench_pipeline_6' => q{
+-- startpipeline only
+\startpipeline
+}
+	});
+
 # Working \startpipeline in prepared query mode with serializable
 $node->pgbench(
 	'-c4 -t 10 -n -M prepared',
-- 
2.39.3 (Apple Git-145)