from __future__ import annotations
import sys
import warnings
from collections.abc import Generator, Iterable
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from types import ModuleType
from typing import TYPE_CHECKING
from tortoise import Tortoise, connections
from tortoise.exceptions import DoesNotExist, IntegrityError
from tortoise.log import logger
if TYPE_CHECKING:
from fastapi import FastAPI, Request
if sys.version_info >= (3, 11):
from typing import Self
else:
from typing_extensions import Self
[docs]
def tortoise_exception_handlers() -> dict:
from fastapi.responses import JSONResponse
async def doesnotexist_exception_handler(request: Request, exc: DoesNotExist):
return JSONResponse(status_code=404, content={"detail": str(exc)})
async def integrityerror_exception_handler(request: Request, exc: IntegrityError):
return JSONResponse(
status_code=422,
content={"detail": [{"loc": [], "msg": str(exc), "type": "IntegrityError"}]},
)
return {
DoesNotExist: doesnotexist_exception_handler,
IntegrityError: integrityerror_exception_handler,
}
[docs]
class RegisterTortoise(AbstractAsyncContextManager):
"""
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)``.
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',
}
}
}
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.
use_tz:
A boolean that specifies if datetime will be timezone-aware by default or not.
timezone:
Timezone to use, default is UTC.
Raises
------
ConfigurationError
For any configuration error
"""
def __init__(
self,
app: FastAPI | None = None,
config: dict | None = None,
config_file: str | None = None,
db_url: str | None = None,
modules: dict[str, Iterable[str | ModuleType]] | None = None,
generate_schemas: bool = False,
add_exception_handlers: bool = False,
use_tz: bool = False,
timezone: str = "UTC",
_create_db: 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
self.use_tz = use_tz
self.timezone = timezone
self._create_db = _create_db
if add_exception_handlers and app is not None:
from starlette.middleware.exceptions import ExceptionMiddleware
warnings.warn(
"Setting `add_exception_handlers` to be true is deprecated, "
"use `FastAPI(exception_handlers=tortoise_exception_handlers())` instead."
"See more about it on https://tortoise.github.io/examples/fastapi",
DeprecationWarning,
)
original_call_func = ExceptionMiddleware.__call__
async def wrap_middleware_call(self, *args, **kw) -> None:
if DoesNotExist not in self._exception_handlers:
self._exception_handlers.update(tortoise_exception_handlers())
await original_call_func(self, *args, **kw)
ExceptionMiddleware.__call__ = wrap_middleware_call # type:ignore
[docs]
async def init_orm(self) -> None: # pylint: disable=W0612
await Tortoise.init(
config=self.config,
config_file=self.config_file,
db_url=self.db_url,
modules=self.modules,
use_tz=self.use_tz,
timezone=self.timezone,
_create_db=self._create_db,
)
logger.info("Tortoise-ORM started, %s, %s", connections._get_storage(), Tortoise.apps)
if self.generate_schemas:
logger.info("Tortoise-ORM generating schema")
await Tortoise.generate_schemas()
[docs]
@staticmethod
async def close_orm() -> None: # pylint: disable=W0612
await connections.close_all()
logger.info("Tortoise-ORM shutdown")
def __call__(self, *args, **kwargs) -> Self:
return self
async def __aenter__(self) -> Self:
await self.init_orm()
return self
async def __aexit__(self, *args, **kw) -> None:
await self.close_orm()
def __await__(self) -> Generator[None, None, Self]:
async def _self() -> Self:
return await self.__aenter__()
return _self().__await__()
[docs]
def register_tortoise(
app: FastAPI,
config: dict | None = None,
config_file: str | None = None,
db_url: str | None = None,
modules: dict[str, Iterable[str | ModuleType]] | None = None,
generate_schemas: bool = False,
add_exception_handlers: bool = False,
) -> None:
"""
Registers Tortoise-ORM with set-up at the beginning of FastAPI application's lifespan
(which allow user to read/write data from/to db inside the lifespan function),
and tear-down at the end of that lifespan.
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',
}
}
}
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.
Raises
------
ConfigurationError
For any configuration error
"""
from fastapi.routing import _merge_lifespan_context
# Leave this function here to compare with old versions
# So people can upgrade tortoise-orm in running project without changing any code
@asynccontextmanager
async def orm_lifespan(app_instance: FastAPI):
async with RegisterTortoise(
app_instance,
config,
config_file,
db_url,
modules,
generate_schemas,
):
yield
original_lifespan = app.router.lifespan_context
app.router.lifespan_context = _merge_lifespan_context(orm_lifespan, original_lifespan)
if add_exception_handlers:
for exp_type, endpoint in tortoise_exception_handlers().items():
app.exception_handler(exp_type)(endpoint)