From aa76ccdd25d58c6563a0548ea13f42ffbe46d2fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stefan=20W=C3=B3jcik?= Date: Tue, 9 Jul 2019 12:08:26 +0200 Subject: [PATCH] Fix Document._object_key (#2125) Previous implementation of `Document._object_key` was *pretending* to work on MongoEngine-level fields (e.g. using "pk" instead of "_id" and separating nested field parts by "__" instead of "."), but then it was also attempting to transform field names from the `shard_key` into DB-level fields. This, expectedly, didn't really work well. Most of the test cases added in this commit were failing prior to the code fixes. --- docs/changelog.rst | 3 +- mongoengine/base/document.py | 1 + mongoengine/document.py | 16 ++++---- tests/document/instance.py | 76 +++++++++++++++++++++++++++++++++--- 4 files changed, 82 insertions(+), 14 deletions(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index e9f342e1..7cf74d66 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -15,8 +15,9 @@ Development - If you catch/use ``MongoEngineConnectionError`` in your code, you'll have to rename it. - BREAKING CHANGE: Positional arguments when instantiating a document are no longer supported. #2103 - From now on keyword arguments (e.g. ``Doc(field_name=value)``) are required. -- The codebase is now formatted using ``black``. #2109 +- Fix updating/modifying/deleting/reloading a document that's sharded by a field with ``db_field`` specified. #2125 - ``ListField`` now accepts an optional ``max_length`` parameter. #2110 +- The codebase is now formatted using ``black``. #2109 Changes in 0.18.2 ================= diff --git a/mongoengine/base/document.py b/mongoengine/base/document.py index 928a00c2..f5109b44 100644 --- a/mongoengine/base/document.py +++ b/mongoengine/base/document.py @@ -202,6 +202,7 @@ class BaseDocument(object): self__initialised = self._initialised except AttributeError: self__initialised = False + # Check if the user has created a new instance of a class if ( self._is_document diff --git a/mongoengine/document.py b/mongoengine/document.py index 41166df4..23968f17 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -575,22 +575,24 @@ class Document(six.with_metaclass(TopLevelDocumentMetaclass, BaseDocument)): @property def _object_key(self): - """Get the query dict that can be used to fetch this object from - the database. + """Return a query dict that can be used to fetch this document. Most of the time the dict is a simple PK lookup, but in case of a sharded collection with a compound shard key, it can contain a more complex query. + + Note that the dict returned by this method uses MongoEngine field + names instead of PyMongo field names (e.g. "pk" instead of "_id", + "some__nested__field" instead of "some.nested.field", etc.). """ select_dict = {"pk": self.pk} shard_key = self.__class__._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 = self - for ak in actual_key: - val = getattr(val, ak) - select_dict["__".join(actual_key)] = val + field_parts = k.split(".") + for part in field_parts: + val = getattr(val, part) + select_dict["__".join(field_parts)] = val return select_dict def update(self, **kwargs): diff --git a/tests/document/instance.py b/tests/document/instance.py index 49606cff..9c854f8d 100644 --- a/tests/document/instance.py +++ b/tests/document/instance.py @@ -466,21 +466,33 @@ class InstanceTest(MongoDBTestCase): meta = {"shard_key": ("superphylum",)} Animal.drop_collection() - doc = Animal(superphylum="Deuterostomia") - doc.save() + doc = Animal.objects.create(superphylum="Deuterostomia") mongo_db = get_mongodb_version() CMD_QUERY_KEY = "command" if mongo_db >= MONGODB_36 else "query" - with query_counter() as q: doc.reload() query_op = q.db.system.profile.find({"ns": "mongoenginetest.animal"})[0] self.assertEqual( - set(query_op[CMD_QUERY_KEY]["filter"].keys()), - set(["_id", "superphylum"]), + set(query_op[CMD_QUERY_KEY]["filter"].keys()), {"_id", "superphylum"} ) - Animal.drop_collection() + def test_reload_sharded_with_db_field(self): + class Person(Document): + nationality = StringField(db_field="country") + meta = {"shard_key": ("nationality",)} + + Person.drop_collection() + doc = Person.objects.create(nationality="Poland") + + mongo_db = get_mongodb_version() + CMD_QUERY_KEY = "command" if mongo_db >= MONGODB_36 else "query" + with query_counter() as q: + doc.reload() + query_op = q.db.system.profile.find({"ns": "mongoenginetest.person"})[0] + self.assertEqual( + set(query_op[CMD_QUERY_KEY]["filter"].keys()), {"_id", "country"} + ) def test_reload_sharded_nested(self): class SuperPhylum(EmbeddedDocument): @@ -3616,5 +3628,57 @@ class InstanceTest(MongoDBTestCase): User.objects().select_related() +class ObjectKeyTestCase(MongoDBTestCase): + def test_object_key_simple_document(self): + class Book(Document): + title = StringField() + + book = Book(title="Whatever") + self.assertEqual(book._object_key, {"pk": None}) + + book.pk = ObjectId() + self.assertEqual(book._object_key, {"pk": book.pk}) + + def test_object_key_with_custom_primary_key(self): + class Book(Document): + isbn = StringField(primary_key=True) + title = StringField() + + book = Book(title="Sapiens") + self.assertEqual(book._object_key, {"pk": None}) + + book = Book(pk="0062316117") + self.assertEqual(book._object_key, {"pk": "0062316117"}) + + def test_object_key_in_a_sharded_collection(self): + class Book(Document): + title = StringField() + meta = {"shard_key": ("pk", "title")} + + book = Book() + self.assertEqual(book._object_key, {"pk": None, "title": None}) + book = Book(pk=ObjectId(), title="Sapiens") + self.assertEqual(book._object_key, {"pk": book.pk, "title": "Sapiens"}) + + def test_object_key_with_custom_db_field(self): + class Book(Document): + author = StringField(db_field="creator") + meta = {"shard_key": ("pk", "author")} + + book = Book(pk=ObjectId(), author="Author") + self.assertEqual(book._object_key, {"pk": book.pk, "author": "Author"}) + + def test_object_key_with_nested_shard_key(self): + class Author(EmbeddedDocument): + name = StringField() + + class Book(Document): + author = EmbeddedDocumentField(Author) + meta = {"shard_key": ("pk", "author.name")} + + book = Book(pk=ObjectId(), author=Author(name="Author")) + self.assertEqual(book._object_key, {"pk": book.pk, "author__name": "Author"}) + + if __name__ == "__main__": unittest.main()