src/database/core.py - master - Ulyanoff/mlops1

src/database/core.py
51 строка2 KB
Artem Ivochkin
первый рабочий асинхронный тест с замоканной БД
21 ноя 2025, 17:09
21 ноя 2025, 17:09e7b8312
# -*- coding: utf-8 -*-
import logging

from sqlalchemy import URL
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession, create_async_engine, async_sessionmaker

from config.db_mdl import db_cfg

boot = logging.getLogger('boot')

# куда и к кому подключаться
db_url = URL.create(
    'postgresql+asyncpg',
    database=db_cfg.name, username=db_cfg.user, password=db_cfg.password,
    host=db_cfg.host, port=db_cfg.port,
)
# создать асинхронный движок (открыть пул коннектов)
try:
    async_engine = create_async_engine(
        db_url,
        echo=False,  # логирование команд SQL
        pool_size=db_cfg.pool_size,  # размер пула соединений
        max_overflow=db_cfg.pool_overflow,  # на сколько можно превысить размер пула
        echo_pool=True,  # логировать ведение пула
        hide_parameters=True,  # подавление вывода параметров SQL
        pool_pre_ping=True,  # проверка, жив ли коннект
    )
except Exception as e:
    boot.error(f'Database initialization failed: {e}')
    exit(2)

# создать фабрику сессий
async_session_maker = async_sessionmaker(bind=async_engine, class_=AsyncSession, expire_on_commit=False)

# Dependency для эндпойнтов
async def get_conn() -> AsyncConnection:
    """Получение коннекта"""
    async with async_engine.connect() as conn:
        yield conn  # будет стоять в этой точке до освобождения коннекта, после чего контекстный менеджер его закроет
        # await conn.close()

# Dependency для эндпойнтов
async def get_session() -> AsyncSession:
    """Получение сессии"""
    async with async_session_maker() as session:
        yield session  # будет стоять в этой точке до освобождения сессии, после чего контекстный менеджер её закроет
        # await session.close()

async def db_closed():
    """Закрытие пула коннектов"""
    await async_engine.dispose()