86 lines
2.4 KiB
Python
86 lines
2.4 KiB
Python
import logging
|
|
import os
|
|
|
|
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
|
|
|
|
|
|
TORTOISE_ORM = {
|
|
'connections': {'default': os.getenv('DB__URL')},
|
|
'apps': {
|
|
'models': {
|
|
'models': ['src.domain', 'aerich.models'],
|
|
'default_connection': 'default',
|
|
},
|
|
},
|
|
'use_tz': True,
|
|
'timezone': 'UTC',
|
|
}
|
|
|
|
|
|
class DatabaseBase:
|
|
def __init__(self, config: DatabaseConfig) -> None:
|
|
self.config = config
|
|
|
|
async def connect(self) -> None:
|
|
if self.config.ECHO:
|
|
logging.getLogger('tortoise.db_client').setLevel(logging.DEBUG)
|
|
logging.getLogger('tortoise').setLevel(logging.DEBUG)
|
|
|
|
await Tortoise.init(
|
|
db_url=str(self.config.URL),
|
|
modules={'models': ['src.domain', 'aerich.models']},
|
|
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')
|
|
|
|
# Проверяем существование таблицы aerich
|
|
result = await conn.execute_query(
|
|
"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'aerich')"
|
|
)
|
|
if not result[1][0]['exists']:
|
|
logging.warning('Migrations not applied. Run: aerich upgrade')
|
|
return
|
|
|
|
# Получаем последнюю примененную версию
|
|
try:
|
|
result = await conn.execute_query('SELECT version, app FROM aerich ORDER BY id DESC LIMIT 1')
|
|
if result[1]:
|
|
current_version = result[1][0]['version']
|
|
logging.info(f'Current migration version: {current_version}')
|
|
except Exception as e:
|
|
logging.debug(f'Could not check migration version: {e}')
|