Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@timgraham
timgraham committed Dec 30, 2024
1 parent a52c33e commit 4d9c6c6
Show file tree
Hide file tree
Showing 2 changed files with 194 additions and 11 deletions.
72 changes: 63 additions & 9 deletions django_mongodb/schema.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,6 @@ def add_field(self, model, field):
self._create_model_indexes(
field.embedded_model, parent_model=model, column_prefix=new_path
)

# Add an index or unique, if required.
if self._field_should_be_indexed(model, field):
self._add_field_index(model, field)
Expand Down Expand Up @@ -162,13 +161,65 @@ def remove_field(self, model, field):
self._remove_field_index(model, field)
elif self._field_should_have_unique(field):
self._remove_field_unique(model, field)
if isinstance(field, EmbeddedModelField):
new_path = f"{field.column}."
self._remove_model_indexes(
field.embedded_model, parent_model=model, column_prefix=new_path
)

def _remove_model_indexes(self, model, column_prefix="", parent_model=None):
"""
When removing an EmbeddedModelField, the indexes need to be removed
recursively.
"""
if not model._meta.managed or model._meta.proxy or model._meta.swapped:
return
# Field indexes and uniques
for field in model._meta.local_fields:
if isinstance(field, EmbeddedModelField):
new_path = f"{column_prefix}{field.column}."
self._remove_model_indexes(
field.embedded_model, parent_model=parent_model or model, column_prefix=new_path
)
if self._field_should_be_indexed(model, field):
self._remove_field_index(parent_model or model, field, column_prefix=column_prefix)
elif self._field_should_have_unique(field):
self._remove_field_unique(parent_model or model, field, column_prefix=column_prefix)
# Meta.index_together (RemovedInDjango51Warning)
for field_names in model._meta.index_together:
self._remove_composed_index(
model,
field_names,
{"index": True, "unique": False},
column_prefix=column_prefix,
parent_model=parent_model,
)
# Meta.unique_together
if model._meta.unique_together:
self.alter_unique_together(
model,
model._meta.unique_together,
[],
column_prefix=column_prefix,
parent_model=parent_model,
)
# Meta.constraints
for constraint in model._meta.constraints:
self.remove_constraint(
model, constraint, column_prefix=column_prefix, parent_model=parent_model
)
# Meta.indexes
for index in model._meta.indexes:
self.remove_index(model, index, column_prefix=column_prefix, parent_model=parent_model)

def alter_index_together(self, model, old_index_together, new_index_together, column_prefix=""):
olds = {tuple(fields) for fields in old_index_together}
news = {tuple(fields) for fields in new_index_together}
# Deleted indexes
for field_names in olds.difference(news):
self._remove_composed_index(model, field_names, {"index": True, "unique": False})
self._remove_composed_index(
model, field_names, {"index": True, "unique": False}, column_prefix=""
)
# Created indexes
for field_names in news.difference(olds):
self._add_composed_index(model, field_names, column_prefix=column_prefix)
Expand Down Expand Up @@ -256,16 +307,18 @@ def remove_index(self, model, index):
return
self.get_collection(model._meta.db_table).drop_index(index.name)

def _remove_composed_index(self, model, field_names, constraint_kwargs):
def _remove_composed_index(
self, model, field_names, constraint_kwargs, column_prefix="", parent_model=None
):
"""
Remove the index on the given list of field_names created by
index/unique_together, depending on constraint_kwargs.
"""
meta_constraint_names = {constraint.name for constraint in model._meta.constraints}
meta_index_names = {constraint.name for constraint in model._meta.indexes}
columns = [model._meta.get_field(field).column for field in field_names]
columns = [column_prefix + model._meta.get_field(field).column for field in field_names]
constraint_names = self._constraint_names(
model,
parent_model or model,
columns,
exclude=meta_constraint_names | meta_index_names,
**constraint_kwargs,
Expand All @@ -277,16 +330,17 @@ def _remove_composed_index(self, model, field_names, constraint_kwargs):
f"Found wrong number ({num_found}) of constraints for "
f"{model._meta.db_table}({columns_str})."
)
model = parent_model or model
collection = self.get_collection(model._meta.db_table)
collection.drop_index(constraint_names[0])

def _remove_field_index(self, model, field):
def _remove_field_index(self, model, field, column_prefix=""):
"""Remove a field's db_index=True index."""
collection = self.get_collection(model._meta.db_table)
meta_index_names = {index.name for index in model._meta.indexes}
index_names = self._constraint_names(
model,
[field.column],
[column_prefix + field.column],
index=True,
# Retrieve only BTREE indexes since this is what's created with
# db_index=True.
Expand Down Expand Up @@ -345,12 +399,12 @@ def remove_constraint(self, model, constraint):
)
self.remove_index(model, idx)

def _remove_field_unique(self, model, field):
def _remove_field_unique(self, model, field, column_prefix=""):
# 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],
[column_prefix + field.column],
unique=True,
primary_key=False,
exclude=meta_constraint_names,
Expand Down
133 changes: 131 additions & 2 deletions tests/schema_/test_embedded_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -392,9 +392,128 @@ class Meta:
editor.delete_model(Author)
self.assertTableNotExists(Author)

# SchemaEditor.add_field() tests
# SchemaEditor.add_field() / remove_field() tests
@isolate_apps("schema_")
def test_add_field_db_index(self):
def test_add_remove_field_db_index_and_unique(self):
"""AddField + EmbeddedModelField"""

class Book(models.Model):
name = models.CharField(max_length=100)

class Meta:
app_label = "schema_"

new_field = EmbeddedModelField(Author)
new_field.set_attributes_from_name("author")
with connection.schema_editor() as editor:
# Create the table amd add the field.
editor.create_model(Book)
editor.add_field(Book, new_field)
# Embedded indexes are created.
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.age"]),
["schema__book_author.age_dc08100b"],
)
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.address.zip_code"]),
["schema__book_author.address.zip_code_7b9a9307"],
)
# Embedded uniques
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.employee_id"]),
["schema__book_author.employee_id_7d4d3eff_uniq"],
)
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.address.uid"]),
["schema__book_author.address.uid_8124a01f_uniq"],
)
editor.remove_field(Book, new_field)
# Embedded indexes are removed.
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.age"]),
[],
)
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.address.zip_code"]),
[],
)
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.employee_id"]),
[],
)
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.address.uid"]),
[],
)
editor.delete_model(Book)
self.assertTableNotExists(Author)

@ignore_warnings(category=RemovedInDjango51Warning)
@isolate_apps("schema_")
def test_add_remove_field_index_together(self):
"""Meta.index_together on an embedded model."""

class Address(models.Model):
index_together_one = models.CharField(max_length=10)
index_together_two = models.CharField(max_length=10)

class Meta:
app_label = "schema_"
index_together = [("index_together_one", "index_together_two")]

class Author(models.Model):
address = EmbeddedModelField(Address)
index_together_three = models.CharField(max_length=10)
index_together_four = models.CharField(max_length=10)

class Meta:
app_label = "schema_"
index_together = [("index_together_three", "index_together_four")]

class Book(models.Model):
class Meta:
app_label = "schema_"

new_field = EmbeddedModelField(Author)
new_field.set_attributes_from_name("author")
with connection.schema_editor() as editor:
# Create the table amd add the field.
editor.create_model(Book)
editor.add_field(Book, new_field)
# Embedded index_togethers are created.
self.assertEqual(
self.get_constraints_for_columns(
Book, ["author.address.index_together_one", "author.address.index_together_two"]
),
["schema__add_index_t_efa93e_idx"],
)
self.assertEqual(
self.get_constraints_for_columns(
Book,
["author.index_together_three", "author.index_together_four"],
),
["schema__aut_index_t_df32aa_idx"],
)
editor.remove_field(Book, new_field)
# Embedded indexes are removed.
self.assertEqual(
self.get_constraints_for_columns(
Book, ["author.address.index_together_one", "author.address.index_together_two"]
),
[],
)
self.assertEqual(
self.get_constraints_for_columns(
Book,
["author.index_together_three", "author.index_together_four"],
),
[],
)
editor.delete_model(Book)
self.assertTableNotExists(Book)

@isolate_apps("schema_")
def _test_add_field_db_index(self):
"""AddField + EmbeddedModelField"""

class Book(models.Model):
Expand All @@ -419,5 +538,15 @@ class Meta:
self.get_constraints_for_columns(Book, ["author.address.zip_code"]),
["schema__book_author.address.zip_code_7b9a9307"],
)
editor.remove_field(Book, new_field)
# Embedded indexes are removed.
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.age"]),
[],
)
self.assertEqual(
self.get_constraints_for_columns(Book, ["author.address.zip_code"]),
[],
)
editor.delete_model(Book)
self.assertTableNotExists(Author)

0 comments on commit 4d9c6c6

Please sign in to comment.