Handle UUID as string and define embeddedfield.db_type()

Author: WaVEV

django_mongodb_backend/fields/embedded_model.py

@@ -21,6 +21,9 @@ def __init__(self, embedded_model, *args, **kwargs):
         self.embedded_model = embedded_model
         super().__init__(*args, **kwargs)
 
+    def db_type(self, connection):
+        return "embeddedDocuments"
+
     def check(self, **kwargs):
         from ..models import EmbeddedModel
 

django_mongodb_backend/indexes.py

@@ -125,18 +125,12 @@ def check(self, model, connection):
             )
         return errors
 
-    def search_index_data_types(self, field, db_type):
+    def search_index_data_types(self, db_type):
         """
         Map a model field's internal type to search index type.
         Reference: https://www.mongodb.com/docs/atlas/atlas-search/define-field-mappings/#data-types
         """
-        if field.get_internal_type() == "UUIDField":
-            return "uuid"
-        if field.get_internal_type() in {"ObjectIdAutoField", "ObjectIdField"}:
-            return "ObjectId"
-        if field.get_internal_type() == "EmbeddedModelField":
-            return "embeddedDocuments"
-        if db_type in {"int", "long"}:
+        if db_type in {"double", "int", "long"}:
             return "number"
         if db_type == "binData":
             return "string"
@@ -153,8 +147,8 @@ def get_pymongo_index_model(
             return None
         fields = {}
         for field_name, _ in self.fields_orders:
-            field_ = model._meta.get_field(field_name)
-            type_ = self.search_index_data_types(field_, field_.db_type(schema_editor.connection))
+            field = model._meta.get_field(field_name)
+            type_ = self.search_index_data_types(field.db_type(schema_editor.connection))
             field_path = column_prefix + model._meta.get_field(field_name).column
             fields[field_path] = {"type": type_}
         return SearchIndexModel(
@@ -212,8 +206,7 @@ def check(self, model, connection):
                         )
                     )
             else:
-                field_type = field_.db_type(connection)
-                search_type = self.search_index_data_types(field_, field_type)
+                search_type = self.search_index_data_types(field_.db_type(connection))
                 # filter - for fields that contain boolean, date, objectId,
                 # numeric, string, or UUID values. Reference:
                 # https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/#atlas-vector-search-index-fields
@@ -253,7 +246,7 @@ def get_pymongo_index_model(
             return None
         similarities = (
             itertools.cycle([self.similarities])
-            if isinstance(self.similarities, str)
+            if not self._multiple_similarities
             else iter(self.similarities)
         )
         fields = []

tests/indexes_/models.py

@@ -1,6 +1,6 @@
 from django.db import models
 
-from django_mongodb_backend.fields import ArrayField, EmbeddedModelField
+from django_mongodb_backend.fields import ArrayField, EmbeddedModelField, ObjectIdField
 from django_mongodb_backend.models import EmbeddedModel
 
 
@@ -10,6 +10,7 @@ class Data(EmbeddedModel):
 
 class Article(models.Model):
     headline = models.CharField(max_length=100)
+    object_id = ObjectIdField()
     number = models.IntegerField()
     body = models.TextField()
     data = models.JSONField()

tests/indexes_/test_search_indexes.py

@@ -79,43 +79,30 @@ def test_define_field_twice(self):
             )
 
 
-class SchemaAssertionMixin:
-    def assertAddRemoveIndex(self, editor, model, index):
-        with self.assertNumQueries(1):
-            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)
-        self.assertNotIn(
-            index.name,
-            connection.introspection.get_constraints(
-                cursor=None,
-                table_name=model._meta.db_table,
-            ),
-        )
-
-
 @skipUnlessDBFeature("supports_atlas_search")
-class SearchIndexSchemaTests(SchemaAssertionMixin, TestCase):
+class SearchIndexSchemaTests(TestCase):
     def test_simple(self):
         with connection.schema_editor() as editor:
             index = SearchIndex(
                 name="recent_article_idx",
                 fields=["number"],
             )
             editor.add_index(index=index, model=Article)
-            self.assertAddRemoveIndex(editor, Article, index)
+            editor.remove_index(index=index, model=Article)
 
     def test_multiple_fields(self):
         with connection.schema_editor() as editor:
             index = SearchIndex(
                 name="recent_article_idx",
-                fields=["headline", "number", "body", "data", "embedded", "created_at"],
+                fields=[
+                    "headline",
+                    "number",
+                    "body",
+                    "data",
+                    "embedded",
+                    "created_at",
+                    "object_id",
+                ],
             )
             editor.add_index(index=index, model=Article)
             index_info = connection.introspection.get_constraints(
@@ -146,20 +133,21 @@ def test_multiple_fields(self):
                         "representation": "double",
                         "type": "number",
                     },
+                    "object_id": {"type": "objectId"},
                 },
             }
             self.assertCountEqual(index_info[index.name]["columns"], index.fields)
             self.assertEqual(index_info[index.name]["options"], expected_options)
-            self.assertAddRemoveIndex(editor, Article, index)
+            editor.remove_index(index=index, model=Article)
 
 
 @skipUnlessDBFeature("supports_atlas_search")
-class VectorSearchIndexSchemaTests(SchemaAssertionMixin, TestCase):
+class VectorSearchIndexSchemaTests(TestCase):
     def test_simple_vector_search(self):
         with connection.schema_editor() as editor:
             index = VectorSearchIndex(name="recent_article_idx", fields=["number"])
             editor.add_index(index=index, model=Article)
-            self.assertAddRemoveIndex(editor, Article, index)
+            editor.remove_index(index=index, model=Article)
 
     def test_multiple_fields(self):
         with connection.schema_editor() as editor:
@@ -195,4 +183,4 @@ def test_multiple_fields(self):
             index_info[index.name]["options"].pop("id")
             index_info[index.name]["options"].pop("status")
             self.assertEqual(index_info[index.name]["options"], expected_options)
-            self.assertAddRemoveIndex(editor, Article, index)
+            editor.remove_index(index=index, model=Article)