Alembic: Extend autogenerate outside upgrade() and downgrade()

[olirice]

Following along with ReplaceableObject cookbook

https://alembic.sqlalchemy.org/en/latest/cookbook.html#create-initial-migrations

we end up with a migration like:

from alembic import op
import sqlalchemy as sa
from foo import ReplaceableObject

some_sp = ReplaceableObject(
    "to_upper(some_text text)",
    " RETURNS text AS $$ SELECT upper(text) $$ LANGUAGE sql;")

def upgrade():
    op.create_sp(some_sp)


def downgrade():
    op.drop_sp(some_sp)

and can refer back to that for a replace operation in a later migration with

op.replace_sp(some_updated_sp, replaces="28af9800143f.some_sp")

In order for some_sp to be importable for use in the replaces parameter, it must be defined outside of the upgrade and downgrade functions.

The autogenerate api

@renderers.dispatch_for(ReplaceFunctionOp)

can be used to render ReplaceableObjects within an upgrade or downgrade function but I haven't found the appropriate place to hook to render ReplaceableObject in the migration's header.

In other words, can easily autogenerate

def upgrade():
    some_sp = ReplaceableObject(
        "to_upper(some_text text)",
        " RETURNS text AS $$ SELECT upper(text) $$ LANGUAGE sql;")

    op.create_sp(some_sp)
...

but it isn't clear how to generate

some_sp = ReplaceableObject(
        "to_upper(some_text text)",
        " RETURNS text AS $$ SELECT upper(text) $$ LANGUAGE sql;")

def upgrade():    
    op.create_sp(some_sp)
...

Is that the expected approach? If not, I'd really appreciate a nudge in the right direction

[zzzeek]

OK, making an autogeneration path for stored procedures is pretty ambitious, so congrats on getting this far, I'm not sure how popular "replaceable object" is because for most people it is quite esoteric, however I think the basic idea can be used to build up a fully functioning stored procedure management system.

Within the scope of the question, how to render module-level things, what we have when we deal with @renderers is a function that renders Python code that corresponds to the UpgradeOps and DowngradeOps structures, where we see UpgradeOps here: https://alembic.sqlalchemy.org/en/latest/api/operations.html?highlight=upgradeops#alembic.operations.ops.UpgradeOps

these two structures are stuck onto a MigrationScript object which represents the whole "script" that would be generated.

The MigrationScript currently has fields for "upgrade" and "downgrade" ops, as well as "imports". These three sections as well as a few others match up to tokens that are in the pre-fabricated script.py.mako template.

for this API to suit what you want to do directly, MigrationScript would need a new section called "ModulePython" or something like that. I'm not opposed to trying to make this a first class feature of MigrationScript. However, we don't have that right now, so the only way to get new Python code in there is first to add a new template variable to your script.py.mako template where you would want "module level' code to go, that is, right above the upgrade_ops(), then get the content that you need to be sent into that template variable at script production time.

The next thing I wanted to say was that you could populate these template args from within your render functions. Unfortunately, that is not connected right now; it certainly could be connected as the autogen_context and the template_args are both right there, but in the current released version there's not a direct path to the template_args dictionary, unless you were to use stackframe inspection to walk up and find it :(. If you want to contribute to Alembic to help make this connection, we can go that route.

There is a way that this effect can be hacked around for now, which is to put a mutable structure into the template_args at configuration time, then access that structure from within the autogen context, as the config level template_args are passed around everywhere, it's just that they are copied into a new dictionary used at render time, but if we put a mutable object in there, we can change it.

So step 1, put this in your env.py in the context.configure() call:

def run_migrations_online():
    # ...
        context.configure(
            connection=connection, target_metadata=target_metadata,
            template_args={"mutable_structure": {}}
        )
   # ...

step 2, inside your render function, put Python code as desired into this mutable structure like this:

autogen_context.opts['template_args']['mutable_structure']['helloworld'] = 'print("hi")'

step 3, refer to the above Python code in your script.py.mako as ${mutable_structure['helloworld']}:

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


${mutable_structure['helloworld']}

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


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

I added that autogen context line into my render_create_table_op and got:

# revision identifiers, used by Alembic.
revision = "bc7e28ed1604"
down_revision = None
branch_labels = None
depends_on = None


print("hi")


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table("t1", sa.Column("foo", sa.Integer(), nullable=True))
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("t1")
    # ### end Alembic commands ###

that was long. what do you think?

[olirice]

Thanks for such a thought out answer!

Summary of options:

• Extend MigrationScript with a ModuleOps property (like UpgradeOps)
• Make template args available in render functions
• Work around template args not being available in render functions with shared refs to mutable containers

The plan is for replaceable object stuff to get packaged up as an extension for use in multiple projects. With that in mind, ease of setup for end users is a priority. The last two options require a lot of project-by-project config in env.py and mako template which makes me lean towards option 1 (although who wouldn't prefer first class support :) ).

To get this off the ground, I've added a RevertOp which retrieves the existing definition of a function from a live database (in pg_proc) and uses that in the downgrade statement instead of importing a definition from a previous migration.

e.g.

def upgrade() -> None:
    # this definition was local
    public_to_uppercase_f4ed96 = PGFunction(
        schema="public",
        signature="to_uppercase(some_text text)",
        definition="""returns text as
    $$
        SELECT upper(some_text) || 'xyz';
    $$ LANGUAGE SQL;
    """,
    )
    op.replace_function(public_to_uppercase_f4ed96)


def downgrade() -> None:
    # this definition came from postgres at the time the revision was generated
    public_to_uppercase_f4ed96 = PGFunction(
        schema="public",
        signature="to_uppercase(some_text text)",
        definition="""returns text
     LANGUAGE sql
    AS $function$
            SELECT upper(some_text) || 'abc';
        $function$""",
    )
    op.replace_function(public_to_uppercase_f4ed96)

Which skirts the issue. Since --autogenerate will error if the db is not up to date, that should be fine, right? (please correct me if I'm wrong about that)

If you have any thoughts on how that could be better please let me know!

After a few weeks, if it ends up feeling gross to work with or introduces unexpected issues, I'll take a swing at implementing ModuleOps.

Thanks again for your answer above! It really helped me understand what was going on

[zzzeek]

I wrote "Replaceable object" a long time ago. At the moment I'm not sure what the full rationale is that it has to look into historical migrations in order to know what it's going to "drop". If this is just named stored procedures, you don't need the old definition, just the name so you can emit "DROP". Which means you can just render the stored proceduires inside of upgrade and downgrade as you are doing. after all the "replaceable object" recipe states right at the top that it's "theoretical". am I missing something ?

[zzzeek]

oh, for downgrades, that's right. Downgrades are not even a thing in most of the real world.

[zzzeek]

but yes, as you are rendering the "old" name entirely in the "drop", I think if autogenerate is in play here then you can simplify "replaceable object" to not need the part where it loads up a former migration. this is not any different from how create_table / drop_table renders in an autogenerate script, tables that are dropped in "upgrade" are rendered fully in "downgrade".

I think a key aspect of ReplaceableObject was that I wanted to remove the need to write out the full SP definition multiple times. But autogenerate changes that calculus entirely and perhaps this should be noted in the recipe's introduction.

[olirice]

In the unlikely event that someone winds up here and want to see how adding autogenerate support for functions and views turned out:

Its now available in alembic_utils via pip install alembic_utils

Source:
https://github.com/olirice/alembic_utils

Docs:
https://olirice.github.io/alembic_utils/

[zzzeek]

I will gladly add a link to your project in this section of the docs, do you have a pypi release yet?

[olirice]

That'd be awesome, thanks!!

Yes, its live on pypi AO this morning

Please link to github rather than the documentation b/c I'll probably switch to readthedocs if the project grows

[zzzeek]

that's pushed, should build on the site in a little while.

[olirice]

very cool, thanks!