From 4c3091a23f826563318ac093562d8b731a6f62fe Mon Sep 17 00:00:00 2001
From: Serj Shevchenko <Serj.Shevchenko@gmail.com>
Date: Wed, 13 Oct 2021 00:23:29 +0300
Subject: [PATCH] change the way for logging integration

---
 .../opencensus/ext/logging/trace.py           | 11 +++---
 .../tests/test_logging_integration.py         | 34 ++++++++++++++++---
 contrib/opencensus-ext-logging/version.py     |  2 +-
 opencensus/log/__init__.py                    | 17 ++++++++++
 4 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/contrib/opencensus-ext-logging/opencensus/ext/logging/trace.py b/contrib/opencensus-ext-logging/opencensus/ext/logging/trace.py
index 13d95e846..86cb21e16 100644
--- a/contrib/opencensus-ext-logging/opencensus/ext/logging/trace.py
+++ b/contrib/opencensus-ext-logging/opencensus/ext/logging/trace.py
@@ -14,17 +14,20 @@
 
 import logging
 
-from opencensus.log import TraceLogger
+from opencensus.log import decorate_log_record_factory
 from opencensus.trace import integrations
 
 
 def trace_integration(tracer=None):
-    """Replace the global default logging class with `TraceLogger`.
+    """Customize LogRecord with opencensus trace data.
+    https://docs.python.org/3/howto/logging-cookbook.html#customizing-logrecord
 
-    Loggers created after the integration will produce `LogRecord`s
+    LogRecordFactory created after the integration will produce `LogRecord`s
     with extra traceId, spanId, and traceSampled attributes from the opencensus
     context.
     """
-    logging.setLoggerClass(TraceLogger)
+    logging.setLogRecordFactory(
+        decorate_log_record_factory(logging.getLogRecordFactory()),
+    )
     # pylint: disable=protected-access
     integrations.add_integration(integrations._Integrations.LOGGING)
diff --git a/contrib/opencensus-ext-logging/tests/test_logging_integration.py b/contrib/opencensus-ext-logging/tests/test_logging_integration.py
index 768f1a2c1..e38d47269 100644
--- a/contrib/opencensus-ext-logging/tests/test_logging_integration.py
+++ b/contrib/opencensus-ext-logging/tests/test_logging_integration.py
@@ -14,6 +14,10 @@
 
 import logging
 import unittest
+try:
+    from io import StringIO
+except ImportError:
+    from StringIO import StringIO
 
 from opencensus.trace import config_integration
 
@@ -21,13 +25,35 @@
 class TestLoggingIntegration(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        cls._old_logger_class = logging.getLoggerClass()
+        cls.log_stream = StringIO()
+        logging.basicConfig(
+            format="%(message)s traceId=%(traceId)s",
+            stream=cls.log_stream,
+            level=logging.INFO,
+        )
+        cls._old_logger_factory = logging.getLogRecordFactory()
 
     @classmethod
     def tearDownClass(cls):
-        logging.setLoggerClass(cls._old_logger_class)
+        logging.setLogRecordFactory(cls._old_logger_factory)
 
     def test_integration(self):
-        self.assertEqual(self._old_logger_class, logging.getLoggerClass())
+        self.assertEqual(self._old_logger_factory, logging.getLogRecordFactory())
         config_integration.trace_integrations(['logging'])
-        self.assertNotEqual(self._old_logger_class, logging.getLoggerClass())
+        self.assertNotEqual(self._old_logger_factory, logging.getLogRecordFactory())
+
+    def test_logger(self):
+        log_msg_before_integration = "catch logger_before_integration"
+        log1 = logging.getLogger("log1")
+
+        config_integration.trace_integrations(['logging'])
+
+        log_after_integration = "catch logger_after_integration"
+        log2 = logging.getLogger("log2")
+
+        log1.info(log_msg_before_integration)
+        log2.info(log_after_integration)
+
+        all_logs = self.log_stream.getvalue()
+        assert "{} traceId=".format(log_msg_before_integration) in all_logs
+        assert "{} traceId=".format(log_after_integration) in all_logs
diff --git a/contrib/opencensus-ext-logging/version.py b/contrib/opencensus-ext-logging/version.py
index 00ffd3fd7..63922d675 100644
--- a/contrib/opencensus-ext-logging/version.py
+++ b/contrib/opencensus-ext-logging/version.py
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = '0.1.1'
+__version__ = '0.2.0'
diff --git a/opencensus/log/__init__.py b/opencensus/log/__init__.py
index ce4ce5f57..7783a3a29 100644
--- a/opencensus/log/__init__.py
+++ b/opencensus/log/__init__.py
@@ -113,3 +113,20 @@ def makeRecord(self, *args, **kwargs):
                 kwargs['extra'] = extra
         _set_extra_attrs(extra)
         return super(TraceLogger, self).makeRecord(*args, **kwargs)
+
+
+def set_default_attr(obj, attr, def_value):
+    if not hasattr(obj, attr):
+        setattr(obj, attr, def_value)
+
+
+def decorate_log_record_factory(old_factory):
+    def _new_factory(*args, **kwargs):
+        record = old_factory(*args, **kwargs)
+        trace_id, span_id, sampling_decision = get_log_attrs()
+        set_default_attr(record, TRACE_ID_KEY, trace_id)
+        set_default_attr(record, SPAN_ID_KEY, span_id)
+        set_default_attr(record, SAMPLING_DECISION_KEY, sampling_decision)
+        return record
+
+    return _new_factory