Send trace context with logs from web applications

logging/google/cloud/logging/handlers/_helpers.py

@@ -17,6 +17,13 @@
 import math
 import json
 
+try:
+    import flask
+except ImportError:
+    flask = None
+
+from google.cloud.logging.handlers.middleware.request import get_request
+
 
 def format_stackdriver_json(record, message):
     """Helper to format a LogRecord in in Stackdriver fluentd format.
@@ -37,3 +44,46 @@ def format_stackdriver_json(record, message):
     }
 
     return json.dumps(payload)
+
+
+def get_trace_id_from_flask():
+    """Get trace_id from flask request headers.
+
+    :rtype: str
+    :return: Trace_id in HTTP request headers.
+    """
+    try:
+        trace_id = flask.request.headers['X_CLOUD_TRACE_CONTEXT'].split('/')[0]
+    except Exception:
+        trace_id = None
+    return trace_id
+
+
+def get_trace_id_from_django():
+    """Get trace_id from django request headers.
+
+    :rtype: str
+    :return: Trace_id in HTTP request headers.
+    """
+    try:
+        request = get_request()
+        trace_id = request.META['HTTP_X_CLOUD_TRACE_CONTEXT'].split('/')[0]
+    except Exception:
+        trace_id = None
+    return trace_id
+
+
+def get_trace_id():
+    """Helper to get trace_id from web application request header.
+
+    :rtype: str
+    :returns: Trace_id in HTTP request headers.
+    """
+    checkers = [get_trace_id_from_django, get_trace_id_from_flask]
+
+    for checker in checkers:
+        trace_id = checker()
+        if trace_id is not None:
+            return trace_id
+
+    return trace_id

logging/google/cloud/logging/handlers/app_engine.py

@@ -20,6 +20,7 @@
 
 import os
 
+from google.cloud.logging.handlers._helpers import get_trace_id
 from google.cloud.logging.handlers.handlers import CloudLoggingHandler
 from google.cloud.logging.handlers.transports import BackgroundThreadTransport
 from google.cloud.logging.resource import Resource
@@ -50,7 +51,8 @@ def __init__(self, client,
             client,
             name=_DEFAULT_GAE_LOGGER_NAME,
             transport=transport,
-            resource=self.get_gae_resource())
+            resource=self.get_gae_resource(),
+            labels=self.get_gae_labels())
 
     def get_gae_resource(self):
         """Return the GAE resource using the environment variables.
@@ -67,3 +69,17 @@ def get_gae_resource(self):
             },
         )
         return gae_resource
+
+    def get_gae_labels(self):
+        """Return the labels for GAE app which includes trace_id.
+
+        :rtype: dict
+        :returns: Labels for GAE app.
+        """
+        trace_id = get_trace_id()
+        if trace_id is None:
+            trace_id = 'unknown'
+        gae_labels = {
+            'appengine.googleapis.com/trace_id': trace_id,
+        }
+        return gae_labels

logging/google/cloud/logging/handlers/handlers.py

@@ -57,6 +57,9 @@ class CloudLoggingHandler(logging.StreamHandler):
     :param resource: (Optional) Monitored resource of the entry, defaults
                      to the global resource type.
 
+    :type labels: dict
+    :param labels: (Optional) Mapping of labels for the entry.
+
     Example:
 
     .. code-block:: python
@@ -79,12 +82,14 @@ class CloudLoggingHandler(logging.StreamHandler):
     def __init__(self, client,
                  name=DEFAULT_LOGGER_NAME,
                  transport=BackgroundThreadTransport,
-                 resource=_GLOBAL_RESOURCE):
+                 resource=_GLOBAL_RESOURCE,
+                 labels=None):
         super(CloudLoggingHandler, self).__init__()
         self.name = name
         self.client = client
         self.transport = transport(client, name)
         self.resource = resource
+        self.labels = labels
 
     def emit(self, record):
         """Actually log the specified logging record.
@@ -97,7 +102,11 @@ def emit(self, record):
         :param record: The record to be logged.
         """
         message = super(CloudLoggingHandler, self).format(record)
-        self.transport.send(record, message, resource=self.resource)
+        self.transport.send(
+            record,
+            message,
+            resource=self.resource,
+            labels=self.labels)
 
 
 def setup_logging(handler, excluded_loggers=EXCLUDED_LOGGER_DEFAULTS,

logging/google/cloud/logging/handlers/middleware/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud.logging.handlers.middleware.request import RequestMiddleware
+
+__all__ = ['RequestMiddleware']

logging/google/cloud/logging/handlers/middleware/request.py

@@ -0,0 +1,39 @@
+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from threading import local
+
+
+_thread_locals = local()
+
+
+def get_request():
+    """Get Django request from thread local.
+
+    :rtype: str
+    :returns: Django request.
+    """
+    return getattr(_thread_locals, 'request', None)
+
+
+class RequestMiddleware(object):
+    """Saves the request in thread local"""
+
+    def process_request(self, request):
+        """Called on each request, before Django decides which view to execute.
+
+        :type request: :class:`~django.http.request.HttpRequest`
+        :param request: Django http request.
+        """
+        _thread_locals.request = request

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

@@ -203,7 +203,7 @@ def _main_thread_terminated(self):
         else:
             print('Failed to send %d pending logs.' % (self._queue.qsize(),))
 
-    def enqueue(self, record, message, resource=None):
+    def enqueue(self, record, message, resource=None, labels=None):
         """Queues a log entry to be written by the background thread.
 
         :type record: :class:`logging.LogRecord`
@@ -215,6 +215,9 @@ def enqueue(self, record, message, resource=None):
 
         :type resource: :class:`~google.cloud.logging.resource.Resource`
         :param resource: (Optional) Monitored resource of the entry
+
+        :type labels: dict
+        :param labels: (Optional) Mapping of labels for the entry.
         """
         self._queue.put_nowait({
             'info': {
@@ -223,6 +226,7 @@ def enqueue(self, record, message, resource=None):
             },
             'severity': record.levelname,
             'resource': resource,
+            'labels': labels,
         })
 
     def flush(self):
@@ -257,7 +261,7 @@ def __init__(self, client, name, grace_period=_DEFAULT_GRACE_PERIOD,
         self.worker = _Worker(logger)
         self.worker.start()
 
-    def send(self, record, message, resource=None):
+    def send(self, record, message, resource=None, labels=None):
         """Overrides Transport.send().
 
         :type record: :class:`logging.LogRecord`
@@ -269,8 +273,11 @@ def send(self, record, message, resource=None):
 
         :type resource: :class:`~google.cloud.logging.resource.Resource`
         :param resource: (Optional) Monitored resource of the entry.
+
+        :type labels: dict
+        :param labels: (Optional) Mapping of labels for the entry.
         """
-        self.worker.enqueue(record, message, resource=resource)
+        self.worker.enqueue(record, message, resource=resource, labels=labels)
 
     def flush(self):
         """Submit any pending log records."""

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

@@ -22,7 +22,7 @@ class Transport(object):
     client and name object, and must override :meth:`send`.
     """
 
-    def send(self, record, message, resource=None):
+    def send(self, record, message, resource=None, labels=None):
         """Transport send to be implemented by subclasses.
 
         :type record: :class:`logging.LogRecord`
@@ -34,6 +34,9 @@ def send(self, record, message, resource=None):
 
         :type resource: :class:`~google.cloud.logging.resource.Resource`
         :param resource: (Optional) Monitored resource of the entry.
+
+        :type labels: dict
+        :param labels: (Optional) Mapping of labels for the entry.
         """
         raise NotImplementedError
 

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

@@ -29,7 +29,7 @@ class SyncTransport(Transport):
     def __init__(self, client, name):
         self.logger = client.logger(name)
 
-    def send(self, record, message, resource=None):
+    def send(self, record, message, resource=None, labels=None):
         """Overrides transport.send().
 
         :type record: :class:`logging.LogRecord`
@@ -38,8 +38,15 @@ def send(self, record, message, resource=None):
         :type message: str
         :param message: The message from the ``LogRecord`` after being
                         formatted by the associated log formatters.
+
+        :type resource: :class:`~google.cloud.logging.resource.Resource`
+        :param resource: (Optional) Monitored resource of the entry.
+
+        :type labels: dict
+        :param labels: (Optional) Mapping of labels for the entry.
         """
         info = {'message': message, 'python_logger': record.name}
         self.logger.log_struct(info,
                                severity=record.levelname,
-                               resource=resource)
+                               resource=resource,
+                               labels=labels)

logging/tests/unit/handlers/test_app_engine.py

@@ -34,6 +34,7 @@ def test_constructor(self):
         from google.cloud.logging.handlers.app_engine import _GAE_VERSION_ENV
 
         client = mock.Mock(project=self.PROJECT, spec=['project'])
+
         with mock.patch('os.environ', new={_GAE_PROJECT_ENV: 'test_project',
                                            _GAE_SERVICE_ENV: 'test_service',
                                            _GAE_VERSION_ENV: 'test_version'}):
@@ -43,13 +44,15 @@ def test_constructor(self):
         self.assertEqual(handler.resource.labels['project_id'], 'test_project')
         self.assertEqual(handler.resource.labels['module_id'], 'test_service')
         self.assertEqual(handler.resource.labels['version_id'], 'test_version')
+        self.assertEqual(handler.labels['appengine.googleapis.com/trace_id'], 'unknown')
 
     def test_emit(self):
         import mock
 
         client = mock.Mock(project=self.PROJECT, spec=['project'])
         handler = self._make_one(client, transport=_Transport)
         gae_resource = handler.get_gae_resource()
+        gae_labels = handler.get_gae_labels()
         logname = 'app'
         message = 'hello world'
         record = logging.LogRecord(logname, logging, None, None, message,
@@ -58,7 +61,9 @@ def test_emit(self):
 
         self.assertIs(handler.transport.client, client)
         self.assertEqual(handler.transport.name, logname)
-        self.assertEqual(handler.transport.send_called_with, (record, message, gae_resource))
+        self.assertEqual(
+            handler.transport.send_called_with,
+            (record, message, gae_resource, gae_labels))
 
 
 class _Transport(object):
@@ -67,5 +72,5 @@ def __init__(self, client, name):
         self.client = client
         self.name = name
 
-    def send(self, record, message, resource):
-        self.send_called_with = (record, message, resource)
+    def send(self, record, message, resource, labels):
+        self.send_called_with = (record, message, resource, labels)

logging/tests/unit/handlers/test_handlers.py

@@ -45,7 +45,7 @@ def test_emit(self):
                                    None, None)
         handler.emit(record)
 
-        self.assertEqual(handler.transport.send_called_with, (record, message, _GLOBAL_RESOURCE))
+        self.assertEqual(handler.transport.send_called_with, (record, message, _GLOBAL_RESOURCE, None))
 
 
 class TestSetupLogging(unittest.TestCase):
@@ -110,5 +110,5 @@ class _Transport(object):
     def __init__(self, client, name):
         pass
 
-    def send(self, record, message, resource):
-        self.send_called_with = (record, message, resource)
+    def send(self, record, message, resource, labels=None):
+        self.send_called_with = (record, message, resource, labels)

logging/tests/unit/handlers/transports/test_background_thread.py

@@ -61,9 +61,9 @@ def test_send(self):
             python_logger_name, logging.INFO,
             None, None, message, None, None)
 
-        transport.send(record, message, _GLOBAL_RESOURCE)
+        transport.send(record, message, _GLOBAL_RESOURCE, None)
 
-        transport.worker.enqueue.assert_called_once_with(record, message, _GLOBAL_RESOURCE)
+        transport.worker.enqueue.assert_called_once_with(record, message, _GLOBAL_RESOURCE, None)
 
     def test_flush(self):
         client = _Client(self.PROJECT)
@@ -287,13 +287,13 @@ def __init__(self):
         self.commit_called = False
         self.commit_count = None
 
-    def log_struct(self, info, severity=logging.INFO, resource=None):
+    def log_struct(self, info, severity=logging.INFO, resource=None, labels=None):
         from google.cloud.logging.logger import _GLOBAL_RESOURCE
 
         assert resource is None
         resource = _GLOBAL_RESOURCE
 
-        self.log_struct_called_with = (info, severity, resource)
+        self.log_struct_called_with = (info, severity, resource, labels)
         self.entries.append(info)
 
     def commit(self):

logging/tests/unit/handlers/transports/test_sync.py

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