src.scripts.db_recover

  1"""Восстановление БД из дампа"""
  2
  3import asyncio
  4import logging
  5import logging.config
  6import os
  7import platform
  8import sys
  9from typing import Literal
 10
 11sys.path.insert(1, os.path.dirname(sys.path[0]))
 12
 13
 14from core.config import Base, settings
 15from core.path import PATH
 16from db import models as _  # COMMENT TRIGGER
 17from db.database import async_engine
 18
 19logging.config.fileConfig("log.ini", disable_existing_loggers=False)
 20logging.getLogger("sqlalchemy.engine.Engine").setLevel(logging.INFO)
 21logger = logging.getLogger()
 22
 23
 24DUMPNAME = "dump_12_11_2024-16_20.sql"
 25PLATFORM: Literal["Windows", "Linux"] = platform.system()
 26
 27
 28async def postgresql_recover_tables():
 29    """Восстанавливает таблицы PostgreSQL из дампа.
 30
 31    Эта функция выполняет последовательные команды для завершения всех подключений к базе данных,
 32    удаления существующей базы данных, создания новой базы данных и восстановления данных из дампа.
 33
 34    Raises:
 35        Exception: Если возникла ошибка при выполнении команд.
 36    """
 37    terminate = (
 38        f"psql -U {settings.DB_USER} "
 39        f"-h {settings.DB_HOST} "
 40        "-d nullbase "
 41        f"-c \"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{settings.DB_NAME}';\""
 42    )
 43    drop = (
 44        f"psql -U {settings.DB_USER} "
 45        f"-h {settings.DB_HOST} "
 46        "-d nullbase "
 47        f'-c "DROP DATABASE IF EXISTS {settings.DB_NAME};"'
 48    )
 49
 50    create = (
 51        f"psql -U {settings.DB_USER} "
 52        f"-h {settings.DB_HOST} "
 53        "-d nullbase "
 54        f"""-c \"CREATE DATABASE {settings.DB_NAME}
 55                    WITH
 56                    OWNER = {settings.DB_USER}
 57                    ENCODING = 'UTF8'
 58                    LC_COLLATE = 'en_US.utf8'
 59                    LC_CTYPE = 'en_US.utf8'
 60                    TABLESPACE = pg_default
 61                    CONNECTION LIMIT = -1
 62                    IS_TEMPLATE = False;\"
 63                        """
 64    )
 65
 66    recover = (
 67        f"psql -U {settings.DB_USER} "
 68        f"-h {settings.DB_HOST} "
 69        f"-d {settings.DB_NAME} "
 70        f'-f {os.path.join(PATH, "src", "db", "dumps", DUMPNAME)}'
 71    )
 72    if PLATFORM == "Windows":
 73        logger.error(
 74            "ЭТА ПЛАТФОРМА НЕ ПОДДЕРЖИВАЕТСЯ",
 75        )
 76        return
 77
 78    elif PLATFORM == "Linux":
 79        login = f"export PGPASSWORD={settings.DB_PASS.get_secret_value()}"
 80
 81    for cmd in (terminate, drop, create, recover):
 82        cmd = f"{login} && {cmd}"
 83
 84        logger.info(f"Executing command: {cmd}")
 85
 86        proc = await asyncio.create_subprocess_shell(
 87            cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
 88        )
 89        stdout, stderr = await proc.communicate()
 90
 91        logger.info(f"pg_dump exited with {proc.returncode}")
 92
 93        if proc.returncode != 0:
 94            error_message = stderr.decode("utf-8", errors="ignore")
 95            logger.error(f"Ошибка создания дампа БД: {error_message}")
 96
 97        else:
 98            print(stdout.decode("utf-8", errors="ignore"))
 99
100
101async def create_new_tables():
102    """Создает новые таблицы в базе данных.
103
104    Эта функция создает все таблицы, определенные в метаданных базы данных,
105    если они еще не существуют.
106
107    Raises:
108        Exception: Если возникла ошибка при создании таблиц.
109    """
110    async with async_engine.connect() as conn:
111        async_engine.echo = True
112        await conn.run_sync(Base.metadata.create_all)
113        await conn.commit()
114
115
116if __name__ == "__main__":
117
118    async def start():
119        """Запускает процесс восстановления таблиц и их создания.
120
121        Эта функция собирает задачи для восстановления таблиц и создания новых,
122        обрабатывает возможные исключения и ведет журнал.
123
124        Raises:
125            Exception: Если возникла ошибка при подключении к БД.
126        """
127        bases = await asyncio.gather(
128            *(postgresql_recover_tables(),),
129            return_exceptions=True,
130        )
131
132        for base in bases:
133            try:
134                if isinstance(base, Exception):
135                    raise base
136            except Exception:
137                logger.exception("Возникло исключение при подключении к БД")
138                raise
139
140        await create_new_tables()
141
142    loop = asyncio.get_event_loop()
143    loop.run_until_complete(start())