Merge pull request #2854 from bagerard/Winedays-feature/query_set_more_support_for_collation_hint_comment

set more support for collation hint comment (Winedays feature/query )

Author: bagerard

docs/changelog.rst

@@ -15,6 +15,7 @@ Development
 - Further to the deprecation warning, remove ability to use an unpacked list to `Queryset.aggregate(*pipeline)`, a plain list must be provided instead `Queryset.aggregate(pipeline)`, as it's closer to pymongo interface
 - Further to the deprecation warning, remove `full_response` from `QuerySet.modify` as it wasn't supported with Pymongo 3+
 - Fixed stacklevel of many warnings (to point places emitting the warning more accurately)
+- Add support for collation/hint/comment to delete/update and aggregate #2842
 
 Changes in 0.29.0
 =================

tests/document/test_instance.py

@@ -42,7 +42,12 @@
     PickleSignalsTest,
     PickleTest,
 )
-from tests.utils import MongoDBTestCase, get_as_pymongo
+from tests.utils import (
+    MongoDBTestCase,
+    db_ops_tracker,
+    get_as_pymongo,
+    requires_mongodb_gte_44,
+)
 
 TEST_IMAGE_PATH = os.path.join(os.path.dirname(__file__), "../fields/mongoengine.png")
 
@@ -1871,6 +1876,53 @@ class User(self.Person):
         person = self.Person.objects.get()
         assert not person.comments_dict["first_post"].published
 
+    @requires_mongodb_gte_44
+    def test_update_propagates_hint_collation_and_comment(self):
+        """Make sure adding a hint/comment/collation to the query gets added to the query"""
+        mongo_ver = get_mongodb_version()
+
+        base = {"locale": "en", "strength": 2}
+        index_name = "name_1"
+
+        class AggPerson(Document):
+            name = StringField()
+            meta = {
+                "indexes": [{"fields": ["name"], "name": index_name, "collation": base}]
+            }
+
+        AggPerson.drop_collection()
+        _ = AggPerson.objects.first()
+
+        comment = "test_comment"
+
+        if PYMONGO_VERSION >= (4, 1):
+            with db_ops_tracker() as q:
+                _ = AggPerson.objects.comment(comment).update_one(name="something")
+                query_op = q.db.system.profile.find(
+                    {"ns": "mongoenginetest.agg_person"}
+                )[0]
+                CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+                assert "hint" not in query_op[CMD_QUERY_KEY]
+                assert query_op[CMD_QUERY_KEY]["comment"] == comment
+                assert "collation" not in query_op[CMD_QUERY_KEY]
+
+        with db_ops_tracker() as q:
+            _ = AggPerson.objects.hint(index_name).update_one(name="something")
+            query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
+            CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+
+            assert query_op[CMD_QUERY_KEY]["hint"] == {"$hint": index_name}
+            assert "comment" not in query_op[CMD_QUERY_KEY]
+            assert "collation" not in query_op[CMD_QUERY_KEY]
+
+        with db_ops_tracker() as q:
+            _ = AggPerson.objects.collation(base).update_one(name="something")
+            query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
+            CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+            assert "hint" not in query_op[CMD_QUERY_KEY]
+            assert "comment" not in query_op[CMD_QUERY_KEY]
+            assert query_op[CMD_QUERY_KEY]["collation"] == base
+
     def test_delete(self):
         """Ensure that document may be deleted using the delete method."""
         person = self.Person(name="Test User", age=30)
@@ -1879,6 +1931,53 @@ def test_delete(self):
         person.delete()
         assert self.Person.objects.count() == 0
 
+    @requires_mongodb_gte_44
+    def test_delete_propagates_hint_collation_and_comment(self):
+        """Make sure adding a hint/comment/collation to the query gets added to the query"""
+        mongo_ver = get_mongodb_version()
+
+        base = {"locale": "en", "strength": 2}
+        index_name = "name_1"
+
+        class AggPerson(Document):
+            name = StringField()
+            meta = {
+                "indexes": [{"fields": ["name"], "name": index_name, "collation": base}]
+            }
+
+        AggPerson.drop_collection()
+        _ = AggPerson.objects.first()
+
+        comment = "test_comment"
+
+        if PYMONGO_VERSION >= (4, 1):
+            with db_ops_tracker() as q:
+                _ = AggPerson.objects().comment(comment).delete()
+                query_op = q.db.system.profile.find(
+                    {"ns": "mongoenginetest.agg_person"}
+                )[0]
+                CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+                assert "hint" not in query_op[CMD_QUERY_KEY]
+                assert query_op[CMD_QUERY_KEY]["comment"] == comment
+                assert "collation" not in query_op[CMD_QUERY_KEY]
+
+        with db_ops_tracker() as q:
+            _ = AggPerson.objects.hint(index_name).delete()
+            query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
+            CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+
+            assert query_op[CMD_QUERY_KEY]["hint"] == {"$hint": index_name}
+            assert "comment" not in query_op[CMD_QUERY_KEY]
+            assert "collation" not in query_op[CMD_QUERY_KEY]
+
+        with db_ops_tracker() as q:
+            _ = AggPerson.objects.collation(base).delete()
+            query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
+            CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+            assert "hint" not in query_op[CMD_QUERY_KEY]
+            assert "comment" not in query_op[CMD_QUERY_KEY]
+            assert query_op[CMD_QUERY_KEY]["collation"] == base
+
     def test_save_custom_id(self):
         """Ensure that a document may be saved with a custom _id."""
 

tests/queryset/test_queryset.py

@@ -27,21 +27,13 @@
 )
 from mongoengine.queryset.base import BaseQuerySet
 from tests.utils import (
+    db_ops_tracker,
     requires_mongodb_gte_42,
     requires_mongodb_gte_44,
     requires_mongodb_lt_42,
 )
 
 
-class db_ops_tracker(query_counter):
-    def get_ops(self):
-        ignore_query = dict(self._ignored_query)
-        ignore_query["command.count"] = {
-            "$ne": "system.profile"
-        }  # Ignore the query issued by query_counter
-        return list(self.db.system.profile.find(ignore_query))
-
-
 def get_key_compat(mongo_ver):
     ORDER_BY_KEY = "sort"
     CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"

tests/queryset/test_queryset_aggregation.py

@@ -3,8 +3,12 @@
 import pytest
 from pymongo.read_preferences import ReadPreference
 
-from mongoengine import *
-from tests.utils import MongoDBTestCase
+from mongoengine import Document, IntField, PointField, StringField
+from mongoengine.mongodb_support import (
+    MONGODB_36,
+    get_mongodb_version,
+)
+from tests.utils import MongoDBTestCase, db_ops_tracker
 
 
 class TestQuerysetAggregate(MongoDBTestCase):
@@ -87,6 +91,49 @@ class Person(Document):
             {"_id": p3.pk, "name": "SANDRA MARA"},
         ]
 
+    def test_aggregation_propagates_hint_collation_and_comment(self):
+        """Make sure adding a hint/comment/collation to the query gets added to the query"""
+        mongo_ver = get_mongodb_version()
+
+        base = {"locale": "en", "strength": 2}
+        index_name = "name_1"
+
+        class AggPerson(Document):
+            name = StringField()
+            meta = {
+                "indexes": [{"fields": ["name"], "name": index_name, "collation": base}]
+            }
+
+        AggPerson.drop_collection()
+        _ = AggPerson.objects.first()
+
+        pipeline = [{"$project": {"name": {"$toUpper": "$name"}}}]
+        comment = "test_comment"
+
+        with db_ops_tracker() as q:
+            _ = list(AggPerson.objects.comment(comment).aggregate(pipeline))
+            query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
+            CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+            assert "hint" not in query_op[CMD_QUERY_KEY]
+            assert query_op[CMD_QUERY_KEY]["comment"] == comment
+            assert "collation" not in query_op[CMD_QUERY_KEY]
+
+        with db_ops_tracker() as q:
+            _ = list(AggPerson.objects.hint(index_name).aggregate(pipeline))
+            query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
+            CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+            assert query_op[CMD_QUERY_KEY]["hint"] == "name_1"
+            assert "comment" not in query_op[CMD_QUERY_KEY]
+            assert "collation" not in query_op[CMD_QUERY_KEY]
+
+        with db_ops_tracker() as q:
+            _ = list(AggPerson.objects.collation(base).aggregate(pipeline))
+            query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
+            CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
+            assert "hint" not in query_op[CMD_QUERY_KEY]
+            assert "comment" not in query_op[CMD_QUERY_KEY]
+            assert query_op[CMD_QUERY_KEY]["collation"] == base
+
     def test_queryset_aggregation_with_limit(self):
         class Person(Document):
             name = StringField()

tests/utils.py

@@ -2,12 +2,16 @@
 import operator
 import unittest
 
+import pymongo
 import pytest
 
 from mongoengine import connect
 from mongoengine.connection import disconnect_all, get_db
+from mongoengine.context_managers import query_counter
 from mongoengine.mongodb_support import get_mongodb_version
 
+PYMONGO_VERSION = tuple(pymongo.version_tuple[:2])
+
 MONGO_TEST_DB = "mongoenginetest"  # standard name for the test database
 
 
@@ -92,3 +96,12 @@ def _inner(*args, **kwargs):
             pytest.skip(f"Needs MongoDB {oper.__name__} v{pretty_version}")
 
     return _inner
+
+
+class db_ops_tracker(query_counter):
+    def get_ops(self):
+        ignore_query = dict(self._ignored_query)
+        ignore_query["command.count"] = {
+            "$ne": "system.profile"
+        }  # Ignore the query issued by query_counter
+        return list(self.db.system.profile.find(ignore_query))