feat(aiohttp): Apply data_collection filtering to URL query strings

sentry_sdk/integrations/aiohttp.py

@@ -4,7 +4,10 @@
 
 import sentry_sdk
 from sentry_sdk.api import continue_trace
 from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
+from sentry_sdk.data_collection import (
+    _apply_data_collection_filtering_to_query_string,
+)
 from sentry_sdk.integrations import (
     _DEFAULT_FAILED_REQUEST_STATUS_CODES,
     DidNotEnable,
@@ -46,6 +49,7 @@
     capture_internal_exceptions,
     ensure_integration_enabled,
     event_from_exception,
+    has_data_collection_enabled,
     logger,
     parse_url,
     parse_version,
@@ -161,7 +165,26 @@
                             )
 
                         url_attributes = {}
-                        if should_send_default_pii():
+                        if has_data_collection_enabled(client.options):
+                            url_attributes["url.full"] = "%s://%s%s" % (
+                                request.scheme,
+                                request.host,
+                                request.path,
+                            )
+                            url_attributes["url.path"] = request.path
+
+                            if request.query_string:
+                                filtered_query_string = (
+                                    _apply_data_collection_filtering_to_query_string(
+                                        query_string=request.query_string,
+                                        behaviour=client.options["data_collection"][
+                                            "url_query_params"
+                                        ],
+                                    )
+                                )
+                                if filtered_query_string:
+                                    url_attributes["url.query"] = filtered_query_string
+                        elif should_send_default_pii():
                             url_attributes["url.full"] = "%s://%s%s" % (
                                 request.scheme,
                                 request.host,
@@ -362,14 +385,33 @@
                     "sentry.origin": AioHttpIntegration.origin,
                     "http.request.method": method,
                 }
-                if parsed_url is not None and should_send_default_pii():
-                    attributes["url.full"] = parsed_url.url
-                    attributes["url.path"] = params.url.path
-
-                    if parsed_url.query:
-                        attributes["url.query"] = parsed_url.query
-                    if parsed_url.fragment:
-                        attributes["url.fragment"] = parsed_url.fragment
+                if parsed_url is not None:
+                    if has_data_collection_enabled(client.options):
+                        attributes["url.full"] = parsed_url.url
+                        attributes["url.path"] = params.url.path
+
+                        if parsed_url.fragment:
+                            attributes["url.fragment"] = parsed_url.fragment
+
+                        if parsed_url.query:
+                            filtered_query = (
+                                _apply_data_collection_filtering_to_query_string(
+                                    query_string=parsed_url.query,
+                                    behaviour=client.options["data_collection"][
+                                        "url_query_params"
+                                    ],
+                                )
+                            )
+                            if filtered_query:
+                                attributes["url.query"] = filtered_query
+                    elif should_send_default_pii():
+                        attributes["url.full"] = parsed_url.url
+                        attributes["url.path"] = params.url.path
+
+                        if parsed_url.query:
+                            attributes["url.query"] = parsed_url.query
+                        if parsed_url.fragment:
+                            attributes["url.fragment"] = parsed_url.fragment
 
                 span = sentry_sdk.traces.start_span(
                     name=span_name, attributes=attributes

tests/integrations/aiohttp/test_aiohttp.py

@@ -1795,3 +1795,164 @@ async def hello(request):
     else:
         assert "user.ip_address" not in server_span["attributes"]
         assert "user.ip_address" not in child_span["attributes"]
+
+
+_QUERY_PARAM_DATA_COLLECTION_CASES = [
+    pytest.param(
+        {"send_default_pii": True},
+        "toy=tennisball&color=red&auth=secret",
+        id="send_default_pii_true",
+    ),
+    pytest.param(
+        {"send_default_pii": False},
+        None,
+        id="send_default_pii_false",
+    ),
+    pytest.param(
+        {},
+        None,
+        id="defaults",
+    ),
+    pytest.param(
+        {"_experiments": {"data_collection": {}}},
+        "toy=tennisball&color=red&auth=%5BFiltered%5D",
+        id="data_collection_denylist_default",
+    ),
+    pytest.param(
+        {
+            "_experiments": {
+                "data_collection": {
+                    "url_query_params": {"mode": "denylist", "terms": ["toy"]}
+                }
+            }
+        },
+        "toy=%5BFiltered%5D&color=red&auth=%5BFiltered%5D",
+        id="data_collection_denylist_custom_terms",
+    ),
+    pytest.param(
+        {
+            "_experiments": {
+                "data_collection": {
+                    "url_query_params": {"mode": "allowlist", "terms": ["toy"]}
+                }
+            }
+        },
+        "toy=tennisball&color=%5BFiltered%5D&auth=%5BFiltered%5D",
+        id="data_collection_allowlist",
+    ),
+    pytest.param(
+        {
+            "_experiments": {
+                "data_collection": {
+                    "url_query_params": {"mode": "allowlist", "terms": ["auth"]}
+                }
+            }
+        },
+        "toy=%5BFiltered%5D&color=%5BFiltered%5D&auth=%5BFiltered%5D",
+        id="data_collection_allowlist_sensitive_term",
+    ),
+    pytest.param(
+        {"_experiments": {"data_collection": {"url_query_params": {"mode": "off"}}}},
+        None,
+        id="data_collection_off",
+    ),
+    pytest.param(
+        {
+            "send_default_pii": True,
+            "_experiments": {"data_collection": {"url_query_params": {"mode": "off"}}},
+        },
+        None,
+        id="data_collection_wins_over_send_default_pii",
+    ),
+]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "init_kwargs, expected_query", _QUERY_PARAM_DATA_COLLECTION_CASES
+)
+async def test_server_url_query_data_collection_span_streaming(
+    sentry_init, aiohttp_client, capture_items, init_kwargs, expected_query
+):
+    init_kwargs = dict(init_kwargs)
+    sentry_init(
+        integrations=[AioHttpIntegration()],
+        traces_sample_rate=1.0,
+        _experiments={
+            "trace_lifecycle": "stream",
+            **init_kwargs.pop("_experiments", {}),
+        },
+        **init_kwargs,
+    )
+
+    async def hello(request):
+        return web.Response(text="hello")
+
+    app = web.Application()
+    app.router.add_get("/", hello)
+
+    items = capture_items("span")
+
+    client = await aiohttp_client(app)
+    resp = await client.get("/?toy=tennisball&color=red&auth=secret")
+    assert resp.status == 200
+
+    sentry_sdk.flush()
+
+    (server_span,) = [item.payload for item in items]
+
+    if expected_query is None:
+        assert "url.query" not in server_span["attributes"]
+    else:
+        assert server_span["attributes"]["url.query"] == expected_query
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "init_kwargs, expected_query", _QUERY_PARAM_DATA_COLLECTION_CASES
+)
+async def test_client_url_query_data_collection_span_streaming(
+    sentry_init,
+    aiohttp_raw_server,
+    aiohttp_client,
+    capture_items,
+    init_kwargs,
+    expected_query,
+):
+    init_kwargs = dict(init_kwargs)
+    sentry_init(
+        integrations=[AioHttpIntegration()],
+        traces_sample_rate=1.0,
+        _experiments={
+            "trace_lifecycle": "stream",
+            **init_kwargs.pop("_experiments", {}),
+        },
+        **init_kwargs,
+    )
+
+    async def handler(request):
+        return web.Response(text="OK")
+
+    raw_server = await aiohttp_raw_server(handler)
+
+    async def hello(request):
+        span_client = await aiohttp_client(raw_server)
+        await span_client.get("/?toy=tennisball&color=red&auth=secret")
+        return web.Response(text="hello")
+
+    app = web.Application()
+    app.router.add_get(r"/", hello)
+
+    items = capture_items("span")
+
+    client = await aiohttp_client(app)
+    await client.get("/")
+
+    sentry_sdk.flush()
+
+    inner_client_span = items[0].payload
+
+    if expected_query is None:
+        assert "url.query" not in inner_client_span["attributes"]
+    else:
+        assert inner_client_span["attributes"]["url.query"] == expected_query