Fix failing unittests and update compare_fk method

2 files changed, 36 insertions, 24 deletions

diff --git a/alembic/autogenerate/compare.py b/alembic/autogenerate/compare.py
index 5e7a303..db2e9c5 100644
--- a/alembic/autogenerate/compare.py
+++ b/alembic/autogenerate/compare.py
@@ -570,29 +570,40 @@ FKInfo = collections.namedtuple('fk_info', ['constrained_columns',
 
 def _compare_foreign_keys(schema, tname, object_filters, conn_table,
                           metadata_table, diffs, autogen_context, inspector):
-        # This methods checks foreign keys that tables contain in models with
-        # foreign keys that are in db.
-            # Get all necessary information about key of current table from db
-        fk_db = dict((_get_fk_info_from_db(i), i['name']) for i in
-                     inspector.get_foreign_keys(tname))
-        fk_db_set = set(fk_db.keys())
-        # Get all necessary information about key of current table from
-        # models
-        fk_models = dict((_get_fk_info_from_model(fk), fk) for fk in
-                         metadata_table.foreign_keys)
-        fk_models_set = set(fk_models.keys())
-        for key in (fk_db_set - fk_models_set):
-                diffs.append(('drop_fk', fk_db[key], conn_table, key))
-                log.info(("Detected removed foreign key %(fk)r on "
-                          "table %(table)r"), {'fk': fk_db[key],
-                                               'table': conn_table})
-        for key in (fk_models_set - fk_db_set):
-                diffs.append(('add_fk', fk_models[key], key))
-                log.info((
-                    "Detected added foreign key for column %(fk)r on table "
-                    "%(table)r"), {'fk': fk_models[key].column.name,
-                                   'table': conn_table})
-        return diffs
+
+    # This methods checks foreign keys that tables contain in models with
+    # foreign keys that are in db.
+    # Get all necessary information about key of current table from db
+    if conn_table is None:
+        return
+
+    fk_db = {}
+    if hasattr(inspector, "get_foreign_keys"):
+        try:
+            fk_db = dict((_get_fk_info_from_db(i), i['name']) for i in
+                         inspector.get_foreign_keys(tname, schema=schema))
+        except NotImplementedError:
+            pass
+
+
+    # Get all necessary information about key of current table from
+    # models
+    fk_models = dict((_get_fk_info_from_model(fk), fk) for fk in
+                     metadata_table.foreign_keys)
+    fk_models_set = set(fk_models.keys())
+    fk_db_set = set(fk_db.keys())
+    for key in (fk_db_set - fk_models_set):
+            diffs.append(('drop_fk', fk_db[key], conn_table, key))
+            log.info(("Detected removed foreign key %(fk)r on "
+                      "table %(table)r"), {'fk': fk_db[key],
+                                           'table': conn_table})
+    for key in (fk_models_set - fk_db_set):
+            diffs.append(('add_fk', fk_models[key], key))
+            log.info((
+                "Detected added foreign key for column %(fk)r on table "
+                "%(table)r"), {'fk': fk_models[key].column.name,
+                               'table': conn_table})
+    return diffs
 
 
 def _get_fk_info_from_db(fk):
diff --git a/tests/test_autogenerate.py b/tests/test_autogenerate.py
index 780d44c..19a0d1d 100644
--- a/tests/test_autogenerate.py
+++ b/tests/test_autogenerate.py
@@ -462,7 +462,6 @@ class AutogenerateDiffTest(ModelOne, AutogenTest, TestBase):
         template_args = {}
         autogenerate._produce_migration_diffs(
             self.context, template_args, set())
-
         eq_(re.sub(r"u'", "'", template_args['upgrades']),
             """### commands auto generated by Alembic - please adjust! ###
     op.create_table('item',
@@ -482,6 +481,7 @@ nullable=True))
                type_=sa.Numeric(precision=10, scale=2),
                nullable=True,
                existing_server_default=sa.text('0'))
+    op.create_foreign_key(None, order, user, ['user_id'], ['id'], schema=None)
     op.drop_column('user', 'pw')
     op.alter_column('user', 'a1',
                existing_type=sa.TEXT(),
@@ -503,6 +503,7 @@ nullable=True))
                existing_nullable=True)
     op.add_column('user', sa.Column('pw', sa.VARCHAR(length=50), \
 nullable=True))
+    op.drop_constraint(None, order, type_=foreignkey)
     op.alter_column('order', 'amount',
                existing_type=sa.Numeric(precision=10, scale=2),
                type_=sa.NUMERIC(precision=8, scale=2),