How to link an event DDL with a migration?

[mberdyshev]

What is the right way to set an SQLAlchemy custom event DDL so that it would be added to the autogenerated migration file?
For example:

event.listen(
    issue_table,
    "after_create",
    DDL("ALTER TABLE v1_issue MODIFY COLUMN id INTEGER NOT NULL AUTO_INCREMENT").execute_if(("mysql", "mariadb")),
    once=True,
)

The query doesn't appear automatically with the alembic revision --autogenerate command.
My current solution is to manually copy this DDL to the migration. Is there a better, automatic way?

P. S. I found this SO question, but I could not make it work with the proposed answers.

[zzzeek]

There's multiple levels of "why?" to this question. so here are those:

1. the exact thing being asked for

as far as it being autogenerated directly into the migration file, not really, you'd need to write custom render functions to do that.

2. events directly in a migration file, part 1

I wouldnt use events inside a migration file, I'd have code that calls either the DDL() operation directly using op.execute(), or in this case just an op.alter_column() operation, just written out explicitly. for the mysql/mariadb thing I'd just have a conditional that looks at op.get_bind().dialect.name

3. events directly in a migration file, part 2

if you want a certain thing to happen for all create table operations, I'd use the event but it would be defined at the env.py level so that it modifies how create_table behaves, which is what events are for in general

4. what is this migration for?

the actual modify here looks like what SQLAlchemy does already for an integer primary key column under MySQL, and it's not stated here what you are changing it from. if this is for the datatype you could more easily use with_variant in your model (autogen should render that also).

overall I would need more context to advise on the best way to solve your issue

[mberdyshev]

I'm going to answer your points starting from the end:

4. what is this migration for?

It is an attempt to create an autoincrement non-primary integer field for MySQL within SQLAlchemy. I can find that this definition is not supported by SQLAlchemy as of now - sqlalchemy/sqlalchemy#9622. So what I attempt instead is to create the id field with the unique index and then execute manual DDL via events to add the autoincrement property to the field. At least it works as expected checking metadata.create_all(mock_engine, checkfirst=False).
That's the reasoning behind the custom DDL. And I was interested if it is possible to add the DDL to an alembic migration with the --autogenerate option.

I'd use the event but it would be defined at the env.py level so that it modifies how create_table behaves, which is what events are for in general

Probably I misunderstood the concept of events. But I looked at examples given in the docs:

from sqlalchemy import DDL
event.listen(
    some_table,
    "after_create",
    DDL("ALTER TABLE %(table)s SET name=foo_%(table)s")
)

event.listen(
    users,
    "after_create",
    DDL(
        "ALTER TABLE users ADD CONSTRAINT "
        "cst_user_name_length CHECK (length(user_name) >= 8)"
    ).execute_if(callable_=should_create),
)

I wouldnt use events inside a migration file, I'd have code that calls either the DDL() operation directly using op.execute()

I don't use events inside a migration file. The event is defined in another file tables.py which is then imported within the env.py file.
Currently, in the migration file, I have added the call sa.DDL("ALTER TABLE v1_issue MODIFY COLUMN id INTEGER NOT NULL AUTO_INCREMENT").execute_if(("mysql", "mariadb"))(target=None, bind=op.get_bind()). It was taken from this SO answer.

[zzzeek]

I don't use events inside a migration file. The event is defined in another file tables.py which is then imported within the env.py file.

What's the target for the event in env.py ? it would have to be Table here in order to pick up on an op.create_table() statement in a migration. Also this would not require any code to be added to a migration file if you are using an event that was already set up.

My current solution is to manually add this DDL to the migration.

It's not totally clear what you are doing here, can you illustrate a sample of how you "manually add this DDL to the migration" ?

[mberdyshev]

Ok, I'm providing the full example and the sequence of actions.

tables.py

from sqlalchemy import Column, DDL, event, Integer, MetaData, String, Table


metadata = MetaData()

test_table = Table(
    "test",
    metadata,
    Column("counter", Integer, primary_key=False, autoincrement=True, nullable=False, unique=True, index=True),
    Column("uuid", String(36), primary_key=True),
)

if __name__ == "__main__":
    # Taken from https://docs.sqlalchemy.org/en/20/faq/metadata_schema.html#how-can-i-get-the-create-table-drop-table-output-as-a-string
    from sqlalchemy import create_mock_engine

    def dump(sql, *multiparams, **params):
        print(sql.compile(dialect=mock_engine.dialect))

    mock_engine = create_mock_engine("mysql+pymysql://", dump)
    metadata.create_all(mock_engine, checkfirst=False)
# Outputs:
# CREATE TABLE test (
#         counter INTEGER NOT NULL,
#         uuid VARCHAR(36) NOT NULL,
#         PRIMARY KEY (uuid)
# )
# CREATE UNIQUE INDEX ix_test_counter ON test (counter)

In order to fulfill the mentioned purpose of an autoincrement non-primary integer field with the given restrictions I've added the event.listen call to the file based on this SO answer :

Updated tables.py

from sqlalchemy import Column, DDL, event, Integer, MetaData, String, Table


metadata = MetaData()

test_table = Table(
    "test",
    metadata,
    Column("counter", Integer, primary_key=False, autoincrement=True, nullable=False, unique=True, index=True),
    Column("uuid", String(36), primary_key=True),
)

event.listen(
    test_table,
    "after_create",
    DDL("ALTER TABLE test MODIFY COLUMN counter INT NOT NULL AUTO_INCREMENT").execute_if(("mysql", "mariadb")),
    once=True,
)

if __name__ == "__main__":
    # Taken from https://docs.sqlalchemy.org/en/20/faq/metadata_schema.html#how-can-i-get-the-create-table-drop-table-output-as-a-string
    from sqlalchemy import create_mock_engine

    def dump(sql, *multiparams, **params):
        print(sql.compile(dialect=mock_engine.dialect))

    mock_engine = create_mock_engine("mysql+pymysql://", dump)
    metadata.create_all(mock_engine, checkfirst=False)
# Outputs:
# CREATE TABLE test (
#         counter INTEGER NOT NULL,
#         uuid VARCHAR(36) NOT NULL,
#         PRIMARY KEY (uuid)
# )
# CREATE UNIQUE INDEX ix_test_counter ON test (counter)
# ALTER TABLE test MODIFY COLUMN counter INT NOT NULL AUTO_INCREMENT

Now I'd like to generate an alembic migration with --autogenerate option. I've provided MySQL database URL to alembic.ini and have added 2 lines to env.py:

from db.tables import *
target_metadata = metadata

Now alembic revision --autogenerate -m "Initial migration" generates this:

5adaac4b74cc_initial_migration.py

"""Initial migration

Revision ID: 5adaac4b74cc
Revises: 
Create Date: 2024-10-10 04:43:55.183718

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '5adaac4b74cc'
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('test',
    sa.Column('counter', sa.Integer(), autoincrement=True, nullable=False),
    sa.Column('uuid', sa.String(length=36), nullable=False),
    sa.PrimaryKeyConstraint('uuid')
    )
    op.create_index(op.f('ix_test_counter'), 'test', ['counter'], unique=True)
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f('ix_test_counter'), table_name='test')
    op.drop_table('test')
    # ### end Alembic commands ###

As you can see there is no ALTER command in this migration. And the command alembic upgrade head --sql confirms that the AUTOINCREMENT property is not being added:

CREATE TABLE test (
    counter INTEGER NOT NULL,
    uuid VARCHAR(36) NOT NULL,
    PRIMARY KEY (uuid)
);
CREATE UNIQUE INDEX ix_test_counter ON test (counter);
INSERT INTO alembic_version (version_num) VALUES ('5adaac4b74cc');

So, now based on the SO answer I've manually copied the DDL to the migration:

Updated 5adaac4b74cc_initial_migration.py

"""Initial migration

Revision ID: 5adaac4b74cc
Revises: 
Create Date: 2024-10-10 04:43:55.183718

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '5adaac4b74cc'
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('test',
    sa.Column('counter', sa.Integer(), autoincrement=True, nullable=False),
    sa.Column('uuid', sa.String(length=36), nullable=False),
    sa.PrimaryKeyConstraint('uuid')
    )
    op.create_index(op.f('ix_test_counter'), 'test', ['counter'], unique=True)
    sa.DDL("ALTER TABLE test MODIFY COLUMN counter INT NOT NULL AUTO_INCREMENT").execute_if(("mysql", "mariadb"))(
        target=None, bind=op.get_bind()
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f('ix_test_counter'), table_name='test')
    op.drop_table('test')
    # ### end Alembic commands ###

Now it works as I'd like it to. alembic upgrade head --sql expectedly outputs:

CREATE TABLE test (
    counter INTEGER NOT NULL,
    uuid VARCHAR(36) NOT NULL,
    PRIMARY KEY (uuid)
);
CREATE UNIQUE INDEX ix_test_counter ON test (counter);
ALTER TABLE test MODIFY COLUMN counter INT NOT NULL AUTO_INCREMENT;
INSERT INTO alembic_version (version_num) VALUES ('5adaac4b74cc');

The copied DDL can be changed to the op.execute("ALTER TABLE test MODIFY COLUMN counter INT NOT NULL AUTO_INCREMENT") as you mentioned before. But my question is the next - is it somehow possible to add the ALTER command from the event.listen call within alembic revision --autogenerate command? Without copying a query as I did. Or what is the right approach to manage this?

[zzzeek]

OK yes you have the event in your model code and the ask here is that autogenerate would somehow pick up on selected events and determine that ready-to-go objects like DDL() are the actual target, and not arbitrary code.

So no, autogenerate could never do this from event hooks like that, and you would have to as mentioned write those DDL() instructions directly in your migration files as coming right after an op.create_table() operation. You could in theory extend autogenerate yourself here but it would be very brittle code as event listeners dont have any public introspectability like that. they are imperative, not declarative.

to get autogenerate to do something you'd need a declarative indicator of some kind. which here I'd recommend as a adding a clue to the .info dictionary of the column like this:

test_table = Table(
    "test",
    metadata,
    Column("counter", Integer, primary_key=False, autoincrement=True, nullable=False, unique=True, index=True, info={"mysql_auto_increment": True}),
    Column("uuid", String(35), primary_key=True),
)

then we'd make use of Rewriter to intercept that and we can use the ExecuteSQLOp to write out an ALTER:

from alembic.autogenerate import rewriter
from alembic.operations import ops

writer = rewriter.Rewriter()

@writer.rewrites(ops.CreateTableOp)
def create_table(context, revision, op):
    output = [op]
    for col in op.columns:
        if col.info.get("mysql_auto_increment", False):
            output.append(
                ops.ExecuteSQLOp(
                    f"ALTER TABLE {op.table_name} MODIFY COLUMN "
                    f"{col.name} INTEGER NOT NULL AUTO_INCREMENT"
                )
            )
    return output

link that to the run like this:

def run_migrations_online() -> None:

   #...
    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            process_revision_directives=writer,
        )

        with context.begin_transaction():
            context.run_migrations()

so you could also use that same info in your event hook, and make a generalized event hook that you make against Table and does this same DDL for a plain metadata.create_all() call as well.

there's not currently a way to get the execute SQL op to be conditional for MySQL/MariaDB only from autogenerate though but if your app relies on autoincrement like that it would appear to be a mysql-specific application in any case.

[mberdyshev]

Thank you a lot for your reaction and informative answers!