Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add support for partial indexes #159

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/test-python.yml
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ jobs:
from_db_value
generic_relations
generic_relations_regress
indexes
introspection
known_related_objects
lookup
Expand Down
2 changes: 2 additions & 0 deletions django_mongodb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,14 @@
from .expressions import register_expressions # noqa: E402
from .fields import register_fields # noqa: E402
from .functions import register_functions # noqa: E402
from .indexes import register_indexes # noqa: E402
from .lookups import register_lookups # noqa: E402
from .query import register_nodes # noqa: E402

register_aggregates()
register_expressions()
register_fields()
register_functions()
register_indexes()
register_lookups()
register_nodes()
3 changes: 2 additions & 1 deletion django_mongodb/expressions.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,8 @@ def case(self, compiler, connection):

def col(self, compiler, connection): # noqa: ARG001
# Add the column's collection's alias for columns in joined collections.
prefix = f"{self.alias}." if self.alias != compiler.collection_name else ""
has_alias = self.alias and self.alias != compiler.collection_name
prefix = f"{self.alias}." if has_alias else ""
return f"${prefix}{self.target.column}"


Expand Down
44 changes: 4 additions & 40 deletions django_mongodb/features.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
supports_collation_on_charfield = False
supports_column_check_constraints = False
supports_date_lookup_using_string = False
supports_deferrable_unique_constraints = False
supports_explaining_query_execution = True
supports_expression_defaults = False
supports_expression_indexes = False
Expand Down Expand Up @@ -69,28 +70,9 @@ class DatabaseFeatures(BaseDatabaseFeatures):
"backends.tests.ThreadTests.test_pass_connection_between_threads",
"backends.tests.ThreadTests.test_closing_non_shared_connections",
"backends.tests.ThreadTests.test_default_connection_thread_local",
# AddField
"schema.tests.SchemaTests.test_add_unique_charfield",
# AlterField
"schema.tests.SchemaTests.test_alter_field_fk_to_o2o",
"schema.tests.SchemaTests.test_alter_field_o2o_keeps_unique",
"schema.tests.SchemaTests.test_alter_field_o2o_to_fk",
"schema.tests.SchemaTests.test_alter_int_pk_to_int_unique",
# AlterField (unique)
"schema.tests.SchemaTests.test_indexes",
"schema.tests.SchemaTests.test_unique",
"schema.tests.SchemaTests.test_unique_and_reverse_m2m",
# alter_unique_together
"migrations.test_operations.OperationTests.test_alter_unique_together",
"schema.tests.SchemaTests.test_unique_together",
# add/remove_constraint
"introspection.tests.IntrospectionTests.test_get_constraints",
"migrations.test_operations.OperationTests.test_add_partial_unique_constraint",
"migrations.test_operations.OperationTests.test_create_model_with_partial_unique_constraint",
"migrations.test_operations.OperationTests.test_remove_partial_unique_constraint",
"schema.tests.SchemaTests.test_composed_constraint_with_fk",
"schema.tests.SchemaTests.test_remove_ignored_unique_constraint_not_create_fk_index",
"schema.tests.SchemaTests.test_unique_constraint",
# TODO:
"indexes.tests.PartialIndexTests.test_is_null_condition",
"indexes.tests.PartialIndexTests.test_multiple_conditions",
}
# $bitAnd, #bitOr, and $bitXor are new in MongoDB 6.3.
_django_test_expected_failures_bitwise = {
Expand Down Expand Up @@ -181,24 +163,6 @@ def django_test_expected_failures(self):
"model_fields.test_autofield.SmallAutoFieldTests",
"queries.tests.TestInvalidValuesRelation.test_invalid_values",
},
"MongoDB does not enforce UNIQUE constraints.": {
"auth_tests.test_basic.BasicTestCase.test_unicode_username",
"auth_tests.test_migrations.ProxyModelWithSameAppLabelTests.test_migrate_with_existing_target_permission",
"constraints.tests.UniqueConstraintTests.test_database_constraint",
"contenttypes_tests.test_operations.ContentTypeOperationsTests.test_content_type_rename_conflict",
"contenttypes_tests.test_operations.ContentTypeOperationsTests.test_existing_content_type_rename",
"custom_pk.tests.CustomPKTests.test_unique_pk",
"force_insert_update.tests.ForceInsertInheritanceTests.test_force_insert_with_existing_grandparent",
"get_or_create.tests.GetOrCreateTestsWithManualPKs.test_create_with_duplicate_primary_key",
"get_or_create.tests.GetOrCreateTestsWithManualPKs.test_savepoint_rollback",
"get_or_create.tests.GetOrCreateThroughManyToMany.test_something",
"get_or_create.tests.UpdateOrCreateTests.test_manual_primary_key_test",
"get_or_create.tests.UpdateOrCreateTestsWithManualPKs.test_create_with_duplicate_primary_key",
"introspection.tests.IntrospectionTests.test_get_constraints_unique_indexes_orders",
"model_fields.test_filefield.FileFieldTests.test_unique_when_same_filename",
"one_to_one.tests.OneToOneTests.test_multiple_o2o",
"queries.test_bulk_update.BulkUpdateTests.test_database_routing_batch_atomicity",
},
"MongoDB does not enforce PositiveIntegerField constraint.": {
"model_fields.test_integerfield.PositiveIntegerFieldTests.test_negative_values",
},
Expand Down
23 changes: 23 additions & 0 deletions django_mongodb/indexes.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
from django.db.models import Index
from django.db.models.sql.query import Query


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)
mql_ = where.as_mql(compiler, schema_editor.connection)
# Transform aggregate() query syntax into find() syntax.
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This stinks and doesn't work for multiple conditions / null conditions (but maybe could be done). I wonder if it would be better to have a separate path from as_mql() to generate find() compatible MQL rather than trying to convert from the aggregate syntax.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As I understand it, we're pulling the index conditions from a model's field definitions using the compiler built by the schema editor? Is this schema_editor compiler different from our Mongo one?

can self.condition ever be a list of items?

Speaking more broadly, I'd say, yes, let's include the find compatible conversions, but how many nodes/expressions would we need that for? Also, since find is much more limited there would still be gaps.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

query.get_compiler() is the same compiler we use for other queries.

condition is a Q object: "For example, condition=Q(pages__gt=400) indexes records with more than 400 pages."

What doesn't currently work is multiple conditions like:

Q(
    pub_date__gt=datetime.datetime(
        year=2015,
        month=1,
        day=1,
        tzinfo=timezone.get_current_timezone(),
    )
)
& Q(headline__contains="China")

Maybe that can be translated without too much difficulty. I didn't think about it too much.

Obviously, we can only support conditions that MongoDB supports.

mql = {}
for key in mql_:
col, value = mql_[key]
# multiple conditions don't work yet
if isinstance(col, dict):
return {}
mql[col.lstrip("$")] = {key: value}
return mql


def register_indexes():
Index._get_condition_mql = _get_condition_mql
132 changes: 121 additions & 11 deletions django_mongodb/schema.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from collections import defaultdict

from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.models import Index
from django.db.models import Index, UniqueConstraint
from pymongo import ASCENDING, DESCENDING
from pymongo.operations import IndexModel

Expand Down Expand Up @@ -29,18 +31,26 @@ def create_model(self, model):

def _create_model_indexes(self, model):
"""
Create all indexes (field indexes, index_together, Meta.indexes) for
the specified model.
Create all indexes (field indexes & uniques, Meta.index_together,
Meta.unique_together, Meta.constraints, Meta.indexes) for the model.
"""
if not model._meta.managed or model._meta.proxy or model._meta.swapped:
return
# Field indexes
# Field indexes and uniques
for field in model._meta.local_fields:
if self._field_should_be_indexed(model, field):
self._add_field_index(model, field)
elif self._field_should_have_unique(field):
self._add_field_unique(model, field)
# Meta.index_together (RemovedInDjango51Warning)
for field_names in model._meta.index_together:
self._add_composed_index(model, field_names)
# Meta.unique_together
if model._meta.unique_together:
self.alter_unique_together(model, [], model._meta.unique_together)
# Meta.constraints
for constraint in model._meta.constraints:
self.add_constraint(model, constraint)
# Meta.indexes
for index in model._meta.indexes:
self.add_index(model, index)
Expand All @@ -62,9 +72,11 @@ def add_field(self, model, field):
self.get_collection(model._meta.db_table).update_many(
{}, [{"$set": {column: self.effective_default(field)}}]
)
# Add an index, if required.
# Add an index or unique, if required.
if self._field_should_be_indexed(model, field):
self._add_field_index(model, field)
elif self._field_should_have_unique(field):
self._add_field_unique(model, field)

def _alter_field(
self,
Expand All @@ -78,6 +90,11 @@ def _alter_field(
strict=False,
):
collection = self.get_collection(model._meta.db_table)
# Has unique been removed?
old_field_unique = self._field_should_have_unique(old_field)
new_field_unique = self._field_should_have_unique(new_field)
if old_field_unique and not new_field_unique:
self._remove_field_unique(model, old_field)
# Removed an index?
old_field_indexed = self._field_should_be_indexed(model, old_field)
new_field_indexed = self._field_should_be_indexed(model, new_field)
Expand All @@ -90,6 +107,10 @@ def _alter_field(
if old_field_indexed and new_field_indexed:
self._remove_field_index(model, old_field)
self._add_field_index(model, new_field)
# Move unique to the new field, if needed.
if old_field_unique and new_field_unique:
self._remove_field_unique(model, old_field)
self._add_field_unique(model, new_field)
# Replace NULL with the field default if the field and was changed from
# NULL to NOT NULL.
if new_field.has_default() and old_field.null and not new_field.null:
Expand All @@ -99,6 +120,9 @@ def _alter_field(
# Added an index?
if not old_field_indexed and new_field_indexed:
self._add_field_index(model, new_field)
# Added a unique?
if not old_field_unique and new_field_unique:
self._add_field_unique(model, new_field)

def remove_field(self, model, field):
# Remove implicit M2M tables.
Expand All @@ -110,6 +134,8 @@ def remove_field(self, model, field):
self.get_collection(model._meta.db_table).update_many({}, {"$unset": {column: ""}})
if self._field_should_be_indexed(model, field):
self._remove_field_index(model, field)
elif self._field_should_have_unique(field):
self._remove_field_unique(model, field)

def alter_index_together(self, model, old_index_together, new_index_together):
olds = {tuple(fields) for fields in old_index_together}
Expand All @@ -122,11 +148,43 @@ def alter_index_together(self, model, old_index_together, new_index_together):
self._add_composed_index(model, field_names)

def alter_unique_together(self, model, old_unique_together, new_unique_together):
pass
olds = {tuple(fields) for fields in old_unique_together}
news = {tuple(fields) for fields in new_unique_together}
# Deleted uniques
for field_names in olds.difference(news):
self._remove_composed_index(
model,
field_names,
{"unique": True, "primary_key": False},
)
# Created uniques
for field_names in news.difference(olds):
columns = [model._meta.get_field(field).column for field in field_names]
name = str(self._unique_constraint_name(model._meta.db_table, columns))
constraint = UniqueConstraint(fields=field_names, name=name)
self.add_constraint(model, constraint)

def add_index(self, model, index, field=None):
def add_index(self, model, index, field=None, unique=False):
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
if field:
filter_expression[field.column].update({"$type": field.db_type(self.connection)})
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need to maintain type? if so, why?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment is below: "Use partialFilterExpression to allow multiple null values for a unique constraint." It's not really in the best spot and should be amended to reference $type since partialFilterExpression is also used for condition.

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)}
)
# Use partialFilterExpression to allow multiple null values for a
# unique constraint.
if filter_expression:
kwargs["partialFilterExpression"] = filter_expression
index_orders = (
[(field.column, ASCENDING)]
if field
Expand All @@ -137,7 +195,7 @@ def add_index(self, model, index, field=None):
for field_name, order in index.fields_orders
]
)
idx = IndexModel(index_orders, name=index.name)
idx = IndexModel(index_orders, name=index.name, **kwargs)
self.get_collection(model._meta.db_table).create_indexes([idx])

def _add_composed_index(self, model, field_names):
Expand Down Expand Up @@ -202,13 +260,65 @@ def _remove_field_index(self, model, field):
)
collection.drop_index(index_names[0])

def add_constraint(self, model, constraint):
pass
def add_constraint(self, model, constraint, field=None):
if isinstance(constraint, UniqueConstraint) and self._unique_supported(
condition=constraint.condition,
deferrable=constraint.deferrable,
include=constraint.include,
expressions=constraint.expressions,
nulls_distinct=constraint.nulls_distinct,
):
idx = Index(
fields=constraint.fields,
condition=constraint.condition,
name=constraint.name,
)
self.add_index(model, idx, field=field, unique=True)

def _add_field_unique(self, model, field):
name = str(self._unique_constraint_name(model._meta.db_table, [field.column]))
constraint = UniqueConstraint(fields=[field.name], name=name)
self.add_constraint(model, constraint, field=field)

def remove_constraint(self, model, constraint):
pass
if isinstance(constraint, UniqueConstraint) and self._unique_supported(
condition=constraint.condition,
deferrable=constraint.deferrable,
include=constraint.include,
expressions=constraint.expressions,
nulls_distinct=constraint.nulls_distinct,
):
idx = Index(
fields=constraint.fields,
condition=constraint.condition,
name=constraint.name,
)
self.remove_index(model, idx)

def _remove_field_unique(self, model, field):
# Find the unique constraint for this field
meta_constraint_names = {constraint.name for constraint in model._meta.constraints}
constraint_names = self._constraint_names(
model,
[field.column],
unique=True,
primary_key=False,
exclude=meta_constraint_names,
)
if len(constraint_names) != 1:
num_found = len(constraint_names)
raise ValueError(
f"Found wrong number ({num_found}) of unique constraints for "
f"{model._meta.db_table}.{field.column}."
)
self.get_collection(model._meta.db_table).drop_index(constraint_names[0])

def alter_db_table(self, model, old_db_table, new_db_table):
if old_db_table == new_db_table:
return
self.get_collection(old_db_table).rename(new_db_table)

def _field_should_have_unique(self, field):
db_type = field.db_type(self.connection)
# The _id column is automatically unique.
return db_type and field.unique and field.column != "_id"
Loading