Support 'trace' attribute of log entries (#5878)

salrashid123 · tseaver · commit 396d079c2159 · 2018-09-01T12:46:32.000-04:00
See: #5505
diff --git a/google/cloud/logging/entries.py b/google/cloud/logging/entries.py
@@ -76,9 +76,13 @@ class _BaseEntry(object):
 
     :type resource: :class:`~google.cloud.logging.resource.Resource`
     :param resource: (Optional) Monitored resource of the entry
+
+    :type trace: str
+    :param trace: (optional) traceid to apply to the entry.
     """
     def __init__(self, payload, logger, insert_id=None, timestamp=None,
-                 labels=None, severity=None, http_request=None, resource=None):
+                 labels=None, severity=None, http_request=None, resource=None,
+                 trace=None):
         self.payload = payload
         self.logger = logger
         self.insert_id = insert_id
@@ -87,6 +91,7 @@ def __init__(self, payload, logger, insert_id=None, timestamp=None,
         self.severity = severity
         self.http_request = http_request
         self.resource = resource
+        self.trace = trace
 
     @classmethod
     def from_api_repr(cls, resource, client, loggers=None):
@@ -123,6 +128,7 @@ def from_api_repr(cls, resource, client, loggers=None):
         labels = resource.get('labels')
         severity = resource.get('severity')
         http_request = resource.get('httpRequest')
+        trace = resource.get('trace')
 
         monitored_resource_dict = resource.get('resource')
         monitored_resource = None
@@ -131,7 +137,7 @@ def from_api_repr(cls, resource, client, loggers=None):
 
         return cls(payload, logger, insert_id=insert_id, timestamp=timestamp,
                    labels=labels, severity=severity, http_request=http_request,
-                   resource=monitored_resource)
+                   resource=monitored_resource, trace=trace)
 
 
 class TextEntry(_BaseEntry):
@@ -185,15 +191,19 @@ class ProtobufEntry(_BaseEntry):
 
     :type resource: :class:`~google.cloud.logging.resource.Resource`
     :param resource: (Optional) Monitored resource of the entry
+
+    :type trace: str
+    :param trace: (optional) traceid to apply to the entry.
     """
     _PAYLOAD_KEY = 'protoPayload'
 
     def __init__(self, payload, logger, insert_id=None, timestamp=None,
-                 labels=None, severity=None, http_request=None, resource=None):
+                 labels=None, severity=None, http_request=None, resource=None,
+                 trace=None):
         super(ProtobufEntry, self).__init__(
             payload, logger, insert_id=insert_id, timestamp=timestamp,
             labels=labels, severity=severity, http_request=http_request,
-            resource=resource)
+            resource=resource, trace=trace)
         if isinstance(self.payload, any_pb2.Any):
             self.payload_pb = self.payload
             self.payload = None
diff --git a/google/cloud/logging/logger.py b/google/cloud/logging/logger.py
@@ -96,7 +96,7 @@ def batch(self, client=None):
     def _make_entry_resource(self, text=None, info=None, message=None,
                              labels=None, insert_id=None, severity=None,
                              http_request=None, timestamp=None,
-                             resource=_GLOBAL_RESOURCE):
+                             resource=_GLOBAL_RESOURCE, trace=None):
         """Return a log entry resource of the appropriate type.
 
         Helper for :meth:`log_text`, :meth:`log_struct`, and :meth:`log_proto`.
@@ -131,6 +131,9 @@ def _make_entry_resource(self, text=None, info=None, message=None,
         :type resource: :class:`~google.cloud.logging.resource.Resource`
         :param resource: (Optional) Monitored resource of the entry
 
+        :type trace: str
+        :param trace: (optional) traceid to apply to the entry.
+
         :rtype: dict
         :returns: The JSON resource created.
         """
@@ -172,11 +175,14 @@ def _make_entry_resource(self, text=None, info=None, message=None,
         if timestamp is not None:
             entry['timestamp'] = _datetime_to_rfc3339(timestamp)
 
+        if trace is not None:
+            entry['trace'] = trace
+
         return entry
 
     def log_text(self, text, client=None, labels=None, insert_id=None,
                  severity=None, http_request=None, timestamp=None,
-                 resource=_GLOBAL_RESOURCE):
+                 resource=_GLOBAL_RESOURCE, trace=None):
         """API call:  log a text message via a POST request
 
         See
@@ -203,22 +209,26 @@ def log_text(self, text, client=None, labels=None, insert_id=None,
         :param http_request: (optional) info about HTTP request associated with
                              the entry
 
+        :type timestamp: :class:`datetime.datetime`
+        :param timestamp: (optional) timestamp of event being logged.
+
         :type resource: :class:`~google.cloud.logging.resource.Resource`
         :param resource: Monitored resource of the entry, defaults
                          to the global resource type.
 
-        :type timestamp: :class:`datetime.datetime`
-        :param timestamp: (optional) timestamp of event being logged.
+        :type trace: str
+        :param trace: (optional) traceid to apply to the entry.
         """
         client = self._require_client(client)
         entry_resource = self._make_entry_resource(
             text=text, labels=labels, insert_id=insert_id, severity=severity,
-            http_request=http_request, timestamp=timestamp, resource=resource)
+            http_request=http_request, timestamp=timestamp, resource=resource,
+            trace=trace)
         client.logging_api.write_entries([entry_resource])
 
     def log_struct(self, info, client=None, labels=None, insert_id=None,
                    severity=None, http_request=None, timestamp=None,
-                   resource=_GLOBAL_RESOURCE):
+                   resource=_GLOBAL_RESOURCE, trace=None):
         """API call:  log a structured message via a POST request
 
         See
@@ -245,22 +255,26 @@ def log_struct(self, info, client=None, labels=None, insert_id=None,
         :param http_request: (optional) info about HTTP request associated with
                              the entry.
 
+        :type timestamp: :class:`datetime.datetime`
+        :param timestamp: (optional) timestamp of event being logged.
+
         :type resource: :class:`~google.cloud.logging.resource.Resource`
         :param resource: Monitored resource of the entry, defaults
                          to the global resource type.
 
-        :type timestamp: :class:`datetime.datetime`
-        :param timestamp: (optional) timestamp of event being logged.
+        :type trace: str
+        :param trace: (optional) traceid to apply to the entry.
         """
         client = self._require_client(client)
         entry_resource = self._make_entry_resource(
             info=info, labels=labels, insert_id=insert_id, severity=severity,
-            http_request=http_request, timestamp=timestamp, resource=resource)
+            http_request=http_request, timestamp=timestamp, resource=resource,
+            trace=trace)
         client.logging_api.write_entries([entry_resource])
 
     def log_proto(self, message, client=None, labels=None, insert_id=None,
                   severity=None, http_request=None, timestamp=None,
-                  resource=_GLOBAL_RESOURCE):
+                  resource=_GLOBAL_RESOURCE, trace=None):
         """API call:  log a protobuf message via a POST request
 
         See
@@ -287,18 +301,21 @@ def log_proto(self, message, client=None, labels=None, insert_id=None,
         :param http_request: (optional) info about HTTP request associated with
                              the entry.
 
+        :type timestamp: :class:`datetime.datetime`
+        :param timestamp: (optional) timestamp of event being logged.
+
         :type resource: :class:`~google.cloud.logging.resource.Resource`
         :param resource: Monitored resource of the entry, defaults
                          to the global resource type.
 
-        :type timestamp: :class:`datetime.datetime`
-        :param timestamp: (optional) timestamp of event being logged.
+        :type trace: str
+        :param trace: (optional) traceid to apply to the entry.
         """
         client = self._require_client(client)
         entry_resource = self._make_entry_resource(
             message=message, labels=labels, insert_id=insert_id,
             severity=severity, http_request=http_request, timestamp=timestamp,
-            resource=resource)
+            resource=resource, trace=trace)
         client.logging_api.write_entries([entry_resource])
 
     def delete(self, client=None):
@@ -392,7 +409,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             self.commit()
 
     def log_text(self, text, labels=None, insert_id=None, severity=None,
-                 http_request=None, timestamp=None, resource=_GLOBAL_RESOURCE):
+                 http_request=None, timestamp=None, resource=_GLOBAL_RESOURCE,
+                 trace=None):
         """Add a text entry to be logged during :meth:`commit`.
 
         :type text: str
@@ -420,14 +438,17 @@ def log_text(self, text, labels=None, insert_id=None, severity=None,
                          resource of the batch is used for this entry. If
                          both this resource and the Batch resource are None,
                          the API will return an error.
+
+        :type trace: str
+        :param trace: (optional) traceid to apply to the entry.
         """
         self.entries.append(
             ('text', text, labels, insert_id, severity, http_request,
-             timestamp, resource))
+             timestamp, resource, trace))
 
     def log_struct(self, info, labels=None, insert_id=None, severity=None,
                    http_request=None, timestamp=None,
-                   resource=_GLOBAL_RESOURCE):
+                   resource=_GLOBAL_RESOURCE, trace=None):
         """Add a struct entry to be logged during :meth:`commit`.
 
         :type info: dict
@@ -455,14 +476,17 @@ def log_struct(self, info, labels=None, insert_id=None, severity=None,
                          resource of the batch is used for this entry. If
                          both this resource and the Batch resource are None,
                          the API will return an error.
+
+        :type trace: str
+        :param trace: (optional) traceid to apply to the entry.
         """
         self.entries.append(
             ('struct', info, labels, insert_id, severity, http_request,
-             timestamp, resource))
+             timestamp, resource, trace))
 
     def log_proto(self, message, labels=None, insert_id=None, severity=None,
                   http_request=None, timestamp=None,
-                  resource=_GLOBAL_RESOURCE):
+                  resource=_GLOBAL_RESOURCE, trace=None):
         """Add a protobuf entry to be logged during :meth:`commit`.
 
         :type message: protobuf message
@@ -490,10 +514,13 @@ def log_proto(self, message, labels=None, insert_id=None, severity=None,
                          resource of the batch is used for this entry. If
                          both this resource and the Batch resource are None,
                          the API will return an error.
+
+        :type trace: str
+        :param trace: (optional) traceid to apply to the entry.
         """
         self.entries.append(
             ('proto', message, labels, insert_id, severity, http_request,
-             timestamp, resource))
+             timestamp, resource, trace))
 
     def commit(self, client=None):
         """Send saved log entries as a single API call.
@@ -517,7 +544,7 @@ def commit(self, client=None):
 
         entries = []
         for (entry_type, entry, labels, iid, severity, http_req,
-             timestamp, resource) in self.entries:
+             timestamp, resource, trace) in self.entries:
             if entry_type == 'text':
                 info = {'textPayload': entry}
             elif entry_type == 'struct':
@@ -544,6 +571,8 @@ def commit(self, client=None):
                 info['httpRequest'] = http_req
             if timestamp is not None:
                 info['timestamp'] = _datetime_to_rfc3339(timestamp)
+            if trace is not None:
+                info['trace'] = trace
             entries.append(info)
 
         client.logging_api.write_entries(entries, **kwargs)
diff --git a/tests/unit/test_entries.py b/tests/unit/test_entries.py
@@ -68,6 +68,7 @@ def test_ctor_defaults(self):
         self.assertIsNone(entry.severity)
         self.assertIsNone(entry.http_request)
         self.assertIsNone(entry.resource)
+        self.assertIsNone(entry.trace)
 
     def test_ctor_explicit(self):
         import datetime
@@ -87,6 +88,7 @@ def test_ctor_explicit(self):
             'status': STATUS,
         }
         resource = Resource(type='global', labels={})
+        TRACE = '12345678-1234-5678-1234-567812345678'
 
         logger = _Logger(self.LOGGER_NAME, self.PROJECT)
         entry = self._make_one(PAYLOAD, logger,
@@ -95,7 +97,8 @@ def test_ctor_explicit(self):
                                labels=LABELS,
                                severity=SEVERITY,
                                http_request=REQUEST,
-                               resource=resource)
+                               resource=resource,
+                               trace=TRACE)
         self.assertEqual(entry.payload, PAYLOAD)
         self.assertIs(entry.logger, logger)
         self.assertEqual(entry.insert_id, IID)
@@ -106,6 +109,7 @@ def test_ctor_explicit(self):
         self.assertEqual(entry.http_request['requestUrl'], URI)
         self.assertEqual(entry.http_request['status'], STATUS)
         self.assertEqual(entry.resource, resource)
+        self.assertEqual(entry.trace, TRACE)
 
     def test_from_api_repr_missing_data_no_loggers(self):
         client = _Client(self.PROJECT)
@@ -122,6 +126,7 @@ def test_from_api_repr_missing_data_no_loggers(self):
         self.assertIsNone(entry.timestamp)
         self.assertIsNone(entry.severity)
         self.assertIsNone(entry.http_request)
+        self.assertIsNone(entry.trace)
         logger = entry.logger
         self.assertIsInstance(logger, _Logger)
         self.assertIs(logger.client, client)
@@ -154,6 +159,7 @@ def test_from_api_repr_w_loggers_no_logger_match(self):
             }
         )
         STATUS = '500'
+        TRACE = '12345678-1234-5678-1234-567812345678'
         API_REPR = {
             'dummyPayload': PAYLOAD,
             'logName': LOG_NAME,
@@ -167,6 +173,7 @@ def test_from_api_repr_w_loggers_no_logger_match(self):
                 'status': STATUS,
             },
             'resource': RESOURCE._to_dict(),
+            'trace': TRACE
         }
         loggers = {}
         entry = klass.from_api_repr(API_REPR, client, loggers=loggers)
@@ -184,6 +191,7 @@ def test_from_api_repr_w_loggers_no_logger_match(self):
         self.assertEqual(logger.name, self.LOGGER_NAME)
         self.assertEqual(loggers, {LOG_NAME: logger})
         self.assertEqual(entry.resource, RESOURCE)
+        self.assertEqual(entry.trace, TRACE)
 
     def test_from_api_repr_w_loggers_w_logger_match(self):
         from datetime import datetime
@@ -196,12 +204,14 @@ def test_from_api_repr_w_loggers_w_logger_match(self):
         TIMESTAMP = _datetime_to_rfc3339_w_nanos(NOW)
         LOG_NAME = 'projects/%s/logs/%s' % (self.PROJECT, self.LOGGER_NAME)
         LABELS = {'foo': 'bar', 'baz': 'qux'}
+        TRACE = '12345678-1234-5678-1234-567812345678'
         API_REPR = {
             'dummyPayload': PAYLOAD,
             'logName': LOG_NAME,
             'insertId': IID,
             'timestamp': TIMESTAMP,
             'labels': LABELS,
+            'trace': TRACE
         }
         LOGGER = object()
         loggers = {LOG_NAME: LOGGER}
@@ -211,6 +221,7 @@ def test_from_api_repr_w_loggers_w_logger_match(self):
         self.assertEqual(entry.insert_id, IID)
         self.assertEqual(entry.timestamp, NOW)
         self.assertEqual(entry.labels, LABELS)
+        self.assertEqual(entry.trace, TRACE)
         self.assertIs(entry.logger, LOGGER)
 
 
@@ -239,6 +250,7 @@ def test_constructor_basic(self):
         self.assertIsNone(pb_entry.labels)
         self.assertIsNone(pb_entry.severity)
         self.assertIsNone(pb_entry.http_request)
+        self.assertIsNone(pb_entry.trace)
 
     def test_constructor_with_any(self):
         from google.protobuf.any_pb2 import Any
@@ -253,6 +265,7 @@ def test_constructor_with_any(self):
         self.assertIsNone(pb_entry.labels)
         self.assertIsNone(pb_entry.severity)
         self.assertIsNone(pb_entry.http_request)
+        self.assertIsNone(pb_entry.trace)
 
     def test_parse_message(self):
         import json
diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py
