Add starlette ServerErrorMiddleware instrumentation (#1674)

* Add starlette ServerErrorMiddleware instrumentation

* Set creates_transactions to True so the instrumentation works

* Add a test

* CHANGELOG

Author: basepi

CHANGELOG.asciidoc

@@ -47,6 +47,7 @@ endif::[]
 * Fix Django's `manage.py check` when agent is disabled {pull}1632[#1632]
 * Fix an issue with long body truncation for Starlette {pull}1635[#1635]
 * Fix an issue with transaction outcomes in Flask for uncaught exceptions {pull}1637[#1637]
+* FIx Starlette instrumentation to make sure transaction information is still present during exception handling {pull}1674[#1674]
 
 
 [[release-notes-6.x]]

elasticapm/contrib/starlette/__init__.py

@@ -185,6 +185,8 @@ async def request_receive() -> Message:
         request = Request(scope, receive=_mocked_receive or receive)
         await self._request_started(request)
 
+        # We don't end the transaction here, we rely on the starlette
+        # instrumentation of ServerErrorMiddleware to end the transaction
         try:
             await self.app(scope, _request_receive or receive, wrapped_send)
             elasticapm.set_transaction_outcome(constants.OUTCOME.SUCCESS, override=False)
@@ -197,8 +199,6 @@ async def request_receive() -> Message:
             elasticapm.set_context({"status_code": 500}, "response")
 
             raise
-        finally:
-            self.client.end_transaction()
 
     async def capture_exception(self, *args, **kwargs):
         """Captures your exception.

elasticapm/instrumentation/packages/asyncio/starlette.py

@@ -0,0 +1,52 @@
+#  BSD 3-Clause License
+#
+#  Copyright (c) 2019, Elasticsearch BV
+#  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  * Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+#
+#  * Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+#  * Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from elasticapm import get_client
+from elasticapm.instrumentation.packages.asyncio.base import AsyncAbstractInstrumentedModule
+
+
+class StarletteServerErrorMiddlewareInstrumentation(AsyncAbstractInstrumentedModule):
+    name = "starlette"
+
+    instrument_list = [("starlette.middleware.errors", "ServerErrorMiddleware.__call__")]
+
+    # This instrumentation doesn't actually create transactions. However, it
+    # does wrap a context outside of the normal starlette transaction, so we
+    # need to make sure it always calls the wrapped version even if a sampled
+    # transaction is not active.
+    creates_transactions = True
+
+    async def call(self, module, method, wrapped, instance, args, kwargs):
+        try:
+            return await wrapped(*args, **kwargs)
+        finally:
+            client = get_client()
+            if client:
+                client.end_transaction()

elasticapm/instrumentation/register.py

@@ -88,6 +88,7 @@
             "elasticapm.instrumentation.packages.asyncio.aioredis.RedisConnectionInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.aiomysql.AioMySQLInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.aiobotocore.AioBotocoreInstrumentation",
+            "elasticapm.instrumentation.packages.asyncio.starlette.StarletteServerErrorMiddlewareInstrumentation",
         ]
     )
 

tests/contrib/asyncio/starlette_tests.py

@@ -28,10 +28,13 @@
 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-import pytest  # isort:skip
+from shutil import ExecError
 
 from tests.fixtures import TempStoreClient
 
+import pytest  # isort:skip
+
+
 starlette = pytest.importorskip("starlette")  # isort:skip
 
 import os
@@ -65,6 +68,12 @@ def app(elasticapm_client):
     app.mount("/sub", sub)
     sub.mount("/subsub", subsub)
 
+    @app.exception_handler(Exception)
+    async def handle_exception(request, exc):
+        transaction_id = elasticapm.get_transaction_id()
+        exc.transaction_id = transaction_id
+        return PlainTextResponse(f"{transaction_id}", status_code=500)
+
     @app.route("/", methods=["GET", "POST"])
     async def hi(request):
         body = await request.body()
@@ -88,6 +97,11 @@ async def raise_exception(request):
         await request.body()
         raise ValueError()
 
+    @app.route("/raise-base-exception", methods=["GET", "POST"])
+    async def raise_base_exception(request):
+        await request.body()
+        raise Exception()
+
     @app.route("/hi/{name}/with/slash/", methods=["GET", "POST"])
     async def with_slash(request):
         return PlainTextResponse("Hi {}".format(request.path_params["name"]))
@@ -508,3 +522,15 @@ def test_websocket(app, elasticapm_client):
         assert data == "Hello, world!"
 
     assert len(elasticapm_client.events[constants.TRANSACTION]) == 0
+
+
+def test_transaction_active_in_base_exception_handler(app, elasticapm_client):
+    client = TestClient(app)
+    try:
+        response = client.get("/raise-base-exception")
+    except Exception as exc:
+        # This is set by the exception handler -- we want to make sure the
+        # handler has access to the transaction.
+        assert exc.transaction_id
+
+    assert len(elasticapm_client.events[constants.TRANSACTION]) == 1