Getting a db session outside of a route handler

[aranvir]

Hi, I was wondering what the best way is to get a db session outside of a route handler. I know that with the SQLAlchemyInitPlugin installed I can get references to the engine and to a session via dependency injection (default: db_engine, db_session).

If I am using a security backend like the session backend, I might want to get the User instance from the database - thus, I need a session.

Before, I would just create a session maker from sqlalchemy_config.create_session_maker() that I made available globally. However, while writing tests I realized that I set this up in a way that I could not guarantee that this session maker was configured for the same database as my tests were.

One solution would be to change that and ensure that the global reference is always correctly configured. However, I think the safer way is to instead have the app instance (e.g., as set up by the test client) provide the session maker. I produced something like this but it is not working, e.g. when using it in the retrieve_user_handler of the session backend. It can't find the session_maker_class key within the connection.state.

What am I doing wrong? Is the connection state not yet fully populated when the retrieve_user_handler is called? Are there other (better) ways to get a db session from the connection.app?

Snippet below for example usage.

async def session_maker(connection: ASGIConnection) -> AsyncSession:
    """Fetch a session for the current connection."""
    session_provider = connection.app.dependencies.get("db_session")
    session = await session_provider(state=connection.state, scope=connection.scope)
    return session

...

async def retrieve_user_handler(session: dict[str, Any], connection: ASGIConnection) -> User:
    """Fetches the current user state from the database and returns it as User instance."""
    db_session: AsyncSession = await session_maker(connection)
    user_id = session.get(SESSION_USER_ID_KEY)
    async with db_session.begin():
        user = await db_session.get(
            entity=User,
            ident=(user_id,),
        )
...

[Alc-Alc]

Can you try session = await session_provider(state=connection.scope['app'].state, scope=connection.scope)

[Alc-Alc]

cc: @peterschutt

[aranvir]

Thanks @Alc-Alc your suggestion works! What also works is state=connection.app.state.

Had a bit of tunnel vision when I put together this workaround so I jumped on the first state attribute I could find - was not really considering which state is needed here (connection or app).

[aranvir]

I have a follow-up question on this.. as this workaround might occasionally break later use of the session in this request. I'm still learning about the details of async sqlalchemy so maybe there is a simple way of solving this that I am not aware of.

I use the approach above to provide a session from a connection to functions that do not have access to dependency injection like the retrieve_user_handler. My User object has a relationship (list of elements) that I load with selectinload within the retrieve_user_handler, looks a bit like this:

async def retrieve_user_handler(session: dict[str, Any], connection: ASGIConnection) -> User | Admin:
    """Fetches the current user state from the database and returns it as User instance."""
    db_session: AsyncSession = await session_maker(connection)
    user_id = session.get(SESSION_USER_ID_KEY)
    async with db_session.begin():
        user = await db_session.get(
            entity=User,
            ident=(user_id,),
            options=(selectinload(User.projects),)
        )
        if user is None:
            raise exceptions.HTTPException(status_code=status_codes.HTTP_400_BAD_REQUEST, detail="User not found.")
        db_session.expunge(user)
    return user

Now, for some route handlers where I do similar selectinload queries I get the error

sqlalchemy.exc.MissingGreenlet: greenlet_spawn has not been called; can't call await_only() here. Was IO attempted in an unexpected place?

And I dont get it: The primary object I'm querying is still returned correctly, but any relationship objects are not. At one point I also got the message that the "current session was already used", so I'm wondering if the error actually comes from my retrieve_user_handler code and that I use the same db_session twice?

[peterschutt]

Have you configured your session_maker with expire_on_commit=False? This is worth reading if you're using the ORM with async sqlalchemy.

[aranvir]

Yep, it is. This is my plugin setup

    session_config = AsyncSessionConfig(expire_on_commit=False)
    sqlalchemy_config = SQLAlchemyAsyncConfig(
        connection_string=get_connection_string(),
        session_config=session_config
    )
    sqlalchemy_init_plugin = SQLAlchemyInitPlugin(config=sqlalchemy_config)

Before, I would get the session maker directly from SQLAlchemyAsyncConfig and that worked. So, I adjusted my helper function to get the session maker from the config instead of the dependency - and this works better. All my tests pass again, since now 2 different sessions are used.

async def session_maker(connection: ASGIConnection) -> AsyncSession:
    """Fetch a session for the current connection."""
    plugin: SQLAlchemyInitPlugin = connection.app.plugins.get(SQLAlchemyInitPlugin)
    return plugin._config.create_session_maker()()

Still a bit confused why this happens. E.g., I did check the docs you mentioned when I ran into some lazy loading issues when working with DTOs. I then solved it with selectinload where needed. But here it really seems that the session is "partially expired"? I still could use it to query other elements but just not to resolve any relationships via selectinload.