khushalps
3c49354c76
git-subtree-dir: utils_v2 git-subtree-split: 83dcddc9c108ac692991d595b7392e5581296e20
359 lines
13 KiB
Python
359 lines
13 KiB
Python
"""
|
|
|
|
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
|