28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140 | class PythonObjectJsonToMongo(PythonObjectJson):
"""PythonObjectJson subclass with built-in save/load functionality to/from MongoDB."""
def __init__(self, mongo_host: str, mongo_port: int, mongo_database: str, mongo_user: str, mongo_password: str):
super().__init__(excluded_attributes=["(^mongo_[A-Za-z]*)"])
self.mongo_host: str = mongo_host
self.mongo_port: int = mongo_port
self.mongo_database: str = mongo_database
self.mongo_user: str = mongo_user
self.mongo_password: str = mongo_password
def _get_mongo_client(self):
"""Create a MongoDB connection.
Returns:
MongoClient: A pymongo MongoClient instance.
"""
return MongoClient(
f"mongodb://{quote_plus(self.mongo_user)}:{quote_plus(self.mongo_password)}"
f"@{self.mongo_host}:{self.mongo_port}/{self.mongo_database}"
f"?authSource=admin",
serverSelectionTimeoutMS=5000,
)
def _validate_or_create_collection(self, mongo_collection: str) -> Collection:
"""Create a pymongo Database instance from a pymongo MongoClient, check if a given MongoDB collection exists,
and create the collection if it does not exist.
Args:
mongo_collection (str): The name of the MongoDB collection for which to check existence or create.
Returns:
Collection: A pymongo Collection instance.
"""
db = self._get_mongo_client()[self.mongo_database]
try:
db.validate_collection(mongo_collection)
except ServerSelectionTimeoutError:
logger.warning(f'Unable to connect to MongoDB server at "{self.mongo_host}:{self.mongo_port}".')
sys.exit(1)
except OperationFailure:
logger.debug(f'MongoDB collection "{mongo_collection}" does not exist.')
db.create_collection(mongo_collection)
return db.get_collection(mongo_collection)
@staticmethod
def _validate_document_id(mongo_document_id: Union[ObjectId, bytes, str]) -> None:
"""This method checks to see if a given MongoDB document ID is valid.
Args:
mongo_document_id (Union[ObjectId, bytes, str]): The MongoDB document ID to validate.
Returns:
None
"""
if not ObjectId.is_valid(mongo_document_id):
logger.error(
f'Invalid MongoDb document ID "{mongo_document_id}". MongoDB requires document ObjectId values to '
f"be either 12 bytes long or a 24-character hexadecimal string."
)
sys.exit(1)
def save_to_mongo(
self, mongo_collection: str, mongo_document_id: Optional[Union[ObjectId, bytes, str]] = None
) -> ObjectId:
"""Save the custom Python object to a specified MongoDB collection.
Args:
mongo_collection (str): The name of the MongoDB collection into which to save the custom Python object.
mongo_document_id (Optional[ObjectId, bytes, str], optional): MongoDB document ID. Defaults to None, which
will result in MongoDB automatically generating a unique document ID.
Returns:
ObjectId: The MongoDB document ID to which the custom Python object JSON was saved.
"""
# only validate MongoDB document ID if one is provided
if mongo_document_id:
self._validate_document_id(mongo_document_id)
collection = self._validate_or_create_collection(mongo_collection)
document: Dict[str, Any] = collection.find_one_and_update(
{"_id": ObjectId(mongo_document_id) if mongo_document_id else ObjectId()},
{"$set": {"custom_class": self.serialize()}},
projection={"_id": True}, # filter out all fields besides the document ID
upsert=True, # create a new document if it does not exist, otherwise update the existing document
return_document=ReturnDocument.AFTER, # return the updated or created document after the update/creation
)
return document["_id"]
def load_from_mongo(self, mongo_collection: str, mongo_document_id: Union[ObjectId, bytes, str]) -> None:
"""Load the JSON values from a specified MongoDB document ID to the custom Python object from a specified
MongoDB collection.
Args:
mongo_collection (str): The name of the MongoDB collection from which to load the custom Python object data.
mongo_document_id (Union[ObjectId, bytes, str]): The MongoDB document ID from which the custom Python object
JSON was loaded.
Returns:
None
"""
self._validate_document_id(mongo_document_id)
# get MongoDb collection
collection = self._validate_or_create_collection(mongo_collection)
self.deserialize(collection.find_one({"_id": ObjectId(mongo_document_id)}).get("custom_class"))
|