def test_update_document(self):
mock_mongo_client = MagicMock()
with patch('qilib.data_set.mongo_data_set_io.MongoClient',
return_value={'qilib': {
'data_sets': mock_mongo_client
}}):
mongo_data_set_io = MongoDataSetIO(name='test_data_set')
mongo_data_set_io.update_document(
{'array_updates': ('(2,2)', {
'test': 5
})})
mock_mongo_client.update_one.called_once_with(
{'name': 'test_data_set'}, {
'$set': {
'array_updates': ('(2,2)', {
'test': 5
})
},
"$currentDate": {
"lastModified": True
}
})
class MongoDataSetIOWriter(DataSetIOWriter):
""" Allow a DataSet to store changes, and complete DataSet, to a mongodb."""
def __init__(
self,
name: Optional[str] = None,
document_id: Optional[str] = None,
database: str = MongoDataSetIO.DEFAULT_DATABASE_NAME,
collection: str = MongoDataSetIO.DEFAULT_COLLECTION_NAME) -> None:
""" Construct a new instance of MongoDataSetIOWriter. If name is provided, but not found in the database
a new document is created with that name.
Args:
name: DataSet name.
document_id: _id of the DataSet in the database.
database: Name of the database.
collection: Name of the collections.
Raises:
DocumentNotFoundError: If document_id is provided but not found in the database.
"""
super().__init__()
self._mongo_data_set_io = MongoDataSetIO(name,
document_id,
database=database,
collection=collection)
def sync_metadata_to_storage(self, field_name: str, value: Any) -> None:
""" Update or add metadata field to database.
Args:
field_name: Field that changed.
value: The new value.
"""
self._is_finalized()
update_data = {
"{}.{}".format(DataSetIOReader.METADATA, field_name): value
}
self._mongo_data_set_io.update_document(update_data)
def sync_data_to_storage(self, index_or_slice: Union[Tuple[int, ...], int],
data: Dict[str, Any]) -> None:
""" Registers a DataArray update to the database. The change is registered as a change event and is
applied on finalize().
Args:
index_or_slice: The indices of the DataArray to update.
data: Name of the DataArray to be updated and the new value.
"""
self._is_finalized()
update_data = {DataSetIOReader.ARRAY_UPDATES: (index_or_slice, data)}
self._mongo_data_set_io.append_to_document(update_data)
def sync_add_data_array_to_storage(self, data_array: DataArray) -> None:
""" Add or update a DataArray in the database.
Args:
data_array: The DataArray to be updated or added.
"""
self._is_finalized()
update_data = {
"{}.{}".format(DataSetIOReader.DATA_ARRAYS, data_array.name): {
"name": data_array.name,
"label": data_array.label,
"unit": data_array.unit,
"is_setpoint": data_array.is_setpoint,
"set_arrays": [array.name for array in data_array.set_arrays],
"preset_data": MongoDataSetIO.encode_numpy_array(data_array)
}
}
self._mongo_data_set_io.update_document(update_data)
def finalize(self) -> None:
""" Update the underlying DataSet and close the connection to the database."""
self._mongo_data_set_io.update_document(
{DataSetIOReader.ARRAY_UPDATES: []})
self._mongo_data_set_io.finalize()
self._finalized = True