Add support for logging the trace-id in webapp2 apps.

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()

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.

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