|
1 | | -from functools import singledispatchmethod |
2 | | - |
3 | 1 | from django.db.backends.base.schema import BaseDatabaseSchemaEditor |
4 | 2 | from django.db.models import Index, UniqueConstraint |
5 | 3 | from pymongo.operations import IndexModel, SearchIndexModel |
@@ -261,18 +259,13 @@ def alter_unique_together( |
261 | 259 | model, constraint, parent_model=parent_model, column_prefix=column_prefix |
262 | 260 | ) |
263 | 261 |
|
264 | | - @singledispatchmethod |
265 | 262 | def _add_index(self, index, model): |
| 263 | + if isinstance(index, SearchIndexModel): |
| 264 | + return self.get_collection(model._meta.db_table).create_search_index(index) |
| 265 | + if isinstance(index, IndexModel): |
| 266 | + return self.get_collection(model._meta.db_table).create_indexes([index]) |
266 | 267 | raise ValueError(f"{type(index)} isn't a supported index type") |
267 | 268 |
|
268 | | - @_add_index.register |
269 | | - def _(self, index: IndexModel, model): |
270 | | - return self.get_collection(model._meta.db_table).create_indexes([index]) |
271 | | - |
272 | | - @_add_index.register |
273 | | - def _(self, index: SearchIndexModel, model): |
274 | | - return self.get_collection(model._meta.db_table).create_search_index(index) |
275 | | - |
276 | 269 | @ignore_embedded_models |
277 | 270 | def add_index( |
278 | 271 | self, model, index, *, field=None, unique=False, column_prefix="", parent_model=None |
@@ -301,18 +294,13 @@ def _add_field_index(self, model, field, *, column_prefix=""): |
301 | 294 | index.name = self._create_index_name(model._meta.db_table, [column_prefix + field.column]) |
302 | 295 | self.add_index(model, index, field=field, column_prefix=column_prefix) |
303 | 296 |
|
304 | | - @singledispatchmethod |
305 | 297 | def _remove_index(self, index, model): |
| 298 | + if isinstance(index, AtlasSearchIndex | AtlasVectorSearchIndex): |
| 299 | + return self.get_collection(model._meta.db_table).drop_search_index(index.name) |
| 300 | + if isinstance(index, Index): |
| 301 | + return self.get_collection(model._meta.db_table).drop_index(index.name) |
306 | 302 | raise ValueError(f"{type(index)} isn't a supported index type") |
307 | 303 |
|
308 | | - @_remove_index.register |
309 | | - def _(self, index: Index, model): |
310 | | - return self.get_collection(model._meta.db_table).drop_index(index.name) |
311 | | - |
312 | | - @_remove_index.register |
313 | | - def _(self, index: AtlasSearchIndex | AtlasVectorSearchIndex, model): |
314 | | - return self.get_collection(model._meta.db_table).drop_search_index(index.name) |
315 | | - |
316 | 304 | @ignore_embedded_models |
317 | 305 | def remove_index(self, model, index): |
318 | 306 | if index.contains_expressions: |
|
0 commit comments