Fix strawberry integration (#3874)

getsentry · Dec 18, 2024 · 92f5391 · 92f5391
1 parent 1984a5c
commit 92f5391
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 41 deletions.
diff --git a/sentry_sdk/integrations/strawberry.py b/sentry_sdk/integrations/strawberry.py
@@ -179,59 +179,50 @@ def on_operation(self):
             },
         )
 
-        span = sentry_sdk.get_current_span()
-        if span:
-            self.graphql_span = span.start_child(
-                op=op,
-                name=description,
-                origin=StrawberryIntegration.origin,
-            )
-        else:
-            self.graphql_span = sentry_sdk.start_span(
-                op=op,
-                name=description,
-                origin=StrawberryIntegration.origin,
-                only_if_parent=True,
-            )
+        with sentry_sdk.start_span(
+            op=op,
+            name=description,
+            origin=StrawberryIntegration.origin,
+            only_if_parent=True,
+        ) as graphql_span:
+            graphql_span.set_data("graphql.operation.type", operation_type)
+            graphql_span.set_data("graphql.document", self.execution_context.query)
+            graphql_span.set_data("graphql.resource_name", self._resource_name)
+
+            yield
 
-        self.graphql_span.set_data("graphql.operation.type", operation_type)
-        self.graphql_span.set_data("graphql.operation.name", self._operation_name)
-        self.graphql_span.set_data("graphql.document", self.execution_context.query)
-        self.graphql_span.set_data("graphql.resource_name", self._resource_name)
+            # we might have a more accurate operation_name after the parsing
+            self._operation_name = self.execution_context.operation_name
 
-        yield
+            if self._operation_name is not None:
+                graphql_span.set_data("graphql.operation.name", self._operation_name)
 
-        transaction = self.graphql_span.containing_transaction
-        if transaction and self.execution_context.operation_name:
-            transaction.name = self.execution_context.operation_name
-            transaction.source = TRANSACTION_SOURCE_COMPONENT
-            transaction.op = op
+                sentry_sdk.get_current_scope().set_transaction_name(
+                    self._operation_name,
+                    source=TRANSACTION_SOURCE_COMPONENT,
+                )
 
-        self.graphql_span.finish()
+            root_span = graphql_span.root_span
+            if root_span:
+                root_span.op = op
 
     def on_validate(self):
         # type: () -> Generator[None, None, None]
-        self.validation_span = self.graphql_span.start_child(
+        with sentry_sdk.start_span(
             op=OP.GRAPHQL_VALIDATE,
             name="validation",
             origin=StrawberryIntegration.origin,
-        )
-
-        yield
-
-        self.validation_span.finish()
+        ):
+            yield
 
     def on_parse(self):
         # type: () -> Generator[None, None, None]
-        self.parsing_span = self.graphql_span.start_child(
+        with sentry_sdk.start_span(
             op=OP.GRAPHQL_PARSE,
             name="parsing",
             origin=StrawberryIntegration.origin,
-        )
-
-        yield
-
-        self.parsing_span.finish()
+        ):
+            yield
 
     def should_skip_tracing(self, _next, info):
         # type: (Callable[[Any, GraphQLResolveInfo, Any, Any], Any], GraphQLResolveInfo) -> bool
@@ -253,7 +244,7 @@ async def resolve(self, _next, root, info, *args, **kwargs):
 
         field_path = "{}.{}".format(info.parent_type, info.field_name)
 
-        with self.graphql_span.start_child(
+        with sentry_sdk.start_span(
             op=OP.GRAPHQL_RESOLVE,
             name="resolving {}".format(field_path),
             origin=StrawberryIntegration.origin,
@@ -274,7 +265,7 @@ def resolve(self, _next, root, info, *args, **kwargs):
 
         field_path = "{}.{}".format(info.parent_type, info.field_name)
 
-        with self.graphql_span.start_child(
+        with sentry_sdk.start_span(
             op=OP.GRAPHQL_RESOLVE,
             name="resolving {}".format(field_path),
             origin=StrawberryIntegration.origin,

diff --git a/tests/integrations/strawberry/test_strawberry.py b/tests/integrations/strawberry/test_strawberry.py
@@ -505,7 +505,7 @@ def test_transaction_no_operation_name(
     query_span = query_spans[0]
     assert query_span["description"] == "query"
     assert query_span["data"]["graphql.operation.type"] == "query"
-    assert query_span["data"]["graphql.operation.name"] is None
+    assert "graphql.operation.name" not in query_span["data"]
     assert query_span["data"]["graphql.document"] == query
     assert query_span["data"]["graphql.resource_name"]
 
@@ -582,7 +582,7 @@ def test_transaction_mutation(
     query_span = query_spans[0]
     assert query_span["description"] == "mutation"
     assert query_span["data"]["graphql.operation.type"] == "mutation"
-    assert query_span["data"]["graphql.operation.name"] is None
+    assert query_span["data"]["graphql.operation.name"] == "Change"
     assert query_span["data"]["graphql.document"] == query
     assert query_span["data"]["graphql.resource_name"]