Send trace context with logs from web applications

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

@@ -17,6 +17,23 @@
 import math
 import json
 
+try:
+    import django
+except ImportError:
+    _USE_DJANGO = False
+else:
+    _USE_DJANGO = True
+
+try:
+    import flask
+except ImportError:
+    _USE_FLASK = False
+    flask = None
+else:
+    _USE_FLASK = True
+
+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 +54,25 @@ def format_stackdriver_json(record, message):
     }
 
     return json.dumps(payload)
+
+
+def get_trace_id_from_request_header():
+    """Helper to get trace_id from web application request header.
+
+    :rtype: str
+    :returns: Trace_id in HTTP request headers.
+    """
+    if _USE_FLASK:
+        try:
+            trace_id = flask.request.headers['X_CLOUD_TRACE_CONTEXT'].split('/')[0]
+        except Exception:
+            trace_id = None
+    elif _USE_DJANGO:
+        try:
+            request = get_request()
+            trace_id = request.META['HTTP_X_CLOUD_TRACE_CONTEXT'].split('/')[0]
+        except Exception:
+            trace_id = None
+    else:
+        trace_id = None
+    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_request_header
 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_from_request_header()
+        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,10 @@ 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 2016 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,40 @@
+# Copyright 2016 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.
+
+try:
+    from threading import local
+except ImportError:
+    from django.utils._threading_local 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)