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

@@ -21,7 +21,17 @@
 from . import transaction
 
 
-class Connection:
+class _ConnectionProxy:
+ pass
+
+
+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 +42,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 +80,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 +489,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 +537,9 @@ def _get_reset_query(self):
 
  return _reset_query
 
+ def _set_proxy(self, proxy):
+ self._proxy = proxy
+
 
 async def connect(dsn=None, *,
  host=None, port=None,
@@ -526,7 +549,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 +587,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 +601,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,9 +644,9 @@ 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
 

asyncpg/pool.py

@@ -6,11 +6,81 @@
 
 
 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):
+ super().__init__(name, bases, dct)
+
+
+class PooledConnectionProxy(connection._ConnectionProxy,
+ metaclass=PooledConnectionProxyMeta,
+ wrap=True):
+
+ __slots__ = ('_in_pool', '_con', '_owner')
+
+ def __init__(self, owner: 'Pool', con: connection.Connection):
+ self._in_pool = True
+ self._con = con
+ self._owner = owner
+ con._set_proxy(self)
+
+ def _unwrap(self) -> connection.Connection:
+ if not self._in_pool:
+ raise exceptions.InterfaceError(
+ 'internal asyncpg error: cannot unwrap pooled connection')
+
+ self._in_pool = False
+ con, self._con = self._con, None
+ con._set_proxy(None)
+ return con
+
+ def _dispatch(self, meth, args, kwargs):
+ if not self._in_pool:
+ 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 +238,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 +251,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 +411,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))

tests/test_pool.py

@@ -6,6 +6,8 @@
 
 
 import asyncio
+import asyncpg
+import inspect
 import platform
 import os
 import unittest
@@ -35,7 +37,7 @@ async def reset(self):
 class SlowResetConnectionPool(pg_pool.Pool):
  async def _connect(self, *args, **kwargs):
  return await pg_connection.connect(
- *args, connection_class=SlowResetConnection, **kwargs)
+ *args, __connection_class__=SlowResetConnection, **kwargs)
 
 
 class TestPool(tb.ConnectedTestCase):
@@ -88,7 +90,7 @@ async def test_pool_04(self):
  con.terminate()
  await pool.release(con)
 
- async with pool.acquire(timeout=POOL_NOMINAL_TIMEOUT):
+ async with pool.acquire(timeout=POOL_NOMINAL_TIMEOUT) as con:
  con.terminate()
 
  con = await pool.acquire(timeout=POOL_NOMINAL_TIMEOUT)
@@ -127,7 +129,7 @@ async def test_pool_07(self):
  cons = set()
 
  async def setup(con):
- if con not in cons:
+ if con._con not in cons: # `con` is `PooledConnectionProxy`.
  raise RuntimeError('init was not called before setup')
 
  async def init(con):
@@ -137,7 +139,7 @@ async def init(con):
 
  async def user(pool):
  async with pool.acquire() as con:
- if con not in cons:
+ if con._con not in cons: # `con` is `PooledConnectionProxy`.
  raise RuntimeError('init was not called')
 
  async with self.create_pool(database='postgres',
@@ -150,6 +152,79 @@ async def user(pool):
 
  self.assertEqual(len(cons), 5)
 
+ async def test_pool_08(self):
+ pool = await self.create_pool(database='postgres',
+ min_size=1, max_size=1)
+
+ con = await pool.acquire(timeout=POOL_NOMINAL_TIMEOUT)
+ with self.assertRaisesRegex(asyncpg.InterfaceError, 'is not a member'):
+ await pool.release(con._con)
+
+ async def test_pool_09(self):
+ pool1 = await self.create_pool(database='postgres',
+ min_size=1, max_size=1)
+
+ pool2 = await self.create_pool(database='postgres',
+ min_size=1, max_size=1)
+
+ con = await pool1.acquire(timeout=POOL_NOMINAL_TIMEOUT)
+ with self.assertRaisesRegex(asyncpg.InterfaceError, 'is not a member'):
+ await pool2.release(con)
+
+ await pool1.close()
+ await pool2.close()
+
+ async def test_pool_10(self):
+ pool = await self.create_pool(database='postgres',
+ min_size=1, max_size=1)
+
+ con = await pool.acquire()
+ await pool.release(con)
+ await pool.release(con)
+
+ await pool.close()
+
+ async def test_pool_11(self):
+ pool = await self.create_pool(database='postgres',
+ min_size=1, max_size=1)
+
+ async with pool.acquire() as con:
+ self.assertIn(repr(con._con), repr(con)) # Test __repr__.
+
+ self.assertIn('[released]', repr(con))
+
+ with self.assertRaisesRegex(
+ asyncpg.InterfaceError,
+ r'cannot call Connection\.execute.*released back to the pool'):
+
+ con.execute('select 1')
+
+ await pool.close()
+
+ async def test_pool_12(self):
+ pool = await self.create_pool(database='postgres',
+ min_size=1, max_size=1)
+
+ async with pool.acquire() as con:
+ self.assertTrue(isinstance(con, pg_connection.Connection))
+ self.assertFalse(isinstance(con, list))
+
+ await pool.close()
+
+ async def test_pool_13(self):
+ pool = await self.create_pool(database='postgres',
+ min_size=1, max_size=1)
+
+ async with pool.acquire() as con:
+ self.assertIn('Execute an SQL command', con.execute.__doc__)
+ self.assertEqual(con.execute.__name__, 'execute')
+
+ self.assertIn(
+ str(inspect.signature(con.execute))[1:],
+ str(inspect.signature(pg_connection.Connection.execute)))
+
+ await pool.close()
+
  async def test_pool_auth(self):
  if not self.cluster.is_managed():
  self.skipTest('unmanaged cluster')