Feat #4329: Add on_subscription_result hook to SchemaExtension

RELEASE.md

@@ -0,0 +1,7 @@
+Release type: minor
+
+Adds a new `on_subscription_result` hook to `SchemaExtension` that allows extensions to interact with and mutate the stream of events yielded by GraphQL subscriptions.
+
+Previously, extensions were only triggered during the initial setup phase of a subscription, meaning transport layers (like WebSockets) bypassed them during the actual data streaming phase. This new hook solves this by executing right before each result is yielded to the client.
+
+As part of this architectural update, the built-in `MaskErrors` extension has been updated to use this new hook, ensuring that sensitive exceptions are now correctly masked during WebSocket subscriptions.

strawberry/extensions/base_extension.py

@@ -9,7 +9,7 @@
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
 
-    from strawberry.types import ExecutionContext
+    from strawberry.types import ExecutionContext, ExecutionResult
 
 
 class LifecycleStep(Enum):
@@ -51,6 +51,14 @@ def on_execute(  # type: ignore
         """Called before and after the execution step."""
         yield None
 
+    def on_subscription_result(
+        self, result: ExecutionResult
+    ) -> None | AwaitableOrValue[None]:
+        """Called exactly once for every event/result yielded by a GraphQL subscription.
+
+        Extensions can mutate the `result` object directly (e.g., masking errors).
+        """
+
     def resolve(
         self,
         _next: Callable,

strawberry/extensions/mask_errors.py

@@ -59,3 +59,7 @@ def on_operation(self) -> Iterator[None]:
             self._process_result(result)
         elif result:
             self._process_result(result.initial_result)
+
+    def on_subscription_result(self, result: StrawberryExecutionResult) -> None:
+        """Mask errors on streaming subscription results."""
+        self._process_result(result)

strawberry/extensions/runner.py

@@ -12,7 +12,7 @@
 from strawberry.utils.await_maybe import await_maybe
 
 if TYPE_CHECKING:
-    from strawberry.types import ExecutionContext
+    from strawberry.types import ExecutionContext, ExecutionResult
 
     from . import SchemaExtension
 
@@ -40,6 +40,14 @@ def parsing(self) -> ParsingContextManager:
     def executing(self) -> ExecutingContextManager:
         return ExecutingContextManager(self.extensions)
 
+    async def on_subscription_result(self, result: ExecutionResult) -> None:
+        """Run the subscription result hook across all active extensions."""
+        for extension in self.extensions:
+            # Check if the extension implemented the new hook
+            hook = getattr(extension, "on_subscription_result", None)
+            if hook:
+                await await_maybe(hook(result))
+
     def get_extensions_results_sync(self) -> dict[str, Any]:
         data: dict[str, Any] = {}
         for extension in self.extensions:

strawberry/schema/schema.py

@@ -875,20 +875,29 @@ async def _subscribe(
                     try:
                         async with aclosing(aiter_or_result):
                             async for result in aiter_or_result:
-                                yield await self._handle_execution_result(
+                                extension_result = await self._handle_execution_result(
                                     execution_context,
                                     result,
                                     extensions_runner,
                                 )
+
+                                await extensions_runner.on_subscription_result(
+                                    extension_result
+                                )
+
+                                yield extension_result
+
                     # graphql-core doesn't handle exceptions raised while executing.
                     except Exception as exc:  # noqa: BLE001
-                        yield await self._handle_execution_result(
+                        execution_result = await self._handle_execution_result(
                             execution_context,
                             OriginalExecutionResult(
                                 data=None, errors=[_coerce_error(exc)]
                             ),
                             extensions_runner,
                         )
+                        await extensions_runner.on_subscription_result(execution_result)
+                        yield execution_result
             # catch exceptions raised in `on_execute` hook.
             except Exception as exc:  # noqa: BLE001
                 origin_result = OriginalExecutionResult(

tests/extensions/test_subscription_hook.py

@@ -0,0 +1,87 @@
+from collections.abc import AsyncGenerator
+
+import pytest
+
+import strawberry
+from strawberry.extensions import SchemaExtension
+from strawberry.extensions.mask_errors import MaskErrors
+from strawberry.types import ExecutionResult
+
+
+# Dummy extension that uses the new hook
+class StreamModifierExtension(SchemaExtension):
+    def on_subscription_result(self, result: ExecutionResult) -> None:
+        if result.data and "count" in result.data:
+            # Mutate the outgoing data stream
+            result.data["count"] = f"Modified: {result.data['count']}"
+
+
+# Create a basic schema with a subscription
+@strawberry.type
+class Query:
+    hello: str = "world"
+
+
+@strawberry.type
+class Subscription:
+    @strawberry.subscription
+    async def count(self) -> AsyncGenerator[int, None]:
+        yield 1
+        yield 2
+
+    @strawberry.subscription
+    async def dangerous_stream(self) -> AsyncGenerator[int, None]:
+        yield 1
+        raise ValueError("Secret database credentials leaked!")
+
+
+# Tests
+@pytest.mark.asyncio
+async def test_extension_modifies_subscription_stream():
+    schema = strawberry.Schema(
+        query=Query, subscription=Subscription, extensions=[StreamModifierExtension]
+    )
+
+    query = "subscription { count }"
+
+    sub_generator = await schema.subscribe(query)
+
+    # Get all yielded results
+    results = [result async for result in sub_generator]
+
+    assert len(results) == 2
+
+    # Assert that the extension successfully intercepted and modified the stream
+    assert results[0].data["count"] == "Modified: 1"
+    assert results[1].data["count"] == "Modified: 2"
+
+
+@pytest.mark.asyncio
+async def test_mask_errors_scrubs_subscription_exceptions():
+    # Initialize schema with the MaskErrors extension
+    schema = strawberry.Schema(
+        query=Query, subscription=Subscription, extensions=[MaskErrors()]
+    )
+
+    query = "subscription { dangerousStream }"
+
+    sub_generator = await schema.subscribe(query)
+
+    # Get all yielded results
+    results = [result async for result in sub_generator]
+
+    # We expect 2 results: the successful yield, and the error
+    assert len(results) == 2
+
+    # Assert the first yield worked normally
+    assert results[0].data["dangerousStream"] == 1
+    assert not results[0].errors
+
+    # Assert the error was caught and MASKED
+    assert results[1].data is None
+    assert len(results[1].errors) == 1
+
+    # The crucial check: The raw exception message MUST NOT be exposed
+    error_message = results[1].errors[0].message
+    assert error_message == "Unexpected error."
+    assert "Secret database credentials" not in error_message