Skip to content

Commit

Permalink
Make PoolConnectionProxy more dynamic. Fixes #155.
Browse files Browse the repository at this point in the history
Specifically, allow proxied Connection objects to use overridden
public methods.
  • Loading branch information
1st1 committed Jun 30, 2017
1 parent 550c408 commit 6ca1f28
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 17 deletions.
35 changes: 20 additions & 15 deletions asyncpg/pool.py
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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,
Expand All @@ -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(
Expand All @@ -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(
Expand Down
10 changes: 8 additions & 2 deletions tests/test_pool.py
Expand Up @@ -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)
Expand Down

0 comments on commit 6ca1f28

Please sign in to comment.