Add starlette ServerErrorMiddleware instrumentation

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