338 lines
15 KiB
Python
338 lines
15 KiB
Python
import os
|
|
|
|
import motor.motor_asyncio
|
|
|
|
from datetime import datetime
|
|
from dotenv import load_dotenv
|
|
from typing import Optional
|
|
from pymongo import ReturnDocument
|
|
|
|
from schemas.base_schemas import User, Visited, Card, Comment
|
|
from schemas.api_schemas import BaseResponse
|
|
from logger import logger
|
|
|
|
|
|
class MongoWorker:
|
|
"""Worker class for handling all MongoDB database operations."""
|
|
def __init__(self):
|
|
"""Initializes the MongoDB connection and sets up collection references."""
|
|
load_dotenv()
|
|
self.client = motor.motor_asyncio.AsyncIOMotorClient(
|
|
host=os.getenv('MONGO_HOST'),
|
|
port=int(os.getenv('MONGO_PORT', 27017)),
|
|
username=os.getenv('MONGO_USER'),
|
|
password=os.getenv('MONGO_PASS'),
|
|
serverSelectionTimeoutMS=5000,
|
|
connectTimeoutMS=5000,
|
|
maxPoolSize=50,
|
|
minPoolSize=5,
|
|
maxIdleTimeMS=60000,
|
|
waitQueueTimeoutMS=5000
|
|
)
|
|
logger.info("MongoDB connection established.")
|
|
self.db = self.client["data"]
|
|
self.users_data = self.db["users"]
|
|
self.visited_data = self.db["visited"]
|
|
self.counters = self.db["counters"]
|
|
self.game_data = self.db["cards"]
|
|
self.comments_data = self.db["comments"]
|
|
|
|
async def create_indexes(self) -> None:
|
|
"""Creates indexes on application startup."""
|
|
await self.users_data.create_index("user_id", unique=True)
|
|
await self.game_data.create_index("card_id", unique=True)
|
|
await self.game_data.create_index("active_status")
|
|
await self.visited_data.create_index("user_id", unique=True)
|
|
await self.comments_data.create_index("comment_id", unique=True)
|
|
logger.info("MongoDB indexes created.")
|
|
|
|
|
|
async def check_user(self, user_id: int) -> bool:
|
|
"""Checks if a user exists in the database by their user_id."""
|
|
document = await self.users_data.find_one({"user_id": user_id}, {"_id": 1})
|
|
return document is not None
|
|
|
|
async def add_user(
|
|
self, user_id: int, username: str, first_name: str, last_name: str, photo_url: str) -> User:
|
|
"""Creates a new user record in the database."""
|
|
new_user = User(
|
|
user_id=user_id,
|
|
username=username,
|
|
first_name=first_name,
|
|
last_name=last_name,
|
|
photo_url=photo_url,
|
|
registration_date=datetime.now().isoformat(),
|
|
)
|
|
await self.users_data.insert_one(new_user.model_dump())
|
|
logger.debug("User added: user_id={}, username={}", user_id, username)
|
|
return new_user
|
|
|
|
async def get_user(self, user_id: int) -> User:
|
|
"""Retrieves a user's details from the database."""
|
|
document = await self.users_data.find_one({"user_id": user_id})
|
|
return User.model_validate(document)
|
|
|
|
|
|
async def get_and_update_counter(self, counter_name: str) -> int:
|
|
"""Atomically increments the counter and returns the new value."""
|
|
counter = await self.counters.find_one_and_update(
|
|
{"counter_name": counter_name},
|
|
{"$inc": {"counter": 1}},
|
|
upsert=True,
|
|
return_document=ReturnDocument.AFTER,
|
|
)
|
|
return counter["counter"]
|
|
|
|
|
|
async def get_visited_cards(self, user_id: int) -> BaseResponse:
|
|
"""Retrieves the set of card IDs that a user has already visited."""
|
|
document = await self.visited_data.find_one({"user_id": user_id})
|
|
if not document:
|
|
if await self.check_user(user_id):
|
|
return BaseResponse(result=Visited(user_id=user_id, cards_visited=set()))
|
|
return BaseResponse(result="User doesn't exist", error=True)
|
|
return BaseResponse(result=Visited.model_validate(document))
|
|
|
|
async def update_visited_cards(self, user_id: int, visited_card_id: int) -> Visited:
|
|
"""Adds a specific card ID to the user's set of visited cards."""
|
|
updated = await self.visited_data.find_one_and_update(
|
|
{"user_id": user_id},
|
|
{"$addToSet": {"cards_visited": visited_card_id}},
|
|
upsert=True,
|
|
return_document=ReturnDocument.AFTER,
|
|
)
|
|
return Visited.model_validate(updated)
|
|
|
|
async def try_mark_visited(self, user_id: int, card_id: int) -> bool:
|
|
"""
|
|
Atomically marks a card as visited for the user.
|
|
|
|
Returns True if the card was newly marked (was not visited before).
|
|
Returns False if the card was already in the visited set.
|
|
|
|
Uses a conditional update filter (cards_visited: {$ne: card_id}) so that
|
|
only one concurrent request can "win" the mark — eliminating the TOCTOU
|
|
race condition between checking and writing.
|
|
"""
|
|
result = await self.visited_data.update_one(
|
|
{"user_id": user_id, "cards_visited": {"$ne": card_id}},
|
|
{"$addToSet": {"cards_visited": card_id}},
|
|
)
|
|
if result.modified_count == 1:
|
|
return True
|
|
|
|
# No document matched: either the visited doc doesn't exist yet,
|
|
# or the card is already in the set.
|
|
doc = await self.visited_data.find_one({"user_id": user_id}, {"cards_visited": 1})
|
|
if doc is None:
|
|
# First vote ever for this user — create the visited document.
|
|
await self.visited_data.update_one(
|
|
{"user_id": user_id},
|
|
{"$addToSet": {"cards_visited": card_id}},
|
|
upsert=True,
|
|
)
|
|
return True
|
|
|
|
# Card is already present in the visited set.
|
|
return False
|
|
|
|
|
|
async def get_card(self, card_id: int) -> Optional[Card]:
|
|
"""Retrieves a card's details from the database by its card_id."""
|
|
document = await self.game_data.find_one({"card_id": card_id})
|
|
if document:
|
|
return Card.model_validate(document)
|
|
return None
|
|
|
|
async def get_random_cards(self, amount: int,active_status: bool,exclude_ids: Optional[set[int]] = None,) -> Optional[list[Card]]:
|
|
"""Returns random cards, excluding already visited ones (in a single query)."""
|
|
match_filter: dict = {"active_status": active_status}
|
|
if exclude_ids:
|
|
match_filter["card_id"] = {"$nin": list(exclude_ids)}
|
|
|
|
pipeline = [
|
|
{"$match": match_filter},
|
|
{"$sample": {"size": amount}},
|
|
]
|
|
raw_items = await self.game_data.aggregate(pipeline).to_list(length=amount)
|
|
|
|
if raw_items:
|
|
return [Card.model_validate(item) for item in raw_items]
|
|
return None
|
|
|
|
def filter_cards(self, random_cards: list[Card], cards_visited: set) -> tuple[list[Card], list[int]]:
|
|
"""Filters a list of random cards to exclude those already visited by the user."""
|
|
filtered_cards = [card for card in random_cards if card.card_id not in cards_visited]
|
|
filtered_cards_id = [card.card_id for card in filtered_cards]
|
|
return filtered_cards, filtered_cards_id
|
|
|
|
async def add_card_by_api(self, choice_A: str, choice_B: str, author_id: int) -> Card:
|
|
"""Creates a new card in the database with data received from the API."""
|
|
new_card = Card(
|
|
card_id=await self.get_and_update_counter(counter_name="card"),
|
|
choice_A=choice_A,
|
|
choice_B=choice_B,
|
|
author_id=author_id,
|
|
creation_date=datetime.now().isoformat(),
|
|
)
|
|
await self.game_data.insert_one(new_card.model_dump())
|
|
logger.debug("Card created by API: card_id={}, author_id={}", new_card.card_id, author_id)
|
|
return new_card
|
|
|
|
async def add_card_by_base_model(self, new_card: Card) -> Optional[Card]:
|
|
"""Inserts a Card model directly into the database."""
|
|
new_card.card_id = await self.get_and_update_counter(counter_name="card")
|
|
try:
|
|
await self.game_data.insert_one(new_card.model_dump())
|
|
return new_card
|
|
except Exception as exc:
|
|
logger.error("Failed to insert card: {}", exc)
|
|
raise
|
|
|
|
async def accept_card(self, card_id: int) -> BaseResponse:
|
|
"""Accepts a card: sets active_status=True and moderation_date=now."""
|
|
result = await self.game_data.find_one_and_update(
|
|
{"card_id": card_id},
|
|
{"$set": {
|
|
"active_status": True,
|
|
"moderation_date": datetime.now().isoformat(),
|
|
}},
|
|
return_document=ReturnDocument.AFTER,
|
|
)
|
|
if not result:
|
|
logger.debug("Attempted to accept non-existent card: card_id={}", card_id)
|
|
return BaseResponse(result="Card doesn't exist", error=True)
|
|
logger.debug("Card accepted: card_id={}", card_id)
|
|
return BaseResponse(result=Card.model_validate(result))
|
|
|
|
async def reject_card(self, card_id: int) -> BaseResponse:
|
|
"""Rejects a card: deletes it from the database."""
|
|
result = await self.game_data.delete_one({"card_id": card_id})
|
|
if result.deleted_count == 0:
|
|
logger.debug("Attempted to reject non-existent card: card_id={}", card_id)
|
|
return BaseResponse(result="Card doesn't exist", error=True)
|
|
logger.debug("Card rejected and deleted: card_id={}", card_id)
|
|
return BaseResponse(result=f"Card {card_id} rejected and deleted")
|
|
|
|
async def select_choice(self, card_id: int, choice: str) -> BaseResponse:
|
|
"""Increments the vote count for the selected choice (A or B) and total votes on a card."""
|
|
if choice == "A":
|
|
count_field = "count_choice_A"
|
|
elif choice == "B":
|
|
count_field = "count_choice_B"
|
|
else:
|
|
return BaseResponse(result="Wrong choice", error=True)
|
|
|
|
result = await self.game_data.find_one_and_update(
|
|
{"card_id": card_id},
|
|
{"$inc": {"count_total": 1, count_field: 1}},
|
|
)
|
|
if not result:
|
|
return BaseResponse(result="Card doesn't exist", error=True)
|
|
return BaseResponse(result=True, error=False)
|
|
|
|
|
|
async def like_card(self, card_id: int, user_id: int) -> BaseResponse:
|
|
"""Atomically adds a like to a card and records the user's like action."""
|
|
if not await self.check_user(user_id):
|
|
return BaseResponse(result="User doesn't exist", error=True)
|
|
|
|
# Atomically add card_id to liked_card_ids ONLY IF it is not already
|
|
# present in liked_card_ids OR disliked_card_ids.
|
|
# Using a conditional filter makes this a single, race-condition-free
|
|
# test-and-set: if modified_count == 0, another request already won.
|
|
user_update = await self.users_data.find_one_and_update(
|
|
{
|
|
"user_id": user_id,
|
|
"liked_card_ids": {"$ne": card_id},
|
|
"disliked_card_ids": {"$ne": card_id},
|
|
},
|
|
{"$addToSet": {"liked_card_ids": card_id}},
|
|
projection={"_id": 1},
|
|
)
|
|
if not user_update:
|
|
return BaseResponse(result="Card already liked or disliked", error=True)
|
|
|
|
updated_card = await self.game_data.find_one_and_update(
|
|
{"card_id": card_id},
|
|
{"$inc": {"count_likes": 1}},
|
|
)
|
|
if not updated_card:
|
|
# Card doesn't exist — roll back the user update (best effort).
|
|
await self.users_data.update_one(
|
|
{"user_id": user_id},
|
|
{"$pull": {"liked_card_ids": card_id}},
|
|
)
|
|
return BaseResponse(result="Card doesn't exist", error=True)
|
|
|
|
logger.debug("Card liked: card_id={}, user_id={}", card_id, user_id)
|
|
return BaseResponse(result=True, error=False)
|
|
|
|
async def dislike_card(self, card_id: int, user_id: int) -> BaseResponse:
|
|
"""Atomically adds a dislike to a card and records the user's dislike action."""
|
|
if not await self.check_user(user_id):
|
|
return BaseResponse(result="User doesn't exist", error=True)
|
|
|
|
# Same atomic test-and-set pattern as like_card.
|
|
user_update = await self.users_data.find_one_and_update(
|
|
{
|
|
"user_id": user_id,
|
|
"liked_card_ids": {"$ne": card_id},
|
|
"disliked_card_ids": {"$ne": card_id},
|
|
},
|
|
{"$addToSet": {"disliked_card_ids": card_id}},
|
|
projection={"_id": 1},
|
|
)
|
|
if not user_update:
|
|
return BaseResponse(result="Card already liked or disliked", error=True)
|
|
|
|
updated_card = await self.game_data.find_one_and_update(
|
|
{"card_id": card_id},
|
|
{"$inc": {"count_dislikes": 1}},
|
|
)
|
|
if not updated_card:
|
|
# Card doesn't exist — roll back the user update (best effort).
|
|
await self.users_data.update_one(
|
|
{"user_id": user_id},
|
|
{"$pull": {"disliked_card_ids": card_id}},
|
|
)
|
|
return BaseResponse(result="Card doesn't exist", error=True)
|
|
|
|
logger.debug("Card disliked: card_id={}, user_id={}", card_id, user_id)
|
|
return BaseResponse(result=True, error=False)
|
|
|
|
|
|
async def add_comment(self, user_id: int, card_id: int, comment_text: str) -> BaseResponse:
|
|
"""Adds a new comment to a card and links it to the user."""
|
|
if not await self.check_user(user_id):
|
|
return BaseResponse(result="User doesn't exist", error=True)
|
|
if not await self.get_card(card_id):
|
|
return BaseResponse(result="Card doesn't exist", error=True)
|
|
|
|
new_comment = Comment(
|
|
comment_id=await self.get_and_update_counter(counter_name="comment"),
|
|
author_id=user_id,
|
|
card_id=card_id,
|
|
comment_text=comment_text,
|
|
creation_date=datetime.now().isoformat(),
|
|
)
|
|
await self.comments_data.insert_one(new_comment.model_dump())
|
|
|
|
updated_user = await self.users_data.find_one_and_update(
|
|
{"user_id": user_id},
|
|
{"$addToSet": {"comments_ids": new_comment.comment_id}},
|
|
return_document=ReturnDocument.AFTER,
|
|
)
|
|
if not updated_user:
|
|
return BaseResponse(result="Difficulty adding comment_id to user", error=True)
|
|
|
|
logger.debug("Comment added: comment_id={}, card_id={}, author_id={}", new_comment.comment_id, card_id, user_id)
|
|
return BaseResponse(result=new_comment)
|
|
|
|
async def get_comments(self, card_id: int) -> BaseResponse:
|
|
"""Retrieves all comments associated with a specific card_id."""
|
|
if not await self.get_card(card_id):
|
|
return BaseResponse(result="Card doesn't exist", error=True)
|
|
comments = await self.comments_data.find({"card_id": card_id}).sort("creation_date", -1).to_list(length=None)
|
|
comments = [Comment.model_validate(comment) for comment in comments]
|
|
return BaseResponse(result=comments) |