Send trace context with logs from web applications

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

@@ -17,6 +17,16 @@
 import math
 import json
 
+try:
+    import flask
+except ImportError:
+    flask = None
+
+from google.cloud.logging.handlers.middleware.request import _get_django_request
+
+_FLASK_TRACE_HEADER = 'X_CLOUD_TRACE_CONTEXT'
+_DJANGO_TRACE_HEADER = 'HTTP_X_CLOUD_TRACE_CONTEXT'
+
 
 def format_stackdriver_json(record, message):
     """Helper to format a LogRecord in in Stackdriver fluentd format.
@@ -37,3 +47,61 @@ 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.
+    """
+    if flask is None or not flask.request:
+        return None
+
+    header = flask.request.headers.get(_FLASK_TRACE_HEADER)
+
+    if header is None:
+        return None
+
+    trace_id = header.split('/')[0]
+
+    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.
+    """
+    request = _get_django_request()
+
+    if request is None:
+        return None
+
+    try:
+        header = request.META.get(_DJANGO_TRACE_HEADER)
+    except KeyError:
+        return None
+
+    if header is None:
+        return None
+    trace_id = header.split('/')[0]
+
+    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 None

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
@@ -30,6 +31,8 @@
 _GAE_SERVICE_ENV = 'GAE_SERVICE'
 _GAE_VERSION_ENV = 'GAE_VERSION'
 
+_TRACE_ID_LABEL = 'appengine.googleapis.com/trace_id'
+
 
 class AppEngineHandler(CloudLoggingHandler):
     """A logging handler that sends App Engine-formatted logs to Stackdriver.
@@ -50,7 +53,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 +71,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.
+        """
+        gae_labels = {}
+
+        trace_id = get_trace_id()
+        if trace_id is not None:
+            gae_labels[_TRACE_ID_LABEL] = 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.
+
+import threading
+
+
+_thread_locals = threading.local()
+
+
+def _get_django_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/nox.py

@@ -31,7 +31,9 @@ def unit_tests(session, python_version):
     session.interpreter = 'python{}'.format(python_version)
 
     # Install all test dependencies, then install this package in-place.
-    session.install('mock', 'pytest', 'pytest-cov', *LOCAL_DEPS)
+    session.install(
+        'mock', 'pytest', 'pytest-cov',
+        'flask', 'django', *LOCAL_DEPS)
     session.install('-e', '.')
 
     # Run py.test against the unit tests.

logging/tests/unit/handlers/middleware/test_request.py

@@ -0,0 +1,50 @@
+# 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.
+
+import unittest
+
+
+class TestRequestMiddleware(unittest.TestCase):
+
+    def _get_target_class(self):
+        from google.cloud.logging.handlers.middleware.request import RequestMiddleware
+
+        return RequestMiddleware
+
+    def _make_one(self, *args, **kw):
+        return self._get_target_class()(*args, **kw)
+
+    @classmethod
+    def setUpClass(cls):
+        from django.conf import settings
+        from django.test.utils import setup_test_environment
+
+        if not settings.configured:
+            settings.configure()
+        setup_test_environment()
+
+    @classmethod
+    def tearDownClass(cls):
+        from django.test.utils import teardown_test_environment
+
+        teardown_test_environment()
+
+    def test_get_django_request(self):
+        from django.test import RequestFactory
+        from google.cloud.logging.handlers.middleware.request import _get_django_request
+
+        middleware = self._make_one()
+        request = RequestFactory().get('/')
+        middleware.process_request(request)
+        self.assertEqual(_get_django_request(), request)