TypeError when running database upgrade with generated migration from updated table comment

[FranzForstmayr]

Describe the bug
The generated schema_upgrades() function raises a TypeError when the model's comment is updated.

To Reproduce
I'm using this project to start:
https://github.com/litestar-org/litestar-fullstack

When I add a simple model like this:

class MyNewModel(TimestampedDatabaseModel):
    __table_args__ = {"comment": "Comment Version 1"}

    something: Mapped[str]

I get the following snippet in my migration script, which is fine:

def schema_upgrades() -> None:
    """schema upgrade migrations go here."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('my_new_model',
    sa.Column('id', sa.GUID(length=16), nullable=False),
    sa.Column('something', sa.String(), nullable=False),
    sa.Column('sa_orm_sentinel', sa.Integer(), nullable=True),
    sa.Column('created_at', sa.DateTimeUTC(timezone=True), nullable=False),
    sa.Column('updated_at', sa.DateTimeUTC(timezone=True), nullable=False),
    sa.PrimaryKeyConstraint('id', name=op.f('pk_my_new_model'))
    )
    # ### end Alembic commands ###

However when I change the comment afterwards:

class MyNewModel(TimestampedDatabaseModel):
    __table_args__ = {"comment": "Comment Version 2"}

    something: Mapped[str]

I get this schema_upgrade:

def schema_upgrades() -> None:
    """schema upgrade migrations go here."""
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('my_new_model', schema=None) as batch_op:
        batch_op.create_table_comment(
        'my_new_model',
        'Comment Version 2',
        existing_comment=None,
        schema=None
    )

    # ### end Alembic commands ###

which fails, as the signature of drop_table_comment does not allow schema parameter and table name.

        def drop_table_comment(
            self, *, existing_comment: Optional[str] = None
        ) -> None:

Error

Traceback (most recent call last):
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/bin/app", line 8, in <module>
    sys.exit(run_cli())
             ^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/src/app/__main__.py", line 23, in run_cli
    run_litestar_cli()
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/litestar/__main__.py", line 6, in run_cli
    litestar_group()
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/rich_click/rich_command.py", line 126, in main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/click/decorators.py", line 33, in new_func
    return f(get_current_context(), *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/litestar/cli/_utils.py", line 280, in wrapped
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/click/decorators.py", line 33, in new_func
    return f(get_current_context(), *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/litestar/cli/_utils.py", line 280, in wrapped
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/advanced_alchemy/extensions/litestar/cli.py", line 109, in upgrade_database
    alembic_commands.upgrade(revision=revision, sql=sql, tag=tag)
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/advanced_alchemy/alembic/commands.py", line 92, in upgrade
    return migration_command.upgrade(config=self.config, revision=revision, tag=tag, sql=sql)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/alembic/command.py", line 398, in upgrade
    script.run_env()
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/alembic/script/base.py", line 579, in run_env
    util.load_python_file(self.dir, "env.py")
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/alembic/util/pyfiles.py", line 93, in load_python_file
    module = load_module_py(module_id, path)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/alembic/util/pyfiles.py", line 109, in load_module_py
    spec.loader.exec_module(module)  # type: ignore
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib._bootstrap_external>", line 940, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/home/franz/Workspaces/python/litestar-fullstack/src/app/lib/db/migrations/env.py", line 142, in <module>
    asyncio.run(run_migrations_online())
  File "/usr/lib/python3.11/asyncio/runners.py", line 190, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/asyncio/base_events.py", line 653, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/src/app/lib/db/migrations/env.py", line 134, in run_migrations_online
    await connection.run_sync(do_run_migrations)
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/sqlalchemy/ext/asyncio/engine.py", line 886, in run_sync
    return await greenlet_spawn(fn, self._proxied, *arg, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/sqlalchemy/util/_concurrency_py3k.py", line 192, in greenlet_spawn
    result = context.switch(value)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/home/franz/Workspaces/python/litestar-fullstack/src/app/lib/db/migrations/env.py", line 105, in do_run_migrations
    context.run_migrations()
  File "<string>", line 8, in run_migrations
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/alembic/runtime/environment.py", line 938, in run_migrations
    self.get_context().run_migrations(**kw)
  File "/home/franz/Workspaces/python/litestar-fullstack/.venv/lib/python3.11/site-packages/alembic/runtime/migration.py", line 624, in run_migrations
    step.migration_fn(**kw)
  File "/home/franz/Workspaces/python/litestar-fullstack/src/app/lib/db/migrations/versions/2023-11-17_test_734075cb77dd.py", line 39, in upgrade
    schema_upgrades()
  File "/home/franz/Workspaces/python/litestar-fullstack/src/app/lib/db/migrations/versions/2023-11-17_test_734075cb77dd.py", line 53, in schema_upgrades
    batch_op.create_table_comment(
TypeError: create_table_comment() got an unexpected keyword argument 'schema'

Versions.

• OS: Ubuntu 20.04.6 LTS (WSL2)
• Python: 3.11.4
• Alembic: 1.12.1
• SQLAlchemy: 2.0.23
• Database: postgres 16 (docker)
• DBAPI: asyncpg 0.29.0

Thanks!

[sqla-tester]

Mike Bayer has proposed a fix for this issue in the main branch:

add batch template for render create_table_comment https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/4951

[zzzeek]

thanks for reporting!

[FranzForstmayr]

Awesome reaction time :)

[FranzForstmayr]

drop_table_comment has the same issue I just found out.

[zzzeek]

yah I should have checked on that, will add

[FranzForstmayr]

Thanks