Working with beanie/motor transactions/sessions

[leonballoniMakes]

I have been working with beanie for some time and there are a few operations in my code that i would like to run like a transaction to avoid inserting half completed information to certain collections (all or nothing gets in). Is there someone that have worked with this? I have tried a few things using fastapi, beanie and motor but all seem to failure and found no tutorial to replicate.

Here is my MongoDbClient that i built to run the necessary Datalayer services

class MongodbClient:
    def __init__(self,
                 db: str = None,
                 username: str = None,
                 password: str = None,
                 cluster: str = 'localhost',
                 port: int = 27017,
                 conn_string: str = None,
                 models: List[beanie.Document] = __all__):
        self.db = db if db else 'test'
        self.username = username if username else ''
        self.password = password if password else ''
        self.cluster = cluster
        self.port = port
        self.models = models
        if conn_string:
            self.conn_string = conn_string
        elif self.cluster == 'localhost':
            self.conn_string = f"mongodb://{self.cluster}:{self.port}/{self.db}"
        else:
            self.conn_string = f"mongodb+srv://{self.username}:{self.password}@{self.cluster}"
        self.client = motor.motor_asyncio.AsyncIOMotorClient(self.conn_string, maxPoolSize=100, minPoolSize=0)
        self.database_beanied_list = []

    async def beanie_connection(self, database):
        try:
            if database not in self.database_beanied_list:
                await beanie.init_beanie(database=self.client[database or self.db], document_models=self.models or __all__)
                self.database_beanied_list.append(database)
        except ConnectionError as error:
            logger.error(f'Error to init beanie for database {database}: {error}')
            raise error 

    def close_connection(self):
        self.client.close()

    # @asynccontextmanager
    async def get_db(self):
        database = self.client[self.db]
        await self.beanie_connection(self.db)
        # await beanie.init_beanie(database, document_models=self.models)
        session = await self.client.start_session()
        async with session.start_transaction():
            try:
                yield database
                await session.commit_transaction()
            except Exception as e:
                await session.abort_transaction() 
                raise e
            finally:
                await session.end_session()
                # self.client.close()

At the current moment I made an object of my mongoclient at the start of my controllers and added the get_db method as a Dependency injection to my fastapi route:

Example:

chat_client = MongodbClient(db='prototype', conn_string=MONGO_STRING, models=[Chat])
def init_routes(app: FastAPI):
    @app.get('/', status_code=200)
    def health_check():
        return {'status': 'ok'}
    

    app.include_router(
        ChatApi(
            ChatService(
                ChatDomain()
                    )
                ).router,
        tags=["chat"],
        prefix="/chat"
       dependencies=[Depends(clients.get_db)] 
 )

The problem here is that my transaction is been commited before the session.commit_transaction() is run (at the yield database)

Could anyone help me out?

[leonballoniMakes]

Found a way to make it work. Here is a code example

import logging
from contextlib import asynccontextmanager
from datetime import datetime
from functools import wraps
from typing import Annotated, Sequence, Type

import motor.motor_asyncio
import uvicorn
from beanie import Document, init_beanie
from fastapi import Depends, FastAPI
from pydantic import BaseModel, Field


class ItemModel(BaseModel):
    name: str
    description: str
    date: datetime = Field(default_factory=lambda: datetime.now())


class Item(Document, ItemModel):

    class Settings:
        collection = "items"


client = motor.motor_asyncio.AsyncIOMotorClient('mongodb://localhost:27017')


async def get_client():
    async with await client.start_session() as session:
        async with session.start_transaction():
            try:
                yield session
                session.commit_transaction()
            except Exception as exp:
                session.abort_transaction()
                raise exp
            finally:
                logging.info(session.session_id)
                # session.end_session()


app = FastAPI(dependencies=[Depends(get_client)])


@app.on_event("startup")
async def startup_event():
    database = client['test_db']
    await init_beanie(database, document_models=[Item])


@app.post("/items/")
async def create_item(item: ItemModel,
                      session: Type[motor.motor_asyncio.AsyncIOMotorClientSession] = Depends(get_client)):
    try:
        new_item = Item(**item.model_dump())
        new_item = await Item.insert_one(new_item, session=session)
    except:
        raise Exception
    return new_item


app_base_configs = {
    "host": "0.0.0.0",
    "port": 5000,
    "workers": 3,
    "access_log": True,
    "reload": True
}
if __name__ == "__main__":
    uvicorn.run("test:app", **app_base_configs)

[yousafsabir]

Is there a way to initialize the session in the body of the route?
by using await in the expression:

session = await get_client()

I get

object async_generator can't be used in 'await' expression.

and without it, I get

'async_generator' object has no attribute 'in_transaction'.

I need it to make it work in some cron jobs where I can use Depends