74 lines
2.6 KiB
Python
74 lines
2.6 KiB
Python
from typing import List, Optional
|
|
from datetime import datetime
|
|
from bson import ObjectId
|
|
from motor.motor_asyncio import AsyncIOMotorClient
|
|
from models.Environment import Environment
|
|
|
|
|
|
class EnvironmentRepo:
|
|
def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
|
|
self.collection = client[db_name]["environments"]
|
|
|
|
async def create_env(self, env: Environment) -> Environment:
|
|
env_dict = env.model_dump(exclude={"id"})
|
|
res = await self.collection.insert_one(env_dict)
|
|
env.id = str(res.inserted_id)
|
|
return env
|
|
|
|
async def get_env(self, env_id: str) -> Optional[Environment]:
|
|
res = await self.collection.find_one({"_id": ObjectId(env_id)})
|
|
if not res:
|
|
return None
|
|
res["id"] = str(res.pop("_id"))
|
|
return Environment(**res)
|
|
|
|
async def get_character_envs(self, character_id: str) -> List[Environment]:
|
|
cursor = self.collection.find({"character_id": character_id})
|
|
envs = []
|
|
async for doc in cursor:
|
|
doc["id"] = str(doc.pop("_id"))
|
|
envs.append(Environment(**doc))
|
|
return envs
|
|
|
|
async def update_env(self, env_id: str, update_data: dict) -> bool:
|
|
update_data["updated_at"] = datetime.utcnow()
|
|
res = await self.collection.update_one(
|
|
{"_id": ObjectId(env_id)},
|
|
{"$set": update_data}
|
|
)
|
|
return res.modified_count > 0
|
|
|
|
async def delete_env(self, env_id: str) -> bool:
|
|
res = await self.collection.delete_one({"_id": ObjectId(env_id)})
|
|
return res.deleted_count > 0
|
|
|
|
async def add_asset(self, env_id: str, asset_id: str) -> bool:
|
|
res = await self.collection.update_one(
|
|
{"_id": ObjectId(env_id)},
|
|
{
|
|
"$addToSet": {"asset_ids": asset_id},
|
|
"$set": {"updated_at": datetime.utcnow()}
|
|
}
|
|
)
|
|
return res.modified_count > 0
|
|
|
|
async def add_assets(self, env_id: str, asset_ids: List[str]) -> bool:
|
|
res = await self.collection.update_one(
|
|
{"_id": ObjectId(env_id)},
|
|
{
|
|
"$addToSet": {"asset_ids": {"$each": asset_ids}},
|
|
"$set": {"updated_at": datetime.utcnow()}
|
|
}
|
|
)
|
|
return res.modified_count > 0
|
|
|
|
async def remove_asset(self, env_id: str, asset_id: str) -> bool:
|
|
res = await self.collection.update_one(
|
|
{"_id": ObjectId(env_id)},
|
|
{
|
|
"$pull": {"asset_ids": asset_id},
|
|
"$set": {"updated_at": datetime.utcnow()}
|
|
}
|
|
)
|
|
return res.modified_count > 0
|