feat(pymongo): Respect data_collection.database_query_data option

sentry_sdk/integrations/pymongo.py

@@ -8,7 +8,7 @@
 from sentry_sdk.traces import SpanStatus, StreamedSpan
 from sentry_sdk.tracing import Span
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
-from sentry_sdk.utils import capture_internal_exceptions
+from sentry_sdk.utils import capture_internal_exceptions, has_data_collection_enabled
 
 try:
     from pymongo import monitoring
@@ -146,7 +146,10 @@ def started(self, event: "CommandStartedEvent") -> None:
             db_name = event.database_name
 
             lsid = command.pop("lsid", None)
-            if not should_send_default_pii():
+            if has_data_collection_enabled(client.options):
+                if not client.options["data_collection"]["database_query_data"]:
+                    command = _strip_pii(command)
+            elif not should_send_default_pii():
                 command = _strip_pii(command)
 
             query = json.dumps(command, default=str)

tests/integrations/pymongo/test_pymongo.py

@@ -110,6 +110,87 @@ def test_transactions(sentry_init, capture_events, mongo_server, with_pii):
     assert insert_fail["tags"]["status"] == "internal_error"
 
 
+DATA_COLLECTION_DATABASE_QUERY_DATA_USE_CASES = [
+    pytest.param(
+        {"_experiments": {"data_collection": {"database_query_data": True}}},
+        True,
+        id="query_data_enabled",
+    ),
+    pytest.param(
+        {"_experiments": {"data_collection": {"database_query_data": False}}},
+        False,
+        id="query_data_disabled",
+    ),
+    pytest.param(
+        {"_experiments": {"data_collection": {}}},
+        True,
+        id="query_data_default",
+    ),
+    pytest.param(
+        {
+            "send_default_pii": False,
+            "_experiments": {"data_collection": {"database_query_data": True}},
+        },
+        True,
+        id="data_collection_overrides_pii_off",
+    ),
+    pytest.param(
+        {
+            "send_default_pii": True,
+            "_experiments": {"data_collection": {"database_query_data": False}},
+        },
+        False,
+        id="data_collection_overrides_pii_on",
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "init_kwargs,expect_query_values", DATA_COLLECTION_DATABASE_QUERY_DATA_USE_CASES
+)
+def test_transactions_with_data_collection(
+    sentry_init, capture_events, mongo_server, init_kwargs, expect_query_values
+):
+    sentry_init(
+        integrations=[PyMongoIntegration()],
+        traces_sample_rate=1.0,
+        **init_kwargs,
+    )
+    events = capture_events()
+
+    connection = MongoClient(mongo_server.uri)
+
+    with start_transaction():
+        list(
+            connection["test_db"]["test_collection"].find({"foobar": 1})
+        )  # force query execution
+        connection["test_db"]["test_collection"].insert_one({"foo": 2})
+        try:
+            connection["test_db"]["erroneous"].insert_many([{"bar": 3}, {"baz": 4}])
+            pytest.fail("Request should raise")
+        except Exception:
+            pass
+
+    (event,) = events
+    (find, insert_success, insert_fail) = event["spans"]
+
+    assert find["description"].startswith('{"find')
+    assert insert_success["description"].startswith('{"insert')
+    assert insert_fail["description"].startswith('{"insert')
+
+    if expect_query_values:
+        assert "1" in find["description"]
+        assert "2" in insert_success["description"]
+        assert "3" in insert_fail["description"] and "4" in insert_fail["description"]
+    else:
+        assert "1" not in find["description"]
+        assert "2" not in insert_success["description"]
+        assert (
+            "3" not in insert_fail["description"]
+            and "4" not in insert_fail["description"]
+        )
+
+
 @pytest.mark.parametrize("with_pii", [False, True])
 def test_segment_span_streaming(sentry_init, capture_items, mongo_server, with_pii):
     sentry_init(
@@ -179,6 +260,57 @@ def test_segment_span_streaming(sentry_init, capture_items, mongo_server, with_p
     assert insert_fail["status"] == "error"
 
 
+@pytest.mark.parametrize(
+    "init_kwargs,expect_query_values", DATA_COLLECTION_DATABASE_QUERY_DATA_USE_CASES
+)
+def test_segment_span_streaming_with_data_collection(
+    sentry_init, capture_items, mongo_server, init_kwargs, expect_query_values
+):
+    sentry_init(
+        integrations=[PyMongoIntegration()],
+        traces_sample_rate=1.0,
+        trace_lifecycle="stream",
+        **init_kwargs,
+    )
+    items = capture_items("span")
+
+    connection = MongoClient(mongo_server.uri)
+
+    with sentry_sdk.traces.start_span(name="test_segment"):
+        list(
+            connection["test_db"]["test_collection"].find({"foobar": 1})
+        )  # force query execution
+        connection["test_db"]["test_collection"].insert_one({"foo": 2})
+        try:
+            connection["test_db"]["erroneous"].insert_many([{"bar": 3}, {"baz": 4}])
+            pytest.fail("Request should raise")
+        except Exception:
+            pass
+    sentry_sdk.flush()
+
+    spans = [item.payload for item in items]
+    assert len(spans) == 4
+
+    (find, insert_success, insert_fail, segment) = spans
+    assert segment["name"] == "test_segment"
+
+    assert find["name"].startswith('{"find')
+    assert insert_success["name"].startswith('{"insert')
+    assert insert_fail["name"].startswith('{"insert')
+
+    for span in find, insert_success, insert_fail:
+        assert span["attributes"][SPANDATA.DB_QUERY_TEXT] == span["name"]
+
+    if expect_query_values:
+        assert "1" in find["name"]
+        assert "2" in insert_success["name"]
+        assert "3" in insert_fail["name"] and "4" in insert_fail["name"]
+    else:
+        assert "1" not in find["name"]
+        assert "2" not in insert_success["name"]
+        assert "3" not in insert_fail["name"] and "4" not in insert_fail["name"]
+
+
 @pytest.mark.parametrize("with_pii", [False, True])
 def test_breadcrumbs(sentry_init, capture_events, mongo_server, with_pii):
     sentry_init(
@@ -216,6 +348,38 @@ def test_breadcrumbs(sentry_init, capture_events, mongo_server, with_pii):
     }
 
 
+@pytest.mark.parametrize(
+    "init_kwargs,expect_query_values", DATA_COLLECTION_DATABASE_QUERY_DATA_USE_CASES
+)
+def test_breadcrumbs_with_data_collection(
+    sentry_init, capture_events, mongo_server, init_kwargs, expect_query_values
+):
+    sentry_init(
+        integrations=[PyMongoIntegration()],
+        traces_sample_rate=1.0,
+        **init_kwargs,
+    )
+    events = capture_events()
+
+    connection = MongoClient(mongo_server.uri)
+
+    list(
+        connection["test_db"]["test_collection"].find({"foobar": 1})
+    )  # force query execution
+    capture_message("hi")
+
+    (event,) = events
+    (crumb,) = event["breadcrumbs"]["values"]
+
+    assert crumb["category"] == "query"
+    assert crumb["message"].startswith('{"find')
+    if expect_query_values:
+        assert "1" in crumb["message"]
+    else:
+        assert "1" not in crumb["message"]
+    assert crumb["type"] == "db"
+
+
 @pytest.mark.parametrize("with_pii", [False, True])
 def test_breadcrumbs_span_streaming(sentry_init, capture_items, mongo_server, with_pii):
     sentry_init(
@@ -257,6 +421,40 @@ def test_breadcrumbs_span_streaming(sentry_init, capture_items, mongo_server, wi
     assert data[SPANDATA.SERVER_PORT] == mongo_server.port
 
 
+@pytest.mark.parametrize(
+    "init_kwargs,expect_query_values", DATA_COLLECTION_DATABASE_QUERY_DATA_USE_CASES
+)
+def test_breadcrumbs_span_streaming_with_data_collection(
+    sentry_init, capture_items, mongo_server, init_kwargs, expect_query_values
+):
+    sentry_init(
+        integrations=[PyMongoIntegration()],
+        traces_sample_rate=1.0,
+        trace_lifecycle="stream",
+        **init_kwargs,
+    )
+    items = capture_items("event")
+
+    connection = MongoClient(mongo_server.uri)
+
+    list(
+        connection["test_db"]["test_collection"].find({"foobar": 1})
+    )  # force query execution
+    capture_message("hi")
+
+    event = items[0].payload
+    (crumb,) = event["breadcrumbs"]["values"]
+
+    assert crumb["category"] == "query"
+    assert crumb["message"].startswith('{"find')
+    if expect_query_values:
+        assert "1" in crumb["message"]
+    else:
+        assert "1" not in crumb["message"]
+    assert crumb["type"] == "db"
+    assert crumb["data"]["db.query.text"] == crumb["message"]
+
+
 @pytest.mark.parametrize(
     "testcase",
     [