feat(redis): support data_collection option

sentry_sdk/integrations/redis/utils.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING
 
+import sentry_sdk
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.redis.consts import (
     _COMMANDS_INCLUDING_SENSITIVE_DATA,
@@ -11,7 +12,7 @@
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import Span
-from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE
+from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE, has_data_collection_enabled
 
 if TYPE_CHECKING:
     from typing import Any, Optional, Sequence, Union
@@ -22,6 +23,7 @@ def _get_safe_command(name: str, args: "Sequence[Any]") -> str:
 
     name_low = name.lower()
     send_default_pii = should_send_default_pii()
+    client_options = sentry_sdk.get_client().options
 
     for i, arg in enumerate(args):
         if i > _MAX_NUM_ARGS:
@@ -35,7 +37,10 @@ def _get_safe_command(name: str, args: "Sequence[Any]") -> str:
         if arg_is_the_key:
             command_parts.append(repr(arg))
         else:
-            if send_default_pii:
+            if has_data_collection_enabled(client_options):
+                if client_options["data_collection"]["database_query_data"]:
+                    command_parts.append(repr(arg))
+            elif send_default_pii:
                 command_parts.append(repr(arg))
             else:
                 command_parts.append(SENSITIVE_DATA_SUBSTITUTE)

tests/integrations/redis/test_redis.py

@@ -112,6 +112,72 @@ def test_redis_pipeline(
         }
 
 
+@pytest.mark.parametrize("span_streaming", [True, False])
+@pytest.mark.parametrize(
+    "data_collection, expected_first_ten",
+    [
+        (
+            {"database_query_data": False},
+            ["GET 'foo'", "SET 'bar'", "SET 'baz'"],
+        ),
+        (
+            {"database_query_data": True},
+            ["GET 'foo'", "SET 'bar' 1", "SET 'baz' 2"],
+        ),
+    ],
+)
+def test_redis_pipeline_data_collection(
+    sentry_init,
+    capture_events,
+    capture_items,
+    data_collection,
+    expected_first_ten,
+    span_streaming,
+):
+    sentry_init(
+        integrations=[RedisIntegration()],
+        traces_sample_rate=1.0,
+        trace_lifecycle="stream" if span_streaming else "static",
+        _experiments={"data_collection": data_collection},
+    )
+
+    connection = FakeStrictRedis()
+
+    if span_streaming:
+        items = capture_items("span")
+        with sentry_sdk.traces.start_span(name="custom parent"):
+            pipeline = connection.pipeline(transaction=False)
+            pipeline.get("foo")
+            pipeline.set("bar", 1)
+            pipeline.set("baz", 2)
+            pipeline.execute()
+        sentry_sdk.flush()
+
+        assert len(items) == 2
+        pipeline_span, parent_span = items[0].payload, items[1].payload
+
+        assert parent_span["name"] == "custom parent"
+        assert pipeline_span["name"] == "redis.pipeline.execute"
+        assert pipeline_span["attributes"]["sentry.op"] == "db.redis"
+    else:
+        events = capture_events()
+        with start_transaction():
+            pipeline = connection.pipeline(transaction=False)
+            pipeline.get("foo")
+            pipeline.set("bar", 1)
+            pipeline.set("baz", 2)
+            pipeline.execute()
+
+        (event,) = events
+        (span,) = event["spans"]
+        assert span["op"] == "db.redis"
+        assert span["description"] == "redis.pipeline.execute"
+        assert span["data"]["redis.commands"] == {
+            "count": 3,
+            "first_ten": expected_first_ten,
+        }
+
+
 @pytest.mark.parametrize("span_streaming", [True, False])
 def test_sensitive_data(sentry_init, capture_events, capture_items, span_streaming):
     # fakeredis does not support the AUTH command, so we need to mock it
@@ -201,6 +267,117 @@ def test_pii_data_redacted(sentry_init, capture_events, capture_items, span_stre
         assert spans[3]["description"] == "DEL 'somekey1' [Filtered]"
 
 
+@pytest.mark.parametrize("span_streaming", [True, False])
+@pytest.mark.parametrize(
+    "data_collection, expected_description",
+    [
+        ({"database_query_data": False}, "SET 'somekey1'"),
+        ({"database_query_data": True}, "SET 'somekey1' 'my secret string1'"),
+        ({}, "SET 'somekey1' 'my secret string1'"),
+    ],
+    ids=[
+        "database_query_data_disabled",
+        "database_query_data_enabled",
+        "database_query_data_not_provided_uses_defaults",
+    ],
+)
+def test_data_collection_database_query_data(
+    sentry_init,
+    capture_events,
+    capture_items,
+    span_streaming,
+    data_collection,
+    expected_description,
+):
+    sentry_init(
+        integrations=[RedisIntegration()],
+        traces_sample_rate=1.0,
+        trace_lifecycle="stream" if span_streaming else "static",
+        _experiments={"data_collection": data_collection},
+    )
+
+    connection = FakeStrictRedis()
+
+    if span_streaming:
+        items = capture_items("span")
+        with sentry_sdk.traces.start_span(name="custom parent"):
+            connection.set("somekey1", "my secret string1")
+        sentry_sdk.flush()
+
+        assert len(items) == 2
+        set_span, parent = [item.payload for item in items]
+
+        assert parent["name"] == "custom parent"
+        assert set_span["name"] == expected_description
+        assert set_span["attributes"][SPANDATA.DB_QUERY_TEXT] == expected_description
+        assert set_span["attributes"]["sentry.op"] == "db.redis"
+    else:
+        events = capture_events()
+        with start_transaction():
+            connection.set("somekey1", "my secret string1")
+
+        (event,) = events
+        spans = event["spans"]
+        assert spans[0]["op"] == "db.redis"
+        assert spans[0]["description"] == expected_description
+
+
+@pytest.mark.parametrize("span_streaming", [True, False])
+@pytest.mark.parametrize(
+    "data_collection, send_default_pii, expected_description",
+    [
+        ({"database_query_data": False}, True, "SET 'somekey1'"),
+        (
+            {"database_query_data": True},
+            False,
+            "SET 'somekey1' 'my secret string1'",
+        ),
+    ],
+)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_database_query_data_takes_precedence_over_send_default_pii(
+    sentry_init,
+    capture_events,
+    capture_items,
+    span_streaming,
+    data_collection,
+    send_default_pii,
+    expected_description,
+):
+    sentry_init(
+        integrations=[RedisIntegration()],
+        traces_sample_rate=1.0,
+        send_default_pii=send_default_pii,
+        trace_lifecycle="stream" if span_streaming else "static",
+        _experiments={"data_collection": data_collection},
+    )
+
+    connection = FakeStrictRedis()
+
+    if span_streaming:
+        items = capture_items("span")
+        with sentry_sdk.traces.start_span(name="custom parent"):
+            connection.set("somekey1", "my secret string1")
+        sentry_sdk.flush()
+
+        assert len(items) == 2
+        set_span, parent = [item.payload for item in items]
+
+        assert parent["name"] == "custom parent"
+        assert set_span["name"] == expected_description
+        assert set_span["attributes"][SPANDATA.DB_QUERY_TEXT] == expected_description
+        assert set_span["attributes"]["sentry.op"] == "db.redis"
+    else:
+        events = capture_events()
+        with start_transaction():
+            connection.set("somekey1", "my secret string1")
+
+        (event,) = events
+        spans = event["spans"]
+        assert spans[0]["op"] == "db.redis"
+        assert spans[0]["description"] == expected_description
+
+
 @pytest.mark.parametrize("span_streaming", [True, False])
 def test_pii_data_sent(sentry_init, capture_events, capture_items, span_streaming):
     sentry_init(