Create manager.py

database/manager.py

@@ -0,0 +1,123 @@
+# filename: database/manager.py
+
+import logging
+from datetime import datetime, timedelta
+from .connection import db
+
+logger = logging.getLogger(__name__)
+
+# --- Index Creation ---
+async def setup_database_indexes():
+    """Create indexes on collections for faster queries. Should be run on startup."""
+    try:
+        # Indexes for the primary database (cache, groups)
+        await db.primary_db.file_cache.create_index("short_id", unique=True)
+        await db.primary_db.authorized_groups.create_index("group_id", unique=True)
+        
+        # Indexes for all user databases
+        for user_db in db.databases:
+            await user_db.users.create_index("user_id", unique=True)
+            await user_db.users.create_index("is_premium")
+            await user_db.users.create_index("is_banned")
+        
+        logger.info("Database indexes have been configured.")
+    except Exception as e:
+        logger.error(f"Error setting up database indexes: {e}", exc_info=True)
+
+
+# --- User Management ---
+async def add_or_update_user(user_id: int, username: str | None, first_name: str):
+    """Adds a new user or updates the last_seen time for an existing one."""
+    user_db = await db.find_user_db(user_id)
+    
+    if user_db:
+        # User exists, update their info in their assigned DB
+        await user_db.users.update_one(
+            {"user_id": user_id},
+            {"$set": {"username": username, "first_name": first_name, "last_seen": datetime.utcnow()}}
+        )
+    else:
+        # New user, assign them to a DB using round-robin
+        user_db = db.get_next_db_for_new_user()
+        await user_db.users.insert_one({
+            "user_id": user_id,
+            "username": username,
+            "first_name": first_name,
+            "is_premium": False,
+            "premium_expiry_date": None,
+            "is_banned": False,
+            "join_date": datetime.utcnow(),
+            "last_seen": datetime.utcnow()
+        })
+    return await get_user(user_id)
+
+async def get_user(user_id: int) -> dict | None:
+    """Finds a user across all available databases."""
+    for user_db in db.databases:
+        user_doc = await user_db.users.find_one({"user_id": user_id})
+        if user_doc:
+            return user_doc
+    return None
+
+async def grant_premium(user_id: int, duration: timedelta):
+    """Grants or extends premium for a user."""
+    user = await get_user(user_id)
+    if not user: return False
+
+    user_db = await db.get_user_db(user['user_id']) # Get the correct DB for this user
+    
+    # If user is already premium, extend their current subscription. Otherwise, start from now.
+    current_expiry = user.get("premium_expiry_date")
+    start_time = current_expiry if current_expiry and current_expiry > datetime.utcnow() else datetime.utcnow()
+    
+    new_expiry_date = start_time + duration
+    
+    await user_db.users.update_one(
+        {"user_id": user_id},
+        {"$set": {"is_premium": True, "premium_expiry_date": new_expiry_date}}
+    )
+    return True
+
+async def ban_user(user_id: int, ban_status: bool = True):
+    """Bans or unbans a user."""
+    user_db = await db.find_user_db(user_id)
+    if not user_db: return False
+    
+    await user_db.users.update_one({"user_id": user_id}, {"$set": {"is_banned": ban_status}})
+    return True
+
+# --- Cache Management (uses primary DB) ---
+async def get_cached_file(short_id: str) -> dict | None:
+    """Retrieves a cached file's data from the primary database."""
+    return await db.primary_db.file_cache.find_one({"short_id": short_id})
+
+async def add_to_cache(short_id: str, file_id: str, file_name: str, file_size: int):
+    """Adds a new file to the cache in the primary database."""
+    await db.primary_db.file_cache.update_one(
+        {"short_id": short_id},
+        {"$set": {
+            "file_id": file_id,
+            "file_name": file_name,
+            "file_size": file_size,
+            "cached_at": datetime.utcnow()
+        }},
+        upsert=True # Creates the document if it doesn't exist
+    )
+
+# --- Group Management (uses primary DB) ---
+async def is_group_authorized(group_id: int) -> bool:
+    """Checks if a group is authorized."""
+    doc = await db.primary_db.authorized_groups.find_one({"group_id": group_id})
+    return bool(doc)
+
+async def authorize_group(group_id: int, status: bool = True):
+    """Authorizes or de-authorizes a group."""
+    if status:
+        await db.primary_db.authorized_groups.update_one(
+            {"group_id": group_id},
+            {"$set": {"authorized_at": datetime.utcnow()}},
+            upsert=True
+        )
+    else:
+        await db.primary_db.authorized_groups.delete_one({"group_id": group_id})
+    return True