Closed farzbood closed 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}")
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
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