From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 10 Apr 2024 20:37:24 +0000 (+0200)
Subject: refactor: clearer cancel_safe implementation
X-Git-Tag: 3.2.0~44^2~1
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=0bbbe4bf02a391589c65e88601f9474191d786bd;p=thirdparty%2Fpsycopg.git

refactor: clearer cancel_safe implementation

Avoid catching NotSupported, just check for the libpq version.

Also avoid the half exception handler in `_cancel_gen()`: as in the
legacy branch, warn and ignore any error happening in the outermost
method, without adding an ignore exception policy in an implementation
method.

Close #778.
---

diff --git a/psycopg/psycopg/_connection_base.py b/psycopg/psycopg/_connection_base.py
index 4944f1d31..3d876f11e 100644
--- a/psycopg/psycopg/_connection_base.py
+++ b/psycopg/psycopg/_connection_base.py
@@ -320,13 +320,9 @@ class BaseConnection(Generic[Row]):
         return True
 
     def _cancel_gen(self) -> PQGenConn[None]:
-        try:
-            cancel_conn = self.pgconn.cancel_conn()
-        except e.OperationalError as ex:  # if the connection is closed
-            logger.warning("couldn't create a cancel connection: %s", ex)
-        else:
-            cancel_conn.start()
-            yield from generators.cancel(cancel_conn)
+        cancel_conn = self.pgconn.cancel_conn()
+        cancel_conn.start()
+        yield from generators.cancel(cancel_conn)
 
     def add_notice_handler(self, callback: NoticeHandler) -> None:
         """
diff --git a/psycopg/psycopg/connection.py b/psycopg/psycopg/connection.py
index 80e2b66da..ec334d122 100644
--- a/psycopg/psycopg/connection.py
+++ b/psycopg/psycopg/connection.py
@@ -271,11 +271,17 @@ class Connection(BaseConnection[Row]):
         If the underlying libpq is older than version 17, the method will fall
         back to using the same implementation of `!cancel()`.
         """
-        if self._should_cancel():
+        if not self._should_cancel():
+            return
+
+        # TODO: replace with capabilities.has_safe_cancel after merging #782
+        if pq.__build_version__ >= 170000:
             try:
                 waiting.wait_conn(self._cancel_gen(), interval=_WAIT_INTERVAL)
-            except e.NotSupportedError:
-                self.cancel()
+            except Exception as ex:
+                logger.warning("couldn't try to cancel query: %s", str(ex))
+        else:
+            self.cancel()
 
     @contextmanager
     def transaction(
diff --git a/psycopg/psycopg/connection_async.py b/psycopg/psycopg/connection_async.py
index 61444a31c..8f36bf63e 100644
--- a/psycopg/psycopg/connection_async.py
+++ b/psycopg/psycopg/connection_async.py
@@ -288,16 +288,23 @@ class AsyncConnection(BaseConnection[Row]):
         If the underlying libpq is older than version 17, the method will fall
         back to using the same implementation of `!cancel()`.
         """
-        if self._should_cancel():
+        if not self._should_cancel():
+            return
+
+        # TODO: replace with capabilities.has_safe_cancel after merging #782
+        if pq.__build_version__ >= 170000:
             try:
                 await waiting.wait_conn_async(
                     self._cancel_gen(), interval=_WAIT_INTERVAL
                 )
-            except e.NotSupportedError:
-                if True:  # ASYNC
-                    await to_thread(self.cancel)
-                else:
-                    self.cancel()
+            except Exception as ex:
+                logger.warning("couldn't try to cancel query: %s", str(ex))
+
+        else:
+            if True:  # ASYNC
+                await to_thread(self.cancel)
+            else:
+                self.cancel()
 
     @asynccontextmanager
     async def transaction(