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,11 @@ 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:
+            await await_maybe(extension.on_subscription_result(result))
+
     def get_extensions_results_sync(self) -> dict[str, Any]:
         data: dict[str, Any] = {}
         for extension in self.extensions:

strawberry/schema/schema.py

@@ -833,9 +833,11 @@ async def _subscribe(
                 initial_error.extensions = (
                     await extensions_runner.get_extensions_results(execution_context)
                 )
-                yield await self._handle_execution_result(
+                execution_result = await self._handle_execution_result(
                     execution_context, initial_error, extensions_runner
                 )
+                await extensions_runner.on_subscription_result(execution_result)
+                yield execution_result
             try:
                 async with extensions_runner.executing():
                     assert execution_context.graphql_document is not None
@@ -866,39 +868,52 @@ async def _subscribe(
 
                 # Handle pre-execution errors.
                 if isinstance(aiter_or_result, OriginalExecutionResult):
-                    yield await self._handle_execution_result(
+                    execution_result = await self._handle_execution_result(
                         execution_context,
                         PreExecutionError(data=None, errors=aiter_or_result.errors),
                         extensions_runner,
                     )
+                    await extensions_runner.on_subscription_result(execution_result)
+                    yield execution_result
                 else:
                     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(
                     data=None, errors=[_coerce_error(exc)]
                 )
-                yield await self._handle_execution_result(
+                execution_result = await self._handle_execution_result(
                     execution_context,
                     origin_result,
                     extensions_runner,
                 )
+                await extensions_runner.on_subscription_result(execution_result)
+                yield execution_result
 
     async def subscribe(
         self,

tests/extensions/test_subscription_hook.py

@@ -0,0 +1,157 @@
+import asyncio
+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
+
+
+class AsyncStreamModifierExtension(SchemaExtension):
+    side_effect_ran = False
+
+    async def _side_effect(self) -> None:
+        # Simulate an async dependency / side effect
+        await asyncio.sleep(0)
+        AsyncStreamModifierExtension.side_effect_ran = True
+
+    async def on_subscription_result(self, result: ExecutionResult) -> None:
+        # This should be awaited by extensions_runner.on_subscription_result
+        await self._side_effect()
+
+        if result.data and "count" in result.data:
+            # Mutate the outgoing data stream after the async side effect
+            result.data["count"] = f"Modified: {result.data['count']}"
+
+
+@pytest.mark.asyncio
+async def test_async_on_subscription_result_is_awaited() -> None:
+    # Reset class-level flag before running the subscription
+    AsyncStreamModifierExtension.side_effect_ran = False
+
+    schema = strawberry.Schema(
+        query=Query,
+        subscription=Subscription,
+        extensions=[AsyncStreamModifierExtension()],
+    )
+
+    query = "subscription { count }"
+    sub_generator = await schema.subscribe(query)
+
+    # Consume first result from the async iterator
+    first_result = await sub_generator.__anext__()
+
+    assert first_result.errors is None
+    assert first_result.data == {"count": "Modified: 1"}
+    assert AsyncStreamModifierExtension.side_effect_ran is True
+
+
+@pytest.mark.asyncio
+async def test_mask_errors_scrubs_pre_execution_errors():
+    # Initialize schema with MaskErrors
+    schema = strawberry.Schema(
+        query=Query, subscription=Subscription, extensions=[MaskErrors()]
+    )
+
+    # Querying a field that doesn't exist triggers Validation errors BEFORE execution
+    query = "subscription { fieldThatDoesNotExist }"
+
+    # Run the subscription
+    sub_generator = await schema.subscribe(query)
+
+    # Exhaust the generator
+    results = [result async for result in sub_generator]
+
+    # Strawberry yields 1 or more errors depending on the validation layers it hits.
+    # We must ensure that EVERY result yielded was successfully intercepted and masked.
+    assert len(results) > 0
+
+    for result in results:
+        assert result.data is None
+        assert len(result.errors) == 1
+
+        # The crucial check: MaskErrors successfully intercepted and masked it!
+        error_message = result.errors[0].message
+        assert error_message == "Unexpected error."
+        assert "fieldThatDoesNotExist" not in error_message