Little fixes, add exception handler

birthday_pool_bot/repositories/repositories/subscriptions.py

@@ -54,12 +54,16 @@ class SubscriptionsRepository(BaseRepository[Subscription, SubscriptionFilter]):
     async def get_to_users_subscriptions(
             self,
             to_users_ids: Iterable[uuid.UUID],
+            have_pool: bool | None = None,
+            pagination: BasePagination | None = None,
             with_from_user: bool = False,
             with_to_user: bool = False,
             with_pool: bool = False,
             with_pool_owner: bool = False,
-    ) -> list[Subscription]:
+    ) -> AsyncGenerator[Subscription, None]:
         filter_ = SubscriptionFilter(to_user_id=set(to_users_ids))
+        if have_pool is not None:
+            filter_.pool_id__null = not have_pool
         options = []
         if with_from_user:
             options.append(joinedload(DBSubscription.from_user))
@@ -72,10 +76,12 @@ class SubscriptionsRepository(BaseRepository[Subscription, SubscriptionFilter]):
 
         subscriptions_generator = self.get_items(
             filter_=filter_,
+            pagination=pagination,
             options=options,
         )
 
-        return [subscription async for subscription in subscriptions_generator]
+        async for subscription in subscriptions_generator:
+            yield subscription
 
     async def get_subscription(
             self,
@@ -134,3 +140,20 @@ class SubscriptionsRepository(BaseRepository[Subscription, SubscriptionFilter]):
             to_user_id={to_user_id},
         )
         await self.delete_items(filter_=filter_)
+
+    async def set_pool_for_subscription(
+            self,
+            from_user_id: uuid.UUID,
+            to_user_id: uuid.UUID,
+            pool_id: uuid.UUID,
+    ) -> Subscription | None:
+        filter_ = SubscriptionFilter(
+            from_user_id={from_user_id},
+            to_user_id={to_user_id},
+        )
+        subscriptions_generator = self.update_items(
+            filter_=filter_,
+            pool_id=pool_id,
+        )
+        async for subscription in subscriptions_generator:
+            return subscription

birthday_pool_bot/repositories/repositories/users.py

@@ -1,8 +1,8 @@
 import uuid
 from datetime import date
-from typing import Container, Sequence, TypeVar
+from typing import AsyncGenerator, Sequence, TypeVar
 
-from pydantic_filters import BasePagination, OffsetPagination
+from pydantic_filters import OffsetPagination
 from sqlalchemy import Column, func, inspect, or_, select, update
 from sqlmodel import SQLModel
 
@@ -38,22 +38,12 @@ class UsersRepository(BaseRepository[User, UserFilter]):
         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]:
+    ) -> AsyncGenerator[User, None]:
         filters = []
         if user_id is not None:
             filters.append(DBUser.id == user_id)
@@ -62,24 +52,24 @@ class UsersRepository(BaseRepository[User, UserFilter]):
         if phone is not None:
             filters.append(DBUser.phone == phone)
         if not filters:
-            return []
+            return
 
-        statement = select(DBUser).where(or_(filters))
+        statement = select(DBUser).where(or_(*filters))
         result = await self.session.exec(statement)
-
-        return [db_user.to_item() for db_user in result.all()]
+        for db_user in result.all():
+            yield db_user.to_item()
 
     async def get_users_by_birthdays(
             self,
             birthday: date | None = None,
-    ) -> list[User]:
+    ) -> AsyncGenerator[User, None]:
         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()]
+        for (db_user,) in result.all():
+            yield db_user.to_item()
 
     async def create_user(self, user: User) -> User:
         return await self.create_item(item=user)