feat(graphene): Add span streaming support (#6476)

Add support for OpenTelemetry span streaming in the Graphene GraphQL
integration. When span streaming is enabled via the trace_lifecycle
experiment, use sentry_sdk.traces.start_span() with attributes instead
of the traditional start_span() + set_data() approach.

Also improve operation_name handling to provide a sensible default when
the operation name is not specified in the GraphQL request.

Fixes PY-2328
Fixes #6026

Author: ericapisani

sentry_sdk/integrations/graphene.py

@@ -4,6 +4,7 @@
 from sentry_sdk.consts import OP
 from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
@@ -122,7 +123,7 @@ def _event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event":
 def graphql_span(
     schema: "GraphQLSchema", source: "Union[str, Source]", kwargs: "Dict[str, Any]"
 ) -> "Generator[None, None, None]":
-    operation_name = kwargs.get("operation_name")
+    operation_name = kwargs.get("operation_name") or "<unknown graphql operation>"
 
     operation_type = "query"
     op = OP.GRAPHQL_QUERY
@@ -143,16 +144,38 @@ def graphql_span(
         },
     )
 
-    _graphql_span = sentry_sdk.start_span(op=op, name=operation_name)
+    is_span_streaming_enabled = has_span_streaming_enabled(
+        sentry_sdk.get_client().options
+    )
 
-    if 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)
+    if is_span_streaming_enabled:
+        additional_attributes = {}
+        if should_send_default_pii():
+            additional_attributes["graphql.document"] = source
+
+        _graphql_span = sentry_sdk.traces.start_span(
+            name=operation_name,
+            attributes={
+                "sentry.op": op,
+                "graphql.operation.name": operation_name,
+                "graphql.operation.type": operation_type,
+                **additional_attributes,
+            },
+        )
+    else:
+        _graphql_span = sentry_sdk.start_span(op=op, name=operation_name)
+
+        if 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)
 
-    _graphql_span.__enter__()
+        _graphql_span.__enter__()
 
     try:
         yield
     finally:
-        _graphql_span.__exit__(None, None, None)
+        if is_span_streaming_enabled:
+            _graphql_span.end()  # type: ignore
+        else:
+            _graphql_span.__exit__(None, None, None)

tests/integrations/graphene/test_graphene.py

@@ -4,6 +4,7 @@
 from flask import Flask, jsonify, request
 from graphene import ObjectType, Schema, String
 
+import sentry_sdk
 from sentry_sdk.consts import OP
 from sentry_sdk.integrations.fastapi import FastApiIntegration
 from sentry_sdk.integrations.flask import FlaskIntegration
@@ -254,6 +255,63 @@ def graphql_server_sync():
         assert "graphql.document" not in span["data"]
 
 
+@pytest.mark.parametrize(
+    "send_default_pii",
+    [True, False],
+)
+def test_graphql_streamed_span_holds_query_information(
+    sentry_init, capture_items, send_default_pii
+):
+    sentry_init(
+        integrations=[GrapheneIntegration(), FlaskIntegration()],
+        traces_sample_rate=1.0,
+        default_integrations=False,
+        send_default_pii=send_default_pii,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    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 send_default_pii is True:
+        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=[
@@ -293,3 +351,50 @@ def graphql_server_sync():
     assert breadcrumb["data"]["operation_name"] == query["operationName"]
     assert breadcrumb["data"]["operation_type"] == "query"
     assert breadcrumb["type"] == "default"
+
+
+def test_breadcrumbs_hold_query_information_on_error_with_span_streaming(
+    sentry_init, capture_items
+):
+    sentry_init(
+        integrations=[
+            GrapheneIntegration(),
+        ],
+        default_integrations=False,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    items = capture_items("span", "event")
+
+    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 ErrorQuery { goodbye }",
+        "operationName": "ErrorQuery",
+    }
+    client = sync_app.test_client()
+    client.post("/graphql", json=query)
+
+    sentry_sdk.get_client().flush()
+
+    events = [item.payload for item in items if item.type == "event"]
+    assert len(events) == 1
+
+    (event,) = events
+    assert len(event["breadcrumbs"]) == 1
+
+    breadcrumbs = event["breadcrumbs"]["values"]
+    assert len(breadcrumbs) == 1
+
+    (breadcrumb,) = breadcrumbs
+    assert breadcrumb["category"] == "graphql.operation"
+    assert breadcrumb["data"]["operation_name"] == query["operationName"]
+    assert breadcrumb["data"]["operation_type"] == "query"
+    assert breadcrumb["type"] == "default"