Repository: scoutapp/scout_apm_python

Issue:
Instrument Starlette authentication
Starlette has a basic authentication framework: https://www.starlette.io/authentication/

We should probably copy what we do for Django here, tagging user ID's/ display names on traces.


Discussion:

Resolving patch:
diff --git a/src/scout_apm/async_/starlette.py b/src/scout_apm/async_/starlette.py
--- a/src/scout_apm/async_/starlette.py
+++ b/src/scout_apm/async_/starlette.py
@@ -55,7 +55,7 @@ async def __call__(self, scope, receive, send):
             amazon_queue_time = request.headers.get("x-amzn-trace-id", default="")
             track_amazon_request_queue_time(amazon_queue_time, tracked_request)
 
-        def rename_controller_span_from_endpoint():
+        def grab_extra_data():
             if "endpoint" in scope:
                 # Rename top span
                 endpoint = scope["endpoint"]
@@ -64,13 +64,23 @@ def rename_controller_span_from_endpoint():
                 )
                 tracked_request.is_real_request = True
 
+            # From AuthenticationMiddleware - bypass request.user because it
+            # throws AssertionError if 'user' is not in Scope, and we need a
+            # try/except already
+            try:
+                username = scope["user"].display_name
+            except (KeyError, AttributeError):
+                pass
+            else:
+                tracked_request.tag("username", username)
+
         async def wrapped_send(data):
             # Finish HTTP span when body finishes sending, not later (e.g.
             # after background tasks)
             if data.get("type", None) == "http.response.body" and not data.get(
                 "more_body", False
             ):
-                rename_controller_span_from_endpoint()
+                grab_extra_data()
                 tracked_request.stop_span()
             return await send(data)
 
@@ -81,7 +91,7 @@ async def wrapped_send(data):
             raise exc
         finally:
             if tracked_request.end_time is None:
-                rename_controller_span_from_endpoint()
+                grab_extra_data()
                 tracked_request.stop_span()
 
 


Test patch:
diff --git a/tests/integration/test_starlette_py36plus.py b/tests/integration/test_starlette_py36plus.py
--- a/tests/integration/test_starlette_py36plus.py
+++ b/tests/integration/test_starlette_py36plus.py
@@ -8,8 +8,10 @@
 import pytest
 from asgiref.testing import ApplicationCommunicator
 from starlette.applications import Starlette
+from starlette.authentication import AuthCredentials, AuthenticationBackend, SimpleUser
 from starlette.background import BackgroundTasks
 from starlette.endpoints import HTTPEndpoint
+from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.responses import PlainTextResponse
 
 from scout_apm.api import Config
@@ -24,7 +26,7 @@
 
 
 @contextmanager
-def app_with_scout(*, scout_config=None):
+def app_with_scout(*, app=None, scout_config=None):
     """
     Context manager that configures and installs the Scout plugin for a basic
     Starlette application.
@@ -35,7 +37,8 @@ def app_with_scout(*, scout_config=None):
     scout_config["core_agent_launch"] = False
     scout_config.setdefault("monitor", True)
 
-    app = Starlette()
+    if app is None:
+        app = Starlette()
 
     @app.exception_handler(500)
     async def error(request, exc):
@@ -329,3 +332,49 @@ async def test_background_jobs(tracked_requests):
         "Job/tests.integration.test_starlette_py36plus."
         + "app_with_scout.<locals>.background_jobs.<locals>.async_noop"
     )
+
+
+@async_test
+async def test_username(tracked_requests):
+    base_app = Starlette()
+
+    class DummyBackend(AuthenticationBackend):
+        async def authenticate(self, request):
+            return AuthCredentials(), SimpleUser("dummy")
+
+    base_app.add_middleware(AuthenticationMiddleware, backend=DummyBackend())
+
+    with app_with_scout(app=base_app) as app:
+        communicator = ApplicationCommunicator(app, get_scope(path="/"))
+        await communicator.send_input({"type": "http.request"})
+        response_start = await communicator.receive_output()
+        await communicator.receive_output()
+
+    assert response_start["type"] == "http.response.start"
+    assert response_start["status"] == 200
+    assert len(tracked_requests) == 1
+    tracked_request = tracked_requests[0]
+    assert tracked_request.tags["username"] == "dummy"
+
+
+@async_test
+async def test_username_bad_user(tracked_requests):
+    base_app = Starlette()
+
+    class BadUserBackend(AuthenticationBackend):
+        async def authenticate(self, request):
+            return AuthCredentials(), object()
+
+    base_app.add_middleware(AuthenticationMiddleware, backend=BadUserBackend())
+
+    with app_with_scout(app=base_app) as app:
+        communicator = ApplicationCommunicator(app, get_scope(path="/"))
+        await communicator.send_input({"type": "http.request"})
+        response_start = await communicator.receive_output()
+        await communicator.receive_output()
+
+    assert response_start["type"] == "http.response.start"
+    assert response_start["status"] == 200
+    assert len(tracked_requests) == 1
+    tracked_request = tracked_requests[0]
+    assert "username" not in tracked_request.tags