Access to app state in CLI plugins

[WulkanHefajstos]

Hi!

I'm trying to build quite simple CLI plugin similar to flask's shell, that would allow us to easily play around with SQLAlchemy models & queries by using injected app object that already has access to DB. Our app relies on couple of start up hooks / plugins, including SQLAlchemyInitPlugin. Some of these hooks modify app.state in functions that are added to AppConfig.lifespan field ( like in SQLAlchemyInitPlugin).

I might be completely wrong, but it seems these lifespan hooks are only called within lifespan method of Litestar class (i.e. here). This method is not called / context manager it returns is never entered before the function defined as part of the CLI plugin is called, hence app state is empty in the shell that we create.

I'm wondering if someone could help me to find out if app.lifespan is part of the public API and whether i can simply use it on my own (see example below) 🙂

Here is a simplified implementation of the plugin that injects app and session_maker variables into IPython terminal namespace (Ipython supports async code execution). It seems to work just fine, but i'm not sure that's a right way

import asyncio
import sys
from functools import wraps

from click import Group
from litestar import Litestar
from litestar.plugins import CLIPlugin
from sqlalchemy.ext.asyncio import async_sessionmaker


def _make_coroutine(f):
    @wraps
    def wrapper(*args, **kwargs):
        return asyncio.run(f(*args, **kwargs))
    
    return wrapper
    
    
class ShellPlugin(CLIPlugin):
    def on_cli_init(self, cli: Group) -> None:
        @cli.command(name="shell", help="Start IPython repl with async support")
        @make_coroutine
        async def shell(app: Litestar) -> None:
            async with app.lifespan():
                session_maker: type[async_sessionmaker] = app.state["session_maker_class"]
                await asyncio.gather(
                    asyncio.to_thread(_start_shell, app=app, session_maker=session_maker)
                )


def _start_shell(app: Litestar, session_maker: type[async_sessionmaker]) -> None:
    # args accepted by function will be injected by IPython.embed into shell
    import IPython
    
    IPython.embed(
        using="asyncio",
        banner2="The coconut nut is a giant nut; If you eat too much, you'll get very fat",
    )

Cheers!

[peterschutt]

Litestar.lifespan() is documented: https://docs.litestar.dev/2/reference/app.html#litestar.app.Litestar.lifespan

The lifespan is only entered by the application in response to the ASGI lifespan startup message from the server, which doesn't happen for cli operations (except litestar run of course).

Your approach looks reasonable to me. Are you encountering any issues?

Another way you can approach it is to have your own lifespan context that is separate from, but called from within the litestar lifespan context. This way you can enter resource contexts separate from the application.

[WulkanHefajstos]

Your approach looks reasonable to me. Are you encountering any issues?

Nope, at least not yet 🙃

I guess i'll share it with my team mates and we'll see how it goes. Thanks a lot!

[peterschutt]

Your approach looks reasonable to me. Are you encountering any issues?

Nope, at least not yet 🙃

I guess i'll share it with my team mates and we'll see how it goes. Thanks a lot!

No worries! Feel free to come back with any feedback 👍