From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 28 Feb 2021 15:19:49 +0000 (+0100)
Subject: Add Connection.broken
X-Git-Tag: 3.0.dev0~87^2~33
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=6ea57644db36812c2f0a1678484c79ca85663768;p=thirdparty%2Fpsycopg.git

Add Connection.broken
---

diff --git a/docs/connection.rst b/docs/connection.rst
index c82c7bf65..ff54bbdc6 100644
--- a/docs/connection.rst
+++ b/docs/connection.rst
@@ -64,6 +64,9 @@ The `!Connection` class
     .. autoattribute:: closed
         :annotation: bool
 
+    .. autoattribute:: broken
+        :annotation: bool
+
 
     .. method:: cursor(*, binary: bool = False, row_factory: Optional[RowFactory] = None) -> Cursor
     .. method:: cursor(name: str, *, binary: bool = False, row_factory: Optional[RowFactory] = None) -> ServerCursor
diff --git a/psycopg3/psycopg3/connection.py b/psycopg3/psycopg3/connection.py
index d8ca89750..c0cd4d555 100644
--- a/psycopg3/psycopg3/connection.py
+++ b/psycopg3/psycopg3/connection.py
@@ -118,10 +118,10 @@ class BaseConnection(AdaptContext):
         # only a begin/commit and not a savepoint.
         self._savepoints: List[str] = []
 
+        self._closed = False  # closed by an explicit close()
         self._prepared: PrepareManager = PrepareManager()
 
         wself = ref(self)
-
         pgconn.notice_handler = partial(BaseConnection._notice_handler, wself)
         pgconn.notify_handler = partial(BaseConnection._notify_handler, wself)
 
@@ -158,9 +158,19 @@ class BaseConnection(AdaptContext):
 
     @property
     def closed(self) -> bool:
-        """`True` if the connection is closed."""
+        """`!True` if the connection is closed."""
         return self.pgconn.status == ConnStatus.BAD
 
+    @property
+    def broken(self) -> bool:
+        """
+        `!True` if the connection was interrupted.
+
+        A broken connection is always `closed`, but wasn't closed in a clean
+        way, such as using `close()` or a ``with`` block.
+        """
+        return self.pgconn.status == ConnStatus.BAD and not self._closed
+
     @property
     def autocommit(self) -> bool:
         """The autocommit state of the connection."""
@@ -480,6 +490,9 @@ class Connection(BaseConnection):
 
     def close(self) -> None:
         """Close the database connection."""
+        if self.closed:
+            return
+        self._closed = True
         self.pgconn.finish()
 
     @overload
@@ -652,6 +665,9 @@ class AsyncConnection(BaseConnection):
             await self.close()
 
     async def close(self) -> None:
+        if self.closed:
+            return
+        self._closed = True
         self.pgconn.finish()
 
     @overload
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 83328f71a..cdfbe2e0d 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -60,8 +60,10 @@ def test_connect_timeout():
 
 def test_close(conn):
     assert not conn.closed
+    assert not conn.broken
     conn.close()
     assert conn.closed
+    assert not conn.broken
     assert conn.pgconn.status == conn.ConnStatus.BAD
 
     cur = conn.cursor()
@@ -74,6 +76,18 @@ def test_close(conn):
         cur.execute("select 1")
 
 
+def test_broken(conn):
+    with pytest.raises(psycopg3.OperationalError):
+        conn.execute(
+            "select pg_terminate_backend(%s)", [conn.pgconn.backend_pid]
+        )
+    assert conn.closed
+    assert conn.broken
+    conn.close()
+    assert conn.closed
+    assert conn.broken
+
+
 def test_connection_warn_close(dsn, recwarn):
     conn = Connection.connect(dsn)
     conn.close()
@@ -110,6 +124,7 @@ def test_context_commit(conn, dsn):
             cur.execute("create table textctx ()")
 
     assert conn.closed
+    assert not conn.broken
 
     with psycopg3.connect(dsn) as conn:
         with conn.cursor() as cur:
@@ -129,6 +144,7 @@ def test_context_rollback(conn, dsn):
                 1 / 0
 
     assert conn.closed
+    assert not conn.broken
 
     with psycopg3.connect(dsn) as conn:
         with conn.cursor() as cur:
diff --git a/tests/test_connection_async.py b/tests/test_connection_async.py
index a25d2d732..2c303e140 100644
--- a/tests/test_connection_async.py
+++ b/tests/test_connection_async.py
@@ -65,8 +65,10 @@ async def test_connect_timeout():
 
 async def test_close(aconn):
     assert not aconn.closed
+    assert not aconn.broken
     await aconn.close()
     assert aconn.closed
+    assert not aconn.broken
     assert aconn.pgconn.status == aconn.ConnStatus.BAD
 
     cur = aconn.cursor()
@@ -79,6 +81,18 @@ async def test_close(aconn):
         await cur.execute("select 1")
 
 
+async def test_broken(aconn):
+    with pytest.raises(psycopg3.OperationalError):
+        await aconn.execute(
+            "select pg_terminate_backend(%s)", [aconn.pgconn.backend_pid]
+        )
+    assert aconn.closed
+    assert aconn.broken
+    await aconn.close()
+    assert aconn.closed
+    assert aconn.broken
+
+
 async def test_connection_warn_close(dsn, recwarn):
     conn = await AsyncConnection.connect(dsn)
     await conn.close()
@@ -115,6 +129,7 @@ async def test_context_commit(aconn, dsn):
             await cur.execute("create table textctx ()")
 
     assert aconn.closed
+    assert not aconn.broken
 
     async with await psycopg3.AsyncConnection.connect(dsn) as aconn:
         async with aconn.cursor() as cur:
@@ -134,6 +149,7 @@ async def test_context_rollback(aconn, dsn):
                 1 / 0
 
     assert aconn.closed
+    assert not aconn.broken
 
     async with await psycopg3.AsyncConnection.connect(dsn) as aconn:
         async with aconn.cursor() as cur: