Atlas index creation

Author: WaVEV

django_mongodb_backend/base.py

@@ -175,7 +175,12 @@ def get_connection_params(self):
 
     @async_unsafe
     def get_new_connection(self, conn_params):
-        return MongoClient(**conn_params, driver=self._driver_info())
+        # import ipdb
+        # ipdb.set_trace()
+        return MongoClient(
+            "mongodb://127.0.0.1:27017/?directConnection=true", driver=self._driver_info()
+        )
+        # return MongoClient(**conn_params, driver=self._driver_info())
 
     def _driver_info(self):
         if not os.environ.get("RUNNING_DJANGOS_TEST_SUITE"):

django_mongodb_backend/indexes.py

@@ -60,7 +60,7 @@ def where_node_idx(self, compiler, connection):
     return mql
 
 
-def create_mongodb_index(self, model, schema_editor, field=None, unique=False, column_prefix=""):
+def create_mongodb_index(self, model, schema_editor, *, field=None, unique=False, column_prefix=""):
     from collections import defaultdict
 
     if self.contains_expressions:
@@ -105,10 +105,14 @@ def __init__(self, *expressions, **kwargs):
         super().__init__(*expressions, **kwargs)
 
     def create_mongodb_index(
-        self, model, schema_editor, field=None, unique=False, column_prefix=""
+        self, model, schema_editor, connection=None, field=None, unique=False, column_prefix=""
     ):
+        fields = {}
+        for field_name, _ in self.fields_orders:
+            field_ = model._meta.get_field(field_name)
+            fields[field_name] = {"type": field_.db_type(connection)}
         return SearchIndexModel(
-            definitions={},
+            definition={"mappings": {"dynamic": False}, "fields": fields}, name=self.name
         )
 
 

django_mongodb_backend/schema.py

@@ -1,5 +1,10 @@
+from functools import singledispatchmethod
+
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
 from django.db.models import Index, UniqueConstraint
+from pymongo.operations import IndexModel, SearchIndexModel
+
+from django_mongodb_backend.indexes import AtlasSearchIndex
 
 from .fields import EmbeddedModelField
 from .query import wrap_database_errors
@@ -256,16 +261,33 @@ def alter_unique_together(
                 model, constraint, parent_model=parent_model, column_prefix=column_prefix
             )
 
+    @singledispatchmethod
+    def _add_index(self, index, model):
+        raise ValueError(f"{type(index)} isn't a supported index type")
+
+    @_add_index.register
+    def _(self, index: IndexModel, model):
+        return self.get_collection(model._meta.db_table).create_indexes([index])
+
+    @_add_index.register
+    def _(self, index: SearchIndexModel, model):
+        return self.get_collection(model._meta.db_table).create_search_index(index)
+
     @ignore_embedded_models
     def add_index(
         self, model, index, *, field=None, unique=False, column_prefix="", parent_model=None
     ):
         idx = index.create_mongodb_index(
-            model, self, field=field, unique=unique, column_prefix=column_prefix
+            model,
+            self,
+            field=field,
+            unique=unique,
+            column_prefix=column_prefix,
+            connection=self.connection,
         )
         if idx:
             model = parent_model or model
-            self.get_collection(model._meta.db_table).create_indexes([idx])
+            self._add_index(idx, model)
 
     def _add_composed_index(self, model, field_names, column_prefix="", parent_model=None):
         """Add an index on the given list of field_names."""
@@ -279,11 +301,23 @@ def _add_field_index(self, model, field, *, column_prefix=""):
         index.name = self._create_index_name(model._meta.db_table, [column_prefix + field.column])
         self.add_index(model, index, field=field, column_prefix=column_prefix)
 
+    @singledispatchmethod
+    def _remove_index(self, index, model):
+        raise ValueError(f"{type(index)} isn't a supported index type")
+
+    @_remove_index.register
+    def _(self, index: Index, model):
+        return self.get_collection(model._meta.db_table).drop_index(index.name)
+
+    @_remove_index.register
+    def _(self, index: AtlasSearchIndex, model):
+        return self.get_collection(model._meta.db_table).drop_search_index(index.name)
+
     @ignore_embedded_models
     def remove_index(self, model, index):
         if index.contains_expressions:
             return
-        self.get_collection(model._meta.db_table).drop_index(index.name)
+        self._remove_index(index, model)
 
     def _remove_composed_index(
         self, model, field_names, constraint_kwargs, column_prefix="", parent_model=None

requirements.txt

@@ -1,2 +1,2 @@
 django>=5.1,<5.2
-pymongo>=4.6,<5.0
+pymongo>=4.11,<5.0