src/database/init.py - master - Ulyanoff/mlops1

src/database/init.py
76 строк3 KB
Artem Ivochkin
первый рабочий асинхронный тест с замоканной БД
21 ноя 2025, 17:09
21 ноя 2025, 17:09e7b8312
# -*- coding: utf-8 -*-
import logging
import sys

from _socket import gaierror
from asyncpg import InternalServerError, InvalidCatalogNameError, InvalidPasswordError, TooManyConnectionsError
from asyncpg.exceptions import ConnectionDoesNotExistError
from sqlalchemy.exc import ProgrammingError

from config.db_mdl import db_cfg
from database.core import db_url, async_engine
from database.utils import get_table_names, async2sync

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

async def db_init(log_int):
    """Пытаемся связаться с базой и логировать параметры успешной попытки"""
    try:
        table_names = await async2sync(async_engine, get_table_names)  # проверим наличие таблиц
    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}\n{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 ConnectionDoesNotExistError as exc:
        sys.tracebacklimit = 0
        boot.error(f'{exc}: {exc.__cause__}')
        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'])) * '*' if log_int <= logging.DEBUG else '********',
            driver=async_engine.driver,
            version=async_engine.dialect.name + ' ' + str(async_engine.dialect.server_version_info),
            founded_tables=f'{len(table_names)} '
                           f'{("(" + ", ".join(table_names[:50]) + "...)") if log_int <= logging.DEBUG else ""}',
        )
        boot.info(f'Database initialized successfully: ' + ''.join([f'\n\t{k:<16} \t= {v}' for k, v in _dict.items()]))

    return None