toolbox/db/schemas/toolkit.py

from datetime import datetime
from sqlalchemy import (
    Boolean,
    Column,
    DateTime,
    ForeignKey,
    Integer,
    String,
    Text,
    select,
)
from sqlalchemy.dialects.postgresql import JSONB
from db import Base, CRUD
import utils
from sqlalchemy import Column, Integer, ForeignKey, UniqueConstraint, CheckConstraint


class ToolkitCompatibility(Base):
    __tablename__ = "toolkit_compatibility"

    id = Column(Integer, primary_key=True)
    toolkit_id = Column(
        Integer, ForeignKey("toolkits.id", ondelete="CASCADE"), nullable=False
    )
    compatible_toolkit_id = Column(
        Integer, ForeignKey("toolkits.id", ondelete="CASCADE"), nullable=False
    )
    __table_args__ = (
        # Запрещаем дубли (A-B и A-B)
        UniqueConstraint(
            "toolkit_id", "compatible_toolkit_id", name="uq_toolkit_compatibility"
        ),
        # Запрещаем связь с самим собой
        CheckConstraint(
            "toolkit_id <> compatible_toolkit_id", name="ck_toolkit_not_self"
        ),
    )

    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

    def toDict(self):
        return {
            "id": self.id,
            "toolkit_id": self.toolkit_id,
            "compatible_toolkit_id": self.compatible_toolkit_id,
        }

    @staticmethod
    def getUnique(
        compatibles: list["ToolkitCompatibility"], originalToolkitId: int
    ) -> dict:
        unique = []
        uniqueData = {}
        for c in compatibles:
            if c.toolkit_id not in unique and c.toolkit_id != originalToolkitId:
                unique.append(c.toolkit_id)
                uniqueData[c.id] = c.toolkit_id
            if (
                c.compatible_toolkit_id not in unique
                and c.compatible_toolkit_id != originalToolkitId
            ):
                unique.append(c.compatible_toolkit_id)
                uniqueData[c.id] = c.compatible_toolkit_id
        return uniqueData

    async def save(self):
        await CRUD.create(self)

    @staticmethod
    async def add_compatibility(a_id: int, b_id: int):
        errorMsg = {
            "status": "error",
        }

        if a_id == b_id:
            utils.logger.error("Невозможно добавить совместимость с самим собой")
            errorMsg["errorMessage"] = "Невозможно добавить совместимость с самим собой"
            return errorMsg

        toolkit_id, compatible_id = sorted([a_id, b_id])
        try:

            await ToolkitCompatibility(
                toolkit_id=toolkit_id, compatible_toolkit_id=compatible_id
            ).save()

            return {"status": "ok"}

        except Exception as e:
            utils.logger.error(f"Ошибка добавления совместимости: {str(e)}")
            errorMsg["errorMessage"] = f"Ошибка добавления совместимости: {str(e)}"
            return errorMsg

    @staticmethod
    async def remove_compatibility(a_id: int, b_id: int):
        errorMsg = {
            "status": "error",
        }
        if a_id == b_id:
            utils.logger.error("Невозможно удалить совместимость с самим собой")
            errorMsg["errorMessage"] = "Невозможно удалить совместимость с самим собой"
            return errorMsg

        toolkit_id, compatible_id = sorted([a_id, b_id])

        try:
            compatibility = await CRUD.read(
                select(ToolkitCompatibility).where(
                    ToolkitCompatibility.toolkit_id == toolkit_id,
                    ToolkitCompatibility.compatible_toolkit_id == compatible_id,
                )
            )

            if not compatibility:
                utils.logger.error("Совместимость не найдена")
                errorMsg["errorMessage"] = "Совместимость не найдена"
                return errorMsg

            await CRUD.delete(compatibility)

            return {"status": "ok"}

        except Exception as e:
            utils.logger.error(f"Ошибка удаления совместимости: {str(e)}")
            errorMsg["errorMessage"] = f"Ошибка удаления совместимости: {str(e)}"
            return errorMsg

    @staticmethod
    async def get_compatibility(toolkit_id: int):
        try:

            result = await CRUD.read(
                select(ToolkitCompatibility).where(
                    (ToolkitCompatibility.toolkit_id == toolkit_id)
                    | (ToolkitCompatibility.compatible_toolkit_id == toolkit_id)
                ),
                True,
            )
            if not result:
                return {"status": "ok", "data": {}}

            return {
                "status": "ok",
                "data": ToolkitCompatibility.getUnique(result, toolkit_id),
            }

        except Exception as e:
            utils.logger.error(f"Ошибка получения совместимости: {str(e)}")
            return {"errorMessage": f"Ошибка получения совместимости: {str(e)}"}


class Toolkit(Base):
    __tablename__ = "toolkits"

    id = Column(Integer, primary_key=True, index=True)
    title = Column(String, unique=True, index=True)
    description = Column(Text, nullable=True)
    specifications = Column(JSONB, default={})
    category_id = Column(Integer, ForeignKey("categories.id", ondelete="CASCADE"))
    image = Column(JSONB)
    quantity_min = Column(Integer, nullable=True)
    quantity_min_extra = Column(Integer, nullable=True)
    external_link = Column(String, nullable=True)
    hidden = Column(Boolean, default=False)
    comment_text = Column(Text, nullable=True)
    comment_user_id = Column(
        Integer, ForeignKey("users.id", ondelete="CASCADE"), nullable=True
    )
    comment_at = Column(DateTime, nullable=True)
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
    refilled_at = Column(DateTime, default=datetime.now)
    moved_at = Column(DateTime, default=datetime.now)

    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

    def toDict(self):
        return utils.toDict(self)

    async def save(self):
        return await CRUD.create(self, refresh=True)

    async def edit(self, **kwargs):
        return await CRUD.update(Toolkit, self.id, **kwargs)