feat(boto3): Support data_collection filtering for URL query params

sentry_sdk/integrations/boto3.py

@@ -3,6 +3,9 @@
 
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.data_collection import (
+    _apply_data_collection_filtering_to_query_string,
+)
 from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.traces import StreamedSpan
@@ -15,6 +18,7 @@
 )
 from sentry_sdk.utils import (
     capture_internal_exceptions,
+    has_data_collection_enabled,
     parse_url,
     parse_version,
 )
@@ -24,6 +28,10 @@
 
     from botocore.model import ServiceId
 
+    from sentry_sdk._types import Attributes
+    from sentry_sdk.client import BaseClient as SentryClient
+    from sentry_sdk.utils import ParsedUrl
+
 try:
     from botocore import __version__ as BOTOCORE_VERSION
     from botocore.awsrequest import AWSRequest
@@ -62,6 +70,40 @@ def sentry_patched_init(
         BaseClient.__init__ = sentry_patched_init  # type: ignore
 
 
+def _get_url_attributes(
+    client: "SentryClient", parsed_url: "Optional[ParsedUrl]"
+) -> "Attributes":
+    attributes: "Attributes" = {}
+    if parsed_url is None:
+        return attributes
+
+    query: "Optional[str]"
+    if has_data_collection_enabled(client.options):
+        query = None
+        if parsed_url.query:
+            query = _apply_data_collection_filtering_to_query_string(
+                query_string=parsed_url.query,
+                behaviour=client.options["data_collection"]["url_query_params"],
+            )
+    elif should_send_default_pii():
+        query = parsed_url.query
+    else:
+        return attributes
+
+    url_full = parsed_url.url
+    if query:
+        attributes[SPANDATA.URL_QUERY] = query
+        url_full += "?" + query
+
+    if parsed_url.fragment:
+        attributes[SPANDATA.URL_FRAGMENT] = parsed_url.fragment
+        url_full += "#" + parsed_url.fragment
+
+    attributes[SPANDATA.URL_FULL] = url_full
+
+    return attributes
+
+
 def _sentry_request_created(
     service_id: "ServiceId", request: "AWSRequest", operation_name: str, **kwargs: "Any"
 ) -> None:
@@ -81,14 +123,8 @@ def _sentry_request_created(
     is_span_streaming_enabled = has_span_streaming_enabled(client.options)
     span: "Union[Span, StreamedSpan, None]" = None
     if is_span_streaming_enabled:
-        if parsed_url and should_send_default_pii():
-            breadcrumb.update(
-                {
-                    SPANDATA.URL_FULL: parsed_url.url,
-                    SPANDATA.URL_QUERY: parsed_url.query,
-                    SPANDATA.URL_FRAGMENT: parsed_url.fragment,
-                }
-            )
+        url_attributes = _get_url_attributes(client, parsed_url)
+        breadcrumb.update(url_attributes)
 
         if request.method is not None:
             breadcrumb[SPANDATA.HTTP_REQUEST_METHOD] = request.method
@@ -102,14 +138,7 @@ def _sentry_request_created(
                     SPANDATA.RPC_METHOD: f"{service_id}/{operation_name}",
                 },
             )
-            if parsed_url and should_send_default_pii():
-                span.set_attributes(
-                    {
-                        SPANDATA.URL_FULL: parsed_url.url,
-                        SPANDATA.URL_QUERY: parsed_url.query,
-                        SPANDATA.URL_FRAGMENT: parsed_url.fragment,
-                    }
-                )
+            span.set_attributes(url_attributes)
 
             if request.method is not None:
                 span.set_attribute(SPANDATA.HTTP_REQUEST_METHOD, request.method)

tests/integrations/boto3/test_s3.py

@@ -131,14 +131,12 @@ def test_streaming(
         }
         if send_default_pii:
             expected_attrs["url.full"] = "https://bucket.s3.amazonaws.com/foo.pdf"
-            expected_attrs["url.fragment"] = ""
-            expected_attrs["url.query"] = ""
         assert span1["attributes"] == ApproxDict(expected_attrs)
 
+        assert "url.fragment" not in span1["attributes"]
+        assert "url.query" not in span1["attributes"]
         if not send_default_pii:
             assert "url.full" not in span1["attributes"]
-            assert "url.fragment" not in span1["attributes"]
-            assert "url.query" not in span1["attributes"]
 
         span2 = spans[1]
         assert span2["attributes"]["sentry.op"] == "http.client.stream"
@@ -426,9 +424,9 @@ def test_breadcrumb_span_streaming(sentry_init, capture_events, send_default_pii
                 SPANDATA.URL_FULL: mock.ANY,
                 SPANDATA.HTTP_REQUEST_METHOD: "GET",
                 SPANDATA.URL_QUERY: mock.ANY,
-                SPANDATA.URL_FRAGMENT: "",
             }
         )
+        assert SPANDATA.URL_FRAGMENT not in crumb["data"]
     else:
         assert crumb["data"] == ApproxDict(
             {
@@ -438,3 +436,159 @@ def test_breadcrumb_span_streaming(sentry_init, capture_events, send_default_pii
         assert SPANDATA.URL_FULL not in crumb["data"]
         assert SPANDATA.URL_QUERY not in crumb["data"]
         assert SPANDATA.URL_FRAGMENT not in crumb["data"]
+
+
+BUCKET_URL = "https://bucket.s3.amazonaws.com/"
+
+# ``expected_query`` of ``None`` means no URL data is recorded at all; ``""``
+# means the URL is recorded without a query string.
+# Structure of the parameters is "init_kwargs, expected_query"
+URL_QUERY_PARAMS = [
+    pytest.param(
+        {"send_default_pii": True},
+        "list-type=2&prefix=foo&continuation-token=abc&encoding-type=url",
+        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": {}}},
+        "list-type=2&prefix=foo&continuation-token=%5BFiltered%5D&encoding-type=url",
+        id="data_collection_denylist_default",
+    ),
+    pytest.param(
+        {
+            "_experiments": {
+                "data_collection": {
+                    "url_query_params": {"mode": "denylist", "terms": ["prefix"]}
+                }
+            }
+        },
+        "list-type=2&prefix=%5BFiltered%5D&continuation-token=%5BFiltered%5D&encoding-type=url",
+        id="data_collection_denylist_custom_terms",
+    ),
+    pytest.param(
+        {
+            "_experiments": {
+                "data_collection": {
+                    "url_query_params": {"mode": "allowlist", "terms": ["prefix"]}
+                }
+            }
+        },
+        "list-type=%5BFiltered%5D&prefix=foo&continuation-token=%5BFiltered%5D&encoding-type=%5BFiltered%5D",
+        id="data_collection_allowlist",
+    ),
+    pytest.param(
+        {
+            "_experiments": {
+                "data_collection": {
+                    "url_query_params": {
+                        "mode": "allowlist",
+                        "terms": ["continuation-token"],
+                    }
+                }
+            }
+        },
+        "list-type=%5BFiltered%5D&prefix=%5BFiltered%5D&continuation-token=%5BFiltered%5D&encoding-type=%5BFiltered%5D",
+        id="data_collection_allowlist_sensitive_term",
+    ),
+    pytest.param(
+        {"_experiments": {"data_collection": {"url_query_params": {"mode": "off"}}}},
+        "",
+        id="data_collection_off",
+    ),
+    pytest.param(
+        {
+            "send_default_pii": True,
+            "_experiments": {"data_collection": {"url_query_params": {"mode": "off"}}},
+        },
+        "",
+        id="data_collection_wins_over_send_default_pii",
+    ),
+]
+
+
+@pytest.mark.parametrize("init_kwargs, expected_query", URL_QUERY_PARAMS)
+def test_url_query_data_collection_span_streaming(
+    sentry_init, capture_items, init_kwargs, expected_query
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[Boto3Integration()],
+        default_integrations=False,
+        trace_lifecycle="stream",
+        **init_kwargs,
+    )
+
+    client = session.client("s3")
+
+    items = capture_items("span")
+
+    with sentry_sdk.traces.start_span(name="custom parent"), MockResponse(
+        client, 200, {}, read_fixture("s3_list.xml")
+    ):
+        client.list_objects_v2(Bucket="bucket", Prefix="foo", ContinuationToken="abc")
+
+    sentry_sdk.flush()
+
+    (span,) = [
+        item.payload
+        for item in items
+        if item.payload["attributes"].get("sentry.op") == "http.client"
+    ]
+
+    if expected_query is None:
+        assert SPANDATA.URL_QUERY not in span["attributes"]
+        assert SPANDATA.URL_FULL not in span["attributes"]
+    elif expected_query == "":
+        assert SPANDATA.URL_QUERY not in span["attributes"]
+        assert span["attributes"][SPANDATA.URL_FULL] == BUCKET_URL
+    else:
+        assert span["attributes"][SPANDATA.URL_QUERY] == expected_query
+        assert (
+            span["attributes"][SPANDATA.URL_FULL] == BUCKET_URL + "?" + expected_query
+        )
+
+
+@pytest.mark.parametrize("init_kwargs, expected_query", URL_QUERY_PARAMS)
+def test_url_query_data_collection_breadcrumb(
+    sentry_init, capture_events, init_kwargs, expected_query
+):
+    sentry_init(
+        integrations=[Boto3Integration()],
+        default_integrations=False,
+        trace_lifecycle="stream",
+        **init_kwargs,
+    )
+
+    client = session.client("s3")
+
+    events = capture_events()
+
+    with sentry_sdk.traces.start_span(name="custom parent"), MockResponse(
+        client, 200, {}, read_fixture("s3_list.xml")
+    ):
+        client.list_objects_v2(Bucket="bucket", Prefix="foo", ContinuationToken="abc")
+
+    capture_message("Testing!")
+
+    (event,) = events
+    (crumb,) = event["breadcrumbs"]["values"]
+
+    if expected_query is None:
+        assert SPANDATA.URL_QUERY not in crumb["data"]
+        assert SPANDATA.URL_FULL not in crumb["data"]
+    elif expected_query == "":
+        assert SPANDATA.URL_QUERY not in crumb["data"]
+        assert crumb["data"][SPANDATA.URL_FULL] == BUCKET_URL
+    else:
+        assert crumb["data"][SPANDATA.URL_QUERY] == expected_query
+        assert crumb["data"][SPANDATA.URL_FULL] == BUCKET_URL + "?" + expected_query