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

This commit is contained in:
Felix Schultheiß 2020-10-27 12:55:35 +01:00
parent 9f82a02ddf
commit f3919dd839

View File

@ -464,8 +464,10 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass):
# insert_one will provoke UniqueError alongside save does not
# therefore, it need to catch and call replace_one.
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(
{"_id": doc["_id"]}, doc
select_dict, doc
)
if raw_object:
return doc["_id"]
@ -489,6 +491,21 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass):
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):
"""Update an existing document.
@ -504,15 +521,7 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass):
select_dict["_id"] = object_id
# 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
select_dict = self._integrate_shard_key(doc, select_dict)
update_doc = self._get_update_doc()
if update_doc: