From 7ded5c7cdbdcdc56d5def36c8cf865abc27c9823 Mon Sep 17 00:00:00 2001 From: Waket Zheng Date: Sat, 27 Apr 2024 14:03:06 +0800 Subject: [PATCH] Enhancement for FastAPI lifespan support (#1371)(#1576) (#1541) --- CHANGELOG.rst | 1 + docs/contrib/fastapi.rst | 2 +- examples/fastapi/README.rst | 2 +- examples/fastapi/main.py | 31 +++-- pyproject.toml | 1 + tortoise/contrib/fastapi/__init__.py | 195 +++++++++++++++++++++++---- 6 files changed, 190 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d6e58603e..951024d60 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,6 +25,7 @@ Added - Add binary compression support for `UUIDField` in `MySQL`. (#1458) - Only `Model`, `Tortoise`, `BaseDBAsyncClient`, `__version__`, and `connections` are now exported from `tortoise` - Add parameter `validators` to `pydantic_model_creator`. (#1471) +- Enhancement for FastAPI lifespan support (#1371) Fixed ^^^^^ diff --git a/docs/contrib/fastapi.rst b/docs/contrib/fastapi.rst index 12b90785b..c518e522a 100644 --- a/docs/contrib/fastapi.rst +++ b/docs/contrib/fastapi.rst @@ -4,7 +4,7 @@ Tortoise-ORM FastAPI integration ================================ -We have a lightweight integration util ``tortoise.contrib.fastapi`` which has a single function ``register_tortoise`` which sets up Tortoise-ORM on startup and cleans up on teardown. +We have a lightweight integration util ``tortoise.contrib.fastapi`` which has a class ``RegisterTortoise`` that can be used to set/clean up Tortoise-ORM in lifespan context. FastAPI is basically Starlette & Pydantic, but in a very specific way. diff --git a/examples/fastapi/README.rst b/examples/fastapi/README.rst index f2b92fddc..3ac1d2bad 100644 --- a/examples/fastapi/README.rst +++ b/examples/fastapi/README.rst @@ -1,7 +1,7 @@ Tortoise-ORM FastAPI example ============================ -We have a lightweight integration util ``tortoise.contrib.fastapi`` which has a single function ``register_tortoise`` which sets up Tortoise-ORM on startup and cleans up on teardown. +We have a lightweight integration util ``tortoise.contrib.fastapi`` which has a class ``RegisterTortoise`` that can be used to set/clean up Tortoise-ORM in lifespan context. Usage ----- diff --git a/examples/fastapi/main.py b/examples/fastapi/main.py index a6d0507f6..a7744579f 100644 --- a/examples/fastapi/main.py +++ b/examples/fastapi/main.py @@ -1,4 +1,5 @@ # pylint: disable=E0611,E0401 +from contextlib import asynccontextmanager from typing import List from fastapi import FastAPI @@ -6,9 +7,26 @@ from pydantic import BaseModel from starlette.exceptions import HTTPException -from tortoise.contrib.fastapi import register_tortoise +from tortoise.contrib.fastapi import RegisterTortoise -app = FastAPI(title="Tortoise ORM FastAPI example") + +@asynccontextmanager +async def lifespan(app: FastAPI): + # app startup + async with RegisterTortoise( + app, + db_url="sqlite://:memory:", + modules={"models": ["models"]}, + generate_schemas=True, + add_exception_handlers=True, + ): + # db connected + yield + # app teardown + # db connections closed + + +app = FastAPI(title="Tortoise ORM FastAPI example", lifespan=lifespan) class Status(BaseModel): @@ -43,12 +61,3 @@ async def delete_user(user_id: int): if not deleted_count: raise HTTPException(status_code=404, detail=f"User {user_id} not found") return Status(message=f"Deleted user {user_id}") - - -register_tortoise( - app, - db_url="sqlite://:memory:", - modules={"models": ["models"]}, - generate_schemas=True, - add_exception_handlers=True, -) diff --git a/pyproject.toml b/pyproject.toml index 136e4199a..47750630f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ classifiers = [ "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", "Programming Language :: PL/SQL", diff --git a/tortoise/contrib/fastapi/__init__.py b/tortoise/contrib/fastapi/__init__.py index 72baa7d60..aa9df8874 100644 --- a/tortoise/contrib/fastapi/__init__.py +++ b/tortoise/contrib/fastapi/__init__.py @@ -1,32 +1,30 @@ +from __future__ import annotations + +import warnings +from contextlib import AbstractAsyncContextManager, asynccontextmanager from types import ModuleType -from typing import Dict, Iterable, Optional, Union +from typing import TYPE_CHECKING, Dict, Iterable, Optional, Union -from fastapi import FastAPI +from fastapi.responses import JSONResponse from pydantic import BaseModel # pylint: disable=E0611 -from starlette.requests import Request -from starlette.responses import JSONResponse +from starlette.routing import _DefaultLifespan from tortoise import Tortoise, connections from tortoise.exceptions import DoesNotExist, IntegrityError from tortoise.log import logger +if TYPE_CHECKING: + from fastapi import FastAPI, Request + class HTTPNotFoundError(BaseModel): detail: str -def register_tortoise( - 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: +class RegisterTortoise(AbstractAsyncContextManager): """ - Registers ``startup`` and ``shutdown`` events to set-up and tear-down Tortoise-ORM - inside a FastAPI application. + Registers Tortoise-ORM with set-up and tear-down + inside a FastAPI application's lifespan. You can configure using only one of ``config``, ``config_file`` and ``(db_url, modules)``. @@ -89,28 +87,167 @@ def register_tortoise( For any configuration error """ - @app.on_event("startup") - async def init_orm() -> None: # pylint: disable=W0612 + 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.app = app + 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: + + @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"}]}, + ) + + async def init_orm(self) -> None: # pylint: disable=W0612 + config, config_file = self.config, self.config_file + db_url, modules = self.db_url, self.modules 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: + if self.generate_schemas: logger.info("Tortoise-ORM generating schema") await Tortoise.generate_schemas() - @app.on_event("shutdown") + @staticmethod async def close_orm() -> None: # pylint: disable=W0612 await connections.close_all() logger.info("Tortoise-ORM shutdown") - if add_exception_handlers: + def __call__(self, *args, **kwargs) -> "RegisterTortoise": + return self + + async def __aenter__(self) -> "RegisterTortoise": + await self.init_orm() + return self + + async def __aexit__(self, *args, **kw): + await self.close_orm() + + +def register_tortoise( + 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: + """ + Registers ``startup`` and ``shutdown`` events to set-up and tear-down Tortoise-ORM + inside a FastAPI application. + + You can configure using only one of ``config``, ``config_file`` + and ``(db_url, modules)``. + + Parameters + ---------- + app: + FastAPI app. + config: + Dict containing config: + + Example + ------- + + .. code-block:: python3 + + { + 'connections': { + # Dict format for connection + 'default': { + 'engine': 'tortoise.backends.asyncpg', + 'credentials': { + 'host': 'localhost', + 'port': '5432', + 'user': 'tortoise', + 'password': 'qwerty123', + 'database': 'test', + } + }, + # Using a DB_URL string + 'default': 'postgres://postgres:qwerty123@localhost:5432/events' + }, + 'apps': { + 'models': { + 'models': ['__main__'], + # If no default_connection specified, defaults to 'default' + 'default_connection': 'default', + } + } + } - @app.exception_handler(DoesNotExist) - async def doesnotexist_exception_handler(request: Request, exc: DoesNotExist): - return JSONResponse(status_code=404, content={"detail": str(exc)}) + config_file: + Path to .json or .yml (if PyYAML installed) file containing config with + same format as above. + db_url: + Use a DB_URL string. See :ref:`db_url` + modules: + Dictionary of ``key``: [``list_of_modules``] that defined "apps" and modules that + should be discovered for models. + generate_schemas: + True to generate schema immediately. Only useful for dev environments + or SQLite ``:memory:`` databases + add_exception_handlers: + True to add some automatic exception handlers for ``DoesNotExist`` & ``IntegrityError``. + This is not recommended for production systems as it may leak data. - @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"}]}, - ) + Raises + ------ + ConfigurationError + For any configuration error + """ + orm = RegisterTortoise( + app, + config, + config_file, + db_url, + modules, + generate_schemas, + add_exception_handlers, + ) + if isinstance(lifespan := app.router.lifespan_context, _DefaultLifespan): + # Leave on_event here to compare with old versions + # So people can upgrade tortoise-orm in running project without changing any code + + @app.on_event("startup") # type: ignore[unreachable] + async def init_orm() -> None: # pylint: disable=W0612 + await orm.init_orm() + + @app.on_event("shutdown") + async def close_orm() -> None: # pylint: disable=W0612 + await orm.close_orm() + + else: + # If custom lifespan was passed to app, register tortoise in it + warnings.warn( + "`register_tortoise` function is deprecated, " + "use the `RegisterTortoise` class instead." + "See more about it on https://tortoise.github.io/examples/fastapi", + DeprecationWarning, + ) + + @asynccontextmanager + async def orm_lifespan(app_instance: "FastAPI"): + async with orm: + async with lifespan(app_instance): + yield + + app.router.lifespan_context = orm_lifespan