change the way for logging integration

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)

contrib/opencensus-ext-logging/tests/test_logging_integration.py

@@ -14,20 +14,46 @@
 
 import logging
 import unittest
+try:
+    from io import StringIO
+except ImportError:
+    from StringIO import StringIO
 
 from opencensus.trace import config_integration
 
 
 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

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'

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