From f3919dd839390e3a989de73c9c5d418a1f872685 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felix=20Schulthei=C3=9F?= Date: Tue, 27 Oct 2020 12:55:35 +0100 Subject: [PATCH 1/6] stripped out integrating shard key from _save_update, use it also in _save_create --- mongoengine/document.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 4a57d511..a3a30bc1 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -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: From 5a8e5e5a40c2453f99fc9882b8bdae61eefefc45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felix=20Schulthei=C3=9F?= Date: Tue, 27 Oct 2020 16:34:57 +0100 Subject: [PATCH 2/6] updated docstring --- mongoengine/document.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index a3a30bc1..2dcd3116 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -492,6 +492,9 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass): return update_doc def _integrate_shard_key(self, doc, select_dict): + """Integrates the collection's shard key to the `select_dict`, which will be used for the query. + The value from the shard key is taken from the `doc` and finally the select_dict is returned. + """ # Need to add shard key to query, or you get an error shard_key = self._meta.get("shard_key", tuple()) @@ -505,7 +508,6 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass): return select_dict - def _save_update(self, doc, save_condition, write_concern): """Update an existing document. From bf1d04e399062a77580385ea452cb19b2a5f4e97 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felix=20Schulthei=C3=9F?= Date: Thu, 29 Oct 2020 14:56:08 +0100 Subject: [PATCH 3/6] black reformatting --- mongoengine/document.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 2dcd3116..801c8df8 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -466,9 +466,7 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass): 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( - select_dict, doc - ) + raw_object = wc_collection.find_one_and_replace(select_dict, doc) if raw_object: return doc["_id"] @@ -930,7 +928,7 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass): @classmethod def list_indexes(cls): - """ Lists all of the indexes that should be created for given + """Lists all of the indexes that should be created for given collection. It includes all the indexes from super- and sub-classes. """ if cls._meta.get("abstract"): @@ -995,7 +993,7 @@ class Document(BaseDocument, metaclass=TopLevelDocumentMetaclass): @classmethod def compare_indexes(cls): - """ Compares the indexes defined in MongoEngine with the ones + """Compares the indexes defined in MongoEngine with the ones existing in the database. Returns any missing/extra indexes. """ From 30a3c6a5b76eb5f350859be3d926351bf060f363 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felix=20Schulthei=C3=9F?= Date: Mon, 2 Nov 2020 17:30:24 +0100 Subject: [PATCH 4/6] added testcase for save create with shard key --- tests/document/test_instance.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py index 8d42d15b..e5ac2756 100644 --- a/tests/document/test_instance.py +++ b/tests/document/test_instance.py @@ -501,7 +501,7 @@ class TestDocumentInstance(MongoDBTestCase): doc.reload() Animal.drop_collection() - def test_update_shard_key_routing(self): + def test_save_update_shard_key_routing(self): """Ensures updating a doc with a specified shard_key includes it in the query. """ @@ -529,6 +529,36 @@ class TestDocumentInstance(MongoDBTestCase): Animal.drop_collection() + def test_save_create_shard_key_routing(self): + """Ensures inserting a doc with a specified shard_key includes it in + the query. + """ + + class Animal(Document): + _id = UUIDField(binary=False, primary_key=True, default=uuid.uuid4) + is_mammal = BooleanField() + name = StringField() + meta = {"shard_key": ("is_mammal",)} + + Animal.drop_collection() + doc = Animal(is_mammal=True, name="Dog") + + mongo_db = get_mongodb_version() + + with query_counter() as q: + doc.save() + query_op = q.db.system.profile.find({"ns": "mongoenginetest.animal"})[0] + assert query_op["op"] == "command" + assert query_op["command"]["findAndModify"] == "animal" + if mongo_db <= MONGODB_34: + assert set(query_op["query"].keys()) == set(["_id", "is_mammal"]) + else: + assert set(query_op["command"]["query"].keys()) == set( + ["_id", "is_mammal"] + ) + + Animal.drop_collection() + def test_reload_with_changed_fields(self): """Ensures reloading will not affect changed fields""" From 0bc18cd6e100747ed5139d1e0774dc0502fb7d06 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felix=20Schulthei=C3=9F?= Date: Tue, 3 Nov 2020 10:00:02 +0100 Subject: [PATCH 5/6] fixed shard test case for old mongodb version --- tests/document/test_instance.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py index e5ac2756..3a077735 100644 --- a/tests/document/test_instance.py +++ b/tests/document/test_instance.py @@ -543,19 +543,12 @@ class TestDocumentInstance(MongoDBTestCase): Animal.drop_collection() doc = Animal(is_mammal=True, name="Dog") - mongo_db = get_mongodb_version() - with query_counter() as q: doc.save() query_op = q.db.system.profile.find({"ns": "mongoenginetest.animal"})[0] assert query_op["op"] == "command" assert query_op["command"]["findAndModify"] == "animal" - if mongo_db <= MONGODB_34: - assert set(query_op["query"].keys()) == set(["_id", "is_mammal"]) - else: - assert set(query_op["command"]["query"].keys()) == set( - ["_id", "is_mammal"] - ) + assert set(query_op["command"]["query"].keys()) == set(["_id", "is_mammal"]) Animal.drop_collection() From 79f9f223d018dea27bea49807094a5df6703aafb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felix=20Schulthei=C3=9F?= Date: Tue, 3 Nov 2020 10:00:15 +0100 Subject: [PATCH 6/6] added to authors --- AUTHORS | 1 + 1 file changed, 1 insertion(+) diff --git a/AUTHORS b/AUTHORS index 02e43955..80fbd99a 100644 --- a/AUTHORS +++ b/AUTHORS @@ -257,3 +257,4 @@ that much better: * Matthew Simpson (https://github.com/mcsimps2) * Leonardo Domingues (https://github.com/leodmgs) * Agustin Barto (https://github.com/abarto) + * Felix Schultheiß (https://github.com/felix-smashdocs)