Add indices after relation creates

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
 

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)

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

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")