updating library for release version

Author: prateeksachan

tracereq/__init__.py

@@ -3,35 +3,46 @@
 _installer_lock = Lock()
 _installed_integrations = {}
 
-def get_default_integrations(wsgi_app):
-    from .customhttplib import CustomHttpLib
-    yield CustomHttpLib()
 
-    from .flasklib import FlaskLib
-    yield FlaskLib(wsgi_app)
+def get_core_integrations():
+    from .customlib import CustomlibIntegration
+    yield CustomlibIntegration()
 
 
-def setup_integrations(wsgi_app):
+def get_custom_integrations(**kwargs):
+    if kwargs.get('flask_wsgi_app'):
+        from .flasklib import FlasklibIntegration
+        yield FlasklibIntegration(kwargs.get('flask_wsgi_app'))
+
+
+def setup_integrations(*args, **kwargs):
     integrations = list()
-    for integration in get_default_integrations(wsgi_app):
-        if not any(isinstance(x, type(integration)) for x in integrations):
-            integrations.append(integration)
+
+    core_integrations = get_core_integrations()
+    custom_integrations = get_custom_integrations(**kwargs)
+    for instance in core_integrations:
+        if not any(isinstance(x, type(instance)) for x in integrations):
+            integrations.append(instance)
+
+    for instance in custom_integrations:
+        if not any(isinstance(x, type(instance)) for x in integrations):
+            integrations.append(instance)
 
     for integration in integrations:
         integration()
 
 
 class Integration(object):
-    identifier = None
+    integration_key = None
 
     def install(self):
         raise NotImplementedError()
 
     def __call__(self, environ=None, start_response=None):
-        assert self.identifier
+        assert self.integration_key
         with _installer_lock:
-            if self.identifier in _installed_integrations:
+            if self.integration_key in _installed_integrations:
                 return
 
             self.install()
-            _installed_integrations[self.identifier] = self
+            _installed_integrations[self.integration_key] = self

tracereq/api.py

@@ -1,15 +1,17 @@
-from .hub import Hub
+from .engine import Engine
 from .client import Client
 from . import setup_integrations
 
-__all__ = ['Hub', 'Client']
+__all__ = ['Engine', 'Client']
+
 
 def public(f):
     __all__.append(f.__name__)
     return f
 
+
 @public
-def init(flask_wsgi_app, *args, **kwargs):
-    client = Client(*args, **kwargs)
-    Hub.master.attach_client(client)
-    setup_integrations(flask_wsgi_app)
+def init(token, *args, **kwargs):
+    client = Client(token, *args, **kwargs)
+    Engine.main.attach_client(client)
+    setup_integrations(*args, **kwargs)

tracereq/client.py

@@ -1,16 +1,20 @@
-from .constants import DESTINATION_URL
+from .constants import DESTINATION_URL, SENSITIVE_HEADERS
 from .transport import HTTPTransport
 
+
 class Client(object):
     def __init__(
             self,
             token: str = '',
-            dest: str = DESTINATION_URL,
             *args,
             **kwargs
     ):
         self._token = token
-        self._dest = dest
+        self._dest = DESTINATION_URL
+        self._exclude_url_paths = kwargs.get('exclude_url_paths', [])
+        self._exclude_urls_regex = kwargs.get('exclude_urls_regex', [])
+        self._exclude_functions = kwargs.get('exclude_functions', [])
+        self._exclude_sensitive_headers = kwargs.get('exclude_sensitive_headers', SENSITIVE_HEADERS)
         self._transport = HTTPTransport(self._token, self._dest)
 
     @property

tracereq/constants.py

@@ -1,5 +1,9 @@
-HEADER_NAME = 'TRACEREQ_TOKEN'
+HEADER_NAME = 'TRACEREQ-TOKEN'
 REQUEST = 'request'
 RESPONSE = 'response'
 TRANSPORT_SEND_TIMEOUT = 1000.0
-DESTINATION_URL = 'https://app.tracereq.com'
+DESTINATION_URL = 'https://ingest.tracereq.com'
+SENSITIVE_HEADERS = [
+    'Authorization',
+    'Cookie'
+]

tracereq/customhttplib.py

@@ -0,0 +1,63 @@
+from . import Integration
+from .engine import Engine
+from .constants import HEADER_NAME, DESTINATION_URL
+from urllib.parse import urlsplit
+
+
+try:
+    from httplib import HTTPConnection
+except ImportError:
+    from http.client import HTTPConnection
+
+
+class CustomlibIntegration(Integration):
+    integration_key = "customlib"
+
+    def __init__(self):
+        self.custom_httplib_connection = HTTPConnection
+
+    def install_httplib(self):
+        original_putrequest = self.custom_httplib_connection.putrequest
+        original_getresponse = self.custom_httplib_connection.getresponse
+
+        def custom_putrequest(self, method, url, *args, **kwargs):
+            self.engine = Engine.current
+            req = original_putrequest(self, method, url, *args, **kwargs)
+            if DESTINATION_URL in self.host:
+                return req
+
+            parsed_url = urlsplit(url)
+
+            self._trace = self.engine.get_trace()
+
+            span = self._trace.start_new_span(span_type='http.client')
+            self._trace.set_span_data(span, 'http.method', method)
+            self._trace.set_span_data(span, 'http.path', parsed_url.path)
+            self._trace.set_span_data(span, 'http.query', parsed_url.query)
+
+            if self._trace:
+                self.putheader(HEADER_NAME, span.construct_trace_header_info())
+
+            self._span = span
+            return req
+
+        def custom_getresponse(self, *args, **kwargs):
+            res = original_getresponse(self, *args, **kwargs)
+            if DESTINATION_URL in self.host:
+                return res
+
+            self._trace = self.engine.get_trace()
+
+            current_span_event = self._span
+            self._trace.mark_finish(current_span_event)
+            self._trace.set_span_data(current_span_event, 'reason', res.reason)
+            self._trace.set_span_data(current_span_event, 'status_code', res.status)
+
+            self.engine.client.transport.send(current_span_event.to_json())
+            return res
+
+        HTTPConnection.putrequest = custom_putrequest
+        HTTPConnection.getresponse = custom_getresponse
+
+    def install(self):
+        self.install_httplib()

tracereq/customlib.py

@@ -0,0 +1,61 @@
+from threading import local
+
+_local = local()
+
+
+class ParentEngine(type):
+    @property
+    def current(cls):
+        try:
+            existing_engine = _local.engine
+        except AttributeError:
+            _local.engine = existing_engine = Engine(MAIN_ENGINE)
+        return existing_engine
+
+    @property
+    def main(cls):
+        return MAIN_ENGINE
+
+
+def with_metaclass(meta, *bases):
+    class metaclass(type):
+        def __new__(cls, name, this_bases, d):
+            return meta(name, bases, d)
+
+    return type.__new__(metaclass, "temporary_class", (), {})
+
+
+class Engine(with_metaclass(ParentEngine)):
+    def __init__(self, _engine=None):
+        if isinstance(_engine, Engine):
+            engine = _engine
+            client = engine._client
+        else:
+            client = _engine
+
+        self._client = client
+        self._trace = None
+        self._request_in_process = None
+
+    def set_trace(self, trace):
+        self._trace = trace
+
+    def get_trace(self):
+        return self._trace
+
+    def set_request_data(self, request_data):
+        self._request_in_process = request_data
+
+    def get_request_data(self):
+        return self._request_in_process
+
+    @property
+    def client(self):
+        return self._client
+
+    def attach_client(self, client):
+        self._client = client
+
+
+MAIN_ENGINE = Engine()
+_local.engine = MAIN_ENGINE

tracereq/engine.py

@@ -1,21 +1,22 @@
-from flask import Flask, request
+from flask import Flask, _request_ctx_stack
 from flask.signals import request_started, request_finished
 from . import Integration
-from .hub import Hub
-from .tracing import generate_trace_req_packet, Trace
+from .engine import Engine
+from .tracing import generate_trace_event, Trace, generalize_request
 
 
-class FlaskLibIntegration(Integration):
-    identifier = "flasklib"
+class FlasklibIntegration(Integration):
+    integration_key = "flasklib"
 
     def install(self):
         request_started.connect(_request_started)
+        request_finished.connect(_request_finished)
 
         def tracereq_wsgi_app(self, environ, start_response):
-            hub = Hub.current
+            engine = Engine.current
 
-            trace_id = Trace.get_trace_from_headers(environ)
-            hub.set_trace(trace_id)
+            trace = Trace.get_trace_from_environ(environ, span_type='http.server')
+            engine.set_trace(trace)
             return self.wsgi_app(environ, start_response)
 
         Flask.__call__ = tracereq_wsgi_app
@@ -25,7 +26,18 @@ def __init__(self, wsgi_app):
 
 
 def _request_started(app, **kwargs):
-    hub = Hub.current
-    current_trace = hub.get_trace()
-    current_req_packet = generate_trace_req_packet(current_trace, request)
-    hub.client.transport.send(current_req_packet)
+    engine = Engine.current
+    current_request_data = generalize_request(_request_ctx_stack.top.request)
+
+    engine.set_request_data(current_request_data)
+
+
+def _request_finished(app, response, **kwargs):
+    engine = Engine.current
+    current_trace = engine.get_trace()
+    Trace.mark_finish(current_trace)
+
+    current_request_data = engine.get_request_data()
+    current_trace_event = generate_trace_event(current_trace, current_request_data, response)
+
+    engine.client.transport.send(current_trace_event)