tortoise / tortoise-orm

Familiar asyncio ORM for python, built with relations in mind
https://tortoise.github.io
Apache License 2.0
4.39k stars 359 forks source link

FastAPI app's lifespan() event-handler conflicts with TortoiseORM way of initialisation by app.on_event() #1450

Closed farzbood closed 10 months ago

farzbood commented 11 months ago

Describe the bug When set the "lifespan" of an FastAPI app, "register_tortoise" function in "tortoise.contrib.fastapi" "init.py" would fail, because can't set both "lifespan()" and "on_event()" function for an app, simultaneously (according to starlette inline docs.) https://github.com/encode/starlette/blob/d6007d7198c35c1a7ed81e678a81c3bca86bee5e/starlette/applications.py#L63

To Reproduce Set the "lifespan()" of a FastaAPI app and run the server.

Expected behavior That the TortoiseORM be compatible with the situation that the "lifespan" is set for the app and add its handler in addition to the existing one.

Additional context An implementaion of such a lifespan-manager exist here. https://github.com/uriyyo/fastapi-lifespan-manager/tree/main#fastapi-lifespanmanager

All The Best V

waketzheng commented 10 months ago

You can do it like this way(examples/fastapi/main.py):

import asyncio
from types import ModuleType
from typing import AsyncIterator, Dict, Iterable, List, Optional, Union

from fastapi import FastAPI
from fastapi_lifespan_manager import LifespanManager, State
from models import User_Pydantic, UserIn_Pydantic, Users
from pydantic import BaseModel  # pylint: disable=E0611
from starlette.exceptions import HTTPException
from starlette.requests import Request
from starlette.responses import JSONResponse

from tortoise import Tortoise, connections
from tortoise.exceptions import DoesNotExist, IntegrityError
from tortoise.log import logger

manager = LifespanManager()

class RegisterTortoise:
    @staticmethod
    async def close() -> None:
        await connections.close_all()
        logger.info("Tortoise-ORM shutdown")

    @staticmethod
    def add_handler(app) -> None:
        @app.exception_handler(DoesNotExist)
        async def doesnotexist_exception_handler(request: Request, exc: DoesNotExist):
            return JSONResponse(status_code=404, content={"detail": str(exc)})

        @app.exception_handler(IntegrityError)
        async def integrityerror_exception_handler(request: Request, exc: IntegrityError):
            return JSONResponse(
                status_code=422,
                content={"detail": [{"loc": [], "msg": str(exc), "type": "IntegrityError"}]},
            )

    def __init__(
        self,
        app: FastAPI,
        config: Optional[dict] = None,
        config_file: Optional[str] = None,
        db_url: Optional[str] = None,
        modules: Optional[Dict[str, Iterable[Union[str, ModuleType]]]] = None,
        generate_schemas: bool = False,
        add_exception_handlers: bool = False,
    ) -> None:
        self.config = config
        self.config_file = config_file
        self.db_url = db_url
        self.modules = modules
        self.generate_schemas = generate_schemas
        if add_exception_handlers:
            self.add_handler(app)

    def __await__(self):
        yield from asyncio.create_task(self.init())
        return self

    async def init(self) -> None:
        config = self.config
        config_file = self.config_file
        db_url = self.db_url
        modules = self.modules
        generate_schemas = self.generate_schemas
        await Tortoise.init(config=config, config_file=config_file, db_url=db_url, modules=modules)
        logger.info("Tortoise-ORM started, %s, %s", connections._get_storage(), Tortoise.apps)
        if generate_schemas:
            logger.info("Tortoise-ORM generating schema")
            await Tortoise.generate_schemas()

@manager.add
async def setup_db(app: FastAPI) -> AsyncIterator[State]:
    orm = await RegisterTortoise(
        app,
        db_url="sqlite://:memory:",
        modules={"models": ["models"]},
        generate_schemas=True,
        add_exception_handlers=True,
    )
    yield {"db": orm}
    await orm.close()

app = FastAPI(title="Tortoise ORM FastAPI example", lifespan=manager)

class Status(BaseModel):
    message: str

@app.get("/users", response_model=List[User_Pydantic])
async def get_users():
    return await User_Pydantic.from_queryset(Users.all())

@app.post("/users", response_model=User_Pydantic)
async def create_user(user: UserIn_Pydantic):
    user_obj = await Users.create(**user.model_dump(exclude_unset=True))
    return await User_Pydantic.from_tortoise_orm(user_obj)

@app.get("/user/{user_id}", response_model=User_Pydantic)
async def get_user(user_id: int):
    return await User_Pydantic.from_queryset_single(Users.get(id=user_id))

@app.put("/user/{user_id}", response_model=User_Pydantic)
async def update_user(user_id: int, user: UserIn_Pydantic):
    await Users.filter(id=user_id).update(**user.model_dump(exclude_unset=True))
    return await User_Pydantic.from_queryset_single(Users.get(id=user_id))

@app.delete("/user/{user_id}", response_model=Status)
async def delete_user(user_id: int):
    deleted_count = await Users.filter(id=user_id).delete()
    if not deleted_count:
        raise HTTPException(status_code=404, detail=f"User {user_id} not found")
    return Status(message=f"Deleted user {user_id}")
farzbood commented 10 months ago

Thank you for the response with a clean code example. Although this technique will cover the case in hand perfectly, I think such a Lifespan-Management mechanism would be more intuitive as a built-in feature (by FastAPI of course) and I opened a discussion there already, but got no answer yet! https://github.com/tiangolo/fastapi/discussions/10083

So, once again thanks for the response and hoping the people at FastAPI consider this a priority since many others would soon meet such a situation. PEACE V