Repository: strawberry-graphql/strawberry

Issue:
Surprising FIFO behaviour of lifecycle hooks
## Describe the (maybe) Bug

I'm surprised that various lifecycle hooks (`on_operation`, `on_parse`, `on_validate`, `on_execute`) are completed in a FIFO fashion, rather than LIFO.

I would expect that if we're wrapping an operation with `on_operation` with 2 extension the following will happen (LIFO):

* First extension starts (before `yield` part)
* Second extension starts (before `yield` part)
* Second extension completes (after `yield` part)
* First extension completes (after `yield` part)

However, the order I'm _actually_ seeing is the following (FIFO):

* First extension starts (before `yield` part)
* Second extension starts (before `yield` part)
* First extension completes (after `yield` part)
* Second extension completes (after `yield` part)

I'm concerned about it because extension can mutate state, so it would be good for them to behave like a context manager. [Example of state mutation.](https://strawberry.rocks/docs/guides/custom-extensions#execution-context)

In fact, I do find it surprising that this is how things work out. Notably, overriding `resolve` doesn't have the same problem – but it also happens in a slightly different way.

## Repro details

Here're some toy extensions I built to investigate things:

```python
class MyCustomExtension(SchemaExtension):
    id = "?"

    @override
    async def on_validate(self) -> AsyncGenerator[None, None]:
        print(f"GraphQL validation start ({self.id})")
        yield
        print(f"GraphQL validation end ({self.id})")

    @override
    def on_parse(self) -> Generator[None, None, None]:
        print(f"GraphQL parsing start ({self.id})")
        yield
        print(f"GraphQL parsing end ({self.id})")

    @override
    def on_execute(self) -> Generator[None, None, None]:
        print(f"GraphQL execution start ({self.id})")
        yield
        print(f"GraphQL execution end ({self.id})")

    @override
    def on_operation(self) -> Generator[None, None, None]:
        print(f"GraphQL operation start ({self.id})")
        yield
        print(f"GraphQL operation end ({self.id})")

    @override
    async def resolve(
        self,
        _next: Callable[..., object],
        root: object,
        info: GraphQLResolveInfo,
        *args,
        **kwargs,
    ) -> AwaitableOrValue[object]:
        random_id = randint(0, 1000)

        print(f"GraphQL resolver {random_id} start ({self.id})")
        result = await await_maybe(_next(root, info, *args, **kwargs))
        print(f"GraphQL resolver {random_id} end ({self.id})")
        return result


class MyCustomExtensionA(MyCustomExtension):
    id = "A"


class MyCustomExtensionB(MyCustomExtension):
    id = "B"
```

I'm testing it by running a simple query against a GraphQL Schema:

```python
@strawberry.type
class Me:
    id: str

@strawberry.type
class Query:
    @strawberry.field
    @staticmethod
    async def me() -> Me:
        return Me(id="foo")

schema = MySchema(
    query=Query,
    extensions=[MyCustomExtensionA, MyCustomExtensionB],
)
```

When running a simple GraphQL query against this schema:

```graphql
query {
  me { id }
}
```

I see the following lines being printed:

```
GraphQL operation start (A)
GraphQL operation start (B)
GraphQL parsing start (A)
GraphQL parsing start (B)
GraphQL parsing end (A)
GraphQL parsing end (B)
GraphQL validation start (A)
GraphQL validation start (B)
GraphQL validation end (A)
GraphQL validation end (B)
GraphQL execution start (A)
GraphQL execution start (B)
GraphQL resolver 598 start (B)
GraphQL resolver 975 start (A)
GraphQL resolver 975 end (A)
GraphQL resolver 598 end (B)
GraphQL resolver 196 start (B)
GraphQL resolver 638 start (A)
GraphQL resolver 638 end (A)
GraphQL resolver 196 end (B)
GraphQL execution end (A)
GraphQL execution end (B)
GraphQL operation end (A)
GraphQL operation end (B)
```

## System Information

 - Strawberry version (if applicable): `0.220.0`

<!-- POLAR PLEDGE BADGE START -->
## Upvote & Fund

- We're using [Polar.sh](https://polar.sh/strawberry-graphql) so you can upvote and help fund this issue.
- We receive the funding once the issue is completed & confirmed by you.
- Thank you in advance for helping prioritize & fund our backlog.

<a href="https://polar.sh/strawberry-graphql/strawberry/issues/3413">
<picture>
  <source media="(prefers-color-scheme: dark)" srcset="https://polar.sh/api/github/strawberry-graphql/strawberry/issues/3413/pledge.svg?darkmode=1">
  <img alt="Fund with Polar" src="https://polar.sh/api/github/strawberry-graphql/strawberry/issues/3413/pledge.svg">
</picture>
</a>
<!-- POLAR PLEDGE BADGE END -->



Discussion:
I'm not sure if this is actually a bug... Perhaps this is a natural consequence of the Python pattern Strawberry is leveraging?

However, it would be nice to get a second pair of eyes and opinion on this!
It is actually tested as the wanted behaviour for some reason
https://github.com/strawberry-graphql/strawberry/blob/22838db26389c28450820e0f0e9771e015403e8c/tests/schema/extensions/test_extensions.py#L420

Resolving patch:
diff --git a/strawberry/extensions/context.py b/strawberry/extensions/context.py
--- a/strawberry/extensions/context.py
+++ b/strawberry/extensions/context.py
@@ -2,14 +2,15 @@
 
 import contextlib
 import inspect
+import types
 import warnings
 from asyncio import iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncIterator,
+    AsyncContextManager,
     Callable,
-    Iterator,
+    ContextManager,
     List,
     NamedTuple,
     Optional,
@@ -28,12 +29,18 @@
 
 class WrappedHook(NamedTuple):
     extension: SchemaExtension
-    initialized_hook: Union[AsyncIterator[None], Iterator[None]]
+    hook: Callable[..., Union[AsyncContextManager[None], ContextManager[None]]]
     is_async: bool
 
 
 class ExtensionContextManagerBase:
-    __slots__ = ("hooks", "deprecation_message", "default_hook")
+    __slots__ = (
+        "hooks",
+        "deprecation_message",
+        "default_hook",
+        "async_exit_stack",
+        "exit_stack",
+    )
 
     def __init_subclass__(cls):
         cls.DEPRECATION_MESSAGE = (
@@ -73,10 +80,20 @@ def get_hook(self, extension: SchemaExtension) -> Optional[WrappedHook]:
 
         if hook_fn:
             if inspect.isgeneratorfunction(hook_fn):
-                return WrappedHook(extension, hook_fn(extension), False)
+                context_manager = contextlib.contextmanager(
+                    types.MethodType(hook_fn, extension)
+                )
+                return WrappedHook(
+                    extension=extension, hook=context_manager, is_async=False
+                )
 
             if inspect.isasyncgenfunction(hook_fn):
-                return WrappedHook(extension, hook_fn(extension), True)
+                context_manager_async = contextlib.asynccontextmanager(
+                    types.MethodType(hook_fn, extension)
+                )
+                return WrappedHook(
+                    extension=extension, hook=context_manager_async, is_async=True
+                )
 
             if callable(hook_fn):
                 return self.from_callable(extension, hook_fn)
@@ -96,27 +113,31 @@ def from_legacy(
     ) -> WrappedHook:
         if iscoroutinefunction(on_start) or iscoroutinefunction(on_end):
 
+            @contextlib.asynccontextmanager
             async def iterator():
                 if on_start:
                     await await_maybe(on_start())
+
                 yield
+
                 if on_end:
                     await await_maybe(on_end())
 
-            hook = iterator()
-            return WrappedHook(extension, hook, True)
+            return WrappedHook(extension=extension, hook=iterator, is_async=True)
 
         else:
 
-            def iterator():
+            @contextlib.contextmanager
+            def iterator_async():
                 if on_start:
                     on_start()
+
                 yield
+
                 if on_end:
                     on_end()
 
-            hook = iterator()
-            return WrappedHook(extension, hook, False)
+            return WrappedHook(extension=extension, hook=iterator_async, is_async=False)
 
     @staticmethod
     def from_callable(
@@ -125,59 +146,34 @@ def from_callable(
     ) -> WrappedHook:
         if iscoroutinefunction(func):
 
-            async def async_iterator():
+            @contextlib.asynccontextmanager
+            async def iterator():
                 await func(extension)
                 yield
 
-            hook = async_iterator()
-            return WrappedHook(extension, hook, True)
+            return WrappedHook(extension=extension, hook=iterator, is_async=True)
         else:
 
+            @contextlib.contextmanager
             def iterator():
                 func(extension)
                 yield
 
-            hook = iterator()
-            return WrappedHook(extension, hook, False)
+            return WrappedHook(extension=extension, hook=iterator, is_async=False)
 
-    def run_hooks_sync(self, is_exit: bool = False) -> None:
-        """Run extensions synchronously."""
-        ctx = (
-            contextlib.suppress(StopIteration, StopAsyncIteration)
-            if is_exit
-            else contextlib.nullcontext()
-        )
-        for hook in self.hooks:
-            with ctx:
-                if hook.is_async:
-                    raise RuntimeError(
-                        f"SchemaExtension hook {hook.extension}.{self.HOOK_NAME} "
-                        "failed to complete synchronously."
-                    )
-                else:
-                    hook.initialized_hook.__next__()  # type: ignore[union-attr]
-
-    async def run_hooks_async(self, is_exit: bool = False) -> None:
-        """Run extensions asynchronously with support for sync lifecycle hooks.
-
-        The ``is_exit`` flag is required as a `StopIteration` cannot be raised from
-        within a coroutine.
-        """
-        ctx = (
-            contextlib.suppress(StopIteration, StopAsyncIteration)
-            if is_exit
-            else contextlib.nullcontext()
-        )
+    def __enter__(self) -> None:
+        self.exit_stack = contextlib.ExitStack()
 
-        for hook in self.hooks:
-            with ctx:
-                if hook.is_async:
-                    await hook.initialized_hook.__anext__()  # type: ignore[union-attr]
-                else:
-                    hook.initialized_hook.__next__()  # type: ignore[union-attr]
+        self.exit_stack.__enter__()
 
-    def __enter__(self):
-        self.run_hooks_sync()
+        for hook in self.hooks:
+            if hook.is_async:
+                raise RuntimeError(
+                    f"SchemaExtension hook {hook.extension}.{self.HOOK_NAME} "
+                    "failed to complete synchronously."
+                )
+            else:
+                self.exit_stack.enter_context(hook.hook())  # type: ignore
 
     def __exit__(
         self,
@@ -185,10 +181,18 @@ def __exit__(
         exc_val: Optional[BaseException],
         exc_tb: Optional[TracebackType],
     ):
-        self.run_hooks_sync(is_exit=True)
+        self.exit_stack.__exit__(exc_type, exc_val, exc_tb)
 
-    async def __aenter__(self):
-        await self.run_hooks_async()
+    async def __aenter__(self) -> None:
+        self.async_exit_stack = contextlib.AsyncExitStack()
+
+        await self.async_exit_stack.__aenter__()
+
+        for hook in self.hooks:
+            if hook.is_async:
+                await self.async_exit_stack.enter_async_context(hook.hook())  # type: ignore
+            else:
+                self.async_exit_stack.enter_context(hook.hook())  # type: ignore
 
     async def __aexit__(
         self,
@@ -196,7 +200,7 @@ async def __aexit__(
         exc_val: Optional[BaseException],
         exc_tb: Optional[TracebackType],
     ):
-        await self.run_hooks_async(is_exit=True)
+        await self.async_exit_stack.__aexit__(exc_type, exc_val, exc_tb)
 
 
 class OperationContextManager(ExtensionContextManagerBase):


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -68,11 +68,9 @@ jobs:
             .nox
           key:
             ${{ runner.os }}-nox-${{ matrix.session.session }}-${{
-            hashFiles('**/poetry.lock') }}-${{ hashFiles('**/noxfile.py') }}-2
+            hashFiles('**/poetry.lock') }}-${{ hashFiles('**/noxfile.py') }}-3
 
       - run: pip install poetry nox nox-poetry
-      - run: poetry --help
-      - run: which poetry
       - run: nox -r -t tests -s "${{ matrix.session.session }}"
       - uses: actions/upload-artifact@v4
         if: ${{ always() }}
diff --git a/tests/schema/extensions/test_extensions.py b/tests/schema/extensions/test_extensions.py
--- a/tests/schema/extensions/test_extensions.py
+++ b/tests/schema/extensions/test_extensions.py
@@ -407,18 +407,18 @@ async def on_execute(self):
         "ExtensionB, on_operation Entered",
         "ExtensionA, on_parse Entered",
         "ExtensionB, on_parse Entered",
-        "ExtensionA, on_parse Exited",
         "ExtensionB, on_parse Exited",
+        "ExtensionA, on_parse Exited",
         "ExtensionA, on_validate Entered",
         "ExtensionB, on_validate Entered",
-        "ExtensionA, on_validate Exited",
         "ExtensionB, on_validate Exited",
+        "ExtensionA, on_validate Exited",
         "ExtensionA, on_execute Entered",
         "ExtensionB, on_execute Entered",
-        "ExtensionA, on_execute Exited",
         "ExtensionB, on_execute Exited",
-        "ExtensionA, on_operation Exited",
+        "ExtensionA, on_execute Exited",
         "ExtensionB, on_operation Exited",
+        "ExtensionA, on_operation Exited",
     ]
 
 
@@ -684,18 +684,19 @@ class ExtensionB(SchemaExtension):
         def on_execute(self):
             execution_order.append(type(self))
             yield
+            execution_order.append(type(self))
 
     class ExtensionC(SchemaExtension):
         def on_execute(self):
             execution_order.append(type(self))
             yield
+            execution_order.append(type(self))
 
     @strawberry.type
     class Query:
         food: str = "strawberry"
 
-    extensions = [ExtensionB, ExtensionC]
-    schema = strawberry.Schema(query=Query, extensions=extensions)
+    schema = strawberry.Schema(query=Query, extensions=[ExtensionB, ExtensionC])
 
     query = """
         query TestQuery {
@@ -707,7 +708,7 @@ class Query:
 
     assert not result.errors
     assert result.data == {"food": "strawberry"}
-    assert execution_order == extensions
+    assert execution_order == [ExtensionB, ExtensionC, ExtensionC, ExtensionB]
 
 
 def test_async_extension_in_sync_context():
@@ -1024,3 +1025,20 @@ def hi(self) -> str:
         ValueError, match="Hook on_operation on <(.*)> must be callable, received 'ABC'"
     ):
         schema.execute_sync(query)
+
+
+@pytest.mark.asyncio
+async def test_calls_hooks_when_there_are_errors(async_extension):
+    @strawberry.type
+    class Query:
+        @strawberry.field
+        def hi(self) -> str:
+            raise Exception("This is an error")
+
+    schema = strawberry.Schema(query=Query, extensions=[async_extension])
+
+    query = "{ hi }"
+
+    result = await schema.execute(query)
+    assert result.errors
+    async_extension.perform_test()