BeanieODM/beanie

[BUG] (some) type annotations for session wrong

entropymatters opened this issue · 3 comments

Describe the bug
A comment in #536 gives an example of using database sessions and transactions with beanie. However, when one wants to factor out the session into a separate function and annotates it, it does not properly type check using mypy anymore.
Interestingly, this is not the case for all Document methods, despite beanie annotating session as pymongo.client_session.ClientSession everywhere.

I'd be glad to create a PR if the issue is really only to replace pymongo.client_session.ClientSession with motor.motor_asyncio.AsyncIOMotorClientSession, and would have a look in any other hints you can give for a fix.

Note that just casting the session to a pymongo ClientSession does not work since its start_session method is no coroutine.

To Reproduce

import asyncio
from contextlib import asynccontextmanager
from typing import AsyncGenerator
from motor import motor_asyncio
from beanie import Document, init_beanie


class Sample(Document):
    name: str

@asynccontextmanager
async def get_db_session(client: motor_asyncio.AsyncIOMotorClient) -> AsyncGenerator[motor_asyncio.AsyncIOMotorClientSession, None]:
    async with await client.start_session() as session:
        yield session

async def main() -> None:
    client = motor_asyncio.AsyncIOMotorClient('mongodb://mongo')
    database = client['mydatabase']
    await init_beanie(database=database, document_models=[Sample])

    async with get_db_session(client) as session:
        async with session.start_transaction():
            await Sample(name='alice').insert(session=session)  # OK
            await Sample.find(session=session).to_list()        # not OK

if __name__ == '__main__':
    asyncio.run(main())

Current behavior
mypy shows an error for find() (but not for insert()):

error: No overload variant of "find" of "FindInterface" matches argument type "AsyncIOMotorClientSession"  [call-overload]
note: Possible overload variants:
note:     def find(cls, *args: Mapping[str, Any] | bool, projection_model: None = ..., skip: int | None = ..., limit: int | None = ..., sort: str | list[tuple[str, SortDirection]] | None = ..., session: ClientSession | None = ..., ignore_cache: bool = ..., fetch_links: bool = ..., with_children: bool = ..., lazy_parse: bool = ..., nesting_depth: int | None = ..., nesting_depths_per_field: dict[str, int] | None = ..., **pymongo_kwargs: Any) -> FindMany[Document]
note:     def [DocumentProjectionType <: BaseModel] find(cls, *args: Mapping[str, Any] | bool, projection_model: type[DocumentProjectionType], skip: int | None = ..., limit: int | None = ..., sort: str | list[tuple[str, SortDirection]] | None = ..., session: ClientSession | None = ..., ignore_cache: bool = ..., fetch_links: bool = ..., with_children: bool = ..., lazy_parse: bool = ..., nesting_depth: int | None = ..., nesting_depths_per_field: dict[str, int] | None = ..., **pymongo_kwargs: Any) -> FindMany[DocumentProjectionType]

Expected behavior
Code type checks using mypy

Additional context
none

Hi! Thank you for the report. I'll pick it up during the next bug fixing session

That would be nice.
This is basically the same problem that I reported some time ago but unfortunately it was closed.
see Issue: #706

I can confirm that this only is a typing issue. Passing the motor.motor_asyncio.AsyncIOMotorClientSession to the session parameter works fine and the transaction is acutally used. I worked around the typing issue with # type: ignore

async with await db.client.start_session() as session:
    async with session.start_transaction():
        await mydocument.insert(session=session)  # type: ignore