""" AUTHOR: Khushal P Soonderji DATE: Tuesday, 28th May, 2024 OBJECTIVE: To have one central place from where all async database connectivity happens. REFERENCES: 01. https://motor.readthedocs.io/en/stable/ DOWNLOADS: N/A """ # ***************************************************************************************************************** # ***** **** # *** IMPORT *** # ***** **** # ***************************************************************************************************************** # To make sibling directories accessible for imports: # --- import sys sys.path.append(".") sys.path.append("..") # For system-level activity: import io # For async behaviour: import asyncio # MongoDB for File Storage: from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket from bson.objectid import ObjectId # For debugging: from icecream import IceCreamDebugger # ***************************************************************************************************************** # ***** **** # *** MACROS / ONE-TIME INIT *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** VARIABLES *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** FUNCTIONS *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** CLASSES *** # ***** **** # ***************************************************************************************************************** class AsyncMongoStorage: __db = None __db_name = None __client = None __store = None def __init__( self, connection_string = None, max_connections = 5, host_name = "localhost", port = 27017, database_name = "fileStore", debug = True, debug_prefix = "GridFS (M) | ", debug_only_errors = True ): # Database Initialization: self.__host_name = host_name self.__port = port self.__db_name = database_name self.__connection_string = connection_string self.__max_connections = max_connections # Debugging: self.__debug_only_errors = debug_only_errors self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True) if not debug: self.__printer.disable() @staticmethod def generate_id(): """ Just generates an '_id' in MongoDB style. :return: The '_id' in MongoDB style. """ return str(ObjectId()) async def connect(self): """ Initialize the database connection. :return: Nothing. """ if self.__connection_string is None: self.__client = AsyncIOMotorClient( self.__host_name, self.__port, maxPoolSize = self.__max_connections, minPoolSize = self.__max_connections ) else: self.__client = AsyncIOMotorClient( self.__connection_string, maxPoolSize = self.__max_connections, minPoolSize = self.__max_connections ) self.__db = self.__client.get_database(self.__db_name) self.__store = AsyncIOMotorGridFSBucket(self.__db) @property def fs(self): """ To access the features that have not been wrapped in this reportlab directly. This could include things like streaming files chunk-by-chunk. :return: The file-store instance. """ return self.__store async def write_from_memory(self, file_name, file_data, metadata_json = None): """ Save a file (from RAM) to Mongo. Suitable for smaller files (a few MBs max.). :param file_name: The name of the file. :param file_data: The data of the file (held in RAM). :param metadata_json: A JSON of metadata information that can later be used to search files (RECOMMENDED). :return: The file's id as a string (if it gets saved) or None. """ if self.__store is None: await self.connect() file_data.seek(0) file_size = file_data.__sizeof__() file_id = None try: file_id = await self.__store.upload_from_stream(file_name, file_data, metadata = metadata_json) except Exception as exception: self.__printer(exception, file_name, file_size, file_id) if not self.__debug_only_errors: self.__printer(file_name, file_size, file_id) return str(file_id) async def read_to_memory(self, file_identifier, by_id = True): """ To retrieve a file (in RAM) based on the provided identifier. Suitable for smaller files (a few MBs max.). :param file_identifier: Either the name or the "_id" of the file. :param by_id: Set to True if you are fetching by the "_id" of the file. :return: Either the file (in RAM) or None. """ if self.__store is None: await self.connect() file_data = None try: if by_id: grid_out = await self.__store.open_download_stream(ObjectId(file_identifier)) else: grid_out = await self.__store.open_download_stream_by_name(file_identifier) file_data = io.BytesIO(await grid_out.read()) file_data.seek(0) except Exception as exception: file_data = None self.__printer(exception, file_identifier, by_id) if not self.__debug_only_errors: self.__printer(file_identifier, by_id) return file_data async def delete_file_by_id(self, file_id): """ Tries to delete one file by the id. :param file_id: The id of the file in the database. :return: True or False based on the success of the operation. """ if self.__store is None: await self.connect() deleted = False try: response = await self.__store.delete(file_id = ObjectId(file_id)) deleted = True except Exception as exception: self.__printer(exception, file_id, deleted) return deleted def __format_metadata_json(self, metadata_json): """ NOTE: ONLY USE WHEN SEARCHING FILES BY METADATA. MongoDB expects dot-notation while searching for files by the metadata. We are making a function to search files assuming that the conditions are to be applied to the metadata itself. So this function add the dot-notation to the right places to conduct a successful search. :param metadata_json: The JSON to format. :return: The formatted JSON that has the right dot-notation. """ formatted_metadata_json = {} for key, value in metadata_json.items(): if not key.startswith("$"): key = f"metadata.{key}" else: if type(value) is dict: value = self.__format_metadata_json(value) if type(value) is list: value = [self.__format_metadata_json(item) for item in value] formatted_metadata_json[key] = value return formatted_metadata_json async def find_file_by_metadata(self, metadata_json, limit = None, skip = None, sort = None): """ This method only lists the files that match the criteria mentioned in the metadata JSON. :param metadata_json: The JSON that describes what you want to find. :param limit: Max. no. of records to retrieve. :param skip: No. of starting results to skip. Useful for pagination. :param sort: The sorting conditions to follow. :return: A list of (JSONs of) files that match the conditions. The list can be empty. """ if self.__store is None: await self.connect() files_list = [] try: limit = limit or 10 skip = skip or 0 sort = {"_id": -1} if not isinstance(sort, dict) else sort formatted_metadata_json = self.__format_metadata_json(metadata_json) return await self.__store.find( formatted_metadata_json ).sort(sort).skip(skip).limit(limit).to_list(None) except Exception as exception: self.__printer(exception, metadata_json, len(files_list)) if not self.__debug_only_errors: self.__printer(metadata_json, len(files_list)) return files_list async def find_file_by_id(self, file_id): """ This method allows you to get the file's info from the id of the file. :param file_id: The id that was assigned by Mongo during upload. :return: The file's info or None if the file doesn't exist. """ if self.__store is None: await self.connect() file_info = None try: formatted_metadata_json = {"_id": ObjectId(file_id)} file_info = (await self.__store.find(formatted_metadata_json).to_list(1))[0] except Exception as exception: self.__printer(exception, file_id, file_info) return file_info async def get_file_name(self, file_id): """ Returns the file name if the id of the file is known. :param file_id: The id of the file as assigned by MongoDB when the file was stored. :return: The file's name (if it exists), or None. """ # Ensure that we are connected: if self.__store is None: await self.connect() # Ensure that the input given is of 'ObjectId' type: if type(file_id) is not ObjectId: file_id = ObjectId(str(file_id)) # Fetch and return the file name: files_list = await self.__store.find( {"_id": file_id}, {"filename": True} ).sort({"_id": -1}).limit(1).to_list(None) try: file_name = files_list[0]["filename"] except: file_name = None return file_name async def aggregate( self, collection, pipeline, limit = None ): """ Perform an advance query on the data. :param collection: The collection to perform the query on. :param pipeline: The pipeline of actions to take. Must be a list. :param limit: The max. no. of records to retrieve. :return: The array of matching records or null if there was an exception. """ if self.__store is None: await self.connect() results = None try: results = await self.__db[collection].aggregate(pipeline).to_list(limit) except Exception as exception: self.__printer(exception) return results # ***************************************************************************************************************** # ***** **** # *** MAIN PROGRAM *** # ***** **** # ***************************************************************************************************************** if __name__ == "__main__": pass