migrate custom schema extensions to the new design

strawberry/extensions/add_validation_rules.py

@@ -7,6 +7,8 @@
 if TYPE_CHECKING:
     from graphql import ASTValidationRule
 
+    from strawberry.types.execution import ExecutionContext
+
 
 class AddValidationRules(SchemaExtension):
     """Add graphql-core validation rules.
@@ -42,9 +44,9 @@ def enter_field(self, node, *args) -> None:
     def __init__(self, validation_rules: List[Type[ASTValidationRule]]) -> None:
         self.validation_rules = validation_rules
 
-    def on_operation(self) -> Iterator[None]:
-        self.execution_context.validation_rules = (
-            self.execution_context.validation_rules + tuple(self.validation_rules)
+    def on_operation(self, execution_context: ExecutionContext) -> Iterator[None]:
+        execution_context.validation_rules = execution_context.validation_rules + tuple(
+            self.validation_rules
         )
         yield
 

strawberry/extensions/base_extension.py

@@ -19,11 +19,12 @@ class LifecycleStep(Enum):
 
 
 class SchemaExtension:
-    # to support extensions that still use the old signature
-    # we have an optional argument here for ease of initialization.
-    def __init__(
-        self, *, execution_context: ExecutionContext | None = None
-    ) -> None: ...
+    if not TYPE_CHECKING:
+        # to support extensions that still use the old signature
+        # we have an optional argument here for ease of initialization.
+        def __init__(
+            self, *, execution_context: ExecutionContext | None = None
+        ) -> None: ...
 
     def on_operation(  # type: ignore
         self, execution_context: ExecutionContext

strawberry/extensions/disable_validation.py

@@ -1,6 +1,7 @@
 from typing import Iterator
 
 from strawberry.extensions.base_extension import SchemaExtension
+from strawberry.types.execution import ExecutionContext
 
 
 class DisableValidation(SchemaExtension):
@@ -26,8 +27,8 @@ def __init__(self) -> None:
         # some in the future
         pass
 
-    def on_operation(self) -> Iterator[None]:
-        self.execution_context.validation_rules = ()  # remove all validation_rules
+    def on_operation(self, execution_context: ExecutionContext) -> Iterator[None]:
+        execution_context.validation_rules = ()  # remove all validation_rules
         yield
 
 

strawberry/extensions/mask_errors.py

@@ -3,6 +3,7 @@
 from graphql.error import GraphQLError
 
 from strawberry.extensions.base_extension import SchemaExtension
+from strawberry.types.execution import ExecutionContext
 
 
 def default_should_mask_error(_: GraphQLError) -> bool:
@@ -32,9 +33,9 @@ def anonymise_error(self, error: GraphQLError) -> GraphQLError:
             original_error=None,
         )
 
-    def on_operation(self) -> Iterator[None]:
+    def on_operation(self, execution_context: ExecutionContext) -> Iterator[None]:
         yield
-        result = self.execution_context.result
+        result = execution_context.result
         if result and result.errors:
             processed_errors: List[GraphQLError] = []
             for error in result.errors:

strawberry/extensions/max_tokens.py

@@ -1,6 +1,7 @@
 from typing import Iterator
 
 from strawberry.extensions.base_extension import SchemaExtension
+from strawberry.types.execution import ExecutionContext
 
 
 class MaxTokensLimiter(SchemaExtension):
@@ -34,8 +35,8 @@ def __init__(
         """
         self.max_token_count = max_token_count
 
-    def on_operation(self) -> Iterator[None]:
-        self.execution_context.parse_options["max_tokens"] = self.max_token_count
+    def on_operation(self, execution_context: ExecutionContext) -> Iterator[None]:
+        execution_context.parse_options["max_tokens"] = self.max_token_count
         yield
 
 

strawberry/extensions/parser_cache.py

@@ -3,6 +3,7 @@
 
 from strawberry.extensions.base_extension import SchemaExtension
 from strawberry.schema.execute import parse_document
+from strawberry.types.execution import ExecutionContext
 
 
 class ParserCache(SchemaExtension):
@@ -33,9 +34,7 @@ def __init__(self, maxsize: Optional[int] = None) -> None:
         """
         self.cached_parse_document = lru_cache(maxsize=maxsize)(parse_document)
 
-    def on_parse(self) -> Iterator[None]:
-        execution_context = self.execution_context
-
+    def on_parse(self, execution_context: ExecutionContext) -> Iterator[None]:
         execution_context.graphql_document = self.cached_parse_document(
             execution_context.query, **execution_context.parse_options
         )

strawberry/extensions/pyinstrument.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Iterator
+from typing import TYPE_CHECKING, Iterator
 
 from pyinstrument import Profiler
 
 from strawberry.extensions.base_extension import SchemaExtension
 
+if TYPE_CHECKING:
+    from strawberry.types.execution import ExecutionContext
+
 
 class PyInstrument(SchemaExtension):
     """Extension to profile the execution time of resolvers using PyInstrument."""
@@ -17,7 +20,7 @@ def __init__(
     ) -> None:
         self._report_path = report_path
 
-    def on_operation(self) -> Iterator[None]:
+    def on_operation(self, execution_context: ExecutionContext) -> Iterator[None]:
         profiler = Profiler()
         profiler.start()
 

strawberry/extensions/tracing/apollo.py

@@ -14,6 +14,8 @@
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
 
+    from strawberry.types.execution import ExecutionContext
+
 DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 
 if TYPE_CHECKING:
@@ -80,23 +82,30 @@ def to_json(self) -> Dict[str, Any]:
 
 
 class ApolloTracingExtension(SchemaExtension):
-    def __init__(self, execution_context: ExecutionContext) -> None:
-        self._resolver_stats: List[ApolloResolverStats] = []
-        self.execution_context = execution_context
-
-    def on_operation(self) -> Generator[None, None, None]:
+    def __init__(self) -> None:
+        self._resolver_stats: List[
+            ApolloResolverStats
+        ] = []  # TODO: this is probably a bug just as using self.execution_context was a bug.
+
+    def on_operation(
+        self, execution_context: ExecutionContext
+    ) -> Generator[None, None, None]:
         self.start_timestamp = self.now()
         self.start_time = datetime.now(timezone.utc)
         yield
         self.end_timestamp = self.now()
         self.end_time = datetime.now(timezone.utc)
 
-    def on_parse(self) -> Generator[None, None, None]:
+    def on_parse(
+        self, execution_context: ExecutionContext
+    ) -> Generator[None, None, None]:
         self._start_parsing = self.now()
         yield
         self._end_parsing = self.now()
 
-    def on_validate(self) -> Generator[None, None, None]:
+    def on_validate(
+        self, execution_context: ExecutionContext
+    ) -> Generator[None, None, None]:
         self._start_validation = self.now()
         yield
         self._end_validation = self.now()
@@ -121,7 +130,9 @@ def stats(self) -> ApolloTracingStats:
             ),
         )
 
-    def get_results(self) -> Dict[str, Dict[str, Any]]:
+    def get_results(
+        self, execution_context: ExecutionContext
+    ) -> Dict[str, Dict[str, Any]]:
         return {"tracing": self.stats.to_json()}
 
     async def resolve(

strawberry/extensions/tracing/datadog.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import hashlib
-from functools import cached_property
+from functools import lru_cache
 from inspect import isawaitable
 from typing import TYPE_CHECKING, Any, Callable, Generator, Iterator, Optional
 
@@ -25,15 +25,15 @@ def __init__(
         if execution_context:
             self.execution_context = execution_context
 
-    @cached_property
-    def _resource_name(self) -> str:
-        if self.execution_context.query is None:
+    @lru_cache
+    def _resource_name(self, execution_context: ExecutionContext) -> str:
+        if execution_context.query is None:
             return "query_missing"
 
-        query_hash = self.hash_query(self.execution_context.query)
+        query_hash = self.hash_query(execution_context.query)
 
-        if self.execution_context.operation_name:
-            return f"{self.execution_context.operation_name}:{query_hash}"
+        if execution_context.operation_name:
+            return f"{execution_context.operation_name}:{query_hash}"
 
         return query_hash
 
@@ -54,7 +54,7 @@ class CustomExtension(DatadogTracingExtension):
             def create_span(self, lifecycle_step, name, **kwargs):
                 span = super().create_span(lifecycle_step, name, **kwargs)
                 if lifecycle_step == LifeCycleStep.OPERATION:
-                    span.set_tag("graphql.query", self.execution_context.query)
+                    span.set_tag("graphql.query", execution_context.query)
                 return span
         ```
         """
@@ -67,21 +67,21 @@ def create_span(self, lifecycle_step, name, **kwargs):
     def hash_query(self, query: str) -> str:
         return hashlib.md5(query.encode("utf-8")).hexdigest()
 
-    def on_operation(self) -> Iterator[None]:
-        self._operation_name = self.execution_context.operation_name
+    def on_operation(self, execution_context: ExecutionContext) -> Iterator[None]:
+        self._operation_name = execution_context.operation_name
         span_name = (
             f"{self._operation_name}" if self._operation_name else "Anonymous Query"
         )
 
         self.request_span = self.create_span(
             LifecycleStep.OPERATION,
             span_name,
-            resource=self._resource_name,
+            resource=self._resource_name(execution_context),
             service="strawberry",
         )
         self.request_span.set_tag("graphql.operation_name", self._operation_name)
 
-        query = self.execution_context.query
+        query = execution_context.query
 
         if query is not None:
             query = query.strip()
@@ -100,15 +100,19 @@ def on_operation(self) -> Iterator[None]:
 
         self.request_span.finish()
 
-    def on_validate(self) -> Generator[None, None, None]:
+    def on_validate(
+        self, execution_context: ExecutionContext
+    ) -> Generator[None, None, None]:
         self.validation_span = self.create_span(
             lifecycle_step=LifecycleStep.VALIDATION,
             name="Validation",
         )
         yield
         self.validation_span.finish()
 
-    def on_parse(self) -> Generator[None, None, None]:
+    def on_parse(
+        self, execution_context: ExecutionContext
+    ) -> Generator[None, None, None]:
         self.parsing_span = self.create_span(
             lifecycle_step=LifecycleStep.PARSE,
             name="Parsing",
@@ -150,6 +154,9 @@ async def resolve(
 
             return result
 
+    def __hash__(self) -> int:
+        return id(self)
+
 
 class DatadogTracingExtensionSync(DatadogTracingExtension):
     def resolve(

strawberry/extensions/tracing/opentelemetry.py

@@ -51,8 +51,10 @@ def __init__(
         if execution_context:
             self.execution_context = execution_context
 
-    def on_operation(self) -> Generator[None, None, None]:
-        self._operation_name = self.execution_context.operation_name
+    def on_operation(
+        self, execution_context: ExecutionContext
+    ) -> Generator[None, None, None]:
+        self._operation_name = execution_context.operation_name
         span_name = (
             f"GraphQL Query: {self._operation_name}"
             if self._operation_name
@@ -64,9 +66,9 @@ def on_operation(self) -> Generator[None, None, None]:
         )
         self._span_holder[LifecycleStep.OPERATION].set_attribute("component", "graphql")
 
-        if self.execution_context.query:
+        if execution_context.query:
             self._span_holder[LifecycleStep.OPERATION].set_attribute(
-                "query", self.execution_context.query
+                "query", execution_context.query
             )
 
         yield
@@ -75,12 +77,14 @@ def on_operation(self) -> Generator[None, None, None]:
         # operation but we don't know until the parsing stage has finished. If
         # that's the case we want to update the span name so that we have a more
         # useful name in our trace.
-        if not self._operation_name and self.execution_context.operation_name:
-            span_name = f"GraphQL Query: {self.execution_context.operation_name}"
+        if not self._operation_name and execution_context.operation_name:
+            span_name = f"GraphQL Query: {execution_context.operation_name}"
             self._span_holder[LifecycleStep.OPERATION].update_name(span_name)
         self._span_holder[LifecycleStep.OPERATION].end()
 
-    def on_validate(self) -> Generator[None, None, None]:
+    def on_validate(
+        self, execution_context: ExecutionContext
+    ) -> Generator[None, None, None]:
         ctx = trace.set_span_in_context(self._span_holder[LifecycleStep.OPERATION])
         self._span_holder[LifecycleStep.VALIDATION] = self._tracer.start_span(
             "GraphQL Validation",
@@ -89,7 +93,9 @@ def on_validate(self) -> Generator[None, None, None]:
         yield
         self._span_holder[LifecycleStep.VALIDATION].end()
 
-    def on_parse(self) -> Generator[None, None, None]:
+    def on_parse(
+        self, execution_context: ExecutionContext
+    ) -> Generator[None, None, None]:
         ctx = trace.set_span_in_context(self._span_holder[LifecycleStep.OPERATION])
         self._span_holder[LifecycleStep.PARSE] = self._tracer.start_span(
             "GraphQL Parsing", context=ctx