From dced53f79657fa287825038ea606e696a72d5dd2 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Wed, 3 Jul 2024 22:39:29 -0500
Subject: [PATCH] Better error message when reading past end of batch

https://github.com/jackc/pgx/issues/1801
---
 batch.go      | 34 ++++++++++++++++++++--------------
 batch_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/batch.go b/batch.go
index 3540f57f..54cfd00c 100644
--- a/batch.go
+++ b/batch.go
@@ -128,7 +128,7 @@ func (br *batchResults) Exec() (pgconn.CommandTag, error) {
 	if !br.mrr.NextResult() {
 		err := br.mrr.Close()
 		if err == nil {
-			err = errors.New("no result")
+			err = errors.New("no more results in batch")
 		}
 		if br.conn.batchTracer != nil {
 			br.conn.batchTracer.TraceBatchQuery(br.ctx, br.conn, TraceBatchQueryData{
@@ -180,7 +180,7 @@ func (br *batchResults) Query() (Rows, error) {
 	if !br.mrr.NextResult() {
 		rows.err = br.mrr.Close()
 		if rows.err == nil {
-			rows.err = errors.New("no result")
+			rows.err = errors.New("no more results in batch")
 		}
 		rows.closed = true
 
@@ -287,7 +287,10 @@ func (br *pipelineBatchResults) Exec() (pgconn.CommandTag, error) {
 		return pgconn.CommandTag{}, br.err
 	}
 
-	query, arguments, _ := br.nextQueryAndArgs()
+	query, arguments, err := br.nextQueryAndArgs()
+	if err != nil {
+		return pgconn.CommandTag{}, err
+	}
 
 	results, err := br.pipeline.GetResults()
 	if err != nil {
@@ -330,9 +333,9 @@ func (br *pipelineBatchResults) Query() (Rows, error) {
 		return &baseRows{err: br.err, closed: true}, br.err
 	}
 
-	query, arguments, ok := br.nextQueryAndArgs()
-	if !ok {
-		query = "batch query"
+	query, arguments, err := br.nextQueryAndArgs()
+	if err != nil {
+		return &baseRows{err: err, closed: true}, err
 	}
 
 	rows := br.conn.getRows(br.ctx, query, arguments)
@@ -421,13 +424,16 @@ func (br *pipelineBatchResults) earlyError() error {
 	return br.err
 }
 
-func (br *pipelineBatchResults) nextQueryAndArgs() (query string, args []any, ok bool) {
-	if br.b != nil && br.qqIdx < len(br.b.QueuedQueries) {
-		bi := br.b.QueuedQueries[br.qqIdx]
-		query = bi.SQL
-		args = bi.Arguments
-		ok = true
-		br.qqIdx++
+func (br *pipelineBatchResults) nextQueryAndArgs() (query string, args []any, err error) {
+	if br.b == nil {
+		return "", nil, errors.New("no reference to batch")
 	}
-	return
+
+	if br.qqIdx >= len(br.b.QueuedQueries) {
+		return "", nil, errors.New("no more results in batch")
+	}
+
+	bi := br.b.QueuedQueries[br.qqIdx]
+	br.qqIdx++
+	return bi.SQL, bi.Arguments, nil
 }
diff --git a/batch_test.go b/batch_test.go
index 9ff2417f..eb560e06 100644
--- a/batch_test.go
+++ b/batch_test.go
@@ -290,6 +290,45 @@ func TestConnSendBatchMany(t *testing.T) {
 	})
 }
 
+// https://github.com/jackc/pgx/issues/1801#issuecomment-2203784178
+func TestConnSendBatchReadResultsWhenNothingQueued(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	pgxtest.RunWithQueryExecModes(ctx, t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		batch := &pgx.Batch{}
+		br := conn.SendBatch(ctx, batch)
+		commandTag, err := br.Exec()
+		require.Equal(t, "", commandTag.String())
+		require.EqualError(t, err, "no more results in batch")
+		err = br.Close()
+		require.NoError(t, err)
+	})
+}
+
+func TestConnSendBatchReadMoreResultsThanQueriesSent(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	pgxtest.RunWithQueryExecModes(ctx, t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		batch := &pgx.Batch{}
+		batch.Queue("select 1")
+		br := conn.SendBatch(ctx, batch)
+		commandTag, err := br.Exec()
+		require.Equal(t, "SELECT 1", commandTag.String())
+		require.NoError(t, err)
+		commandTag, err = br.Exec()
+		require.Equal(t, "", commandTag.String())
+		require.EqualError(t, err, "no more results in batch")
+		err = br.Close()
+		require.NoError(t, err)
+	})
+}
+
 func TestConnSendBatchWithPreparedStatement(t *testing.T) {
 	t.Parallel()