From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 27 Mar 2022 15:35:06 +0000 (+0200)
Subject: fix: better error on exit of a pipeline with broken connection
X-Git-Tag: 3.1~146^2~4
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=0b1ccd4c79b46839a5249f7ee9997604d16f3de0;p=thirdparty%2Fpsycopg.git

fix: better error on exit of a pipeline with broken connection

Raise an OperationalError with a message "the connection is closed" rather
than "connection pointer is NULL". The exception class was right already.
---

diff --git a/psycopg/psycopg/_pipeline.py b/psycopg/psycopg/_pipeline.py
index 76b6d7d20..429b2c75d 100644
--- a/psycopg/psycopg/_pipeline.py
+++ b/psycopg/psycopg/_pipeline.py
@@ -9,7 +9,7 @@ from typing import Any, List, Optional, Union, Tuple, Type, TYPE_CHECKING
 
 from . import pq
 from . import errors as e
-from .pq import ExecStatus
+from .pq import ConnStatus, ExecStatus
 from .abc import PipelineCommand, PQGen
 from ._compat import Deque, TypeAlias
 from ._cmodule import _psycopg
@@ -63,7 +63,8 @@ class BasePipeline:
         self.pgconn.enter_pipeline_mode()
 
     def _exit(self) -> None:
-        self.pgconn.exit_pipeline_mode()
+        if self.pgconn.status != ConnStatus.BAD:
+            self.pgconn.exit_pipeline_mode()
 
     def _communicate_gen(self) -> PQGen[None]:
         """Communicate with pipeline to send commands and possibly fetch
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index 331db905e..a19e4bdae 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -30,6 +30,16 @@ def test_pipeline_reenter(conn: psycopg.Connection[Any]) -> None:
     assert p1.status == pq.PipelineStatus.OFF
 
 
+def test_pipeline_broken_conn_exit(conn: psycopg.Connection[Any]) -> None:
+    with pytest.raises(e.OperationalError):
+        with conn.pipeline():
+            conn.execute("select 1")
+            conn.close()
+            closed = True
+
+    assert closed
+
+
 def test_cursor_stream(conn):
     with conn.pipeline(), conn.cursor() as cur:
         with pytest.raises(psycopg.ProgrammingError):
diff --git a/tests/test_pipeline_async.py b/tests/test_pipeline_async.py
index bb78d0822..bd079c925 100644
--- a/tests/test_pipeline_async.py
+++ b/tests/test_pipeline_async.py
@@ -33,6 +33,16 @@ async def test_pipeline_reenter(aconn: psycopg.AsyncConnection[Any]) -> None:
     assert p1.status == pq.PipelineStatus.OFF
 
 
+async def test_pipeline_broken_conn_exit(aconn: psycopg.AsyncConnection[Any]) -> None:
+    with pytest.raises(e.OperationalError):
+        async with aconn.pipeline():
+            await aconn.execute("select 1")
+            await aconn.close()
+            closed = True
+
+    assert closed
+
+
 async def test_cursor_stream(aconn):
     async with aconn.pipeline(), aconn.cursor() as cur:
         with pytest.raises(psycopg.ProgrammingError):