litestar-org / litestar

Production-ready, Light, Flexible and Extensible ASGI API framework | Effortlessly Build Performant APIs
https://litestar.dev/
MIT License
5.06k stars 347 forks source link

Enhancement: Context Managers as Dependencies #3528

Open cofin opened 1 month ago

cofin commented 1 month ago

Summary

We currently support generator based dependencies, but not context managers. For instance the following works:

     async def provide_connection(
        self,
        state: State,
        scope: Scope,
    ) -> AsyncGenerator[Union[PoolConnectionProxy,Connection], None]:  # noqa: UP007
        """Create a connection instance.

        Args:
            state: The ``Litestar.state`` instance.
            scope: The current connection's scope.

        Returns:
            A connection instance.
        """
        connection = cast("Optional[Union[PoolConnectionProxy,Connection]]", get_scope_state(scope, CONNECTION_SCOPE_KEY))
        if connection is None:
            pool = cast("Pool", state.get(self.pool_app_state_key))

            async with pool.acquire() as connection:
                set_scope_state(scope, CONNECTION_SCOPE_KEY, connection)
                yield connection

But adding an context manager decorator like this:

    @asynccontextmanager
    async def provide_connection(
        self,
        state: State,
        scope: Scope,
    ) -> AsyncGenerator[Union[PoolConnectionProxy,Connection], None]:  # noqa: UP007
        """Create a connection instance.

        Args:
            state: The ``Litestar.state`` instance.
            scope: The current connection's scope.

        Returns:
            A connection instance.
        """
        connection = cast("Optional[Union[PoolConnectionProxy,Connection]]", get_scope_state(scope, CONNECTION_SCOPE_KEY))
        if connection is None:
            pool = cast("Pool", state.get(self.pool_app_state_key))

            async with pool.acquire() as connection:
                set_scope_state(scope, CONNECTION_SCOPE_KEY, connection)
                yield connection

produces

msgspec.ValidationError: Unsupported type: <class 'contextlib._AsyncGeneratorContextManager'> - at `$.db_connection`

Basic Example

No response

Drawbacks and Impact

No response

Unresolved questions

No response


[!NOTE]
While we are open for sponsoring on GitHub Sponsors and OpenCollective, we also utilize Polar.sh to engage in pledge-based sponsorship.

Check out all issues funded or available for funding on our Polar.sh dashboard

  • If you would like to see an issue prioritized, make a pledge towards it!
  • We receive the pledge once the issue is completed & verified
  • This, along with engagement in the community, helps us know which features are a priority to our users.

Fund with Polar

provinzkraut commented 2 weeks ago

I don't think we can reasonably support this without requiring to add some abstractions that make it rather useless as a feature. The issue here is that we cannot distinguish between a dependency that is meant to be used as a contextmanager, and a dependency that just happens to be a contextmanager.

E.g.

from sqlalchemy.orm import Session

def provide_session() -> Session:
 session = get_session() # gets session from somewhere
 return session

app = Litestar(depenencies={"session": provide_session})

this wouldn't work, because Session is a context manager, and Litestar would try to call its __enter__ and __exit__ methods, causing unintended side-effects.

We could solve this by providing some wrapper that indicates that this dependency provider should be treated as a context manager, but then again you can achieve the same thing with

async def provide_connection() -> Union[PoolConnectionProxy, Connection]:
 async with thing.provide_connection() as conn:
  yield conn