Remove the query logger when the query_logger() block raises (#1373)

query_logger() removed its callback after a bare yield, so any exception
in the with-block, such as a failing query, left the logger attached to
the connection. Pooled connections then kept reporting later, unrelated
queries to it.

Author: kratos0718

asyncpg/connection.py

@@ -1932,8 +1932,10 @@ def __call__(self, record):
         .. versionadded:: 0.29.0
         """
         self.add_query_logger(callback)
-        yield
-        self.remove_query_logger(callback)
+        try:
+            yield
+        finally:
+            self.remove_query_logger(callback)
 
     @contextlib.contextmanager
     def _time_and_log(self, query, args, timeout):

tests/test_logging.py

@@ -49,3 +49,15 @@ async def test_error_logging(self):
             type(log.records[0].exception),
             exceptions.UndefinedColumnError
         )
+
+    async def test_logging_context_removed_on_error(self):
+        log = LogCollector()
+        with self.assertRaises(exceptions.UndefinedColumnError):
+            with self.con.query_logger(log):
+                await self.con.execute("SELECT x")
+
+        self.assertEqual(len(self.con._query_loggers), 0)
+
+        await self.con.execute("SELECT 1")
+        await asyncio.sleep(0)  # wait for logging
+        self.assertEqual(len(log.records), 1)