pub/mongoengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

stripped out integrating shard key from _save_update, use it also in _save_create

Browse Source
This commit is contained in:
Felix Schultheiß 2020-10-27 12:55:35 +01:00
parent 9f82a02ddf
commit f3919dd839
1 changed files with 19 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
mongoengine/document.py
View File

@ -464,8 +464,10 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass):
# insert_one will provoke UniqueError alongside save does not # insert_one will provoke UniqueError alongside save does not
# therefore, it need to catch and call replace_one. # therefore, it need to catch and call replace_one.
if "_id" in doc: if "_id" in doc:
select_dict = {"_id": doc["_id"]}
select_dict = self._integrate_shard_key(doc, select_dict)
raw_object = wc_collection.find_one_and_replace( raw_object = wc_collection.find_one_and_replace(
{"_id": doc["_id"]}, doc select_dict, doc
) )
if raw_object: if raw_object:
return doc["_id"] return doc["_id"]
@ -489,6 +491,21 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass):
return update_doc return update_doc
def _integrate_shard_key(self, doc, select_dict):
# Need to add shard key to query, or you get an error
shard_key = self._meta.get("shard_key", tuple())
for k in shard_key:
path = self._lookup_field(k.split("."))
actual_key = [p.db_field for p in path]
val = doc
for ak in actual_key:
val = val[ak]
select_dict[".".join(actual_key)] = val
return select_dict
def _save_update(self, doc, save_condition, write_concern): def _save_update(self, doc, save_condition, write_concern):
"""Update an existing document. """Update an existing document.
@ -504,15 +521,7 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass):
select_dict["_id"] = object_id select_dict["_id"] = object_id
# Need to add shard key to query, or you get an error select_dict = self._integrate_shard_key(doc, select_dict)
shard_key = self._meta.get("shard_key", tuple())
for k in shard_key:
path = self._lookup_field(k.split("."))
actual_key = [p.db_field for p in path]
val = doc
for ak in actual_key:
val = val[ak]
select_dict[".".join(actual_key)] = val
update_doc = self._get_update_doc() update_doc = self._get_update_doc()
if update_doc: if update_doc: