Skip to content

Commit bf70cf1

Browse files
committed
implement SchemaEditor.add/remove_index()
And also creating indexes in create_model() and add_field(), and removing them in remove_field().
1 parent 5c29684 commit bf70cf1

File tree

3 files changed

+80
-31
lines changed

3 files changed

+80
-31
lines changed

django_mongodb/features.py

Lines changed: 2 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -70,29 +70,13 @@ class DatabaseFeatures(BaseDatabaseFeatures):
7070
"backends.tests.ThreadTests.test_closing_non_shared_connections",
7171
"backends.tests.ThreadTests.test_default_connection_thread_local",
7272
# AddField
73-
"schema.tests.SchemaTests.test_add_indexed_charfield",
7473
"schema.tests.SchemaTests.test_add_unique_charfield",
75-
# Add/RemoveIndex
76-
"migrations.test_operations.OperationTests.test_add_index",
77-
"migrations.test_operations.OperationTests.test_alter_field_with_index",
78-
"migrations.test_operations.OperationTests.test_remove_index",
79-
"migrations.test_operations.OperationTests.test_rename_index",
80-
"migrations.test_operations.OperationTests.test_rename_index_unknown_unnamed_index",
81-
"migrations.test_operations.OperationTests.test_rename_index_unnamed_index",
82-
"schema.tests.SchemaTests.test_add_remove_index",
83-
"schema.tests.SchemaTests.test_composed_desc_index_with_fk",
84-
"schema.tests.SchemaTests.test_composed_index_with_fk",
85-
"schema.tests.SchemaTests.test_create_index_together",
86-
"schema.tests.SchemaTests.test_order_index",
87-
"schema.tests.SchemaTests.test_text_field_with_db_index",
8874
# AlterField
8975
"schema.tests.SchemaTests.test_alter_field_add_index_to_integerfield",
90-
"schema.tests.SchemaTests.test_alter_field_fk_keeps_index",
9176
"schema.tests.SchemaTests.test_alter_field_fk_to_o2o",
9277
"schema.tests.SchemaTests.test_alter_field_o2o_keeps_unique",
9378
"schema.tests.SchemaTests.test_alter_field_o2o_to_fk",
9479
"schema.tests.SchemaTests.test_alter_int_pk_to_int_unique",
95-
"schema.tests.SchemaTests.test_alter_not_unique_field_to_primary_key",
9680
# AlterField (db_index)
9781
"schema.tests.SchemaTests.test_alter_renames_index",
9882
"schema.tests.SchemaTests.test_indexes",
@@ -114,9 +98,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
11498
"schema.tests.SchemaTests.test_composed_constraint_with_fk",
11599
"schema.tests.SchemaTests.test_remove_ignored_unique_constraint_not_create_fk_index",
116100
"schema.tests.SchemaTests.test_unique_constraint",
117-
# subclasses of BaseDatabaseIntrospection may require a get_constraints() method
118-
"migrations.test_operations.OperationTests.test_add_func_unique_constraint",
119-
"migrations.test_operations.OperationTests.test_remove_func_unique_constraint",
120101
}
121102
# $bitAnd, #bitOr, and $bitXor are new in MongoDB 6.3.
122103
_django_test_expected_failures_bitwise = {
@@ -220,6 +201,7 @@ def django_test_expected_failures(self):
220201
"get_or_create.tests.GetOrCreateThroughManyToMany.test_something",
221202
"get_or_create.tests.UpdateOrCreateTests.test_manual_primary_key_test",
222203
"get_or_create.tests.UpdateOrCreateTestsWithManualPKs.test_create_with_duplicate_primary_key",
204+
"introspection.tests.IntrospectionTests.test_get_constraints_unique_indexes_orders",
223205
"model_fields.test_filefield.FileFieldTests.test_unique_when_same_filename",
224206
"one_to_one.tests.OneToOneTests.test_multiple_o2o",
225207
"queries.test_bulk_update.BulkUpdateTests.test_database_routing_batch_atomicity",
@@ -611,14 +593,8 @@ def django_test_expected_failures(self):
611593
"introspection.tests.IntrospectionTests.test_get_table_description_types",
612594
"introspection.tests.IntrospectionTests.test_smallautofield",
613595
},
614-
"DatabaseIntrospection.get_constraints() not implemented.": {
615-
"introspection.tests.IntrospectionTests.test_get_constraints",
616-
"introspection.tests.IntrospectionTests.test_get_constraints_index_types",
617-
"introspection.tests.IntrospectionTests.test_get_constraints_indexes_orders",
618-
"introspection.tests.IntrospectionTests.test_get_constraints_unique_indexes_orders",
619-
"introspection.tests.IntrospectionTests.test_get_primary_key_column",
620-
},
621596
"MongoDB can't introspect primary key.": {
597+
"introspection.tests.IntrospectionTests.test_get_primary_key_column",
622598
"schema.tests.SchemaTests.test_alter_primary_key_the_same_name",
623599
"schema.tests.SchemaTests.test_primary_key",
624600
},

django_mongodb/introspection.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,31 @@
11
from django.db.backends.base.introspection import BaseDatabaseIntrospection
2+
from django.db.models import Index
23

34

45
class DatabaseIntrospection(BaseDatabaseIntrospection):
6+
ORDER_DIR = {1: "ASC", -1: "DESC"}
7+
58
def table_names(self, cursor=None, include_views=False):
69
return sorted([x["name"] for x in self.connection.database.list_collections()])
10+
11+
def get_constraints(self, cursor, table_name):
12+
indexes = self.connection.database[table_name].index_information()
13+
constraints = {}
14+
for name, details in indexes.items():
15+
# Remove underscore prefix from "_id" columns in primary key index.
16+
if is_primary_key := name == "_id_":
17+
name = "id"
18+
details["key"] = [("id", 1)]
19+
constraints[name] = {
20+
"check": False,
21+
"columns": [field for field, order in details["key"]],
22+
"definition": None,
23+
"foreign_key": None,
24+
"index": True,
25+
"orders": [self.ORDER_DIR[order] for field, order in details["key"]],
26+
"primary_key": is_primary_key,
27+
"type": Index.suffix,
28+
"unique": details.get("unique", False),
29+
"options": {},
30+
}
31+
return constraints

django_mongodb/schema.py

Lines changed: 53 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
2+
from django.db.models import Index
3+
from pymongo.operations import IndexModel
24

35
from .query import wrap_database_errors
46

@@ -7,11 +9,30 @@ class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
79
@wrap_database_errors
810
def create_model(self, model):
911
self.connection.database.create_collection(model._meta.db_table)
12+
self._create_model_indexes(model)
1013
# Make implicit M2M tables.
1114
for field in model._meta.local_many_to_many:
1215
if field.remote_field.through._meta.auto_created:
1316
self.create_model(field.remote_field.through)
1417

18+
def _create_model_indexes(self, model):
19+
"""
20+
Create all indexes (field indexes, index_together, Meta.indexes) for
21+
the specified model.
22+
"""
23+
if not model._meta.managed or model._meta.proxy or model._meta.swapped:
24+
return
25+
# Field indexes
26+
for field in model._meta.local_fields:
27+
if self._field_should_be_indexed(model, field):
28+
self._add_field_index(model, field)
29+
# Meta.index_together (RemovedInDjango51Warning)
30+
for field_names in model._meta.index_together:
31+
self._add_composed_index(model, field_names)
32+
# Meta.indexes
33+
for index in model._meta.indexes:
34+
self.add_index(model, index)
35+
1536
def delete_model(self, model):
1637
# Delete implicit M2m tables.
1738
for field in model._meta.local_many_to_many:
@@ -29,6 +50,9 @@ def add_field(self, model, field):
2950
self.connection.database[model._meta.db_table].update_many(
3051
{}, [{"$set": {column: self.effective_default(field)}}]
3152
)
53+
# Add an index, if required.
54+
if self._field_should_be_indexed(model, field):
55+
self._add_field_index(model, field)
3256

3357
def _alter_field(
3458
self,
@@ -60,21 +84,45 @@ def remove_field(self, model, field):
6084
# Unset field on existing documents.
6185
if column := field.column:
6286
self.connection.database[model._meta.db_table].update_many({}, {"$unset": {column: ""}})
87+
if self._field_should_be_indexed(model, field):
88+
self._remove_field_index(model, field)
6389

6490
def alter_index_together(self, model, old_index_together, new_index_together):
6591
pass
6692

6793
def alter_unique_together(self, model, old_unique_together, new_unique_together):
6894
pass
6995

70-
def add_index(self, model, index):
71-
pass
96+
def add_index(self, model, index, field=None):
97+
if index.contains_expressions:
98+
return
99+
index_orders = (
100+
[(field.column, 1)]
101+
if field
102+
else [
103+
(model._meta.get_field(field_name).column, 1 if order == "" else -1)
104+
for field_name, order in index.fields_orders
105+
]
106+
)
107+
idx = IndexModel(index_orders, name=index.name)
108+
self.connection.database[model._meta.db_table].create_indexes([idx])
72109

73-
def rename_index(self, model, old_index, new_index):
74-
pass
110+
def _add_composed_index(self, model, field_names):
111+
"""Add an index on the given list of field_names."""
112+
idx = Index(fields=field_names)
113+
idx.set_name_with_model(model)
114+
self.add_index(model, idx)
115+
116+
def _add_field_index(self, model, field):
117+
"""Add an index on a field with db_index=True."""
118+
index = Index(fields=[field.name])
119+
index.name = self._create_index_name(model._meta.db_table, [field.column])
120+
self.add_index(model, index, field=field)
75121

76122
def remove_index(self, model, index):
77-
pass
123+
if index.contains_expressions:
124+
return
125+
self.connection.database[model._meta.db_table].drop_index(index.name)
78126

79127
def add_constraint(self, model, constraint):
80128
pass

0 commit comments

Comments
 (0)