fix(sanic): support latest version (#2240)

ddtrace/contrib/sanic/patch.py

@@ -6,6 +6,7 @@
 from ddtrace import config
 from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
 from ddtrace.ext import SpanTypes
+from ddtrace.pin import Pin
 from ddtrace.utils.wrappers import unwrap as _u
 from ddtrace.vendor import wrapt
 from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
@@ -18,34 +19,37 @@
 
 config._add("sanic", dict(_default_service="sanic", distributed_tracing=True))
 
+SANIC_PRE_21 = None
+
+
+def update_span(span, response):
+    if isinstance(response, sanic.response.BaseHTTPResponse):
+        status_code = response.status
+        response_headers = response.headers
+    else:
+        # invalid response causes ServerError exception which must be handled
+        status_code = 500
+        response_headers = None
+    trace_utils.set_http_meta(span, config.sanic, status_code=status_code, response_headers=response_headers)
+
 
 def _wrap_response_callback(span, callback):
-    # wrap response callbacks (either sync or async function) to set span tags
-    # based on response and finish span before returning response
-
-    def update_span(response):
-        if isinstance(response, sanic.response.BaseHTTPResponse):
-            status_code = response.status
-            response_headers = response.headers
-        else:
-            # invalid response causes ServerError exception which must be handled
-            status_code = 500
-            response_headers = None
-        trace_utils.set_http_meta(span, config.sanic, status_code=status_code, response_headers=response_headers)
-        span.finish()
+    # Only for sanic 20 and older
+    # Wrap response callbacks (either sync or async function) to set HTTP
+    # response span tags
 
     @wrapt.function_wrapper
     def wrap_sync(wrapped, instance, args, kwargs):
         r = wrapped(*args, **kwargs)
         response = args[0]
-        update_span(response)
+        update_span(span, response)
         return r
 
     @wrapt.function_wrapper
     async def wrap_async(wrapped, instance, args, kwargs):
         r = await wrapped(*args, **kwargs)
         response = args[0]
-        update_span(response)
+        update_span(span, response)
         return r
 
     if asyncio.iscoroutinefunction(callback):
@@ -54,6 +58,18 @@ async def wrap_async(wrapped, instance, args, kwargs):
     return wrap_sync(callback)
 
 
+async def patch_request_respond(wrapped, instance, args, kwargs):
+    # Only for sanic 21 and newer
+    # Wrap the framework response to set HTTP response span tags
+    response = await wrapped(*args, **kwargs)
+    pin = Pin._find(instance.ctx)
+    if pin is not None and pin.enabled():
+        span = pin.tracer.current_span()
+        if span is not None:
+            update_span(span, response)
+    return response
+
+
 def _get_path(request):
     """Get path and replace path parameter values with names if route exists."""
     path = request.path
@@ -66,57 +82,83 @@ def _get_path(request):
     return path
 
 
+async def patch_run_request_middleware(wrapped, instance, args, kwargs):
+    # Set span resource from the framework request
+    request = args[0]
+    pin = Pin._find(request.ctx)
+    if pin is not None and pin.enabled():
+        span = pin.tracer.current_span()
+        if span is not None:
+            span.resource = "{} {}".format(request.method, _get_path(request))
+    return await wrapped(*args, **kwargs)
+
+
 def patch():
     """Patch the instrumented methods."""
+    global SANIC_PRE_21
+
     if getattr(sanic, "__datadog_patch", False):
         return
     setattr(sanic, "__datadog_patch", True)
+
+    SANIC_PRE_21 = sanic.__version__[:2] < "21"
+
     _w("sanic", "Sanic.handle_request", patch_handle_request)
+    _w("sanic", "Sanic._run_request_middleware", patch_run_request_middleware)
+    if not SANIC_PRE_21:
+        _w(sanic.request, "Request.respond", patch_request_respond)
 
 
 def unpatch():
     """Unpatch the instrumented methods."""
     _u(sanic.Sanic, "handle_request")
+    _u(sanic.Sanic, "_run_request_middleware")
+    if not SANIC_PRE_21:
+        _u(sanic.request.Request, "respond")
     if not getattr(sanic, "__datadog_patch", False):
         return
     setattr(sanic, "__datadog_patch", False)
 
 
 async def patch_handle_request(wrapped, instance, args, kwargs):
     """Wrapper for Sanic.handle_request"""
-    request = kwargs.get("request", args[0])
-    write_callback = kwargs.get("write_callback", args[1])
-    stream_callback = kwargs.get("stream_callback", args[2])
+
+    def unwrap(request, write_callback=None, stream_callback=None, **kwargs):
+        return request, write_callback, stream_callback, kwargs
+
+    request, write_callback, stream_callback, new_kwargs = unwrap(*args, **kwargs)
 
     if request.scheme not in ("http", "https"):
-        return await wrapped(request, write_callback, stream_callback, **kwargs)
+        return await wrapped(*args, **kwargs)
 
-    resource = "{} {}".format(request.method, _get_path(request))
+    pin = Pin()
+    pin.onto(request.ctx)
 
     headers = request.headers.copy()
 
     trace_utils.activate_distributed_headers(ddtrace.tracer, int_config=config.sanic, request_headers=headers)
 
-    span = ddtrace.tracer.trace(
+    with pin.tracer.trace(
         "sanic.request",
         service=trace_utils.int_service(None, config.sanic),
-        resource=resource,
         span_type=SpanTypes.WEB,
-    )
-    sample_rate = config.sanic.get_analytics_sample_rate(use_global_config=True)
-    if sample_rate is not None:
-        span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, sample_rate)
-
-    method = request.method
-    url = "{scheme}://{host}{path}".format(scheme=request.scheme, host=request.host, path=request.path)
-    query_string = request.query_string
-    if isinstance(query_string, bytes):
-        query_string = query_string.decode()
-    trace_utils.set_http_meta(span, config.sanic, method=method, url=url, query=query_string, request_headers=headers)
-
-    if write_callback is not None:
-        write_callback = _wrap_response_callback(span, write_callback)
-    if stream_callback is not None:
-        stream_callback = _wrap_response_callback(span, stream_callback)
-
-    return await wrapped(request, write_callback, stream_callback, **kwargs)
+    ) as span:
+        sample_rate = config.sanic.get_analytics_sample_rate(use_global_config=True)
+        if sample_rate is not None:
+            span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, sample_rate)
+
+        method = request.method
+        url = "{scheme}://{host}{path}".format(scheme=request.scheme, host=request.host, path=request.path)
+        query_string = request.query_string
+        if isinstance(query_string, bytes):
+            query_string = query_string.decode()
+        trace_utils.set_http_meta(
+            span, config.sanic, method=method, url=url, query=query_string, request_headers=headers
+        )
+
+        if write_callback is not None:
+            new_kwargs["write_callback"] = _wrap_response_callback(span, write_callback)
+        if stream_callback is not None:
+            new_kwargs["stream_callback"] = _wrap_response_callback(span, stream_callback)
+
+        return await wrapped(request, **new_kwargs)

tests/contrib/sanic/test_sanic.py

@@ -4,6 +4,7 @@
 
 import pytest
 from sanic import Sanic
+from sanic.config import DEFAULT_CONFIG
 from sanic.exceptions import ServerError
 from sanic.response import json
 from sanic.response import stream
@@ -38,8 +39,11 @@ async def _response_text(response):
     return resp_text
 
 
-@pytest.yield_fixture
+@pytest.fixture
 def app(tracer):
+    # Sanic 20.12 and newer prevent loading multiple applications
+    # with the same name if register is True.
+    DEFAULT_CONFIG["REGISTER"] = False
     app = Sanic(__name__)
 
     @tracer.wrap()
@@ -283,7 +287,7 @@ async def test_multiple_requests(tracer, client, test_spans):
 async def test_invalid_response_type_str(tracer, client, test_spans):
     response = await client.get("/invalid")
     assert _response_status(response) == 500
-    assert await _response_text(response) == "Invalid response type"
+    assert (await _response_text(response)).startswith("Invalid response type")
 
     spans = test_spans.pop_traces()
     assert len(spans) == 1
@@ -301,7 +305,7 @@ async def test_invalid_response_type_str(tracer, client, test_spans):
 async def test_invalid_response_type_empty(tracer, client, test_spans):
     response = await client.get("/empty")
     assert _response_status(response) == 500
-    assert await _response_text(response) == "Invalid response type"
+    assert (await _response_text(response)).startswith("Invalid response type")
 
     spans = test_spans.pop_traces()
     assert len(spans) == 1

tox.ini

@@ -93,7 +93,7 @@ envlist =
     pyramid_contrib{,_autopatch}-py{27,35,36,37,38,39}-pyramid{17,18,19,110,}-webtest
     redis_contrib-py{27,35,36,37,38,39}-redis{210,30,32,33,34,35,}
     rediscluster_contrib-py{27,35,36,37,38,39}-rediscluster{135,136,200,}-redis210
-    sanic_contrib-py{37,38,39}-sanic{1906,1909,1912,2003,2006}
+    sanic_contrib-py{37,38,39}-sanic{1906,1909,1912,2003,2006,2103}
     sqlite3_contrib-py{27,35,36,37,38,39}-sqlite3
     tornado_contrib-py{27,35,36,37,38,39}-tornado{44,45}
     tornado_contrib-py{37,38,39}-tornado{50,51,60,}
@@ -311,6 +311,8 @@ deps =
     sanic1912: sanic~=19.12.0
     sanic2003: sanic~=20.3.0
     sanic2006: sanic~=20.6.0
+    sanic2103: sanic~=21.3.0
+    sanic2103: httpx
     sanic: sanic
     sqlalchemy: sqlalchemy
     sslmodules3: aiohttp