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()