Add API for LISTEN; document get_settings()

Author: 1st1

asyncpg/_testbase.py

@@ -17,6 +17,7 @@
 
 
 from asyncpg import cluster as pg_cluster
+from asyncpg import pool as pg_pool
 
 
 @contextlib.contextmanager
@@ -127,6 +128,11 @@ def setUp(self):
             'log_connections': 'on'
         })
 
+    def create_pool(self, **kwargs):
+        addr = self.cluster.get_connection_addr()
+        return pg_pool.create_pool(host=addr[0], port=addr[1],
+                                   loop=self.loop, **kwargs)
+
 
 class ConnectedTestCase(ClusterTestCase):
 

asyncpg/connection.py

@@ -29,7 +29,7 @@ class Connection:
     __slots__ = ('_protocol', '_transport', '_loop', '_types_stmt',
                  '_type_by_name_stmt', '_top_xact', '_uid', '_aborted',
                  '_stmt_cache_max_size', '_stmt_cache', '_stmts_to_close',
-                 '_addr', '_opts', '_command_timeout')
+                 '_addr', '_opts', '_command_timeout', '_listeners')
 
     def __init__(self, protocol, transport, loop, addr, opts, *,
                  statement_cache_size, command_timeout):
@@ -51,7 +51,44 @@ def __init__(self, protocol, transport, loop, addr, opts, *,
 
         self._command_timeout = command_timeout
 
+        self._listeners = {}
+
+    async def add_listener(self, channel, callback):
+        """Add a listener for Postgres notifications.
+
+        :param str channel: Channel to listen on.
+        :param callable callback:
+                A callable receiving the following arguments:
+                **connection**: a Connection the callback is registered with;
+                **pid**: PID of the Postgres server that sent the notification;
+                **channel**: name of the channel the notification was sent to;
+                **payload**: the payload.
+        """
+        if channel not in self._listeners:
+            await self.fetch('LISTEN {}'.format(channel))
+            self._listeners[channel] = set()
+        self._listeners[channel].add(callback)
+
+    async def remove_listener(self, channel, callback):
+        """Remove a listening callback on the specified channel."""
+        if channel not in self._listeners:
+            return
+        if callback not in self._listeners[channel]:
+            return
+        self._listeners[channel].remove(callback)
+        if not self._listeners[channel]:
+            del self._listeners[channel]
+            await self.fetch('UNLISTEN {}'.format(channel))
+
+    def get_server_pid(self):
+        """Return the PID of the Postgres server the connection is bound to."""
+        return self._protocol.get_server_pid()
+
     def get_settings(self):
+        """Return connection settings.
+
+        :return: :class:`~asyncpg.ConnectionSettings`.
+        """
         return self._protocol.get_settings()
 
     def transaction(self, *, isolation='read_committed', readonly=False,
@@ -269,17 +306,20 @@ async def close(self):
         if self.is_closed():
             return
         self._close_stmts()
+        self._listeners = {}
         self._aborted = True
         protocol = self._protocol
         await protocol.close()
 
     def terminate(self):
         """Terminate the connection without waiting for pending data."""
         self._close_stmts()
+        self._listeners = {}
         self._aborted = True
         self._protocol.abort()
 
     async def reset(self):
+        self._listeners = {}
         await self.execute('''
             SET SESSION AUTHORIZATION DEFAULT;
             RESET ALL;
@@ -351,6 +391,20 @@ async def cancel():
 
         self._loop.create_task(cancel())
 
+    def _notify(self, pid, channel, payload):
+        if channel not in self._listeners:
+            return
+
+        for cb in self._listeners[channel]:
+            try:
+                cb(self, pid, channel, payload)
+            except Exception as ex:
+                self._loop.call_exception_handler({
+                    'message': 'Unhandled exception in asyncpg notification '
+                               'listener callback {!r}'.format(cb),
+                    'exception': ex
+                })
+
 
 async def connect(dsn=None, *,
                   host=None, port=None,

asyncpg/protocol/coreproto.pxd

@@ -101,6 +101,7 @@ cdef class CoreProtocol:
 
     cdef _parse_msg_authentication(self)
     cdef _parse_msg_parameter_status(self)
+    cdef _parse_msg_notification(self)
     cdef _parse_msg_backend_key_data(self)
     cdef _parse_msg_ready_for_query(self)
     cdef _parse_data_msgs(self)
@@ -140,5 +141,6 @@ cdef class CoreProtocol:
     cdef _decode_row(self, const char* buf, int32_t buf_len)
 
     cdef _on_result(self)
+    cdef _on_notification(self, pid, channel, payload)
     cdef _set_server_parameter(self, name, val)
     cdef _on_connection_lost(self, exc)

asyncpg/protocol/coreproto.pyx

@@ -46,6 +46,10 @@ cdef class CoreProtocol:
                     # ParameterStatus
                     self._parse_msg_parameter_status()
                     continue
+                elif mtype == b'A':
+                    # NotificationResponse
+                    self._parse_msg_notification()
+                    continue
 
                 if state == PROTOCOL_AUTH:
                     self._process__auth(mtype)
@@ -304,6 +308,12 @@ cdef class CoreProtocol:
 
         self._set_server_parameter(name, val)
 
+    cdef _parse_msg_notification(self):
+        pid = self.buffer.read_int32()
+        channel = self.buffer.read_cstr().decode(self.encoding)
+        payload = self.buffer.read_cstr().decode(self.encoding)
+        self._on_notification(pid, channel, payload)
+
     cdef _parse_msg_authentication(self):
         cdef:
             int32_t status
@@ -617,6 +627,9 @@ cdef class CoreProtocol:
     cdef _on_result(self):
         pass
 
+    cdef _on_notification(self, pid, channel, payload):
+        pass
+
     cdef _on_connection_lost(self, exc):
         pass
 

asyncpg/protocol/protocol.pyx

@@ -113,6 +113,9 @@ cdef class BaseProtocol(CoreProtocol):
     def set_connection(self, connection):
         self.connection = connection
 
+    def get_server_pid(self):
+        return self.backend_pid
+
     def get_settings(self):
         return self.settings
 
@@ -445,6 +448,9 @@ cdef class BaseProtocol(CoreProtocol):
             self.last_query = None
             self.return_extra = False
 
+    cdef _on_notification(self, pid, channel, payload):
+        self.connection._notify(pid, channel, payload)
+
     cdef _on_connection_lost(self, exc):
         if self.closing:
             # The connection was lost because

asyncpg/protocol/settings.pyx

@@ -32,13 +32,13 @@ cdef class ConnectionSettings:
     cpdef inline register_data_types(self, types):
         self._data_codecs.add_types(types)
 
-    cpdef inline add_python_codec(self, typeoid, typename, typeschema, typekind,
-                                 encoder, decoder, binary):
+    cpdef inline add_python_codec(self, typeoid, typename, typeschema,
+                                  typekind, encoder, decoder, binary):
         self._data_codecs.add_python_codec(typeoid, typename, typeschema,
                                            typekind, encoder, decoder, binary)
 
-    cpdef inline set_builtin_type_codec(self, typeoid, typename, typeschema, typekind,
-                                 alias_to):
+    cpdef inline set_builtin_type_codec(self, typeoid, typename, typeschema,
+                                        typekind, alias_to):
         self._data_codecs.set_builtin_type_codec(typeoid, typename, typeschema,
                                           typekind, alias_to)
 

docs/api/index.rst

@@ -301,6 +301,24 @@ of values either by a numeric index or by a field name:
 
       Return an iterator over ``(field, value)`` pairs.
 
+
+.. class:: ConnectionSettings()
+
+    A read-only collection of Connection settings.
+
+    .. describe:: settings.setting_name
+
+       Return the value of the "setting_name" setting.  Raises an
+       ``AttributeError`` if the setting is not defined.
+
+       Example:
+
+       .. code-block:: pycon
+
+           >>> connection.get_settings().client_encoding
+           'UTF8'
+
+
 Introspection
 =============
 

tests/test_connect.py

@@ -22,6 +22,14 @@ async def test_connect_1(self):
             await asyncpg.connect(user="__does_not_exist__", loop=self.loop)
 
 
+class TestSettings(tb.ConnectedTestCase):
+
+    async def test_get_settings_01(self):
+        self.assertEqual(
+            self.con.get_settings().client_encoding,
+            'UTF8')
+
+
 class TestAuthentication(tb.ConnectedTestCase):
     def setUp(self):
         super().setUp()

tests/test_listeners.py

@@ -0,0 +1,76 @@
+# Copyright (C) 2016-present the ayncpg authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of asyncpg and is released under
+# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
+
+
+import asyncio
+
+from asyncpg import _testbase as tb
+
+
+class TestListeners(tb.ClusterTestCase):
+
+    async def test_listen_01(self):
+        async with self.create_pool(database='postgres') as pool:
+            async with pool.acquire() as con:
+
+                q1 = asyncio.Queue(loop=self.loop)
+                q2 = asyncio.Queue(loop=self.loop)
+
+                def listener1(*args):
+                    q1.put_nowait(args)
+
+                def listener2(*args):
+                    q2.put_nowait(args)
+
+                await con.add_listener('test', listener1)
+                await con.add_listener('test', listener2)
+
+                await con.execute("NOTIFY test, 'aaaa'")
+
+                self.assertEqual(
+                    await q1.get(),
+                    (con, con.get_server_pid(), 'test', 'aaaa'))
+                self.assertEqual(
+                    await q2.get(),
+                    (con, con.get_server_pid(), 'test', 'aaaa'))
+
+                await con.remove_listener('test', listener2)
+
+                await con.execute("NOTIFY test, 'aaaa'")
+
+                self.assertEqual(
+                    await q1.get(),
+                    (con, con.get_server_pid(), 'test', 'aaaa'))
+                with self.assertRaises(asyncio.TimeoutError):
+                    await asyncio.wait_for(q2.get(),
+                                           timeout=0.05, loop=self.loop)
+
+                await con.reset()
+                await con.remove_listener('test', listener1)
+                await con.execute("NOTIFY test, 'aaaa'")
+
+                with self.assertRaises(asyncio.TimeoutError):
+                    await asyncio.wait_for(q1.get(),
+                                           timeout=0.05, loop=self.loop)
+                with self.assertRaises(asyncio.TimeoutError):
+                    await asyncio.wait_for(q2.get(),
+                                           timeout=0.05, loop=self.loop)
+
+    async def test_listen_02(self):
+        async with self.create_pool(database='postgres') as pool:
+            async with pool.acquire() as con1, pool.acquire() as con2:
+
+                q1 = asyncio.Queue(loop=self.loop)
+
+                def listener1(*args):
+                    q1.put_nowait(args)
+
+                await con1.add_listener('ipc', listener1)
+                await con2.execute("NOTIFY ipc, 'hello'")
+
+                self.assertEqual(
+                    await q1.get(),
+                    (con1, con2.get_server_pid(), 'ipc', 'hello'))

tests/test_pool.py

@@ -6,7 +6,6 @@
 
 
 import asyncio
-import asyncpg
 
 from asyncpg import _testbase as tb
 
@@ -16,11 +15,8 @@ class TestPool(tb.ClusterTestCase):
     async def test_pool_01(self):
         for n in {1, 3, 5, 10, 20, 100}:
             with self.subTest(tasksnum=n):
-                addr = self.cluster.get_connection_addr()
-                pool = await asyncpg.create_pool(host=addr[0], port=addr[1],
-                                                 database='postgres',
-                                                 loop=self.loop, min_size=5,
-                                                 max_size=10)
+                pool = await self.create_pool(database='postgres',
+                                              min_size=5, max_size=10)
 
                 async def worker():
                     con = await pool.acquire()
@@ -34,11 +30,8 @@ async def worker():
     async def test_pool_02(self):
         for n in {1, 3, 5, 10, 20, 100}:
             with self.subTest(tasksnum=n):
-                addr = self.cluster.get_connection_addr()
-                async with asyncpg.create_pool(host=addr[0], port=addr[1],
-                                               database='postgres',
-                                               loop=self.loop, min_size=5,
-                                               max_size=5) as pool:
+                async with self.create_pool(database='postgres',
+                                            min_size=5, max_size=5) as pool:
 
                     async def worker():
                         con = await pool.acquire(timeout=1)
@@ -49,11 +42,8 @@ async def worker():
                     await asyncio.gather(*tasks, loop=self.loop)
 
     async def test_pool_03(self):
-        addr = self.cluster.get_connection_addr()
-        pool = await asyncpg.create_pool(host=addr[0], port=addr[1],
-                                         database='postgres',
-                                         loop=self.loop, min_size=1,
-                                         max_size=1)
+        pool = await self.create_pool(database='postgres',
+                                      min_size=1, max_size=1)
 
         con = await pool.acquire(timeout=1)
         with self.assertRaises(asyncio.TimeoutError):
@@ -63,11 +53,8 @@ async def test_pool_03(self):
         del con
 
     async def test_pool_04(self):
-        addr = self.cluster.get_connection_addr()
-        pool = await asyncpg.create_pool(host=addr[0], port=addr[1],
-                                         database='postgres',
-                                         loop=self.loop, min_size=1,
-                                         max_size=1)
+        pool = await self.create_pool(database='postgres',
+                                      min_size=1, max_size=1)
 
         con = await pool.acquire(timeout=0.1)
         con.terminate()
@@ -84,11 +71,8 @@ async def test_pool_04(self):
     async def test_pool_05(self):
         for n in {1, 3, 5, 10, 20, 100}:
             with self.subTest(tasksnum=n):
-                addr = self.cluster.get_connection_addr()
-                pool = await asyncpg.create_pool(host=addr[0], port=addr[1],
-                                                 database='postgres',
-                                                 loop=self.loop, min_size=5,
-                                                 max_size=10)
+                pool = await self.create_pool(database='postgres',
+                                              min_size=5, max_size=10)
 
                 async def worker():
                     async with pool.acquire() as con: