- more foreign key stuffrel_0_4_0

2 files changed, 74 insertions, 7 deletions

diff --git a/alembic/autogenerate.py b/alembic/autogenerate.py
index b6af35c..7383bae 100644
--- a/alembic/autogenerate.py
+++ b/alembic/autogenerate.py
@@ -603,16 +603,35 @@ def _render_primary_key(constraint, autogen_context):
         ),
     }
 
+def _fk_colspec(fk, metadata_schema):
+    """Implement a 'safe' version of ForeignKey._get_colspec() that
+    never tries to resolve the remote table.
+
+    """
+    if metadata_schema is None:
+        return fk._get_colspec()
+    else:
+        # need to render schema breaking up tokens by hand, since the
+        # ForeignKeyConstraint here may not actually have a remote
+        # Table present
+        tokens = fk._colspec.split(".")
+        # no schema in the colspec, render it
+        if len(tokens) == 2:
+            return "%s.%s" % (metadata_schema, fk._colspec)
+        else:
+            return fk._colspec
+
 def _render_foreign_key(constraint, autogen_context):
     opts = []
     if constraint.name:
         opts.append(("name", repr(constraint.name)))
+    apply_metadata_schema = constraint.parent.metadata.schema
     # TODO: deferrable, initially, etc.
     return "%(prefix)sForeignKeyConstraint([%(cols)s], "\
             "[%(refcols)s], %(args)s)" % {
         "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
         "cols": ", ".join("'%s'" % f.parent.key for f in constraint.elements),
-        "refcols": ", ".join(repr(f._get_colspec())
+        "refcols": ", ".join(repr(_fk_colspec(f, apply_metadata_schema))
                             for f in constraint.elements),
         "args": ", ".join(
             ["%s=%s" % (kwname, val) for kwname, val in opts]
diff --git a/tests/test_autogenerate.py b/tests/test_autogenerate.py
index 18807e0..9cc94b9 100644
--- a/tests/test_autogenerate.py
+++ b/tests/test_autogenerate.py
@@ -36,14 +36,9 @@ def _model_one(schema=None):
                 server_default="0"),
     )
 
-    if schema is not None:
-        schema_arg = schema + "."
-    else:
-        schema_arg = ""
-
     Table('extra', m,
         Column("x", CHAR),
-        Column('uid', Integer, ForeignKey('%suser.id' % schema_arg))
+        Column('uid', Integer, ForeignKey('user.id'))
     )
 
     return m
@@ -764,12 +759,65 @@ class AutogenRenderTest(TestCase):
         m = MetaData()
         t = Table('test', m,
             Column('id', Integer, primary_key=True),
+            Column('q', Integer, ForeignKey('address.id')),
             schema='foo'
         )
         eq_ignore_whitespace(
             autogenerate._add_table(t, self.autogen_context),
             "op.create_table('test',"
             "sa.Column('id', sa.Integer(), nullable=False),"
+            "sa.Column('q', sa.Integer(), nullable=True),"
+            "sa.ForeignKeyConstraint(['q'], ['address.id'], ),"
+            "sa.PrimaryKeyConstraint('id'),"
+            "schema='foo'"
+            ")"
+        )
+
+    def test_render_table_w_fk_schema(self):
+        m = MetaData()
+        t = Table('test', m,
+            Column('id', Integer, primary_key=True),
+            Column('q', Integer, ForeignKey('foo.address.id')),
+        )
+        eq_ignore_whitespace(
+            autogenerate._add_table(t, self.autogen_context),
+            "op.create_table('test',"
+            "sa.Column('id', sa.Integer(), nullable=False),"
+            "sa.Column('q', sa.Integer(), nullable=True),"
+            "sa.ForeignKeyConstraint(['q'], ['foo.address.id'], ),"
+            "sa.PrimaryKeyConstraint('id')"
+            ")"
+        )
+
+    def test_render_table_w_metadata_schema(self):
+        m = MetaData(schema="foo")
+        t = Table('test', m,
+            Column('id', Integer, primary_key=True),
+            Column('q', Integer, ForeignKey('address.id')),
+        )
+        eq_ignore_whitespace(
+            autogenerate._add_table(t, self.autogen_context),
+            "op.create_table('test',"
+            "sa.Column('id', sa.Integer(), nullable=False),"
+            "sa.Column('q', sa.Integer(), nullable=True),"
+            "sa.ForeignKeyConstraint(['q'], ['foo.address.id'], ),"
+            "sa.PrimaryKeyConstraint('id'),"
+            "schema='foo'"
+            ")"
+        )
+
+    def test_render_table_w_metadata_schema_override(self):
+        m = MetaData(schema="foo")
+        t = Table('test', m,
+            Column('id', Integer, primary_key=True),
+            Column('q', Integer, ForeignKey('bar.address.id')),
+        )
+        eq_ignore_whitespace(
+            autogenerate._add_table(t, self.autogen_context),
+            "op.create_table('test',"
+            "sa.Column('id', sa.Integer(), nullable=False),"
+            "sa.Column('q', sa.Integer(), nullable=True),"
+            "sa.ForeignKeyConstraint(['q'], ['bar.address.id'], ),"
             "sa.PrimaryKeyConstraint('id'),"
             "schema='foo'"
             ")"