fix(tornado.py): fix mishandling of 404s (#317)

Author: ranrib

epsagon/events/redis.py

@@ -8,6 +8,8 @@
 
 from ..event import BaseEvent
 from ..trace import trace_factory
+from ..utils import add_data_if_needed
+
 
 MAX_VALUE_SIZE = 1024
 MAX_CMD_PIPELINE = 10
@@ -110,6 +112,13 @@ def __init__(self, wrapped, instance, args, kwargs, start_time, response,
         self.resource['operation'] = operation
         self.resource['metadata']['Redis Key'] = key
 
+        if response:
+            add_data_if_needed(
+                self.resource['metadata'],
+                'redis.response',
+                response
+            )
+
 
 class RedisMultiExecutionEvent(BaseRedisEvent):
     """

epsagon/modules/logging.py

@@ -11,7 +11,7 @@
 import wrapt
 
 from ..trace import trace_factory
-from ..utils import print_debug
+from ..utils import print_debug, get_trace_log_config
 
 LOGGING_FUNCTIONS = (
     'info',
@@ -99,7 +99,7 @@ def patch():
     wrapt.wrap_function_wrapper('logging', 'Logger.exception', _wrapper)
 
     # Instrument logging with Epsagon trace ID
-    if trace_factory.is_logging_tracing_enabled():
+    if get_trace_log_config():
         wrapt.wrap_function_wrapper(
             'logging',
             'Logger.log',

epsagon/modules/redis.py

@@ -54,3 +54,9 @@ def patch():
         'Pipeline.execute',
         _multi_wrapper
     )
+    # Version < 3.0
+    wrapt.wrap_function_wrapper(
+        'redis',
+        'StrictRedis.execute_command',
+        _single_wrapper
+    )

epsagon/modules/tornado.py

@@ -11,7 +11,7 @@
 import epsagon.trace
 from epsagon.runners.tornado import TornadoRunner
 from epsagon.http_filters import ignore_request, is_ignored_endpoint
-from epsagon.utils import collect_container_metadata
+from epsagon.utils import collect_container_metadata, print_debug
 
 TORNADO_TRACE_ID = 'epsagon_tornado_trace_key'
 
@@ -31,6 +31,7 @@ def before_request(cls, wrapped, instance, args, kwargs):
         :param args: wrapt's args
         :param kwargs: wrapt's kwargs
         """
+        print_debug('before_request Tornado request')
         try:
             ignored = ignore_request('', instance.request.path)
             if not ignored and not is_ignored_endpoint(instance.request.path):
@@ -48,6 +49,7 @@ def before_request(cls, wrapped, instance, args, kwargs):
                 )
 
                 trace.set_runner(cls.RUNNERS[unique_id])
+                print_debug('Created Tornado Runner')
 
                 # Collect metadata in case this is a container.
                 collect_container_metadata(
@@ -69,9 +71,18 @@ def after_request(cls, wrapped, instance, args, kwargs):
         :param args: wrapt's args
         :param kwargs: wrapt's kwargs
         """
-        response_body = getattr(instance, '_write_buffer', None)
-        if response_body and isinstance(response_body, list):
-            response_body = b''.join(response_body)
+        print_debug('after_request Tornado request')
+        response_body = None
+        try:
+            response_body = getattr(instance, '_write_buffer', None)
+            if response_body and isinstance(response_body, list):
+                response_body = b''.join(response_body)
+        except Exception as instrumentation_exception:  # pylint: disable=W0703
+            epsagon.trace.trace_factory.add_exception(
+                instrumentation_exception,
+                traceback.format_exc()
+            )
+
         res = wrapped(*args, **kwargs)
         try:
             unique_id = getattr(instance, TORNADO_TRACE_ID, None)
@@ -84,6 +95,12 @@ def after_request(cls, wrapped, instance, args, kwargs):
                 unique_id
             )
 
+            # Ignoring 404s
+            if getattr(instance, '_status_code', None) == 404:
+                epsagon.trace.trace_factory.pop_trace(trace=trace)
+                print_debug('Ignoring 404 Tornado request')
+                return res
+
             if trace:
                 content = (
                     instance._headers.get(  # pylint: disable=protected-access
@@ -111,6 +128,7 @@ def collect_exception(cls, wrapped, instance, args, kwargs):
         :param args: wrapt's args
         :param kwargs: wrapt's kwargs
         """
+        print_debug('collect_exception Tornado request')
         try:
             unique_id = getattr(instance, TORNADO_TRACE_ID, None)
 

epsagon/runners/aiohttp.py

@@ -39,7 +39,7 @@ def __init__(self, start_time, request, body, handler):
             'Base URL': request.url.host,
             'Path': request.url.path,
             'User Agent': request.headers.get('User-Agent', 'N/A'),
-            'Endpoint': handler.__name__,
+            'Endpoint': getattr(handler, '__name__', ''),
         })
 
         if request.query_string:

epsagon/runners/tornado.py

@@ -6,7 +6,7 @@
 from __future__ import absolute_import
 import uuid
 from ..event import BaseEvent
-from ..utils import add_data_if_needed
+from ..utils import add_data_if_needed, print_debug
 from ..constants import EPSAGON_HEADER_TITLE
 
 MAX_PAYLOAD_BYTES = 2000
@@ -69,6 +69,19 @@ def __init__(self, start_time, request):
                 request_headers
             )
 
+        try:
+            if request.body:
+                body = request.body
+                if isinstance(body, bytes):
+                    body = body.decode('utf-8')
+                add_data_if_needed(
+                    self.resource['metadata'],
+                    'Request Data',
+                    body
+                )
+        except Exception as exception:  # pylint: disable=broad-except
+            print_debug('Could not extract request body: {}'.format(exception))
+
     def update_response(self, response, response_body=None):
         """
         Adds response data to event.
@@ -83,10 +96,16 @@ def update_response(self, response, response_body=None):
         )
 
         if response_body:
+            body = response_body
+            if isinstance(body, list):
+                body = body[0]
+            if isinstance(body, bytes):
+                body = body.decode('utf-8')
+
             add_data_if_needed(
                 self.resource['metadata'],
                 'Response Body',
-                str(response_body)[:MAX_PAYLOAD_BYTES]
+                str(body)[:MAX_PAYLOAD_BYTES]
             )
 
         self.resource['metadata']['status_code'] = response._status_code

epsagon/utils.py

@@ -95,6 +95,19 @@ def get_tc_url(use_ssl):
     return TRACE_COLLECTOR_URL.format(protocol=protocol, region=REGION)
 
 
+def get_trace_log_config():
+    """Returns the trace log correlation configuration"""
+    # In case we're running on AWS Lambda, logging correlation is disabled
+    if is_lambda_env():
+        return False
+
+    # If EPSAGON_LOGGING_TRACING_ENABLED exists as an env var - use it
+    if os.getenv('EPSAGON_LOGGING_TRACING_ENABLED'):
+        return os.getenv('EPSAGON_LOGGING_TRACING_ENABLED').upper() == 'TRUE'
+
+    return True
+
+
 def init(
     token='',
     app_name='Application',
@@ -174,15 +187,7 @@ def init(
     if os.getenv('EPSAGON_METADATA'):
         metadata_only = (os.getenv('EPSAGON_METADATA') or '').upper() == 'TRUE'
 
-    # If EPSAGON_LOGGING_TRACING_ENABLED exists as an env var - use it
-    if os.getenv('EPSAGON_LOGGING_TRACING_ENABLED'):
-        logging_tracing_enabled = (
-            os.getenv('EPSAGON_LOGGING_TRACING_ENABLED') or ''
-        ).upper() == 'TRUE'
-
-    # In case we're running on AWS Lambda, logging correlation is disabled
-    if is_lambda_env():
-        logging_tracing_enabled = False
+    logging_tracing_enabled = get_trace_log_config()
 
     if os.getenv('EPSAGON_SAMPLE_RATE'):
         sample_rate = float(os.getenv('EPSAGON_SAMPLE_RATE'))

epsagon/wrappers/aiohttp.py

@@ -6,6 +6,7 @@
 import time
 import warnings
 from aiohttp.web import middleware
+from aiohttp.web_exceptions import HTTPNotFound
 
 import epsagon.trace
 import epsagon.triggers.http
@@ -16,6 +17,7 @@
     get_traceback_data_from_exception
 )
 from ..http_filters import ignore_request
+from ..utils import print_debug
 
 
 @middleware
@@ -26,9 +28,11 @@ async def AiohttpMiddleware(request, handler):
     :param handler: original handler
     :return: response data from the handler
     """
+    print_debug('[aiohttp] started middleware')
     epsagon.trace.trace_factory.switch_to_async_tracer()
 
     if ignore_request('', request.path.lower()):
+        print_debug('[aiohttp] ignoring request')
         return await handler(request)
 
     trace = epsagon.trace.trace_factory.get_or_create_trace()
@@ -38,15 +42,22 @@ async def AiohttpMiddleware(request, handler):
 
     try:
         body = await request.text()
+        print_debug('[aiohttp] got body')
         runner = AiohttpRunner(time.time(), request, body, handler)
         trace.set_runner(runner)
         collect_container_metadata(runner.resource['metadata'])
+        print_debug('[aiohttp] initialized runner')
     except Exception as exception: # pylint: disable=W0703
         warnings.warn('Could not extract request', EpsagonWarning)
 
     raised_err = None
     try:
         response = await handler(request)
+        print_debug('[aiohttp] got response')
+    except HTTPNotFound:
+        # Ignoring 404s
+        epsagon.trace.trace_factory.pop_trace(trace)
+        raise
     except Exception as exception:  # pylint: disable=W0703
         raised_err = exception
         traceback_data = get_traceback_data_from_exception(exception)
@@ -59,7 +70,10 @@ async def AiohttpMiddleware(request, handler):
         runner.update_response(response)
 
     if runner:
+        print_debug('[aiohttp] sending trace')
         epsagon.trace.trace_factory.send_traces()
     if raised_err:
+        print_debug('[aiohttp] raising error')
         raise raised_err
+    print_debug('[aiohttp] middleware done')
     return response

tests/wrappers/test_aiohttp_wrapper.py

@@ -53,3 +53,15 @@ async def test_aiohttp_exception(_, trace_transport, aiohttp_client):
     assert runner.error_code == ErrorCode.EXCEPTION
     assert runner.exception['type'] == 'CustomAioHttpException'
     assert runner.exception['message'] == 'test'
+
+
+@asynctest.patch('epsagon.trace.trace_factory.use_async_tracer')
+async def test_aiohttp_no_endpoint(_, trace_transport, aiohttp_client):
+    """Test when no route exists (404)."""
+    client = await aiohttp_client(create_app)
+    response = await client.get('/not/exists')
+    assert response.status == 404
+    # Trace not sent
+    trace_transport.send.assert_not_called()
+    # Trace removed
+    assert trace_transport.last_trace is None