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

@@ -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
 =================

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

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):

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()