Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
- spec: dynamic field (dict-like), can have any structure
- tag: str
- base_results: tuple (required), first value is the class type
of the result, {'results' or 'procedure'). The second value is
the ID of the result in the DB. Example:
"base_result": ('results', result_id)
Returns
-------
dict (data and meta)
'data' is a list of the IDs of the tasks IN ORDER, including
duplicates. An errored task has 'None' in its ID
meta['duplicates'] has the duplicate tasks
"""
meta = add_metadata_template()
results = []
with self.session_scope() as session:
for task_num, record in enumerate(data):
try:
task_dict = record.dict(exclude={"id"})
# # for compatibility with mongoengine
# if isinstance(task_dict['base_result'], dict):
# task_dict['base_result'] = task_dict['base_result']['id']
task = TaskQueueORM(**task_dict)
task.priority = task.priority.value # Must be an integer for sorting
session.add(task)
session.commit()
results.append(str(task.id))
meta["n_inserted"] += 1
except IntegrityError as err: # rare case
def _copy_users(self, record_list: Dict):
"""
copy the given users as-is to the DB. Used for data migration
Parameters
----------
record_list : list of dict of managers data
Returns
-------
Dict with keys: data, meta
Data is the ids of the inserted/updated/existing docs
"""
meta = add_metadata_template()
user_names = []
with self.session_scope() as session:
for user in record_list:
doc = session.query(UserORM).filter_by(username=user["username"])
if get_count_fast(doc) == 0:
doc = UserORM(**user)
if isinstance(doc.password, str): # TODO, for mongo
doc.password = doc.password.encode("ascii")
session.add(doc)
session.commit()
user_names.append(doc.username)
meta["n_inserted"] += 1
else:
name = doc.first().username
def add_wavefunction_store(self, blobs_list: List[Dict[str, Any]]):
"""
Adds to the wavefunction key/value store table.
Parameters
----------
blobs_list : List[Dict[str, Any]]
A list of wavefunction data blobs to add.
Returns
-------
TYPE
Description
"""
meta = add_metadata_template()
blob_ids = []
with self.session_scope() as session:
for blob in blobs_list:
if blob is None:
blob_ids.append(None)
continue
doc = WavefunctionStoreORM(**blob)
session.add(doc)
session.commit()
blob_ids.append(str(doc.id))
meta["n_inserted"] += 1
meta["success"] = True
return {"data": blob_ids, "meta": meta}
"""
Adds to the key/value store table.
Parameters
----------
blobs_list : List[Any]
A list of data blobs to add.
Returns
-------
TYPE
Description
"""
meta = add_metadata_template()
blob_ids = []
with self.session_scope() as session:
for blob in blobs_list:
if blob is None:
blob_ids.append(None)
continue
doc = KVStoreORM(value=blob)
session.add(doc)
session.commit()
blob_ids.append(str(doc.id))
meta["n_inserted"] += 1
meta["success"] = True
return {"data": blob_ids, "meta": meta}
"""
Adds to the key/value store table.
Parameters
----------
blobs_list : List[Any]
A list of data blobs to add.
Returns
-------
A dict with keys: 'data' and 'meta'
(see get_metadata_template())
The 'data' part is an object of the result or None if not found
"""
meta = add_metadata_template()
blob_ids = []
for blob in blobs_list:
if blob is None:
blob_ids.append(None)
continue
doc = KVStoreORM(value=blob)
doc.save()
blob_ids.append(str(doc.id))
meta['n_inserted'] += 1
meta["success"] = True
return {"data": blob_ids, "meta": meta}
Parameters
----------
data : list of dict
Each dict must have:
program, driver, method, basis, options, molecule
Where molecule is the molecule id in the DB
In addition, it should have the other attributes that it needs
to store
Returns
-------
Dict with keys: data, meta
Data is the ids of the inserted/updated/existing docs
"""
meta = add_metadata_template()
result_ids = []
with self.session_scope() as session:
for result in record_list:
doc = session.query(ResultORM).filter_by(
program=result.program,
driver=result.driver,
method=result.method,
basis=result.basis,
keywords=result.keywords,
molecule=result.molecule,
)
if get_count_fast(doc) == 0:
doc = ResultORM(**result.dict(exclude={"id"}))
A list of KeywordSets to be inserted.
Returns
-------
A dict with keys: 'data' and 'meta'
(see add_metadata_template())
The 'data' part is a list of ids of the inserted options
data['duplicates'] has the duplicate entries
Notes
------
Duplicates are not considered errors.
"""
meta = add_metadata_template()
keywords = []
for kw in keyword_sets:
kw_dict = kw.json_dict(exclude={"id"})
# search by index keywords not by all keys, much faster
found = KeywordsORM.objects(hash_index=kw_dict['hash_index']).first()
if not found:
doc = KeywordsORM(**kw_dict).save()
keywords.append(str(doc.id))
meta['n_inserted'] += 1
else:
meta['duplicates'].append(str(found.id)) # TODO
keywords.append(str(found.id))
def add_molecules(self, molecules: List[Molecule]):
"""
Adds molecules to the database.
Parameters
----------
data : dict of molecule-like JSON objects
A {key: molecule} dictionary of molecules to input.
Returns
-------
bool
Whether the operation was successful.
"""
meta = add_metadata_template()
results = []
for dmol in molecules:
mol_dict = dmol.json_dict(exclude={"id"})
mol_dict["fix_com"] = True
mol_dict["fix_orientation"] = True
# Build fresh indices
mol_dict["molecule_hash"] = dmol.get_hash()
mol_dict["molecular_formula"] = dmol.get_molecular_formula()
mol_dict["identifiers"] = {}
mol_dict["identifiers"]["molecule_hash"] = mol_dict["molecule_hash"]
mol_dict["identifiers"]["molecular_formula"] = mol_dict["molecular_formula"]
Returns
-------
A dict with keys: 'data' and 'meta'
(see add_metadata_template())
The 'data' part is the id of the inserted document or none
Notes
-----
** Change: The data doesn't have to include the ID, the document
is identified by the (collection, name) pairs.
** Change: New fields will be added to the collection, but existing won't
be removed.
"""
meta = add_metadata_template()
col_id = None
# try:
# if ("id" in data) and (data["id"] == "local"):
# data.pop("id", None)
if "id" in data: # remove the ID in any case
data.pop("id", None)
lname = data.get("name").lower()
collection = data.pop("collection").lower()
# Get collection class if special type is implemented
collection_class = get_collection_class(collection)
update_fields = {}
for field in collection_class._all_col_names():
if field in data: