pool: Track connections and prohibit using them after release.

Connection pool now wraps all connections in `PooledConnectionProxy`
objects to raise `InterfaceError` if they are used after being
released back to the pool.  We also check if connection passed
to `pool.release` actually belong to the pool and correctly handle
multiple calls to `pool.release` with the same connection object.

`PooledConnectionProxy` transparently wraps Connection instances,
exposing all Connection public API.

`isinstance(asyncpg.connection.Connection)` is `True` for Instances
of `PooledConnectionProxy` class.

Author: 1st1

asyncpg/connection.py

@@ -14,14 +14,22 @@
 import urllib.parse
 
 from . import cursor
+from . import exceptions
 from . import introspection
 from . import prepared_stmt
 from . import protocol
 from . import serverversion
 from . import transaction
 
 
-class Connection:
+class ConnectionMeta(type):
+
+    def __instancecheck__(cls, instance):
+        mro = type(instance).__mro__
+        return Connection in mro or _ConnectionProxy in mro
+
+
+class Connection(metaclass=ConnectionMeta):
     """A representation of a database session.
 
     Connections are created by calling :func:`~asyncpg.connection.connect`.
@@ -32,7 +40,7 @@ class Connection:
                  '_stmt_cache_max_size', '_stmt_cache', '_stmts_to_close',
                  '_addr', '_opts', '_command_timeout', '_listeners',
                  '_server_version', '_server_caps', '_intro_query',
-                 '_reset_query')
+                 '_reset_query', '_proxy')
 
     def __init__(self, protocol, transport, loop, addr, opts, *,
                  statement_cache_size, command_timeout):
@@ -70,6 +78,7 @@ def __init__(self, protocol, transport, loop, addr, opts, *,
             self._intro_query = introspection.INTRO_LOOKUP_TYPES
 
         self._reset_query = None
+        self._proxy = None
 
     async def add_listener(self, channel, callback):
         """Add a listener for Postgres notifications.
@@ -478,9 +487,18 @@ def _notify(self, pid, channel, payload):
         if channel not in self._listeners:
             return
 
+        if self._proxy is None:
+            con_ref = self
+        else:
+            # `_proxy` is not None when the connection is a member
+            # of a connection pool.  Which means that the user is working
+            # with a PooledConnectionProxy instance, and expects to see it
+            # (and not the actual Connection) in their event callbacks.
+            con_ref = self._proxy
+
         for cb in self._listeners[channel]:
             try:
-                cb(self, pid, channel, payload)
+                cb(con_ref, pid, channel, payload)
             except Exception as ex:
                 self._loop.call_exception_handler({
                     'message': 'Unhandled exception in asyncpg notification '
@@ -517,6 +535,14 @@ def _get_reset_query(self):
 
         return _reset_query
 
+    def _set_proxy(self, proxy):
+        if self._proxy is not None and proxy is not None:
+            # Should not happen unless there is a bug in `Pool`.
+            raise exceptions.InterfaceError(
+                'internal asyncpg error: connection is already proxied')
+
+        self._proxy = proxy
+
 
 async def connect(dsn=None, *,
                   host=None, port=None,
@@ -526,7 +552,7 @@ async def connect(dsn=None, *,
                   timeout=60,
                   statement_cache_size=100,
                   command_timeout=None,
-                  connection_class=Connection,
+                  __connection_class__=Connection,
                   **opts):
     """A coroutine to establish a connection to a PostgreSQL server.
 
@@ -564,11 +590,7 @@ async def connect(dsn=None, *,
     :param float command_timeout: the default timeout for operations on
                           this connection (the default is no timeout).
 
-    :param builtins.type connection_class: A class used to represent
-                the connection.
-                Defaults to :class:`~asyncpg.connection.Connection`.
-
-    :return: A *connection_class* instance.
+    :return: A :class:`~asyncpg.connection.Connection` instance.
 
     Example:
 
@@ -582,10 +604,6 @@ async def connect(dsn=None, *,
         ...     print(types)
         >>> asyncio.get_event_loop().run_until_complete(run())
         [<Record typname='bool' typnamespace=11 ...
-
-
-    .. versionadded:: 0.10.0
-       *connection_class* argument.
     """
     if loop is None:
         loop = asyncio.get_event_loop()
@@ -629,13 +647,18 @@ async def connect(dsn=None, *,
         tr.close()
         raise
 
-    con = connection_class(pr, tr, loop, addr, opts,
-                           statement_cache_size=statement_cache_size,
-                           command_timeout=command_timeout)
+    con = __connection_class__(pr, tr, loop, addr, opts,
+                               statement_cache_size=statement_cache_size,
+                               command_timeout=command_timeout)
     pr.set_connection(con)
     return con
 
 
+class _ConnectionProxy:
+    # Base class to enable `isinstance(Connection)` check.
+    __slots__ = ()
+
+
 def _parse_connect_params(*, dsn, host, port, user,
                           password, database, opts):
 

asyncpg/pool.py

@@ -6,11 +6,80 @@
 
 
 import asyncio
+import functools
 
 from . import connection
 from . import exceptions
 
 
+class PooledConnectionProxyMeta(type):
+
+    def __new__(mcls, name, bases, dct, *, wrap=False):
+        if wrap:
+            def get_wrapper(methname):
+                meth = getattr(connection.Connection, methname)
+
+                def wrapper(self, *args, **kwargs):
+                    return self._dispatch(meth, args, kwargs)
+
+                return wrapper
+
+            for attrname in dir(connection.Connection):
+                if attrname.startswith('_') or attrname in dct:
+                    continue
+                wrapper = get_wrapper(attrname)
+                wrapper = functools.update_wrapper(
+                    wrapper, getattr(connection.Connection, attrname))
+                dct[attrname] = wrapper
+
+            if '__doc__' not in dct:
+                dct['__doc__'] = connection.Connection.__doc__
+
+        return super().__new__(mcls, name, bases, dct)
+
+    def __init__(cls, name, bases, dct, *, wrap=False):
+        # Needed for Python 3.5 to handle `wrap` class keyword argument.
+        super().__init__(name, bases, dct)
+
+
+class PooledConnectionProxy(connection._ConnectionProxy,
+                            metaclass=PooledConnectionProxyMeta,
+                            wrap=True):
+
+    __slots__ = ('_con', '_owner')
+
+    def __init__(self, owner: 'Pool', con: connection.Connection):
+        self._con = con
+        self._owner = owner
+        con._set_proxy(self)
+
+    def _unwrap(self) -> connection.Connection:
+        if self._con is None:
+            raise exceptions.InterfaceError(
+                'internal asyncpg error: cannot unwrap pooled connection')
+
+        con, self._con = self._con, None
+        con._set_proxy(None)
+        return con
+
+    def _dispatch(self, meth, args, kwargs):
+        if self._con is None:
+            raise exceptions.InterfaceError(
+                'cannot call Connection.{}(): '
+                'connection has been released back to the pool'.format(
+                    meth.__name__))
+
+        return meth(self._con, *args, **kwargs)
+
+    def __repr__(self):
+        if self._con is None:
+            return '<{classname} [released] {id:#x}>'.format(
+                classname=self.__class__.__name__, id=id(self))
+        else:
+            return '<{classname} {con!r} {id:#x}>'.format(
+                classname=self.__class__.__name__, con=self._con, id=id(self))
+
+
 class Pool:
     """A connection pool.
 
@@ -168,6 +237,8 @@ async def _acquire_impl(self):
             else:
                 con = await self._queue.get()
 
+        con = PooledConnectionProxy(self, con)
+
         if self._setup is not None:
             try:
                 await self._setup(con)
@@ -179,6 +250,20 @@ async def _acquire_impl(self):
 
     async def release(self, connection):
         """Release a database connection back to the pool."""
+
+        if (connection.__class__ is not PooledConnectionProxy or
+                connection._owner is not self):
+            raise exceptions.InterfaceError(
+                'Pool.release() received invalid connection: '
+                '{connection!r} is not a member of this pool'.format(
+                    connection=connection))
+
+        if connection._con is None:
+            # Already released, do nothing.
+            return
+
+        connection = connection._unwrap()
+
         # Use asyncio.shield() to guarantee that task cancellation
         # does not prevent the connection from being returned to the
         # pool properly.
@@ -325,6 +410,10 @@ def create_pool(dsn=None, *,
     :param loop: An asyncio event loop instance.  If ``None``, the default
                  event loop will be used.
     :return: An instance of :class:`~asyncpg.pool.Pool`.
+
+    .. versionchanged:: 0.10.0
+       An :exc:`~asyncpg.exceptions.InterfaceError` will be raised on any
+       attempted operation on a released connection.
     """
     return Pool(dsn,
                 min_size=min_size, max_size=max_size,

tests/test_connect.py

@@ -13,7 +13,7 @@
 
 import asyncpg
 from asyncpg import _testbase as tb
-from asyncpg.connection import _parse_connect_params
+from asyncpg import connection
 from asyncpg.serverversion import split_server_version_string
 
 _system = platform.uname().system
@@ -355,7 +355,7 @@ def run_testcase(self, testcase):
             if expected_error:
                 es.enter_context(self.assertRaisesRegex(*expected_error))
 
-            result = _parse_connect_params(
+            result = connection._parse_connect_params(
                 dsn=dsn, host=host, port=port, user=user, password=password,
                 database=database, opts=opts)
 
@@ -411,3 +411,11 @@ def test_test_connect_params_run_testcase(self):
     def test_connect_params(self):
         for testcase in self.TESTS:
             self.run_testcase(testcase)
+
+
+class TestConnection(tb.ConnectedTestCase):
+
+    async def test_connection_isinstance(self):
+        self.assertTrue(isinstance(self.con, connection.Connection))
+        self.assertTrue(isinstance(self.con, object))
+        self.assertFalse(isinstance(self.con, list))