diff --git a/pony/migrate/diagram_ops.py b/pony/migrate/diagram_ops.py index 2b31b5936..e73338004 100644 --- a/pony/migrate/diagram_ops.py +++ b/pony/migrate/diagram_ops.py @@ -164,6 +164,19 @@ def apply(op, db): # op.reverse_name, _clone_attr(op.reverse_attr)) +class AddIndex(Operation): + def __init__(self, entity_name, attrs, name, is_unique, **kwargs): + Operation.__init__(self, **kwargs) + self.entity_name = entity_name + self.attrs = attrs + self.name = name + self.is_unique = is_unique + + def apply(self, db): + entity = db.entities[self.entity_name] + attrs = [entity._adict_[attr_name] for attr_name in self.attrs] + entity._add_index_(attrs, self.name, self.is_unique) + class Custom(Operation): is_custom = True diff --git a/pony/migrate/writer.py b/pony/migrate/writer.py index 1c45bf7e1..879d9b3c5 100644 --- a/pony/migrate/writer.py +++ b/pony/migrate/writer.py @@ -212,10 +212,9 @@ def ename(attr): bases = [c.__name__ for c in entities[ename].__bases__] regular = [(k, v) for k, v in attrs if not v.reverse] - # Patch by andgein. Add _table_, _discriminator_ and _table_options_ fields to entity attributes - # TODO: add composite_index, composite_key, PrimaryKey here? + # Patch by andgein. Add _table_, _discriminator_ and _table_options_ fields to entity attributes. special_fields = ["_table_", "_discriminator_", "_table_options_"] - # Some special fields can't been redefined in child entity + # Some special fields can't been redefined in child entity. special_fields_blocked_in_child = ["_table_", "_table_options_"] entity = self.db.entities[ename] @@ -227,8 +226,9 @@ def ename(attr): if field_value is not None: regular.append((field, field_value)) - self_indexes = [index for index in entity._indexes_ if index.entity == entity] - regular.append(("_indexes_", self_indexes)) + own_indexes = [index for index in entity._indexes_ if index.entity == entity] + composite_indexes = [index for index in own_indexes if len(index.attrs) > 1 and not index.is_pk] + regular.append(("_indexes_", list(set(own_indexes) - set(composite_indexes)))) # End of patch by andgein result.append( @@ -257,6 +257,14 @@ def ename(attr): result.append(ops.AddRelation( ename, name(attr), attr, name(attr.reverse), attr.reverse )) + + # Patch by andgein. Composite indices should be added after relations + for index in composite_indexes: + result.append( + ops.AddIndex(ename, [attr.name for attr in index.attrs], index.name, index.is_unique) + ) + # End of patch + for ename, attrs in sorted(eremoved.items()): result.append( ops.RemoveEntity(ename) diff --git a/pony/orm/core.py b/pony/orm/core.py index 5fd814f98..4ad66c0c9 100644 --- a/pony/orm/core.py +++ b/pony/orm/core.py @@ -4230,6 +4230,10 @@ def _remove_relation_(entity, attr_name): if schema is None: return m2m_table = schema.tables[m2m_table_name] m2m_table.remove() + + def _add_index_(entity, attrs, name, is_unique): + entity._indexes_.append(Index(*attrs, name=name, is_unique=is_unique)) + def _erase_bits_(entity): entity._bits_ = entity._bits_except_volatile_ = entity._all_bits_ = entity._all_bits_except_volatile_ = None entity._offset_counter_ = 0 diff --git a/pony/tests/test_migrations/test_cli.py b/pony/tests/test_migrations/test_cli.py index 44a8b2efd..72b39005f 100644 --- a/pony/tests/test_migrations/test_cli.py +++ b/pony/tests/test_migrations/test_cli.py @@ -1093,7 +1093,7 @@ class Child(Parent): _discriminator_ = 2 another_field = orm.Required(str) - orm.composite_index("type_field", another_field) + orm.composite_key("type_field", another_field) command.migrate(db, "make -v") command.migrate(db, "apply -v")