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