| import re |
| import secrets |
| import string |
| from asyncio import create_task |
| from datetime import datetime, timedelta, timezone |
| from typing import Any, Dict, List, Optional, Tuple |
|
|
| import motor.motor_asyncio |
| from bson import ObjectId |
| from pydantic import ValidationError |
| from pymongo import ASCENDING, DESCENDING |
|
|
| from Backend.config import Telegram |
| from Backend.helper.encrypt import decode_string, encode_string |
| from Backend.helper.modal import Episode, MovieSchema, QualityDetail, QualityPart, Season, TVShowSchema |
| from Backend.helper.settings_manager import SettingsManager |
| from Backend.helper.task_manager import delete_message |
| from Backend.logger import LOGGER |
|
|
|
|
|
|
| def convert_objectid_to_str(document: Dict[str, Any]) -> Dict[str, Any]: |
| for key, value in document.items(): |
| if isinstance(value, ObjectId): |
| document[key] = str(value) |
| elif isinstance(value, list): |
| document[key] = [convert_objectid_to_str(item) if isinstance(item, dict) else item for item in value] |
| elif isinstance(value, dict): |
| document[key] = convert_objectid_to_str(value) |
| return document |
|
|
|
|
| class Database: |
| def __init__(self, db_name: str = "dbFyvio"): |
| self.db_uris = Telegram.DATABASE |
| self.db_name = db_name |
|
|
| if len(self.db_uris) < 2: |
| raise ValueError("At least 2 database URIs are required (1 for tracking + 1 for storage).") |
|
|
| self.clients: Dict[str, motor.motor_asyncio.AsyncIOMotorClient] = {} |
| self.dbs: Dict[str, motor.motor_asyncio.AsyncIOMotorDatabase] = {} |
|
|
| self.current_db_index = 1 |
|
|
| async def connect(self): |
| try: |
| for index, uri in enumerate(self.db_uris): |
| client = motor.motor_asyncio.AsyncIOMotorClient(uri) |
| db_key = "tracking" if index == 0 else f"storage_{index}" |
| self.clients[db_key] = client |
| self.dbs[db_key] = client[self.db_name] |
| db_type = "Tracking" if index == 0 else f"Storage {index}" |
|
|
| masked_uri = re.sub(r"://(.*?):.*?@", r"://\1:*****@", uri) |
| masked_uri = masked_uri.split('?')[0] |
| |
| LOGGER.info(f"{db_type} Database connected successfully: {masked_uri}") |
|
|
| state = await self.dbs["tracking"]["state"].find_one({"_id": "db_index"}) |
| if not state: |
| await self.dbs["tracking"]["state"].insert_one({"_id": "db_index", "current_index": 1}) |
| self.current_db_index = 1 |
| else: |
| self.current_db_index = state["current_index"] |
|
|
| LOGGER.info(f"Active storage DB: storage_{self.current_db_index}") |
|
|
| except Exception as e: |
| LOGGER.error(f"Database connection error: {e}") |
|
|
| async def disconnect(self): |
| for client in self.clients.values(): |
| client.close() |
| LOGGER.info("All database connections closed.") |
|
|
| async def update_current_db_index(self): |
| await self.dbs["tracking"]["state"].update_one( |
| {"_id": "db_index"}, |
| {"$set": {"current_index": self.current_db_index}}, |
| upsert=True |
| ) |
|
|
|
|
| async def get_settings(self) -> dict: |
| try: |
| doc = await self.dbs["tracking"]["settings"].find_one({"_id": "app_settings"}) |
| return doc or {} |
| except Exception as e: |
| LOGGER.error(f"Database.get_settings error: {e}") |
| return {} |
|
|
| async def save_settings(self, settings: dict) -> bool: |
| try: |
| clean = {k: v for k, v in settings.items() if k != "_id"} |
| await self.dbs["tracking"]["settings"].update_one( |
| {"_id": "app_settings"}, |
| {"$set": clean}, |
| upsert=True, |
| ) |
| return True |
| except Exception as e: |
| LOGGER.error(f"Database.save_settings error: {e}") |
| return False |
|
|
|
|
|
|
| async def connect_storage_db(self, uri: str, index: int) -> bool: |
| try: |
| client = motor.motor_asyncio.AsyncIOMotorClient(uri) |
| await client.admin.command("ping") |
|
|
| db_key = "tracking" if index == 0 else f"storage_{index}" |
| self.clients[db_key] = client |
| self.dbs[db_key] = client[self.db_name] |
|
|
| db_type = "Tracking" if index == 0 else f"Storage {index}" |
| masked_uri = re.sub(r"://(.*?):.*?@", r"://\1:*****@", uri).split('?')[0] |
| LOGGER.info(f"{db_type} Database connected successfully: {masked_uri}") |
| return True |
| except Exception as e: |
| LOGGER.error(f"Failed to connect database at index {index}: {e}") |
| return False |
|
|
| async def disconnect_storage_db(self, index: int) -> None: |
| db_key = f"storage_{index}" |
| client = self.clients.pop(db_key, None) |
| self.dbs.pop(db_key, None) |
| if client: |
| client.close() |
| LOGGER.info(f"Disconnected {db_key}.") |
|
|
| def get_database_list(self) -> List[Dict[str, Any]]: |
| result = [] |
| for index, uri in enumerate(self.db_uris): |
| masked = re.sub(r"://(.*?):.*?@", r"://\1:*****@", uri).split('?')[0] |
| db_key = "tracking" if index == 0 else f"storage_{index}" |
| entry = { |
| "index": index, |
| "uri_masked": masked, |
| "locked": index <= 1, |
| "type": "tracking" if index == 0 else "storage", |
| "connected": db_key in self.clients, |
| } |
| if index > 1: |
| entry["full_uri"] = uri |
| result.append(entry) |
| return result |
|
|
| async def reload_extra_databases(self, new_extra_uris: List[str]) -> Dict[str, Any]: |
| old_extra = self.db_uris[2:] |
| new_extra = [u.strip() for u in (new_extra_uris or []) if u and u.strip()] |
|
|
| common_len = min(len(old_extra), len(new_extra)) |
| for i in range(common_len): |
| if old_extra[i] != new_extra[i]: |
| raise ValueError( |
| f"Cannot modify storage_{i + 2} in place — existing media " |
| f"documents reference it by position. Only appending new " |
| f"databases or removing the last one(s) is supported." |
| ) |
|
|
| added = 0 |
| removed = 0 |
|
|
| if len(new_extra) > len(old_extra): |
| for offset, uri in enumerate(new_extra[len(old_extra):]): |
| index = len(old_extra) + 2 + offset |
| ok = await self.connect_storage_db(uri, index) |
| if not ok: |
| raise ValueError( |
| f"Failed to connect new database at storage_{index}. " |
| f"Check the URI and try again — no changes were saved." |
| ) |
| added += 1 |
|
|
| elif len(new_extra) < len(old_extra): |
| for index in range(len(old_extra) + 1, len(new_extra) + 1, -1): |
| await self.disconnect_storage_db(index) |
| removed += 1 |
|
|
| self.db_uris = self.db_uris[:2] + new_extra |
|
|
| message = f"{added} database(s) added, {removed} database(s) removed." |
| LOGGER.info(f"reload_extra_databases: {message}") |
| return {"added": added, "removed": removed, "message": message} |
|
|
| |
| |
| |
| async def get_user(self, user_id: int) -> Optional[dict]: |
| return await self.dbs["tracking"]["users"].find_one({"_id": user_id}) |
|
|
| |
| @staticmethod |
| def is_subscription_active(user: Optional[dict], now: datetime = None) -> bool: |
| if not user or user.get("subscription_status") != "active": |
| return False |
| expiry = user.get("subscription_expiry") |
| if not expiry: |
| return False |
| reference = now or datetime.utcnow() |
| try: |
| if expiry.tzinfo is not None: |
| reference = datetime.now(timezone.utc) |
| except AttributeError: |
| pass |
| return expiry > reference |
|
|
| |
| @staticmethod |
| def content_totals(db_stats: List[dict]) -> Tuple[int, int]: |
| total_movies = sum(stat.get("movie_count", 0) for stat in db_stats) |
| total_tv = sum(stat.get("tv_count", 0) for stat in db_stats) |
| return total_movies, total_tv |
|
|
| async def update_user_interaction(self, user_id: int, first_name: str, username: str): |
| await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| {"$set": {"first_name": first_name, "username": username, "last_interaction": datetime.utcnow()}}, |
| upsert=True |
| ) |
|
|
| async def set_pending_payment(self, user_id: int, plan_duration: int, msg_id: int, price=0, admin_messages: list = None): |
| update_data = { |
| "pending_payment": { |
| "duration": plan_duration, |
| "price": price, |
| "msg_id": msg_id, |
| "date": datetime.utcnow(), |
| } |
| } |
| if admin_messages is not None: |
| update_data["pending_payment"]["admin_messages"] = admin_messages |
| await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| {"$set": update_data}, |
| upsert=True |
| ) |
|
|
| async def approve_payment(self, user_id: int) -> Optional[dict]: |
| user = await self.get_user(user_id) |
| if not user or "pending_payment" not in user: |
| return None |
|
|
| duration = user["pending_payment"]["duration"] |
| |
| |
| current_expiry = user.get("subscription_expiry") |
| now = datetime.utcnow() |
| if current_expiry and current_expiry > now: |
| new_expiry = current_expiry + timedelta(days=duration) |
| else: |
| new_expiry = now + timedelta(days=duration) |
|
|
| await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| { |
| "$set": {"subscription_expiry": new_expiry, "subscription_status": "active"}, |
| "$unset": {"pending_payment": ""} |
| } |
| ) |
| return await self.get_user(user_id) |
|
|
| async def reject_payment(self, user_id: int) -> bool: |
| result = await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| {"$unset": {"pending_payment": ""}} |
| ) |
| return result.modified_count > 0 |
|
|
| async def get_expired_users(self) -> List[dict]: |
| cursor = self.dbs["tracking"]["users"].find({ |
| "subscription_expiry": {"$lt": datetime.utcnow()}, |
| "subscription_status": "active" |
| }) |
| return await cursor.to_list(None) |
|
|
| async def mark_user_expired(self, user_id: int): |
| await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| {"$set": {"subscription_status": "expired"}} |
| ) |
|
|
| async def get_expiring_users(self, hours: int = 24) -> List[dict]: |
| now = datetime.utcnow() |
| target_time = now + timedelta(hours=hours) |
| cursor = self.dbs["tracking"]["users"].find({ |
| "subscription_expiry": {"$gt": now, "$lte": target_time}, |
| "reminder_sent": {"$ne": True}, |
| "subscription_status": "active" |
| }) |
| return await cursor.to_list(None) |
| |
| async def mark_reminder_sent(self, user_id: int): |
| await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| {"$set": {"reminder_sent": True}} |
| ) |
|
|
| |
| |
| |
| async def get_subscription_plans(self) -> List[dict]: |
| cursor = self.dbs["tracking"]["sub_plans"].find().sort("days", ASCENDING) |
| plans = await cursor.to_list(None) |
| return [convert_objectid_to_str(plan) for plan in plans] |
|
|
| async def add_subscription_plan(self, days: int, price: float) -> Optional[str]: |
| result = await self.dbs["tracking"]["sub_plans"].insert_one({ |
| "days": days, |
| "price": price, |
| "created_at": datetime.utcnow() |
| }) |
| return str(result.inserted_id) |
|
|
| async def update_subscription_plan(self, plan_id: str, days: int, price: float) -> bool: |
| try: |
| result = await self.dbs["tracking"]["sub_plans"].update_one( |
| {"_id": ObjectId(plan_id)}, |
| {"$set": {"days": days, "price": price, "updated_at": datetime.utcnow()}} |
| ) |
| return result.modified_count > 0 |
| except Exception: |
| return False |
|
|
| async def delete_subscription_plan(self, plan_id: str) -> bool: |
| try: |
| result = await self.dbs["tracking"]["sub_plans"].delete_one({"_id": ObjectId(plan_id)}) |
| return result.deleted_count > 0 |
| except Exception: |
| return False |
|
|
| async def get_all_subscribers(self) -> List[dict]: |
| cursor = self.dbs["tracking"]["users"].find({ |
| "subscription_status": {"$in": ["active", "expired"]} |
| }).sort("subscription_expiry", DESCENDING) |
| users = await cursor.to_list(None) |
| return [convert_objectid_to_str(u) for u in users] |
|
|
| async def manage_subscriber(self, user_id: int, action: str, days: int = 0) -> bool: |
| now = datetime.utcnow() |
| user = await self.get_user(user_id) |
|
|
| if action in ("extend", "reduce"): |
| current_expiry = user.get("subscription_expiry") if user else None |
|
|
| if action == "extend": |
| base = current_expiry if (current_expiry and current_expiry > now) else now |
| new_expiry = base + timedelta(days=days) |
| else: |
| if current_expiry: |
| new_expiry = current_expiry - timedelta(days=days) |
| if new_expiry < now: |
| new_expiry = now |
| else: |
| new_expiry = now |
|
|
| status = "active" if new_expiry > now else "expired" |
|
|
| await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| { |
| "$set": {"subscription_expiry": new_expiry, "subscription_status": status}, |
| "$setOnInsert": { |
| "first_name": f"User {user_id}", |
| "username": None, |
| "created_at": now, |
| }, |
| }, |
| upsert=True |
| ) |
| return True |
|
|
| elif action == "delete": |
| await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| {"$set": {"subscription_status": "expired", "subscription_expiry": now}} |
| ) |
| return True |
|
|
| return False |
|
|
| async def assign_subscription(self, user_id: int, days: int) -> dict: |
| |
| now = datetime.utcnow() |
|
|
| user = await self.get_user(user_id) |
| if user: |
| current_expiry = user.get("subscription_expiry") |
| if current_expiry and current_expiry > now: |
| new_expiry = current_expiry + timedelta(days=days) |
| else: |
| new_expiry = now + timedelta(days=days) |
| else: |
| new_expiry = now + timedelta(days=days) |
|
|
| await self.dbs["tracking"]["users"].update_one( |
| {"_id": user_id}, |
| { |
| "$set": { |
| "subscription_expiry": new_expiry, |
| "subscription_status": "active", |
| }, |
| "$setOnInsert": { |
| "_id": user_id, |
| "first_name": f"User {user_id}", |
| "username": None, |
| "created_at": now, |
| } |
| }, |
| upsert=True |
| ) |
| return { |
| "user_id": user_id, |
| "subscription_expiry": new_expiry.isoformat(), |
| "subscription_status": "active", |
| "days_assigned": days, |
| } |
|
|
| |
| |
| |
| async def create_custom_catalog(self, name: str, visible: bool = True) -> Optional[str]: |
| name = (name or "").strip() |
| if not name: |
| return None |
|
|
| now = datetime.utcnow() |
| result = await self.dbs["tracking"]["custom_catalogs"].insert_one({ |
| "name": name, |
| "visible": bool(visible), |
| "items": [], |
| "created_at": now, |
| "updated_at": now, |
| }) |
| return str(result.inserted_id) |
|
|
| async def get_custom_catalogs(self, visible_only: bool = False) -> List[dict]: |
| query = {"visible": True} if visible_only else {} |
| cursor = self.dbs["tracking"]["custom_catalogs"].find(query).sort("updated_at", DESCENDING) |
| catalogs = await cursor.to_list(None) |
| return [convert_objectid_to_str(catalog) for catalog in catalogs] |
|
|
| async def get_custom_catalog(self, catalog_id: str) -> Optional[dict]: |
| try: |
| catalog = await self.dbs["tracking"]["custom_catalogs"].find_one({"_id": ObjectId(catalog_id)}) |
| return convert_objectid_to_str(catalog) if catalog else None |
| except Exception: |
| return None |
|
|
| async def update_custom_catalog(self, catalog_id: str, name: Optional[str] = None, visible: Optional[bool] = None) -> bool: |
| update_data = {"updated_at": datetime.utcnow()} |
| if name is not None: |
| clean_name = name.strip() |
| if clean_name: |
| update_data["name"] = clean_name |
| if visible is not None: |
| update_data["visible"] = bool(visible) |
|
|
| try: |
| result = await self.dbs["tracking"]["custom_catalogs"].update_one( |
| {"_id": ObjectId(catalog_id)}, |
| {"$set": update_data} |
| ) |
| return result.modified_count > 0 |
| except Exception: |
| return False |
|
|
| async def delete_custom_catalog(self, catalog_id: str) -> bool: |
| try: |
| result = await self.dbs["tracking"]["custom_catalogs"].delete_one({"_id": ObjectId(catalog_id)}) |
| return result.deleted_count > 0 |
| except Exception: |
| return False |
|
|
| async def add_item_to_custom_catalog( |
| self, catalog_id: str, tmdb_id: int, db_index: int, media_type: str |
| ) -> bool: |
| media_type = self._collection_for(media_type) |
| item = { |
| "tmdb_id": int(tmdb_id), |
| "db_index": int(db_index), |
| "media_type": media_type, |
| "added_at": datetime.utcnow(), |
| } |
| try: |
| result = await self.dbs["tracking"]["custom_catalogs"].update_one( |
| { |
| "_id": ObjectId(catalog_id), |
| "items": { |
| "$not": { |
| "$elemMatch": { |
| "tmdb_id": int(tmdb_id), |
| "db_index": int(db_index), |
| "media_type": media_type, |
| } |
| } |
| }, |
| }, |
| { |
| "$push": {"items": {"$each": [item], "$position": 0}}, |
| "$set": {"updated_at": datetime.utcnow()}, |
| } |
| ) |
| return result.modified_count > 0 |
| except Exception: |
| return False |
|
|
| async def remove_item_from_custom_catalog( |
| self, catalog_id: str, tmdb_id: int, db_index: int, media_type: str |
| ) -> bool: |
| media_type = self._collection_for(media_type) |
| try: |
| result = await self.dbs["tracking"]["custom_catalogs"].update_one( |
| {"_id": ObjectId(catalog_id)}, |
| { |
| "$pull": { |
| "items": { |
| "tmdb_id": int(tmdb_id), |
| "db_index": int(db_index), |
| "media_type": media_type, |
| } |
| }, |
| "$set": {"updated_at": datetime.utcnow()}, |
| } |
| ) |
| return result.modified_count > 0 |
| except Exception: |
| return False |
|
|
| async def find_media_doc(self, media_type: str, tmdb_id: int) -> Optional[Tuple[dict, int]]: |
| |
| collection_name = self._collection_for(media_type) |
| try: |
| tmdb_id = int(tmdb_id) |
| except (TypeError, ValueError): |
| return None |
|
|
| for i in range(1, self.current_db_index + 1): |
| db_key = f"storage_{i}" |
| if db_key not in self.dbs: |
| continue |
| doc = await self.dbs[db_key][collection_name].find_one({"tmdb_id": tmdb_id}) |
| if doc: |
| doc["db_index"] = i |
| return doc, i |
| return None |
|
|
| async def purge_media_from_catalogs(self, tmdb_id: int, media_type: str) -> int: |
| |
| if tmdb_id in (None, "", 0): |
| return 0 |
| try: |
| tmdb_id = int(tmdb_id) |
| except (TypeError, ValueError): |
| return 0 |
|
|
| media_type = self._collection_for(media_type) |
| collection = self.dbs["tracking"]["custom_catalogs"] |
| now = datetime.utcnow() |
| try: |
| result = await collection.update_many( |
| {"items": {"$elemMatch": {"tmdb_id": tmdb_id, "media_type": media_type}}}, |
| [ |
| { |
| "$set": { |
| "items": { |
| "$filter": { |
| "input": "$items", |
| "as": "it", |
| "cond": { |
| "$not": [ |
| { |
| "$and": [ |
| {"$eq": ["$$it.tmdb_id", tmdb_id]}, |
| {"$eq": ["$$it.media_type", media_type]}, |
| ] |
| } |
| ] |
| }, |
| } |
| } |
| } |
| }, |
| {"$set": {"item_count": {"$size": "$items"}, "updated_at": now}}, |
| ], |
| ) |
| if result.modified_count: |
| LOGGER.info( |
| f"Purged {media_type} tmdb_id {tmdb_id} from " |
| f"{result.modified_count} catalog(s)." |
| ) |
| return result.modified_count |
| except Exception as e: |
| LOGGER.error(f"Failed to purge tmdb_id {tmdb_id} from catalogs: {e}") |
| return 0 |
|
|
| async def custom_catalog_contains_item( |
| self, catalog_id: str, tmdb_id: int, db_index: int, media_type: str |
| ) -> bool: |
| media_type = self._collection_for(media_type) |
| try: |
| catalog = await self.dbs["tracking"]["custom_catalogs"].find_one({ |
| "_id": ObjectId(catalog_id), |
| "items": { |
| "$elemMatch": { |
| "tmdb_id": int(tmdb_id), |
| "db_index": int(db_index), |
| "media_type": media_type, |
| } |
| } |
| }) |
| return bool(catalog) |
| except Exception: |
| return False |
|
|
| async def get_custom_catalog_items( |
| self, catalog_id: str, media_type: Optional[str] = None, page: int = 1, page_size: int = 24 |
| ) -> dict: |
| catalog = await self.get_custom_catalog(catalog_id) |
| if not catalog: |
| return {"catalog": None, "items": [], "total_count": 0, "current_page": page, "total_pages": 0} |
|
|
| db_media_type = None |
| if media_type: |
| db_media_type = self._collection_for(media_type) |
|
|
| raw_items = catalog.get("items", []) or [] |
| if db_media_type: |
| raw_items = [item for item in raw_items if item.get("media_type") == db_media_type] |
|
|
| total_count = len(raw_items) |
| skip = (page - 1) * page_size |
| selected_items = raw_items[skip:skip + page_size] |
|
|
| hydrated_items = [] |
| for item in selected_items: |
| doc = await self.get_document( |
| item.get("media_type", "movie"), |
| int(item.get("tmdb_id")), |
| int(item.get("db_index", 1)) |
| ) |
| if doc: |
| hydrated_items.append(doc) |
|
|
| total_pages = (total_count + page_size - 1) // page_size if total_count else 0 |
| return { |
| "catalog": catalog, |
| "items": hydrated_items, |
| "total_count": total_count, |
| "current_page": page, |
| "total_pages": total_pages, |
| } |
|
|
|
|
| |
| |
| |
| def _get_sort_dict(self, sort_params: List[Tuple[str, str]]) -> Dict[str, int]: |
| if sort_params: |
| sort_field, sort_direction = sort_params[0] |
| return {sort_field: DESCENDING if sort_direction.lower() == "desc" else ASCENDING} |
| return {"updated_on": DESCENDING} |
|
|
| async def _paginate_collection( |
| self, |
| collection_name: str, |
| sort_dict: Dict[str, int], |
| page: int, |
| page_size: int, |
| filter_dict: Optional[dict] = None |
| ): |
| filter_dict = filter_dict or {} |
| skip = (page - 1) * page_size |
| results = [] |
| dbs_checked = [] |
| total_count = 0 |
|
|
| db_counts = [] |
| for i in range(1, self.current_db_index + 1): |
| db_key = f"storage_{i}" |
| db = self.dbs[db_key] |
| count = await db[collection_name].count_documents(filter_dict) |
| db_counts.append((i, count)) |
| total_count += count |
|
|
| start_db_index = None |
| for db_index, count in reversed(db_counts): |
| if skip < count: |
| start_db_index = db_index |
| break |
| skip -= count |
|
|
| if not start_db_index: |
| return [], [], total_count |
|
|
| for db_index, count in reversed(db_counts): |
| if db_index < start_db_index: |
| continue |
|
|
| db_key = f"storage_{db_index}" |
| db = self.dbs[db_key] |
| dbs_checked.append(db_index) |
|
|
| cursor = ( |
| db[collection_name] |
| .find(filter_dict) |
| .sort(sort_dict) |
| .skip(skip if db_index == start_db_index else 0) |
| .limit(page_size - len(results)) |
| ) |
|
|
| docs = await cursor.to_list(None) |
| results.extend(docs) |
|
|
| if len(results) >= page_size: |
| break |
|
|
| return results, dbs_checked, total_count |
|
|
| async def _move_document( |
| self, collection_name: str, document: dict, old_db_index: int |
| ) -> bool: |
| current_db_key = f"storage_{self.current_db_index}" |
| old_db_key = f"storage_{old_db_index}" |
| document["db_index"] = self.current_db_index |
| try: |
| await self.dbs[current_db_key][collection_name].insert_one(document) |
| await self.dbs[old_db_key][collection_name].delete_one({"_id": document["_id"]}) |
| LOGGER.info(f"✅ Moved document {document.get('tmdb_id')} from {old_db_key} to {current_db_key}") |
| return True |
| except Exception as e: |
| LOGGER.error(f"Error moving document to {current_db_key}: {e}") |
| return False |
|
|
| async def _handle_storage_error(self, func, *args, total_storage_dbs: int) -> Optional[Any]: |
| next_db_index = (self.current_db_index % total_storage_dbs) + 1 |
| if next_db_index == 1: |
| LOGGER.warning("⚠️ All storage databases are full! Add more.") |
| return None |
| self.current_db_index = next_db_index |
| await self.update_current_db_index() |
| LOGGER.info(f"Switched to storage_{self.current_db_index}") |
| return await func(*args) |
|
|
| |
| @staticmethod |
| def _collection_for(media_type: str) -> str: |
| return "tv" if str(media_type).lower() in ("tv", "series") else "movie" |
|
|
| |
| async def _edit_media_doc(self, collection_name: str, tmdb_id: int, db_index: int, mutate) -> bool: |
| db_key = f"storage_{db_index}" |
| doc = await self.dbs[db_key][collection_name].find_one({"tmdb_id": tmdb_id}) |
| if not doc: |
| return False |
| if not await mutate(doc): |
| return False |
| doc["updated_on"] = datetime.utcnow() |
| result = await self.dbs[db_key][collection_name].replace_one({"tmdb_id": tmdb_id}, doc) |
| return result.modified_count > 0 |
|
|
| |
| async def _find_existing_media( |
| self, collection_name: str, imdb_id, tmdb_id, title, release_year, total_storage_dbs: int |
| ) -> Tuple[Optional[dict], Optional[str], Optional[int]]: |
| for db_index in range(1, total_storage_dbs + 1): |
| col = self.dbs[f"storage_{db_index}"][collection_name] |
| doc = None |
| if imdb_id: |
| doc = await col.find_one({"imdb_id": imdb_id}) |
| if not doc and tmdb_id: |
| doc = await col.find_one({"tmdb_id": tmdb_id}) |
| if not doc and title and release_year: |
| doc = await col.find_one({"title": title, "release_year": release_year}) |
| if doc: |
| return doc, f"storage_{db_index}", db_index |
| return None, None, None |
|
|
|
|
| |
| |
| |
|
|
| async def _build_part_id_and_size(self, parts: List[dict]) -> Tuple[str, str]: |
| sorted_parts = sorted(parts, key=lambda p: p.get("part_number", 0)) |
| payload = {"parts": [{"chat_id": p["chat_id"], "msg_id": p["msg_id"]} for p in sorted_parts]} |
| encoded = await encode_string(payload) |
| total_bytes = sum(p.get("size_bytes", 0) for p in sorted_parts) |
| from Backend.helper.pyro import get_readable_file_size |
| size_str = get_readable_file_size(total_bytes) |
| return encoded, size_str |
|
|
| async def remove_media_part(self, channel: int, msg_id: int) -> bool: |
| try: |
| legacy_hash = await encode_string({"chat_id": channel, "msg_id": msg_id}) |
| if await self.delete_media_by_stream_id(legacy_hash): |
| return True |
| except Exception as e: |
| LOGGER.error(f"remove_media_part: legacy lookup failed: {e}") |
|
|
| for i in range(1, self.current_db_index + 1): |
| db = self.dbs[f"storage_{i}"] |
|
|
| movie = await db["movie"].find_one( |
| {"telegram.parts": {"$elemMatch": {"chat_id": channel, "msg_id": msg_id}}} |
| ) |
| if movie: |
| new_telegram = [] |
| for q in movie.get("telegram", []): |
| parts = q.get("parts") |
| if parts and any(p.get("chat_id") == channel and p.get("msg_id") == msg_id for p in parts): |
| remaining = [p for p in parts if not (p.get("chat_id") == channel and p.get("msg_id") == msg_id)] |
| if not remaining: |
| continue |
| new_id, new_size = await self._build_part_id_and_size(remaining) |
| q["parts"] = remaining |
| q["id"] = new_id |
| q["size"] = new_size |
| new_telegram.append(q) |
|
|
| if len(new_telegram) == 0: |
| await db["movie"].delete_one({"_id": movie["_id"]}) |
| await self.purge_media_from_catalogs(movie.get("tmdb_id"), "movie") |
| else: |
| movie["telegram"] = new_telegram |
| movie["updated_on"] = datetime.utcnow() |
| await db["movie"].replace_one({"_id": movie["_id"]}, movie) |
| return True |
|
|
| tv = await db["tv"].find_one( |
| {"seasons.episodes.telegram.parts": {"$elemMatch": {"chat_id": channel, "msg_id": msg_id}}} |
| ) |
| if tv: |
| for season in tv.get("seasons", []): |
| for episode in season.get("episodes", []): |
| new_telegram = [] |
| for q in episode.get("telegram", []): |
| parts = q.get("parts") |
| if parts and any(p.get("chat_id") == channel and p.get("msg_id") == msg_id for p in parts): |
| remaining = [p for p in parts if not (p.get("chat_id") == channel and p.get("msg_id") == msg_id)] |
| if not remaining: |
| continue |
| new_id, new_size = await self._build_part_id_and_size(remaining) |
| q["parts"] = remaining |
| q["id"] = new_id |
| q["size"] = new_size |
| new_telegram.append(q) |
| episode["telegram"] = new_telegram |
| season["episodes"] = [e for e in season.get("episodes", []) if e.get("telegram")] |
| tv["seasons"] = [s for s in tv.get("seasons", []) if s.get("episodes")] |
|
|
| if len(tv["seasons"]) == 0: |
| await db["tv"].delete_one({"_id": tv["_id"]}) |
| await self.purge_media_from_catalogs(tv.get("tmdb_id"), "tv") |
| else: |
| tv["updated_on"] = datetime.utcnow() |
| await db["tv"].replace_one({"_id": tv["_id"]}, tv) |
| return True |
|
|
| return False |
|
|
| async def insert_media( |
| self, metadata_info: dict, |
| channel: int, msg_id: int, size: str, name: str, raw_size: int = 0 |
| ) -> Optional[ObjectId]: |
|
|
| group_key = metadata_info.get("group_key") |
| part_number = metadata_info.get("part_number") |
|
|
| if group_key: |
| part = { |
| "part_number": part_number or 1, |
| "chat_id": channel, |
| "msg_id": msg_id, |
| "size_bytes": raw_size, |
| } |
| part_id, part_size = await self._build_part_id_and_size([part]) |
| quality_detail = QualityDetail( |
| quality=metadata_info['quality'], |
| id=part_id, |
| name=name, |
| size=part_size, |
| group_key=group_key, |
| parts=[QualityPart(**part)], |
| ) |
| else: |
| quality_detail = QualityDetail( |
| quality=metadata_info['quality'], |
| id=metadata_info['encoded_string'], |
| name=name, |
| size=size, |
| ) |
|
|
| if metadata_info['media_type'] == "movie": |
| media = MovieSchema( |
| tmdb_id=metadata_info['tmdb_id'], |
| imdb_id=metadata_info['imdb_id'], |
| db_index=self.current_db_index, |
| title=metadata_info['title'], |
| genres=metadata_info['genres'], |
| description=metadata_info['description'], |
| rating=metadata_info['rate'], |
| release_year=metadata_info['year'], |
| poster=metadata_info['poster'], |
| backdrop=metadata_info['backdrop'], |
| logo=metadata_info['logo'], |
| cast=metadata_info['cast'], |
| runtime=metadata_info['runtime'], |
| media_type=metadata_info['media_type'], |
| is_anime=metadata_info.get('is_anime', False), |
| original_language=metadata_info.get('original_language'), |
| origin_country=metadata_info.get('origin_country', []) or [], |
| telegram=[quality_detail] |
| ) |
| return await self.update_movie(media) |
| else: |
| tv_show = TVShowSchema( |
| tmdb_id=metadata_info['tmdb_id'], |
| imdb_id=metadata_info['imdb_id'], |
| db_index=self.current_db_index, |
| title=metadata_info['title'], |
| genres=metadata_info['genres'], |
| description=metadata_info['description'], |
| rating=metadata_info['rate'], |
| release_year=metadata_info['year'], |
| poster=metadata_info['poster'], |
| backdrop=metadata_info['backdrop'], |
| logo=metadata_info['logo'], |
| cast=metadata_info['cast'], |
| runtime=metadata_info['runtime'], |
| media_type=metadata_info['media_type'], |
| is_anime=metadata_info.get('is_anime', False), |
| original_language=metadata_info.get('original_language'), |
| origin_country=metadata_info.get('origin_country', []) or [], |
| seasons=[Season( |
| season_number=metadata_info['season_number'], |
| episodes=[Episode( |
| episode_number=metadata_info['episode_number'], |
| title=metadata_info['episode_title'], |
| episode_backdrop=metadata_info['episode_backdrop'], |
| overview=metadata_info['episode_overview'], |
| released=metadata_info['episode_released'], |
| telegram=[quality_detail] |
| )] |
| )] |
| ) |
| return await self.update_tv_show(tv_show) |
|
|
| async def _delete_split_part(self, part: dict) -> None: |
| try: |
| chat_id = int(f"-100{part['chat_id']}") |
| msg_id = int(part["msg_id"]) |
| create_task(delete_message(chat_id, msg_id)) |
| except Exception as e: |
| LOGGER.error(f"Failed to delete split part message: {e}") |
|
|
| async def _merge_split_part(self, qualities: List[dict], quality_to_update: dict) -> List[dict]: |
| group_key = quality_to_update.get("group_key") |
| incoming_parts = quality_to_update.get("parts") or [] |
| if not incoming_parts: |
| return qualities + [quality_to_update] |
| new_part = incoming_parts[0] |
| replace_mode = SettingsManager.current().replace_mode |
|
|
| merged = False |
| result = [] |
| for q in qualities: |
| if group_key is not None and q.get("group_key") == group_key: |
| existing_parts = q.get("parts") or [] |
|
|
| if replace_mode: |
| for old_part in existing_parts: |
| if old_part.get("part_number") != new_part.get("part_number"): |
| continue |
| if ( |
| old_part.get("chat_id") == new_part.get("chat_id") |
| and old_part.get("msg_id") == new_part.get("msg_id") |
| ): |
| continue |
| await self._delete_split_part(old_part) |
|
|
| existing_parts = [ |
| p for p in existing_parts if p.get("part_number") != new_part.get("part_number") |
| ] |
| existing_parts.append(new_part) |
| new_id, new_size = await self._build_part_id_and_size(existing_parts) |
| q["parts"] = existing_parts |
| q["id"] = new_id |
| q["size"] = new_size |
| q["name"] = quality_to_update.get("name", q.get("name")) |
| merged = True |
| result.append(q) |
|
|
| if not merged: |
| result.append(quality_to_update) |
| return result |
|
|
| async def _apply_quality_update( |
| self, existing_qualities: List[dict], quality_to_update: dict |
| ) -> List[dict]: |
| target_quality = quality_to_update.get("quality") |
| incoming_group_key = quality_to_update.get("group_key") |
| replace_mode = SettingsManager.current().replace_mode |
|
|
| if incoming_group_key: |
| |
| if replace_mode: |
| stale = [ |
| q for q in existing_qualities |
| if q.get("quality") == target_quality |
| and q.get("group_key") != incoming_group_key |
| ] |
| for q in stale: |
| await self._queue_quality_deletion(q) |
| existing_qualities = [ |
| q for q in existing_qualities |
| if not ( |
| q.get("quality") == target_quality |
| and q.get("group_key") != incoming_group_key |
| ) |
| ] |
| return await self._merge_split_part(existing_qualities, quality_to_update) |
|
|
| |
| if replace_mode: |
| stale = [q for q in existing_qualities if q.get("quality") == target_quality] |
| for q in stale: |
| await self._queue_quality_deletion(q) |
| existing_qualities = [ |
| q for q in existing_qualities if q.get("quality") != target_quality |
| ] |
| existing_qualities.append(quality_to_update) |
| return existing_qualities |
|
|
| |
| existing_qualities.append(quality_to_update) |
| return existing_qualities |
|
|
| async def update_movie(self, movie_data: MovieSchema) -> Optional[ObjectId]: |
| try: |
| movie_dict = movie_data.dict() |
| except ValidationError as e: |
| LOGGER.error(f"Validation error: {e}") |
| return None |
|
|
| imdb_id = movie_dict["imdb_id"] |
| tmdb_id = movie_dict["tmdb_id"] |
| title = movie_dict["title"] |
| release_year = movie_dict["release_year"] |
|
|
| quality_to_update = movie_dict["telegram"][0] |
|
|
| current_db_key = f"storage_{self.current_db_index}" |
| total_storage_dbs = len(self.dbs) - 1 |
|
|
| existing_movie, existing_db_key, existing_db_index = await self._find_existing_media( |
| "movie", imdb_id, tmdb_id, title, release_year, total_storage_dbs |
| ) |
|
|
| |
| if not existing_movie: |
| try: |
| movie_dict["db_index"] = self.current_db_index |
| result = await self.dbs[current_db_key]["movie"].insert_one(movie_dict) |
| return result.inserted_id |
| except Exception as e: |
| LOGGER.error(f"Insertion failed in {current_db_key}: {e}") |
| if any(keyword in str(e).lower() for keyword in ["storage", "quota"]): |
| return await self._handle_storage_error(self.update_movie, movie_data, total_storage_dbs=total_storage_dbs) |
| return None |
|
|
| |
| movie_id = existing_movie["_id"] |
|
|
| if imdb_id and not existing_movie.get("imdb_id"): |
| existing_movie["imdb_id"] = imdb_id |
| if tmdb_id and not existing_movie.get("tmdb_id"): |
| existing_movie["tmdb_id"] = tmdb_id |
| if movie_dict.get("is_anime"): |
| existing_movie["is_anime"] = True |
|
|
| existing_qualities = existing_movie.get("telegram", []) |
|
|
| existing_qualities = await self._apply_quality_update(existing_qualities, quality_to_update) |
|
|
| existing_movie["telegram"] = existing_qualities |
| existing_movie["updated_on"] = datetime.utcnow() |
|
|
| if existing_db_index != self.current_db_index: |
| try: |
| if await self._move_document("movie", existing_movie, existing_db_index): |
| return movie_id |
| except Exception as e: |
| LOGGER.error(f"Error moving movie to {current_db_key}: {e}") |
| if any(keyword in str(e).lower() for keyword in ["storage", "quota"]): |
| return await self._handle_storage_error(self.update_movie, movie_data, total_storage_dbs=total_storage_dbs) |
|
|
| try: |
| await self.dbs[existing_db_key]["movie"].replace_one({"_id": movie_id}, existing_movie) |
| return movie_id |
| except Exception as e: |
| LOGGER.error(f"Failed to update movie {tmdb_id} in {existing_db_key}: {e}") |
| if any(keyword in str(e).lower() for keyword in ["storage", "quota"]): |
| return await self._handle_storage_error(self.update_movie, movie_data, total_storage_dbs=total_storage_dbs) |
|
|
| async def update_tv_show(self, tv_show_data: TVShowSchema) -> Optional[ObjectId]: |
| try: |
| tv_show_dict = tv_show_data.dict() |
| except ValidationError as e: |
| LOGGER.error(f"Validation error: {e}") |
| return None |
|
|
| imdb_id = tv_show_dict.get("imdb_id") |
| tmdb_id = tv_show_dict.get("tmdb_id") |
| title = tv_show_dict["title"] |
| release_year = tv_show_dict["release_year"] |
|
|
| current_db_key = f"storage_{self.current_db_index}" |
| total_storage_dbs = len(self.dbs) - 1 |
|
|
| existing_tv, existing_db_key, existing_db_index = await self._find_existing_media( |
| "tv", imdb_id, tmdb_id, title, release_year, total_storage_dbs |
| ) |
|
|
| |
| if not existing_tv: |
| try: |
| tv_show_dict["db_index"] = self.current_db_index |
| result = await self.dbs[current_db_key]["tv"].insert_one(tv_show_dict) |
| return result.inserted_id |
| except Exception as e: |
| LOGGER.error(f"Insertion failed in {current_db_key}: {e}") |
| if any(keyword in str(e).lower() for keyword in ["storage", "quota"]): |
| return await self._handle_storage_error(self.update_tv_show, tv_show_data, total_storage_dbs=total_storage_dbs) |
| return None |
|
|
| |
| tv_id = existing_tv["_id"] |
|
|
| if imdb_id and not existing_tv.get("imdb_id"): |
| existing_tv["imdb_id"] = imdb_id |
| if tmdb_id and not existing_tv.get("tmdb_id"): |
| existing_tv["tmdb_id"] = tmdb_id |
| if tv_show_dict.get("is_anime"): |
| existing_tv["is_anime"] = True |
|
|
| for season in tv_show_dict["seasons"]: |
| existing_season = next( |
| (s for s in existing_tv["seasons"] |
| if s["season_number"] == season["season_number"]), |
| None |
| ) |
|
|
| if not existing_season: |
| existing_tv["seasons"].append(season) |
| continue |
|
|
| for episode in season["episodes"]: |
| existing_episode = next( |
| (e for e in existing_season["episodes"] |
| if e["episode_number"] == episode["episode_number"]), |
| None |
| ) |
|
|
| if not existing_episode: |
| existing_season["episodes"].append(episode) |
| continue |
|
|
| existing_episode.setdefault("telegram", []) |
|
|
| for quality in episode["telegram"]: |
| existing_episode["telegram"] = await self._apply_quality_update( |
| existing_episode["telegram"], quality |
| ) |
|
|
| existing_tv["updated_on"] = datetime.utcnow() |
|
|
| |
| if existing_db_index != self.current_db_index: |
| try: |
| if await self._move_document("tv", existing_tv, existing_db_index): |
| return tv_id |
| except Exception as e: |
| LOGGER.error(f"Error moving TV show to {current_db_key}: {e}") |
| if any(keyword in str(e).lower() for keyword in ["storage", "quota"]): |
| return await self._handle_storage_error(self.update_tv_show, tv_show_data, total_storage_dbs=total_storage_dbs) |
| return tv_id |
|
|
| try: |
| await self.dbs[existing_db_key]["tv"].replace_one({"_id": tv_id}, existing_tv) |
| return tv_id |
| except Exception as e: |
| LOGGER.error(f"Failed to update TV show {tmdb_id} in {existing_db_key}: {e}") |
| if any(keyword in str(e).lower() for keyword in ["storage", "quota"]): |
| return await self._handle_storage_error(self.update_tv_show, tv_show_data, total_storage_dbs=total_storage_dbs) |
| |
| async def sort_movies(self, sort_params, page, page_size, genre_filter=None): |
| sort_dict = self._get_sort_dict(sort_params) |
| filter_dict = {"genres": {"$in": [genre_filter]}} if genre_filter else {} |
| results, dbs_checked, total_count = await self._paginate_collection( |
| "movie", sort_dict, page, page_size, filter_dict=filter_dict |
| ) |
| total_pages = (total_count + page_size - 1) // page_size |
| return { |
| "total_count": total_count, |
| "total_pages": total_pages, |
| "databases_checked": dbs_checked, |
| "current_page": page, |
| "movies": [convert_objectid_to_str(result) for result in results], |
| } |
|
|
| async def sort_tv_shows(self, sort_params, page, page_size, genre_filter=None): |
| sort_dict = self._get_sort_dict(sort_params) |
| filter_dict = {"genres": {"$in": [genre_filter]}} if genre_filter else {} |
| results, dbs_checked, total_count = await self._paginate_collection( |
| "tv", sort_dict, page, page_size, filter_dict=filter_dict |
| ) |
| total_pages = (total_count + page_size - 1) // page_size |
| return { |
| "total_count": total_count, |
| "total_pages": total_pages, |
| "databases_checked": dbs_checked, |
| "current_page": page, |
| "tv_shows": [convert_objectid_to_str(result) for result in results], |
| } |
|
|
| async def search_documents( |
| self, |
| query: str, |
| page: int, |
| page_size: int |
| ) -> dict: |
|
|
| skip = (page - 1) * page_size |
| |
| words = query.split() |
| regex_query = { |
| '$regex': '.*' + '.*'.join(words) + '.*', |
| '$options': 'i' |
| } |
| |
| tv_pipeline = [ |
| {"$match": {"$or": [ |
| {"title": regex_query}, |
| {"seasons.episodes.telegram.name": regex_query} |
| ]}}, |
| {"$project": { |
| "_id": 1, "tmdb_id": 1, "title": 1, "genres": 1, "rating": 1, "imdb_id": 1, |
| "release_year": 1, "poster": 1, "backdrop": 1, "description": 1, "logo": 1, |
| "media_type": 1, "db_index": 1 |
| }} |
| ] |
| |
| movie_pipeline = [ |
| {"$match": {"$or": [ |
| {"title": regex_query}, |
| {"telegram.name": regex_query} |
| ]}}, |
| {"$project": { |
| "_id": 1, "tmdb_id": 1, "title": 1, "genres": 1, "rating": 1, |
| "release_year": 1, "poster": 1, "backdrop": 1, "description": 1, |
| "media_type": 1, "db_index": 1, "imdb_id": 1, "logo": 1 |
| }} |
| ] |
| |
| results = [] |
| dbs_checked = [] |
| |
| active_db_key = f"storage_{self.current_db_index}" |
| active_db = self.dbs[active_db_key] |
| dbs_checked.append(self.current_db_index) |
| |
| tv_results = await active_db["tv"].aggregate(tv_pipeline).to_list(None) |
| movie_results = await active_db["movie"].aggregate(movie_pipeline).to_list(None) |
| combined = tv_results + movie_results |
| results.extend(combined) |
| |
| if len(results) < page_size: |
| previous_db_index = self.current_db_index - 1 |
| while previous_db_index > 0 and len(results) < page_size: |
| prev_db_key = f"storage_{previous_db_index}" |
| prev_db = self.dbs[prev_db_key] |
| tv_results_prev = await prev_db["tv"].aggregate(tv_pipeline).to_list(None) |
| movie_results_prev = await prev_db["movie"].aggregate(movie_pipeline).to_list(None) |
| combined_prev = tv_results_prev + movie_results_prev |
| results.extend(combined_prev) |
| dbs_checked.append(previous_db_index) |
| previous_db_index -= 1 |
|
|
| total_count = 0 |
| for db_index in dbs_checked: |
| key = f"storage_{db_index}" |
| db = self.dbs[key] |
| tv_count = await db["tv"].count_documents({ |
| "$or": [ |
| {"title": regex_query}, |
| {"seasons.episodes.telegram.name": regex_query} |
| ] |
| }) |
| movie_count = await db["movie"].count_documents({ |
| "$or": [ |
| {"title": regex_query}, |
| {"telegram.name": regex_query} |
| ] |
| }) |
| total_count += (tv_count + movie_count) |
| |
| paged_results = results[skip:skip + page_size] |
|
|
| return { |
| "total_count": total_count, |
| "results": [convert_objectid_to_str(doc) for doc in paged_results] |
| } |
|
|
|
|
| async def get_media_details( |
| self, |
| imdb_id: str, |
| season_number: Optional[int] = None, |
| episode_number: Optional[int] = None |
| ) -> Optional[dict]: |
|
|
| for db_idx in range(self.current_db_index, 0, -1): |
| db_key = f"storage_{db_idx}" |
| |
| if episode_number is not None and season_number is not None: |
| tv_show = await self.dbs[db_key]["tv"].find_one({"imdb_id": imdb_id}) |
| if tv_show: |
| for season in tv_show.get("seasons", []): |
| if season.get("season_number") == season_number: |
| for episode in season.get("episodes", []): |
| if episode.get("episode_number") == episode_number: |
| details = convert_objectid_to_str(episode) |
| details.update({ |
| "imdb_id": imdb_id, |
| "type": "tv", |
| "season_number": season_number, |
| "episode_number": episode_number, |
| "backdrop": episode.get("episode_backdrop"), |
| "db_index": db_idx |
| }) |
| return details |
| |
| elif season_number is not None: |
| tv_show = await self.dbs[db_key]["tv"].find_one({"imdb_id": imdb_id}) |
| if tv_show: |
| for season in tv_show.get("seasons", []): |
| if season.get("season_number") == season_number: |
| details = convert_objectid_to_str(season) |
| details.update({ |
| "imdb_id": imdb_id, |
| "type": "tv", |
| "season_number": season_number, |
| "db_index": db_idx |
| }) |
| return details |
| |
| else: |
| tv_doc = await self.dbs[db_key]["tv"].find_one({"imdb_id": imdb_id}) |
| if tv_doc: |
| tv_doc = convert_objectid_to_str(tv_doc) |
| tv_doc["type"] = "tv" |
| tv_doc["db_index"] = db_idx |
| return tv_doc |
| |
| movie_doc = await self.dbs[db_key]["movie"].find_one({"imdb_id": imdb_id}) |
| if movie_doc: |
| movie_doc = convert_objectid_to_str(movie_doc) |
| movie_doc["type"] = "movie" |
| movie_doc["db_index"] = db_idx |
| return movie_doc |
| |
| return None |
|
|
| |
| |
| |
|
|
| async def get_document(self, media_type: str, tmdb_id: int, db_index: int) -> Optional[Dict[str, Any]]: |
| db_key = f"storage_{db_index}" |
| collection_name = self._collection_for(media_type) |
| document = await self.dbs[db_key][collection_name].find_one({"tmdb_id": int(tmdb_id)}) |
| return convert_objectid_to_str(document) if document else None |
|
|
| async def update_document( |
| self, media_type: str, tmdb_id: int, db_index: int, update_data: Dict[str, Any] |
| ): |
| update_data.pop('_id', None) |
| db_key = f"storage_{db_index}" |
| collection_name = self._collection_for(media_type) |
| collection = self.dbs[db_key][collection_name] |
|
|
| try: |
| result = await collection.update_one({"tmdb_id": int(tmdb_id)}, {"$set": update_data}) |
|
|
| return result.modified_count > 0 |
|
|
| except Exception as e: |
| err_str = str(e).lower() |
| LOGGER.error(f"Error updating document in {db_key}: {e}") |
| if "storage" in err_str or "quota" in err_str: |
| total_storage_dbs = len(self.dbs) - 1 |
| db_index_int = int(db_index) |
| next_db_index = (db_index_int % total_storage_dbs) + 1 |
| if next_db_index == 1: |
| LOGGER.warning("⚠️ All storage databases are full! Add more.") |
| return False |
|
|
| new_db_key = f"storage_{next_db_index}" |
| LOGGER.info(f"Switching from {db_key} to {new_db_key} due to storage error.") |
|
|
| try: |
| old_doc = await self.dbs[db_key][collection_name].find_one({"tmdb_id": int(tmdb_id)}) |
| if not old_doc: |
| LOGGER.error(f"Document with tmdb_id {tmdb_id} not found in {db_key} during migration.") |
| return False |
|
|
| old_doc.update(update_data) |
| old_doc["db_index"] = next_db_index |
| old_doc.pop("_id", None) |
| insert_result = await self.dbs[new_db_key][collection_name].insert_one(old_doc) |
| LOGGER.info(f"Inserted document {insert_result.inserted_id} into {new_db_key}") |
| await self.dbs[db_key][collection_name].delete_one({"tmdb_id": int(tmdb_id)}) |
| LOGGER.info(f"Deleted document tmdb_id {tmdb_id} from {db_key}") |
| self.current_db_index = next_db_index |
| await self.update_current_db_index() |
| LOGGER.info(f"Switched to {new_db_key} and document migrated successfully.") |
| return True |
|
|
| except Exception as migrate_error: |
| LOGGER.error(f"Error migrating document tmdb_id {tmdb_id} to {new_db_key}: {migrate_error}") |
| return False |
| raise |
|
|
| |
| async def _queue_quality_deletion(self, quality: dict) -> None: |
| parts = quality.get("parts") |
| if parts: |
| for part in parts: |
| await self._delete_split_part(part) |
| return |
|
|
| old_id = quality.get("id") |
| if not old_id: |
| return |
| try: |
| decoded = await decode_string(old_id) |
| if isinstance(decoded, dict) and decoded.get("parts"): |
| for part in decoded["parts"]: |
| await self._delete_split_part(part) |
| return |
| chat_id = int(f"-100{decoded['chat_id']}") |
| msg_id = int(decoded["msg_id"]) |
| create_task(delete_message(chat_id, msg_id)) |
| except Exception as e: |
| LOGGER.error(f"Failed to queue file for deletion: {e}") |
|
|
| async def delete_document(self, media_type: str, tmdb_id: int, db_index: int) -> bool: |
| db_key = f"storage_{db_index}" |
| collection_name = self._collection_for(media_type) |
|
|
| doc = await self.dbs[db_key][collection_name].find_one({"tmdb_id": tmdb_id}) |
| if doc: |
| if collection_name == "movie": |
| for quality in doc.get("telegram", []): |
| await self._queue_quality_deletion(quality) |
| else: |
| for season in doc.get("seasons", []): |
| for episode in season.get("episodes", []): |
| for quality in episode.get("telegram", []): |
| await self._queue_quality_deletion(quality) |
|
|
| result = await self.dbs[db_key][collection_name].delete_one({"tmdb_id": tmdb_id}) |
| if result.deleted_count > 0: |
| await self.purge_media_from_catalogs(tmdb_id, collection_name) |
| LOGGER.info(f"{media_type} with tmdb_id {tmdb_id} deleted successfully.") |
| return True |
| LOGGER.info(f"No document found with tmdb_id {tmdb_id}.") |
| return False |
|
|
| async def get_title_by_stream_id(self, stream_id_hash: str) -> Optional[str]: |
| for i in range(1, self.current_db_index + 1): |
| db = self.dbs[f"storage_{i}"] |
| |
| |
| movie = await db["movie"].find_one({"telegram.id": stream_id_hash}) |
| if movie and "telegram" in movie: |
| for t in movie["telegram"]: |
| if t.get("id") == stream_id_hash: |
| return movie.get("title") |
|
|
| |
| tv = await db["tv"].find_one({"seasons.episodes.telegram.id": stream_id_hash}) |
| if tv and "seasons" in tv: |
| title = tv.get("title", "Unknown Series") |
| for season in tv.get("seasons", []): |
| for episode in season.get("episodes", []): |
| for t in episode.get("telegram", []): |
| if t.get("id") == stream_id_hash: |
| s_num = season.get("season_number", 0) |
| e_num = episode.get("episode_number", 0) |
| return f"{title} S{s_num:02d}E{e_num:02d}" |
|
|
| return None |
|
|
| async def delete_media_by_stream_id(self, stream_id_hash: str) -> bool: |
| for i in range(1, self.current_db_index + 1): |
| db = self.dbs[f"storage_{i}"] |
| |
| |
| movie = await db["movie"].find_one({"telegram.id": stream_id_hash}) |
| if movie: |
| movie["telegram"] = [q for q in movie.get("telegram", []) if q.get("id") != stream_id_hash] |
| if len(movie["telegram"]) == 0: |
| await db["movie"].delete_one({"_id": movie["_id"]}) |
| await self.purge_media_from_catalogs(movie.get("tmdb_id"), "movie") |
| else: |
| movie['updated_on'] = datetime.utcnow() |
| await db["movie"].replace_one({"_id": movie["_id"]}, movie) |
| return True |
|
|
| |
| tv = await db["tv"].find_one({"seasons.episodes.telegram.id": stream_id_hash}) |
| if tv: |
| for season in tv.get("seasons", []): |
| for episode in season.get("episodes", []): |
| for q in episode.get("telegram", []): |
| if q.get("id") == stream_id_hash: |
| episode["telegram"] = [t for t in episode.get("telegram", []) if t.get("id") != stream_id_hash] |
| if len(episode["telegram"]) == 0: |
| season["episodes"] = [e for e in season.get("episodes", []) if e.get("episode_number") != episode.get("episode_number")] |
| if len(season["episodes"]) == 0: |
| tv["seasons"] = [s for s in tv.get("seasons", []) if s.get("season_number") != season.get("season_number")] |
| if len(tv["seasons"]) == 0: |
| await db["tv"].delete_one({"_id": tv["_id"]}) |
| await self.purge_media_from_catalogs(tv.get("tmdb_id"), "tv") |
| return True |
| tv['updated_on'] = datetime.utcnow() |
| await db["tv"].replace_one({"_id": tv["_id"]}, tv) |
| return True |
| return False |
|
|
| async def delete_movie_quality(self, tmdb_id: int, db_index: int, id: str) -> bool: |
| async def mutate(movie): |
| qualities = movie.get("telegram") or [] |
| for q in qualities: |
| if q.get("id") == id: |
| await self._queue_quality_deletion(q) |
| break |
| original_len = len(qualities) |
| movie["telegram"] = [q for q in qualities if q.get("id") != id] |
| return len(movie["telegram"]) != original_len |
| return await self._edit_media_doc("movie", tmdb_id, db_index, mutate) |
|
|
| async def delete_tv_quality(self, tmdb_id: int, db_index: int, season_number: int, episode_number: int, id: str) -> bool: |
| async def mutate(tv): |
| for season in tv.get("seasons", []): |
| if season.get("season_number") != season_number: |
| continue |
| for episode in season.get("episodes", []): |
| if episode.get("episode_number") == episode_number and "telegram" in episode: |
| for q in episode["telegram"]: |
| if q.get("id") == id: |
| await self._queue_quality_deletion(q) |
| break |
| original_len = len(episode["telegram"]) |
| episode["telegram"] = [q for q in episode["telegram"] if q.get("id") != id] |
| return original_len > len(episode["telegram"]) |
| return False |
| return await self._edit_media_doc("tv", tmdb_id, db_index, mutate) |
|
|
| async def delete_tv_episode(self, tmdb_id: int, db_index: int, season_number: int, episode_number: int) -> bool: |
| async def mutate(tv): |
| for season in tv.get("seasons", []): |
| if season.get("season_number") != season_number: |
| continue |
| episodes = season.get("episodes", []) |
| for ep in episodes: |
| if ep.get("episode_number") == episode_number: |
| for quality in ep.get("telegram", []): |
| await self._queue_quality_deletion(quality) |
| break |
| original_len = len(episodes) |
| season["episodes"] = [ep for ep in episodes if ep.get("episode_number") != episode_number] |
| return original_len > len(season["episodes"]) |
| return False |
| return await self._edit_media_doc("tv", tmdb_id, db_index, mutate) |
|
|
| async def delete_tv_season(self, tmdb_id: int, db_index: int, season_number: int) -> bool: |
| async def mutate(tv): |
| seasons = tv.get("seasons", []) |
| for season in seasons: |
| if season.get("season_number") == season_number: |
| for episode in season.get("episodes", []): |
| for quality in episode.get("telegram", []): |
| await self._queue_quality_deletion(quality) |
| break |
| original_len = len(seasons) |
| tv["seasons"] = [s for s in seasons if s.get("season_number") != season_number] |
| return len(tv["seasons"]) != original_len |
| return await self._edit_media_doc("tv", tmdb_id, db_index, mutate) |
|
|
|
|
| |
| async def get_database_stats(self): |
| stats = [] |
| for key in self.dbs.keys(): |
| if key.startswith("storage_"): |
| db = self.dbs[key] |
| movie_count = await db["movie"].count_documents({}) |
| tv_count = await db["tv"].count_documents({}) |
| db_stats = await db.command("dbstats") |
| stats.append({ |
| "db_name": key, |
| "movie_count": movie_count, |
| "tv_count": tv_count, |
| "storageSize": db_stats.get("storageSize", 0), |
| "dataSize": db_stats.get("dataSize", 0) |
| }) |
| return stats |
|
|
|
|
|
|
| |
| |
| |
|
|
| async def add_api_token(self, name: str, daily_limit_gb: float = None, monthly_limit_gb: float = None, user_id: int = None) -> dict: |
| |
| if user_id: |
| existing = await self.dbs["tracking"]["api_tokens"].find_one({"user_id": user_id}) |
| if existing: |
| return convert_objectid_to_str(existing) |
|
|
| alphabet = string.ascii_letters + string.digits |
| token = ''.join(secrets.choice(alphabet) for _ in range(32)) |
| |
| token_doc = { |
| "name": name, |
| "token": token, |
| "user_id": user_id, |
| "created_at": datetime.utcnow(), |
| "limits": { |
| "daily_limit_gb": daily_limit_gb if daily_limit_gb else 0, |
| "monthly_limit_gb": monthly_limit_gb if monthly_limit_gb else 0 |
| }, |
| "usage": { |
| "total_bytes": 0, |
| "daily": {"date": datetime.now(timezone.utc).strftime("%Y-%m-%d"), "bytes": 0}, |
| "monthly": {"month": datetime.now(timezone.utc).strftime("%Y-%m"), "bytes": 0} |
| } |
| } |
| |
| await self.dbs["tracking"]["api_tokens"].insert_one(token_doc) |
| return convert_objectid_to_str(token_doc) |
|
|
| async def get_api_token(self, token: str) -> Optional[dict]: |
| doc = await self.dbs["tracking"]["api_tokens"].find_one({"token": token}) |
| return convert_objectid_to_str(doc) if doc else None |
|
|
| async def get_all_api_tokens(self) -> List[dict]: |
| cursor = self.dbs["tracking"]["api_tokens"].find().sort("created_at", DESCENDING) |
| tokens = await cursor.to_list(None) |
| return [convert_objectid_to_str(token) for token in tokens] |
|
|
| async def revoke_api_token(self, token: str) -> bool: |
| result = await self.dbs["tracking"]["api_tokens"].delete_one({"token": token}) |
| return result.deleted_count > 0 |
|
|
| async def link_token_user(self, token: str, user_id: int) -> bool: |
| |
| result = await self.dbs["tracking"]["api_tokens"].update_one( |
| {"token": token}, |
| {"$set": {"user_id": user_id}} |
| ) |
| return result.modified_count > 0 |
|
|
| async def update_token_usage(self, token: str, bytes_delta: int): |
| today_str = datetime.now(timezone.utc).strftime("%Y-%m-%d") |
| month_str = datetime.now(timezone.utc).strftime("%Y-%m") |
| |
| token_doc = await self.dbs["tracking"]["api_tokens"].find_one({"token": token}) |
| if not token_doc: |
| return |
|
|
| current_daily = token_doc.get("usage", {}).get("daily", {}) |
| if current_daily.get("date") != today_str: |
| await self.dbs["tracking"]["api_tokens"].update_one( |
| {"token": token}, |
| {"$set": {"usage.daily": {"date": today_str, "bytes": 0}}} |
| ) |
|
|
| current_monthly = token_doc.get("usage", {}).get("monthly", {}) |
| if current_monthly.get("month") != month_str: |
| await self.dbs["tracking"]["api_tokens"].update_one( |
| {"token": token}, |
| {"$set": {"usage.monthly": {"month": month_str, "bytes": 0}}} |
| ) |
|
|
| await self.dbs["tracking"]["api_tokens"].update_one( |
| {"token": token}, |
| { |
| "$inc": { |
| "usage.total_bytes": bytes_delta, |
| "usage.daily.bytes": bytes_delta, |
| "usage.monthly.bytes": bytes_delta |
| } |
| } |
| ) |
|
|
| async def update_api_token_limits(self, token: str, daily_limit_gb: float, monthly_limit_gb: float) -> bool: |
| result = await self.dbs["tracking"]["api_tokens"].update_one( |
| {"token": token}, |
| {"$set": { |
| "limits": { |
| "daily_limit_gb": daily_limit_gb if daily_limit_gb else 0, |
| "monthly_limit_gb": monthly_limit_gb if monthly_limit_gb else 0 |
| } |
| }} |
| ) |
| return result.modified_count > 0 |
|
|
| |
| |
| |
| async def flag_dead_link(self, media_type: str, tmdb_id: int, db_index: int, quality_id: str) -> bool: |
| |
| db_key = f"storage_{db_index}" |
| |
| if media_type == "movie": |
| |
| result = await self.dbs[db_key]["movie"].update_one( |
| {"tmdb_id": tmdb_id, "telegram.id": quality_id}, |
| {"$set": {"telegram.$.is_dead": True, "updated_on": datetime.utcnow()}} |
| ) |
| return result.modified_count > 0 |
| |
| elif media_type == "tv": |
| |
| |
| tv = await self.dbs[db_key]["tv"].find_one({"tmdb_id": tmdb_id}) |
| if not tv or "seasons" not in tv: |
| return False |
| |
| found = False |
| for s_idx, season in enumerate(tv["seasons"]): |
| for e_idx, episode in enumerate(season.get("episodes", [])): |
| for q_idx, quality in enumerate(episode.get("telegram", [])): |
| if quality.get("id") == quality_id: |
| tv["seasons"][s_idx]["episodes"][e_idx]["telegram"][q_idx]["is_dead"] = True |
| found = True |
| break |
| if found: break |
| if found: break |
| |
| if found: |
| tv["updated_on"] = datetime.utcnow() |
| result = await self.dbs[db_key]["tv"].replace_one({"tmdb_id": tmdb_id}, tv) |
| return result.modified_count > 0 |
| |
| return False |
|
|
| async def get_all_dead_links(self) -> List[dict]: |
| |
| dead_links = [] |
| |
| for i in range(1, self.current_db_index + 1): |
| db_key = f"storage_{i}" |
| db = self.dbs[db_key] |
| |
| |
| |
| movie_cursor = db["movie"].find({"telegram.is_dead": True}) |
| async for movie in movie_cursor: |
| for quality in movie.get("telegram", []): |
| if quality.get("is_dead"): |
| dead_links.append({ |
| "type": "movie", |
| "tmdb_id": movie.get("tmdb_id"), |
| "db_index": movie.get("db_index", i), |
| "title": movie.get("title"), |
| "year": movie.get("year"), |
| "poster": movie.get("poster"), |
| "quality_id": quality.get("id"), |
| "quality": quality.get("quality"), |
| "size": quality.get("size"), |
| "date_added": quality.get("date_added") |
| }) |
| |
| |
| |
| tv_cursor = db["tv"].find({"seasons.episodes.telegram.is_dead": True}) |
| async for tv in tv_cursor: |
| title = tv.get("title") |
| year = tv.get("year") |
| poster = tv.get("poster") |
| for season in tv.get("seasons", []): |
| s_num = season.get("season_number") |
| for ep in season.get("episodes", []): |
| e_num = ep.get("episode_number") |
| for quality in ep.get("telegram", []): |
| if quality.get("is_dead"): |
| dead_links.append({ |
| "type": "tv", |
| "tmdb_id": tv.get("tmdb_id"), |
| "db_index": tv.get("db_index", i), |
| "title": f"{title} (S{s_num:02d}E{e_num:02d})", |
| "year": year, |
| "poster": poster, |
| "season": s_num, |
| "episode": e_num, |
| "quality_id": quality.get("id"), |
| "quality": quality.get("quality"), |
| "size": quality.get("size"), |
| "date_added": quality.get("date_added") |
| }) |
| |
| return dead_links |
|
|
| |
| |
| |
|
|
| async def log_stream_stats(self, stats: dict) -> None: |
| |
| try: |
| record = { |
| "stream_id": stats.get("stream_id"), |
| "msg_id": stats.get("msg_id"), |
| "chat_id": stats.get("chat_id"), |
| "dc_id": stats.get("dc_id"), |
| "title": stats.get("meta", {}).get("title"), |
| "client_index": stats.get("client_index"), |
| "total_bytes": stats.get("total_bytes", 0), |
| "duration_sec": round(stats.get("duration", 0.0), 2), |
| "avg_mbps": round(stats.get("avg_mbps", 0.0), 3), |
| "peak_mbps": round(stats.get("peak_mbps", 0.0), 3), |
| "status": stats.get("status", "finished"), |
| "parallelism": stats.get("parallelism"), |
| "chunk_size": stats.get("chunk_size"), |
| "logged_at": datetime.utcnow(), |
| } |
| await self.dbs["tracking"]["stream_analytics"].insert_one(record) |
| except Exception as e: |
| LOGGER.warning(f"Stream analytics log failed: {e}") |
|
|
| async def get_stream_analytics(self, limit: int = 200) -> dict: |
| |
| try: |
| col = self.dbs["tracking"]["stream_analytics"] |
|
|
| |
| pipeline = [ |
| {"$group": { |
| "_id": None, |
| "total_streams": {"$sum": 1}, |
| "total_bytes": {"$sum": "$total_bytes"}, |
| "avg_speed": {"$avg": "$avg_mbps"}, |
| "peak_speed": {"$max": "$peak_mbps"}, |
| "avg_duration": {"$avg": "$duration_sec"}, |
| }}, |
| ] |
| agg = await col.aggregate(pipeline).to_list(1) |
| summary = agg[0] if agg else {} |
| summary.pop("_id", None) |
|
|
| |
| per_client_pipeline = [ |
| {"$group": { |
| "_id": "$client_index", |
| "streams": {"$sum": 1}, |
| "avg_mbps": {"$avg": "$avg_mbps"}, |
| "peak_mbps": {"$max": "$peak_mbps"}, |
| "total_bytes": {"$sum": "$total_bytes"}, |
| }}, |
| {"$sort": {"_id": 1}}, |
| ] |
| per_client = await col.aggregate(per_client_pipeline).to_list(None) |
| for row in per_client: |
| row["client_index"] = row.pop("_id") |
| row["avg_mbps"] = round(row.get("avg_mbps", 0), 3) |
| row["peak_mbps"] = round(row.get("peak_mbps", 0), 3) |
|
|
| |
| recent_cursor = col.find( |
| {}, |
| {"_id": 0, "stream_id": 1, "client_index": 1, "dc_id": 1, |
| "total_bytes": 1, "duration_sec": 1, "avg_mbps": 1, |
| "peak_mbps": 1, "status": 1, "logged_at": 1, "title": 1} |
| ).sort("logged_at", DESCENDING).limit(limit) |
| recent = await recent_cursor.to_list(None) |
| for r in recent: |
| if "logged_at" in r: |
| r["logged_at"] = r["logged_at"].isoformat() |
|
|
| return { |
| "summary": summary, |
| "per_client": per_client, |
| "recent": recent, |
| } |
| except Exception as e: |
| LOGGER.error(f"get_stream_analytics error: {e}") |
| return {"summary": {}, "per_client": [], "recent": []} |
|
|
|
|
|
|
| @staticmethod |
| def _merge_telegram_lists(primary: List[dict], secondary: List[dict]) -> List[dict]: |
| merged = list(primary or []) |
| existing_ids = {q.get("id") for q in merged if q.get("id")} |
| existing_groups = {q.get("group_key") for q in merged if q.get("group_key")} |
| for q in (secondary or []): |
| group_key = q.get("group_key") |
| if group_key and group_key in existing_groups: |
| continue |
| if q.get("id") and q.get("id") in existing_ids: |
| continue |
| merged.append(q) |
| if q.get("id"): |
| existing_ids.add(q.get("id")) |
| if group_key: |
| existing_groups.add(group_key) |
| return merged |
|
|
| def _merge_season_lists(self, primary: List[dict], secondary: List[dict]) -> List[dict]: |
| merged = list(primary or []) |
| season_map = {s.get("season_number"): s for s in merged} |
| for season in (secondary or []): |
| season_number = season.get("season_number") |
| target_season = season_map.get(season_number) |
| if not target_season: |
| merged.append(season) |
| season_map[season_number] = season |
| continue |
|
|
| target_season.setdefault("episodes", []) |
| episode_map = {e.get("episode_number"): e for e in target_season["episodes"]} |
| for episode in season.get("episodes", []): |
| episode_number = episode.get("episode_number") |
| target_episode = episode_map.get(episode_number) |
| if not target_episode: |
| target_season["episodes"].append(episode) |
| episode_map[episode_number] = episode |
| continue |
| target_episode["telegram"] = self._merge_telegram_lists( |
| target_episode.get("telegram", []), episode.get("telegram", []) |
| ) |
| return merged |
|
|
| async def replace_media_metadata( |
| self, |
| media_type: str, |
| tmdb_id: int, |
| db_index: int, |
| metadata: Dict[str, Any] |
| ) -> Optional[dict]: |
| db_key = f"storage_{db_index}" |
| collection_name = self._collection_for(media_type) |
| collection = self.dbs[db_key][collection_name] |
|
|
| current_doc = await collection.find_one({"tmdb_id": int(tmdb_id)}) |
| if not current_doc: |
| return None |
|
|
| source_id = current_doc["_id"] |
|
|
| def _pick(key: str): |
| new_value = metadata.get(key) |
| if new_value in (None, "", [], {}, 0, 0.0, "0"): |
| return current_doc.get(key) |
| return new_value |
|
|
| new_tmdb_id = int(metadata.get("tmdb_id") or tmdb_id) |
| new_imdb_id = _pick("imdb_id") |
|
|
| if collection_name == "movie": |
| current_doc.update({ |
| "tmdb_id": new_tmdb_id, |
| "imdb_id": new_imdb_id, |
| "title": _pick("title"), |
| "release_year": _pick("release_year"), |
| "rating": _pick("rating"), |
| "description": _pick("description"), |
| "poster": _pick("poster"), |
| "backdrop": _pick("backdrop"), |
| "logo": _pick("logo"), |
| "genres": _pick("genres"), |
| "cast": _pick("cast"), |
| "runtime": _pick("runtime"), |
| "media_type": "movie", |
| "telegram": current_doc.get("telegram", []), |
| "updated_on": datetime.utcnow(), |
| }) |
| else: |
| current_doc.update({ |
| "tmdb_id": new_tmdb_id, |
| "imdb_id": new_imdb_id, |
| "title": _pick("title"), |
| "release_year": _pick("release_year"), |
| "rating": _pick("rating"), |
| "description": _pick("description"), |
| "poster": _pick("poster"), |
| "backdrop": _pick("backdrop"), |
| "logo": _pick("logo"), |
| "genres": _pick("genres"), |
| "cast": _pick("cast"), |
| "runtime": _pick("runtime"), |
| "media_type": "tv", |
| "seasons": current_doc.get("seasons", []), |
| "updated_on": datetime.utcnow(), |
| }) |
|
|
| identity_filters = [] |
| if new_imdb_id: |
| identity_filters.append({"imdb_id": new_imdb_id}) |
| identity_filters.append({"tmdb_id": new_tmdb_id}) |
|
|
| existing_other = await collection.find_one({ |
| "$and": [{"$or": identity_filters}, {"_id": {"$ne": source_id}}] |
| }) |
|
|
| if existing_other: |
| if collection_name == "movie": |
| existing_other["telegram"] = self._merge_telegram_lists( |
| existing_other.get("telegram", []), current_doc.get("telegram", []) |
| ) |
| else: |
| existing_other["seasons"] = self._merge_season_lists( |
| existing_other.get("seasons", []), current_doc.get("seasons", []) |
| ) |
|
|
| for field in ( |
| "tmdb_id", "imdb_id", "title", "release_year", "rating", |
| "description", "poster", "backdrop", "logo", "genres", |
| "cast", "runtime", "media_type", |
| ): |
| if field in current_doc: |
| existing_other[field] = current_doc[field] |
| existing_other["updated_on"] = datetime.utcnow() |
|
|
| await collection.delete_one({"_id": source_id}) |
| await collection.replace_one({"_id": existing_other["_id"]}, existing_other) |
|
|
| updated_doc = await collection.find_one({"_id": existing_other["_id"]}) |
| return convert_objectid_to_str(updated_doc) if updated_doc else None |
| await collection.replace_one({"_id": source_id}, current_doc) |
|
|
| updated_doc = await collection.find_one({"_id": source_id}) |
| return convert_objectid_to_str(updated_doc) if updated_doc else None |
|
|
