1754 lines
59 KiB
Python
1754 lines
59 KiB
Python
"""
|
|
|
|
AUTHOR:
|
|
|
|
Khushal P Soonderji
|
|
|
|
DATE:
|
|
|
|
Original: Tuesday, 28th May, 2024
|
|
Modified: Tuesday, 17th Sept., 2024
|
|
|
|
OBJECTIVE:
|
|
|
|
To have one central place from where all async MongoDB activity 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
|
|
from bson.json_util import dumps, loads
|
|
|
|
# My utils:
|
|
from utils_v2.string import json
|
|
|
|
# For datetime handling:
|
|
import pytz
|
|
import datetime
|
|
|
|
# For debugging:
|
|
from icecream import IceCreamDebugger
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MACROS / ONE-TIME INIT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** VARIABLES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** EXCEPTIONS ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** FUNCTIONS ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** CLASSES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
class AsyncMongoBase:
|
|
|
|
def __init__(
|
|
self,
|
|
connection_string = None,
|
|
max_connections = 5,
|
|
host_name = "localhost",
|
|
port = 27017,
|
|
database_name = "myDb",
|
|
debug = True,
|
|
debug_prefix = "Mongo | ",
|
|
debug_only_errors = True
|
|
):
|
|
|
|
# Basic variables that will be needed later:
|
|
self._client = None
|
|
self._db = None
|
|
self._fs = None
|
|
|
|
# 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()
|
|
|
|
def enable_debug(self):
|
|
self._printer.enable()
|
|
|
|
def disable_debug(self):
|
|
self._printer.disable()
|
|
|
|
async def connect(self) -> bool:
|
|
|
|
"""
|
|
Initialize the database connection.
|
|
:return: None.
|
|
"""
|
|
|
|
# Start by assuming success:
|
|
success = True
|
|
|
|
try:
|
|
|
|
self._printer("Connecting...")
|
|
|
|
# If a connection string is given,
|
|
# give preference to it:
|
|
if self._connection_string is None:
|
|
self._client = AsyncIOMotorClient(
|
|
self._host_name,
|
|
self._port,
|
|
maxPoolSize = self._max_connections,
|
|
minPoolSize = self._max_connections
|
|
)
|
|
|
|
# In the absense of a connection string,
|
|
# try to connect through the other credentials:
|
|
else:
|
|
self._client = AsyncIOMotorClient(
|
|
self._connection_string,
|
|
maxPoolSize = self._max_connections,
|
|
minPoolSize = self._max_connections
|
|
)
|
|
|
|
# Debugging print:
|
|
if not self._debug_only_errors:
|
|
server_info = await self._client.server_info()
|
|
self._printer(server_info)
|
|
|
|
# Now we connect to the database and the file-store:
|
|
self._db = self._client.get_database(self._db_name)
|
|
self._fs = AsyncIOMotorGridFSBucket(self._db)
|
|
|
|
self._printer("Connected!")
|
|
|
|
# If something goes wrong:
|
|
except Exception as exception:
|
|
self._printer("NOT Connected")
|
|
self._printer(exception)
|
|
success = False
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def ensure_connection(self):
|
|
|
|
"""
|
|
Call this at the start of every function to automatically connect to the database when the user of the library
|
|
forgets to explicitly connect to the database.
|
|
:return: None.
|
|
"""
|
|
|
|
if self._client is None: await self.connect()
|
|
|
|
@staticmethod
|
|
def generate_id(as_str = False):
|
|
|
|
"""
|
|
Just generates an '_id' in MongoDB style.
|
|
:param as_str: Set to True to convert the generated id to a string; and False to receive it as an instance of
|
|
'ObjectId'. This is useful when you need to pre-assign ids to files.
|
|
:return: The '_id' in MongoDB style.
|
|
"""
|
|
|
|
generated_id = ObjectId()
|
|
return str(generated_id) if as_str else generated_id
|
|
|
|
@property
|
|
async def client(self):
|
|
|
|
"""
|
|
returns the client to perform activities that have not been explicitly wrapped in the class.
|
|
:return: The client's instance.
|
|
"""
|
|
|
|
await self.ensure_connection()
|
|
return self._client
|
|
|
|
@property
|
|
async def db(self):
|
|
|
|
"""
|
|
returns the database connector to perform activities that have not been explicitly wrapped in the class.
|
|
:return: The database's connection instance.
|
|
"""
|
|
|
|
await self.ensure_connection()
|
|
return self._fs
|
|
|
|
@property
|
|
async def fs(self):
|
|
|
|
"""
|
|
returns the file-store to perform activities that have not been explicitly wrapped in the class.
|
|
:return: The file-store instance.
|
|
"""
|
|
|
|
await self.ensure_connection()
|
|
return self._fs
|
|
|
|
@staticmethod
|
|
def from_json_string(json_data):
|
|
|
|
"""
|
|
Converts from a JSON string to BSON.
|
|
:param json_data: The JSON string to convert to BSON.
|
|
:return: The BSON interpretation of the input JSON string.
|
|
"""
|
|
|
|
return loads(json_data)
|
|
|
|
@staticmethod
|
|
def to_json_string(data, indent = 4, default = None):
|
|
|
|
"""
|
|
Converts from a BSON to JSON string.
|
|
:param data: The input BSON data.
|
|
:param indent: The no. of spaces to put into the string for pretty print.
|
|
:param default: The default function to apply to data that cannot be converted directly.
|
|
:return: The JSON string from the input data.
|
|
"""
|
|
|
|
return dumps(data, indent = indent, default = default)
|
|
|
|
@staticmethod
|
|
def dict_to_dot_notation(input_dict, pk = "", s = "."):
|
|
|
|
"""
|
|
Converts an input dict to dot notation format. Can be used as a utility to perform searches.
|
|
:param input_dict: The dict that you want to convert to dot notation.
|
|
:param pk: Parent Key. DO NOT TOUCH (meant to be used during recursion).
|
|
:param s: Separator. DO NOT TOUCH (meant to be used during recursion).
|
|
:return: The dot notation representation of the input dict.
|
|
"""
|
|
|
|
items = []
|
|
for k, v in input_dict.items():
|
|
new_key = f"{pk}{s}{k}" if pk else k
|
|
if isinstance(v, dict) and v: items.extend(AsyncMongoBase.dict_to_dot_notation(v, new_key, s = s).items())
|
|
else: items.append((new_key, v))
|
|
return dict(items)
|
|
|
|
@staticmethod
|
|
def normalize_date_time(document):
|
|
|
|
"""
|
|
MongoDB doesn't support timezones. A good strategy would be to convert everything to UTC format and store it.
|
|
This method does exactly that. Any datetime object is converted to UTC timezone. If the datetime object was
|
|
timezone naive, UTC timezone will be applied to it without changing the time value.
|
|
:param document: The document that you want to normalize the date-time in.
|
|
:return: The document with normalized datetime.
|
|
"""
|
|
|
|
if isinstance(document, datetime.datetime):
|
|
utc_tz = pytz.timezone("UTC")
|
|
if document.tzinfo is None: document = utc_tz.localize(document)
|
|
else: document = document.astimezone(utc_tz)
|
|
|
|
if type(document) is list:
|
|
document = [AsyncMongoBase.normalize_date_time(item) for item in document]
|
|
|
|
if type(document) is dict:
|
|
document = {
|
|
AsyncMongoBase.normalize_date_time(k): AsyncMongoBase.normalize_date_time(v)
|
|
for k, v in document.items()
|
|
}
|
|
|
|
return document
|
|
|
|
@staticmethod
|
|
def read_to_ram(file_path):
|
|
|
|
"""
|
|
Reads a file into a BytesIO object in RAM.
|
|
:param file_path: The path to the file on disk.
|
|
:return: The file in a BytesIO object.
|
|
"""
|
|
|
|
with open(file_path, "rb") as file: file_data = file.read()
|
|
file_in_ram = io.BytesIO(file_data)
|
|
file_in_ram.seek(0)
|
|
return file_in_ram
|
|
|
|
|
|
# ---------------------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
class AsyncMongo(AsyncMongoBase):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
async def list_indexes(
|
|
self,
|
|
collection,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Lists out the indexes of a collection.
|
|
:param collection: The collection whose indexes you want to list out,
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The list of indexes or None if the action fails.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
indexes = None
|
|
|
|
# Try to list the indexes:
|
|
try:
|
|
responses = await self._db[collection].list_indexes(session = session).to_list(None)
|
|
indexes = [{key: value for key, value in response.items()} for response in responses]
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return indexes
|
|
|
|
async def create_index(
|
|
self,
|
|
collection,
|
|
keys,
|
|
options = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Creates an index on a collection.
|
|
:param collection: The collection to create the index on.
|
|
:param keys: The keys (and sorting) to implement the index on.
|
|
:param options: Additional config.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the execution.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# Try to make the insertion:
|
|
try:
|
|
options = options or {}
|
|
keys = [(k, v) for k, v in keys.items()]
|
|
response = await self._db[collection].create_index(keys, session = session, **options)
|
|
if response: success = True
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def count(
|
|
self,
|
|
collection,
|
|
filter,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Counts the no. of documents that match the given filter condition.
|
|
:param collection: The name of the collection to count in.
|
|
:param filter: The filter criteria that the documents must satisfy to be counted.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The count if the operation was performed successfully or None is something went wrong and the exception
|
|
was suppressed.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
count = None
|
|
|
|
# try to query the data:
|
|
try: count = await self._db[collection].count_documents(filter, session = session)
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return count
|
|
|
|
async def insert_one(
|
|
self,
|
|
collection,
|
|
document,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Insert data into a collection.
|
|
:param collection: The collection you want to feed the data into.
|
|
:param document: The data to be stored.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The id of the inserted data, or null if the action fails.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
inserted_id = None
|
|
|
|
# Try to make the insertion:
|
|
try:
|
|
response = await self._db[collection].insert_one(document.copy(), session = session)
|
|
inserted_id = response.inserted_id
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return inserted_id
|
|
|
|
async def insert_many(
|
|
self,
|
|
collection,
|
|
documents,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Insert a lot of data into a collection.
|
|
:param collection: The collection you want to feed the data into.
|
|
:param documents: The list of documents to be stored.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The id of the inserted data, or null if the action fails.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
inserted_ids = []
|
|
|
|
# Try to make the insertion:
|
|
try:
|
|
response = await self._db[collection].insert_many(documents, session = session)
|
|
inserted_ids = response.inserted_ids
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return inserted_ids
|
|
|
|
async def find_one(
|
|
self,
|
|
collection,
|
|
filter,
|
|
projection = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Finds one record that matches the given conditions.
|
|
:param collection: The name of the collection to perform the search in.
|
|
:param filter: The filter criteria.
|
|
:param projection: What parts of the matching data you want to fetch.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The matching record or null if there was an exception.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
results = None
|
|
|
|
# try to query the data:
|
|
try: results = await self._db[collection].find_one(filter, projection, session = session)
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Check results and return:
|
|
return results
|
|
|
|
async def find_many(
|
|
self,
|
|
collection,
|
|
filter,
|
|
projection = None,
|
|
skip = 0,
|
|
limit = None,
|
|
sort = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Finds one or more records that match the given conditions.
|
|
:param collection: The name of the collection to perform the search in.
|
|
:param filter: The filter criteria.
|
|
:param projection: What parts of the matching data you want to fetch.
|
|
:param skip: The no. of records to skip before picking next ones. Needed for pagination.
|
|
:param limit: The max. no. of records you want to fetch.
|
|
:param sort: The sorting rules to apply.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The array of matching records or null if there was an exception.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
results = None
|
|
|
|
# Assume defaults:
|
|
if sort is None: sort = {"_id": -1}
|
|
if limit is None: limit = 10
|
|
|
|
# try to query the data:
|
|
try:
|
|
results = await self._db[collection].find(
|
|
filter,
|
|
projection,
|
|
sort = sort,
|
|
skip = skip,
|
|
limit = limit,
|
|
session = session
|
|
).to_list(None)
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return results
|
|
|
|
async def update_one(
|
|
self,
|
|
collection,
|
|
filter = None,
|
|
update = None,
|
|
upsert = False,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Update one document.
|
|
:param collection: The collection you want to update.
|
|
:param filter: The selection criteria to locate the document to update.
|
|
:param update: The values you want to update.
|
|
:param upsert: If you want to insert if the document doesn't already exist.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# Try to make the update:
|
|
try:
|
|
response = await self._db[collection].update_one(
|
|
filter,
|
|
update,
|
|
upsert = upsert,
|
|
session = session
|
|
)
|
|
success = True if response.modified_count > 0 or response.upserted_id else False
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def update_many(
|
|
self,
|
|
collection,
|
|
filter = None,
|
|
update = None,
|
|
upsert = False,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Update many documents.
|
|
:param collection: The collection you want to update.
|
|
:param filter: The selection criteria to locate the document to update.
|
|
:param update: The values you want to update.
|
|
:param upsert: If you want to insert if the document doesn't already exist.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The no. of records that were affected.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
update_count = 0
|
|
|
|
# Try to make the insertion:
|
|
try:
|
|
response = await self._db[collection].update_many(
|
|
filter,
|
|
update,
|
|
upsert = upsert,
|
|
session = session
|
|
)
|
|
update_count = response.modified_count
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return update_count
|
|
|
|
async def find_one_and_update(
|
|
self,
|
|
collection,
|
|
filter = None,
|
|
update = None,
|
|
projection = None,
|
|
return_updated = False,
|
|
upsert = False,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Update one document.
|
|
:param collection: The collection you want to update.
|
|
:param filter: The selection criteria to locate the document to update.
|
|
:param update: The values you want to update.
|
|
:param projection: What parts of the matching data you want to fetch.
|
|
:param return_updated: To choose whether you want to retrieve the original document or the updated document.
|
|
:param upsert: If you want to insert if the document doesn't already exist.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The fetched document.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
document = None
|
|
|
|
# Try to make the update:
|
|
try:
|
|
document = await self._db[collection].find_one_and_update(
|
|
filter,
|
|
update,
|
|
projection = projection,
|
|
return_document = return_updated,
|
|
upsert = upsert,
|
|
session = session
|
|
)
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return document
|
|
|
|
async def replace_one(
|
|
self,
|
|
collection,
|
|
filter,
|
|
replacement,
|
|
upsert = False,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
To delete one document from a collection.
|
|
:param collection: The collection from which you want to delete many records.
|
|
:param filter: The filter criteria.
|
|
:param replacement: The data to put in place of the existing document.
|
|
:param upsert: If you want to insert if the document doesn't already exist.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# try to query the data:
|
|
try:
|
|
result = await self._db[collection].replace_one(
|
|
filter,
|
|
replacement,
|
|
upsert = upsert,
|
|
session = session
|
|
)
|
|
if result.modified_count or result.upserted_id: success = True
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def find_one_and_replace(
|
|
self,
|
|
collection,
|
|
filter,
|
|
replacement,
|
|
projection = None,
|
|
return_replaced = False,
|
|
upsert = False,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
To delete one document from a collection.
|
|
:param collection: The collection from which you want to delete many records.
|
|
:param filter: The filter criteria.
|
|
:param replacement: The data to put in place of the existing document.
|
|
:param projection: What parts of the matching data you want to fetch.
|
|
:param return_replaced: To choose whether you want to retrieve the original document or the updated document.
|
|
:param upsert: If you want to insert if the document doesn't already exist.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
document = None
|
|
|
|
# try to query the data:
|
|
try:
|
|
document = await self._db[collection].find_one_and_replace(
|
|
filter,
|
|
replacement,
|
|
projection = projection,
|
|
return_document = return_replaced,
|
|
upsert = upsert,
|
|
session = session
|
|
)
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return document
|
|
|
|
async def delete_one(
|
|
self,
|
|
collection,
|
|
filter,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
To delete one document from a collection.
|
|
:param collection: The collection from which you want to delete many records.
|
|
:param filter: The filter criteria.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# try to query the data:
|
|
try:
|
|
result = await self._db[collection].delete_one(filter, session = session)
|
|
success = True if result.deleted_count else False
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def delete_many(
|
|
self,
|
|
collection,
|
|
filter,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
To delete many documents from a collection.
|
|
WARNING: sending {} in the filter would mean deleting ALL the documents.
|
|
:param collection: The collection from which you want to delete many records.
|
|
:param filter: The filter criteria.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
count = 0
|
|
|
|
# try to query the data:
|
|
try:
|
|
result = await self._db[collection].delete_many(filter, session = session)
|
|
count = result.deleted_count
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return count
|
|
|
|
async def find_one_and_delete(
|
|
self,
|
|
collection,
|
|
filter,
|
|
projection = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
To delete one document from a collection.
|
|
:param collection: The collection from which you want to delete a record.
|
|
:param filter: The filter criteria.
|
|
:param projection: What parts of the matching data you want to fetch.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The document that matched your criteria or None.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
document = None
|
|
|
|
# try to query the data:
|
|
try:
|
|
document = await self._db[collection].find_one_and_delete(
|
|
filter,
|
|
projection = projection,
|
|
session = session
|
|
)
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return document
|
|
|
|
async def bulk_write(
|
|
self,
|
|
collection,
|
|
requests,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
To perform various individual operations in one go. You will have to import individual actions like "UpdateOne"
|
|
and "InsertMany" from PyMongo and pass them as an array of requests (operations) to this method.
|
|
:param collection: The collection you want to run the requests on.
|
|
:param requests: The array of requests (operations) to be performed.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The no of operations done.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
count = 0
|
|
|
|
# try to query the data:
|
|
try:
|
|
response = await self._db[collection].bulk_write(requests, session = session)
|
|
count = response.modified_count + response.inserted_count + response.upserted_count + response.deleted_count
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return count
|
|
|
|
async def aggregate(
|
|
self,
|
|
collection,
|
|
pipeline,
|
|
limit = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
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. preferably apply the limit from within the pipeline.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The array of matching records or null if there was an exception.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
results = None
|
|
|
|
# try to perform the aggregation action:
|
|
try: results = await self._db[collection].aggregate(pipeline, session = session).to_list(limit)
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return results
|
|
|
|
|
|
# ---------------------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
class AsyncMongoStorage(AsyncMongo):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
async def easy_upload(
|
|
self,
|
|
source,
|
|
file_name: str,
|
|
file_metadata: dict = None,
|
|
file_id = None,
|
|
chunk_size: int = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Easily write one file to MongoDB's GridFS. Ideal for directly uploading small files without having to worry
|
|
about any internal mechanisms.
|
|
:param source: The actual file, supplied as either a path string or a file-like object, to be written to the
|
|
database.
|
|
:param file_name: The name of the file as it will be stored on (and retrieved from) GridFS.
|
|
:param file_metadata: Any metadata to later search the file by.
|
|
:param file_id: Any custom id to be given to the file. TRY STICKING TO THE ID GENERATED BY 'generate_id'.
|
|
:param chunk_size: The chunk size (in bytes) to use for storing this file.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The id of the inserted file or None if the upload failed.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
inserted_id = None
|
|
|
|
# Try to write the file:
|
|
try:
|
|
|
|
# Input pre-processing:
|
|
if isinstance(source, io.BytesIO): source.seek(0)
|
|
elif isinstance(source, str): source = open(source, mode = "rb")
|
|
|
|
# If no file id was supplied:
|
|
if file_id is None:
|
|
inserted_id = await self._fs.upload_from_stream(
|
|
filename = file_name,
|
|
source = source,
|
|
metadata = file_metadata,
|
|
chunk_size_bytes = chunk_size,
|
|
session = session
|
|
)
|
|
|
|
# If a file id was supplied:
|
|
else:
|
|
await self._fs.upload_from_stream_with_id(
|
|
file_id = file_id,
|
|
filename = file_name,
|
|
source = source,
|
|
metadata = file_metadata,
|
|
chunk_size_bytes = chunk_size,
|
|
session = session
|
|
)
|
|
inserted_id = file_id
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return inserted_id
|
|
|
|
async def easy_download(
|
|
self,
|
|
destination,
|
|
file_id = None,
|
|
file_name = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Easily read one file from MongoDB's GridFS. Ideal for directly downloading small files without having to worry
|
|
about any internal mechanisms.
|
|
:param destination: The path on the local disk or a buffer in RAM to save the downloaded data to.
|
|
:param file_id: (RECOMMENDED) the id of the save file.
|
|
:param file_name: The name of the saved file. NOT RECOMMENDED because you could have many files with the same
|
|
name. The best way to tell files apart if from the id.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation. The contents of the stored file are written
|
|
directly to the destination.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# Try to read the file:
|
|
try:
|
|
|
|
# Input pre-processing:
|
|
if isinstance(destination, str): destination = open(destination, mode = "wb")
|
|
|
|
# If a file id is supplied (preferred way):
|
|
if file_id is not None:
|
|
await self._fs.download_to_stream(
|
|
destination = destination,
|
|
file_id = file_id,
|
|
session = session
|
|
)
|
|
if isinstance(destination, io.BytesIO): destination.seek(0)
|
|
success = True
|
|
|
|
# If a file name is supplied:
|
|
elif file_name is not None:
|
|
await self._fs.download_to_stream_by_name(
|
|
destination = destination,
|
|
filename = file_name,
|
|
session = session
|
|
)
|
|
if isinstance(destination, io.BytesIO): destination.seek(0)
|
|
success = True
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def count_files(
|
|
self,
|
|
filter,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Counts the no. of documents that match the given filter condition.
|
|
:param filter: The filter criteria that the documents must satisfy to be counted.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The count if the operation was performed successfully or None is something went wrong and the exception
|
|
was suppressed.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
count = None
|
|
|
|
# try to query the data:
|
|
try: count = await self._db["fs.files"].count_documents(filter, session = session)
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return count
|
|
|
|
async def find_one_file(
|
|
self,
|
|
filter,
|
|
projection = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Finds one record that matches the given conditions. This does NOT return the file itself, it returns the record
|
|
that describes the file.
|
|
:param filter: The filter criteria.
|
|
:param projection: What parts of the matching data you want to fetch.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The matching record or null if there was an exception.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
results = None
|
|
|
|
# try to query the data:
|
|
try: results = await self._db["fs.files"].find_one(filter, projection, session = session)
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Check results and return:
|
|
return results
|
|
|
|
async def find_many_files(
|
|
self,
|
|
filter,
|
|
projection = None,
|
|
skip = 0,
|
|
limit = None,
|
|
sort = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Finds one or more records that match the given conditions. This doesn't return any actual files directly, it
|
|
returns the records that describe the files.
|
|
:param filter: The filter criteria.
|
|
:param projection: What parts of the matching data you want to fetch.
|
|
:param skip: The no. of records to skip before picking next ones. Needed for pagination.
|
|
:param limit: The max. no. of records you want to fetch.
|
|
:param sort: The sorting rules to apply.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The array of matching records or null if there was an exception.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
results = None
|
|
|
|
# Assume defaults:
|
|
if sort is None: sort = {"_id": -1}
|
|
if limit is None: limit = 10
|
|
|
|
# try to query the data:
|
|
try:
|
|
results = await self._db["fs.files"].find(
|
|
filter,
|
|
projection,
|
|
session = session
|
|
).sort(sort).skip(skip).limit(limit).to_list(None)
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return results
|
|
|
|
async def update_metadata_for_one_file(
|
|
self,
|
|
filter: dict,
|
|
unset_data: dict = None,
|
|
set_data: dict = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Updates the metadata for one file.
|
|
:param filter: The conditions to filter the files by.
|
|
:param unset_data: The fields that you want to discard.
|
|
:param set_data: The fields that you want to add or update.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# Try to make the update:
|
|
try:
|
|
|
|
# Input pre-processing:
|
|
update = {}
|
|
if unset_data is not None: update["$unset"] = self.dict_to_dot_notation({"metadata": unset_data})
|
|
if set_data is not None: update["$set"] = self.dict_to_dot_notation({"metadata": set_data})
|
|
|
|
# Actual update happens here:
|
|
if update:
|
|
response = await self._db["fs.files"].update_one(
|
|
filter,
|
|
update,
|
|
session = session
|
|
)
|
|
success = False if response.modified_count == 0 else True
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def update_metadata_for_many_files(
|
|
self,
|
|
filter: dict,
|
|
unset_data: dict = None,
|
|
set_data: dict = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Updates the metadata for many files. Practically the same as 'update_metadata_for_one' except that the scope of
|
|
the modifications is far wider.
|
|
:param filter: The conditions to filter the files by.
|
|
:param unset_data: The fields that you want to discard.
|
|
:param set_data: The fields that you want to add or update.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
count = None
|
|
|
|
# Try to make the update:
|
|
try:
|
|
|
|
# Input pre-processing:
|
|
update = {}
|
|
if unset_data is not None: update["$unset"] = self.dict_to_dot_notation({"metadata": unset_data})
|
|
if set_data is not None: update["$set"] = self.dict_to_dot_notation({"metadata": set_data})
|
|
|
|
# Actual update happens here:
|
|
if update:
|
|
response = await self._db["fs.files"].update_many(
|
|
filter,
|
|
update,
|
|
session = session
|
|
)
|
|
count = response.modified_count
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return count
|
|
|
|
async def replace_metadata_for_one_file(
|
|
self,
|
|
filter: dict,
|
|
replacement,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Updates the metadata for one file.
|
|
:param filter: The conditions to filter the files by.
|
|
:param replacement: The new metadata to put inplace of the old one.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# Try to make the update:
|
|
try:
|
|
|
|
# Actual update happens here:
|
|
result = await self._db["fs.files"].update_one(
|
|
filter,
|
|
{"$set": {"metadata": replacement}},
|
|
upsert = False,
|
|
session = session
|
|
)
|
|
if result.modified_count or result.upserted_id: success = True
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def replace_metadata_for_many_files(
|
|
self,
|
|
filter: dict,
|
|
replacement,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Updates the metadata for one file.
|
|
:param filter: The conditions to filter the files by.
|
|
:param replacement: The new metadata to put inplace of the old one.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
count = None
|
|
|
|
# Try to make the update:
|
|
try:
|
|
|
|
# Actual update happens here:
|
|
response = await self._db["fs.files"].update_many(
|
|
filter,
|
|
{"$set": {"metadata": replacement}},
|
|
session = session
|
|
)
|
|
count = response.modified_count
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return count
|
|
|
|
async def delete_file_by_id(
|
|
self,
|
|
file_id,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Deletes one file by the file's id. Deleting files is different from deleting simple documents because, in files,
|
|
you also need to clear out the chunks (which are stored in a separate collection). The built-in mechanism of
|
|
Motor only provides support to delete by the file's id.
|
|
:param file_id: The id of the file that you want to delete.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# Try to read the file:
|
|
try:
|
|
|
|
# Delete the file:
|
|
await self._fs.delete(file_id = file_id, session = session)
|
|
success = True
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def delete_one_file(
|
|
self,
|
|
filter,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
To delete one document from a collection.
|
|
:param filter: The filter criteria.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: True or False based on the success of the operation.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
success = False
|
|
|
|
# Find one file that matches the given filter:
|
|
file = await self.find_one(
|
|
collection = "fs.files",
|
|
filter = filter,
|
|
session = session,
|
|
raise_exception = raise_exception
|
|
)
|
|
|
|
# If we have a list of files to work with:
|
|
if file: success = await self.delete_file_by_id(file["_id"])
|
|
|
|
# Done here:
|
|
return success
|
|
|
|
async def delete_many_files(
|
|
self,
|
|
filter,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
To delete many documents from a collection.
|
|
WARNING: sending {} in the filter would mean deleting ALL the documents.
|
|
:param filter: The filter criteria.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The number of files deleted (can be zero) or None if something failed while searching the files.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
count = 0
|
|
|
|
# Find all the files that match the given filter:
|
|
files = await self.find_many(
|
|
collection = "fs.files",
|
|
filter = filter,
|
|
session = session,
|
|
raise_exception = raise_exception
|
|
)
|
|
|
|
# If we have a list of files to work with:
|
|
if files is not None:
|
|
tasks = [self.delete_file_by_id(file["_id"]) for file in files]
|
|
results = await asyncio.gather(*tasks)
|
|
count = sum(results)
|
|
|
|
# Done here:
|
|
return count
|
|
|
|
async def get_upload_stream(
|
|
self,
|
|
file_name: str,
|
|
file_metadata: dict = None,
|
|
file_id = None,
|
|
chunk_size: int = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Returns a GridIn object so that you can perform your own upload using the built-in writing methods. You must use
|
|
the 'write' method to write data to the file by passing it either a string of bytes or a file-like object. When
|
|
the file has been fully written, you must call the 'close' method to finish the operation. In case you need to
|
|
cancel the operation, you can call the 'abort' method to delete all the already written data and stop uploading
|
|
new data.
|
|
:param file_name: The name of the file as it will be stored on (and retrieved from) GridFS.
|
|
:param file_metadata: Any metadata to later search the file by.
|
|
:param file_id: Any custom id to be given to the file. TRY STICKING TO THE ID GENERATED BY 'generate_id'.
|
|
:param chunk_size: The chunk size (in bytes) to use for storing this file.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The upload stream that implements the 'write', 'close', and 'abort' methods, or None if something
|
|
failed.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
stream = None
|
|
|
|
# Try to open the upload stream:
|
|
try:
|
|
|
|
# If no file id was supplied:
|
|
if file_id is None:
|
|
stream = self._fs.open_upload_stream(
|
|
filename = file_name,
|
|
metadata = file_metadata,
|
|
chunk_size_bytes = chunk_size,
|
|
session = session
|
|
)
|
|
|
|
# If a file id was supplied:
|
|
else:
|
|
stream = self._fs.open_upload_stream_with_id(
|
|
file_id = file_id,
|
|
filename = file_name,
|
|
metadata = file_metadata,
|
|
chunk_size_bytes = chunk_size,
|
|
session = session
|
|
)
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return stream
|
|
|
|
async def get_download_stream(
|
|
self,
|
|
file_id = None,
|
|
file_name = None,
|
|
session = None,
|
|
raise_exception = False
|
|
):
|
|
|
|
"""
|
|
Returns a GridOut object so that you can implement your own download logic using the built-in 'read' method.
|
|
Once the reading is done, use the 'close' method to release the resources used by the stream.
|
|
:param file_id: (RECOMMENDED) the id of the save file.
|
|
:param file_name: The name of the saved file. NOT RECOMMENDED because you could have many files with the same
|
|
name. The best way to tell files apart if from the id.
|
|
:param session: The session if you need to do this in a transaction.
|
|
:param raise_exception: Whether, or not, you want to raise an exception when something fails.
|
|
:return: The download stream that implements the 'read' and 'close' methods, or None if something failed.
|
|
"""
|
|
|
|
# Ensure you are connected:
|
|
await self.ensure_connection()
|
|
|
|
# Assume failure:
|
|
stream = None
|
|
|
|
# Try to open a download stream:
|
|
try:
|
|
|
|
# If a file id is supplied (preferred way):
|
|
if file_id is not None:
|
|
stream = await self._fs.open_download_stream(
|
|
file_id = file_id,
|
|
session = session
|
|
)
|
|
|
|
# If a file name is supplied:
|
|
elif file_name is not None:
|
|
stream = await self._fs.open_download_stream_by_name(
|
|
filename = file_name,
|
|
session = session
|
|
)
|
|
|
|
# When something goes wrong:
|
|
except Exception as exception:
|
|
self._printer(exception)
|
|
if raise_exception: raise
|
|
|
|
# Done here:
|
|
return stream
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MAIN PROGRAM ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
from shared import constants
|
|
|
|
async def main():
|
|
|
|
# Create an instance of the database connector:
|
|
my_db = AsyncMongo(
|
|
connection_string = constants.MONGO_DATA_CONNECTION_STRING,
|
|
database_name = constants.MONGO_DATA_DATABASE_NAME,
|
|
max_connections = 10,
|
|
debug = True
|
|
)
|
|
|
|
# Connect to the database:
|
|
await my_db.connect()
|
|
|
|
# # Get the documents to migrate:
|
|
# documents = await my_db.find_many(
|
|
# collection = "scriptData",
|
|
# filter = {},
|
|
# limit = 50,
|
|
# projection = {"_id": False}
|
|
# )
|
|
# # print(json.to_string(documents, default = str))
|
|
#
|
|
# # Adjust them:
|
|
# adjusted_documents = []
|
|
# for document in documents:
|
|
# script_id = document.pop("scriptId")
|
|
# adjusted_document = {
|
|
# "scriptId": script_id,
|
|
# "desc": "no desc",
|
|
# "content": document
|
|
# }
|
|
# adjusted_documents.append(adjusted_document)
|
|
# print(json.to_string(adjusted_documents, default = str))
|
|
#
|
|
# # Insert the adjusted ones to the new collection:
|
|
# response = await my_db.insert_many(
|
|
# collection = "_scriptData",
|
|
# documents = adjusted_documents
|
|
# )
|
|
# print("RESPONSE:", response)
|
|
|
|
|
|
asyncio.run(main())
|