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

Close #972

diff --git a/docs/news.rst b/docs/news.rst
index 7ffea445067ecf710655b0086a5b3134754bb489..b81b1f70c25ce11d2bdb25d912d00a28c09afecd 100644 (file)
--- 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 682d7ce4e991126c775ec5ad720bf9795680d7b4..5d6554f36fb1585f55c0f20ff11c0d2846b4edf4 100644 (file)
--- 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 6bce3f9b2455fffe550058941a62d9069b7802c7..157d64b0dbc5dbf61079f3a8680d18fb0e05f251 100644 (file)
--- 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 aca0a7eb61da8ad0afdd3045feeef3f235c052cd..68ffd9463d73a98fc059898faf2050d6c7a57be1 100644 (file)
--- 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