implement SchemaEditor.add/remove_constraint() and creating uniques in create_model()

Author: timgraham

django_mongodb/features.py

@@ -21,6 +21,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     # BSON Date type doesn't support microsecond precision.
     supports_microsecond_precision = False
     supports_paramstyle_pyformat = False
+    # Not implemented.
+    supports_partial_indexes = False
     supports_select_difference = False
     supports_select_intersection = False
     supports_sequence_reset = False
@@ -92,9 +94,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "schema.tests.SchemaTests.test_alter_field_add_index_to_integerfield",
         "schema.tests.SchemaTests.test_alter_field_default_dropped",
         "schema.tests.SchemaTests.test_alter_field_fk_to_o2o",
-        "schema.tests.SchemaTests.test_alter_field_o2o_keeps_unique",
         "schema.tests.SchemaTests.test_alter_field_o2o_to_fk",
-        "schema.tests.SchemaTests.test_alter_int_pk_to_int_unique",
         "schema.tests.SchemaTests.test_alter_null_to_not_null",
         "schema.tests.SchemaTests.test_alter_null_to_not_null_keeping_default",
         "schema.tests.SchemaTests.test_alter_primary_key_the_same_name",
@@ -121,12 +121,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "schema.tests.SchemaTests.test_m2m_through_alter_custom",
         "schema.tests.SchemaTests.test_m2m_through_alter_inherited",
         "schema.tests.SchemaTests.test_m2m_through_remove",
-        # add/remove_constraint
-        "schema.tests.SchemaTests.test_composed_constraint_with_fk",
-        "schema.tests.SchemaTests.test_remove_ignored_unique_constraint_not_create_fk_index",
-        "schema.tests.SchemaTests.test_unique_constraint",
-        # constraints not fully implemented.
-        "introspection.tests.IntrospectionTests.test_get_constraints",
     }
     # $bitAnd, #bitOr, and $bitXor are new in MongoDB 6.3.
     _django_test_expected_failures_bitwise = {
@@ -196,24 +190,6 @@ def django_test_expected_failures(self):
             "model_fields.test_autofield.SmallAutoFieldTests",
             "queries.tests.TestInvalidValuesRelation.test_invalid_values",
         },
-        "MongoDB does not enforce UNIQUE constraints.": {
-            "auth_tests.test_basic.BasicTestCase.test_unicode_username",
-            "auth_tests.test_migrations.ProxyModelWithSameAppLabelTests.test_migrate_with_existing_target_permission",
-            "constraints.tests.UniqueConstraintTests.test_database_constraint",
-            "contenttypes_tests.test_operations.ContentTypeOperationsTests.test_content_type_rename_conflict",
-            "contenttypes_tests.test_operations.ContentTypeOperationsTests.test_existing_content_type_rename",
-            "custom_pk.tests.CustomPKTests.test_unique_pk",
-            "force_insert_update.tests.ForceInsertInheritanceTests.test_force_insert_with_existing_grandparent",
-            "get_or_create.tests.GetOrCreateTestsWithManualPKs.test_create_with_duplicate_primary_key",
-            "get_or_create.tests.GetOrCreateTestsWithManualPKs.test_savepoint_rollback",
-            "get_or_create.tests.GetOrCreateThroughManyToMany.test_something",
-            "get_or_create.tests.UpdateOrCreateTests.test_manual_primary_key_test",
-            "get_or_create.tests.UpdateOrCreateTestsWithManualPKs.test_create_with_duplicate_primary_key",
-            "introspection.tests.IntrospectionTests.test_get_constraints_unique_indexes_orders",
-            "model_fields.test_filefield.FileFieldTests.test_unique_when_same_filename",
-            "one_to_one.tests.OneToOneTests.test_multiple_o2o",
-            "queries.test_bulk_update.BulkUpdateTests.test_database_routing_batch_atomicity",
-        },
         "MongoDB does not enforce PositiveIntegerField constraint.": {
             "model_fields.test_integerfield.PositiveIntegerFieldTests.test_negative_values",
         },

django_mongodb/query.py

@@ -9,14 +9,17 @@
 from django.db.models.sql.constants import INNER
 from django.db.models.sql.datastructures import Join
 from django.db.models.sql.where import AND, OR, XOR, NothingNode, WhereNode
-from pymongo.errors import DuplicateKeyError, PyMongoError
+from pymongo.errors import BulkWriteError, DuplicateKeyError, PyMongoError
 
 
 def wrap_database_errors(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
+        except BulkWriteError as e:
+            if "E11000 duplicate key error" in str(e):
+                raise IntegrityError from e
         except DuplicateKeyError as e:
             raise IntegrityError from e
         except PyMongoError as e:

django_mongodb/schema.py

@@ -1,5 +1,5 @@
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
-from django.db.models import Index
+from django.db.models import Index, UniqueConstraint
 from pymongo.operations import IndexModel
 
 
@@ -14,8 +14,8 @@ def create_model(self, model):
 
     def _create_model_indexes(self, model):
         """
-        Create all indexes (field indexes, index_together, Meta.indexes) for
-        the specified model.
+        Create all indexes (field indexes & uniques, index_together,
+        Meta.constraints, Meta.indexes) for the specified model.
         """
         if not model._meta.managed or model._meta.proxy or model._meta.swapped:
             return
@@ -25,11 +25,19 @@ def _create_model_indexes(self, model):
                 index = Index(fields=[field.name])
                 index.set_name_with_model(model)
                 self.add_index(model, index)
+        # Field uniques
+        for field in model._meta.local_fields:
+            if field.unique and field.column != "_id":
+                constraint = UniqueConstraint(fields=[field.name], name="%s_uniq" % field.column)
+                self.add_constraint(model, constraint)
         # Meta.index_together (RemovedInDjango51Warning)
         for field_names in model._meta.index_together:
             index = Index(fields=field_names)
             index.set_name_with_model(model)
             self.add_index(model, index)
+        # Meta.constraints
+        for constraint in model._meta.constraints:
+            self.add_constraint(model, constraint)
         # Meta.indexes
         for index in model._meta.indexes:
             self.add_index(model, index)
@@ -81,7 +89,7 @@ def alter_index_together(self, model, old_index_together, new_index_together):
     def alter_unique_together(self, model, old_unique_together, new_unique_together):
         pass
 
-    def add_index(self, model, index):
+    def add_index(self, model, index, unique=False):
         if index.contains_expressions:
             return
         idx = IndexModel(
@@ -90,6 +98,7 @@ def add_index(self, model, index):
                 for field_name, order in index.fields_orders
             ],
             name=index.name,
+            unique=unique,
         )
         self.connection.database[model._meta.db_table].create_indexes([idx])
 
@@ -99,10 +108,18 @@ def remove_index(self, model, index):
         self.connection.database[model._meta.db_table].drop_index(index.name)
 
     def add_constraint(self, model, constraint):
-        pass
+        if constraint.expressions:
+            return
+        if isinstance(constraint, UniqueConstraint):
+            idx = Index(fields=constraint.fields, name=constraint.name)
+            self.add_index(model, idx, unique=True)
 
     def remove_constraint(self, model, constraint):
-        pass
+        if constraint.expressions:
+            return
+        if isinstance(constraint, UniqueConstraint):
+            idx = Index(fields=constraint.fields, name=constraint.name)
+            self.remove_index(model, idx)
 
     def alter_db_table(self, model, old_db_table, new_db_table):
         self.connection.database[old_db_table].rename(new_db_table)