How to use imports in script.py.mako

[luolingchun]

I checked the document and issue and found two useful messages:

• Does not add imports of external modules like geoalchemy2 in the migrations file #703 (comment)
• https://alembic.sqlalchemy.org/en/latest/autogenerate.html#comparing-and-rendering-types

How to use the imports in script.py.mako to dynamically add custom import instead of manually editing script.py.mako.

This is just a proposal that will help automate the migration.

"""${message}

Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}

"""
from alembic import op
import sqlalchemy as sa
${imports if imports else ""}

# revision identifiers, used by Alembic.
revision = ${repr(up_revision)}
down_revision = ${repr(down_revision)}
branch_labels = ${repr(branch_labels)}
depends_on = ${repr(depends_on)}


def upgrade() -> None:
    ${upgrades if upgrades else "pass"}


def downgrade() -> None:
    ${downgrades if downgrades else "pass"}

[zzzeek]

well the documented hook for this is in the context of using custom types, where the render_item hook grants access to the context that provides this list of imports.

there's a another way to affect imports directly by using a rewriter around ops.MigrationScript, and adding elements to the 'imports' collection:

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

writer = rewriter.Rewriter()

@writer.rewrites(ops.MigrationScript)
def add_imports(context, revision, op):
    op.imports.add("from myprogram import something_important")
    return [op]

you add this to the run in env.py like this:

def run_migrations_online() -> None:
    connectable = ...

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

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

i just tested to confirm this is working.

[luolingchun]

@zzzeek Thank you for your answer!

I'm looking for a simpler way, I found that just add imports to the revision command.

However, this requires upgrading the source code. I don't know if it will have other bad effects.

It will be easier to use if it works.

https://github.com/sqlalchemy/alembic/blob/rel_1_9_2/alembic/command.py#L116

[zzzeek]

if you want to fork the source that's fine. I think using the published APIs is ultimately easier but either way.

[luolingchun]

Can I submit PR as a new feature? I hope this feature can be used in future versions.

[zzzeek]

we already support customizing the imports for a script in a way that is consistent with how everything else about a revision script is optimized. it is inconsistent to carve out a special block for "imports". if you have imports you want to add to all revision scripts, just put them into the script.py.mako. There's not any rationale for "imports" to be part of the revision command on the command line.

[zzzeek]

there is also template_args:

https://alembic.sqlalchemy.org/en/latest/api/runtime.html#alembic.runtime.environment.EnvironmentContext.configure.params.template_args

with this, you can add any arbitrary variables to templates from the env.py script. you can pass custom things from the command line using the -x option. you can then reference them in your template.

so that's three ways to do this, I think that's a lot