mirror of
https://github.com/trushildhokiya/allininx-2.git
synced 2025-03-15 18:48:40 +00:00
78 lines
2.0 KiB
Python
78 lines
2.0 KiB
Python
|
from typing import Awaitable, Callable
|
||
|
|
||
|
from fastapi import FastAPI
|
||
|
from sqlalchemy.ext.asyncio import async_sessionmaker
|
||
|
|
||
|
from reworkd_platform.db.meta import meta
|
||
|
from reworkd_platform.db.models import load_all_models
|
||
|
from reworkd_platform.db.utils import create_engine
|
||
|
from reworkd_platform.services.tokenizer.lifetime import init_tokenizer
|
||
|
|
||
|
|
||
|
def _setup_db(app: FastAPI) -> None: # pragma: no cover
|
||
|
"""
|
||
|
Creates connection to the database.
|
||
|
|
||
|
This function creates SQLAlchemy engine instance,
|
||
|
session_factory for creating sessions
|
||
|
and stores them in the application's state property.
|
||
|
|
||
|
:param app: fastAPI application.
|
||
|
"""
|
||
|
engine = create_engine()
|
||
|
session_factory = async_sessionmaker(
|
||
|
engine,
|
||
|
expire_on_commit=False,
|
||
|
)
|
||
|
app.state.db_engine = engine
|
||
|
app.state.db_session_factory = session_factory
|
||
|
|
||
|
|
||
|
async def _create_tables() -> None: # pragma: no cover
|
||
|
"""Populates tables in the database."""
|
||
|
load_all_models()
|
||
|
|
||
|
engine = create_engine()
|
||
|
async with engine.begin() as connection:
|
||
|
await connection.run_sync(meta.create_all)
|
||
|
await engine.dispose()
|
||
|
|
||
|
|
||
|
def register_startup_event(
|
||
|
app: FastAPI,
|
||
|
) -> Callable[[], Awaitable[None]]: # pragma: no cover
|
||
|
"""
|
||
|
Actions to run on application startup.
|
||
|
|
||
|
This function uses fastAPI app to store data
|
||
|
in the state, such as db_engine.
|
||
|
|
||
|
:param app: the fastAPI application.
|
||
|
:return: function that actually performs actions.
|
||
|
"""
|
||
|
|
||
|
@app.on_event("startup")
|
||
|
async def _startup() -> None: # noqa: WPS430
|
||
|
_setup_db(app)
|
||
|
init_tokenizer(app)
|
||
|
# await _create_tables()
|
||
|
|
||
|
return _startup
|
||
|
|
||
|
|
||
|
def register_shutdown_event(
|
||
|
app: FastAPI,
|
||
|
) -> Callable[[], Awaitable[None]]: # pragma: no cover
|
||
|
"""
|
||
|
Actions to run on application's shutdown.
|
||
|
|
||
|
:param app: fastAPI application.
|
||
|
:return: function that actually performs actions.
|
||
|
"""
|
||
|
|
||
|
@app.on_event("shutdown")
|
||
|
async def _shutdown() -> None: # noqa: WPS430
|
||
|
await app.state.db_engine.dispose()
|
||
|
|
||
|
return _shutdown
|