feat: handle geohaystack removal

Author: terencehonles

mongoengine/base/document.py

@@ -28,6 +28,11 @@
 
 NON_FIELD_ERRORS = "__all__"
 
+try:
+    GEOHAYSTACK = pymongo.GEOHAYSTACK
+except AttributeError:
+    GEOHAYSTACK = None
+
 
 class BaseDocument:
     # TODO simplify how `_changed_fields` is used.
@@ -898,7 +903,10 @@ def _build_index_spec(cls, spec):
             elif key.startswith("("):
                 direction = pymongo.GEOSPHERE
             elif key.startswith(")"):
-                direction = pymongo.GEOHAYSTACK
+                try:
+                    direction = pymongo.GEOHAYSTACK
+                except AttributeError:
+                    raise NotImplementedError
             elif key.startswith("*"):
                 direction = pymongo.GEO2D
             if key.startswith(("+", "-", "*", "$", "#", "(", ")")):
@@ -923,10 +931,10 @@ def _build_index_spec(cls, spec):
             index_list.append((key, direction))
 
         # Don't add cls to a geo index
-        if include_cls and direction not in (
-            pymongo.GEO2D,
-            pymongo.GEOHAYSTACK,
-            pymongo.GEOSPHERE,
+        if (
+            include_cls
+            and direction not in (pymongo.GEO2D, pymongo.GEOSPHERE)
+            and (GEOHAYSTACK is None or direction != GEOHAYSTACK)
         ):
             index_list.insert(0, ("_cls", 1))
 

mongoengine/pymongo_support.py

@@ -4,12 +4,8 @@
 import pymongo
 from pymongo.errors import OperationFailure
 
-_PYMONGO_37 = (3, 7)
-
 PYMONGO_VERSION = tuple(pymongo.version_tuple[:2])
 
-IS_PYMONGO_GTE_37 = PYMONGO_VERSION >= _PYMONGO_37
-
 
 def count_documents(
     collection, filter, skip=None, limit=None, hint=None, collation=None
@@ -29,7 +25,7 @@ def count_documents(
         kwargs["collation"] = collation
 
     # count_documents appeared in pymongo 3.7
-    if IS_PYMONGO_GTE_37:
+    if PYMONGO_VERSION >= (3, 7):
         try:
             return collection.count_documents(filter=filter, **kwargs)
         except OperationFailure:
@@ -49,7 +45,7 @@ def count_documents(
 
 def list_collection_names(db, include_system_collections=False):
     """Pymongo>3.7 deprecates collection_names in favour of list_collection_names"""
-    if IS_PYMONGO_GTE_37:
+    if PYMONGO_VERSION >= (3, 7):
         collections = db.list_collection_names()
     else:
         collections = db.collection_names()

tests/document/test_indexes.py

@@ -11,6 +11,7 @@
     MONGODB_42,
     get_mongodb_version,
 )
+from mongoengine.pymongo_support import PYMONGO_VERSION
 
 
 class TestIndexes(unittest.TestCase):
@@ -247,11 +248,9 @@ class Place(Document):
 
     def test_explicit_geohaystack_index(self):
         """Ensure that geohaystack indexes work when created via meta[indexes]"""
-        pytest.skip(
-            "GeoHaystack index creation is not supported for now"
-            "from meta, as it requires a bucketSize parameter."
-        )
-
+        # This test can be removed when pymongo 3.x is no longer supported
+        if PYMONGO_VERSION >= (4,):
+            pytest.skip('GEOHAYSTACK has been removed in pymongo 4.0')
         class Place(Document):
             location = DictField()
             name = StringField()
@@ -261,22 +260,32 @@ class Place(Document):
             {"fields": [("location.point", "geoHaystack"), ("name", 1)]}
         ] == Place._meta["index_specs"]
 
-        Place.ensure_indexes()
-        info = Place._get_collection().index_information()
-        info = [value["key"] for key, value in info.items()]
-        assert [("location.point", "geoHaystack")] in info
+        # GeoHaystack index creation is not supported for now from meta, as it
+        # requires a bucketSize parameter.
+        if False:
+            Place.ensure_indexes()
+            info = Place._get_collection().index_information()
+            info = [value["key"] for key, value in info.items()]
+            assert [("location.point", "geoHaystack")] in info
 
     def test_create_geohaystack_index(self):
         """Ensure that geohaystack indexes can be created"""
-
         class Place(Document):
             location = DictField()
             name = StringField()
 
-        Place.create_index({"fields": (")location.point", "name")}, bucketSize=10)
-        info = Place._get_collection().index_information()
-        info = [value["key"] for key, value in info.items()]
-        assert [("location.point", "geoHaystack"), ("name", 1)] in info
+        # This test can be removed when pymongo 3.x is no longer supported
+        if PYMONGO_VERSION >= (4,):
+            with pytest.raises(NotImplementedError):
+                Place.create_index(
+                    {"fields": (")location.point", "name")},
+                    bucketSize=10,
+                )
+        else:
+            Place.create_index({"fields": (")location.point", "name")}, bucketSize=10)
+            info = Place._get_collection().index_information()
+            info = [value["key"] for key, value in info.items()]
+            assert [("location.point", "geoHaystack"), ("name", 1)] in info
 
     def test_dictionary_indexes(self):
         """Ensure that indexes are used when meta[indexes] contains

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = pypy3-{mg34,mg36,mg39,mg311,mg312}
+envlist = pypy3-{mg34,mg36,mg39,mg311,mg312,mg4}
 
 [testenv]
 commands =