Add context manager to set correlation id.

.circleci/config.yml

@@ -18,16 +18,13 @@ jobs:
     docker:
       - image: cimg/python:3.8
     steps:
-      - attach_workspace:
-          at: .
+      - checkout
       - run:
           name: Run unit tests
           command: |
             py=$(python3 --version | sed -r 's|.* 3\.([0-9]+)\..*|3.\1|')
-            virtualenv -p python3 venv$py
-            . venv$py/bin/activate
-            pip3 install $(ls ./dist/*.whl)[test,all]
-            python3 -m reconplogger_tests coverage xml coverage_py$py.xml
+            pip3 install .[test,all]
+            TEST_COVERAGE_XML=coverage_py$py.xml ./setup.py test_coverage
       - persist_to_workspace:
           root: .
           paths:
@@ -67,6 +64,17 @@ jobs:
                 --flags py$py \
                 --file coverage_py$py.xml
             done
+  test-py310-installed:
+    docker:
+      - image: cimg/python:3.10
+    steps:
+      - attach_workspace:
+          at: .
+      - run:
+          name: Run unit tests
+          command: |
+            pip3 install $(ls ./dist/*.whl)[test,all]
+            python3 -m reconplogger_tests
   publish-pypi:
     docker:
       - image: cimg/python:3.10
@@ -87,31 +95,33 @@ workflows:
           filters: &tagfilter
             tags:
               only: /^v\d+\.\d+\.\d+.*$/
-      - test-py38: &testreq
+      - test-py38:
           <<: *buildreq
-          requires:
-            - build
       - test-py312:
-          <<: *testreq
+          <<: *buildreq
       - test-py311:
-          <<: *testreq
+          <<: *buildreq
       - test-py310:
-          <<: *testreq
+          <<: *buildreq
       - test-py39:
-          <<: *testreq
+          <<: *buildreq
       - codecov:
-          <<: *testreq
           requires:
             - test-py312
             - test-py311
             - test-py310
             - test-py39
             - test-py38
+      - test-py310-installed:
+          <<: *buildreq
+          requires:
+            - build
       - publish-pypi:
           filters:
             branches:
               ignore: /.*/
             <<: *tagfilter
           context: pypi-upload-context
           requires:
+            - build
             - codecov

README.rst

@@ -15,12 +15,13 @@ ease the standardization of logging within omni:us. The main design decision of
 reconplogger is to allow total freedom to reconfigure loggers without hard
 coding anything.
 
-The package contains essentially three things:
+The package contains essentially the following things:
 
 - A default logging configuration.
 - A function for loading logging configuration for regular python code.
 - A function for loading logging configuration for flask-based microservices.
 - An inheritable class to add a logger property.
+- A context manager to set and get the correlation id.
 - Lower level functions for:
 
   - Loading logging configuration from any of: config file, environment variable, or default.

reconplogger.py

@@ -1,8 +1,10 @@
-
 import os
 import yaml
 import logging
 import logging.config
+from contextlib import contextmanager
+from contextvars import ContextVar
+from importlib.util import find_spec
 from logging import CRITICAL, ERROR, WARNING, INFO, DEBUG, NOTSET
 from typing import Optional, Union
 import uuid
@@ -264,6 +266,9 @@ def logger_setup(
             if not isinstance(handler, logging.FileHandler):
                 handler.setLevel(level)
 
+    # Add correlation id filter
+    logger.addFilter(_CorrelationIdLoggingFilter())
+
     # Log configured done and test logger
     if init_messages:
         logger.info('reconplogger (v'+__version__+') logger configured.')
@@ -338,13 +343,8 @@ def _flask_logging_after_request(response):
         return response
     flask_app.after_request_funcs.setdefault(None, []).append(_flask_logging_after_request)
 
-    # Add logging filter to augment the logs
-    class FlaskLoggingFilter(logging.Filter):
-        def filter(self, record):
-            if has_request_context():
-                record.correlation_id = g.correlation_id
-            return True
-    flask_app.logger.addFilter(FlaskLoggingFilter())
+    # Add correlation id filter
+    flask_app.logger.addFilter(_CorrelationIdLoggingFilter())
 
     # Setup werkzeug logger at least at WARNING level in case its server is used
     # since it also logs at INFO level after each request creating redundancy
@@ -365,13 +365,18 @@ def get_correlation_id() -> str:
         ImportError: When flask package not available.
         RuntimeError: When run outside an application context or if flask app has not been setup.
     """
-    from flask import g
+    correlation_id = current_correlation_id.get()
+    if correlation_id is not None:
+        return correlation_id
+    if find_spec("flask") is None:
+        raise RuntimeError("get_correlation_id used outside correlation_id_context.")
+
     try:
-        has_correlation_id = hasattr(g, 'correlation_id')
+        has_correlation_id = hasattr(g, "correlation_id")
     except RuntimeError:
-        raise RuntimeError('get_correlation_id only intended to be used inside an application context.')
+        raise RuntimeError("get_correlation_id used outside correlation_id_context or flask app context.")
     if not has_correlation_id:
-        raise RuntimeError('correlation_id not found in flask.g, probably flask app not yet setup.')
+        raise RuntimeError("correlation_id not found in flask.g, probably flask app not yet setup.")
     return g.correlation_id
 
 
@@ -390,6 +395,36 @@ def set_correlation_id(correlation_id: str):
     g.correlation_id = str(correlation_id)  # pylint: disable=assigning-non-slot
 
 
+current_correlation_id: ContextVar[Optional[str]] = ContextVar('current_correlation_id', default=None)
+
+
+@contextmanager
+def correlation_id_context(correlation_id: Optional[str]):
+    """Context manager to set the correlation id for the current application context.
+
+    Use as `with correlation_id_context(correlation_id): ...`. Calls to
+    `get_correlation_id()` will return the correlation id set for the context.
+
+    Args:
+        correlation_id: The correlation id to set in the context.
+    """
+    token = current_correlation_id.set(correlation_id)
+    try:
+        yield
+    finally:
+        current_correlation_id.reset(token)
+
+
+class _CorrelationIdLoggingFilter(logging.Filter):
+    def filter(self, record):
+        correlation_id = current_correlation_id.get()
+        if correlation_id is not None:
+            record.correlation_id = correlation_id
+        elif find_spec("flask") and has_request_context():
+            record.correlation_id = g.correlation_id
+        return True
+
+
 class RLoggerProperty:
     """Class designed to be inherited by other classes to add an rlogger property."""
 

reconplogger_tests.py

@@ -10,8 +10,11 @@
 import threading
 import unittest
 import uuid
+from contextlib import ExitStack, contextmanager
+from io import StringIO
 from unittest.mock import patch
 from testfixtures import LogCapture, compare, Comparison
+from typing import Iterator
 
 try:
     from flask import Flask, request
@@ -24,6 +27,16 @@
     requests = None
 
 
+@contextmanager
+def capture_logs(logger: logging.Logger) -> Iterator[StringIO]:
+    with ExitStack() as stack:
+        captured = StringIO()
+        for handler in logger.handlers:
+            if isinstance(handler, logging.StreamHandler):
+                stack.enter_context(patch.object(handler, "stream", captured))
+        yield captured
+
+
 class TestReconplogger(unittest.TestCase):
 
     def setUp(self):
@@ -153,6 +166,23 @@ def test_logger_setup_invalid_level(self):
         with self.assertRaises(ValueError):
             reconplogger.logger_setup(level=True, reload=True)
 
+    @patch.dict(os.environ, {'LOGGER_NAME': 'json_logger'})
+    def test_correlation_id_context(self):
+        logger = reconplogger.logger_setup()
+        correlation_id = str(uuid.uuid4())
+        with reconplogger.correlation_id_context(correlation_id):
+            self.assertEqual(correlation_id, reconplogger.get_correlation_id())
+            with capture_logs(logger) as logs:
+                logger.error('error message')
+                self.assertIn(correlation_id, logs.getvalue())
+
+    def test_get_correlation_id_outside_of_context(self):
+        with patch("reconplogger.find_spec", return_value=None):
+            self.assertIsNone(reconplogger.find_spec("flask"))
+            with self.assertRaises(RuntimeError) as ctx:
+                reconplogger.get_correlation_id()
+            self.assertIn("used outside correlation_id_context", str(ctx.exception))
+
     @unittest.skipIf(not Flask, "flask package is required")
     @patch.dict(os.environ, {
         'RECONPLOGGER_CFG': 'reconplogger_default_cfg',
@@ -315,37 +345,5 @@ def run_tests():
         sys.exit(True)
 
 
-def reimport_reconplogger():
-    del sys.modules['reconplogger']
-    if requests:
-        requests.sessions.Session.request = requests.sessions.Session.request_orig
-    import reconplogger
-
-
-def run_test_coverage():
-    try:
-        import coverage
-    except:
-        print('error: coverage package not found, run_test_coverage requires it.')
-        sys.exit(True)
-    cov = coverage.Coverage(source=['reconplogger'])
-    cov.start()
-    reimport_reconplogger()
-    run_tests()
-    cov.stop()
-    cov.save()
-    cov.report()
-    if 'xml' in sys.argv:
-        outfile = sys.argv[sys.argv.index('xml')+1]
-        cov.xml_report(outfile=outfile)
-        print('\nSaved coverage report to '+outfile+'.')
-    else:
-        cov.html_report(directory='htmlcov')
-        print('\nSaved html coverage report to htmlcov directory.')
-
-
 if __name__ == '__main__':
-    if 'coverage' in sys.argv:
-        run_test_coverage()
-    else:
-        run_tests()
+    run_tests()

setup.py

@@ -1,11 +1,12 @@
 #!/usr/bin/env python3
 
 from setuptools import setup, Command
+import os
 import re
 import sys
+import unittest
 
 
-NAME_TESTS = next(filter(lambda x: x.startswith('test_suite = '), open('setup.cfg').readlines())).strip().split()[-1]
 LONG_DESCRIPTION = re.sub(':class:|:func:|:ref:', '', open('README.rst').read())
 CMDCLASS = {}
 
@@ -17,7 +18,26 @@ class CoverageCommand(Command):
     def initialize_options(self): pass
     def finalize_options(self): pass
     def run(self):
-        __import__(NAME_TESTS).run_test_coverage()
+        try:
+            import coverage
+        except ImportError:
+            print('error: coverage package not found, run_test_coverage requires it.')
+            sys.exit(True)
+        cov = coverage.Coverage(source=['reconplogger'])
+        cov.start()
+        tests = unittest.defaultTestLoader.loadTestsFromName('reconplogger_tests')
+        if not unittest.TextTestRunner(verbosity=2).run(tests).wasSuccessful():
+            sys.exit(True)
+        cov.stop()
+        cov.save()
+        cov.report()
+        if 'TEST_COVERAGE_XML' in os.environ:
+            outfile = os.environ['TEST_COVERAGE_XML']
+            cov.xml_report(outfile=outfile)
+            print('\nSaved coverage report to '+outfile+'.')
+        else:
+            cov.html_report(directory='htmlcov')
+            print('\nSaved html coverage report to htmlcov directory.')
 
 CMDCLASS['test_coverage'] = CoverageCommand