118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
import logging
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
import pydantic
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
|
|
|
|
class DatabaseConfig(pydantic.BaseModel):
|
|
URL: pydantic.PostgresDsn
|
|
ECHO: bool = False
|
|
ECHO_POOL: bool = False
|
|
POOL_SIZE: int = 50
|
|
MAX_OVERFLOW: int = 10
|
|
|
|
|
|
class DatabaseHelper:
|
|
def __init__(self, config: DatabaseConfig, migrations_path: str = 'migration/versions') -> None:
|
|
self.config = config
|
|
self.migrations_path = migrations_path
|
|
self.engine: AsyncEngine | None = None
|
|
self.session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
|
|
async def connect(self) -> None:
|
|
if self.engine is not None:
|
|
logging.warning('Database already connected')
|
|
return
|
|
|
|
self.engine = create_async_engine(
|
|
url=str(self.config.URL),
|
|
echo=self.config.ECHO,
|
|
echo_pool=self.config.ECHO_POOL,
|
|
pool_size=self.config.POOL_SIZE,
|
|
max_overflow=self.config.MAX_OVERFLOW,
|
|
)
|
|
|
|
self.session_factory = async_sessionmaker(
|
|
bind=self.engine,
|
|
autoflush=False,
|
|
autocommit=False,
|
|
expire_on_commit=False,
|
|
)
|
|
|
|
await self.ping()
|
|
await self._check_migrations()
|
|
|
|
logging.info('Database initialized')
|
|
|
|
async def ping(self) -> None:
|
|
try:
|
|
async with self.session_getter() as session:
|
|
await session.execute(text('SELECT 1'))
|
|
except Exception as e:
|
|
logging.error(f'Database ping failed: {e}')
|
|
|
|
async def close(self) -> None:
|
|
if self.engine is None:
|
|
logging.warning('Database already closed')
|
|
return
|
|
|
|
await self.engine.dispose()
|
|
self.engine = None
|
|
self.session_factory = None
|
|
|
|
logging.info('Database closed')
|
|
|
|
@asynccontextmanager
|
|
async def session_getter(self) -> AsyncGenerator[AsyncSession]:
|
|
"""Контекстный менеджер для получения сессии"""
|
|
if self.session_factory is None:
|
|
raise RuntimeError('Database not connected. Call connect() first.')
|
|
|
|
async with self.session_factory() as session:
|
|
yield session
|
|
|
|
async def _check_migrations(self) -> None:
|
|
migrations_dir = Path(self.migrations_path)
|
|
if not migrations_dir.exists():
|
|
return
|
|
|
|
latest_file = next(
|
|
(
|
|
f
|
|
for f in sorted(migrations_dir.glob('*.py'), reverse=True)
|
|
if f.name != '__init__.py'
|
|
),
|
|
None,
|
|
)
|
|
if not latest_file:
|
|
return
|
|
|
|
latest_version = latest_file.stem.split('_')[0]
|
|
|
|
async with self.session_getter() as session:
|
|
result = await session.execute(
|
|
text(
|
|
"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'alembic_version')"
|
|
)
|
|
)
|
|
if not result.scalar():
|
|
logging.error('Migrations not applied. Run: alembic upgrade head')
|
|
return
|
|
|
|
result = await session.execute(text('SELECT version_num FROM alembic_version'))
|
|
current_version = result.scalar()
|
|
|
|
if current_version != latest_version:
|
|
logging.warning(
|
|
f'Database version {current_version} is outdated (latest: {latest_version}). Run: alembic upgrade head'
|
|
)
|