Add rudimentary server capability detection.

Add basic server capability detection mechanism based on server version
and parameters reported by the server through ParameterStatus messages.
This allows altering certain asyncpg behaviour based on the connected
server.

Specifically, this allows asyncpg to connect to CochroachDB servers.

Fixes #87.

Author: elprans

asyncpg/connection.py

@@ -19,6 +19,7 @@
 from . import protocol
 from . import serverversion
 from . import transaction
+from . import types
 
 
 class Connection:
@@ -31,7 +32,8 @@ class Connection:
                  '_type_by_name_stmt', '_top_xact', '_uid', '_aborted',
                  '_stmt_cache_max_size', '_stmt_cache', '_stmts_to_close',
                  '_addr', '_opts', '_command_timeout', '_listeners',
-                 '_server_version', '_intro_query')
+                 '_server_version', '_server_caps', '_intro_query',
+                 '_reset_query')
 
     def __init__(self, protocol, transport, loop, addr, opts, *,
                  statement_cache_size, command_timeout):
@@ -55,15 +57,21 @@ def __init__(self, protocol, transport, loop, addr, opts, *,
 
         self._listeners = {}
 
-        ver_string = self._protocol.get_settings().server_version
+        settings = self._protocol.get_settings()
+        ver_string = settings.server_version
         self._server_version = \
             serverversion.split_server_version_string(ver_string)
 
+        self._server_caps = _detect_server_capabilities(
+            self._server_version, settings)
+
         if self._server_version < (9, 2):
             self._intro_query = introspection.INTRO_LOOKUP_TYPES_91
         else:
             self._intro_query = introspection.INTRO_LOOKUP_TYPES
 
+        self._reset_query = None
+
     async def add_listener(self, channel, callback):
         """Add a listener for Postgres notifications.
 
@@ -107,9 +115,26 @@ def get_server_version(self):
             ServerVersion(major=9, minor=6, micro=1,
                           releaselevel='final', serial=0)
 
+        .. versionadded:: 0.8.0
         """
         return self._server_version
 
+    def get_server_capabilities(self):
+        """Return the capabilities supported by the server as detected.
+
+        The returned value is a named tuple:
+
+        .. code-block:: pycon
+
+           >>> con.get_server_capabilities()
+           ServerCapabilities(advisory_locks=True, cursors=True,
+                              notifications=True, plpgsql=True,
+                              sql_reset=True)
+
+        .. versionadded:: 0.10.0
+        """
+        return self._server_caps
+
     def get_settings(self):
         """Return connection settings.
 
@@ -394,22 +419,10 @@ def terminate(self):
         self._protocol.abort()
 
     async def reset(self):
-        self._listeners = {}
-
-        await self.execute('''
-            DO $$
-            BEGIN
-                PERFORM * FROM pg_listening_channels() LIMIT 1;
-                IF FOUND THEN
-                    UNLISTEN *;
-                END IF;
-            END;
-            $$;
-            SET SESSION AUTHORIZATION DEFAULT;
-            RESET ALL;
-            CLOSE ALL;
-            SELECT pg_advisory_unlock_all();
-        ''')
+        self._listeners.clear()
+        reset_query = self._get_reset_query()
+        if reset_query:
+            await self.execute(reset_query)
 
     def _get_unique_id(self):
         self._uid += 1
@@ -492,6 +505,35 @@ def _notify(self, pid, channel, payload):
                     'exception': ex
                 })
 
+    def _get_reset_query(self):
+        if self._reset_query is not None:
+            return self._reset_query
+
+        caps = self.get_server_capabilities()
+
+        _reset_query = ''
+        if caps.advisory_locks:
+            _reset_query += 'SELECT pg_advisory_unlock_all();\n'
+        if caps.cursors:
+            _reset_query += 'CLOSE ALL;\n'
+        if caps.notifications and caps.plpgsql:
+            _reset_query += '''
+                DO $$
+                BEGIN
+                    PERFORM * FROM pg_listening_channels() LIMIT 1;
+                    IF FOUND THEN
+                        UNLISTEN *;
+                    END IF;
+                END;
+                $$;
+            '''
+        if caps.sql_reset:
+            _reset_query += 'RESET ALL;\n'
+
+        self._reset_query = _reset_query
+
+        return _reset_query
+
 
 async def connect(dsn=None, *,
                   host=None, port=None,
@@ -730,3 +772,20 @@ def _create_future(loop):
         return asyncio.Future(loop=loop)
     else:
         return create_future()
+
+
+def _detect_server_capabilities(server_version, connection_settings):
+    if hasattr(connection_settings, 'crdb_version'):
+        # CocroachDB detected.
+        advisory_locks = cursors = notifications = plpgsql = sql_reset = False
+    else:
+        # Standard PostgreSQL server assumed.
+        advisory_locks = cursors = notifications = plpgsql = sql_reset = True
+
+    return types.ServerCapabilities(
+        advisory_locks=advisory_locks,
+        cursors=cursors,
+        notifications=notifications,
+        plpgsql=plpgsql,
+        sql_reset=sql_reset
+    )

asyncpg/protocol/settings.pyx

@@ -60,3 +60,6 @@ cdef class ConnectionSettings:
                 raise AttributeError(name) from None
 
         return object.__getattr__(self, name)
+
+    def __repr__(self):
+        return '<ConnectionSettings {!r}>'.format(self._settings)

asyncpg/types.py

@@ -10,7 +10,8 @@
 
 __all__ = (
     'Type', 'Attribute', 'Range', 'BitString', 'Point', 'Path', 'Polygon',
-    'Box', 'Line', 'LineSegment', 'Circle', 'ServerVersion'
+    'Box', 'Line', 'LineSegment', 'Circle', 'ServerVersion',
+    'ServerCapabilities'
 )
 
 
@@ -34,6 +35,12 @@
 ServerVersion.__doc__ = 'PostgreSQL server version tuple.'
 
 
+ServerCapabilities = collections.namedtuple(
+    'ServerCapabilities',
+    ['advisory_locks', 'cursors', 'notifications', 'plpgsql', 'sql_reset'])
+ServerCapabilities.__doc__ = 'PostgreSQL server capabilities.'
+
+
 class Range:
     """Immutable representation of PostgreSQL `range` type."""