- have the operations related to createtable still be within a

modify table block

4 files changed, 33 insertions, 1 deletions

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index d4084ad..00fe9de 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -136,6 +136,7 @@ def render_python_code(
     up_or_down_op,
     sqlalchemy_module_prefix='sa.',
     alembic_module_prefix='op.',
+    render_as_batch=False,
     imports=(),
     render_item=None,
 ):
@@ -151,6 +152,7 @@ def render_python_code(
             'sqlalchemy_module_prefix': sqlalchemy_module_prefix,
             'alembic_module_prefix': alembic_module_prefix,
             'render_item': render_item,
+            'render_as_batch': render_as_batch,
         },
         'imports': set(imports)
     }
diff --git a/alembic/autogenerate/compare.py b/alembic/autogenerate/compare.py
index 5bb2964..4001453 100644
--- a/alembic/autogenerate/compare.py
+++ b/alembic/autogenerate/compare.py
@@ -106,11 +106,14 @@ def _compare_tables(conn_table_names, metadata_table_names,
             upgrade_ops.ops.append(
                 ops.CreateTableOp.from_table(metadata_table))
             log.info("Detected added table %r", name)
+            modify_table_ops = ops.ModifyTableOps(tname, [], schema=s)
             _compare_indexes_and_uniques(s, tname, object_filters,
                                          None,
                                          metadata_table,
-                                         upgrade_ops,
+                                         modify_table_ops,
                                          autogen_context, inspector)
+            if not modify_table_ops.is_empty():
+                upgrade_ops.ops.append(modify_table_ops)
 
     removal_metadata = sa_schema.MetaData()
     for s, tname in conn_table_names.difference(metadata_table_names):
diff --git a/tests/test_autogen_composition.py b/tests/test_autogen_composition.py
index b1717ab..ff516f6 100644
--- a/tests/test_autogen_composition.py
+++ b/tests/test_autogen_composition.py
@@ -326,3 +326,4 @@ name='extra_uid_fkey'),
     )
     op.drop_table('item', schema='%(schema)s')
     ### end Alembic commands ###""" % {"schema": self.schema})
+
diff --git a/tests/test_autogen_render.py b/tests/test_autogen_render.py
index 4a49d5c..293c1c5 100644
--- a/tests/test_autogen_render.py
+++ b/tests/test_autogen_render.py
@@ -1457,3 +1457,29 @@ class RenderNamingConventionTest(TestBase):
             "sa.CheckConstraint(!U'im a constraint', name=op.f('ck_t_cc1'))"
         )
 
+    def test_create_table_plus_add_index_in_modify(self):
+        uo = ops.UpgradeOps(ops=[
+            ops.CreateTableOp(
+                "sometable",
+                [Column('x', Integer), Column('y', Integer)]
+            ),
+            ops.ModifyTableOps(
+                "sometable", ops=[
+                    ops.CreateIndexOp('ix1', 'sometable', ['x', 'y'])
+                ]
+            )
+        ])
+
+        eq_(
+            autogenerate.render_python_code(uo, render_as_batch=True),
+            "### commands auto generated by Alembic - please adjust! ###\n"
+            "    op.create_table('sometable',\n"
+            "    sa.Column('x', sa.Integer(), nullable=True),\n"
+            "    sa.Column('y', sa.Integer(), nullable=True)\n"
+            "    )\n"
+            "    with op.batch_alter_table('sometable', schema=None) "
+            "as batch_op:\n"
+            "        batch_op.create_index("
+            "'ix1', ['x', 'y'], unique=False)\n\n"
+            "    ### end Alembic commands ###"
+        )