Log query errors, clean up context managers

Author: dcwatson

asyncpg/connection.py

@@ -9,6 +9,7 @@
 import asyncpg
 import collections
 import collections.abc
+import contextlib
 import functools
 import itertools
 import inspect
@@ -226,10 +227,9 @@ def add_query_logger(self, callback):
         """Add a logger that will be called when queries are executed.
 
         :param callable callback:
-            A callable or a coroutine function receiving two arguments:
-            **connection**: a Connection the callback is registered with.
-            **query**: a LoggedQuery containing the query, args, timeout, and
-                       elapsed.
+            A callable or a coroutine function receiving one argument:
+            **record**: a LoggedQuery containing `query`, `args`, `timeout`,
+                        `elapsed`, `addr`, `params`, and `exception`.
 
         .. versionadded:: 0.29.0
         """
@@ -339,9 +339,8 @@ async def execute(self, query: str, *args, timeout: float=None) -> str:
         self._check_open()
 
         if not args:
-            with utils.timer() as t:
+            with self._time_and_log(query, args, timeout):
                 result = await self._protocol.query(query, timeout)
-            self._log_query(query, args, timeout, t.elapsed)
             return result
 
         _, status, _ = await self._execute(
@@ -1412,6 +1411,7 @@ def _cleanup(self):
         self._mark_stmts_as_closed()
         self._listeners.clear()
         self._log_listeners.clear()
+        self._query_loggers.clear()
         self._clean_tasks()
 
     def _clean_tasks(self):
@@ -1695,15 +1695,15 @@ async def _execute(
             )
         return result
 
+    @contextlib.contextmanager
     def logger(self, callback):
         """Context manager that adds `callback` to the list of query loggers,
         and removes it upon exit.
 
         :param callable callback:
-            A callable or a coroutine function receiving two arguments:
-            **connection**: a Connection the callback is registered with.
-            **query**: a LoggedQuery containing the query, args, timeout, and
-                       elapsed.
+            A callable or a coroutine function receiving one argument:
+            **record**: a LoggedQuery containing `query`, `args`, `timeout`,
+                        `elapsed`, `addr`, and `params`.
 
         Example:
 
@@ -1721,18 +1721,35 @@ def __call__(self, conn, record):
 
         .. versionadded:: 0.29.0
         """
-        return _LoggingContext(self, callback)
-
-    def _log_query(self, query, args, timeout, elapsed):
-        if not self._query_loggers:
-            return
-        con_ref = self._unwrap()
-        record = LoggedQuery(query, args, timeout, elapsed)
-        for cb in self._query_loggers:
-            if cb.is_async:
-                self._loop.create_task(cb.cb(con_ref, record))
-            else:
-                self._loop.call_soon(cb.cb, con_ref, record)
+        self.add_query_logger(callback)
+        yield callback
+        self.remove_query_logger(callback)
+
+    @contextlib.contextmanager
+    def _time_and_log(self, query, args, timeout):
+        start = time.monotonic()
+        exception = None
+        try:
+            yield
+        except Exception as ex:
+            exception = ex
+            raise
+        finally:
+            elapsed = time.monotonic() - start
+            record = LoggedQuery(
+                query=query,
+                args=args,
+                timeout=timeout,
+                elapsed=elapsed,
+                addr=self._addr,
+                params=self._params,
+                exception=exception,
+            )
+            for cb in self._query_loggers:
+                if cb.is_async:
+                    self._loop.create_task(cb.cb(record))
+                else:
+                    self._loop.call_soon(cb.cb, record)
 
     async def __execute(
         self,
@@ -1748,25 +1765,23 @@ async def __execute(
         executor = lambda stmt, timeout: self._protocol.bind_execute(
             stmt, args, '', limit, return_status, timeout)
         timeout = self._protocol._get_timeout(timeout)
-        with utils.timer() as t:
+        with self._time_and_log(query, args, timeout):
             result, stmt = await self._do_execute(
                 query,
                 executor,
                 timeout,
                 record_class=record_class,
                 ignore_custom_codec=ignore_custom_codec,
             )
-        self._log_query(query, args, timeout, t.elapsed)
         return result, stmt
 
     async def _executemany(self, query, args, timeout):
         executor = lambda stmt, timeout: self._protocol.bind_execute_many(
             stmt, args, '', timeout)
         timeout = self._protocol._get_timeout(timeout)
         with self._stmt_exclusive_section:
-            with utils.timer() as t:
+            with self._time_and_log(query, args, timeout):
                 result, _ = await self._do_execute(query, executor, timeout)
-        self._log_query(query, args, timeout, t.elapsed)
         return result
 
     async def _do_execute(
@@ -2401,25 +2416,10 @@ class _ConnectionProxy:
 
 LoggedQuery = collections.namedtuple(
     'LoggedQuery',
-    ['query', 'args', 'timeout', 'elapsed'])
+    ['query', 'args', 'timeout', 'elapsed', 'exception', 'addr', 'params'])
 LoggedQuery.__doc__ = 'Log record of an executed query.'
 
 
-class _LoggingContext:
-    __slots__ = ('_conn', '_cb')
-
-    def __init__(self, conn, callback):
-        self._conn = conn
-        self._cb = callback
-
-    def __enter__(self):
-        self._conn.add_query_logger(self._cb)
-        return self._cb
-
-    def __exit__(self, *exc_info):
-        self._conn.remove_query_logger(self._cb)
-
-
 ServerCapabilities = collections.namedtuple(
     'ServerCapabilities',
     ['advisory_locks', 'notifications', 'plpgsql', 'sql_reset',

asyncpg/utils.py

@@ -6,7 +6,6 @@
 
 
 import re
-import time
 
 
 def _quote_ident(ident):
@@ -44,28 +43,3 @@ async def _mogrify(conn, query, args):
     # Finally, replace $n references with text values.
     return re.sub(
         r'\$(\d+)\b', lambda m: textified[int(m.group(1)) - 1], query)
-
-
-class timer:
-    __slots__ = ('start', 'elapsed')
-
-    def __init__(self):
-        self.start = time.monotonic()
-        self.elapsed = None
-
-    @property
-    def current(self):
-        return time.monotonic() - self.start
-
-    def restart(self):
-        self.start = time.monotonic()
-
-    def stop(self):
-        self.elapsed = self.current
-
-    def __enter__(self):
-        self.restart()
-        return self
-
-    def __exit__(self, *exc):
-        self.stop()

tests/test_logging.py

@@ -1,30 +1,48 @@
 import asyncio
 
 from asyncpg import _testbase as tb
+from asyncpg import exceptions
+
+
+class LogCollector:
+    def __init__(self):
+        self.records = []
+
+    def __call__(self, record):
+        self.records.append(record)
 
 
 class TestQueryLogging(tb.ConnectedTestCase):
 
     async def test_logging_context(self):
         queries = asyncio.Queue()
 
-        def query_saver(conn, record):
+        def query_saver(record):
             queries.put_nowait(record)
 
-        class QuerySaver:
-            def __init__(self):
-                self.queries = []
-
-            def __call__(self, conn, record):
-                self.queries.append(record.query)
-
         with self.con.logger(query_saver):
             self.assertEqual(len(self.con._query_loggers), 1)
-            with self.con.logger(QuerySaver()) as log:
+            await self.con.execute("SELECT 1")
+            with self.con.logger(LogCollector()) as log:
                 self.assertEqual(len(self.con._query_loggers), 2)
-                await self.con.execute("SELECT 1")
-
-        record = await queries.get()
-        self.assertEqual(record.query, "SELECT 1")
-        self.assertEqual(log.queries, ["SELECT 1"])
+                await self.con.execute("SELECT 2")
+
+        r1 = await queries.get()
+        r2 = await queries.get()
+        self.assertEqual(r1.query, "SELECT 1")
+        self.assertEqual(r2.query, "SELECT 2")
+        self.assertEqual(len(log.records), 1)
+        self.assertEqual(log.records[0].query, "SELECT 2")
         self.assertEqual(len(self.con._query_loggers), 0)
+
+    async def test_error_logging(self):
+        with self.con.logger(LogCollector()) as log:
+            with self.assertRaises(exceptions.UndefinedColumnError):
+                await self.con.execute("SELECT x")
+
+        await asyncio.sleep(0)  # wait for logging
+        self.assertEqual(len(log.records), 1)
+        self.assertEqual(
+            type(log.records[0].exception),
+            exceptions.UndefinedColumnError
+        )