feat(tornado_client.py): support tornado http client calls (#318)

Author: ranrib

epsagon/events/tornado_client.py

@@ -0,0 +1,162 @@
+"""
+Tornado AsyncHTTPClient events module.
+"""
+
+from __future__ import absolute_import
+
+import functools
+try:
+    from urllib.parse import urlparse, urlunparse
+except ImportError:
+    from urlparse import urlparse, urlunparse
+import traceback
+from uuid import uuid4
+
+from epsagon.utils import add_data_if_needed
+from ..trace import trace_factory
+from ..event import BaseEvent
+from ..http_filters import (
+    is_blacklisted_url,
+    is_payload_collection_blacklisted
+)
+from ..utils import update_http_headers, normalize_http_url
+from ..constants import HTTP_ERR_CODE, EPSAGON_HEADER_TITLE
+
+
+class TornadoAsyncHTTPClientEvent(BaseEvent):
+    """
+    Represents base request event.
+    """
+
+    ORIGIN = 'tornado_client'
+    RESOURCE_TYPE = 'http'
+
+    # pylint: disable=W0613
+    def __init__(self, wrapped, instance, args, kwargs, start_time, response,
+                 exception):
+        """
+        Initialize.
+        :param wrapped: wrapt's wrapped
+        :param instance: wrapt's instance
+        :param args: wrapt's args
+        :param kwargs: wrapt's kwargs
+        :param start_time: Start timestamp (epoch)
+        :param response: response data
+        :param exception: Exception (if happened)
+        """
+        super(TornadoAsyncHTTPClientEvent, self).__init__(start_time)
+        self.event_id = 'tornado-client-{}'.format(str(uuid4()))
+
+        request = args[0]
+        headers = dict(request.headers)
+        if headers:
+            # Make sure trace ID is present in case headers will be removed.
+            epsagon_trace_id = headers.get(EPSAGON_HEADER_TITLE)
+            if epsagon_trace_id:
+                self.resource['metadata']['http_trace_id'] = epsagon_trace_id
+
+        parsed_url = urlparse(request.url)
+
+        host_url = parsed_url.netloc.split(':')[0]
+        full_url = urlunparse((
+            parsed_url.scheme,
+            host_url,
+            parsed_url.path,
+            parsed_url.params,
+            parsed_url.query,
+            parsed_url.fragment,
+        ))
+
+        self.resource['name'] = normalize_http_url(request.url)
+        self.resource['operation'] = request.method
+        self.resource['metadata']['url'] = request.url
+
+        if not is_payload_collection_blacklisted(full_url):
+            add_data_if_needed(
+                self.resource['metadata'],
+                'request_headers',
+                headers
+            )
+            body = request.body
+            if isinstance(body, bytes):
+                body = body.decode('utf-8')
+            if body:
+                add_data_if_needed(
+                    self.resource['metadata'],
+                    'request_body',
+                    body
+                )
+
+        if response is not None:
+            callback = functools.partial(self.update_response)
+            response.add_done_callback(callback)
+
+        if exception is not None:
+            self.set_exception(exception, traceback.format_exc())
+
+    def update_response(self, future):
+        """
+        Adds response data to event.
+        :param future: Future response object
+        :return: None
+        """
+        response = future.result()
+
+        self.resource['metadata']['status_code'] = response.code
+        self.resource = update_http_headers(
+            self.resource,
+            dict(response.headers)
+        )
+
+        full_url = self.resource['metadata']['url']
+
+        if not is_payload_collection_blacklisted(full_url):
+            add_data_if_needed(
+                self.resource['metadata'],
+                'response_headers',
+                dict(response.headers)
+            )
+            body = response.body
+            if isinstance(body, bytes):
+                try:
+                    body = body.decode('utf-8')
+                except UnicodeDecodeError:
+                    body = str(body)
+            if body:
+                add_data_if_needed(
+                    self.resource['metadata'],
+                    'response_body',
+                    body
+                )
+
+        # Detect errors based on status code
+        if response.code >= HTTP_ERR_CODE:
+            self.set_error()
+
+
+class TornadoClientEventFactory(object):
+    """
+    Factory class, generates AsyncHTTPClient event.
+    """
+
+    @staticmethod
+    def create_event(wrapped, instance, args, kwargs, start_time, response,
+                     exception):
+        """
+        Create an AsyncHTTPClient event.
+        """
+        # Detect if URL is blacklisted, and ignore.
+        if is_blacklisted_url(args[0].url):
+            return
+
+        event = TornadoAsyncHTTPClientEvent(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            start_time,
+            response,
+            exception
+        )
+
+        trace_factory.add_event(event)

epsagon/modules/tornado.py

@@ -8,10 +8,20 @@
 import uuid
 from functools import partial
 import wrapt
+from tornado.httpclient import HTTPRequest
+from tornado.httputil import HTTPHeaders
 import epsagon.trace
+from epsagon.modules.general_wrapper import wrapper
 from epsagon.runners.tornado import TornadoRunner
 from epsagon.http_filters import ignore_request, is_ignored_endpoint
-from epsagon.utils import collect_container_metadata, print_debug
+from epsagon.utils import (
+    collect_container_metadata,
+    print_debug,
+    get_epsagon_http_trace_id
+)
+from ..constants import EPSAGON_HEADER
+from ..events.tornado_client import TornadoClientEventFactory
+
 
 TORNADO_TRACE_ID = 'epsagon_tornado_trace_key'
 
@@ -229,6 +239,50 @@ def thread_pool_execution(cls, unique_id, fn, args, kwargs):
         return res
 
 
+def _prepare_http_request(request, raise_error=True, **kwargs):
+    """
+    Copying parameters from original `AsyncHTTPClient.fetch` function
+    :return: HTTPRequest, raise_error bool
+    """
+    # request can be a URL string or a HTTPRequest object
+    if not isinstance(request, HTTPRequest):
+        request = HTTPRequest(request, **kwargs)
+
+    return request, raise_error
+
+
+def _wrapper(wrapped, instance, args, kwargs):
+    """
+    General wrapper for AsyncHTTPClient instrumentation.
+    :param wrapped: wrapt's wrapped
+    :param instance: wrapt's instance
+    :param args: wrapt's args
+    :param kwargs: wrapt's kwargs
+    :return: None
+    """
+    try:
+        request, raise_error = _prepare_http_request(*args, **kwargs)
+    except Exception:  # pylint: disable=W0703
+        return wrapped(*args, **kwargs)
+
+    trace_header = get_epsagon_http_trace_id()
+
+    if isinstance(request.headers, HTTPHeaders):
+        if not request.headers.get(EPSAGON_HEADER):
+            request.headers.add(EPSAGON_HEADER, trace_header)
+    elif isinstance(request.headers, dict):
+        if EPSAGON_HEADER not in request.headers:
+            request.headers[EPSAGON_HEADER] = trace_header
+
+    return wrapper(
+        TornadoClientEventFactory,
+        wrapped,
+        instance,
+        (request, ),  # new args
+        {'raise_error': raise_error}  # new kwargs
+    )
+
+
 def patch():
     """
     Patch module.
@@ -267,3 +321,9 @@ def patch():
     except Exception:  # pylint: disable=broad-except
         # Can happen in different Tornado versions.
         pass
+
+    wrapt.wrap_function_wrapper(
+        'tornado.httpclient',
+        'AsyncHTTPClient.fetch',
+        _wrapper
+    )

epsagon/runners/aiohttp.py

@@ -7,7 +7,7 @@
 from aiohttp.payload import StringPayload
 from ..event import BaseEvent
 from ..utils import add_data_if_needed, normalize_http_url
-from ..constants import EPSAGON_HEADER
+from ..constants import EPSAGON_HEADER_TITLE
 
 
 class AiohttpRunner(BaseEvent):
@@ -53,9 +53,9 @@ def __init__(self, start_time, request, body, handler):
             )
 
         request_headers = dict(request.headers)
-        if request_headers.get(EPSAGON_HEADER):
+        if request_headers.get(EPSAGON_HEADER_TITLE):
             self.resource['metadata']['http_trace_id'] = request_headers.get(
-                EPSAGON_HEADER
+                EPSAGON_HEADER_TITLE
             )
         if request_headers:
             add_data_if_needed(

epsagon/utils.py

@@ -5,6 +5,7 @@
 from __future__ import absolute_import, print_function
 import os
 import collections
+import uuid
 import socket
 import sys
 import traceback
@@ -85,6 +86,18 @@ def update_http_headers(resource_data, response_headers):
     return resource_data
 
 
+def get_epsagon_http_trace_id():
+    """Returns an Epsagon trace ID to inject over HTTP."""
+    trace_id = uuid.uuid4().hex
+    span_id = uuid.uuid4().hex[16:]
+    parent_span_id = uuid.uuid4().hex[16:]
+    return '{trace_id}:{span_id}:{parent_span_id}:1'.format(
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id
+    )
+
+
 def get_tc_url(use_ssl):
     """
     Get the TraceCollector URL.

requirements-dev.txt

@@ -19,3 +19,4 @@ pytest-aiohttp; python_version >= '3.5'
 httpx; python_version >= '3.5'
 asynctest; python_version >= '3.5'
 moto
+tornado

tests/events/test_tornado_client.py

@@ -0,0 +1,25 @@
+import epsagon.wrappers.python_function
+import epsagon.runners.python_function
+import epsagon.constants
+import mock
+from tornado.httpclient import AsyncHTTPClient
+
+TEST_URL = 'https://example.test/'
+
+
+@mock.patch('epsagon.trace.TraceFactory.add_event')
+def test_sanity(add_event_mock):
+    retval = 'success'
+
+    @epsagon.wrappers.python_function.python_wrapper
+    def wrapped_function():
+        http_client = AsyncHTTPClient()
+        http_client.fetch(TEST_URL)
+        return retval
+    assert wrapped_function() == retval
+    add_event_mock.assert_called()
+    event = add_event_mock.call_args_list[0].args[0]
+    assert event.resource['name'] == 'example.test'
+    assert event.resource['operation'] == 'GET'
+    assert event.resource['type'] == 'http'
+    assert 'http_trace_id' in event.resource['metadata']