82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
import pydantic
|
||
from tortoise import Tortoise
|
||
|
||
|
||
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) -> None:
|
||
self.config = config
|
||
|
||
async def connect(self) -> None:
|
||
await Tortoise.init(
|
||
db_url=str(self.config.URL),
|
||
modules={'models': ['src.domain']},
|
||
use_tz=True,
|
||
timezone='UTC',
|
||
)
|
||
|
||
await self.ping()
|
||
await self._check_migrations()
|
||
|
||
async def close(self) -> None:
|
||
await Tortoise.close_connections()
|
||
|
||
async def ping(self) -> None:
|
||
from tortoise import connections
|
||
|
||
conn = connections.get('default')
|
||
try:
|
||
await conn.execute_query('SELECT 1')
|
||
except Exception as e:
|
||
import logging
|
||
|
||
logging.error(f'Database ping failed: {e}')
|
||
raise
|
||
|
||
async def _check_migrations(self) -> None:
|
||
import logging
|
||
|
||
from tortoise import connections
|
||
|
||
conn = connections.get('default')
|
||
|
||
# Проверяем существование таблицы alembic_version
|
||
result = await conn.execute_query(
|
||
"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'alembic_version')"
|
||
)
|
||
if not result[1][0]['exists']:
|
||
logging.error('Migrations not applied. Run: alembic upgrade head')
|
||
return
|
||
|
||
# Получаем текущую версию
|
||
result = await conn.execute_query('SELECT version_num FROM alembic_version')
|
||
if not result[1]:
|
||
logging.error('No migration version found in database. Run: alembic upgrade head')
|
||
return
|
||
|
||
current_version = result[1][0]['version_num']
|
||
|
||
# Сравниваем с head версией
|
||
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}')
|