-
Notifications
You must be signed in to change notification settings - Fork 6
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
base: main
Are you sure you want to change the base?
Changes from all commits
8e696cb
cf3e82a
85d32a5
575957c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
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. | ||
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 |
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 | ||
|
||
|
@@ -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) | ||
|
@@ -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, | ||
|
@@ -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) | ||
|
@@ -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: | ||
|
@@ -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. | ||
|
@@ -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} | ||
|
@@ -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)}) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. do we need to maintain type? if so, why? There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
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 | ||
|
@@ -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): | ||
|
@@ -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" |
There was a problem hiding this comment.
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 generatefind()
compatible MQL rather than trying to convert from the aggregate syntax.There was a problem hiding this comment.
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, sincefind
is much more limited there would still be gaps.There was a problem hiding this comment.
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 aQ
object: "For example,condition=Q(pages__gt=400)
indexes records with more than 400 pages."What doesn't currently work is multiple conditions like:
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.