src/core/db_base_map.py - master - planver/planver-main

src/core/db_base_map.py
396 строк12 KB
waldemarX
event plans
21 ноя 2025, 12:39
21 ноя 2025, 12:398c6f829
from contextvars import ContextVar
from typing import Type, TypeVar, Optional, Any
from sqlalchemy import and_, func, select, exists, delete

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from sqlalchemy.exc import IntegrityError


Self = TypeVar("Self", bound="BaseDBMap")

session_context: ContextVar[AsyncSession | None] = ContextVar("session", default=None)


class BaseDBMap:
    """
    Базовый mixin с CRUD операциями для SQLAlchemy моделей
    """

    @classmethod
    def _get_db(cls) -> AsyncSession:
        """Получение сессии из контекста"""
        session = session_context.get()
        if session is None:
            raise RuntimeError("No database session found in context")
        return session

    @classmethod
    async def get_by_id(
        cls: Type[Self],
        id: int,
        load: Optional[list[str]] = None,
    ) -> Optional[Self]:
        """
        Получить объект по ID

        Args:
            id (int): ID объекта
            load: Поля для загрузки связанных моделей

        Returns:
            Optional[Self]: Объект или None если не найден
        """
        db = cls._get_db()
        stmt = select(cls).where(cls.id == id)

        if load:
            load_objects = [selectinload(getattr(cls, option)) for option in load]
            stmt = stmt.options(*load_objects)

        result = await db.execute(stmt)
        return result.scalar_one_or_none()

    @classmethod
    async def get_one_by_kwargs(
        cls: Type[Self],
        load: Optional[list[str]] = None,
        **kwargs,
    ) -> Optional[Self]:
        """
        Получить один объект по ключевым аргументам

        Args:
            load: Поля для загрузки связанных моделей
            **kwargs: Аргументы для фильтрации

        Returns:
            Optional[Self]: Один объект или None
        """
        db = cls._get_db()
        stmt = select(cls).filter_by(**kwargs)

        if load:
            load_objects = [selectinload(getattr(cls, option)) for option in load]
            stmt = stmt.options(*load_objects)

        result = await db.execute(stmt)
        return result.scalar_one_or_none()

    @classmethod
    async def get_by_kwargs(
        cls: Type[Self],
        load: Optional[list[str]] = None,
        **kwargs,
    ) -> list[Self]:
        """
        Получить список объектов по ключевым аргументам

        Args:
            load: Поля для загрузки связанных моделей
            **kwargs: Аргументы для фильтрации

        Returns:
            list[Self]: Список объектов
        """
        db = cls._get_db()
        stmt = select(cls).filter_by(**kwargs)

        if load:
            load_objects = [selectinload(getattr(cls, option)) for option in load]
            stmt = stmt.options(*load_objects)

        result = await db.execute(stmt)
        return list(result.scalars().all())

    @classmethod
    async def get_all(
        cls: Type[Self],
        limit: int = 100,
        offset: int = 0,
        load: Optional[list[str]] = None,
    ) -> list[Self]:
        """
        Получить все объекты с пагинацией

        Args:
            skip (int): Пропустить записей
            limit (int): Лимит записей
            load: Поля для загрузки связанных моделей

        Returns:
            list[Self]: Список объектов
        """
        db = cls._get_db()
        stmt = select(cls).offset(offset).limit(limit)

        if load:
            load_objects = [selectinload(getattr(cls, option)) for option in load]
            stmt = stmt.options(*load_objects)

        result = await db.execute(stmt)
        return list(result.scalars().all())

    @classmethod
    async def create(cls: Type[Self], **kwargs) -> Self:
        """
        Создать новый объект

        Args:
            **kwargs: Атрибуты объекта

        Returns:
            Self: Созданный объект
        """
        db = cls._get_db()
        obj = cls(**kwargs)
        db.add(obj)
        await db.commit()
        await db.flush(obj)
        return obj

    @classmethod
    async def get_sorted(
        cls: Type[Self],
        sort_key: str,
        **kwargs: dict[str, Any],
    ) -> None:
        """
        Получение и сортировка объектов по заданному полю

        Args:
            kwargs: Атрибуты объекта.
            sort_key (str): Поле для сортировки.

        Returns:
            list[Self]: Список объектов

        Raises:
            AttributeError: Если поле sort_key  не существует в модели.
        """
        objects = await cls.get_by_kwargs(**kwargs)

        # Проверить, что sort_key существует
        if not hasattr(cls, sort_key):
            raise AttributeError(f"Field '{sort_key}' does not exist in {cls.__name__}")

        # Сортировать по sort_key
        objects.sort(key=lambda obj: getattr(obj, sort_key))

        return objects

    def update(self, **kwargs) -> None:
        """
        Обновить текущий объект

        Args:
            **kwargs: Атрибуты объекта

        Returns:
            None
        """
        for key, value in kwargs.items():
            setattr(self, key, value)

    @classmethod
    async def update_by_id(cls: Type[Self], object_id: int, **kwargs) -> Optional[Self]:
        """
        Обновить объект по ID

        Args:
            object_id (int): ID объекта
            **kwargs: Поля для обновления

        Returns:
            Optional[Self]: Обновленный объект

        Raises:
            ValueError: Если объект не найден
        """
        db = cls._get_db()
        obj = await cls.get_by_id(object_id)
        if not obj:
            raise ValueError(f"{cls.__name__} with id {object_id} not found")

        for key, value in kwargs.items():
            if hasattr(obj, key):
                setattr(obj, key, value)

        db.add(obj)
        await db.commit()
        await db.flush(obj)
        return obj

    @classmethod
    async def update_or_create(
        cls: Type[Self], lookup_fields: dict, **create_kwargs
    ) -> Self:
        """
        Обновить объект, если существует по lookup_fields, иначе создать новый.

        Args:
            lookup_fields (dict): Поля для поиска существующего объекта.
            **create_kwargs: Поля для создания или обновления объекта.

        Returns:
            Self: Обновленный или созданный объект.

        Raises:
            ValueError: Если найдено несколько объектов по lookup_fields.
        """
        db = cls._get_db()

        # Найти существующие объекты по lookup_fields
        existing_objs = await cls.get_by_kwargs(**lookup_fields)

        if len(existing_objs) > 1:
            raise ValueError(
                f"Multiple {cls.__name__} objects found for {lookup_fields}."
            )

        if existing_objs:
            # Обновить существующий объект
            obj = existing_objs[0]
            for key, value in create_kwargs.items():
                if hasattr(obj, key):
                    setattr(obj, key, value)
            db.add(obj)
        else:
            # Создать новый объект
            obj = cls(**lookup_fields, **create_kwargs)
            db.add(obj)

        await db.commit()
        await db.flush(obj)
        return obj

    @classmethod
    async def delete_by_id(cls: Type[Self], id: int) -> bool:
        """
        Удалить объект по ID

        Args:
            id (int): ID объекта

        Returns:
            bool: True если удален, False если не найден
        """
        db = cls._get_db()
        obj = await cls.get_by_id(id)
        if not obj:
            return False

        await db.delete(obj)
        await db.commit()
        return True

    @classmethod
    async def delete_by_kwargs(cls: Type[Self], **kwargs) -> int:
        """
        Удалить объекты по фильтрам

        Args:
            **kwargs: Фильтры для удаления

        Returns:
            int: Количество удаленных объектов
        """
        db = cls._get_db()
        result = await db.execute(delete(cls).filter_by(**kwargs))
        await db.commit()
        return result.rowcount

    @classmethod
    async def get_or_create(
        cls: Type[Self],
        defaults: Optional[dict[str, Any]] = None,
        **kwargs,
    ) -> tuple[Self, bool]:
        """
        Получить или создать объект

        Args:
            defaults: Значения по умолчанию при создании
            **kwargs: Уникальные поля для поиска/создания

        Returns:
            tuple[Self, bool]: (объект, создан_ли)
        """
        db = cls._get_db()
        defaults = defaults or {}

        try:
            # Пытаемся найти существующий
            existing = await cls.get_one_by_kwargs(**kwargs)
            if existing:
                return existing, False

            # Создаем новый
            create_data = {**kwargs, **defaults}
            new_obj = cls(**create_data)

            db.add(new_obj)
            await db.commit()
            await db.flush(new_obj)

            return new_obj, True

        except IntegrityError:
            # Обрабатываем race condition
            await db.rollback()

            # Пытаемся найти снова после отката
            existing = await cls.get_one_by_kwargs(**kwargs)
            if existing:
                return existing, False

            raise  # Если все еще не нашли - пробрасываем исключение

    @classmethod
    async def count(cls: Type[Self], **kwargs) -> int:
        """
        Посчитать количество объектов по фильтрам

        Args:
            **kwargs: Фильтры

        Returns:
            int: Количество объектов
        """
        db = cls._get_db()

        stmt = select(func.count(cls.id))

        filters = []
        for key, value in kwargs.items():
            attr = getattr(cls, key)
            if isinstance(value, list):
                # Явно указываем тип для массива
                filters.append(attr.in_(value))
            else:
                filters.append(attr == value)

        if filters:
            stmt = stmt.where(and_(*filters))

        result = await db.execute(stmt)
        return result.scalar()

    @classmethod
    async def exists(cls, **kwargs) -> bool:
        """
        Проверить существование хотя бы одной записи по фильтрам
        Легковесная проверка без загрузки полных объектов

        Args:
            **kwargs: Фильтры для WHERE (условия равенства)

        Returns:
            bool: True если есть хотя бы одна запись, иначе False
        """
        db = cls._get_db()
        conditions = [getattr(cls, k) == v for k, v in kwargs.items()]
        stmt = select(exists().where(*conditions))

        result = await db.execute(stmt)
        return bool(result.scalar_one_or_none())