include indexes in batch

2 files changed, 50 insertions, 9 deletions

diff --git a/tests/test_autogenerate.py b/tests/test_autogenerate.py
index 60b3a3e..51cae27 100644
--- a/tests/test_autogenerate.py
+++ b/tests/test_autogenerate.py
@@ -594,7 +594,7 @@ nullable=True))
         batch_op.alter_column('name',
                existing_type=sa.VARCHAR(length=50),
                nullable=False)
-        batch_op.drop_index('pw_idx', table_name='user')
+        batch_op.drop_index('pw_idx')
         batch_op.drop_column('pw')
 
     ### end Alembic commands ###""")
@@ -603,7 +603,7 @@ nullable=True))
             """### commands auto generated by Alembic - please adjust! ###
     with op.batch_alter_table('user', schema=None) as batch_op:
         batch_op.add_column(sa.Column('pw', sa.VARCHAR(length=50), nullable=True))
-        batch_op.create_index('pw_idx', 'user', ['pw'], unique=False)
+        batch_op.create_index('pw_idx', ['pw'], unique=False)
         batch_op.alter_column('name',
                existing_type=sa.VARCHAR(length=50),
                nullable=True)
diff --git a/tests/test_batch.py b/tests/test_batch.py
index a387792..3892951 100644
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -10,9 +10,9 @@ from alembic.batch import ApplyBatchImpl
 from alembic.migration import MigrationContext
 
 from sqlalchemy import Integer, Table, Column, String, MetaData, ForeignKey, \
-    UniqueConstraint, ForeignKeyConstraint
+    UniqueConstraint, ForeignKeyConstraint, Index
 from sqlalchemy.sql import column
-from sqlalchemy.schema import CreateTable
+from sqlalchemy.schema import CreateTable, CreateIndex
 
 
 class BatchApplyTest(TestBase):
@@ -42,6 +42,17 @@ class BatchApplyTest(TestBase):
         )
         return ApplyBatchImpl(t, table_args, table_kwargs)
 
+    def _ix_fixture(self, table_args=(), table_kwargs={}):
+        m = MetaData()
+        t = Table(
+            'tname', m,
+            Column('id', Integer, primary_key=True),
+            Column('x', String()),
+            Column('y', Integer),
+            Index('ix1', 'y')
+        )
+        return ApplyBatchImpl(t, table_args, table_kwargs)
+
     def _fk_fixture(self, table_args=(), table_kwargs={}):
         m = MetaData()
         t = Table(
@@ -90,13 +101,23 @@ class BatchApplyTest(TestBase):
             CreateTable(impl.new_table).compile(dialect=context.dialect))
         create_stmt = re.sub(r'[\n\t]', '', create_stmt)
 
+        idx_stmt = ""
+        for idx in impl.new_table.indexes:
+            idx_stmt += str(CreateIndex(idx).compile(dialect=context.dialect))
+        idx_stmt = re.sub(r'[\n\t]', '', idx_stmt)
+
         if ddl_contains:
-            assert ddl_contains in create_stmt
+            assert ddl_contains in create_stmt + idx_stmt
         if ddl_not_contains:
-            assert ddl_not_contains not in create_stmt
+            assert ddl_not_contains not in create_stmt + idx_stmt
 
-        context.assert_(
+        expected = [
             create_stmt,
+        ]
+        if impl.new_table.indexes:
+            expected.append(idx_stmt)
+
+        expected.extend([
             'INSERT INTO _alembic_batch_temp (%(colnames)s) '
             'SELECT %(tname_colnames)s FROM tname' % {
                 "colnames": ", ".join([
@@ -116,7 +137,8 @@ class BatchApplyTest(TestBase):
             },
             'DROP TABLE tname',
             'ALTER TABLE _alembic_batch_temp RENAME TO tname'
-        )
+        ])
+        context.assert_(*expected)
         return impl.new_table
 
     def test_change_type(self):
@@ -250,6 +272,24 @@ class BatchApplyTest(TestBase):
             impl, colnames=['id', 'x', 'y'],
             ddl_not_contains="CONSTRAINT uq1 UNIQUE")
 
+    def test_add_index(self):
+        impl = self._simple_fixture()
+        ix = self.op._index('ix1', 'tname', ['y'])
+
+        impl.add_index(ix)
+        self._assert_impl(
+            impl, colnames=['id', 'x', 'y'],
+            ddl_contains="CREATE INDEX ix1")
+
+    def test_drop_index(self):
+        impl = self._ix_fixture()
+
+        ix = self.op._index('ix1', 'tname', ['y'])
+        impl.drop_index(ix)
+        self._assert_impl(
+            impl, colnames=['id', 'x', 'y'],
+            ddl_not_contains="CONSTRAINT uq1 UNIQUE")
+
     def test_add_table_opts(self):
         impl = self._simple_fixture(table_kwargs={'mysql_engine': 'InnoDB'})
         self._assert_impl(
@@ -304,7 +344,8 @@ class BatchAPITest(TestBase):
                 mock.call(
                     ['x'], ['user.y'],
                     onupdate=None, ondelete=None, name='myfk',
-                    initially=None, deferrable=None, match=None)
+                    initially=None, deferrable=None, match=None,
+                    schema=None)
             ]
         )
         eq_(