feat(graphene): Gate GraphQL data collection behind data_collection option

sentry_sdk/integrations/graphene.py

@@ -9,6 +9,7 @@
     capture_internal_exceptions,
     ensure_integration_enabled,
     event_from_exception,
+    has_data_collection_enabled,
     package_version,
 )
 
@@ -109,10 +110,17 @@ async def _sentry_patched_graphql_async(
 
 
 def _event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event":
-    if should_send_default_pii():
+    client_options = sentry_sdk.get_client().options
+
+    if has_data_collection_enabled(client_options):
+        if client_options["data_collection"]["graphql"]["document"]:
+            request_info = event.setdefault("request", {})
+            request_info["api_target"] = "graphql"
+        elif event.get("request", {}).get("data"):
+            del event["request"]["data"]
+    elif should_send_default_pii():
         request_info = event.setdefault("request", {})
         request_info["api_target"] = "graphql"
-
     elif event.get("request", {}).get("data"):
         del event["request"]["data"]
 
@@ -144,17 +152,19 @@ def graphql_span(
         },
     )
 
-    is_span_streaming_enabled = has_span_streaming_enabled(
-        sentry_sdk.get_client().options
-    )
+    client_options = sentry_sdk.get_client().options
+    is_span_streaming_enabled = has_span_streaming_enabled(client_options)
 
     if is_span_streaming_enabled:
         if sentry_sdk.traces.get_current_span() is None:
             yield
             return
 
         additional_attributes = {}
-        if should_send_default_pii():
+        if has_data_collection_enabled(client_options):
+            if client_options["data_collection"]["graphql"]["document"]:
+                additional_attributes["graphql.document"] = source
+        elif should_send_default_pii():
             additional_attributes["graphql.document"] = source
 
         _graphql_span = sentry_sdk.traces.start_span(
@@ -169,8 +179,12 @@ def graphql_span(
     else:
         _graphql_span = sentry_sdk.start_span(op=op, name=operation_name)
 
-        if should_send_default_pii():
+        if has_data_collection_enabled(client_options):
+            if client_options["data_collection"]["graphql"]["document"]:
+                _graphql_span.set_data("graphql.document", source)
+        elif should_send_default_pii():
             _graphql_span.set_data("graphql.document", source)
+
         _graphql_span.set_data("graphql.operation.name", operation_name)
         _graphql_span.set_data("graphql.operation.type", operation_type)
 

tests/integrations/graphene/test_graphene.py

@@ -11,6 +11,33 @@
 from sentry_sdk.integrations.graphene import GrapheneIntegration
 from sentry_sdk.integrations.starlette import StarletteIntegration
 
+DATA_COLLECTION_GRAPHQL_DOCUMENTS_PARAMS = [
+    pytest.param(
+        {"graphql": {"document": True}},
+        None,
+        True,
+        id="document_on_collects_graphql_data",
+    ),
+    pytest.param(
+        {"graphql": {"document": False}},
+        None,
+        False,
+        id="document_off_omits_graphql_data",
+    ),
+    pytest.param(
+        {"graphql": {"document": False}},
+        True,
+        False,
+        id="data_collection_takes_precedence_over_send_default_pii_on",
+    ),
+    pytest.param(
+        {"graphql": {"document": True}},
+        False,
+        True,
+        id="data_collection_takes_precedence_over_send_default_pii_off",
+    ),
+]
+
 
 class Query(ObjectType):
     hello = String(first_name=String(default_value="stranger"))
@@ -149,6 +176,97 @@ def graphql_server_sync():
     assert "response" not in event["contexts"]
 
 
+@pytest.mark.parametrize(
+    "data_collection,send_default_pii,expect_api_target",
+    DATA_COLLECTION_GRAPHQL_DOCUMENTS_PARAMS,
+)
+def test_event_processor_data_collection_sync(
+    sentry_init, capture_events, data_collection, send_default_pii, expect_api_target
+):
+    init_kwargs = {
+        "integrations": [GrapheneIntegration(), FlaskIntegration()],
+        "_experiments": {"data_collection": data_collection},
+    }
+    if send_default_pii is not None:
+        init_kwargs["send_default_pii"] = send_default_pii
+    sentry_init(**init_kwargs)
+    events = capture_events()
+
+    schema = Schema(query=Query)
+
+    sync_app = Flask(__name__)
+
+    @sync_app.route("/graphql", methods=["POST"])
+    def graphql_server_sync():
+        data = request.get_json()
+        result = schema.execute(data["query"])
+        return jsonify(result.data), 200
+
+    query = {"query": "query ErrorQuery {goodbye}"}
+    client = sync_app.test_client()
+    client.post("/graphql", json=query)
+
+    assert len(events) == 1
+
+    (event,) = events
+    assert event["exception"]["values"][0]["mechanism"]["type"] == "graphene"
+    if expect_api_target:
+        assert event["request"]["api_target"] == "graphql"
+        assert "data" in event.get("request", {})
+    else:
+        assert "api_target" not in event.get("request", {})
+        assert "data" not in event.get("request", {})
+
+
+@pytest.mark.parametrize(
+    "data_collection,send_default_pii,expect_api_target",
+    DATA_COLLECTION_GRAPHQL_DOCUMENTS_PARAMS,
+)
+def test_event_processor_data_collection_async(
+    sentry_init, capture_events, data_collection, send_default_pii, expect_api_target
+):
+    init_kwargs = {
+        "integrations": [
+            GrapheneIntegration(),
+            FastApiIntegration(),
+            StarletteIntegration(),
+        ],
+        "_experiments": {"data_collection": data_collection},
+    }
+
+    if send_default_pii is not None:
+        init_kwargs["send_default_pii"] = send_default_pii
+    sentry_init(**init_kwargs)
+
+    events = capture_events()
+
+    schema = Schema(query=Query)
+
+    async_app = FastAPI()
+
+    @async_app.post("/graphql")
+    async def graphql_server_async(request: Request):
+        data = await request.json()
+        result = await schema.execute_async(data["query"])
+        return result.data
+
+    query = {"query": "query ErrorQuery {goodbye}"}
+    client = TestClient(async_app)
+    client.post("/graphql", json=query)
+
+    assert len(events) == 1
+
+    (event,) = events
+    assert event["exception"]["values"][0]["mechanism"]["type"] == "graphene"
+
+    if expect_api_target:
+        assert event["request"]["api_target"] == "graphql"
+        assert "data" in event.get("request", {})
+    else:
+        assert "api_target" not in event.get("request", {})
+        assert "data" not in event.get("request", {})
+
+
 def test_no_event_if_no_errors_async(sentry_init, capture_events):
     sentry_init(
         integrations=[
@@ -255,6 +373,58 @@ def graphql_server_sync():
         assert "graphql.document" not in span["data"]
 
 
+@pytest.mark.parametrize(
+    "data_collection,send_default_pii,expect_document",
+    DATA_COLLECTION_GRAPHQL_DOCUMENTS_PARAMS,
+)
+def test_graphql_span_data_collection(
+    sentry_init, capture_events, data_collection, send_default_pii, expect_document
+):
+    init_kwargs = {
+        "integrations": [GrapheneIntegration(), FlaskIntegration()],
+        "traces_sample_rate": 1.0,
+        "default_integrations": False,
+        "_experiments": {"data_collection": data_collection},
+    }
+    if send_default_pii is not None:
+        init_kwargs["send_default_pii"] = send_default_pii
+    sentry_init(**init_kwargs)
+    events = capture_events()
+
+    schema = Schema(query=Query)
+
+    sync_app = Flask(__name__)
+
+    @sync_app.route("/graphql", methods=["POST"])
+    def graphql_server_sync():
+        data = request.get_json()
+        result = schema.execute(data["query"], operation_name=data.get("operationName"))
+        return jsonify(result.data), 200
+
+    query = {
+        "query": "query GreetingQuery { hello }",
+        "operationName": "GreetingQuery",
+    }
+    client = sync_app.test_client()
+    client.post("/graphql", json=query)
+
+    assert len(events) == 1
+
+    (event,) = events
+    assert len(event["spans"]) == 1
+
+    (span,) = event["spans"]
+    assert span["op"] == OP.GRAPHQL_QUERY
+    assert span["description"] == query["operationName"]
+    assert span["data"]["graphql.operation.name"] == query["operationName"]
+    assert span["data"]["graphql.operation.type"] == "query"
+
+    if expect_document:
+        assert span["data"]["graphql.document"] == query["query"]
+    else:
+        assert "graphql.document" not in span["data"]
+
+
 @pytest.mark.parametrize(
     "send_default_pii",
     [True, False],
@@ -312,6 +482,66 @@ def graphql_server_sync():
     assert graphql_span["parent_span_id"] == flask_segment["span_id"]
 
 
+@pytest.mark.parametrize(
+    "data_collection,send_default_pii,expect_document",
+    DATA_COLLECTION_GRAPHQL_DOCUMENTS_PARAMS,
+)
+def test_graphql_streamed_span_data_collection(
+    sentry_init, capture_items, data_collection, send_default_pii, expect_document
+):
+    init_kwargs = {
+        "integrations": [GrapheneIntegration(), FlaskIntegration()],
+        "traces_sample_rate": 1.0,
+        "default_integrations": False,
+        "trace_lifecycle": "stream",
+        "_experiments": {"data_collection": data_collection},
+    }
+    if send_default_pii is not None:
+        init_kwargs["send_default_pii"] = send_default_pii
+    sentry_init(**init_kwargs)
+    items = capture_items("span")
+
+    schema = Schema(query=Query)
+
+    sync_app = Flask(__name__)
+
+    @sync_app.route("/graphql", methods=["POST"])
+    def graphql_server_sync():
+        data = request.get_json()
+        result = schema.execute(data["query"], operation_name=data.get("operationName"))
+        return jsonify(result.data), 200
+
+    query = {
+        "query": "query GreetingQuery { hello }",
+        "operationName": "GreetingQuery",
+    }
+    client = sync_app.test_client()
+    client.post("/graphql", json=query)
+
+    sentry_sdk.get_client().flush()
+
+    spans = [item.payload for item in items]
+    assert len(spans) == 2
+
+    graphql_span, flask_segment = spans
+
+    assert graphql_span["name"] == query["operationName"]
+    assert graphql_span["attributes"]["sentry.op"] == OP.GRAPHQL_QUERY
+    assert (
+        graphql_span["attributes"]["graphql.operation.name"] == query["operationName"]
+    )
+    assert graphql_span["attributes"]["graphql.operation.type"] == "query"
+    assert graphql_span["is_segment"] is False
+
+    if expect_document:
+        assert graphql_span["attributes"]["graphql.document"] == query["query"]
+    else:
+        assert "graphql.document" not in graphql_span["attributes"]
+
+    assert flask_segment["is_segment"] is True
+    assert graphql_span["parent_span_id"] == flask_segment["span_id"]
+
+
 def test_breadcrumbs_hold_query_information_on_error(sentry_init, capture_events):
     sentry_init(
         integrations=[