Skip to content

Commit

Permalink
add helpers to add/delete unique
Browse files Browse the repository at this point in the history
  • Loading branch information
timgraham committed Sep 26, 2024
1 parent 096096a commit ad9e91e
Showing 1 changed file with 30 additions and 40 deletions.
70 changes: 30 additions & 40 deletions django_mongodb/schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,7 @@ def _create_model_indexes(self, model):
# Field uniques
for field in model._meta.local_fields:
if field.unique and field.column != "_id":
constraint = UniqueConstraint(
fields=[field.name], name=f"{model._meta.db_table}_{field.column}_key"
)
self.add_constraint(model, constraint)
self._add_unique_for_field(model, field)
# Meta.index_together (RemovedInDjango51Warning)
for field_names in model._meta.index_together:
index = Index(fields=field_names)
Expand Down Expand Up @@ -66,10 +63,7 @@ def add_field(self, model, field):
if self._field_should_be_indexed(model, field):
self._add_index_for_field(model, field)
if field.unique and field.column != "_id":
constraint = UniqueConstraint(
fields=[field.name], name=f"{model._meta.db_table}_{field.column}_key"
)
self.add_constraint(model, constraint, field=field)
self._add_unique_for_field(model, field)

def _add_index_for_field(self, model, field):
new_index = Index(fields=[field.name])
Expand Down Expand Up @@ -97,6 +91,30 @@ def _drop_index_for_field(self, model, field):
# is to look at its name (refs #28053).
collection.drop_index(index_name)

def _add_unique_for_field(self, model, field):
constraint = UniqueConstraint(
fields=[field.name], name=f"{model._meta.db_table}_{field.column}_key"
)
self.add_constraint(model, constraint, field=field)

def _drop_unique_for_field(self, model, field, strict):
# 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 strict and len(constraint_names) != 1:
raise ValueError(
f"Found wrong number ({len(constraint_names)}) of unique "
f"constraints for {model._meta.db_table}.{field.column}"
)
for constraint_name in constraint_names:
self.connection.database[model._meta.db_table].drop_index(constraint_name)

def _alter_field(
self,
model,
Expand All @@ -113,26 +131,7 @@ def _alter_field(
if old_field.unique and (
not new_field.unique or self._field_became_primary_key(old_field, new_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,
[old_field.column],
unique=True,
primary_key=False,
exclude=meta_constraint_names,
)
if strict and len(constraint_names) != 1:
raise ValueError(
f"Found wrong number ({len(constraint_names)}) of unique "
f"constraints for {model._meta.db_table}.{old_field.column}"
)
for constraint_name in constraint_names:
constraint = UniqueConstraint(
fields=[old_field.name],
name=constraint_name,
)
self.remove_constraint(model, constraint)
self._drop_unique_for_field(model, old_field, strict)
# Removed an index? (no strict check, as multiple indexes are possible)
# Remove indexes if db_index switched to False or a unique constraint
# will now be used in lieu of an index. The following lines from the
Expand Down Expand Up @@ -161,14 +160,8 @@ def _alter_field(
# Move unique to the new field, if needed.
if old_field.unique and new_field.unique and new_field.column != "_id":
if not old_field.primary_key:
old_constraint = UniqueConstraint(
fields=[old_field.name], name=f"{model._meta.db_table}_{old_field.column}_key"
)
self.remove_constraint(model, old_constraint)
new_constraint = UniqueConstraint(
fields=[new_field.name], name=f"{model._meta.db_table}_{new_field.column}_key"
)
self.add_constraint(model, new_constraint, field=new_field)
self._drop_unique_for_field(model, old_field, strict)
self._add_unique_for_field(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 @@ -177,10 +170,7 @@ def _alter_field(
collection.update_many({column: {"$eq": None}}, [{"$set": {column: default}}])
# Added a unique?
if self._unique_should_be_added(old_field, new_field):
constraint = UniqueConstraint(
fields=[new_field.name], name=f"{model._meta.db_table}_{new_field.column}_key"
)
self.add_constraint(model, constraint, field=new_field)
self._add_unique_for_field(model, new_field)
# Added an index? Add an index if db_index switched to True or a unique
# constraint will no longer be used in lieu of an index. The following
# lines from the truth table show all True cases; the rest are False:
Expand Down

0 comments on commit ad9e91e

Please sign in to comment.