From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 20 Dec 2024 12:12:54 +0000 (+0100)
Subject: fix: call the notifies callback while the notifies generator is used
X-Git-Tag: 3.2.4~16
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=036d8e2a81dca46bb6cc59c936a230eafc15f242;p=thirdparty%2Fpsycopg.git

fix: call the notifies callback while the notifies generator is used

Close #972
---

diff --git a/docs/news.rst b/docs/news.rst
index 7ffea4450..b81b1f70c 100644
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -7,6 +7,16 @@
 ``psycopg`` release notes
 =========================
 
+Future releases
+---------------
+
+Psycopg 3.2.4 (unreleased)
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- Make sure that the notifies callback is called during the use of the
+  `~Connection.notifies()` generator (:ticket:`#972`).
+
+
 Current release
 ---------------
 
diff --git a/psycopg/psycopg/generators.py b/psycopg/psycopg/generators.py
index 682d7ce4e..5d6554f36 100644
--- a/psycopg/psycopg/generators.py
+++ b/psycopg/psycopg/generators.py
@@ -298,6 +298,8 @@ def notifies(pgconn: PGconn) -> PQGen[list[pq.PGnotify]]:
         n = pgconn.notifies()
         if n:
             ns.append(n)
+            if pgconn.notify_handler:
+                pgconn.notify_handler(n)
         else:
             break
 
diff --git a/tests/test_notify.py b/tests/test_notify.py
index 6bce3f9b2..157d64b0d 100644
--- a/tests/test_notify.py
+++ b/tests/test_notify.py
@@ -219,3 +219,37 @@ def test_notifies_blocking(conn):
         gather(worker)
 
     assert dt > 0.5
+
+
+@pytest.mark.slow
+def test_generator_and_handler(conn, conn_cls, dsn):
+    conn.set_autocommit(True)
+    conn.execute("listen foo")
+
+    n1 = None
+    n2 = None
+
+    def set_n2(n):
+        nonlocal n2
+        n2 = n
+
+    conn.add_notify_handler(set_n2)
+
+    def listener():
+        nonlocal n1
+        for n1 in conn.notifies(timeout=1, stop_after=1):
+            pass
+
+    worker = spawn(listener)
+    try:
+        # Make sure the listener is listening
+        if not conn.lock.locked():
+            sleep(0.01)
+
+        with conn_cls.connect(dsn, autocommit=True) as nconn:
+            nconn.execute("notify foo, '1'")
+    finally:
+        gather(worker)
+
+    assert n1
+    assert n2
diff --git a/tests/test_notify_async.py b/tests/test_notify_async.py
index aca0a7eb6..68ffd9463 100644
--- a/tests/test_notify_async.py
+++ b/tests/test_notify_async.py
@@ -215,3 +215,38 @@ async def test_notifies_blocking(aconn):
         await gather(worker)
 
     assert dt > 0.5
+
+
+@pytest.mark.slow
+async def test_generator_and_handler(aconn, aconn_cls, dsn):
+    await aconn.set_autocommit(True)
+    await aconn.execute("listen foo")
+
+    n1 = None
+    n2 = None
+
+    def set_n2(n):
+        nonlocal n2
+        n2 = n
+
+    aconn.add_notify_handler(set_n2)
+
+    async def listener():
+        nonlocal n1
+        async for n1 in aconn.notifies(timeout=1, stop_after=1):
+            pass
+
+    worker = spawn(listener)
+    try:
+        # Make sure the listener is listening
+        if not aconn.lock.locked():
+            await asleep(0.01)
+
+        async with await aconn_cls.connect(dsn, autocommit=True) as nconn:
+            await nconn.execute("notify foo, '1'")
+
+    finally:
+        await gather(worker)
+
+    assert n1
+    assert n2