Guard against incorrect use of resources associated with a connection

We already check for attempts to use connection methods after the
connection has been returned to the pool.  This adds a similar check
for objects associated with the connection: prepared statements and
cursors.

Issue: MagicStack#190.

Author: elprans

asyncpg/connection.py

@@ -39,10 +39,11 @@ class Connection(metaclass=ConnectionMeta):
 
     __slots__ = ('_protocol', '_transport', '_loop', '_types_stmt',
                  '_type_by_name_stmt', '_top_xact', '_uid', '_aborted',
-                 '_stmt_cache', '_stmts_to_close', '_listeners',
-                 '_server_version', '_server_caps', '_intro_query',
-                 '_reset_query', '_proxy', '_stmt_exclusive_section',
-                 '_config', '_params', '_addr', '_log_listeners')
+                 '_pool_release_ctr', '_stmt_cache', '_stmts_to_close',
+                 '_listeners', '_server_version', '_server_caps',
+                 '_intro_query', '_reset_query', '_proxy',
+                 '_stmt_exclusive_section', '_config', '_params', '_addr',
+                 '_log_listeners')
 
     def __init__(self, protocol, transport, loop,
                  addr: (str, int) or str,
@@ -56,6 +57,10 @@ def __init__(self, protocol, transport, loop,
         self._top_xact = None
         self._uid = 0
         self._aborted = False
+        # Incremented very time the connection is released back to a pool.
+        # Used to catch invalid references to connection-related resources
+        # post-release (e.g. explicit prepared statements).
+        self._pool_release_ctr = 0
 
         self._addr = addr
         self._config = config

asyncpg/connresource.py

@@ -0,0 +1,38 @@
+
+# Copyright (C) 2016-present the asyncpg authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of asyncpg and is released under
+# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
+
+
+import functools
+
+from . import exceptions
+
+
+def guarded(meth):
+    """A decorator to add a sanity check to ConnectionResource methods."""
+
+    @functools.wraps(meth)
+    def _check(self, *args, **kwargs):
+        self._check_conn_validity(meth.__name__)
+        return meth(self, *args, **kwargs)
+
+    return _check
+
+
+class ConnectionResource:
+    __slots__ = ('_connection', '_con_release_ctr')
+
+    def __init__(self, connection):
+        self._connection = connection
+        self._con_release_ctr = getattr(connection, '_pool_release_ctr', None)
+
+    def _check_conn_validity(self, meth_name):
+        con_release_ctr = getattr(self._connection, '_pool_release_ctr', None)
+        if con_release_ctr != self._con_release_ctr:
+            raise exceptions.InterfaceError(
+                'cannot call {}.{}(): '
+                'the underlying connection has been released back '
+                'to the pool'.format(self.__class__.__name__, meth_name))

asyncpg/cursor.py

@@ -8,21 +8,21 @@
 import collections
 
 from . import compat
+from . import connresource
 from . import exceptions
 
 
-class CursorFactory:
+class CursorFactory(connresource.ConnectionResource):
     """A cursor interface for the results of a query.
 
     A cursor interface can be used to initiate efficient traversal of the
     results of a large query.
     """
 
-    __slots__ = ('_state', '_connection', '_args', '_prefetch',
-                 '_query', '_timeout')
+    __slots__ = ('_state', '_args', '_prefetch', '_query', '_timeout')
 
     def __init__(self, connection, query, state, args, prefetch, timeout):
-        self._connection = connection
+        super().__init__(connection)
         self._args = args
         self._prefetch = prefetch
         self._query = query
@@ -32,13 +32,15 @@ def __init__(self, connection, query, state, args, prefetch, timeout):
             state.attach()
 
     @compat.aiter_compat
+    @connresource.guarded
     def __aiter__(self):
         prefetch = 50 if self._prefetch is None else self._prefetch
         return CursorIterator(self._connection,
                               self._query, self._state,
                               self._args, prefetch,
                               self._timeout)
 
+    @connresource.guarded
     def __await__(self):
         if self._prefetch is not None:
             raise exceptions.InterfaceError(
@@ -53,14 +55,13 @@ def __del__(self):
             self._connection._maybe_gc_stmt(self._state)
 
 
-class BaseCursor:
+class BaseCursor(connresource.ConnectionResource):
 
-    __slots__ = ('_state', '_connection', '_args', '_portal_name',
-                 '_exhausted', '_query')
+    __slots__ = ('_state', '_args', '_portal_name', '_exhausted', '_query')
 
     def __init__(self, connection, query, state, args):
+        super().__init__(connection)
         self._args = args
-        self._connection = connection
         self._state = state
         if state is not None:
             state.attach()
@@ -162,9 +163,11 @@ def __init__(self, connection, query, state, args, prefetch, timeout):
         self._timeout = timeout
 
     @compat.aiter_compat
+    @connresource.guarded
     def __aiter__(self):
         return self
 
+    @connresource.guarded
     async def __anext__(self):
         if self._state is None:
             self._state = await self._connection._get_statement(
@@ -199,6 +202,7 @@ async def _init(self, timeout):
         await self._bind(timeout)
         return self
 
+    @connresource.guarded
     async def fetch(self, n, *, timeout=None):
         r"""Return the next *n* rows as a list of :class:`Record` objects.
 
@@ -216,6 +220,7 @@ async def fetch(self, n, *, timeout=None):
             self._exhausted = True
         return recs
 
+    @connresource.guarded
     async def fetchrow(self, *, timeout=None):
         r"""Return the next row.
 
@@ -232,6 +237,7 @@ async def fetchrow(self, *, timeout=None):
             return None
         return recs[0]
 
+    @connresource.guarded
     async def forward(self, n, *, timeout=None) -> int:
         r"""Skip over the next *n* rows.
 

asyncpg/pool.py

@@ -175,6 +175,9 @@ async def release(self):
         assert self._in_use
         self._in_use = False
 
+        # Invalidate external references to the connection.
+        self._con._pool_release_ctr += 1
+
         if self._con.is_closed():
             self._con = None
 

asyncpg/prepared_stmt.py

@@ -7,22 +7,24 @@
 
 import json
 
+from . import connresource
 from . import cursor
 from . import exceptions
 
 
-class PreparedStatement:
+class PreparedStatement(connresource.ConnectionResource):
     """A representation of a prepared statement."""
 
-    __slots__ = ('_connection', '_state', '_query', '_last_status')
+    __slots__ = ('_state', '_query', '_last_status')
 
     def __init__(self, connection, query, state):
-        self._connection = connection
+        super().__init__(connection)
         self._state = state
         self._query = query
         state.attach()
         self._last_status = None
 
+    @connresource.guarded
     def get_query(self) -> str:
         """Return the text of the query for this prepared statement.
 
@@ -33,6 +35,7 @@ def get_query(self) -> str:
         """
         return self._query
 
+    @connresource.guarded
     def get_statusmsg(self) -> str:
         """Return the status of the executed command.
 
@@ -46,6 +49,7 @@ def get_statusmsg(self) -> str:
             return self._last_status
         return self._last_status.decode()
 
+    @connresource.guarded
     def get_parameters(self):
         """Return a description of statement parameters types.
 
@@ -60,9 +64,9 @@ def get_parameters(self):
             #   (Type(oid=23, name='int4', kind='scalar', schema='pg_catalog'),
             #    Type(oid=25, name='text', kind='scalar', schema='pg_catalog'))
         """
-        self._check_open()
         return self._state._get_parameters()
 
+    @connresource.guarded
     def get_attributes(self):
         """Return a description of relation attributes (columns).
 
@@ -85,9 +89,9 @@ def get_attributes(self):
             #       type=Type(oid=26, name='oid', kind='scalar',
             #                 schema='pg_catalog')))
         """
-        self._check_open()
         return self._state._get_attributes()
 
+    @connresource.guarded
     def cursor(self, *args, prefetch=None,
                timeout=None) -> cursor.CursorFactory:
         """Return a *cursor factory* for the prepared statement.
@@ -99,11 +103,11 @@ def cursor(self, *args, prefetch=None,
 
         :return: A :class:`~cursor.CursorFactory` object.
         """
-        self._check_open()
         return cursor.CursorFactory(self._connection, self._query,
                                     self._state, args, prefetch,
                                     timeout)
 
+    @connresource.guarded
     async def explain(self, *args, analyze=False):
         """Return the execution plan of the statement.
 
@@ -145,6 +149,7 @@ async def explain(self, *args, analyze=False):
 
         return json.loads(data)
 
+    @connresource.guarded
     async def fetch(self, *args, timeout=None):
         r"""Execute the statement and return a list of :class:`Record` objects.
 
@@ -157,6 +162,7 @@ async def fetch(self, *args, timeout=None):
         data = await self.__bind_execute(args, 0, timeout)
         return data
 
+    @connresource.guarded
     async def fetchval(self, *args, column=0, timeout=None):
         """Execute the statement and return a value in the first row.
 
@@ -175,6 +181,7 @@ async def fetchval(self, *args, column=0, timeout=None):
             return None
         return data[0][column]
 
+    @connresource.guarded
     async def fetchrow(self, *args, timeout=None):
         """Execute the statement and return the first row.
 
@@ -190,16 +197,21 @@ async def fetchrow(self, *args, timeout=None):
         return data[0]
 
     async def __bind_execute(self, args, limit, timeout):
-        self._check_open()
         protocol = self._connection._protocol
         data, status, _ = await protocol.bind_execute(
             self._state, args, '', limit, True, timeout)
         self._last_status = status
         return data
 
-    def _check_open(self):
+    def _check_open(self, meth_name):
         if self._state.closed:
-            raise exceptions.InterfaceError('prepared statement is closed')
+            raise exceptions.InterfaceError(
+                'cannot call PreparedStmt.{}(): '
+                'the prepared statement is closed'.format(meth_name))
+
+    def _check_conn_validity(self, meth_name):
+        self._check_open(meth_name)
+        super()._check_conn_validity(meth_name)
 
     def __del__(self):
         self._state.detach()

tests/test_pool.py

@@ -193,6 +193,11 @@ async def test_pool_11(self):
         async with pool.acquire() as con:
             self.assertIn(repr(con._con), repr(con))  # Test __repr__.
 
+            ps = await con.prepare('SELECT 1')
+            async with con.transaction():
+                cur = await con.cursor('SELECT 1')
+                ps_cur = await ps.cursor()
+
         self.assertIn('[released]', repr(con))
 
         with self.assertRaisesRegex(
@@ -201,6 +206,32 @@ async def test_pool_11(self):
 
             con.execute('select 1')
 
+        for meth in ('fetchval', 'fetchrow', 'fetch', 'explain',
+                     'get_query', 'get_statusmsg', 'get_parameters',
+                     'get_attributes'):
+            with self.assertRaisesRegex(
+                    asyncpg.InterfaceError,
+                    r'cannot call PreparedStatement\.{meth}.*released '
+                    r'back to the pool'.format(meth=meth)):
+
+                getattr(ps, meth)()
+
+        for c in (cur, ps_cur):
+            for meth in ('fetch', 'fetchrow'):
+                with self.assertRaisesRegex(
+                        asyncpg.InterfaceError,
+                        r'cannot call Cursor\.{meth}.*released '
+                        r'back to the pool'.format(meth=meth)):
+
+                    getattr(c, meth)()
+
+            with self.assertRaisesRegex(
+                    asyncpg.InterfaceError,
+                    r'cannot call Cursor\.forward.*released '
+                    r'back to the pool'.format(meth=meth)):
+
+                c.forward(1)
+
         await pool.close()
 
     async def test_pool_12(self):