62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
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
|