-
Notifications
You must be signed in to change notification settings - Fork 12
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Co-authored-by: Emanuel Lupi <[email protected]>
- Loading branch information
Showing
10 changed files
with
262 additions
and
14 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,74 @@ | ||
| from django.db import NotSupportedError | ||
| from django.db.models import Index | ||
| from django.db.models.fields.related_lookups import In | ||
| 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 .query_utils import process_rhs | ||
|
|
||
| MONGO_INDEX_OPERATORS = { | ||
| "exact": "$eq", | ||
| "gt": "$gt", | ||
| "gte": "$gte", | ||
| "lt": "$lt", | ||
| "lte": "$lte", | ||
| "in": "$in", | ||
| } | ||
|
|
||
|
|
||
| def _get_condition_mql(self, model, schema_editor): | ||
| """Analogous to Index._get_condition_sql().""" | ||
| query = Query(model=model, alias_cols=False) | ||
| where = query.build_where(self.condition) | ||
| compiler = query.get_compiler(connection=schema_editor.connection) | ||
| return where.as_mql_idx(compiler, schema_editor.connection) | ||
|
|
||
|
|
||
| def builtin_lookup_idx(self, compiler, connection): | ||
| lhs_mql = self.lhs.target.column | ||
| value = process_rhs(self, compiler, connection) | ||
| try: | ||
| operator = MONGO_INDEX_OPERATORS[self.lookup_name] | ||
| except KeyError: | ||
| raise NotSupportedError( | ||
| f"MongoDB does not support the '{self.lookup_name}' lookup in indexes." | ||
| ) from None | ||
| return {lhs_mql: {operator: value}} | ||
|
|
||
|
|
||
| def in_idx(self, compiler, connection): | ||
| if not connection.features.supports_in_index_operator: | ||
| raise NotSupportedError("MongoDB < 6.0 does not support the 'in' lookup in indexes.") | ||
| return builtin_lookup_idx(self, compiler, connection) | ||
|
|
||
|
|
||
| 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: | ||
| if not connection.features.supports_in_index_operator: | ||
| raise NotSupportedError("MongoDB < 6.0 does not support the '|' operator in indexes.") | ||
| 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 | ||
| In.as_mql_idx = in_idx | ||
| Index._get_condition_mql = _get_condition_mql | ||
| WhereNode.as_mql_idx = where_node_idx |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,7 @@ | ||
| from django.db import models | ||
|
|
||
|
|
||
| class Article(models.Model): | ||
| headline = models.CharField(max_length=100) | ||
| number = models.IntegerField() | ||
| body = models.TextField() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,126 @@ | ||
| import operator | ||
|
|
||
| from django.db import NotSupportedError, connection | ||
| from django.db.models import Index, Q | ||
| from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature | ||
|
|
||
| from .models import Article | ||
|
|
||
|
|
||
| class PartialIndexTests(TestCase): | ||
| def assertAddRemoveIndex(self, editor, model, index): | ||
| editor.add_index(index=index, model=model) | ||
| self.assertIn( | ||
| index.name, | ||
| connection.introspection.get_constraints( | ||
| cursor=None, | ||
| table_name=model._meta.db_table, | ||
| ), | ||
| ) | ||
| editor.remove_index(index=index, model=model) | ||
|
|
||
| def test_not_supported(self): | ||
| msg = "MongoDB does not support the 'isnull' lookup in indexes." | ||
| with connection.schema_editor() as editor, self.assertRaisesMessage(NotSupportedError, msg): | ||
| Index( | ||
| name="test", | ||
| fields=["headline"], | ||
| condition=Q(pk__isnull=True), | ||
| )._get_condition_mql(Article, schema_editor=editor) | ||
|
|
||
| def test_negated_not_supported(self): | ||
| msg = "MongoDB does not support the '~' operator in indexes." | ||
| with self.assertRaisesMessage(NotSupportedError, msg), connection.schema_editor() as editor: | ||
| Index( | ||
| name="test", | ||
| fields=["headline"], | ||
| condition=~Q(pk=True), | ||
| )._get_condition_mql(Article, schema_editor=editor) | ||
|
|
||
| def test_xor_not_supported(self): | ||
| msg = "MongoDB does not support the '^' operator lookup in indexes." | ||
| with self.assertRaisesMessage(NotSupportedError, msg), connection.schema_editor() as editor: | ||
| Index( | ||
| name="test", | ||
| fields=["headline"], | ||
| condition=Q(pk=True) ^ Q(pk=False), | ||
| )._get_condition_mql(Article, schema_editor=editor) | ||
|
|
||
| @skipIfDBFeature("supports_or_index_operator") | ||
| def test_or_not_supported(self): | ||
| msg = "MongoDB < 6.0 does not support the '|' operator in indexes." | ||
| with self.assertRaisesMessage(NotSupportedError, msg), connection.schema_editor() as editor: | ||
| Index( | ||
| name="test", | ||
| fields=["headline"], | ||
| condition=Q(pk=True) | Q(pk=False), | ||
| )._get_condition_mql(Article, schema_editor=editor) | ||
|
|
||
| @skipIfDBFeature("supports_in_index_operator") | ||
| def test_in_not_supported(self): | ||
| msg = "MongoDB < 6.0 does not support the 'in' lookup in indexes." | ||
| with self.assertRaisesMessage(NotSupportedError, msg), connection.schema_editor() as editor: | ||
| Index( | ||
| name="test", | ||
| fields=["headline"], | ||
| condition=Q(pk__in=[True]), | ||
| )._get_condition_mql(Article, schema_editor=editor) | ||
|
|
||
| def test_operations(self): | ||
| operators = ( | ||
| ("gt", "$gt"), | ||
| ("gte", "$gte"), | ||
| ("lt", "$lt"), | ||
| ("lte", "$lte"), | ||
| ) | ||
| for op, mongo_operator in operators: | ||
| with self.subTest(operator=op), connection.schema_editor() as editor: | ||
| index = Index( | ||
| name="test", | ||
| fields=["headline"], | ||
| condition=Q(**{f"number__{op}": 3}), | ||
| ) | ||
| self.assertEqual( | ||
| {"number": {mongo_operator: 3}}, | ||
| index._get_condition_mql(Article, schema_editor=editor), | ||
| ) | ||
| self.assertAddRemoveIndex(editor, Article, index) | ||
|
|
||
| @skipUnlessDBFeature("supports_in_index_operator") | ||
| def test_composite_index(self): | ||
| with connection.schema_editor() as editor: | ||
| index = Index( | ||
| name="test", | ||
| fields=["headline"], | ||
| condition=Q(number__gte=3) & (Q(body__gt="test1") | Q(body__in=["A", "B"])), | ||
| ) | ||
| self.assertEqual( | ||
| index._get_condition_mql(Article, schema_editor=editor), | ||
| { | ||
| "$and": [ | ||
| {"number": {"$gte": 3}}, | ||
| {"$or": [{"body": {"$gt": "test1"}}, {"body": {"$in": ["A", "B"]}}]}, | ||
| ] | ||
| }, | ||
| ) | ||
| self.assertAddRemoveIndex(editor, Article, index) | ||
|
|
||
| def test_composite_op_index(self): | ||
| operators = ( | ||
| (operator.or_, "$or"), | ||
| (operator.and_, "$and"), | ||
| ) | ||
| if not connection.features.supports_or_index_operator: | ||
| operators = operators[1:] | ||
| for op, mongo_operator in operators: | ||
| with self.subTest(operator=op), connection.schema_editor() as editor: | ||
| index = Index( | ||
| name="test", | ||
| fields=["headline"], | ||
| condition=op(Q(number__gte=3), Q(body__gt="test1")), | ||
| ) | ||
| self.assertEqual( | ||
| {mongo_operator: [{"number": {"$gte": 3}}, {"body": {"$gt": "test1"}}]}, | ||
| index._get_condition_mql(Article, schema_editor=editor), | ||
| ) | ||
| self.assertAddRemoveIndex(editor, Article, index) |