Add composite fk support, add more tests

2 files changed, 159 insertions, 6 deletions

diff --git a/alembic/autogenerate/compare.py b/alembic/autogenerate/compare.py
index c6f70e0..f943926 100644
--- a/alembic/autogenerate/compare.py
+++ b/alembic/autogenerate/compare.py
@@ -596,10 +596,18 @@ def _compare_foreign_keys(schema, tname, object_filters, conn_table,
 
 
 def _get_fk_info_from_db(fk):
-    return FKInfo(tuple(fk['constrained_columns']),
-                       fk['referred_table'],
-                       tuple(fk['referred_columns']))
+    return FKInfo(tuple(fk['constrained_columns']), fk['referred_table'],
+                  tuple(fk['referred_columns']))
+
 
 def _get_fk_info_from_model(fk):
-    return FKInfo((fk.parent.name,), fk.column.table.name,
-                       (fk.column.name,))
\ No newline at end of file
+    constrained_columns = []
+    for column in fk.constraint.columns:
+        if not isinstance(column, basestring):
+            constrained_columns.append(column.name)
+        else:
+            constrained_columns.append(column)
+    return FKInfo(
+        tuple(constrained_columns),
+        fk.column.table.name,
+        tuple(k.column.name for k in fk.constraint._elements.values()))
\ No newline at end of file
diff --git a/tests/test_autogen_fks.py b/tests/test_autogen_fks.py
index ed15eea..d65b1aa 100644
--- a/tests/test_autogen_fks.py
+++ b/tests/test_autogen_fks.py
@@ -78,4 +78,149 @@ class AutogenerateForeignKeysTest(AutogenFixtureTest, TestBase):
         eq_(diffs[0][1].parent.table.name, "user")
         eq_(diffs[0][2].constrained_columns, ('test2',))
         eq_(diffs[0][2].referred_table, 'table')
-        eq_(diffs[0][2].referred_columns, ('test',))
\ No newline at end of file
+        eq_(diffs[0][2].referred_columns, ('test',))
+
+    def test_no_change(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        Table('table', m1,
+              Column('id', Integer, primary_key=True),
+              Column('test', String(10)))
+
+        Table('user', m1,
+              Column('id', Integer, primary_key=True),
+              Column('name', String(50), nullable=False),
+              Column('a1', String(10), server_default="x"),
+              Column('test2', String(10)),
+              ForeignKeyConstraint(['test2'], ['table.test']))
+
+        Table('table', m2,
+              Column('id', Integer, primary_key=True),
+              Column('test', String(10)))
+
+        Table('user', m2,
+              Column('id', Integer, primary_key=True),
+              Column('name', String(50), nullable=False),
+              Column('a1', String(10), server_default="x"),
+              Column('test2', String(10)),
+              ForeignKeyConstraint(['test2'], ['table.test']))
+
+        diffs = self._fixture(m1, m2)
+
+        eq_(diffs, [])
+
+    def test_no_change_composite_fk(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        Table('table', m1,
+              Column('id', Integer, primary_key=True),
+              Column('id_1', String(10)),
+              Column('id_2', String(10)))
+
+        Table('user', m1,
+              Column('id', Integer, primary_key=True),
+              Column('name', String(50), nullable=False),
+              Column('a1', String(10), server_default="x"),
+              Column('other_id_1', String(10)),
+              Column('other_id_2', String(10)),
+              ForeignKeyConstraint(['other_id_1', 'other_id_2'],
+                                   ['table.id_1','table.id_2']))
+
+        Table('table', m2,
+              Column('id', Integer, primary_key=True),
+              Column('id_1', String(10)),
+              Column('id_2', String(10)))
+
+        Table('user', m2,
+              Column('id', Integer, primary_key=True),
+              Column('name', String(50), nullable=False),
+              Column('a1', String(10), server_default="x"),
+              Column('other_id_1', String(10)),
+              Column('other_id_2', String(10)),
+              ForeignKeyConstraint(['other_id_1', 'other_id_2'],
+                                   ['table.id_1','table.id_2']))
+
+        diffs = self._fixture(m1, m2)
+
+        eq_(diffs, [])
+
+    def test_missing_composite_fk_with_name(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        Table('table', m1,
+              Column('id', Integer, primary_key=True),
+              Column('id_1', String(10)),
+              Column('id_2', String(10)))
+
+        Table('user', m1,
+              Column('id', Integer, primary_key=True),
+              Column('name', String(50), nullable=False),
+              Column('a1', String(10), server_default="x"),
+              Column('other_id_1', String(10)),
+              Column('other_id_2', String(10)))
+
+        Table('table', m2,
+              Column('id', Integer, primary_key=True),
+              Column('id_1', String(10)),
+              Column('id_2', String(10)))
+
+        Table('user', m2,
+              Column('id', Integer, primary_key=True),
+              Column('name', String(50), nullable=False),
+              Column('a1', String(10), server_default="x"),
+              Column('other_id_1', String(10)),
+              Column('other_id_2', String(10)),
+              ForeignKeyConstraint(['other_id_1', 'other_id_2'],
+                                   ['table.id_1','table.id_2'],
+                                   name='fk_test_name'))
+
+        diffs = self._fixture(m1, m2)
+
+        eq_(diffs[0][0], "add_fk")
+        eq_(diffs[0][1].parent.table.name, "user")
+        eq_(diffs[0][1].name, "fk_test_name")
+        eq_(diffs[0][2].constrained_columns, ('other_id_1', 'other_id_2'))
+        eq_(diffs[0][2].referred_table, 'table')
+        eq_(diffs[0][2].referred_columns, ('id_1', 'id_2'))
+
+    def test_extra_composite_fk(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        Table('table', m1,
+              Column('id', Integer, primary_key=True),
+              Column('id_1', String(10)),
+              Column('id_2', String(10)))
+
+        Table('user', m1,
+              Column('id', Integer, primary_key=True),
+              Column('name', String(50), nullable=False),
+              Column('a1', String(10), server_default="x"),
+              Column('other_id_1', String(10)),
+              Column('other_id_2', String(10)),
+              ForeignKeyConstraint(['other_id_1', 'other_id_2'],
+                                   ['table.id_1', 'table.id_2'],
+                                   name='fk_test_name'))
+
+        Table('table', m2,
+              Column('id', Integer, primary_key=True),
+              Column('id_1', String(10)),
+              Column('id_2', String(10)))
+
+        Table('user', m2,
+              Column('id', Integer, primary_key=True),
+              Column('name', String(50), nullable=False),
+              Column('a1', String(10), server_default="x"),
+              Column('other_id_1', String(10)),
+              Column('other_id_2', String(10)))
+
+        diffs = self._fixture(m1, m2)
+
+        eq_(diffs[0][0], "drop_fk")
+        eq_(diffs[0][2].name, "user")
+        eq_(diffs[0][3].constrained_columns, ('other_id_1', 'other_id_2'))
+        eq_(diffs[0][3].referred_table, 'table')
+        eq_(diffs[0][3].referred_columns, ('id_1', 'id_2'))
\ No newline at end of file