Init project

birthday_pool_bot/repositories/repositories/__init__.py

@@ -0,0 +1,13 @@
+from .pools import PoolsRepository
+from .subscriptions import SubscriptionsRepository
+from .users import UsersRepository
+
+
+__all__ = (
+    # pools
+    "PoolsRepository",
+    # subscriptions
+    "SubscriptionsRepository",
+    # users
+    "UsersRepository",
+)

birthday_pool_bot/repositories/repositories/pools.py

@@ -0,0 +1,41 @@
+from pydantic_filters import OffsetPagination
+from sqlalchemy.orm import joinedload
+
+from birthday_pool_bot.dto import Pool, PoolFilter
+from birthday_pool_bot.repositories.base import BaseRepository
+from birthday_pool_bot.repositories.tables import BaseSQLModel, Pool as DBPool
+
+
+class PoolsRepository(BaseRepository[Pool, PoolFilter]):
+    def get_db_table(self) -> type[BaseSQLModel]:
+        return DBPool
+
+    async def get_pools_by_birthday_user_id_count(self, birthday_user_id: uuid.UUID) -> int:
+        filter_ = PoolFilter(birthday_user_id={birthday_user_id})
+        return await self.get_items_count(filter_=filter_)
+
+    async def get_pools_by_birthday_user_id(self, birthday_user_id: uuid.UUID) -> list[Pool]:
+        filter_ = PoolFilter(birthday_user_id={birthday_user_id})
+        return [pool async for pool in self.get_items(filter_=filter_)]
+
+    async def get_pool_by_id(
+            self,
+            pool_id: uuid.UUID,
+            with_owner: bool = False,
+    ) -> Pool | None:
+        filter_ = PoolFilter(id={pool_id})
+        pagination = OffsetPagination(limit=1)
+        options = []
+        if with_owner:
+            options.append(joinedload(DBPool.owner))
+
+        pools_generator = self.get_items(
+            filter_=filter_,
+            pagination=pagination,
+            options=options,
+        )
+        async for pool in pools_generator:
+            return pool
+
+    async def create_pool(self, pool: Pool) -> Pool:
+        return await self.create_item(item=pool)

birthday_pool_bot/repositories/repositories/subscriptions.py

@@ -0,0 +1,135 @@
+import uuid
+
+from pydantic_filters import BasePagination, OffsetPagination
+from sqlalchemy.orm import joinedload
+from sqlalchemy import delete, update, select
+
+from birthday_pool_bot.dto import Subscription, SubscriptionFilter
+from birthday_pool_bot.repositories.base import BaseRepository
+from birthday_pool_bot.repositories.tables import (
+    BaseSQLModel,
+    Pool as DBPool,
+    Subscription as DBSubscription,
+)
+
+
+class SubscriptionsRepository(BaseRepository[Subscription, SubscriptionFilter]):
+    def get_db_table(self) -> type[BaseSQLModel]:
+        return DBSubscription
+
+    async def get_user_subscriptions_count(
+            self,
+            user_id: uuid.UUID,
+    ) -> int:
+        filter_ = SubscriptionFilter(from_user_id={user_id})
+        return await self.get_items_count(filter_=filter_)
+
+    def get_user_subscriptions(
+            self,
+            user_id: uuid.UUID,
+            pagination: BasePagination | None = None,
+            with_from_user: bool = False,
+            with_to_user: bool = False,
+            with_pool: bool = False,
+            with_pool_owner: bool = False,
+    ) -> AsyncGenerator[Subscription, None]:
+        filter_ = SubscriptionFilter(from_user_id={user_id})
+        options = []
+        if with_from_user:
+            options.append(joinedload(DBSubscription.from_user))
+        if with_to_user:
+            options.append(joinedload(DBSubscription.to_user))
+        if with_pool:
+            options.append(joinedload(DBSubscription.pool))
+        if with_pool_owner:
+            options.append(joinedload(DBSubscription.pool).joinedload(DBPool.owner))
+
+        return self.get_items(
+            filter_=filter_,
+            pagination=pagination,
+            options=options,
+        )
+
+    async def get_to_users_subscriptions(
+            self,
+            to_users_ids: Iterable[uuid.UUID],
+            with_from_user: bool = False,
+            with_to_user: bool = False,
+            with_pool: bool = False,
+            with_pool_owner: bool = False,
+    ) -> list[Subscription]:
+        filter_ = SubscriptionFilter(to_user_id=set(to_users_ids))
+        options = []
+        if with_from_user:
+            options.append(joinedload(DBSubscription.from_user))
+        if with_to_user:
+            options.append(joinedload(DBSubscription.to_user))
+        if with_pool:
+            options.append(joinedload(DBSubscription.pool))
+        if with_pool_owner:
+            options.append(joinedload(DBSubscription.pool).joinedload(DBPool.owner))
+
+        subscriptions_generator = self.get_items(
+            filter_=filter_,
+            options=options,
+        )
+
+        return [subscription async for subscription in subscriptions_generator]
+
+    async def get_subscription(
+            self,
+            from_user_id: uuid.UUID,
+            to_user_id: uuid.UUID,
+            with_from_user: bool = False,
+            with_to_user: bool = False,
+            with_pool: bool = False,
+            with_pool_owner: bool = False,
+    ) -> Subscription | None:
+        filter_ = SubscriptionFilter(from_user_id={from_user_id}, to_user_id={to_user_id})
+        pagination = OffsetPagination(limit=1)
+        options = []
+        if with_from_user:
+            options.append(joinedload(DBSubscription.from_user))
+        if with_to_user:
+            options.append(joinedload(DBSubscription.to_user))
+        if with_pool:
+            options.append(joinedload(DBSubscription.pool))
+        if with_pool_owner:
+            options.append(joinedload(DBSubscription.pool).joinedload(DBPool.owner))
+
+        subscriptions_generator = self.get_items(
+            filter_=filter_,
+            pagination=pagination,
+            options=options,
+        )
+        async for subscription in subscriptions_generator:
+            return subscription
+
+    async def create_subscription(self, subscription: Subscription) -> Subscription:
+        return await self.create_item(subscription)
+
+    async def delete_subscription(self, from_user_id: uuid.UUID, to_user_id: uuid.UUID):
+        statement = select(DBPool.id).where(
+            DBPool.owner_id == from_user_id,
+            DBPool.birthday_user_id == to_user_id,
+        )
+        result = await self.session.exec(statement)
+        pool_id = result.one_or_none()
+        if pool_id is not None:
+            pool_id = pool_id[0]
+            statement = delete(DBPool).where(DBPool.id == pool_id)
+            await self.session.exec(statement)
+
+        if pool_id is not None:
+            statement = (
+                update(DBSubscription)
+                .where(DBSubscription.pool_id == pool_id)
+                .values(pool_id=None)
+            )
+            await self.session.exec(statement)
+
+        filter_ = SubscriptionFilter(
+            from_user_id={from_user_id},
+            to_user_id={to_user_id},
+        )
+        await self.delete_items(filter_=filter_)

birthday_pool_bot/repositories/repositories/users.py

@@ -0,0 +1,154 @@
+from datetime import date
+from typing import Container, Sequence, TypeVar
+
+from pydantic_filters import OffsetPagination
+from sqlalchemy import Column, func, inspect, select
+from sqlmodel import SQLModel
+
+from birthday_pool_bot.dto import User, UserFilter
+from birthday_pool_bot.repositories.base import BaseRepository
+from birthday_pool_bot.repositories.tables import (
+    BaseSQLModel,
+    Pool as DBPool,
+    Subscription as DBSubscription,
+    User as DBUser,
+)
+
+
+class UsersRepository(BaseRepository[User, UserFilter]):
+    def get_db_table(self) -> type[BaseSQLModel]:
+        return DBUser
+
+    async def get_user_by_id(self, user_id: uuid.UUID) -> User | None:
+        filter_ = UserFilter(id={user_id})
+        pagination = OffsetPagination(limit=1)
+        async for user in self.get_items(filter_=filter_, pagination=pagination):
+            return user
+
+    async def get_user_by_telegram_id(self, telegram_id: int) -> User | None:
+        filter_ = UserFilter(telegram_id={telegram_id})
+        pagination = OffsetPagination(limit=1)
+        async for user in self.get_items(filter_=filter_, pagination=pagination):
+            return user
+
+    async def get_user_by_phone(self, phone: str) -> User | None:
+        filter_ = UserFilter(phone={phone})
+        pagination = OffsetPagination(limit=1)
+        async for user in self.get_items(filter_=filter_, pagination=pagination):
+            return user
+
+    async def get_users_by_ids(
+            self,
+            user_ids: Container[uuid.UUID],
+            pagination: BasePagination | None = None,
+    ) -> list[User]:
+        filter_ = UserFilter(id=set(user_ids))
+        users_generator = self.get_items(filter_=filter_, pagination=pagination)
+        users = [user async for user in users_generator]
+        return users
+
+    async def get_users_by_primary_keys(
+            self,
+            user_id: uuid.UUID | None = None,
+            telegram_id: int | None = None,
+            phone: str | None = None,            
+    ) -> list[User]:
+        filters = []
+        if user_id is not None:
+            filters.append(DBUser.id == user_id)
+        if telegram_id is not None:
+            filters.append(DBUser.telegram_id == telegram_id)
+        if phone is not None:
+            filters.append(DBUser.phone == phone)
+        if not filters:
+            return []
+
+        statement = select(DBUser).where(or_(filters))
+        result = await self.session.exec(statement)
+
+        return [db_user.to_item() for db_user in result.all()]
+
+    async def get_users_by_birthdays(
+            self,
+            birthday: date | None = None,
+    ) -> list[User]:
+        statement = select(DBUser).where(
+            func.extract("month", DBUser.birthday) == birthday.month,
+            func.extract("day", DBUser.birthday) == birthday.day,
+        )
+        result = await self.session.execute(statement)
+
+        return [db_user.to_item() for (db_user,) in result.all()]
+
+    async def create_user(self, user: User) -> User:
+        return await self.create_item(item=user)
+
+    async def update_user(self, user: User) -> User:
+        db_user = self.get_db_table().from_item(item=user)
+        filter_ = UserFilter(id={user.id})
+        users_generator = self.update_items(filter_=filter_, **db_user.to_values())
+        async for user in users_generator:
+            return user
+
+    async def merge_users(self, *users: User) -> User | None:
+        if not users:
+            return None
+
+        merged_user, *users = users
+        merged_user = merge(model=User, first_item=merged_user, items=users)
+
+        if users:
+            users_ids = {user.id for user in users}
+            for model in (DBPool, DBSubscription):
+                foreign_key_columns = get_model_foreign_keys_columns(
+                    model=model,
+                    table_name="users",
+                    column_name="id",
+                )
+                for foreign_key_column in foreign_key_columns:
+                    statement = (
+                        update(model).values(**{foreign_key_column.key: merged_user.id})
+                        .where(foreign_key_column.in_(users_ids))
+                    )
+                    await self.session.exec(statement)
+
+            filter_ = UserFilter(id={user.id for user in users})
+            await self.delete_items(filter_=filter_)
+
+        merged_user = await self.update_user(user=merged_user)
+        return merged_user
+
+
+ModelType = TypeVar("ModelType")
+
+
+def merge(
+        model: type[ModelType],
+        first_item: ModelType,
+        items: Sequence[ModelType],
+) -> ModelType:
+    data = first_item.model_dump()
+    for item in items:
+        item_data = item.model_dump()
+        for field_name, value in data.items():
+            item_value = item_data[field_name]
+            data[field_name] = value or item_value
+    
+    return model(**data)
+
+
+def get_model_foreign_keys_columns(
+        model: type[SQLModel],
+        table_name: str | None = None,
+        column_name: str | None = None,
+) -> list[Column]:
+    mapper = inspect(model)
+    columns = []
+
+    for column in mapper.columns:
+        for foreign_key in column.foreign_keys:
+            table_name = table_name or foreign_key.column.table.name
+            column_name = column_name or foreign_key.column.name
+            if (foreign_key.column.table.name, foreign_key.column.name) == (table_name, column_name):
+                columns.append(column)
+    return columns