Enhancement for FastAPI lifespan support (#1371)(#1576) (#1541)

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
 ^^^^^

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.
 

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
 -----

examples/fastapi/main.py

@@ -1,14 +1,32 @@
 # pylint: disable=E0611,E0401
+from contextlib import asynccontextmanager
 from typing import List
 
 from fastapi import FastAPI
 from models import User_Pydantic, UserIn_Pydantic, Users
 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,
-)

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",

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