WA for FastApi session manager

This commit is contained in:
Dmitry Belousov
2024-04-12 01:20:35 +03:00
parent 28b2bfec50
commit 2aa0a4f1b0

View File

@@ -5,6 +5,7 @@ import logging
import shutil import shutil
import sqlite3 import sqlite3
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from os import PathLike from os import PathLike
from pathlib import Path from pathlib import Path
from typing import Optional, Union from typing import Optional, Union
@@ -188,9 +189,9 @@ async def _backup_task() -> Optional[Path]:
async def _metrics() -> None: async def _metrics() -> None:
"""Create some useful prometheus metrics""" """Create some useful prometheus metrics"""
logger.info("Start metrics collection") logger.info("Start metrics collection")
if __db is None: async with get_session() as session:
raise RuntimeError("DB is not setup.") await filament_metrics(session)
await asyncio.gather(*[filament_metrics(__db.session_maker()), spool_metrics(__db.session_maker())]) await spool_metrics(session)
logger.info("End metrics collection") logger.info("End metrics collection")
@@ -230,3 +231,18 @@ async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
raise exc raise exc
finally: finally:
await session.close() await session.close()
@asynccontextmanager
async def get_session() -> AsyncSession:
if __db is None or __db.session_maker is None:
raise RuntimeError("DB is not setup.")
async with __db.session_maker() as session:
try:
yield session
await session.commit()
except Exception as exc:
await session.rollback()
raise exc
finally:
await session.close()