feat: Add album management functionality with new data model, repository, service, API, and generation integration.

repos/albums_repo.py

@@ -0,0 +1,61 @@
+from typing import List, Optional
+import logging
+from bson import ObjectId
+from motor.motor_asyncio import AsyncIOMotorClient
+
+from models.Album import Album
+
+logger = logging.getLogger(__name__)
+
+class AlbumsRepo:
+    def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
+        self.collection = client[db_name]["albums"]
+
+    async def create_album(self, album: Album) -> str:
+        res = await self.collection.insert_one(album.model_dump())
+        return str(res.inserted_id)
+
+    async def get_album(self, album_id: str) -> Optional[Album]:
+        try:
+            res = await self.collection.find_one({"_id": ObjectId(album_id)})
+            if not res:
+                return None
+            
+            res["id"] = str(res.pop("_id"))
+            return Album(**res)
+        except Exception:
+            return None
+
+    async def get_albums(self, limit: int = 10, offset: int = 0) -> List[Album]:
+        res = await self.collection.find().sort("created_at", -1).skip(offset).limit(limit).to_list(None)
+        albums = []
+        for doc in res:
+            doc["id"] = str(doc.pop("_id"))
+            albums.append(Album(**doc))
+        return albums
+
+    async def update_album(self, album_id: str, album: Album) -> bool:
+        if not album.id:
+            album.id = album_id
+            
+        model_dump = album.model_dump()
+        res = await self.collection.update_one({"_id": ObjectId(album_id)}, {"$set": model_dump})
+        return res.modified_count > 0
+
+    async def delete_album(self, album_id: str) -> bool:
+        res = await self.collection.delete_one({"_id": ObjectId(album_id)})
+        return res.deleted_count > 0
+
+    async def add_generation(self, album_id: str, generation_id: str, cover_asset_id: Optional[str] = None) -> bool:
+        res = await self.collection.update_one(
+            {"_id": ObjectId(album_id)},
+            {"$addToSet": {"generation_ids": generation_id}, "$set": {"cover_asset_id": cover_asset_id}}
+        )
+        return res.modified_count > 0
+
+    async def remove_generation(self, album_id: str, generation_id: str) -> bool:
+        res = await self.collection.update_one(
+            {"_id": ObjectId(album_id)},
+            {"$pull": {"generation_ids": generation_id}}
+        )
+        return res.modified_count > 0

repos/dao.py

@@ -4,6 +4,7 @@ from repos.assets_repo import AssetsRepo
 from repos.char_repo import CharacterRepo
 from repos.generation_repo import GenerationRepo
 from repos.user_repo import UsersRepo
+from repos.albums_repo import AlbumsRepo
 
 
 from typing import Optional
@@ -14,3 +15,4 @@ class DAO:
         self.chars = CharacterRepo(client, db_name)
         self.assets = AssetsRepo(client, s3_adapter, db_name)
         self.generations = GenerationRepo(client, db_name)
+        self.albums = AlbumsRepo(client, db_name)

repos/generation_repo.py

@@ -40,7 +40,7 @@ class GenerationRepo:
             generations.append(Generation(**generation))
         return generations
 
-    async def count_generations(self, character_id: Optional[str] = None, status: Optional[GenerationStatus] = None) -> int:
+    async def count_generations(self, character_id: Optional[str] = None, status: Optional[GenerationStatus] = None, album_id: Optional[str] = None) -> int:
         args = {}
         if character_id is not None:
             args["linked_character_id"] = character_id
@@ -48,5 +48,21 @@ class GenerationRepo:
             args["status"] = status
         return await self.collection.count_documents(args)
 
+    async def get_generations_by_ids(self, generation_ids: List[str]) -> List[Generation]:
+        object_ids = [ObjectId(gen_id) for gen_id in generation_ids if ObjectId.is_valid(gen_id)]
+        res = await self.collection.find({"_id": {"$in": object_ids}}).to_list(None)
+        generations: List[Generation] = []
+        
+        # Maintain order of generation_ids
+        gen_map = {str(doc["_id"]): doc for doc in res}
+        
+        for gen_id in generation_ids:
+            doc = gen_map.get(gen_id)
+            if doc:
+                doc["id"] = str(doc.pop("_id"))
+                generations.append(Generation(**doc))
+                
+        return generations
+
     async def update_generation(self, generation: Generation, ):
         res = await self.collection.update_one({"_id": ObjectId(generation.id)}, {"$set": generation.model_dump()})