Birthday-Pool-Bot/birthday_pool_bot/repositories/tables.py

import uuid
from datetime import date
from typing import Any, Self

import sqlalchemy as sa
from pydantic import BaseModel
from sqlmodel import SQLModel, Field, Relationship

from birthday_pool_bot.dto import (
    PaymentData as DTOPaymentData,
    Pool as DTOPool,
    Subscription as DTOSubscription,
    User as DTOUser,
)


class BaseSQLModel(SQLModel):
    @classmethod
    def from_item(cls, item: BaseModel) -> Self:
        raise NotImplementedError

    def to_item(self) -> BaseModel:
        raise NotImplementedError

    def to_values(self) -> dict[str, Any]:
        return {column.name: getattr(self, column.name) for column in self.__table__.columns}


class User(BaseSQLModel, table=True):
    __tablename__ = "users"

    __table_args__ = (
        sa.Index("ix__users__phone", "phone"),
        sa.Index("ix__users__telegram_id", "telegram_id"),
        sa.UniqueConstraint("phone", name="uq__users__phone"),
        sa.UniqueConstraint("telegram_id", name="uq__users__telegram_id"),
    )

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        nullable=False,
        sa_column_kwargs={"unique": True},
    )
    name: str | None = Field(nullable=True)
    birthday: date | None = Field(nullable=True)
    phone: str | None = Field(default=None, nullable=True)
    telegram_id: int | None = Field(
        default=None,
        sa_column=sa.Column(sa.BigInteger(), nullable=True),
    )
    gift_payment_data: dict | None = Field(
        sa_column=sa.Column(sa.JSON, nullable=True),
        default_factory=dict,
    )

    @classmethod
    def from_item(cls, item: DTOUser) -> Self:
        return cls(
            id=item.id,
            name=item.name,
            birthday=item.birthday,
            phone=item.phone,
            telegram_id=item.telegram_id,
            gift_payment_data=(
                item.gift_payment_data.model_dump_json()
                if item.gift_payment_data is not None else
                None
            ),
        )

    def to_item(self) -> DTOUser:
        return DTOUser(
            id=self.id,
            name=self.name,
            birthday=self.birthday,
            phone=self.phone,
            telegram_id=self.telegram_id,
            gift_payment_data=(
                DTOPaymentData.model_validate_json(self.gift_payment_data)
                if self.gift_payment_data is not None else
                None
            ),
        )


class Pool(BaseSQLModel, table=True):
    __tablename__ = "pools"

    __table_args__ = (
        sa.Index("ix__pools__owner_id", "owner_id"),
        sa.Index("ix__pools__birthday_user_id", "birthday_user_id"),
        sa.CheckConstraint(
            "owner_id <> birthday_user_id",
            name="ck__pools__owner_not_birthday_user",
        ),
    )

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        nullable=False,
        sa_column_kwargs={"unique": True},
    )
    owner_id: uuid.UUID = Field(
        foreign_key="users.id", nullable=False,
    )
    birthday_user_id: uuid.UUID = Field(
        foreign_key="users.id", nullable=False,
    )
    description: str | None = Field(nullable=True) 
    payment_data: dict = Field(
        sa_column=sa.Column(sa.JSON, nullable=False),
        default_factory=dict,
    )

    owner: User = Relationship(
        sa_relationship_kwargs={
            "primaryjoin": "User.id == Pool.owner_id",
            "lazy": None,
        },
    )
    birthday_user: User = Relationship(
        sa_relationship_kwargs={
            "primaryjoin": "User.id == Pool.birthday_user_id",
            "lazy": None,
        },
    )

    @classmethod
    def from_item(cls, item: DTOPool) -> Self:
        return cls(
            id=item.id,
            owner_id=item.owner_id,
            birthday_user_id=item.birthday_user_id,
            description=item.description,
            payment_data=item.payment_data.model_dump_json(),
            owner=None if item.owner is None else DTOUser.from_item(item.owner),
            birthday_user=(
                None
                if item.birthday_user is None else
                DTOUser.from_item(item.birthday_user)
            ),
        )

    def to_item(self) -> DTOPool:
        return DTOPool(
            id=self.id,
            owner_id=self.owner_id,
            birthday_user_id=self.birthday_user_id,
            description=self.description,
            payment_data=DTOPaymentData.model_validate_json(self.payment_data),
            owner=None if self.owner is None else self.owner.to_item(),
            birthday_user=None if self.birthday_user is None else self.birthday_user.to_item(),
        )


class Subscription(BaseSQLModel, table=True):
    __tablename__ = "subscriptions"

    __table_args__ = (
        sa.Index("ix__subscriptions__from_user_id", "from_user_id"),
        sa.Index("ix__subscriptions__to_user_id", "to_user_id"),
        sa.Index("ix__subscriptions__name", "name"),
        sa.Index("ix__subscriptions__pool_id", "pool_id"),
        sa.Index(
            "ix__subscriptions__from_user_id__to_user_id",
            "from_user_id",
            "to_user_id",
        ),
        sa.CheckConstraint(
            "from_user_id <> to_user_id",
            name="ck__subscriptions__from_user_not_to_user",
        ),
        sa.UniqueConstraint(
            "from_user_id",
            "to_user_id",
            name="uq__subscriptions__from_to_user",
        ),
    )

    from_user_id: uuid.UUID = Field(
        foreign_key="users.id", primary_key=True, nullable=False,
    )
    to_user_id: uuid.UUID = Field(
        foreign_key="users.id", primary_key=True, nullable=False,
    )
    name: str = Field(nullable=False)
    pool_id: uuid.UUID = Field(
        foreign_key="pools.id",
        nullable=True,
        ondelete="SET NULL",
    )

    from_user: "User" = Relationship(
        sa_relationship_kwargs={
            "primaryjoin": "User.id == Subscription.from_user_id",
            "lazy": None,
        },
    )
    to_user: "User" = Relationship(
        sa_relationship_kwargs={
            "primaryjoin": "User.id == Subscription.to_user_id",
            "lazy": None,
        },
    )
    pool: "Pool" = Relationship(
        sa_relationship_kwargs={
            "primaryjoin": "Pool.id == Subscription.pool_id",
            "lazy": None,
        },
    )

    @classmethod
    def from_item(cls, item: DTOSubscription) -> Self:
        return cls(
            from_user_id=item.from_user_id,
            to_user_id=item.to_user_id,
            name=item.name,
            pool_id=item.pool_id,
            from_user=None if item.from_user is None else User.from_item(item.from_user),
            to_user=None if item.to_user is None else User.from_item(item.to_user),
            pool=None if item.pool is None else Pool.from_item(item.pool),
        )

    def to_item(self) -> DTOSubscription:
        return DTOSubscription(
            from_user_id=self.from_user_id,
            to_user_id=self.to_user_id,
            name=self.name,
            pool_id=self.pool_id,
            from_user=None if self.from_user is None else self.from_user.to_item(),
            to_user=None if self.to_user is None else self.to_user.to_item(),
            pool=None if self.pool is None else self.pool.to_item(),
        )