Add trace and span_id to logging async API (#5908)

Author: brunoais

google/cloud/logging/handlers/transports/background_thread.py

@@ -229,7 +229,8 @@ def _main_thread_terminated(self):
                 'Failed to send %d pending logs.' % (self._queue.qsize(),),
                 file=sys.stderr)
 
-    def enqueue(self, record, message, resource=None, labels=None):
+    def enqueue(self, record, message, resource=None, labels=None,
+                trace=None, span_id=None):
         """Queues a log entry to be written by the background thread.
 
         :type record: :class:`logging.LogRecord`
@@ -244,6 +245,13 @@ def enqueue(self, record, message, resource=None, labels=None):
 
         :type labels: dict
         :param labels: (Optional) Mapping of labels for the entry.
+
+        :type trace: str
+        :param trace: (optional) traceid to apply to the logging entry.
+
+        :type span_id: str
+        :param span_id: (optional) span_id within the trace for the log entry.
+                        Specify the trace parameter if span_id is set.
         """
         self._queue.put_nowait({
             'info': {
@@ -253,6 +261,8 @@ def enqueue(self, record, message, resource=None, labels=None):
             'severity': record.levelname,
             'resource': resource,
             'labels': labels,
+            'trace': trace,
+            'span_id': span_id,
         })
 
     def flush(self):
@@ -296,7 +306,8 @@ def __init__(self, client, name, grace_period=_DEFAULT_GRACE_PERIOD,
                               max_latency=max_latency)
         self.worker.start()
 
-    def send(self, record, message, resource=None, labels=None):
+    def send(self, record, message, resource=None, labels=None,
+             trace=None, span_id=None):
         """Overrides Transport.send().
 
         :type record: :class:`logging.LogRecord`
@@ -311,8 +322,16 @@ def send(self, record, message, resource=None, labels=None):
 
         :type labels: dict
         :param labels: (Optional) Mapping of labels for the entry.
+
+        :type trace: str
+        :param trace: (optional) traceid to apply to the logging entry.
+
+        :type span_id: str
+        :param span_id: (optional) span_id within the trace for the log entry.
+                        Specify the trace parameter if span_id is set.
         """
-        self.worker.enqueue(record, message, resource=resource, labels=labels)
+        self.worker.enqueue(record, message, resource=resource, labels=labels,
+                            trace=trace, span_id=span_id)
 
     def flush(self):
         """Submit any pending log records."""

google/cloud/logging/handlers/transports/base.py

@@ -22,7 +22,8 @@ class Transport(object):
     client and name object, and must override :meth:`send`.
     """
 
-    def send(self, record, message, resource=None, labels=None):
+    def send(self, record, message, resource=None, labels=None,
+             trace=None, span_id=None):
         """Transport send to be implemented by subclasses.
 
         :type record: :class:`logging.LogRecord`

google/cloud/logging/handlers/transports/sync.py

@@ -29,7 +29,8 @@ class SyncTransport(Transport):
     def __init__(self, client, name):
         self.logger = client.logger(name)
 
-    def send(self, record, message, resource=None, labels=None):
+    def send(self, record, message, resource=None, labels=None,
+             trace=None, span_id=None):
         """Overrides transport.send().
 
         :type record: :class:`logging.LogRecord`
@@ -49,4 +50,6 @@ def send(self, record, message, resource=None, labels=None):
         self.logger.log_struct(info,
                                severity=record.levelname,
                                resource=resource,
-                               labels=labels)
+                               labels=labels,
+                               trace=trace,
+                               span_id=span_id)

tests/unit/handlers/transports/test_background_thread.py

@@ -61,10 +61,55 @@ def test_send(self):
             python_logger_name, logging.INFO,
             None, None, message, None, None)
 
-        transport.send(record, message, _GLOBAL_RESOURCE, None)
+        transport.send(record, message, _GLOBAL_RESOURCE)
 
         transport.worker.enqueue.assert_called_once_with(
-            record, message, _GLOBAL_RESOURCE, None)
+            record, message, _GLOBAL_RESOURCE, None,
+            trace=None, span_id=None)
+
+    def test_trace_send(self):
+        from google.cloud.logging.logger import _GLOBAL_RESOURCE
+
+        client = _Client(self.PROJECT)
+        name = 'python_logger'
+
+        transport, _ = self._make_one(client, name)
+
+        python_logger_name = 'mylogger'
+        message = 'hello world'
+        trace = 'the-project/trace/longlogTraceid'
+
+        record = logging.LogRecord(
+            python_logger_name, logging.INFO,
+            None, None, message, None, None)
+
+        transport.send(record, message, _GLOBAL_RESOURCE, trace=trace)
+
+        transport.worker.enqueue.assert_called_once_with(
+            record, message, _GLOBAL_RESOURCE, None,
+            trace=trace, span_id=None)
+
+    def test_span_send(self):
+        from google.cloud.logging.logger import _GLOBAL_RESOURCE
+
+        client = _Client(self.PROJECT)
+        name = 'python_logger'
+
+        transport, _ = self._make_one(client, name)
+
+        python_logger_name = 'mylogger'
+        message = 'hello world'
+        span_id = 'the-project/trace/longlogTraceid/span/123456789012abbacdac'
+
+        record = logging.LogRecord(
+            python_logger_name, logging.INFO,
+            None, None, message, None, None)
+
+        transport.send(record, message, _GLOBAL_RESOURCE, span_id=span_id)
+
+        transport.worker.enqueue.assert_called_once_with(
+            record, message, _GLOBAL_RESOURCE, None,
+            trace=None, span_id=span_id)
 
     def test_flush(self):
         client = _Client(self.PROJECT)
@@ -381,13 +426,15 @@ def __init__(self):
         self.commit_count = None
 
     def log_struct(
-            self, info, severity=logging.INFO, resource=None, labels=None):
+            self, info, severity=logging.INFO, resource=None, labels=None,
+            trace=None, span_id=None):
         from google.cloud.logging.logger import _GLOBAL_RESOURCE
 
         assert resource is None
         resource = _GLOBAL_RESOURCE
 
-        self.log_struct_called_with = (info, severity, resource, labels)
+        self.log_struct_called_with = (info, severity, resource, labels,
+                                       trace, span_id)
         self.entries.append(info)
 
     def commit(self):

tests/unit/handlers/transports/test_sync.py

@@ -52,7 +52,8 @@ def test_send(self):
             'message': message,
             'python_logger': python_logger_name,
         }
-        EXPECTED_SENT = (EXPECTED_STRUCT, 'INFO', _GLOBAL_RESOURCE, None)
+        EXPECTED_SENT = (EXPECTED_STRUCT, 'INFO', _GLOBAL_RESOURCE, None,
+                         None, None)
         self.assertEqual(
             transport.logger.log_struct_called_with, EXPECTED_SENT)
 
@@ -64,8 +65,10 @@ def __init__(self, name):
         self.name = name
 
     def log_struct(self, message, severity=None,
-                   resource=_GLOBAL_RESOURCE, labels=None):
-        self.log_struct_called_with = (message, severity, resource, labels)
+                   resource=_GLOBAL_RESOURCE, labels=None,
+                   trace=None, span_id=None):
+        self.log_struct_called_with = (message, severity, resource, labels,
+                                       trace, span_id)
 
 
 class _Client(object):