Raise an error if connection is used after being closed

Author: 1st1

asyncpg/connection.py

@@ -118,6 +118,7 @@ async def add_listener(self, channel, callback):
                 **channel**: name of the channel the notification was sent to;
                 **payload**: the payload.
         """
+        self.__check_open()
         if channel not in self._listeners:
             await self.fetch('LISTEN {}'.format(channel))
             self._listeners[channel] = set()
@@ -181,6 +182,7 @@ def transaction(self, *, isolation='read_committed', readonly=False,
         .. _`PostgreSQL documentation`: https://www.postgresql.org/docs/\
                                         current/static/sql-set-transaction.html
         """
+        self.__check_open()
         return transaction.Transaction(self, isolation, readonly, deferrable)
 
     async def execute(self, query: str, *args, timeout: float=None) -> str:
@@ -211,6 +213,8 @@ async def execute(self, query: str, *args, timeout: float=None) -> str:
         .. versionchanged:: 0.5.4
            Made it possible to pass query arguments.
         """
+        self.__check_open()
+
         if not args:
             return await self._protocol.query(query, timeout)
 
@@ -240,6 +244,8 @@ async def executemany(self, command: str, args,
 
         .. versionadded:: 0.7.0
         """
+        self.__check_open()
+
         if 'timeout' in kw:
             timeout = kw.pop('timeout')
         else:
@@ -311,6 +317,7 @@ def cursor(self, query, *args, prefetch=None, timeout=None):
 
         :return: A :class:`~cursor.CursorFactory` object.
         """
+        self.__check_open()
         return cursor.CursorFactory(self, query, None, args,
                                     prefetch, timeout)
 
@@ -322,6 +329,7 @@ async def prepare(self, query, *, timeout=None):
 
         :return: A :class:`~prepared_stmt.PreparedStatement` instance.
         """
+        self.__check_open()
         stmt = await self._get_statement(query, timeout, named=True)
         return prepared_stmt.PreparedStatement(self, query, stmt)
 
@@ -334,6 +342,7 @@ async def fetch(self, query, *args, timeout=None) -> list:
 
         :return list: A list of :class:`Record` instances.
         """
+        self.__check_open()
         return await self._execute(query, args, 0, timeout)
 
     async def fetchval(self, query, *args, column=0, timeout=None):
@@ -350,6 +359,7 @@ async def fetchval(self, query, *args, column=0, timeout=None):
 
         :return: The value of the specified column of the first record.
         """
+        self.__check_open()
         data = await self._execute(query, args, 1, timeout)
         if not data:
             return None
@@ -364,6 +374,7 @@ async def fetchrow(self, query, *args, timeout=None):
 
         :return: The first row as a :class:`Record` instance.
         """
+        self.__check_open()
         data = await self._execute(query, args, 1, timeout)
         if not data:
             return None
@@ -384,6 +395,8 @@ async def set_type_codec(self, typename, *,
                         data.  If ``False`` (the default), the data is
                         expected to be encoded/decoded in text.
         """
+        self.__check_open()
+
         if self._type_by_name_stmt is None:
             self._type_by_name_stmt = await self.prepare(
                 introspection.TYPE_BY_NAME)
@@ -412,6 +425,8 @@ async def set_builtin_type_codec(self, typename, *,
                         (defaults to 'public')
         :param codec_name:  The name of the builtin codec.
         """
+        self.__check_open()
+
         if self._type_by_name_stmt is None:
             self._type_by_name_stmt = await self.prepare(
                 introspection.TYPE_BY_NAME)
@@ -455,11 +470,16 @@ def terminate(self):
         self._protocol.abort()
 
     async def reset(self):
+        self.__check_open()
         self._listeners.clear()
         reset_query = self._get_reset_query()
         if reset_query:
             await self.execute(reset_query)
 
+    def __check_open(self):
+        if self.is_closed():
+            raise exceptions.InterfaceError('connection is closed')
+
     def _get_unique_id(self, prefix):
         self._uid += 1
         return '__asyncpg_{}_{}__'.format(prefix, self._uid)

tests/test_connect.py

@@ -419,3 +419,34 @@ 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))
+
+    async def test_connection_use_after_close(self):
+        def check():
+            return self.assertRaisesRegex(asyncpg.InterfaceError,
+                                          'connection is closed')
+
+        await self.con.close()
+
+        with check():
+            await self.con.add_listener('aaa', lambda: None)
+
+        with check():
+            self.con.transaction()
+
+        with check():
+            await self.con.executemany('SELECT 1', [])
+
+        with check():
+            await self.con.set_type_codec('aaa', encoder=None, decoder=None)
+
+        with check():
+            await self.con.set_builtin_type_codec('aaa', codec_name='aaa')
+
+        for meth in ('execute', 'fetch', 'fetchval', 'fetchrow',
+                     'prepare', 'cursor'):
+
+            with check():
+                await getattr(self.con, meth)('SELECT 1')
+
+        with check():
+            await self.con.reset()