166 lines
5.3 KiB
Python
166 lines
5.3 KiB
Python
import asyncio
|
|
import logging
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import asynccontextmanager
|
|
from contextvars import ContextVar
|
|
|
|
import pydantic
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
|
|
_session_ctx: ContextVar[AsyncSession | None] = ContextVar('session', default=None)
|
|
|
|
|
|
class DatabaseConfig(pydantic.BaseModel):
|
|
URL: pydantic.PostgresDsn
|
|
ECHO: bool = False
|
|
ECHO_POOL: bool = False
|
|
POOL_SIZE: int = 50
|
|
MAX_OVERFLOW: int = 10
|
|
|
|
|
|
class DatabaseBase:
|
|
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
|
|
self._connect_lock = asyncio.Lock()
|
|
|
|
async def connect(self) -> None:
|
|
async with self._connect_lock:
|
|
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:
|
|
if self.session_factory is None:
|
|
raise RuntimeError('Database not connected. Call connect() first.')
|
|
|
|
try:
|
|
async with self.session_factory() as session:
|
|
await session.execute(text('SELECT 1'))
|
|
except Exception as e:
|
|
logging.error(f'Database ping failed: {e}')
|
|
|
|
async def close(self) -> None:
|
|
async with self._connect_lock:
|
|
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')
|
|
|
|
@property
|
|
def session(self) -> AsyncSession:
|
|
session = _session_ctx.get()
|
|
if session is None:
|
|
raise RuntimeError('No active transaction. Use transaction() or begin() first.')
|
|
return session
|
|
|
|
async def begin(self) -> None:
|
|
if _session_ctx.get() is not None:
|
|
raise RuntimeError('Transaction already started')
|
|
|
|
if self.session_factory is None:
|
|
raise RuntimeError('Database not connected. Call connect() first.')
|
|
|
|
session = self.session_factory()
|
|
_session_ctx.set(session)
|
|
|
|
async def commit(self) -> None:
|
|
session = _session_ctx.get()
|
|
if session is None:
|
|
raise RuntimeError('No active transaction')
|
|
|
|
try:
|
|
await session.commit()
|
|
finally:
|
|
await session.close()
|
|
_session_ctx.set(None)
|
|
|
|
async def rollback(self) -> None:
|
|
session = _session_ctx.get()
|
|
if session is None:
|
|
raise RuntimeError('No active transaction')
|
|
|
|
try:
|
|
await session.rollback()
|
|
finally:
|
|
await session.close()
|
|
_session_ctx.set(None)
|
|
|
|
@asynccontextmanager
|
|
async def transaction(self) -> AsyncGenerator[None]:
|
|
await self.begin()
|
|
try:
|
|
yield
|
|
await self.commit()
|
|
except Exception:
|
|
await self.rollback()
|
|
raise
|
|
|
|
async def _check_migrations(self) -> None:
|
|
if self.session_factory is None:
|
|
raise RuntimeError('Database not connected. Call connect() first.')
|
|
|
|
async with self.session_factory() 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 not current_version:
|
|
logging.error('No migration version found in database. Run: alembic upgrade head')
|
|
return
|
|
|
|
try:
|
|
from alembic.config import Config
|
|
from alembic.script import ScriptDirectory
|
|
|
|
alembic_cfg = Config('alembic.ini')
|
|
script = ScriptDirectory.from_config(alembic_cfg)
|
|
head_revision = script.get_current_head()
|
|
|
|
if head_revision and current_version != head_revision:
|
|
logging.warning(
|
|
f'Database version {current_version} is outdated (latest: {head_revision}). '
|
|
f'Run: alembic upgrade head'
|
|
)
|
|
except Exception as e:
|
|
logging.debug(f'Could not check migration head: {e}')
|