Merge branch 'tickets/DM-41730'

python/activator/logger.py

@@ -235,7 +235,7 @@ def _context(self):
     def _context(self, value):
         self._store.context = value
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, name, level, fn, lno, msg, args, exc_info, func=None, sinfo=None, **kwargs):
         """Create a log record from the provided arguments.
 
         See `logging.setLogRecordFactory` for the parameters.
@@ -247,9 +247,14 @@ def __call__(self, *args, **kwargs):
             maps strings to arbitrary values, as determined by any enclosing
             calls to `add_context`.
         """
-        record = self._old_factory(*args, **kwargs)
+        record = self._old_factory(name, level, fn, lno, msg, args, exc_info, func, sinfo, **kwargs)
         # _context is mutable; make sure record can't be changed after the fact.
         record.logging_context = self._context.copy()
+        if exc_info is not None:
+            _, ex, _ = exc_info  # Only care about the exception object passed to the logger
+            if hasattr(ex, "logging_context"):
+                # Context at the point where the exception was raised takes precedence.
+                record.logging_context.update(ex.logging_context)
         return record
 
     @contextmanager
@@ -258,7 +263,8 @@ def add_context(self, **context):
         inside it.
 
         This manager adds key-value pairs to the ``logging_context`` mapping in
-        this factory's log records.
+        this factory's log records. It also adds a mapping of that name to any
+        exceptions that escape.
 
         Parameters
         ----------
@@ -285,6 +291,11 @@ def add_context(self, **context):
         try:
             self._context.update(**context)
             yield
+        except BaseException as e:
+            # In logging, inner context overrules outer context. Need the same for exceptions.
+            inner_context = e.logging_context if hasattr(e, "logging_context") else {}
+            e.logging_context = self._context.copy() | inner_context
+            raise
         finally:
             # This replacement is safe because self._context cannot have been
             # changed by other threads. Changes can only have been made by

tests/test_logger.py

@@ -315,14 +315,64 @@ def test_exception_handling(self):
                 try:
                     with factory.add_context(pid=42):
                         raise RuntimeError("Something failed")
-                except RuntimeError:
+                except RuntimeError as e:
+                    self.assertEqual(e.logging_context, {"color": "blue", "pid": 42})
                     # pid=42 should have been removed here
                     self.log.error("Exception caught")
+                    # Exception should hang on to old state
+                    self.log.exception("Exception detected")
+                    self.log.warning("Implied exception", exc_info=True)
+                    self.log.critical("Explicit exception", exc_info=e)
 
-            self.assertEqual(len(recorder.records), 2)
+            self.assertEqual(len(recorder.records), 5)
             self.assertEqual([dict(rec.logging_context) for rec in recorder.records],
                              [{"color": "blue"},
                               {"color": "blue"},
+                              {"color": "blue", "pid": 42},
+                              {"color": "blue", "pid": 42},
+                              {"color": "blue", "pid": 42},
+                              ])
+
+    def test_overwriting_exception_handling(self):
+        factory = logging.getLogRecordFactory()
+        with self.assertLogs(self.log, "DEBUG") as recorder:
+            try:
+                with factory.add_context(color="blue"):
+                    with factory.add_context(color="red"):
+                        raise RuntimeError("Something failed")
+            except RuntimeError as e:
+                # Was original context preserved while escaping multiple managers?
+                self.assertEqual(e.logging_context, {"color": "red"})
+                # Is original context used *only* when logging the exception?
+                self.log.error("Exception follows")
+                self.log.exception("Exceptional exception")
+                self.log.warning("Nothing to see here")
+
+            self.assertEqual(len(recorder.records), 3)
+            self.assertEqual([dict(rec.logging_context) for rec in recorder.records],
+                             [{},
+                              {"color": "red"},
+                              {},
+                              ])
+
+    def test_bare_exception_handling(self):
+        with self.assertLogs(self.log, "DEBUG") as recorder:
+            self.log.info("This is a log!")
+            try:
+                raise RuntimeError("Something failed")
+            except RuntimeError as e:
+                self.assertFalse(hasattr(e, "logging_context"))
+                # Logger shouldn't choke on lack of logging_context
+                self.log.exception("Exception detected")
+                self.log.warning("Implied exception", exc_info=True)
+                self.log.critical("Explicit exception", exc_info=e)
+
+            self.assertEqual(len(recorder.records), 4)
+            self.assertEqual([dict(rec.logging_context) for rec in recorder.records],
+                             [{},
+                              {},
+                              {},
+                              {},
                               ])
 
     # This decorator works with scons/unittest as well