fix(tornado.py): fix tornado mis-handle of 404s (#314)

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/events/urllib3.py

@@ -113,16 +113,20 @@ def update_response(self, response):
                 'response_headers',
                 headers
             )
-            response_body = response.peek()
+            response_body = getattr(response, 'data', None)
+            if not response_body and getattr(response, 'peek', None):
+                response_body = response.peek()
             if isinstance(response_body, bytes):
                 try:
                     response_body = response_body.decode('utf-8')
                 except UnicodeDecodeError:
                     response_body = str(response_body)
-            add_data_if_needed(
-                self.resource['metadata'],
-                'response_body',
-                response_body)
+            if response_body:
+                add_data_if_needed(
+                    self.resource['metadata'],
+                    'response_body',
+                    response_body
+                )
 
         # Detect errors based on status code
         if response.status >= HTTP_ERR_CODE:

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

@@ -44,6 +44,12 @@ def patch():
         'Redis.execute_command',
         _single_wrapper
     )
+    # Version < 3.0
+    wrapt.wrap_function_wrapper(
+        'redis',
+        'StrictRedis.execute_command',
+        _single_wrapper
+    )
     wrapt.wrap_function_wrapper(
         'redis.client',
         'Pipeline.immediate_execute_command',

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)
@@ -80,6 +91,11 @@ def after_request(cls, wrapped, instance, args, kwargs):
 
             tornado_runner = cls.RUNNERS.pop(unique_id)
 
+            # Ignoring 404s
+            if getattr(instance, '_status_code', None) == 404:
+                print_debug('Ignoring 404 Tornado request')
+                return res
+
             trace = epsagon.trace.trace_factory.switch_active_trace(
                 unique_id
             )
@@ -111,6 +127,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/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 body: {}'.format(exception))
+
     def update_response(self, response, response_body=None):
         """
         Adds response data to event.

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