Update bot/database.py

bot/database.py

@@ -1,6 +1,5 @@
 import os
 from motor.motor_asyncio import AsyncIOMotorClient
-from utils.logger import log
 
 class Database:
     def __init__(self):
@@ -10,95 +9,131 @@ class Database:
     async def connect(self):
         uri = os.environ.get("MONGO_URI")
         if not uri:
-            raise ValueError("MONGO_URI environment variable is missing.")
+            raise ValueError("MONGO_URI not set.")
         self.client = AsyncIOMotorClient(uri)
-        self.db = self.client.get_database("sniper_db")
-        log("✅ Connected to MongoDB.")
+        self.db = self.client.username_sniper
 
-    async def load_settings(self, current_state):
-        if self.db is None: return current_state
+    # --- SETTINGS MANAGER ---
+    async def load_settings(self, state_dict: dict):
+        if self.db is None: return
         doc = await self.db.settings.find_one({"_id": "bot_settings"})
-        if doc:
-            # ✨ NEW: Added mute_minor_alerts to the load array
-            keys = [
-                "check_interval", "snipe_interval", "watchdog_poll_interval", 
-                "autoclaim_on", "db_filter_on", "total_checked", "total_available", 
-                "total_claimed", "channel_claim_priority", "preferred_channel_account",
-                "custom_claim_message", "fragment_workers", "fragment_sniping_active", 
-                "fragment_autoclaim_on", "checking_active", "sniping_active", "cycle_paused",
-                "mute_minor_alerts" 
-            ]
-            for key in keys:
-                if key in doc: current_state[key] = doc[key]
-        return current_state
+        if not doc: return
+        
+        # Ensures all custom settings survive a bot restart
+        keys = [
+            "checking_active", "sniping_active", "cycle_paused",
+            "check_interval", "snipe_interval", "autoclaim_on", "db_filter_on",
+            "total_checked", "total_available", "total_claimed",
+            "channel_claim_priority", "preferred_channel_account", 
+            "watchdog_poll_interval", "fragment_sniping_active", 
+            "fragment_autoclaim_on", "fragment_workers", "custom_claim_message",
+            "mute_minor_alerts"
+        ]
+        for key in keys:
+            if key in doc:
+                state_dict[key] = doc[key]
 
     async def save_setting(self, key: str, value):
-        if self.db is None: return
-        await self.db.settings.update_one({"_id": "bot_settings"}, {"$set": {key: value}}, upsert=True)
+        if self.db is not None:
+            await self.db.settings.update_one(
+                {"_id": "bot_settings"},
+                {"$set": {key: value}},
+                upsert=True
+            )
+
+    # --- ADMIN MANAGER ---
+    async def add_admin(self, user_id: int):
+        if self.db is not None:
+            await self.db.admins.update_one(
+                {"user_id": user_id},
+                {"$set": {"user_id": user_id}},
+                upsert=True
+            )
+
+    async def get_admins(self) -> set:
+        if self.db is None: return set()
+        cursor = self.db.admins.find({}, {"user_id": 1})
+        docs = await cursor.to_list(length=None)
+        return {doc["user_id"] for doc in docs}
 
+    # --- UNAVAILABLE SET ---
     async def get_unavailable_set(self) -> set:
         if self.db is None: return set()
-        docs = await self.db.unavailable.find({}, {"_id": 1}).to_list(length=None)
-        return {doc["_id"] for doc in docs}
+        cursor = self.db.unavailable.find({}, {"username": 1})
+        docs = await cursor.to_list(length=None)
+        return {doc["username"] for doc in docs}
 
     async def add_unavailable(self, username: str):
-        if self.db is None: return
-        await self.db.unavailable.update_one({"_id": username}, {"$set": {"_id": username}}, upsert=True)
+        if self.db is not None:
+            await self.db.unavailable.update_one(
+                {"username": username}, 
+                {"$set": {"username": username}}, 
+                upsert=True
+            )
+
+    async def remove_unavailable(self, username: str):
+        if self.db is not None:
+            await self.db.unavailable.delete_one({"username": username})
 
     async def clear_unavailable(self):
-        if self.db is None: return
-        await self.db.unavailable.delete_many({})
+        if self.db is not None:
+            await self.db.unavailable.delete_many({})
 
+    # --- CUSTOM SNIPES ---
     async def get_custom_snipes(self) -> list:
         if self.db is None: return []
-        docs = await self.db.custom_snipes.find({}, {"_id": 1}).to_list(length=None)
-        return [doc["_id"] for doc in docs]
+        cursor = self.db.custom_snipes.find({}, {"username": 1})
+        docs = await cursor.to_list(length=None)
+        return [doc["username"] for doc in docs]
 
     async def add_custom_snipe(self, username: str):
-        if self.db is None: return
-        await self.db.custom_snipes.update_one({"_id": username}, {"$set": {"_id": username}}, upsert=True)
+        if self.db is not None:
+            await self.db.custom_snipes.update_one(
+                {"username": username}, 
+                {"$set": {"username": username}}, 
+                upsert=True
+            )
 
     async def remove_custom_snipe(self, username: str):
-        if self.db is None: return
-        await self.db.custom_snipes.delete_one({"_id": username})
+        if self.db is not None:
+            await self.db.custom_snipes.delete_one({"username": username})
 
+    # --- WATCHDOG ---
     async def get_watchdog(self) -> dict:
         if self.db is None: return {}
-        docs = await self.db.watchdog.find({}).to_list(length=None)
-        return {doc["_id"]: doc.get("submitter_id") for doc in docs}
-
-    async def add_watchdog(self, username: str, submitter_id: int):
-        if self.db is None: return
-        await self.db.watchdog.update_one({"_id": username}, {"$set": {"submitter_id": submitter_id}}, upsert=True)
+        cursor = self.db.watchdog.find({}, {"username": 1, "added_by": 1})
+        docs = await cursor.to_list(length=None)
+        return {doc["username"]: doc.get("added_by") for doc in docs}
+
+    async def add_watchdog(self, username: str, user_id: int):
+        if self.db is not None:
+            await self.db.watchdog.update_one(
+                {"username": username}, 
+                {"$set": {"username": username, "added_by": user_id}}, 
+                upsert=True
+            )
 
     async def remove_watchdog(self, username: str):
-        if self.db is None: return
-        await self.db.watchdog.delete_one({"_id": username})
-
-    async def get_admins(self) -> set:
-        if self.db is None: return set()
-        docs = await self.db.admins.find({}, {"user_id": 1}).to_list(length=None)
-        return {doc["user_id"] for doc in docs}
-
-    async def add_admin(self, user_id: int):
-        if self.db is None: return
-        await self.db.admins.update_one({"user_id": user_id}, {"$set": {"user_id": user_id}}, upsert=True)
-
-    async def remove_admin(self, user_id: int):
-        if self.db is None: return
-        await self.db.admins.delete_one({"user_id": user_id})
+        if self.db is not None:
+            await self.db.watchdog.delete_one({"username": username})
 
-    async def get_fragcheck(self) -> set:
-        if self.db is None: return set()
-        docs = await self.db.fragcheck.find({}).to_list(length=None)
-        return {doc["_id"] for doc in docs}
+    # --- FRAGMENT CHECK ---
+    async def get_fragcheck(self) -> list:
+        if self.db is None: return []
+        cursor = self.db.fragcheck.find({}, {"username": 1})
+        docs = await cursor.to_list(length=None)
+        return [doc["username"] for doc in docs]
 
     async def add_fragcheck(self, username: str):
-        if self.db is None: return
-        await self.db.fragcheck.update_one({"_id": username}, {"$set": {"_id": username}}, upsert=True)
+        if self.db is not None:
+            await self.db.fragcheck.update_one(
+                {"username": username}, 
+                {"$set": {"username": username}}, 
+                upsert=True
+            )
 
     async def remove_fragcheck(self, username: str):
-        if self.db is None: return
-        await self.db.fragcheck.delete_one({"_id": username})
+        if self.db is not None:
+            await self.db.fragcheck.delete_one({"username": username})
 
 db = Database()