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(), )