From dabf9536a0ec43f852acdff56984f3d5b15c3260 Mon Sep 17 00:00:00 2001 From: Emanuel Lupi Date: Sun, 16 Feb 2025 15:38:37 -0300 Subject: [PATCH 1/4] Refactor index creation --- django_mongodb_backend/indexes.py | 43 +++++++++++++++++++++++++++++ django_mongodb_backend/schema.py | 45 ++++--------------------------- 2 files changed, 48 insertions(+), 40 deletions(-) diff --git a/django_mongodb_backend/indexes.py b/django_mongodb_backend/indexes.py index 567cd407..590f67a8 100644 --- a/django_mongodb_backend/indexes.py +++ b/django_mongodb_backend/indexes.py @@ -3,6 +3,8 @@ from django.db.models.lookups import BuiltinLookup from django.db.models.sql.query import Query from django.db.models.sql.where import AND, XOR, WhereNode +from pymongo import ASCENDING, DESCENDING +from pymongo.operations import IndexModel from .query_utils import process_rhs @@ -58,7 +60,48 @@ def where_node_idx(self, compiler, connection): return mql +def create_mongodb_index(self, model, schema_editor, field=None, unique=False, column_prefix=""): + from collections import defaultdict + + if self.contains_expressions: + return None + kwargs = {} + filter_expression = defaultdict(dict) + if self.condition: + filter_expression.update(self._get_condition_mql(model, schema_editor)) + if unique: + kwargs["unique"] = True + # Indexing on $type matches the value of most SQL databases by + # allowing multiple null values for the unique constraint. + if field: + column = column_prefix + field.column + filter_expression[column].update({"$type": field.db_type(schema_editor.connection)}) + else: + for field_name, _ in self.fields_orders: + field_ = model._meta.get_field(field_name) + filter_expression[field_.column].update( + {"$type": field_.db_type(schema_editor.connection)} + ) + if filter_expression: + kwargs["partialFilterExpression"] = filter_expression + index_orders = ( + [(column_prefix + field.column, ASCENDING)] + if field + else [ + # order is "" if ASCENDING or "DESC" if DESCENDING (see + # django.db.models.indexes.Index.fields_orders). + ( + column_prefix + model._meta.get_field(field_name).column, + ASCENDING if order == "" else DESCENDING, + ) + for field_name, order in self.fields_orders + ] + ) + return IndexModel(index_orders, name=self.name, **kwargs) + + def register_indexes(): BuiltinLookup.as_mql_idx = builtin_lookup_idx Index._get_condition_mql = _get_condition_mql + Index.create_mongodb_index = create_mongodb_index WhereNode.as_mql_idx = where_node_idx diff --git a/django_mongodb_backend/schema.py b/django_mongodb_backend/schema.py index cea89ae2..445b43d3 100644 --- a/django_mongodb_backend/schema.py +++ b/django_mongodb_backend/schema.py @@ -1,9 +1,5 @@ -from collections import defaultdict - from django.db.backends.base.schema import BaseDatabaseSchemaEditor from django.db.models import Index, UniqueConstraint -from pymongo import ASCENDING, DESCENDING -from pymongo.operations import IndexModel from .fields import EmbeddedModelField from .query import wrap_database_errors @@ -264,43 +260,12 @@ def alter_unique_together( def add_index( self, model, index, *, field=None, unique=False, column_prefix="", parent_model=None ): - if index.contains_expressions: - return - kwargs = {} - filter_expression = defaultdict(dict) - if index.condition: - filter_expression.update(index._get_condition_mql(model, self)) - if unique: - kwargs["unique"] = True - # Indexing on $type matches the value of most SQL databases by - # allowing multiple null values for the unique constraint. - if field: - column = column_prefix + field.column - filter_expression[column].update({"$type": field.db_type(self.connection)}) - else: - for field_name, _ in index.fields_orders: - field_ = model._meta.get_field(field_name) - filter_expression[field_.column].update( - {"$type": field_.db_type(self.connection)} - ) - if filter_expression: - kwargs["partialFilterExpression"] = filter_expression - index_orders = ( - [(column_prefix + field.column, ASCENDING)] - if field - else [ - # order is "" if ASCENDING or "DESC" if DESCENDING (see - # django.db.models.indexes.Index.fields_orders). - ( - column_prefix + model._meta.get_field(field_name).column, - ASCENDING if order == "" else DESCENDING, - ) - for field_name, order in index.fields_orders - ] + idx = index.create_mongodb_index( + model, self, field=field, unique=unique, column_prefix=column_prefix ) - idx = IndexModel(index_orders, name=index.name, **kwargs) - model = parent_model or model - self.get_collection(model._meta.db_table).create_indexes([idx]) + if idx: + model = parent_model or model + self.get_collection(model._meta.db_table).create_indexes([idx]) def _add_composed_index(self, model, field_names, column_prefix="", parent_model=None): """Add an index on the given list of field_names.""" From 78e4d7577d1a6c4deee2ad7f878d0b970d0f1443 Mon Sep 17 00:00:00 2001 From: Emanuel Lupi Date: Fri, 21 Mar 2025 20:37:32 -0300 Subject: [PATCH 2/4] Change method name --- django_mongodb_backend/indexes.py | 4 ++-- django_mongodb_backend/schema.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/django_mongodb_backend/indexes.py b/django_mongodb_backend/indexes.py index 590f67a8..143a4959 100644 --- a/django_mongodb_backend/indexes.py +++ b/django_mongodb_backend/indexes.py @@ -60,7 +60,7 @@ def where_node_idx(self, compiler, connection): return mql -def create_mongodb_index(self, model, schema_editor, field=None, unique=False, column_prefix=""): +def get_pymongo_index_model(self, model, schema_editor, field=None, unique=False, column_prefix=""): from collections import defaultdict if self.contains_expressions: @@ -103,5 +103,5 @@ def create_mongodb_index(self, model, schema_editor, field=None, unique=False, c def register_indexes(): BuiltinLookup.as_mql_idx = builtin_lookup_idx Index._get_condition_mql = _get_condition_mql - Index.create_mongodb_index = create_mongodb_index + Index.get_pymongo_index_model = get_pymongo_index_model WhereNode.as_mql_idx = where_node_idx diff --git a/django_mongodb_backend/schema.py b/django_mongodb_backend/schema.py index 445b43d3..8908a6bf 100644 --- a/django_mongodb_backend/schema.py +++ b/django_mongodb_backend/schema.py @@ -260,7 +260,7 @@ def alter_unique_together( def add_index( self, model, index, *, field=None, unique=False, column_prefix="", parent_model=None ): - idx = index.create_mongodb_index( + idx = index.get_pymongo_index_model( model, self, field=field, unique=unique, column_prefix=column_prefix ) if idx: From 86be3d4b66ee0cbeee2d40247d74f8b141620711 Mon Sep 17 00:00:00 2001 From: Emanuel Lupi Date: Fri, 21 Mar 2025 21:35:05 -0300 Subject: [PATCH 3/4] Add docstring --- django_mongodb_backend/indexes.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/django_mongodb_backend/indexes.py b/django_mongodb_backend/indexes.py index 143a4959..81fe17f1 100644 --- a/django_mongodb_backend/indexes.py +++ b/django_mongodb_backend/indexes.py @@ -1,3 +1,5 @@ +from collections import defaultdict + from django.db import NotSupportedError from django.db.models import Index from django.db.models.lookups import BuiltinLookup @@ -61,8 +63,7 @@ def where_node_idx(self, compiler, connection): def get_pymongo_index_model(self, model, schema_editor, field=None, unique=False, column_prefix=""): - from collections import defaultdict - + """Return a pymongo IndexModel for this Django Index.""" if self.contains_expressions: return None kwargs = {} From ca0d966d4a7ce8c0e1ffb88d89f1e7c6d7694c7a Mon Sep 17 00:00:00 2001 From: Emanuel Lupi Date: Fri, 21 Mar 2025 21:39:58 -0300 Subject: [PATCH 4/4] Edits --- django_mongodb_backend/indexes.py | 44 +++++++++++++++---------------- django_mongodb_backend/schema.py | 2 +- 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/django_mongodb_backend/indexes.py b/django_mongodb_backend/indexes.py index 81fe17f1..8ddceda7 100644 --- a/django_mongodb_backend/indexes.py +++ b/django_mongodb_backend/indexes.py @@ -40,28 +40,6 @@ def builtin_lookup_idx(self, compiler, connection): return {lhs_mql: {operator: value}} -def where_node_idx(self, compiler, connection): - if self.connector == AND: - operator = "$and" - elif self.connector == XOR: - raise NotSupportedError("MongoDB does not support the '^' operator lookup in indexes.") - else: - operator = "$or" - if self.negated: - raise NotSupportedError("MongoDB does not support the '~' operator in indexes.") - children_mql = [] - for child in self.children: - mql = child.as_mql_idx(compiler, connection) - children_mql.append(mql) - if len(children_mql) == 1: - mql = children_mql[0] - elif len(children_mql) > 1: - mql = {operator: children_mql} - else: - mql = {} - return mql - - def get_pymongo_index_model(self, model, schema_editor, field=None, unique=False, column_prefix=""): """Return a pymongo IndexModel for this Django Index.""" if self.contains_expressions: @@ -101,6 +79,28 @@ def get_pymongo_index_model(self, model, schema_editor, field=None, unique=False return IndexModel(index_orders, name=self.name, **kwargs) +def where_node_idx(self, compiler, connection): + if self.connector == AND: + operator = "$and" + elif self.connector == XOR: + raise NotSupportedError("MongoDB does not support the '^' operator lookup in indexes.") + else: + operator = "$or" + if self.negated: + raise NotSupportedError("MongoDB does not support the '~' operator in indexes.") + children_mql = [] + for child in self.children: + mql = child.as_mql_idx(compiler, connection) + children_mql.append(mql) + if len(children_mql) == 1: + mql = children_mql[0] + elif len(children_mql) > 1: + mql = {operator: children_mql} + else: + mql = {} + return mql + + def register_indexes(): BuiltinLookup.as_mql_idx = builtin_lookup_idx Index._get_condition_mql = _get_condition_mql diff --git a/django_mongodb_backend/schema.py b/django_mongodb_backend/schema.py index 8908a6bf..8ae60918 100644 --- a/django_mongodb_backend/schema.py +++ b/django_mongodb_backend/schema.py @@ -261,7 +261,7 @@ def add_index( self, model, index, *, field=None, unique=False, column_prefix="", parent_model=None ): idx = index.get_pymongo_index_model( - model, self, field=field, unique=unique, column_prefix=column_prefix + model, schema_editor=self, field=field, unique=unique, column_prefix=column_prefix ) if idx: model = parent_model or model