Events for web frameworks

ddtrace/_events.py

@@ -0,0 +1,104 @@
+from typing import Callable
+from typing import Mapping
+from typing import Optional
+from typing import Type
+from typing import Union
+
+import attr
+
+from ddtrace import Span
+from ddtrace import _hooks
+from ddtrace import config
+from ddtrace.internal import compat
+
+
+_HOOKS = _hooks.Hooks()
+
+
+@attr.s(frozen=True)
+class IntegrationEvent(object):
+    """
+    An IntegrationEvent is emitted by an integration (e.g. the flask framework integration)
+    and is linked to a span.
+    """
+    span = attr.ib(type=Span)
+    integration = attr.ib(type=str)
+
+    def emit(self):
+        # type: () -> None
+        """Alias for emitting this event."""
+        emit(self)
+
+    @classmethod
+    def register(cls, func=None):
+        # type: (Optional[Callable]) -> Optional[Callable]
+        """Alias for registering a listener for this event type."""
+        return register(cls, func=func)
+
+    @classmethod
+    def deregister(cls, func):
+        # type: (Callable) -> None
+        """Alias for deregistering a listener for this event type."""
+        deregister(cls, func)
+
+    if config._raise:
+        @span.validator  # type: ignore
+        def check_span(self, attribute, value):
+            assert isinstance(value, Span)
+
+        @integration.validator  # type: ignore
+        def check_integration(self, attribute, value):
+            assert value in config._config
+
+
+@attr.s(frozen=True)
+class WebRequest(IntegrationEvent):
+    """
+    The WebRequest event is emitted by web framework integrations before the WebResponse event.
+    """
+    method = attr.ib(type=str)
+    url = attr.ib(type=str)
+    headers = attr.ib(type=Mapping[str, str], factory=dict)
+    query = attr.ib(type=Optional[str], default=None)
+
+    if config._raise:
+        @method.validator  # type: ignore
+        def check_method(self, attribute, value):
+            assert value in ("HEAD", "GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "PROPFIND", "TRACE", "CONNECT")
+
+        @url.validator  # type: ignore
+        def check_url(self, attribute, value):
+            compat.parse.urlparse(value)
+
+
+@attr.s(frozen=True)
+class WebResponse(IntegrationEvent):
+    """
+    The WebResponse event is emitted by web frameworks after the WebRequest event.
+    """
+    status_code = attr.ib(type=Union[int, str])
+    status_msg = attr.ib(type=Optional[str], default=None)
+    headers = attr.ib(type=Mapping[str, str], factory=dict)
+
+    if config._raise:
+        @status_code.validator  # type: ignore
+        def check_status_code(self, attribute, value):
+            int(value)
+
+
+def emit(event):
+    # type: (IntegrationEvent) -> None
+    """Notify registered listeners about an event."""
+    _HOOKS.emit(event.__class__, event)
+
+
+def register(event_type, func=None):
+    # type: (Type[IntegrationEvent], Optional[Callable]) -> Optional[Callable]
+    """Register a function for a specific event type."""
+    return _HOOKS.register(event_type, func)
+
+
+def deregister(event_type, func):
+    # type: (Type[IntegrationEvent], Callable) -> None
+    """Deregister a function for an event type."""
+    _HOOKS.deregister(event_type, func)

ddtrace/_hooks.py

@@ -129,5 +129,9 @@ def emit(
         for func in self._hooks.get(hook, ()):
             try:
                 func(*args, **kwargs)
-            except Exception:
+            except Exception as e:
+                from ddtrace import config
+
+                if config._raise:
+                    raise e
                 log.error("Failed to run hook %s function %s", hook, func, exc_info=True)

ddtrace/contrib/aiohttp/middlewares.py

@@ -1,3 +1,4 @@
+from ddtrace import _events
 from ddtrace import config
 
 from .. import trace_utils
@@ -57,6 +58,16 @@ async def attach_context(request):
         request[REQUEST_CONTEXT_KEY] = request_span.context
         request[REQUEST_SPAN_KEY] = request_span
         request[REQUEST_CONFIG_KEY] = app[CONFIG_KEY]
+
+        _events.WebRequest(
+            span=request_span,
+            method=request.method,
+            url=str(request.url),  # DEV: request.url is a yarl's URL object
+            headers=request.headers,
+            query=request.query_string,
+            integration=config.aiohttp.integration_name,
+        ).emit()
+
         try:
             response = await handler(request)
             return response
@@ -98,20 +109,17 @@ async def on_prepare(request, response):
 
     # DEV: aiohttp is special case maintains separate configuration from config api
     trace_query_string = request[REQUEST_CONFIG_KEY].get("trace_query_string")
-    if trace_query_string is None:
-        trace_query_string = config.http.trace_query_string
-    if trace_query_string:
+    if trace_query_string is True:
         request_span.set_tag(http.QUERY_STRING, request.query_string)
+    elif trace_query_string is False:
+        request_span.set_tag(http.QUERY_STRING, None)
 
-    trace_utils.set_http_meta(
-        request_span,
-        config.aiohttp,
-        method=request.method,
-        url=str(request.url),  # DEV: request.url is a yarl's URL object
+    _events.WebResponse(
+        span=request_span,
         status_code=response.status,
-        request_headers=request.headers,
-        response_headers=response.headers,
-    )
+        headers=response.headers,
+        integration=config.aiohttp.integration_name,
+    ).emit()
 
     request_span.finish()
 

ddtrace/contrib/asgi/middleware.py

@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING
 
 import ddtrace
+from ddtrace import _events
 from ddtrace import config
 from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
 from ddtrace.ext import SpanTypes
@@ -138,34 +139,37 @@ async def __call__(self, scope, receive, send):
         else:
             url = None
 
-        if self.integration_config.trace_query_string:
-            query_string = scope.get("query_string")
-            if len(query_string) > 0:
-                query_string = bytes_to_str(query_string)
-        else:
-            query_string = None
+        query_string = scope.get("query_string")
+        if len(query_string) > 0:
+            query_string = bytes_to_str(query_string)
 
-        trace_utils.set_http_meta(
-            span, self.integration_config, method=method, url=url, query=query_string, request_headers=headers
-        )
+        _events.WebRequest(
+            span=span,
+            method=method,
+            url=url,
+            headers=headers,
+            query=query_string,
+            integration=self.integration_config.integration_name,
+        ).emit()
 
         tags = _extract_versions_from_scope(scope, self.integration_config)
         span.set_tags(tags)
 
         async def wrapped_send(message):
-            if span and message.get("type") == "http.response.start" and "status" in message:
-                status_code = message["status"]
-            else:
-                status_code = None
-
-            if "headers" in message:
-                response_headers = message["headers"]
-            else:
-                response_headers = None
-
-            trace_utils.set_http_meta(
-                span, self.integration_config, status_code=status_code, response_headers=response_headers
-            )
+            status = message.get("status")
+            response_headers = message.get("headers")
+            if (
+                span
+                and message.get("type") == "http.response.start"
+                and status is not None
+                and response_headers is not None
+            ):
+                _events.WebResponse(
+                    span=span,
+                    status_code=status,
+                    headers=response_headers,
+                    integration=self.integration_config.integration_name,
+                ).emit()
 
             return await send(message)
 

ddtrace/contrib/bottle/trace.py

@@ -4,6 +4,7 @@
 from bottle import response
 
 import ddtrace
+from ddtrace import _events
 from ddtrace import config
 
 from .. import trace_utils
@@ -52,6 +53,18 @@ def wrapped(*args, **kwargs):
                 # set analytics sample rate with global config enabled
                 s.set_tag(ANALYTICS_SAMPLE_RATE_KEY, config.bottle.get_analytics_sample_rate(use_global_config=True))
 
+                method = request.method
+                url = request.urlparts._replace(query="").geturl()
+
+                _events.WebRequest(
+                    span=s,
+                    method=method,
+                    url=url,
+                    headers=request.headers,
+                    query=request.query_string,
+                    integration=config.bottle.integration_name,
+                ).emit()
+
                 code = None
                 result = None
                 try:
@@ -79,17 +92,11 @@ def wrapped(*args, **kwargs):
                         # will be default
                         response_code = response.status_code
 
-                    method = request.method
-                    url = request.urlparts._replace(query="").geturl()
-                    trace_utils.set_http_meta(
-                        s,
-                        config.bottle,
-                        method=method,
-                        url=url,
+                    _events.WebResponse(
+                        span=s,
                         status_code=response_code,
-                        query=request.query_string,
-                        request_headers=request.headers,
-                        response_headers=response.headers,
-                    )
+                        headers=response.headers,
+                        integration=config.bottle.integration_name,
+                    ).emit()
 
         return wrapped

ddtrace/contrib/cherrypy/middleware.py

@@ -10,6 +10,7 @@
 from cherrypy.lib.httputil import valid_status
 
 # project
+from ddtrace import _events
 from ddtrace import config
 
 from .. import trace_utils
@@ -72,12 +73,22 @@ def _on_start_resource(self):
             self._tracer, int_config=config.cherrypy, request_headers=cherrypy.request.headers
         )
 
-        cherrypy.request._datadog_span = self._tracer.trace(
+        cherrypy.request._datadog_span = span = self._tracer.trace(
             SPAN_NAME,
             service=trace_utils.int_service(None, config.cherrypy, default="cherrypy"),
             span_type=SpanTypes.WEB,
         )
 
+        url = compat.to_unicode(cherrypy.request.base + cherrypy.request.path_info)
+
+        _events.WebRequest(
+            span=span,
+            method=cherrypy.request.method,
+            url=url,
+            headers=cherrypy.request.headers,
+            integration=config.cherrypy.integration_name,
+        ).emit()
+
     def _after_error_response(self):
         span = getattr(cherrypy.request, "_datadog_span", None)
 
@@ -118,18 +129,14 @@ def _close_span(self, span):
             resource = "{} {}".format(cherrypy.request.method, cherrypy.request.path_info)
             span.resource = compat.to_unicode(resource)
 
-        url = compat.to_unicode(cherrypy.request.base + cherrypy.request.path_info)
         status_code, _, _ = valid_status(cherrypy.response.status)
 
-        trace_utils.set_http_meta(
-            span,
-            config.cherrypy,
-            method=cherrypy.request.method,
-            url=url,
+        _events.WebResponse(
+            span=span,
             status_code=status_code,
-            request_headers=cherrypy.request.headers,
-            response_headers=cherrypy.response.headers,
-        )
+            headers=cherrypy.response.headers,
+            integration=config.cherrypy.integration_name,
+        ).emit()
 
         span.finish()