diff --git a/logging/google/cloud/logging/handlers/_helpers.py b/logging/google/cloud/logging/handlers/_helpers.py index 1ebb064ed228..864f0e53617e 100644 --- a/logging/google/cloud/logging/handlers/_helpers.py +++ b/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() diff --git a/logging/nox.py b/logging/nox.py index 068d5ae8d198..ce8d1c0afbce 100644 --- a/logging/nox.py +++ b/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. diff --git a/logging/tests/unit/handlers/test__helpers.py b/logging/tests/unit/handlers/test__helpers.py index 0731c825d32c..516cd93fc2d5 100644 --- a/logging/tests/unit/handlers/test__helpers.py +++ b/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