src.db.database

  1"""Точки входа в БД"""
  2
  3import logging
  4
  5from asyncpg.exceptions import UniqueViolationError
  6from redis.asyncio import ConnectionPool, Redis
  7from redis.asyncio.client import Pipeline
  8from sqlalchemy.exc import IntegrityError
  9from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 10
 11from core.config import settings
 12from core.err import redis_exceptor
 13from core.exceptions import DatabaseError, UniquenessError
 14
 15logger = logging.getLogger("sqlalchemy")
 16"""Logger: Логгер для SQLAlchemy."""
 17
 18rlogger = logging.getLogger("redis")
 19"""Logger: Логгер для Redis."""
 20
 21async_engine = create_async_engine(
 22    url=settings.DATABASE_URL,
 23    echo=False,  # Логирование всех выполняемых операций
 24    pool_size=10,  # Максимальное количество подключений к базе
 25    max_overflow=20,  # Количество дополнительных подключений при переполнении
 26)
 27"""AsyncEngine: Асинхронный движок SQLAlchemy для работы с базой данных."""
 28
 29async_session_factory = async_sessionmaker(async_engine)
 30"""async_sessionmaker: Фабрика асинхронных сессий для работы с базой данных."""
 31
 32__pool = ConnectionPool.from_url(settings.CASHBASE_URL, decode_responses=True)
 33"""ConnectionPool: Пул соединений для Redis."""
 34
 35redis_engine = Redis.from_pool(__pool)
 36"""Redis: Асинхронный клиент Redis, использующий пул соединений."""
 37
 38
 39async def execute_query(query, echo=True):
 40    """Выполняет SQL-запрос в асинхронной сессии.
 41
 42    Args:
 43        query (sqlalchemy.sql.expression): Запрос для выполнения.
 44        echo (bool): Если True, логирует запрос и параметры. По умолчанию True.
 45
 46    Returns:
 47        Result: Результат выполнения запроса.
 48
 49    Raises:
 50        UniquenessError: Если запись уже существует.
 51        DatabaseError: Если произошла ошибка при подключении к БД.
 52    """
 53    async with async_session_factory() as session:
 54        async with session.begin():
 55            try:
 56                if echo:
 57                    q = query.compile()
 58                    if q.params:
 59                        logger.info(f"{q}\nQuery PARAMS: {q.params}")
 60                    else:
 61                        logger.info(str(q))
 62                res = await session.execute(query)
 63                session.expunge_all()
 64                return res
 65            except (UniqueViolationError, IntegrityError):
 66                await session.rollback()
 67                logger.warning("Эта запись в базе уже существует")
 68                raise UniquenessError
 69            except Exception:
 70                await session.rollback()
 71                logger.exception(
 72                    "Ошибка при подключении к БД",
 73                    exc_info=getattr(query, "__dict__", query),
 74                )
 75                raise DatabaseError
 76
 77
 78@redis_exceptor
 79async def execute_redis_query(pipeline: Pipeline):
 80    """Выполняет запросы Redis в рамках асинхронного пайплайна.
 81
 82    Args:
 83        pipeline (Pipeline): Пайплайн Redis с командами для выполнения.
 84
 85    Returns:
 86        list: Результаты выполнения команд в пайплайне.
 87
 88    Raises:
 89        Exception: Если произошла ошибка при выполнении Redis-запроса.
 90    """
 91    if pipeline.command_stack:
 92        rlogger.info(f"TRYING TO REDIS QUERY :: {pipeline.command_stack}")
 93        async with pipeline as pipe:
 94            result = await pipe.execute()
 95            rlogger.info("REDIS QUERY COMPLETED")
 96            return result
 97    return []
 98
 99
100@redis_exceptor
101async def iter_redis_keys(key_pattern):
102    """Итерация по ключам Redis, соответствующим заданному шаблону.
103
104    Args:
105        key_pattern (str): Шаблон ключей для поиска.
106
107    Returns:
108        AsyncIterator: Итератор по найденным ключам.
109    """
110    return redis_engine.scan_iter(key_pattern)