Fix mapping typing issue

WaVEV · WaVEV · commit 60e57b3606aa · 2025-02-25T15:15:38.000-03:00
diff --git a/django_mongodb_backend/base.py b/django_mongodb_backend/base.py
@@ -127,6 +127,38 @@ def _isnull_operator(a, b):
         "iregex": lambda a, b: regex_match(a, b, insensitive=True),
     }
 
+    mongo_data_types = {
+        "AutoField": "number",
+        "BigAutoField": "number",
+        "BinaryField": "string",
+        "BooleanField": "boolean",
+        "CharField": "string",
+        "DateField": "date",
+        "DateTimeField": "date",
+        "DecimalField": "number",
+        "DurationField": "number",
+        "FileField": "string",
+        "FilePathField": "string",
+        "FloatField": "double",
+        "IntegerField": "number",
+        "BigIntegerField": "number",
+        "GenericIPAddressField": "string",
+        "JSONField": "document",
+        "OneToOneField": "number",
+        "PositiveBigIntegerField": "number",
+        "PositiveIntegerField": "number",
+        "PositiveSmallIntegerField": "number",
+        "SlugField": "string",
+        "SmallAutoField": "number",
+        "SmallIntegerField": "number",
+        "TextField": "string",
+        "TimeField": "date",
+        "UUIDField": "uuid",
+        "ObjectIdAutoField": ["objectId", "number"],
+        "ObjectIdField": ["objectId"],
+        "EmbeddedModelField": "embeddedDocuments",
+    }
+
     display_name = "MongoDB"
     vendor = "mongodb"
     Database = Database
@@ -175,12 +207,7 @@ def get_connection_params(self):
 
     @async_unsafe
     def get_new_connection(self, conn_params):
-        # 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())
+        return MongoClient(**conn_params, directConnection=True, driver=self._driver_info())
 
     def _driver_info(self):
         if not os.environ.get("RUNNING_DJANGOS_TEST_SUITE"):
diff --git a/django_mongodb_backend/indexes.py b/django_mongodb_backend/indexes.py
@@ -60,7 +60,16 @@ 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,
+    *,
+    connection=None,  # noqa: ARG001
+    field=None,
+    unique=False,
+    column_prefix="",
+):
     from collections import defaultdict
 
     if self.contains_expressions:
@@ -101,6 +110,8 @@ def create_mongodb_index(self, model, schema_editor, *, field=None, unique=False
 
 
 class AtlasSearchIndex(Index):
+    suffix = "atlas_search"
+
     def __init__(self, *expressions, **kwargs):
         super().__init__(*expressions, **kwargs)
 
@@ -110,9 +121,10 @@ def create_mongodb_index(
         fields = {}
         for field_name, _ in self.fields_orders:
             field_ = model._meta.get_field(field_name)
-            fields[field_name] = {"type": field_.db_type(connection)}
+            type_ = connection.mongo_data_types[field_.get_internal_type()]
+            fields[field_name] = {"type": type_}
         return SearchIndexModel(
-            definition={"mappings": {"dynamic": False}, "fields": fields}, name=self.name
+            definition={"mappings": {"dynamic": False, "fields": fields}}, name=self.name
         )
 
 
diff --git a/django_mongodb_backend/introspection.py b/django_mongodb_backend/introspection.py
@@ -2,14 +2,16 @@
 from django.db.models import Index
 from pymongo import ASCENDING, DESCENDING
 
+from django_mongodb_backend.indexes import AtlasSearchIndex
+
 
 class DatabaseIntrospection(BaseDatabaseIntrospection):
     ORDER_DIR = {ASCENDING: "ASC", DESCENDING: "DESC"}
 
     def table_names(self, cursor=None, include_views=False):
         return sorted([x["name"] for x in self.connection.database.list_collections()])
 
-    def get_constraints(self, cursor, table_name):
+    def _get_index_info(self, table_name):
         indexes = self.connection.get_collection(table_name).index_information()
         constraints = {}
         for name, details in indexes.items():
@@ -30,3 +32,25 @@ def get_constraints(self, cursor, table_name):
                 "options": {},
             }
         return constraints
+
+    def _get_atlas_index_info(self, table_name):
+        constraints = {}
+        indexes = self.connection.get_collection(table_name).list_search_indexes()
+        for details in indexes:
+            columns = list(details["latestDefinition"]["mappings"].get("fields", {}).keys())
+            constraints[details["name"]] = {
+                "check": False,
+                "columns": columns,
+                "definition": None,
+                "foreign_key": None,
+                "index": True,
+                "orders": [],
+                "primary_key": False,
+                "type": AtlasSearchIndex.suffix,
+                "unique": False,
+                "options": details["latestDefinition"]["mappings"],
+            }
+        return constraints
+
+    def get_constraints(self, cursor, table_name):
+        return {**self._get_index_info(table_name), **self._get_atlas_index_info(table_name)}
diff --git a/tests/indexes_/models.py b/tests/indexes_/models.py
@@ -1,7 +1,17 @@
 from django.db import models
 
+from django_mongodb_backend.fields import EmbeddedModelField
+from django_mongodb_backend.models import EmbeddedModel
+
+
+class Data(EmbeddedModel):
+    integer = models.IntegerField()
+
 
 class Article(models.Model):
     headline = models.CharField(max_length=100)
     number = models.IntegerField()
     body = models.TextField()
+    data = models.JSONField()
+    embedded = EmbeddedModelField(Data)
+    auto_now = models.DateTimeField(auto_now=True)
diff --git a/tests/indexes_/test_atlas_indexes.py b/tests/indexes_/test_atlas_indexes.py
@@ -1,31 +1,76 @@
 from django.db import connection
-from django.test import TransactionTestCase
-from django.test.utils import override_settings
+from django.test import TestCase
 
 from django_mongodb_backend.indexes import AtlasSearchIndex
 
 from .models import Article
 
 
-@override_settings(USE_TZ=True)
-class PartialIndexTests(TransactionTestCase):
+class PartialIndexTests(TestCase):
     # Schema editor is used to  create the index to test that it works.
     # available_apps = ["indexes"]
     available_apps = None
 
-    def test_partial_index(self):
+    def assertAddRemoveIndex(self, editor, model, index):
+        editor.add_index(index=index, model=model)
+        self.assertIn(
+            index.name,
+            connection.introspection.get_constraints(
+                cursor=None,
+                table_name=model._meta.db_table,
+            ),
+        )
+        editor.remove_index(index=index, model=model)
+
+    def test_simple_atlas_index(self):
         with connection.schema_editor() as editor:
             index = AtlasSearchIndex(
                 name="recent_article_idx",
                 fields=["number"],
             )
             editor.add_index(index=index, model=Article)
-            # with connection.cursor() as cursor:
-            #     self.assertIn(
-            #         index.name,
-            #         connection.introspection.get_constraints(
-            #             cursor=cursor,
-            #             table_name=Article._meta.db_table,
-            #         ),
-            #     )
-            editor.remove_index(index=index, model=Article)
+            self.assertAddRemoveIndex(editor, Article, index)
+
+    def test_multiple_fields_atlas_index(self):
+        with connection.schema_editor() as editor:
+            index = AtlasSearchIndex(
+                name="recent_article_idx",
+                fields=["headline", "number", "body", "data", "embedded", "auto_now"],
+            )
+            # editor.remove_index(index=index, model=Article)
+            editor.add_index(index=index, model=Article)
+            index_info = connection.introspection.get_constraints(
+                cursor=None,
+                table_name=Article._meta.db_table,
+            )
+            expected_options = {
+                "dynamic": False,
+                "fields": {
+                    "auto_now": {"type": "date"},
+                    "body": {
+                        "indexOptions": "offsets",
+                        "norms": "include",
+                        "store": True,
+                        "type": "string",
+                    },
+                    "data": {"dynamic": False, "fields": {}, "type": "document"},
+                    "embedded": {"dynamic": False, "fields": {}, "type": "embeddedDocuments"},
+                    "headline": {
+                        "indexOptions": "offsets",
+                        "norms": "include",
+                        "store": True,
+                        "type": "string",
+                    },
+                    "number": {
+                        "indexDoubles": True,
+                        "indexIntegers": True,
+                        "representation": "double",
+                        "type": "number",
+                    },
+                },
+            }
+            self.assertCountEqual(index_info[index.name]["columns"], index.fields)
+            self.assertEqual(index_info[index.name]["options"], expected_options)
+
+            self.assertEqual(index_info, {})
+            self.assertAddRemoveIndex(editor, Article, index)
