- Fixed bug in batch mode where the ``batch_op.create_foreign_key()``

directive would be incorrectly rendered with the source table and
schema names in the argument list.
fixes #315

4 files changed, 210 insertions, 9 deletions

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 6f5f96c..422506d 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -212,16 +212,29 @@ def _add_fk_constraint(autogen_context, op):
     args = [
         repr(
             _render_gen_name(autogen_context, op.constraint_name)),
-        repr(_ident(op.source_table)),
-        repr(_ident(op.referent_table)),
-        repr([_ident(col) for col in op.local_cols]),
-        repr([_ident(col) for col in op.remote_cols])
     ]
+    if not autogen_context._has_batch:
+        args.append(
+            repr(_ident(op.source_table))
+        )
+
+    args.extend(
+        [
+            repr(_ident(op.referent_table)),
+            repr([_ident(col) for col in op.local_cols]),
+            repr([_ident(col) for col in op.remote_cols])
+        ]
+    )
+
+    kwargs = [
+        'referent_schema',
+        'onupdate', 'ondelete', 'initially',
+        'deferrable', 'use_alter'
+    ]
+    if not autogen_context._has_batch:
+        kwargs.insert(0, 'source_schema')
 
-    for k in (
-        'source_schema', 'referent_schema',
-        'onupdate', 'ondelete', 'initially', 'deferrable', 'use_alter'
-    ):
+    for k in kwargs:
         if k in op.kw:
             value = op.kw[k]
             if value is not None:
diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst
index e1a5a76..7b657a8 100644
--- a/docs/build/changelog.rst
+++ b/docs/build/changelog.rst
@@ -7,6 +7,14 @@ Changelog
     :version: 0.8.0
 
     .. change::
+      :tags: bug, batch
+      :tickets: 315
+
+      Fixed bug in batch mode where the ``batch_op.create_foreign_key()``
+      directive would be incorrectly rendered with the source table and
+      schema names in the argument list.
+
+    .. change::
       :tags: feature, commands
       :pullreq: bitbucket:46
 
diff --git a/tests/test_autogen_composition.py b/tests/test_autogen_composition.py
index 6d1f55b..96d4856 100644
--- a/tests/test_autogen_composition.py
+++ b/tests/test_autogen_composition.py
@@ -156,7 +156,7 @@ nullable=True))
                type_=sa.Numeric(precision=10, scale=2),
                nullable=True,
                existing_server_default=sa.text('0'))
-        batch_op.create_foreign_key(None, 'order', 'user', ['user_id'], ['id'])
+        batch_op.create_foreign_key(None, 'user', ['user_id'], ['id'])
 
     with op.batch_alter_table('user', schema=None) as batch_op:
         batch_op.alter_column('a1',
diff --git a/tests/test_autogen_render.py b/tests/test_autogen_render.py
index 5d1aca2..989fe4b 100644
--- a/tests/test_autogen_render.py
+++ b/tests/test_autogen_render.py
@@ -77,6 +77,25 @@ class AutogenRenderTest(TestBase):
             "['active', 'code'], unique=False)"
         )
 
+    def test_render_add_index_batch(self):
+        """
+        autogenerate.render._add_index
+        """
+        m = MetaData()
+        t = Table('test', m,
+                  Column('id', Integer, primary_key=True),
+                  Column('active', Boolean()),
+                  Column('code', String(255)),
+                  )
+        idx = Index('test_active_code_idx', t.c.active, t.c.code)
+        op_obj = ops.CreateIndexOp.from_index(idx)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.create_index('test_active_code_idx', "
+                "['active', 'code'], unique=False)"
+            )
+
     def test_render_add_index_schema(self):
         """
         autogenerate.render._add_index using schema
@@ -96,6 +115,26 @@ class AutogenRenderTest(TestBase):
             "['active', 'code'], unique=False, schema='CamelSchema')"
         )
 
+    def test_render_add_index_schema_batch(self):
+        """
+        autogenerate.render._add_index using schema
+        """
+        m = MetaData()
+        t = Table('test', m,
+                  Column('id', Integer, primary_key=True),
+                  Column('active', Boolean()),
+                  Column('code', String(255)),
+                  schema='CamelSchema'
+                  )
+        idx = Index('test_active_code_idx', t.c.active, t.c.code)
+        op_obj = ops.CreateIndexOp.from_index(idx)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.create_index('test_active_code_idx', "
+                "['active', 'code'], unique=False)"
+            )
+
     def test_render_add_index_pg_where(self):
         autogen_context = self.pg_autogen_context
 
@@ -191,6 +230,24 @@ unique=False, """
             "op.drop_index('test_active_code_idx', table_name='test')"
         )
 
+    def test_drop_index_batch(self):
+        """
+        autogenerate.render._drop_index
+        """
+        m = MetaData()
+        t = Table('test', m,
+                  Column('id', Integer, primary_key=True),
+                  Column('active', Boolean()),
+                  Column('code', String(255)),
+                  )
+        idx = Index('test_active_code_idx', t.c.active, t.c.code)
+        op_obj = ops.DropIndexOp.from_index(idx)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.drop_index('test_active_code_idx')"
+            )
+
     def test_drop_index_schema(self):
         """
         autogenerate.render._drop_index using schema
@@ -210,6 +267,25 @@ unique=False, """
             "table_name='test', schema='CamelSchema')"
         )
 
+    def test_drop_index_schema_batch(self):
+        """
+        autogenerate.render._drop_index using schema
+        """
+        m = MetaData()
+        t = Table('test', m,
+                  Column('id', Integer, primary_key=True),
+                  Column('active', Boolean()),
+                  Column('code', String(255)),
+                  schema='CamelSchema'
+                  )
+        idx = Index('test_active_code_idx', t.c.active, t.c.code)
+        op_obj = ops.DropIndexOp.from_index(idx)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.drop_index('test_active_code_idx')"
+            )
+
     def test_add_unique_constraint(self):
         """
         autogenerate.render._add_unique_constraint
@@ -227,6 +303,24 @@ unique=False, """
             "op.create_unique_constraint('uq_test_code', 'test', ['code'])"
         )
 
+    def test_add_unique_constraint_batch(self):
+        """
+        autogenerate.render._add_unique_constraint
+        """
+        m = MetaData()
+        t = Table('test', m,
+                  Column('id', Integer, primary_key=True),
+                  Column('active', Boolean()),
+                  Column('code', String(255)),
+                  )
+        uq = UniqueConstraint(t.c.code, name='uq_test_code')
+        op_obj = ops.AddConstraintOp.from_constraint(uq)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.create_unique_constraint('uq_test_code', ['code'])"
+            )
+
     def test_add_unique_constraint_schema(self):
         """
         autogenerate.render._add_unique_constraint using schema
@@ -246,6 +340,26 @@ unique=False, """
             "['code'], schema='CamelSchema')"
         )
 
+    def test_add_unique_constraint_schema_batch(self):
+        """
+        autogenerate.render._add_unique_constraint using schema
+        """
+        m = MetaData()
+        t = Table('test', m,
+                  Column('id', Integer, primary_key=True),
+                  Column('active', Boolean()),
+                  Column('code', String(255)),
+                  schema='CamelSchema'
+                  )
+        uq = UniqueConstraint(t.c.code, name='uq_test_code')
+        op_obj = ops.AddConstraintOp.from_constraint(uq)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.create_unique_constraint('uq_test_code', "
+                "['code'])"
+            )
+
     def test_drop_unique_constraint(self):
         """
         autogenerate.render._drop_constraint
@@ -294,6 +408,19 @@ unique=False, """
             "op.create_foreign_key('fk_a_id', 'b', 'a', ['a_id'], ['id'])"
         )
 
+    def test_add_fk_constraint_batch(self):
+        m = MetaData()
+        Table('a', m, Column('id', Integer, primary_key=True))
+        b = Table('b', m, Column('a_id', Integer, ForeignKey('a.id')))
+        fk = ForeignKeyConstraint(['a_id'], ['a.id'], name='fk_a_id')
+        b.append_constraint(fk)
+        op_obj = ops.AddConstraintOp.from_constraint(fk)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.create_foreign_key('fk_a_id', 'a', ['a_id'], ['id'])"
+            )
+
     def test_add_fk_constraint_kwarg(self):
         m = MetaData()
         t1 = Table('t', m, Column('c', Integer))
@@ -453,6 +580,26 @@ unique=False, """
             "referent_schema='CamelSchemaTwo')"
         )
 
+    def test_add_fk_constraint_schema_batch(self):
+        m = MetaData()
+        Table(
+            'a', m, Column('id', Integer, primary_key=True),
+            schema="CamelSchemaTwo")
+        b = Table(
+            'b', m, Column('a_id', Integer, ForeignKey('a.id')),
+            schema="CamelSchemaOne")
+        fk = ForeignKeyConstraint(
+            ["a_id"],
+            ["CamelSchemaTwo.a.id"], name='fk_a_id')
+        b.append_constraint(fk)
+        op_obj = ops.AddConstraintOp.from_constraint(fk)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.create_foreign_key('fk_a_id', 'a', ['a_id'], ['id'],"
+                " referent_schema='CamelSchemaTwo')"
+            )
+
     def test_drop_fk_constraint(self):
         m = MetaData()
         Table('a', m, Column('id', Integer, primary_key=True))
@@ -465,6 +612,19 @@ unique=False, """
             "op.drop_constraint('fk_a_id', 'b', type_='foreignkey')"
         )
 
+    def test_drop_fk_constraint_batch(self):
+        m = MetaData()
+        Table('a', m, Column('id', Integer, primary_key=True))
+        b = Table('b', m, Column('a_id', Integer, ForeignKey('a.id')))
+        fk = ForeignKeyConstraint(['a_id'], ['a.id'], name='fk_a_id')
+        b.append_constraint(fk)
+        op_obj = ops.DropConstraintOp.from_constraint(fk)
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.drop_constraint('fk_a_id', type_='foreignkey')"
+            )
+
     def test_drop_fk_constraint_schema(self):
         m = MetaData()
         Table(
@@ -485,6 +645,26 @@ unique=False, """
             "type_='foreignkey')"
         )
 
+    def test_drop_fk_constraint_batch_schema(self):
+        m = MetaData()
+        Table(
+            'a', m, Column('id', Integer, primary_key=True),
+            schema="CamelSchemaTwo")
+        b = Table(
+            'b', m, Column('a_id', Integer, ForeignKey('a.id')),
+            schema="CamelSchemaOne")
+        fk = ForeignKeyConstraint(
+            ["a_id"],
+            ["CamelSchemaTwo.a.id"], name='fk_a_id')
+        b.append_constraint(fk)
+        op_obj = ops.DropConstraintOp.from_constraint(fk)
+
+        with self.autogen_context._within_batch():
+            eq_ignore_whitespace(
+                autogenerate.render_op_text(self.autogen_context, op_obj),
+                "batch_op.drop_constraint('fk_a_id', type_='foreignkey')"
+            )
+
     def test_render_table_upgrade(self):
         m = MetaData()
         t = Table('test', m,