Make PoolConnectionProxy more dynamic. Fixes #155.

Specifically, allow proxied Connection objects to use overridden
public methods.

Author: 1st1

asyncpg/pool.py

@@ -18,11 +18,6 @@ class PoolConnectionProxyMeta(type):
 
     def __new__(mcls, name, bases, dct, *, wrap=False):
         if wrap:
-            def get_wrapper(meth):
-                def wrapper(self, *args, **kwargs):
-                    return self._dispatch_method_call(meth, args, kwargs)
-                return wrapper
-
             for attrname in dir(connection.Connection):
                 if attrname.startswith('_') or attrname in dct:
                     continue
@@ -31,7 +26,7 @@ def wrapper(self, *args, **kwargs):
                 if not inspect.isfunction(meth):
                     continue
 
-                wrapper = get_wrapper(meth)
+                wrapper = mcls._wrap_connection_method(attrname)
                 wrapper = functools.update_wrapper(wrapper, meth)
                 dct[attrname] = wrapper
 
@@ -44,6 +39,21 @@ def __init__(cls, name, bases, dct, *, wrap=False):
         # Needed for Python 3.5 to handle `wrap` class keyword argument.
         super().__init__(name, bases, dct)
 
+    @staticmethod
+    def _wrap_connection_method(meth_name):
+        def call_con_method(self, *args, **kwargs):
+            # This method will be owned by PoolConnectionProxy class.
+            if self._con is None:
+                raise exceptions.InterfaceError(
+                    'cannot call Connection.{}(): '
+                    'connection has been released back to the pool'.format(
+                        meth_name))
+
+            meth = getattr(self._con.__class__, meth_name)
+            return meth(self._con, *args, **kwargs)
+
+        return call_con_method
+
 
 class PoolConnectionProxy(connection._ConnectionProxy,
                           metaclass=PoolConnectionProxyMeta,
@@ -57,6 +67,10 @@ def __init__(self, holder: 'PoolConnectionHolder',
         self._holder = holder
         con._set_proxy(self)
 
+    def __getattr__(self, attr):
+        # Proxy all unresolved attributes to the wrapped Connection object.
+        return getattr(self._con, attr)
+
     def _detach(self):
         if self._con is None:
             raise exceptions.InterfaceError(
@@ -65,15 +79,6 @@ def _detach(self):
         con, self._con = self._con, None
         con._set_proxy(None)
 
-    def _dispatch_method_call(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(

tests/test_pool.py

@@ -356,11 +356,17 @@ async def test_pool_config_persistence(self):
         cons = set()
 
         class MyConnection(asyncpg.Connection):
-            pass
+            async def foo(self):
+                return 42
+
+            async def fetchval(self, query):
+                res = await super().fetchval(query)
+                return res + 1
 
         async def test(pool):
             async with pool.acquire() as con:
-                await con.fetchval('SELECT 1')
+                self.assertEqual(await con.fetchval('SELECT 1'), 2)
+                self.assertEqual(await con.foo(), 42)
                 self.assertTrue(isinstance(con, MyConnection))
                 self.assertEqual(con._con._config.statement_cache_size, 3)
                 cons.add(con)