Add support for logging the trace-id in webapp2 apps. (#3593)

Author: csilvers

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

@@ -22,11 +22,21 @@
 except ImportError:  # pragma: NO COVER
     flask = None
 
+try:
+    import webapp2
+except (ImportError, SyntaxError):  # pragma: NO COVER
+    # If you try to import webapp2 under python3, you'll get a syntax
+    # error (since it hasn't been ported yet).  We just pretend it
+    # doesn't exist.  This is unlikely to hit in real life but does
+    # in the tests.
+    webapp2 = 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'
+_FLASK_TRACE_HEADER = 'X_CLOUD_TRACE_CONTEXT'
+_WEBAPP2_TRACE_HEADER = 'X-CLOUD-TRACE-CONTEXT'
 
 
 def format_stackdriver_json(record, message):
@@ -54,7 +64,7 @@ def get_trace_id_from_flask():
     """Get trace_id from flask request headers.
 
     :rtype: str
-    :return: Trace_id in HTTP request headers.
+    :returns: TraceID in HTTP request headers.
     """
     if flask is None or not flask.request:
         return None
@@ -69,11 +79,38 @@ def get_trace_id_from_flask():
     return trace_id
 
 
+def get_trace_id_from_webapp2():
+    """Get trace_id from webapp2 request headers.
+
+    :rtype: str
+    :returns: TraceID in HTTP request headers.
+    """
+    if webapp2 is None:
+        return None
+
+    try:
+        # get_request() succeeds if we're in the middle of a webapp2
+        # request, or raises an assertion error otherwise:
+        # "Request global variable is not set".
+        req = webapp2.get_request()
+    except AssertionError:
+        return None
+
+    header = req.headers.get(_WEBAPP2_TRACE_HEADER)
+
+    if header is None:
+        return None
+
+    trace_id = header.split('/', 1)[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.
+    :returns: TraceID in HTTP request headers.
     """
     request = _get_django_request()
 
@@ -93,9 +130,11 @@ def get_trace_id():
     """Helper to get trace_id from web application request header.
 
     :rtype: str
-    :returns: Trace_id in HTTP request headers.
+    :returns: TraceID in HTTP request headers.
     """
-    checkers = (get_trace_id_from_django, get_trace_id_from_flask)
+    checkers = (get_trace_id_from_django,
+                get_trace_id_from_flask,
+                get_trace_id_from_webapp2)
 
     for checker in checkers:
         trace_id = checker()

packages/google-cloud-logging/nox.py

@@ -36,7 +36,7 @@ def unit_tests(session, python_version):
     # Install all test dependencies, then install this package in-place.
     session.install(
         'mock', 'pytest', 'pytest-cov',
-        'flask', 'django', *LOCAL_DEPS)
+        'flask', 'webapp2', 'webob', 'django', *LOCAL_DEPS)
     session.install('-e', '.')
 
     # Run py.test against the unit tests.

packages/google-cloud-logging/tests/unit/handlers/test__helpers.py

@@ -12,9 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import unittest
 
 import mock
+import six
+
+try:
+    from webapp2 import RequestHandler
+except SyntaxError:
+    # webapp2 has not been ported to python3, so it will give a syntax
+    # error if we try.  We'll just skip the webapp2 tests in that case.
+    RequestHandler = object
 
 
 class Test_get_trace_id_from_flask(unittest.TestCase):
@@ -37,11 +46,9 @@ def index():
 
         return app
 
-    def setUp(self):
-        self.app = self.create_app()
-
     def test_no_context_header(self):
-        with self.app.test_request_context(
+        app = self.create_app()
+        with app.test_request_context(
                 path='/',
                 headers={}):
             trace_id = self._call_fut()
@@ -53,7 +60,8 @@ def test_valid_context_header(self):
         expected_trace_id = 'testtraceidflask'
         flask_trace_id = expected_trace_id + '/testspanid'
 
-        context = self.app.test_request_context(
+        app = self.create_app()
+        context = app.test_request_context(
             path='/',
             headers={flask_trace_header: flask_trace_id})
 
@@ -63,6 +71,54 @@ def test_valid_context_header(self):
         self.assertEqual(trace_id, expected_trace_id)
 
 
+class _GetTraceId(RequestHandler):
+    def get(self):
+        from google.cloud.logging.handlers import _helpers
+
+        trace_id = _helpers.get_trace_id_from_webapp2()
+        self.response.content_type = 'application/json'
+        self.response.out.write(json.dumps(trace_id))
+
+
+
+@unittest.skipIf(six.PY3, 'webapp2 is Python 2 only')
+class Test_get_trace_id_from_webapp2(unittest.TestCase):
+
+    @staticmethod
+    def create_app():
+        import webapp2
+
+        app = webapp2.WSGIApplication([
+            ('/', _GetTraceId),
+        ])
+
+        return app
+
+    def test_no_context_header(self):
+        import webob
+
+        req = webob.BaseRequest.blank('/')
+        response = req.get_response(self.create_app())
+        trace_id = json.loads(response.body)
+
+        self.assertEquals(None, trace_id)
+
+    def test_valid_context_header(self):
+        import webob
+
+        webapp2_trace_header = 'X-Cloud-Trace-Context'
+        expected_trace_id = 'testtraceidwebapp2'
+        webapp2_trace_id = expected_trace_id + '/testspanid'
+
+        req = webob.BaseRequest.blank(
+            '/',
+            headers={webapp2_trace_header: webapp2_trace_id})
+        response = req.get_response(self.create_app())
+        trace_id = json.loads(response.body)
+
+        self.assertEqual(trace_id, expected_trace_id)
+
+
 class Test_get_trace_id_from_django(unittest.TestCase):
 
     @staticmethod