Add _add_fk_constraint and _drop_fk_constraint to render

3 files changed, 23 insertions, 5 deletions

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 72c22a1..5a5e75a 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -9,7 +9,7 @@ from sqlalchemy.util import OrderedSet
 from .compare import _compare_tables
 from .render import _drop_table, _drop_column, _drop_index, _drop_constraint, \
     _add_table, _add_column, _add_index, _add_constraint, _modify_col, \
-    _add_fk_constraint
+    _add_fk_constraint, _drop_fk_constraint
 from .. import util
 
 log = logging.getLogger(__name__)
@@ -261,7 +261,7 @@ def _invoke_adddrop_command(updown, args, autogen_context):
         "column": (_drop_column, _add_column),
         "index": (_drop_index, _add_index),
         "constraint": (_drop_constraint, _add_constraint),
-        "fk":(_drop_constraint, _add_fk_constraint)
+        "fk":(_drop_fk_constraint, _add_fk_constraint)
     }
 
     cmd_callables = _commands[cmd_type]
diff --git a/alembic/autogenerate/compare.py b/alembic/autogenerate/compare.py
index f943926..5e7a303 100644
--- a/alembic/autogenerate/compare.py
+++ b/alembic/autogenerate/compare.py
@@ -569,7 +569,7 @@ FKInfo = collections.namedtuple('fk_info', ['constrained_columns',
 
 
 def _compare_foreign_keys(schema, tname, object_filters, conn_table,
-            metadata_table, diffs, autogen_context, inspector):
+                          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
@@ -596,7 +596,8 @@ 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'],
+    return FKInfo(tuple(fk['constrained_columns']),
+                  fk['referred_table'],
                   tuple(fk['referred_columns']))
 
 
diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 89dd2f2..a0bc534 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -194,7 +194,24 @@ def _uq_constraint(constraint, autogen_context, alter):
 
 
 def _add_fk_constraint(constraint, fk_info, autogen_context):
-    raise NotImplementedError()
+    args = [repr(_render_gen_name(autogen_context, constraint.name)),
+            constraint.parent.table.name]
+    args.extend([vars(fk_info).values()[1],
+                 str(list(vars(fk_info).values()[0])),
+                 str(list(vars(fk_info).values()[2]))])
+    return "%(prefix)screate_foreign_key(%(args)s)" % {
+        'prefix': _alembic_autogenerate_prefix(autogen_context),
+        'args': ", ".join(args)
+    }
+
+
+def _drop_fk_constraint(constraint, fk_info, autogen_context):
+    args = [repr(_render_gen_name(autogen_context, constraint.name)),
+            constraint.parent.table.name, 'type_=foreignkey']
+    return "%(prefix)sdrop_constraint(%(args)s)" % {
+        'prefix': _alembic_autogenerate_prefix(autogen_context),
+        'args': ", ".join(args)
+    }
 
 
 def _add_pk_constraint(constraint, autogen_context):