11from django .db .backends .base .schema import BaseDatabaseSchemaEditor
2+ from django .db .models import Index
3+ from pymongo import ASCENDING , DESCENDING
4+ from pymongo .operations import IndexModel
25
36from .query import wrap_database_errors
47from .utils import OperationCollector
@@ -18,11 +21,30 @@ def get_database(self):
1821 @wrap_database_errors
1922 def create_model (self , model ):
2023 self .get_database ().create_collection (model ._meta .db_table )
24+ self ._create_model_indexes (model )
2125 # Make implicit M2M tables.
2226 for field in model ._meta .local_many_to_many :
2327 if field .remote_field .through ._meta .auto_created :
2428 self .create_model (field .remote_field .through )
2529
30+ def _create_model_indexes (self , model ):
31+ """
32+ Create all indexes (field indexes, index_together, Meta.indexes) for
33+ the specified model.
34+ """
35+ if not model ._meta .managed or model ._meta .proxy or model ._meta .swapped :
36+ return
37+ # Field indexes
38+ for field in model ._meta .local_fields :
39+ if self ._field_should_be_indexed (model , field ):
40+ self ._add_field_index (model , field )
41+ # Meta.index_together (RemovedInDjango51Warning)
42+ for field_names in model ._meta .index_together :
43+ self ._add_composed_index (model , field_names )
44+ # Meta.indexes
45+ for index in model ._meta .indexes :
46+ self .add_index (model , index )
47+
2648 def delete_model (self , model ):
2749 # Delete implicit M2m tables.
2850 for field in model ._meta .local_many_to_many :
@@ -40,6 +62,9 @@ def add_field(self, model, field):
4062 self .get_collection (model ._meta .db_table ).update_many (
4163 {}, [{"$set" : {column : self .effective_default (field )}}]
4264 )
65+ # Add an index, if required.
66+ if self ._field_should_be_indexed (model , field ):
67+ self ._add_field_index (model , field )
4368
4469 def _alter_field (
4570 self ,
@@ -71,21 +96,68 @@ def remove_field(self, model, field):
7196 # Unset field on existing documents.
7297 if column := field .column :
7398 self .get_collection (model ._meta .db_table ).update_many ({}, {"$unset" : {column : "" }})
99+ if self ._field_should_be_indexed (model , field ):
100+ self ._remove_field_index (model , field )
74101
75102 def alter_index_together (self , model , old_index_together , new_index_together ):
76103 pass
77104
78105 def alter_unique_together (self , model , old_unique_together , new_unique_together ):
79106 pass
80107
81- def add_index (self , model , index ):
82- pass
83-
84- def rename_index (self , model , old_index , new_index ):
85- pass
108+ def add_index (self , model , index , field = None ):
109+ if index .contains_expressions :
110+ return
111+ index_orders = (
112+ [(field .column , ASCENDING )]
113+ if field
114+ else [
115+ # order is "" if ASCENDING or "DESC" if DESCENDING (see
116+ # django.db.models.indexes.Index.fields_orders).
117+ (model ._meta .get_field (field_name ).column , ASCENDING if order == "" else DESCENDING )
118+ for field_name , order in index .fields_orders
119+ ]
120+ )
121+ idx = IndexModel (index_orders , name = index .name )
122+ self .get_collection (model ._meta .db_table ).create_indexes ([idx ])
123+
124+ def _add_composed_index (self , model , field_names ):
125+ """Add an index on the given list of field_names."""
126+ idx = Index (fields = field_names )
127+ idx .set_name_with_model (model )
128+ self .add_index (model , idx )
129+
130+ def _add_field_index (self , model , field ):
131+ """Add an index on a field with db_index=True."""
132+ index = Index (fields = [field .name ])
133+ index .name = self ._create_index_name (model ._meta .db_table , [field .column ])
134+ self .add_index (model , index , field = field )
86135
87136 def remove_index (self , model , index ):
88- pass
137+ if index .contains_expressions :
138+ return
139+ self .get_collection (model ._meta .db_table ).drop_index (index .name )
140+
141+ def _remove_field_index (self , model , field ):
142+ """Remove a field's db_index=True index."""
143+ collection = self .get_collection (model ._meta .db_table )
144+ meta_index_names = {index .name for index in model ._meta .indexes }
145+ index_names = self ._constraint_names (
146+ model ,
147+ [field .column ],
148+ index = True ,
149+ # Retrieve only BTREE indexes since this is what's created with
150+ # db_index=True.
151+ type_ = Index .suffix ,
152+ exclude = meta_index_names ,
153+ )
154+ if len (index_names ) != 1 :
155+ num_found = len (index_names )
156+ raise ValueError (
157+ f"Found wrong number ({ num_found }
158+ f"{ model ._meta .db_table } { field .column }
159+ )
160+ collection .drop_index (index_names [0 ])
89161
90162 def add_constraint (self , model , constraint ):
91163 pass
0 commit comments