GoldyBot.database.database

  1from __future__ import annotations
  2import asyncio
  3from typing import Any, Dict, List
  4import pymongo
  5import motor.motor_asyncio
  6import GoldyBot
  7
  8MODULE_NAME = "DATABASE"
  9
 10class Database():
 11    """Goldy Bot's class to interface with a Mongo Database."""
 12    def __init__(self, database_token_url:str):
 13        self.database_token_url = database_token_url
 14        self.loop = asyncio.get_event_loop()
 15
 16        config = GoldyBot.config.Config(GoldyBot.files.File(GoldyBot.paths.GOLDY_CONFIG_JSON))
 17
 18        try:
 19            if not database_token_url.lower() in ["none", "null"]: 
 20                # Initializing MongoDB database
 21                self.client:pymongo.MongoClient = motor.motor_asyncio.AsyncIOMotorClient(self.database_token_url, serverSelectionTimeoutMS=2000)
 22                self.loop.run_until_complete(self.client.server_info())
 23                self.database = self.client[config.read("database_name")]
 24
 25                # Cache database class.
 26                GoldyBot.cache.main_cache_dict["database"] = self
 27
 28                GoldyBot.logging.log("info_2", f"[{MODULE_NAME}] [CONNECTED]")
 29            
 30            else:
 31                GoldyBot.logging.log("info_5", f"[{MODULE_NAME}] [DATABASE DISABLED]")
 32                GoldyBot.logging.log("info_3", f"[{MODULE_NAME}] Database is disabled because '{database_token_url}' was entered.")
 33
 34        except Exception:
 35            GoldyBot.logging.log("error", f"[{MODULE_NAME}] Couldn't connect to Database! Check the database URL you entered.")
 36            GoldyBot.Goldy().stop("Couldn't connect to Database! Check the database URL you entered.")
 37
 38    async def insert(self, collection:str, data) -> bool:
 39        """Tells database to insert the data provided into a collection."""
 40        await self.database[collection].insert_one(data)
 41        GoldyBot.logging.log(f"[{MODULE_NAME}] Inserted '{data}' into '{collection}.'")
 42        return True
 43
 44    async def edit(self, collection:str, query, data:dict) -> bool:
 45        """Tells database to find and edit a document with the data provided in a collection."""
 46        await self.database[collection].update_one(query, {"$set": data})
 47        GoldyBot.logging.log(f"[{MODULE_NAME}] Edited '{query}' into '{data}.'")
 48        return True
 49
 50    async def remove(self, collection:str, data) -> bool:
 51        """Tells database to find and delete a copy of this data from the collection."""
 52        await self.database[collection].delete_one(data)
 53        GoldyBot.logging.log("INFO_5", f"[{MODULE_NAME}] Deleted '{data}' from '{collection}.'")
 54        return True
 55
 56    async def find(self, collection:str, query, key:str, max_to_find=50) -> List[dict]:
 57        """Searches for documents with the query."""
 58        try:
 59            document_list = []
 60            cursor = self.database[collection].find(query).sort(key)
 61
 62            for document in await cursor.to_list(max_to_find):
 63                document_list.append(document)
 64
 65            return document_list
 66        except KeyError as e:
 67            GoldyBot.logging.log("error", f"[{MODULE_NAME}] Could not find the collection '{collection}'!")
 68            return None
 69
 70    async def find_all(self, collection:str, max_to_find=100) -> List[dict] | None:
 71        """Finds and returns all documents in a collection. This took me a day to make! 😞"""
 72        try:
 73            document_list = []
 74            cursor = self.database[collection].find().sort('_id')
 75
 76            for document in await cursor.to_list(max_to_find):
 77                document_list.append(document)
 78
 79            return document_list
 80        except KeyError as e:
 81            GoldyBot.logging.log("error", f"[{MODULE_NAME}] Could not find the collection '{collection}'!")
 82            return None
 83
 84    async def get_collection(self, collection):
 85        """Returns cursor of the following collection."""
 86        return self.database[collection]
 87
 88    async def list_collection_names(self) -> List[str]:
 89        """Returns list of all collection names."""
 90        return await self.database.list_collection_names()
 91
 92    async def find_one(self, collection:str, query:dict) -> (dict | None):
 93        """Tells database to search for and return specific data from a collection."""
 94        data = await self.database[collection].find_one(query)
 95
 96        if not data == None:
 97            GoldyBot.logging.log(f"[{MODULE_NAME}] Found '{query}' in '{collection}.'")
 98            return data
 99        else:
100            GoldyBot.logging.log(f"[{MODULE_NAME}] '{query}' was not found in '{collection}.'")
101            return None
102        
103    async def create_collection(self, collection_name:str, data):
104        await self.database[collection_name].insert_one(data)
105        GoldyBot.logging.log(f"[{MODULE_NAME}] Database collection '{collection_name}' created.")
106
107    async def delete_collection(self, collection_name:str):
108        await self.database[collection_name].drop()
109        GoldyBot.logging.log("INFO_5", f"[{MODULE_NAME}] Database collection '{collection_name}' dropped.")
110
111    def new_instance(self, database_name:str):
112        """Starts a new database instance the efficient way. 👍"""
113        class NewDatabase(Database):
114            def __init__(self, database_self:Database):
115                self.database = database_self.client[database_name]
116
117        return NewDatabase(self)