s/logger/query_logger/, report BaseException too

Author: elprans

asyncpg/connection.py

@@ -1803,7 +1803,7 @@ async def _execute(
  return result
 
  @contextlib.contextmanager
- def logger(self, callback):
+ def query_logger(self, callback):
  """Context manager that adds `callback` to the list of query loggers,
  and removes it upon exit.
 
@@ -1822,15 +1822,15 @@ def __init__(self):
  self.queries = []
  def __call__(self, record):
  self.queries.append(record.query)
- >>> with con.logger(QuerySaver()) as log:
+ >>> with con.query_logger(QuerySaver()):
  >>> await con.execute("SELECT 1")
  >>> print(log.queries)
  ['SELECT 1']
 
  .. versionadded:: 0.29.0
  """
  self.add_query_logger(callback)
- yield callback
+ yield
  self.remove_query_logger(callback)
 
  @contextlib.contextmanager
@@ -1839,7 +1839,7 @@ def _time_and_log(self, query, args, timeout):
  exception = None
  try:
  yield
- except Exception as ex:
+ except BaseException as ex:
  exception = ex
  raise
  finally:

tests/test_logging.py

@@ -20,10 +20,12 @@ async def test_logging_context(self):
  def query_saver(record):
  queries.put_nowait(record)
 
- with self.con.logger(query_saver):
+ log = LogCollector()
+
+ with self.con.query_logger(query_saver):
  self.assertEqual(len(self.con._query_loggers), 1)
  await self.con.execute("SELECT 1")
- with self.con.logger(LogCollector()) as log:
+ with self.con.query_logger(log):
  self.assertEqual(len(self.con._query_loggers), 2)
  await self.con.execute("SELECT 2")
 
@@ -36,7 +38,8 @@ def query_saver(record):
  self.assertEqual(len(self.con._query_loggers), 0)
 
  async def test_error_logging(self):
- with self.con.logger(LogCollector()) as log:
+ log = LogCollector()
+ with self.con.query_logger(log):
  with self.assertRaises(exceptions.UndefinedColumnError):
  await self.con.execute("SELECT x")