src/database/init.py - main - Ulyanoff/asyncAPI

src/database/init.py
132 строки6 KB
Artem Ivochkin
Обработка ошибок локальной БД
22 окт 2025, 22:18
22 окт 2025, 22:18e17ad11
# -*- coding: utf-8 -*-
import logging
import sys

from _socket import gaierror
from asyncpg import InternalServerError, InvalidCatalogNameError, InvalidPasswordError, TooManyConnectionsError
from sqlalchemy import URL, insert, select
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker

from config.db_mdl import db_cfg
from database.tables import meta, main_menu, users, posts, default_users, default_main_menu, default_posts
from database.utils import get_table_names, async2sync

boot = logging.getLogger('boot')
boot.info('Database initialization started')

# куда и к кому подключаться
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(async_engine, class_=AsyncSession, expire_on_commit=False)


async def db_init():
    """Пытаемся связаться с базой и залогировать параметры успешной попытки"""
    try:
        table_names = await async2sync(async_engine, get_table_names)  # проверим наличие таблиц
        await upload_demo(drop=False, forced=False)  # зальём демо-данные
    except InternalServerError as exc:
        if exc.sqlstate == 'XX000':
            if 'password authentication failed' in getattr(exc, "message", ""):
                boot.info('Check host, user or password...')
                sys.tracebacklimit = 0  # страусов не пугать, пол бетонный
            elif 'database ' + db_cfg.name + ' does not exist' in getattr(exc, "message", "").replace('"', ''):
                boot.info('Check DB name...')
                sys.tracebacklimit = 0
        exit(2)
    except TimeoutError as exc:
        boot.info(f'Timeout with DB-connecting, may be wrong port? Check db_url={db_url} {exc}')
        sys.tracebacklimit = 0
        exit(2)
    except InvalidCatalogNameError as exc:
        sys.tracebacklimit = 0
        boot.error(f'Invalid DB-name: {exc}')
        exit(2)
    except InvalidPasswordError as exc:
        sys.tracebacklimit = 0
        boot.error(f'Invalid user or password: {exc}')
        exit(2)
    except TooManyConnectionsError as exc:
        sys.tracebacklimit = 0
        boot.error(f'Too many DB-connections: {exc}')
        exit(2)
    except ProgrammingError as exc:
        sys.tracebacklimit = 0
        if hasattr(exc, "orig"): _e = getattr(exc.orig, "__context__", 'Unknown')
        _e = getattr(exc.orig, "__context__", 'Unknown') if hasattr(exc, "orig") else 'Unknown'
        boot.error(f'{_e}: {exc}')
        exit(2)
    except gaierror as exc:
        sys.tracebacklimit = 0
        boot.error(f'Invalid host {db_cfg.host}: {exc}')
        exit(2)
    except Exception as exc:
        boot.exception(f'Unexpected exception with DB-init: {exc}')
        exit(2)
    else:
        # красивое логирование параметров базы
        _dict = db_url._asdict()
        _dict.pop('query', None)
        _dict.update(
            password=len(str(_dict['password'])) * '*',
            driver=async_engine.driver,
            version=async_engine.dialect.name + ' ' + str(async_engine.dialect.server_version_info),
            total_tables_found=f'{len(table_names)} ({", ".join(table_names[:5])}...)',
        )
        boot.info(f'Database initialized successfully: ' + ''.join([f'\n\t{k:<16} \t= {v}' for k, v in _dict.items()]))

    return None


async def upload_demo(drop=False, forced=False):
    """ Заполнение таблиц демо-данными
        forced=True приводит к предварительному уничтожению базы drop_all
    """
    boot.info(f'Upload demo data started, drop={drop}, forced={forced}')

    async with async_engine.begin() as conn:
        if drop:
            await conn.run_sync(meta.drop_all)  # Убить все таблицы в базе
            boot.info(f'All database tables dropped...')
        await conn.run_sync(meta.create_all)  # Создать все таблицы в базе (пустые)

        upload = True if forced else False

        if not forced:
            result = await conn.execute(select(users))
            upload = True if result.fetchone() is None else False
        if upload:
            await conn.execute(insert(users), default_users)

        if not forced:
            result = await conn.execute(select(main_menu))
            upload = True if result.fetchone() is None else False
        if upload:
            await conn.execute(insert(main_menu), default_main_menu)

        if not forced:
            result = await conn.execute(select(posts))
            upload = True if result.fetchone() is None else False
        if upload:
            await conn.execute(insert(posts), default_posts)